pdf_parse_by_ocr.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. import json
  2. from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio
  3. from magic_pdf.libs.ocr_dict_merge import merge_spans
  4. def read_json_file(file_path):
  5. with open(file_path, 'r') as f:
  6. data = json.load(f)
  7. return data
  8. def construct_page_component(page_id, text_blocks_preproc):
  9. return_dict = {
  10. 'preproc_blocks': text_blocks_preproc,
  11. 'page_idx': page_id
  12. }
  13. return return_dict
  14. def parse_pdf_by_ocr(
  15. ocr_json_file_path,
  16. start_page_id=0,
  17. end_page_id=None,
  18. ):
  19. ocr_pdf_info = read_json_file(ocr_json_file_path)
  20. pdf_info_dict = {}
  21. end_page_id = end_page_id if end_page_id else len(ocr_pdf_info) - 1
  22. for page_id in range(start_page_id, end_page_id + 1):
  23. ocr_page_info = ocr_pdf_info[page_id]
  24. layout_dets = ocr_page_info['layout_dets']
  25. spans = []
  26. for layout_det in layout_dets:
  27. category_id = layout_det['category_id']
  28. allow_category_id_list = [13, 14, 15]
  29. if category_id in allow_category_id_list:
  30. x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
  31. bbox = [int(x0), int(y0), int(x1), int(y1)]
  32. # 13: 'embedding', # 嵌入公式
  33. # 14: 'isolated', # 单行公式
  34. # 15: 'ocr_text', # ocr识别文本
  35. span = {
  36. 'bbox': bbox,
  37. }
  38. if category_id == 13:
  39. span['content'] = layout_det['latex']
  40. span['type'] = 'inline_equation'
  41. elif category_id == 14:
  42. span['content'] = layout_det['latex']
  43. span['type'] = 'displayed_equation'
  44. elif category_id == 15:
  45. span['content'] = layout_det['text']
  46. span['type'] = 'text'
  47. # print(span)
  48. spans.append(span)
  49. else:
  50. continue
  51. # 合并重叠的spans
  52. for span1 in spans.copy():
  53. for span2 in spans.copy():
  54. if span1 != span2:
  55. overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.8)
  56. if overlap_box is not None:
  57. bbox_to_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
  58. if bbox_to_remove is not None:
  59. spans.remove(bbox_to_remove)
  60. # 将spans合并成line
  61. lines = merge_spans(spans)
  62. # 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
  63. blocks = []
  64. for line in lines:
  65. blocks.append({
  66. "bbox": line['bbox'],
  67. "lines": [line],
  68. })
  69. # 构造pdf_info_dict
  70. page_info = construct_page_component(page_id, blocks)
  71. pdf_info_dict[f"page_{page_id}"] = page_info
  72. return pdf_info_dict