pdf_parse_by_txt_v2.py 9.2 KB

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