pdf_parse_by_ocr.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. import time
  2. from loguru import logger
  3. from magic_pdf.libs.commons import (
  4. fitz,
  5. get_delta_time,
  6. get_docx_model_output,
  7. )
  8. from magic_pdf.libs.convert_utils import dict_to_list
  9. from magic_pdf.libs.coordinate_transform import get_scale_ratio
  10. from magic_pdf.libs.drop_tag import DropTag
  11. from magic_pdf.libs.hash_utils import compute_md5
  12. from magic_pdf.libs.ocr_content_type import ContentType
  13. from magic_pdf.para.para_split import para_split
  14. from magic_pdf.para.para_split_by_model import para_split_by_model
  15. from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component
  16. from magic_pdf.pre_proc.detect_footer_by_model import parse_footers
  17. from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model
  18. from magic_pdf.pre_proc.detect_header import parse_headers
  19. from magic_pdf.pre_proc.detect_page_number import parse_pageNos
  20. from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table
  21. from magic_pdf.pre_proc.ocr_detect_layout import layout_detect
  22. from magic_pdf.pre_proc.ocr_dict_merge import (
  23. merge_spans_to_line_by_layout, merge_lines_to_block,
  24. )
  25. from magic_pdf.pre_proc.ocr_span_list_modify import remove_spans_by_bboxes, remove_overlaps_min_spans, \
  26. adjust_bbox_for_standalone_block, modify_y_axis, modify_inline_equation, get_qa_need_list, \
  27. remove_spans_by_bboxes_dict
  28. from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
  29. def parse_pdf_by_ocr(
  30. pdf_bytes,
  31. pdf_model_output,
  32. imageWriter,
  33. start_page_id=0,
  34. end_page_id=None,
  35. debug_mode=False,
  36. ):
  37. pdf_bytes_md5 = compute_md5(pdf_bytes)
  38. pdf_docs = fitz.open("pdf", pdf_bytes)
  39. # 初始化空的pdf_info_dict
  40. pdf_info_dict = {}
  41. start_time = time.time()
  42. end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
  43. for page_id in range(start_page_id, end_page_id + 1):
  44. # 获取当前页的page对象
  45. page = pdf_docs[page_id]
  46. # 获取当前页的宽高
  47. page_w = page.rect.width
  48. page_h = page.rect.height
  49. if debug_mode:
  50. time_now = time.time()
  51. logger.info(
  52. f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
  53. )
  54. start_time = time_now
  55. # 获取当前页的模型数据
  56. ocr_page_info = get_docx_model_output(
  57. pdf_model_output, page_id
  58. )
  59. """从json中获取每页的页码、页眉、页脚的bbox"""
  60. page_no_bboxes = parse_pageNos(page_id, page, ocr_page_info)
  61. header_bboxes = parse_headers(page_id, page, ocr_page_info)
  62. footer_bboxes = parse_footers(page_id, page, ocr_page_info)
  63. footnote_bboxes = parse_footnotes_by_model(page_id, page, ocr_page_info, debug_mode=debug_mode)
  64. # 构建需要remove的bbox字典
  65. need_remove_spans_bboxes_dict = {
  66. DropTag.PAGE_NUMBER: page_no_bboxes,
  67. DropTag.HEADER: header_bboxes,
  68. DropTag.FOOTER: footer_bboxes,
  69. DropTag.FOOTNOTE: footnote_bboxes,
  70. }
  71. layout_dets = ocr_page_info["layout_dets"]
  72. spans = []
  73. # 计算模型坐标和pymu坐标的缩放比例
  74. horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
  75. ocr_page_info, page
  76. )
  77. for layout_det in layout_dets:
  78. category_id = layout_det["category_id"]
  79. allow_category_id_list = [1, 7, 13, 14, 15]
  80. if category_id in allow_category_id_list:
  81. x0, y0, _, _, x1, y1, _, _ = layout_det["poly"]
  82. bbox = [
  83. int(x0 / horizontal_scale_ratio),
  84. int(y0 / vertical_scale_ratio),
  85. int(x1 / horizontal_scale_ratio),
  86. int(y1 / vertical_scale_ratio),
  87. ]
  88. # 删除高度或者宽度为0的spans
  89. if bbox[2] - bbox[0] == 0 or bbox[3] - bbox[1] == 0:
  90. continue
  91. """要删除的"""
  92. # 3: 'header', # 页眉
  93. # 4: 'page number', # 页码
  94. # 5: 'footnote', # 脚注
  95. # 6: 'footer', # 页脚
  96. """当成span拼接的"""
  97. # 1: 'image', # 图片
  98. # 7: 'table', # 表格
  99. # 13: 'inline_equation', # 行内公式
  100. # 14: 'interline_equation', # 行间公式
  101. # 15: 'text', # ocr识别文本
  102. """layout信息"""
  103. # 11: 'full column', # 单栏
  104. # 12: 'sub column', # 多栏
  105. span = {
  106. "bbox": bbox,
  107. }
  108. if category_id == 1:
  109. span["type"] = ContentType.Image
  110. elif category_id == 7:
  111. span["type"] = ContentType.Table
  112. elif category_id == 13:
  113. span["content"] = layout_det["latex"]
  114. span["type"] = ContentType.InlineEquation
  115. elif category_id == 14:
  116. span["content"] = layout_det["latex"]
  117. span["type"] = ContentType.InterlineEquation
  118. elif category_id == 15:
  119. span["content"] = layout_det["text"]
  120. span["type"] = ContentType.Text
  121. # print(span)
  122. spans.append(span)
  123. else:
  124. continue
  125. '''删除重叠spans中较小的那些'''
  126. spans, dropped_spans_by_span_overlap = remove_overlaps_min_spans(spans)
  127. '''
  128. 删除remove_span_block_bboxes中的bbox
  129. 并增加drop相关数据
  130. '''
  131. spans, dropped_spans_by_removed_bboxes = remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict)
  132. '''对image和table截图'''
  133. spans = ocr_cut_image_and_table(spans, page, page_id, pdf_bytes_md5, imageWriter)
  134. '''行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)'''
  135. displayed_list = []
  136. text_inline_lines = []
  137. modify_y_axis(spans, displayed_list, text_inline_lines)
  138. '''模型识别错误的行间公式, type类型转换成行内公式'''
  139. spans = modify_inline_equation(spans, displayed_list, text_inline_lines)
  140. '''bbox去除粘连'''
  141. spans = remove_overlap_between_bbox(spans)
  142. '''
  143. 对tpye=["interline_equation", "image", "table"]进行额外处理,
  144. 如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
  145. '''
  146. spans = adjust_bbox_for_standalone_block(spans)
  147. '''从ocr_page_info中解析layout信息(按自然阅读方向排序,并修复重叠和交错的bad case)'''
  148. layout_bboxes, layout_tree = layout_detect(ocr_page_info['subfield_dets'], page, ocr_page_info)
  149. '''将spans合并成line(在layout内,从上到下,从左到右)'''
  150. lines, dropped_spans_by_layout = merge_spans_to_line_by_layout(spans, layout_bboxes)
  151. '''将lines合并成block'''
  152. blocks = merge_lines_to_block(lines)
  153. '''获取QA需要外置的list'''
  154. images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
  155. '''drop的span_list合并'''
  156. dropped_spans = []
  157. dropped_spans.extend(dropped_spans_by_span_overlap)
  158. dropped_spans.extend(dropped_spans_by_removed_bboxes)
  159. dropped_spans.extend(dropped_spans_by_layout)
  160. dropped_text_block = []
  161. dropped_image_block = []
  162. dropped_table_block = []
  163. dropped_equation_block = []
  164. for span in dropped_spans:
  165. # drop出的spans进行分类
  166. if span['type'] == ContentType.Text:
  167. dropped_text_block.append(span)
  168. elif span['type'] == ContentType.Image:
  169. dropped_image_block.append(span)
  170. elif span['type'] == ContentType.Table:
  171. dropped_table_block.append(span)
  172. elif span['type'] in [ContentType.InlineEquation, ContentType.InterlineEquation]:
  173. dropped_equation_block.append(span)
  174. '''构造pdf_info_dict'''
  175. page_info = ocr_construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
  176. images, tables, interline_equations, inline_equations,
  177. dropped_text_block, dropped_image_block, dropped_table_block,
  178. dropped_equation_block,
  179. need_remove_spans_bboxes_dict)
  180. pdf_info_dict[f"page_{page_id}"] = page_info
  181. """分段"""
  182. if debug_mode:
  183. para_split_by_model(pdf_info_dict, debug_mode=debug_mode)
  184. else:
  185. para_split(pdf_info_dict, debug_mode=debug_mode)
  186. """dict转list"""
  187. pdf_info_list = dict_to_list(pdf_info_dict)
  188. new_pdf_info_dict = {
  189. "pdf_info": pdf_info_list,
  190. }
  191. return new_pdf_info_dict