pdf_parse_by_ocr.py 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251
  1. import json
  2. import os
  3. import time
  4. from loguru import logger
  5. from magic_pdf.libs.draw_bbox import draw_layout_bbox, draw_text_bbox
  6. from magic_pdf.libs.commons import (
  7. read_file,
  8. join_path,
  9. fitz,
  10. get_img_s3_client,
  11. get_delta_time,
  12. get_docx_model_output,
  13. )
  14. from magic_pdf.libs.coordinate_transform import get_scale_ratio
  15. from magic_pdf.libs.safe_filename import sanitize_filename
  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.ocr_cut_image import 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 construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
  30. images, tables, interline_equations, inline_equations,
  31. dropped_text_block, dropped_image_block, dropped_table_block):
  32. return_dict = {
  33. 'preproc_blocks': blocks,
  34. 'layout_bboxes': layout_bboxes,
  35. 'page_idx': page_id,
  36. 'page_size': [page_w, page_h],
  37. '_layout_tree': layout_tree,
  38. 'images': images,
  39. 'tables': tables,
  40. 'interline_equations': interline_equations,
  41. 'inline_equations': inline_equations,
  42. 'dropped_text_block': dropped_text_block,
  43. 'dropped_image_block': dropped_image_block,
  44. 'dropped_table_block': dropped_table_block,
  45. }
  46. return return_dict
  47. def parse_pdf_by_ocr(
  48. pdf_path,
  49. s3_pdf_profile,
  50. pdf_model_output,
  51. save_path,
  52. book_name,
  53. pdf_model_profile=None,
  54. image_s3_config=None,
  55. start_page_id=0,
  56. end_page_id=None,
  57. debug_mode=False,
  58. ):
  59. pdf_bytes = read_file(pdf_path, s3_pdf_profile)
  60. save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
  61. book_name = sanitize_filename(book_name)
  62. md_bookname_save_path = ""
  63. if debug_mode:
  64. save_path = join_path(save_tmp_path, "md")
  65. pdf_local_path = join_path(save_tmp_path, "download-pdfs", book_name)
  66. if not os.path.exists(os.path.dirname(pdf_local_path)):
  67. # 如果目录不存在,创建它
  68. os.makedirs(os.path.dirname(pdf_local_path))
  69. md_bookname_save_path = join_path(save_tmp_path, "md", book_name)
  70. if not os.path.exists(md_bookname_save_path):
  71. # 如果目录不存在,创建它
  72. os.makedirs(md_bookname_save_path)
  73. with open(pdf_local_path + ".pdf", "wb") as pdf_file:
  74. pdf_file.write(pdf_bytes)
  75. pdf_docs = fitz.open("pdf", pdf_bytes)
  76. # 初始化空的pdf_info_dict
  77. pdf_info_dict = {}
  78. img_s3_client = get_img_s3_client(save_path, image_s3_config)
  79. start_time = time.time()
  80. end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
  81. for page_id in range(start_page_id, end_page_id + 1):
  82. # 获取当前页的page对象
  83. page = pdf_docs[page_id]
  84. # 获取当前页的宽高
  85. page_w = page.rect.width
  86. page_h = page.rect.height
  87. if debug_mode:
  88. time_now = time.time()
  89. logger.info(
  90. f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
  91. )
  92. start_time = time_now
  93. # 获取当前页的模型数据
  94. ocr_page_info = get_docx_model_output(
  95. pdf_model_output, pdf_model_profile, page_id
  96. )
  97. """从json中获取每页的页码、页眉、页脚的bbox"""
  98. page_no_bboxes = parse_pageNos(page_id, page, ocr_page_info)
  99. header_bboxes = parse_headers(page_id, page, ocr_page_info)
  100. footer_bboxes = parse_footers(page_id, page, ocr_page_info)
  101. footnote_bboxes = parse_footnotes_by_model(
  102. page_id, page, ocr_page_info, md_bookname_save_path, debug_mode=debug_mode
  103. )
  104. # 构建需要remove的bbox列表
  105. # need_remove_spans_bboxes = []
  106. # need_remove_spans_bboxes.extend(page_no_bboxes)
  107. # need_remove_spans_bboxes.extend(header_bboxes)
  108. # need_remove_spans_bboxes.extend(footer_bboxes)
  109. # need_remove_spans_bboxes.extend(footnote_bboxes)
  110. # 构建需要remove的bbox字典
  111. need_remove_spans_bboxes_dict = {
  112. "page_no": page_no_bboxes,
  113. "header": header_bboxes,
  114. "footer": footer_bboxes,
  115. "footnote": footnote_bboxes,
  116. }
  117. layout_dets = ocr_page_info["layout_dets"]
  118. spans = []
  119. # 计算模型坐标和pymu坐标的缩放比例
  120. horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
  121. ocr_page_info, page
  122. )
  123. for layout_det in layout_dets:
  124. category_id = layout_det["category_id"]
  125. allow_category_id_list = [1, 7, 13, 14, 15]
  126. if category_id in allow_category_id_list:
  127. x0, y0, _, _, x1, y1, _, _ = layout_det["poly"]
  128. bbox = [
  129. int(x0 / horizontal_scale_ratio),
  130. int(y0 / vertical_scale_ratio),
  131. int(x1 / horizontal_scale_ratio),
  132. int(y1 / vertical_scale_ratio),
  133. ]
  134. """要删除的"""
  135. # 3: 'header', # 页眉
  136. # 4: 'page number', # 页码
  137. # 5: 'footnote', # 脚注
  138. # 6: 'footer', # 页脚
  139. """当成span拼接的"""
  140. # 1: 'image', # 图片
  141. # 7: 'table', # 表格
  142. # 13: 'inline_equation', # 行内公式
  143. # 14: 'displayed_equation', # 行间公式
  144. # 15: 'text', # ocr识别文本
  145. """layout信息"""
  146. # 11: 'full column', # 单栏
  147. # 12: 'sub column', # 多栏
  148. span = {
  149. "bbox": bbox,
  150. }
  151. if category_id == 1:
  152. span["type"] = "image"
  153. elif category_id == 7:
  154. span["type"] = "table"
  155. elif category_id == 13:
  156. span["content"] = layout_det["latex"]
  157. span["type"] = "inline_equation"
  158. elif category_id == 14:
  159. span["content"] = layout_det["latex"]
  160. span["type"] = "displayed_equation"
  161. elif category_id == 15:
  162. span["content"] = layout_det["text"]
  163. span["type"] = "text"
  164. # print(span)
  165. spans.append(span)
  166. else:
  167. continue
  168. # 删除重叠spans中较小的那些
  169. spans = remove_overlaps_min_spans(spans)
  170. # 删除remove_span_block_bboxes中的bbox
  171. # spans = remove_spans_by_bboxes(spans, need_remove_spans_bboxes)
  172. # 按qa要求,增加drop相关数据
  173. spans, dropped_text_block, dropped_image_block, dropped_table_block = remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict)
  174. # 对image和table截图
  175. spans = cut_image_and_table(spans, page, page_id, book_name, save_path)
  176. # 行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)
  177. displayed_list = []
  178. text_inline_lines = []
  179. modify_y_axis(spans, displayed_list, text_inline_lines)
  180. # 模型识别错误的行间公式, type类型转换成行内公式
  181. spans = modify_inline_equation(spans, displayed_list, text_inline_lines)
  182. # bbox去除粘连
  183. spans = remove_overlap_between_bbox(spans)
  184. # 对tpye=["displayed_equation", "image", "table"]进行额外处理,如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
  185. spans = adjust_bbox_for_standalone_block(spans)
  186. # 从ocr_page_info中解析layout信息(按自然阅读方向排序,并修复重叠和交错的bad case)
  187. layout_bboxes, layout_tree = layout_detect(ocr_page_info['subfield_dets'], page, ocr_page_info)
  188. # 将spans合并成line(在layout内,从上到下,从左到右)
  189. lines = merge_spans_to_line_by_layout(spans, layout_bboxes)
  190. # 将lines合并成block
  191. blocks = merge_lines_to_block(lines)
  192. # 根据block合并段落
  193. # 获取QA需要外置的list
  194. images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
  195. # 构造pdf_info_dict
  196. page_info = construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
  197. images, tables, interline_equations, inline_equations,
  198. dropped_text_block, dropped_image_block, dropped_table_block)
  199. pdf_info_dict[f"page_{page_id}"] = page_info
  200. # 在测试时,保存调试信息
  201. if debug_mode:
  202. params_file_save_path = join_path(
  203. save_tmp_path, "md", book_name, "preproc_out.json"
  204. )
  205. with open(params_file_save_path, "w", encoding="utf-8") as f:
  206. json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
  207. # drow_bbox
  208. draw_layout_bbox(pdf_info_dict, pdf_path, md_bookname_save_path)
  209. draw_text_bbox(pdf_info_dict, pdf_path, md_bookname_save_path)
  210. return pdf_info_dict