pdf_parse_by_ocr.py 10 KB

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