pdf_parse_union_core_v2.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839
  1. import copy
  2. import os
  3. import statistics
  4. import time
  5. from typing import List
  6. import torch
  7. from loguru import logger
  8. from magic_pdf.config.drop_reason import DropReason
  9. from magic_pdf.config.enums import SupportedPdfParseMethod
  10. from magic_pdf.config.ocr_content_type import BlockType, ContentType
  11. from magic_pdf.data.dataset import Dataset, PageableData
  12. from magic_pdf.libs.boxbase import calculate_overlap_area_in_bbox1_area_ratio
  13. from magic_pdf.libs.clean_memory import clean_memory
  14. from magic_pdf.libs.commons import fitz, get_delta_time
  15. from magic_pdf.libs.config_reader import get_local_layoutreader_model_dir
  16. from magic_pdf.libs.convert_utils import dict_to_list
  17. from magic_pdf.libs.hash_utils import compute_md5
  18. from magic_pdf.libs.local_math import float_equal
  19. from magic_pdf.libs.pdf_image_tools import cut_image_to_pil_image
  20. from magic_pdf.model.magic_model import MagicModel
  21. os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新
  22. os.environ['YOLO_VERBOSE'] = 'False' # disable yolo logger
  23. try:
  24. import torchtext
  25. if torchtext.__version__ >= "0.18.0":
  26. torchtext.disable_torchtext_deprecation_warning()
  27. except ImportError:
  28. pass
  29. from magic_pdf.model.sub_modules.model_init import AtomModelSingleton
  30. from magic_pdf.para.para_split_v3 import para_split
  31. from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker
  32. from magic_pdf.pre_proc.construct_page_dict import \
  33. ocr_construct_page_component_v2
  34. from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table
  35. from magic_pdf.pre_proc.equations_replace import (
  36. combine_chars_to_pymudict, remove_chars_in_text_blocks,
  37. replace_equations_in_textblock)
  38. from magic_pdf.pre_proc.ocr_detect_all_bboxes import \
  39. ocr_prepare_bboxes_for_layout_split_v2
  40. from magic_pdf.pre_proc.ocr_dict_merge import (fill_spans_in_blocks,
  41. fix_block_spans_v2,
  42. fix_discarded_block)
  43. from magic_pdf.pre_proc.ocr_span_list_modify import (
  44. get_qa_need_list_v2, remove_overlaps_low_confidence_spans,
  45. remove_overlaps_min_spans)
  46. from magic_pdf.pre_proc.resolve_bbox_conflict import \
  47. check_useful_block_horizontal_overlap
  48. def remove_horizontal_overlap_block_which_smaller(all_bboxes):
  49. useful_blocks = []
  50. for bbox in all_bboxes:
  51. useful_blocks.append({'bbox': bbox[:4]})
  52. is_useful_block_horz_overlap, smaller_bbox, bigger_bbox = (
  53. check_useful_block_horizontal_overlap(useful_blocks)
  54. )
  55. if is_useful_block_horz_overlap:
  56. logger.warning(
  57. f'skip this page, reason: {DropReason.USEFUL_BLOCK_HOR_OVERLAP}, smaller bbox is {smaller_bbox}, bigger bbox is {bigger_bbox}'
  58. ) # noqa: E501
  59. for bbox in all_bboxes.copy():
  60. if smaller_bbox == bbox[:4]:
  61. all_bboxes.remove(bbox)
  62. return is_useful_block_horz_overlap, all_bboxes
  63. def __replace_STX_ETX(text_str: str):
  64. """Replace \u0002 and \u0003, as these characters become garbled when extracted using pymupdf. In fact, they were originally quotation marks.
  65. Drawback: This issue is only observed in English text; it has not been found in Chinese text so far.
  66. Args:
  67. text_str (str): raw text
  68. Returns:
  69. _type_: replaced text
  70. """ # noqa: E501
  71. if text_str:
  72. s = text_str.replace('\u0002', "'")
  73. s = s.replace('\u0003', "'")
  74. return s
  75. return text_str
  76. def chars_to_content(span):
  77. # # 先给chars按char['bbox']的x坐标排序
  78. # span['chars'] = sorted(span['chars'], key=lambda x: x['bbox'][0])
  79. # 先给chars按char['bbox']的中心点的x坐标排序
  80. span['chars'] = sorted(span['chars'], key=lambda x: (x['bbox'][0] + x['bbox'][2]) / 2)
  81. content = ''
  82. # 求char的平均宽度
  83. if len(span['chars']) == 0:
  84. span['content'] = content
  85. del span['chars']
  86. return
  87. else:
  88. char_width_sum = sum([char['bbox'][2] - char['bbox'][0] for char in span['chars']])
  89. char_avg_width = char_width_sum / len(span['chars'])
  90. for char in span['chars']:
  91. # 如果下一个char的x0和上一个char的x1距离超过一个字符宽度,则需要在中间插入一个空格
  92. if char['bbox'][0] - span['chars'][span['chars'].index(char) - 1]['bbox'][2] > char_avg_width:
  93. content += ' '
  94. content += char['c']
  95. span['content'] = __replace_STX_ETX(content)
  96. del span['chars']
  97. LINE_STOP_FLAG = ('.', '!', '?', '。', '!', '?', ')', ')', '"', '”', ':', ':', ';', ';', ']', '】', '}', '}', '>', '》', '、', ',', ',', '-', '—', '–',)
  98. def fill_char_in_spans(spans, all_chars):
  99. for char in all_chars:
  100. for span in spans:
  101. # 判断char是否属于LINE_STOP_FLAG
  102. if char['c'] in LINE_STOP_FLAG:
  103. char_is_line_stop_flag = True
  104. else:
  105. char_is_line_stop_flag = False
  106. if calculate_char_in_span(char['bbox'], span['bbox'], char_is_line_stop_flag):
  107. span['chars'].append(char)
  108. break
  109. empty_spans = []
  110. for span in spans:
  111. chars_to_content(span)
  112. if len(span['content']) == 0:
  113. empty_spans.append(span)
  114. return empty_spans
  115. # 使用鲁棒性更强的中心点坐标判断
  116. def calculate_char_in_span(char_bbox, span_bbox, char_is_line_stop_flag):
  117. char_center_x = (char_bbox[0] + char_bbox[2]) / 2
  118. char_center_y = (char_bbox[1] + char_bbox[3]) / 2
  119. span_center_y = (span_bbox[1] + span_bbox[3]) / 2
  120. span_height = span_bbox[3] - span_bbox[1]
  121. if (
  122. span_bbox[0] < char_center_x < span_bbox[2]
  123. and span_bbox[1] < char_center_y < span_bbox[3]
  124. and abs(char_center_y - span_center_y) < span_height / 4 # 字符的中轴和span的中轴高度差不能超过1/4span高度
  125. ):
  126. return True
  127. else:
  128. # 如果char是LINE_STOP_FLAG,就不用中心点判定,换一种方案(左边界在span区域内,高度判定和之前逻辑一致)
  129. # 主要是给结尾符号一个进入span的机会,这个char还应该离span右边界较近
  130. if char_is_line_stop_flag:
  131. if (
  132. (span_bbox[2] - span_height) < char_bbox[0] < span_bbox[2]
  133. and char_center_x > span_bbox[0]
  134. and span_bbox[1] < char_center_y < span_bbox[3]
  135. and abs(char_center_y - span_center_y) < span_height / 4
  136. ):
  137. return True
  138. else:
  139. return False
  140. def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang):
  141. text_blocks = pdf_page.get_text('rawdict', flags=fitz.TEXTFLAGS_TEXT)['blocks']
  142. # @todo: 拿到char之后把倾斜角度较大的先删一遍
  143. all_pymu_chars = []
  144. for block in text_blocks:
  145. for line in block['lines']:
  146. for span in line['spans']:
  147. all_pymu_chars.extend(span['chars'])
  148. useful_spans = []
  149. unuseful_spans = []
  150. for span in spans:
  151. for block in all_bboxes + all_discarded_blocks:
  152. if block[7] in [BlockType.ImageBody, BlockType.TableBody, BlockType.InterlineEquation]:
  153. continue
  154. overlap_ratio = calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], block[0:4])
  155. if overlap_ratio > 0.5:
  156. if block in all_bboxes:
  157. useful_spans.append(span)
  158. else:
  159. unuseful_spans.append(span)
  160. break
  161. new_spans = []
  162. for span in useful_spans + unuseful_spans:
  163. if span['type'] in [ContentType.Text]:
  164. span['chars'] = []
  165. new_spans.append(span)
  166. empty_spans = fill_char_in_spans(new_spans, all_pymu_chars)
  167. if len(empty_spans) > 0:
  168. # 初始化ocr模型
  169. atom_model_manager = AtomModelSingleton()
  170. ocr_model = atom_model_manager.get_atom_model(
  171. atom_model_name="ocr",
  172. ocr_show_log=False,
  173. det_db_box_thresh=0.3,
  174. lang=lang
  175. )
  176. for span in empty_spans:
  177. # 对span的bbox截图再ocr
  178. span_img = cut_image_to_pil_image(span['bbox'], pdf_page, mode="cv2")
  179. ocr_res = ocr_model.ocr(span_img, det=False)
  180. if ocr_res and len(ocr_res) > 0:
  181. if len(ocr_res[0]) > 0:
  182. ocr_text, ocr_score = ocr_res[0][0]
  183. if ocr_score > 0.5 and len(ocr_text) > 0:
  184. span['content'] = ocr_text
  185. return spans
  186. def txt_spans_extract_v1(pdf_page, inline_equations, interline_equations):
  187. text_raw_blocks = pdf_page.get_text('dict', flags=fitz.TEXTFLAGS_TEXT)['blocks']
  188. char_level_text_blocks = pdf_page.get_text('rawdict', flags=fitz.TEXTFLAGS_TEXT)[
  189. 'blocks'
  190. ]
  191. text_blocks = combine_chars_to_pymudict(text_raw_blocks, char_level_text_blocks)
  192. text_blocks = replace_equations_in_textblock(
  193. text_blocks, inline_equations, interline_equations
  194. )
  195. text_blocks = remove_citation_marker(text_blocks)
  196. text_blocks = remove_chars_in_text_blocks(text_blocks)
  197. spans = []
  198. for v in text_blocks:
  199. for line in v['lines']:
  200. for span in line['spans']:
  201. bbox = span['bbox']
  202. if float_equal(bbox[0], bbox[2]) or float_equal(bbox[1], bbox[3]):
  203. continue
  204. if span.get('type') not in (
  205. ContentType.InlineEquation,
  206. ContentType.InterlineEquation,
  207. ):
  208. spans.append(
  209. {
  210. 'bbox': list(span['bbox']),
  211. 'content': __replace_STX_ETX(span['text']),
  212. 'type': ContentType.Text,
  213. 'score': 1.0,
  214. }
  215. )
  216. return spans
  217. def replace_text_span(pymu_spans, ocr_spans):
  218. return list(filter(lambda x: x['type'] != ContentType.Text, ocr_spans)) + pymu_spans
  219. def model_init(model_name: str):
  220. from transformers import LayoutLMv3ForTokenClassification
  221. if torch.cuda.is_available():
  222. device = torch.device('cuda')
  223. if torch.cuda.is_bf16_supported():
  224. supports_bfloat16 = True
  225. else:
  226. supports_bfloat16 = False
  227. else:
  228. device = torch.device('cpu')
  229. supports_bfloat16 = False
  230. if model_name == 'layoutreader':
  231. # 检测modelscope的缓存目录是否存在
  232. layoutreader_model_dir = get_local_layoutreader_model_dir()
  233. if os.path.exists(layoutreader_model_dir):
  234. model = LayoutLMv3ForTokenClassification.from_pretrained(
  235. layoutreader_model_dir
  236. )
  237. else:
  238. logger.warning(
  239. 'local layoutreader model not exists, use online model from huggingface'
  240. )
  241. model = LayoutLMv3ForTokenClassification.from_pretrained(
  242. 'hantian/layoutreader'
  243. )
  244. # 检查设备是否支持 bfloat16
  245. if supports_bfloat16:
  246. model.bfloat16()
  247. model.to(device).eval()
  248. else:
  249. logger.error('model name not allow')
  250. exit(1)
  251. return model
  252. class ModelSingleton:
  253. _instance = None
  254. _models = {}
  255. def __new__(cls, *args, **kwargs):
  256. if cls._instance is None:
  257. cls._instance = super().__new__(cls)
  258. return cls._instance
  259. def get_model(self, model_name: str):
  260. if model_name not in self._models:
  261. self._models[model_name] = model_init(model_name=model_name)
  262. return self._models[model_name]
  263. def do_predict(boxes: List[List[int]], model) -> List[int]:
  264. from magic_pdf.model.sub_modules.reading_oreder.layoutreader.helpers import (
  265. boxes2inputs, parse_logits, prepare_inputs)
  266. inputs = boxes2inputs(boxes)
  267. inputs = prepare_inputs(inputs, model)
  268. logits = model(**inputs).logits.cpu().squeeze(0)
  269. return parse_logits(logits, len(boxes))
  270. def cal_block_index(fix_blocks, sorted_bboxes):
  271. if sorted_bboxes is not None:
  272. # 使用layoutreader排序
  273. for block in fix_blocks:
  274. line_index_list = []
  275. if len(block['lines']) == 0:
  276. block['index'] = sorted_bboxes.index(block['bbox'])
  277. else:
  278. for line in block['lines']:
  279. line['index'] = sorted_bboxes.index(line['bbox'])
  280. line_index_list.append(line['index'])
  281. median_value = statistics.median(line_index_list)
  282. block['index'] = median_value
  283. # 删除图表body block中的虚拟line信息, 并用real_lines信息回填
  284. if block['type'] in [BlockType.ImageBody, BlockType.TableBody]:
  285. block['virtual_lines'] = copy.deepcopy(block['lines'])
  286. block['lines'] = copy.deepcopy(block['real_lines'])
  287. del block['real_lines']
  288. else:
  289. # 使用xycut排序
  290. block_bboxes = []
  291. for block in fix_blocks:
  292. block_bboxes.append(block['bbox'])
  293. # 删除图表body block中的虚拟line信息, 并用real_lines信息回填
  294. if block['type'] in [BlockType.ImageBody, BlockType.TableBody]:
  295. block['virtual_lines'] = copy.deepcopy(block['lines'])
  296. block['lines'] = copy.deepcopy(block['real_lines'])
  297. del block['real_lines']
  298. import numpy as np
  299. from magic_pdf.model.sub_modules.reading_oreder.layoutreader.xycut import \
  300. recursive_xy_cut
  301. random_boxes = np.array(block_bboxes)
  302. np.random.shuffle(random_boxes)
  303. res = []
  304. recursive_xy_cut(np.asarray(random_boxes).astype(int), np.arange(len(block_bboxes)), res)
  305. assert len(res) == len(block_bboxes)
  306. sorted_boxes = random_boxes[np.array(res)].tolist()
  307. for i, block in enumerate(fix_blocks):
  308. block['index'] = sorted_boxes.index(block['bbox'])
  309. # 生成line index
  310. sorted_blocks = sorted(fix_blocks, key=lambda b: b['index'])
  311. line_inedx = 1
  312. for block in sorted_blocks:
  313. for line in block['lines']:
  314. line['index'] = line_inedx
  315. line_inedx += 1
  316. return fix_blocks
  317. def insert_lines_into_block(block_bbox, line_height, page_w, page_h):
  318. # block_bbox是一个元组(x0, y0, x1, y1),其中(x0, y0)是左下角坐标,(x1, y1)是右上角坐标
  319. x0, y0, x1, y1 = block_bbox
  320. block_height = y1 - y0
  321. block_weight = x1 - x0
  322. # 如果block高度小于n行正文,则直接返回block的bbox
  323. if line_height * 3 < block_height:
  324. if (
  325. block_height > page_h * 0.25 and page_w * 0.5 > block_weight > page_w * 0.25
  326. ): # 可能是双列结构,可以切细点
  327. lines = int(block_height / line_height) + 1
  328. else:
  329. # 如果block的宽度超过0.4页面宽度,则将block分成3行(是一种复杂布局,图不能切的太细)
  330. if block_weight > page_w * 0.4:
  331. line_height = (y1 - y0) / 3
  332. lines = 3
  333. elif block_weight > page_w * 0.25: # (可能是三列结构,也切细点)
  334. lines = int(block_height / line_height) + 1
  335. else: # 判断长宽比
  336. if block_height / block_weight > 1.2: # 细长的不分
  337. return [[x0, y0, x1, y1]]
  338. else: # 不细长的还是分成两行
  339. line_height = (y1 - y0) / 2
  340. lines = 2
  341. # 确定从哪个y位置开始绘制线条
  342. current_y = y0
  343. # 用于存储线条的位置信息[(x0, y), ...]
  344. lines_positions = []
  345. for i in range(lines):
  346. lines_positions.append([x0, current_y, x1, current_y + line_height])
  347. current_y += line_height
  348. return lines_positions
  349. else:
  350. return [[x0, y0, x1, y1]]
  351. def sort_lines_by_model(fix_blocks, page_w, page_h, line_height):
  352. page_line_list = []
  353. for block in fix_blocks:
  354. if block['type'] in [
  355. BlockType.Text, BlockType.Title, BlockType.InterlineEquation,
  356. BlockType.ImageCaption, BlockType.ImageFootnote,
  357. BlockType.TableCaption, BlockType.TableFootnote
  358. ]:
  359. if len(block['lines']) == 0:
  360. bbox = block['bbox']
  361. lines = insert_lines_into_block(bbox, line_height, page_w, page_h)
  362. for line in lines:
  363. block['lines'].append({'bbox': line, 'spans': []})
  364. page_line_list.extend(lines)
  365. else:
  366. for line in block['lines']:
  367. bbox = line['bbox']
  368. page_line_list.append(bbox)
  369. elif block['type'] in [BlockType.ImageBody, BlockType.TableBody]:
  370. bbox = block['bbox']
  371. block['real_lines'] = copy.deepcopy(block['lines'])
  372. lines = insert_lines_into_block(bbox, line_height, page_w, page_h)
  373. block['lines'] = []
  374. for line in lines:
  375. block['lines'].append({'bbox': line, 'spans': []})
  376. page_line_list.extend(lines)
  377. if len(page_line_list) > 200: # layoutreader最高支持512line
  378. return None
  379. # 使用layoutreader排序
  380. x_scale = 1000.0 / page_w
  381. y_scale = 1000.0 / page_h
  382. boxes = []
  383. # logger.info(f"Scale: {x_scale}, {y_scale}, Boxes len: {len(page_line_list)}")
  384. for left, top, right, bottom in page_line_list:
  385. if left < 0:
  386. logger.warning(
  387. f'left < 0, left: {left}, right: {right}, top: {top}, bottom: {bottom}, page_w: {page_w}, page_h: {page_h}'
  388. ) # noqa: E501
  389. left = 0
  390. if right > page_w:
  391. logger.warning(
  392. f'right > page_w, left: {left}, right: {right}, top: {top}, bottom: {bottom}, page_w: {page_w}, page_h: {page_h}'
  393. ) # noqa: E501
  394. right = page_w
  395. if top < 0:
  396. logger.warning(
  397. f'top < 0, left: {left}, right: {right}, top: {top}, bottom: {bottom}, page_w: {page_w}, page_h: {page_h}'
  398. ) # noqa: E501
  399. top = 0
  400. if bottom > page_h:
  401. logger.warning(
  402. f'bottom > page_h, left: {left}, right: {right}, top: {top}, bottom: {bottom}, page_w: {page_w}, page_h: {page_h}'
  403. ) # noqa: E501
  404. bottom = page_h
  405. left = round(left * x_scale)
  406. top = round(top * y_scale)
  407. right = round(right * x_scale)
  408. bottom = round(bottom * y_scale)
  409. assert (
  410. 1000 >= right >= left >= 0 and 1000 >= bottom >= top >= 0
  411. ), f'Invalid box. right: {right}, left: {left}, bottom: {bottom}, top: {top}' # noqa: E126, E121
  412. boxes.append([left, top, right, bottom])
  413. model_manager = ModelSingleton()
  414. model = model_manager.get_model('layoutreader')
  415. with torch.no_grad():
  416. orders = do_predict(boxes, model)
  417. sorted_bboxes = [page_line_list[i] for i in orders]
  418. return sorted_bboxes
  419. def get_line_height(blocks):
  420. page_line_height_list = []
  421. for block in blocks:
  422. if block['type'] in [
  423. BlockType.Text, BlockType.Title,
  424. BlockType.ImageCaption, BlockType.ImageFootnote,
  425. BlockType.TableCaption, BlockType.TableFootnote
  426. ]:
  427. for line in block['lines']:
  428. bbox = line['bbox']
  429. page_line_height_list.append(int(bbox[3] - bbox[1]))
  430. if len(page_line_height_list) > 0:
  431. return statistics.median(page_line_height_list)
  432. else:
  433. return 10
  434. def process_groups(groups, body_key, caption_key, footnote_key):
  435. body_blocks = []
  436. caption_blocks = []
  437. footnote_blocks = []
  438. for i, group in enumerate(groups):
  439. group[body_key]['group_id'] = i
  440. body_blocks.append(group[body_key])
  441. for caption_block in group[caption_key]:
  442. caption_block['group_id'] = i
  443. caption_blocks.append(caption_block)
  444. for footnote_block in group[footnote_key]:
  445. footnote_block['group_id'] = i
  446. footnote_blocks.append(footnote_block)
  447. return body_blocks, caption_blocks, footnote_blocks
  448. def process_block_list(blocks, body_type, block_type):
  449. indices = [block['index'] for block in blocks]
  450. median_index = statistics.median(indices)
  451. body_bbox = next((block['bbox'] for block in blocks if block.get('type') == body_type), [])
  452. return {
  453. 'type': block_type,
  454. 'bbox': body_bbox,
  455. 'blocks': blocks,
  456. 'index': median_index,
  457. }
  458. def revert_group_blocks(blocks):
  459. image_groups = {}
  460. table_groups = {}
  461. new_blocks = []
  462. for block in blocks:
  463. if block['type'] in [BlockType.ImageBody, BlockType.ImageCaption, BlockType.ImageFootnote]:
  464. group_id = block['group_id']
  465. if group_id not in image_groups:
  466. image_groups[group_id] = []
  467. image_groups[group_id].append(block)
  468. elif block['type'] in [BlockType.TableBody, BlockType.TableCaption, BlockType.TableFootnote]:
  469. group_id = block['group_id']
  470. if group_id not in table_groups:
  471. table_groups[group_id] = []
  472. table_groups[group_id].append(block)
  473. else:
  474. new_blocks.append(block)
  475. for group_id, blocks in image_groups.items():
  476. new_blocks.append(process_block_list(blocks, BlockType.ImageBody, BlockType.Image))
  477. for group_id, blocks in table_groups.items():
  478. new_blocks.append(process_block_list(blocks, BlockType.TableBody, BlockType.Table))
  479. return new_blocks
  480. def remove_outside_spans(spans, all_bboxes, all_discarded_blocks):
  481. def get_block_bboxes(blocks, block_type_list):
  482. return [block[0:4] for block in blocks if block[7] in block_type_list]
  483. image_bboxes = get_block_bboxes(all_bboxes, [BlockType.ImageBody])
  484. table_bboxes = get_block_bboxes(all_bboxes, [BlockType.TableBody])
  485. other_block_type = []
  486. for block_type in BlockType.__dict__.values():
  487. if not isinstance(block_type, str):
  488. continue
  489. if block_type not in [BlockType.ImageBody, BlockType.TableBody]:
  490. other_block_type.append(block_type)
  491. other_block_bboxes = get_block_bboxes(all_bboxes, other_block_type)
  492. discarded_block_bboxes = get_block_bboxes(all_discarded_blocks, [BlockType.Discarded])
  493. new_spans = []
  494. for span in spans:
  495. span_bbox = span['bbox']
  496. span_type = span['type']
  497. if any(calculate_overlap_area_in_bbox1_area_ratio(span_bbox, block_bbox) > 0.4 for block_bbox in
  498. discarded_block_bboxes):
  499. new_spans.append(span)
  500. continue
  501. if span_type == ContentType.Image:
  502. if any(calculate_overlap_area_in_bbox1_area_ratio(span_bbox, block_bbox) > 0.5 for block_bbox in
  503. image_bboxes):
  504. new_spans.append(span)
  505. elif span_type == ContentType.Table:
  506. if any(calculate_overlap_area_in_bbox1_area_ratio(span_bbox, block_bbox) > 0.5 for block_bbox in
  507. table_bboxes):
  508. new_spans.append(span)
  509. else:
  510. if any(calculate_overlap_area_in_bbox1_area_ratio(span_bbox, block_bbox) > 0.5 for block_bbox in
  511. other_block_bboxes):
  512. new_spans.append(span)
  513. return new_spans
  514. def parse_page_core(
  515. page_doc: PageableData, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, lang
  516. ):
  517. need_drop = False
  518. drop_reason = []
  519. """从magic_model对象中获取后面会用到的区块信息"""
  520. img_groups = magic_model.get_imgs_v2(page_id)
  521. table_groups = magic_model.get_tables_v2(page_id)
  522. """对image和table的区块分组"""
  523. img_body_blocks, img_caption_blocks, img_footnote_blocks = process_groups(
  524. img_groups, 'image_body', 'image_caption_list', 'image_footnote_list'
  525. )
  526. table_body_blocks, table_caption_blocks, table_footnote_blocks = process_groups(
  527. table_groups, 'table_body', 'table_caption_list', 'table_footnote_list'
  528. )
  529. discarded_blocks = magic_model.get_discarded(page_id)
  530. text_blocks = magic_model.get_text_blocks(page_id)
  531. title_blocks = magic_model.get_title_blocks(page_id)
  532. inline_equations, interline_equations, interline_equation_blocks = (
  533. magic_model.get_equations(page_id)
  534. )
  535. page_w, page_h = magic_model.get_page_size(page_id)
  536. """将所有区块的bbox整理到一起"""
  537. # interline_equation_blocks参数不够准,后面切换到interline_equations上
  538. interline_equation_blocks = []
  539. if len(interline_equation_blocks) > 0:
  540. all_bboxes, all_discarded_blocks = ocr_prepare_bboxes_for_layout_split_v2(
  541. img_body_blocks, img_caption_blocks, img_footnote_blocks,
  542. table_body_blocks, table_caption_blocks, table_footnote_blocks,
  543. discarded_blocks,
  544. text_blocks,
  545. title_blocks,
  546. interline_equation_blocks,
  547. page_w,
  548. page_h,
  549. )
  550. else:
  551. all_bboxes, all_discarded_blocks = ocr_prepare_bboxes_for_layout_split_v2(
  552. img_body_blocks, img_caption_blocks, img_footnote_blocks,
  553. table_body_blocks, table_caption_blocks, table_footnote_blocks,
  554. discarded_blocks,
  555. text_blocks,
  556. title_blocks,
  557. interline_equations,
  558. page_w,
  559. page_h,
  560. )
  561. """获取所有的spans信息"""
  562. spans = magic_model.get_all_spans(page_id)
  563. """在删除重复span之前,应该通过image_body和table_body的block过滤一下image和table的span"""
  564. """顺便删除大水印并保留abandon的span"""
  565. spans = remove_outside_spans(spans, all_bboxes, all_discarded_blocks)
  566. """删除重叠spans中置信度较低的那些"""
  567. spans, dropped_spans_by_confidence = remove_overlaps_low_confidence_spans(spans)
  568. """删除重叠spans中较小的那些"""
  569. spans, dropped_spans_by_span_overlap = remove_overlaps_min_spans(spans)
  570. """根据parse_mode,构造spans,主要是文本类的字符填充"""
  571. if parse_mode == SupportedPdfParseMethod.TXT:
  572. """之前的公式替换方案"""
  573. # pymu_spans = txt_spans_extract_v1(page_doc, inline_equations, interline_equations)
  574. # spans = replace_text_span(pymu_spans, spans)
  575. """使用新版本的混合ocr方案"""
  576. spans = txt_spans_extract_v2(page_doc, spans, all_bboxes, all_discarded_blocks, lang)
  577. elif parse_mode == SupportedPdfParseMethod.OCR:
  578. pass
  579. else:
  580. raise Exception('parse_mode must be txt or ocr')
  581. """先处理不需要排版的discarded_blocks"""
  582. discarded_block_with_spans, spans = fill_spans_in_blocks(
  583. all_discarded_blocks, spans, 0.4
  584. )
  585. fix_discarded_blocks = fix_discarded_block(discarded_block_with_spans)
  586. """如果当前页面没有有效的bbox则跳过"""
  587. if len(all_bboxes) == 0:
  588. logger.warning(f'skip this page, not found useful bbox, page_id: {page_id}')
  589. return ocr_construct_page_component_v2(
  590. [],
  591. [],
  592. page_id,
  593. page_w,
  594. page_h,
  595. [],
  596. [],
  597. [],
  598. interline_equations,
  599. fix_discarded_blocks,
  600. need_drop,
  601. drop_reason,
  602. )
  603. """对image和table截图"""
  604. spans = ocr_cut_image_and_table(
  605. spans, page_doc, page_id, pdf_bytes_md5, imageWriter
  606. )
  607. """span填充进block"""
  608. block_with_spans, spans = fill_spans_in_blocks(all_bboxes, spans, 0.5)
  609. """对block进行fix操作"""
  610. fix_blocks = fix_block_spans_v2(block_with_spans)
  611. """获取所有line并计算正文line的高度"""
  612. line_height = get_line_height(fix_blocks)
  613. """获取所有line并对line排序"""
  614. sorted_bboxes = sort_lines_by_model(fix_blocks, page_w, page_h, line_height)
  615. """根据line的中位数算block的序列关系"""
  616. fix_blocks = cal_block_index(fix_blocks, sorted_bboxes)
  617. """将image和table的block还原回group形式参与后续流程"""
  618. fix_blocks = revert_group_blocks(fix_blocks)
  619. """重排block"""
  620. sorted_blocks = sorted(fix_blocks, key=lambda b: b['index'])
  621. """获取QA需要外置的list"""
  622. images, tables, interline_equations = get_qa_need_list_v2(sorted_blocks)
  623. """构造pdf_info_dict"""
  624. page_info = ocr_construct_page_component_v2(
  625. sorted_blocks,
  626. [],
  627. page_id,
  628. page_w,
  629. page_h,
  630. [],
  631. images,
  632. tables,
  633. interline_equations,
  634. fix_discarded_blocks,
  635. need_drop,
  636. drop_reason,
  637. )
  638. return page_info
  639. def pdf_parse_union(
  640. dataset: Dataset,
  641. model_list,
  642. imageWriter,
  643. parse_mode,
  644. start_page_id=0,
  645. end_page_id=None,
  646. debug_mode=False,
  647. lang=None,
  648. ):
  649. pdf_bytes_md5 = compute_md5(dataset.data_bits())
  650. """初始化空的pdf_info_dict"""
  651. pdf_info_dict = {}
  652. """用model_list和docs对象初始化magic_model"""
  653. magic_model = MagicModel(model_list, dataset)
  654. """根据输入的起始范围解析pdf"""
  655. # end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
  656. end_page_id = (
  657. end_page_id
  658. if end_page_id is not None and end_page_id >= 0
  659. else len(dataset) - 1
  660. )
  661. if end_page_id > len(dataset) - 1:
  662. logger.warning('end_page_id is out of range, use pdf_docs length')
  663. end_page_id = len(dataset) - 1
  664. """初始化启动时间"""
  665. start_time = time.time()
  666. for page_id, page in enumerate(dataset):
  667. """debug时输出每页解析的耗时."""
  668. if debug_mode:
  669. time_now = time.time()
  670. logger.info(
  671. f'page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}'
  672. )
  673. start_time = time_now
  674. """解析pdf中的每一页"""
  675. if start_page_id <= page_id <= end_page_id:
  676. page_info = parse_page_core(
  677. page, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, lang
  678. )
  679. else:
  680. page_info = page.get_page_info()
  681. page_w = page_info.w
  682. page_h = page_info.h
  683. page_info = ocr_construct_page_component_v2(
  684. [], [], page_id, page_w, page_h, [], [], [], [], [], True, 'skip page'
  685. )
  686. pdf_info_dict[f'page_{page_id}'] = page_info
  687. """分段"""
  688. para_split(pdf_info_dict)
  689. """dict转list"""
  690. pdf_info_list = dict_to_list(pdf_info_dict)
  691. new_pdf_info_dict = {
  692. 'pdf_info': pdf_info_list,
  693. }
  694. clean_memory()
  695. return new_pdf_info_dict
  696. if __name__ == '__main__':
  697. pass