pdf_parse_by_txt_v2.py 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266
  1. import time
  2. from loguru import logger
  3. from magic_pdf.layout.layout_sort import get_bboxes_layout, LAYOUT_UNPROC, get_columns_cnt_of_layout
  4. from magic_pdf.libs.convert_utils import dict_to_list
  5. from magic_pdf.libs.drop_reason import DropReason
  6. from magic_pdf.libs.hash_utils import compute_md5
  7. from magic_pdf.libs.commons import fitz, get_delta_time
  8. from magic_pdf.model.magic_model import MagicModel
  9. from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2
  10. from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table
  11. from magic_pdf.pre_proc.ocr_detect_all_bboxes import ocr_prepare_bboxes_for_layout_split
  12. from magic_pdf.pre_proc.ocr_dict_merge import (
  13. sort_blocks_by_layout,
  14. fill_spans_in_blocks,
  15. fix_block_spans,
  16. )
  17. from magic_pdf.libs.ocr_content_type import ContentType
  18. from magic_pdf.pre_proc.ocr_span_list_modify import (
  19. remove_overlaps_min_spans,
  20. get_qa_need_list_v2,
  21. )
  22. from magic_pdf.pre_proc.equations_replace import (
  23. combine_chars_to_pymudict,
  24. remove_chars_in_text_blocks,
  25. replace_equations_in_textblock,
  26. )
  27. from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker
  28. from magic_pdf.libs.math import float_equal
  29. from magic_pdf.para.para_split_v2 import para_split
  30. from magic_pdf.pre_proc.resolve_bbox_conflict import check_useful_block_horizontal_overlap
  31. def txt_spans_extract(pdf_page, inline_equations, interline_equations):
  32. text_raw_blocks = pdf_page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"]
  33. char_level_text_blocks = pdf_page.get_text("rawdict", flags=fitz.TEXTFLAGS_TEXT)[
  34. "blocks"
  35. ]
  36. text_blocks = combine_chars_to_pymudict(text_raw_blocks, char_level_text_blocks)
  37. text_blocks = replace_equations_in_textblock(
  38. text_blocks, inline_equations, interline_equations
  39. )
  40. text_blocks = remove_citation_marker(text_blocks)
  41. text_blocks = remove_chars_in_text_blocks(text_blocks)
  42. spans = []
  43. for v in text_blocks:
  44. for line in v["lines"]:
  45. for span in line["spans"]:
  46. bbox = span["bbox"]
  47. if float_equal(bbox[0], bbox[2]) or float_equal(bbox[1], bbox[3]):
  48. continue
  49. if span.get('type') == ContentType.InlineEquation:
  50. spans.append(
  51. {
  52. "bbox": list(span["bbox"]),
  53. "content": span["latex"],
  54. "type": ContentType.InlineEquation,
  55. }
  56. )
  57. elif span.get('type') == ContentType.InterlineEquation:
  58. spans.append(
  59. {
  60. "bbox": list(span["bbox"]),
  61. "content": span["latex"],
  62. "type": ContentType.InterlineEquation,
  63. }
  64. )
  65. else:
  66. spans.append(
  67. {
  68. "bbox": list(span["bbox"]),
  69. "content": span["text"],
  70. "type": ContentType.Text,
  71. }
  72. )
  73. return spans
  74. def replace_text_span(pymu_spans, ocr_spans):
  75. return list(filter(lambda x: x["type"] != ContentType.Text, ocr_spans)) + pymu_spans
  76. def parse_pdf_by_txt(
  77. pdf_bytes,
  78. model_list,
  79. imageWriter,
  80. start_page_id=0,
  81. end_page_id=None,
  82. debug_mode=False,
  83. ):
  84. pdf_bytes_md5 = compute_md5(pdf_bytes)
  85. pdf_docs = fitz.open("pdf", pdf_bytes)
  86. """初始化空的pdf_info_dict"""
  87. pdf_info_dict = {}
  88. """用model_list和docs对象初始化magic_model"""
  89. magic_model = MagicModel(model_list, pdf_docs)
  90. """根据输入的起始范围解析pdf"""
  91. end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
  92. """初始化启动时间"""
  93. start_time = time.time()
  94. for page_id in range(start_page_id, end_page_id + 1):
  95. """debug时输出每页解析的耗时"""
  96. if debug_mode:
  97. time_now = time.time()
  98. logger.info(
  99. f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
  100. )
  101. start_time = time_now
  102. """从magic_model对象中获取后面会用到的区块信息"""
  103. img_blocks = magic_model.get_imgs(page_id)
  104. table_blocks = magic_model.get_tables(page_id)
  105. discarded_blocks = magic_model.get_discarded(page_id)
  106. text_blocks = magic_model.get_text_blocks(page_id)
  107. title_blocks = magic_model.get_title_blocks(page_id)
  108. inline_equations, interline_equations, interline_equation_blocks = (
  109. magic_model.get_equations(page_id)
  110. )
  111. page_w, page_h = magic_model.get_page_size(page_id)
  112. """将所有区块的bbox整理到一起"""
  113. all_bboxes = ocr_prepare_bboxes_for_layout_split(
  114. img_blocks,
  115. table_blocks,
  116. discarded_blocks,
  117. text_blocks,
  118. title_blocks,
  119. interline_equations,
  120. page_w,
  121. page_h,
  122. )
  123. """在切分之前,先检查一下bbox是否有左右重叠的情况,如果有,那么就认为这个pdf暂时没有能力处理好,这种左右重叠的情况大概率是由于pdf里的行间公式、表格没有被正确识别出来造成的 """
  124. useful_blocks = []
  125. for bbox in all_bboxes:
  126. useful_blocks.append({
  127. "bbox": bbox[:4]
  128. })
  129. is_useful_block_horz_overlap = check_useful_block_horizontal_overlap(useful_blocks)
  130. if is_useful_block_horz_overlap:
  131. logger.warning(
  132. f"pdf: {pdf_bytes_md5}, skip this page, page_id: {page_id}, reason: {DropReason.TEXT_BLCOK_HOR_OVERLAP}")
  133. continue
  134. '''根据区块信息计算layout'''
  135. page_boundry = [0, 0, page_w, page_h]
  136. layout_bboxes, layout_tree = get_bboxes_layout(all_bboxes, page_boundry, page_id)
  137. if len(text_blocks) > 0 and len(all_bboxes) > 0 and len(layout_bboxes) == 0:
  138. logger.warning(
  139. f"pdf: {pdf_bytes_md5}, skip this page, page_id: {page_id}, reason: {DropReason.CAN_NOT_DETECT_PAGE_LAYOUT}")
  140. continue
  141. """以下去掉复杂的布局和超过2列的布局"""
  142. if any([lay["layout_label"] == LAYOUT_UNPROC for lay in layout_bboxes]): # 复杂的布局
  143. logger.warning(
  144. f"pdf: {pdf_bytes_md5}, skip this page, page_id: {page_id}, reason: {DropReason.COMPLICATED_LAYOUT}")
  145. continue
  146. layout_column_width = get_columns_cnt_of_layout(layout_tree)
  147. if layout_column_width > 2: # 去掉超过2列的布局pdf
  148. logger.warning(
  149. f"pdf: {pdf_bytes_md5}, skip this page, page_id: {page_id}, reason: {DropReason.TOO_MANY_LAYOUT_COLUMNS}")
  150. continue
  151. """根据layout顺序,对当前页面所有需要留下的block进行排序"""
  152. sorted_blocks = sort_blocks_by_layout(all_bboxes, layout_bboxes)
  153. """ocr 中文本类的 span 用 pymu spans 替换!"""
  154. ocr_spans = magic_model.get_all_spans(page_id)
  155. pymu_spans = txt_spans_extract(
  156. pdf_docs[page_id], inline_equations, interline_equations
  157. )
  158. spans = replace_text_span(pymu_spans, ocr_spans)
  159. """删除重叠spans中较小的那些"""
  160. spans, dropped_spans_by_span_overlap = remove_overlaps_min_spans(spans)
  161. """对image和table截图"""
  162. spans = ocr_cut_image_and_table(
  163. spans, pdf_docs[page_id], page_id, pdf_bytes_md5, imageWriter
  164. )
  165. """将span填入排好序的blocks中"""
  166. block_with_spans = fill_spans_in_blocks(sorted_blocks, spans)
  167. """对block进行fix操作"""
  168. fix_blocks = fix_block_spans(block_with_spans, img_blocks, table_blocks)
  169. """获取QA需要外置的list"""
  170. images, tables, interline_equations = get_qa_need_list_v2(fix_blocks)
  171. """构造pdf_info_dict"""
  172. page_info = ocr_construct_page_component_v2(
  173. fix_blocks,
  174. layout_bboxes,
  175. page_id,
  176. page_w,
  177. page_h,
  178. layout_tree,
  179. images,
  180. tables,
  181. interline_equations,
  182. discarded_blocks,
  183. )
  184. pdf_info_dict[f"page_{page_id}"] = page_info
  185. """分段"""
  186. try:
  187. para_split(pdf_info_dict, debug_mode=debug_mode)
  188. except Exception as e:
  189. logger.exception(e)
  190. raise e
  191. """dict转list"""
  192. pdf_info_list = dict_to_list(pdf_info_dict)
  193. new_pdf_info_dict = {
  194. "pdf_info": pdf_info_list,
  195. }
  196. return new_pdf_info_dict
  197. if __name__ == "__main__":
  198. if 1:
  199. import fitz
  200. import json
  201. with open("/opt/data/pdf/20240418/25536-00.pdf", "rb") as f:
  202. pdf_bytes = f.read()
  203. pdf_docs = fitz.open("pdf", pdf_bytes)
  204. with open("/opt/data/pdf/20240418/25536-00.json") as f:
  205. model_list = json.loads(f.readline())
  206. magic_model = MagicModel(model_list, pdf_docs)
  207. for i in range(7):
  208. print(magic_model.get_imgs(i))
  209. for page_no, page in enumerate(pdf_docs):
  210. inline_equations, interline_equations, interline_equation_blocks = (
  211. magic_model.get_equations(page_no)
  212. )
  213. text_raw_blocks = page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"]
  214. char_level_text_blocks = page.get_text(
  215. "rawdict", flags=fitz.TEXTFLAGS_TEXT
  216. )["blocks"]
  217. text_blocks = combine_chars_to_pymudict(
  218. text_raw_blocks, char_level_text_blocks
  219. )
  220. text_blocks = replace_equations_in_textblock(
  221. text_blocks, inline_equations, interline_equations
  222. )
  223. text_blocks = remove_citation_marker(text_blocks)
  224. text_blocks = remove_chars_in_text_blocks(text_blocks)