pdf_parse_union_core_v2.py 29 KB

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