pdf_parse_union_core_v2.py 30 KB

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