draw_bbox.py 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. from pathlib import Path
  2. from magic_pdf.libs.commons import fitz, join_path # PyMuPDF
  3. from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
  4. import json
  5. import os
  6. def read_json_file(file_path):
  7. with open(file_path, 'r') as f:
  8. data = json.load(f)
  9. return data
  10. # PDF文件路径
  11. pdf_path = "D:\\projects\\Magic-PDF\\ocr_demo\\ocr_0_org.pdf"
  12. doc = fitz.open(pdf_path) # Open the PDF
  13. # 你的数据
  14. data = [[[-2, 0, 603, 80, 24]], [[-3, 0, 602, 80, 24]]]
  15. ocr_json_file_path = r"D:\projects\Magic-PDF\ocr_demo\ocr_0.json"
  16. ocr_pdf_info = read_json_file(ocr_json_file_path)
  17. pth = Path(ocr_json_file_path)
  18. book_name = pth.name
  19. save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
  20. save_path = join_path(save_tmp_path, "md")
  21. pdf_info_dict = parse_pdf_by_ocr(
  22. pdf_path,
  23. None,
  24. ocr_pdf_info,
  25. save_path,
  26. book_name,
  27. debug_mode=True)
  28. data_list = []
  29. for page in pdf_info_dict.values():
  30. page_list = []
  31. blocks = page.get("preproc_blocks")
  32. for block in blocks:
  33. lines = block.get("lines")
  34. for line in lines:
  35. spans = line.get("spans")
  36. for span in spans:
  37. page_list.append(span["bbox"])
  38. data_list.append(page_list)
  39. # 对每个页面进行处理
  40. for i, page in enumerate(doc):
  41. # 获取当前页面的数据
  42. page_data = data_list[i]
  43. for img in page_data:
  44. x0, y0, x1, y1 = img
  45. rect_coords = fitz.Rect(x0, y0, x1, y1) # Define the rectangle
  46. page.draw_rect(rect_coords, color=(1, 0, 0), fill=None, width=1.5, overlay=True) # Draw the rectangle
  47. # Save the PDF
  48. doc.save("D:\\projects\\Magic-PDF\\ocr_demo\\ocr_0_new1.pdf")