magic_model.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709
  1. import json
  2. from magic_pdf.libs.boxbase import (_is_in, _is_part_overlap, bbox_distance,
  3. bbox_relative_pos, calculate_iou,
  4. calculate_overlap_area_in_bbox1_area_ratio)
  5. from magic_pdf.libs.commons import fitz, join_path
  6. from magic_pdf.libs.coordinate_transform import get_scale_ratio
  7. from magic_pdf.libs.local_math import float_gt
  8. from magic_pdf.libs.ModelBlockTypeEnum import ModelBlockTypeEnum
  9. from magic_pdf.libs.ocr_content_type import CategoryId, ContentType
  10. from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
  11. from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
  12. CAPATION_OVERLAP_AREA_RATIO = 0.6
  13. class MagicModel:
  14. """每个函数没有得到元素的时候返回空list."""
  15. def __fix_axis(self):
  16. for model_page_info in self.__model_list:
  17. need_remove_list = []
  18. page_no = model_page_info['page_info']['page_no']
  19. horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
  20. model_page_info, self.__docs[page_no]
  21. )
  22. layout_dets = model_page_info['layout_dets']
  23. for layout_det in layout_dets:
  24. if layout_det.get('bbox') is not None:
  25. # 兼容直接输出bbox的模型数据,如paddle
  26. x0, y0, x1, y1 = layout_det['bbox']
  27. else:
  28. # 兼容直接输出poly的模型数据,如xxx
  29. x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
  30. bbox = [
  31. int(x0 / horizontal_scale_ratio),
  32. int(y0 / vertical_scale_ratio),
  33. int(x1 / horizontal_scale_ratio),
  34. int(y1 / vertical_scale_ratio),
  35. ]
  36. layout_det['bbox'] = bbox
  37. # 删除高度或者宽度小于等于0的spans
  38. if bbox[2] - bbox[0] <= 0 or bbox[3] - bbox[1] <= 0:
  39. need_remove_list.append(layout_det)
  40. for need_remove in need_remove_list:
  41. layout_dets.remove(need_remove)
  42. def __fix_by_remove_low_confidence(self):
  43. for model_page_info in self.__model_list:
  44. need_remove_list = []
  45. layout_dets = model_page_info['layout_dets']
  46. for layout_det in layout_dets:
  47. if layout_det['score'] <= 0.05:
  48. need_remove_list.append(layout_det)
  49. else:
  50. continue
  51. for need_remove in need_remove_list:
  52. layout_dets.remove(need_remove)
  53. def __fix_by_remove_high_iou_and_low_confidence(self):
  54. for model_page_info in self.__model_list:
  55. need_remove_list = []
  56. layout_dets = model_page_info['layout_dets']
  57. for layout_det1 in layout_dets:
  58. for layout_det2 in layout_dets:
  59. if layout_det1 == layout_det2:
  60. continue
  61. if layout_det1['category_id'] in [
  62. 0,
  63. 1,
  64. 2,
  65. 3,
  66. 4,
  67. 5,
  68. 6,
  69. 7,
  70. 8,
  71. 9,
  72. ] and layout_det2['category_id'] in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
  73. if (
  74. calculate_iou(layout_det1['bbox'], layout_det2['bbox'])
  75. > 0.9
  76. ):
  77. if layout_det1['score'] < layout_det2['score']:
  78. layout_det_need_remove = layout_det1
  79. else:
  80. layout_det_need_remove = layout_det2
  81. if layout_det_need_remove not in need_remove_list:
  82. need_remove_list.append(layout_det_need_remove)
  83. else:
  84. continue
  85. else:
  86. continue
  87. for need_remove in need_remove_list:
  88. layout_dets.remove(need_remove)
  89. def __init__(self, model_list: list, docs: fitz.Document):
  90. self.__model_list = model_list
  91. self.__docs = docs
  92. """为所有模型数据添加bbox信息(缩放,poly->bbox)"""
  93. self.__fix_axis()
  94. """删除置信度特别低的模型数据(<0.05),提高质量"""
  95. self.__fix_by_remove_low_confidence()
  96. """删除高iou(>0.9)数据中置信度较低的那个"""
  97. self.__fix_by_remove_high_iou_and_low_confidence()
  98. self.__fix_footnote()
  99. def __fix_footnote(self):
  100. # 3: figure, 5: table, 7: footnote
  101. for model_page_info in self.__model_list:
  102. footnotes = []
  103. figures = []
  104. tables = []
  105. for obj in model_page_info['layout_dets']:
  106. if obj['category_id'] == 7:
  107. footnotes.append(obj)
  108. elif obj['category_id'] == 3:
  109. figures.append(obj)
  110. elif obj['category_id'] == 5:
  111. tables.append(obj)
  112. if len(footnotes) * len(figures) == 0:
  113. continue
  114. dis_figure_footnote = {}
  115. dis_table_footnote = {}
  116. for i in range(len(footnotes)):
  117. for j in range(len(figures)):
  118. pos_flag_count = sum(
  119. list(
  120. map(
  121. lambda x: 1 if x else 0,
  122. bbox_relative_pos(
  123. footnotes[i]['bbox'], figures[j]['bbox']
  124. ),
  125. )
  126. )
  127. )
  128. if pos_flag_count > 1:
  129. continue
  130. dis_figure_footnote[i] = min(
  131. bbox_distance(figures[j]['bbox'], footnotes[i]['bbox']),
  132. dis_figure_footnote.get(i, float('inf')),
  133. )
  134. for i in range(len(footnotes)):
  135. for j in range(len(tables)):
  136. pos_flag_count = sum(
  137. list(
  138. map(
  139. lambda x: 1 if x else 0,
  140. bbox_relative_pos(
  141. footnotes[i]['bbox'], tables[j]['bbox']
  142. ),
  143. )
  144. )
  145. )
  146. if pos_flag_count > 1:
  147. continue
  148. dis_table_footnote[i] = min(
  149. bbox_distance(tables[j]['bbox'], footnotes[i]['bbox']),
  150. dis_table_footnote.get(i, float('inf')),
  151. )
  152. for i in range(len(footnotes)):
  153. if dis_table_footnote.get(i, float('inf')) > dis_figure_footnote[i]:
  154. footnotes[i]['category_id'] = CategoryId.ImageFootnote
  155. def __reduct_overlap(self, bboxes):
  156. N = len(bboxes)
  157. keep = [True] * N
  158. for i in range(N):
  159. for j in range(N):
  160. if i == j:
  161. continue
  162. if _is_in(bboxes[i]['bbox'], bboxes[j]['bbox']):
  163. keep[i] = False
  164. return [bboxes[i] for i in range(N) if keep[i]]
  165. def __tie_up_category_by_distance(
  166. self, page_no, subject_category_id, object_category_id
  167. ):
  168. """假定每个 subject 最多有一个 object (可以有多个相邻的 object 合并为单个 object),每个 object
  169. 只能属于一个 subject."""
  170. ret = []
  171. MAX_DIS_OF_POINT = 10**9 + 7
  172. """
  173. subject 和 object 的 bbox 会合并成一个大的 bbox (named: merged bbox)。
  174. 筛选出所有和 merged bbox 有 overlap 且 overlap 面积大于 object 的面积的 subjects。
  175. 再求出筛选出的 subjects 和 object 的最短距离
  176. """
  177. def may_find_other_nearest_bbox(subject_idx, object_idx):
  178. ret = float('inf')
  179. x0 = min(
  180. all_bboxes[subject_idx]['bbox'][0], all_bboxes[object_idx]['bbox'][0]
  181. )
  182. y0 = min(
  183. all_bboxes[subject_idx]['bbox'][1], all_bboxes[object_idx]['bbox'][1]
  184. )
  185. x1 = max(
  186. all_bboxes[subject_idx]['bbox'][2], all_bboxes[object_idx]['bbox'][2]
  187. )
  188. y1 = max(
  189. all_bboxes[subject_idx]['bbox'][3], all_bboxes[object_idx]['bbox'][3]
  190. )
  191. object_area = abs(
  192. all_bboxes[object_idx]['bbox'][2] - all_bboxes[object_idx]['bbox'][0]
  193. ) * abs(
  194. all_bboxes[object_idx]['bbox'][3] - all_bboxes[object_idx]['bbox'][1]
  195. )
  196. for i in range(len(all_bboxes)):
  197. if (
  198. i == subject_idx
  199. or all_bboxes[i]['category_id'] != subject_category_id
  200. ):
  201. continue
  202. if _is_part_overlap([x0, y0, x1, y1], all_bboxes[i]['bbox']) or _is_in(
  203. all_bboxes[i]['bbox'], [x0, y0, x1, y1]
  204. ):
  205. i_area = abs(
  206. all_bboxes[i]['bbox'][2] - all_bboxes[i]['bbox'][0]
  207. ) * abs(all_bboxes[i]['bbox'][3] - all_bboxes[i]['bbox'][1])
  208. if i_area >= object_area:
  209. ret = min(float('inf'), dis[i][object_idx])
  210. return ret
  211. def expand_bbbox(idxes):
  212. x0s = [all_bboxes[idx]['bbox'][0] for idx in idxes]
  213. y0s = [all_bboxes[idx]['bbox'][1] for idx in idxes]
  214. x1s = [all_bboxes[idx]['bbox'][2] for idx in idxes]
  215. y1s = [all_bboxes[idx]['bbox'][3] for idx in idxes]
  216. return min(x0s), min(y0s), max(x1s), max(y1s)
  217. subjects = self.__reduct_overlap(
  218. list(
  219. map(
  220. lambda x: {'bbox': x['bbox'], 'score': x['score']},
  221. filter(
  222. lambda x: x['category_id'] == subject_category_id,
  223. self.__model_list[page_no]['layout_dets'],
  224. ),
  225. )
  226. )
  227. )
  228. objects = self.__reduct_overlap(
  229. list(
  230. map(
  231. lambda x: {'bbox': x['bbox'], 'score': x['score']},
  232. filter(
  233. lambda x: x['category_id'] == object_category_id,
  234. self.__model_list[page_no]['layout_dets'],
  235. ),
  236. )
  237. )
  238. )
  239. subject_object_relation_map = {}
  240. subjects.sort(
  241. key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2
  242. ) # get the distance !
  243. all_bboxes = []
  244. for v in subjects:
  245. all_bboxes.append(
  246. {
  247. 'category_id': subject_category_id,
  248. 'bbox': v['bbox'],
  249. 'score': v['score'],
  250. }
  251. )
  252. for v in objects:
  253. all_bboxes.append(
  254. {
  255. 'category_id': object_category_id,
  256. 'bbox': v['bbox'],
  257. 'score': v['score'],
  258. }
  259. )
  260. N = len(all_bboxes)
  261. dis = [[MAX_DIS_OF_POINT] * N for _ in range(N)]
  262. for i in range(N):
  263. for j in range(i):
  264. if (
  265. all_bboxes[i]['category_id'] == subject_category_id
  266. and all_bboxes[j]['category_id'] == subject_category_id
  267. ):
  268. continue
  269. dis[i][j] = bbox_distance(all_bboxes[i]['bbox'], all_bboxes[j]['bbox'])
  270. dis[j][i] = dis[i][j]
  271. used = set()
  272. for i in range(N):
  273. # 求第 i 个 subject 所关联的 object
  274. if all_bboxes[i]['category_id'] != subject_category_id:
  275. continue
  276. seen = set()
  277. candidates = []
  278. arr = []
  279. for j in range(N):
  280. pos_flag_count = sum(
  281. list(
  282. map(
  283. lambda x: 1 if x else 0,
  284. bbox_relative_pos(
  285. all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
  286. ),
  287. )
  288. )
  289. )
  290. if pos_flag_count > 1:
  291. continue
  292. if (
  293. all_bboxes[j]['category_id'] != object_category_id
  294. or j in used
  295. or dis[i][j] == MAX_DIS_OF_POINT
  296. ):
  297. continue
  298. left, right, _, _ = bbox_relative_pos(
  299. all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
  300. ) # 由 pos_flag_count 相关逻辑保证本段逻辑准确性
  301. if left or right:
  302. one_way_dis = all_bboxes[i]['bbox'][2] - all_bboxes[i]['bbox'][0]
  303. else:
  304. one_way_dis = all_bboxes[i]['bbox'][3] - all_bboxes[i]['bbox'][1]
  305. if dis[i][j] > one_way_dis:
  306. continue
  307. arr.append((dis[i][j], j))
  308. arr.sort(key=lambda x: x[0])
  309. if len(arr) > 0:
  310. """
  311. bug: 离该subject 最近的 object 可能跨越了其它的 subject。
  312. 比如 [this subect] [some sbuject] [the nearest object of subject]
  313. """
  314. if may_find_other_nearest_bbox(i, arr[0][1]) >= arr[0][0]:
  315. candidates.append(arr[0][1])
  316. seen.add(arr[0][1])
  317. # 已经获取初始种子
  318. for j in set(candidates):
  319. tmp = []
  320. for k in range(i + 1, N):
  321. pos_flag_count = sum(
  322. list(
  323. map(
  324. lambda x: 1 if x else 0,
  325. bbox_relative_pos(
  326. all_bboxes[j]['bbox'], all_bboxes[k]['bbox']
  327. ),
  328. )
  329. )
  330. )
  331. if pos_flag_count > 1:
  332. continue
  333. if (
  334. all_bboxes[k]['category_id'] != object_category_id
  335. or k in used
  336. or k in seen
  337. or dis[j][k] == MAX_DIS_OF_POINT
  338. or dis[j][k] > dis[i][j]
  339. ):
  340. continue
  341. is_nearest = True
  342. for ni in range(i + 1, N):
  343. if ni in (j, k) or ni in used or ni in seen:
  344. continue
  345. if not float_gt(dis[ni][k], dis[j][k]):
  346. is_nearest = False
  347. break
  348. if is_nearest:
  349. nx0, ny0, nx1, ny1 = expand_bbbox(list(seen) + [k])
  350. n_dis = bbox_distance(
  351. all_bboxes[i]['bbox'], [nx0, ny0, nx1, ny1]
  352. )
  353. if float_gt(dis[i][j], n_dis):
  354. continue
  355. tmp.append(k)
  356. seen.add(k)
  357. candidates = tmp
  358. if len(candidates) == 0:
  359. break
  360. # 已经获取到某个 figure 下所有的最靠近的 captions,以及最靠近这些 captions 的 captions 。
  361. # 先扩一下 bbox,
  362. ox0, oy0, ox1, oy1 = expand_bbbox(list(seen) + [i])
  363. ix0, iy0, ix1, iy1 = all_bboxes[i]['bbox']
  364. # 分成了 4 个截取空间,需要计算落在每个截取空间下 objects 合并后占据的矩形面积
  365. caption_poses = [
  366. [ox0, oy0, ix0, oy1],
  367. [ox0, oy0, ox1, iy0],
  368. [ox0, iy1, ox1, oy1],
  369. [ix1, oy0, ox1, oy1],
  370. ]
  371. caption_areas = []
  372. for bbox in caption_poses:
  373. embed_arr = []
  374. for idx in seen:
  375. if (
  376. calculate_overlap_area_in_bbox1_area_ratio(
  377. all_bboxes[idx]['bbox'], bbox
  378. )
  379. > CAPATION_OVERLAP_AREA_RATIO
  380. ):
  381. embed_arr.append(idx)
  382. if len(embed_arr) > 0:
  383. embed_x0 = min([all_bboxes[idx]['bbox'][0] for idx in embed_arr])
  384. embed_y0 = min([all_bboxes[idx]['bbox'][1] for idx in embed_arr])
  385. embed_x1 = max([all_bboxes[idx]['bbox'][2] for idx in embed_arr])
  386. embed_y1 = max([all_bboxes[idx]['bbox'][3] for idx in embed_arr])
  387. caption_areas.append(
  388. int(abs(embed_x1 - embed_x0) * abs(embed_y1 - embed_y0))
  389. )
  390. else:
  391. caption_areas.append(0)
  392. subject_object_relation_map[i] = []
  393. if max(caption_areas) > 0:
  394. max_area_idx = caption_areas.index(max(caption_areas))
  395. caption_bbox = caption_poses[max_area_idx]
  396. for j in seen:
  397. if (
  398. calculate_overlap_area_in_bbox1_area_ratio(
  399. all_bboxes[j]['bbox'], caption_bbox
  400. )
  401. > CAPATION_OVERLAP_AREA_RATIO
  402. ):
  403. used.add(j)
  404. subject_object_relation_map[i].append(j)
  405. for i in sorted(subject_object_relation_map.keys()):
  406. result = {
  407. 'subject_body': all_bboxes[i]['bbox'],
  408. 'all': all_bboxes[i]['bbox'],
  409. 'score': all_bboxes[i]['score'],
  410. }
  411. if len(subject_object_relation_map[i]) > 0:
  412. x0 = min(
  413. [all_bboxes[j]['bbox'][0] for j in subject_object_relation_map[i]]
  414. )
  415. y0 = min(
  416. [all_bboxes[j]['bbox'][1] for j in subject_object_relation_map[i]]
  417. )
  418. x1 = max(
  419. [all_bboxes[j]['bbox'][2] for j in subject_object_relation_map[i]]
  420. )
  421. y1 = max(
  422. [all_bboxes[j]['bbox'][3] for j in subject_object_relation_map[i]]
  423. )
  424. result['object_body'] = [x0, y0, x1, y1]
  425. result['all'] = [
  426. min(x0, all_bboxes[i]['bbox'][0]),
  427. min(y0, all_bboxes[i]['bbox'][1]),
  428. max(x1, all_bboxes[i]['bbox'][2]),
  429. max(y1, all_bboxes[i]['bbox'][3]),
  430. ]
  431. ret.append(result)
  432. total_subject_object_dis = 0
  433. # 计算已经配对的 distance 距离
  434. for i in subject_object_relation_map.keys():
  435. for j in subject_object_relation_map[i]:
  436. total_subject_object_dis += bbox_distance(
  437. all_bboxes[i]['bbox'], all_bboxes[j]['bbox']
  438. )
  439. # 计算未匹配的 subject 和 object 的距离(非精确版)
  440. with_caption_subject = set(
  441. [
  442. key
  443. for key in subject_object_relation_map.keys()
  444. if len(subject_object_relation_map[i]) > 0
  445. ]
  446. )
  447. for i in range(N):
  448. if all_bboxes[i]['category_id'] != object_category_id or i in used:
  449. continue
  450. candidates = []
  451. for j in range(N):
  452. if (
  453. all_bboxes[j]['category_id'] != subject_category_id
  454. or j in with_caption_subject
  455. ):
  456. continue
  457. candidates.append((dis[i][j], j))
  458. if len(candidates) > 0:
  459. candidates.sort(key=lambda x: x[0])
  460. total_subject_object_dis += candidates[0][1]
  461. with_caption_subject.add(j)
  462. return ret, total_subject_object_dis
  463. def get_imgs(self, page_no: int):
  464. with_captions, _ = self.__tie_up_category_by_distance(page_no, 3, 4)
  465. with_footnotes, _ = self.__tie_up_category_by_distance(
  466. page_no, 3, CategoryId.ImageFootnote
  467. )
  468. ret = []
  469. N, M = len(with_captions), len(with_footnotes)
  470. assert N == M
  471. for i in range(N):
  472. record = {
  473. 'score': with_captions[i]['score'],
  474. 'img_caption_bbox': with_captions[i].get('object_body', None),
  475. 'img_body_bbox': with_captions[i]['subject_body'],
  476. 'img_footnote_bbox': with_footnotes[i].get('object_body', None),
  477. }
  478. x0 = min(with_captions[i]['all'][0], with_footnotes[i]['all'][0])
  479. y0 = min(with_captions[i]['all'][1], with_footnotes[i]['all'][1])
  480. x1 = max(with_captions[i]['all'][2], with_footnotes[i]['all'][2])
  481. y1 = max(with_captions[i]['all'][3], with_footnotes[i]['all'][3])
  482. record['bbox'] = [x0, y0, x1, y1]
  483. ret.append(record)
  484. return ret
  485. def get_tables(
  486. self, page_no: int
  487. ) -> list: # 3个坐标, caption, table主体,table-note
  488. with_captions, _ = self.__tie_up_category_by_distance(page_no, 5, 6)
  489. with_footnotes, _ = self.__tie_up_category_by_distance(page_no, 5, 7)
  490. ret = []
  491. N, M = len(with_captions), len(with_footnotes)
  492. assert N == M
  493. for i in range(N):
  494. record = {
  495. 'score': with_captions[i]['score'],
  496. 'table_caption_bbox': with_captions[i].get('object_body', None),
  497. 'table_body_bbox': with_captions[i]['subject_body'],
  498. 'table_footnote_bbox': with_footnotes[i].get('object_body', None),
  499. }
  500. x0 = min(with_captions[i]['all'][0], with_footnotes[i]['all'][0])
  501. y0 = min(with_captions[i]['all'][1], with_footnotes[i]['all'][1])
  502. x1 = max(with_captions[i]['all'][2], with_footnotes[i]['all'][2])
  503. y1 = max(with_captions[i]['all'][3], with_footnotes[i]['all'][3])
  504. record['bbox'] = [x0, y0, x1, y1]
  505. ret.append(record)
  506. return ret
  507. def get_equations(self, page_no: int) -> list: # 有坐标,也有字
  508. inline_equations = self.__get_blocks_by_type(
  509. ModelBlockTypeEnum.EMBEDDING.value, page_no, ['latex']
  510. )
  511. interline_equations = self.__get_blocks_by_type(
  512. ModelBlockTypeEnum.ISOLATED.value, page_no, ['latex']
  513. )
  514. interline_equations_blocks = self.__get_blocks_by_type(
  515. ModelBlockTypeEnum.ISOLATE_FORMULA.value, page_no
  516. )
  517. return inline_equations, interline_equations, interline_equations_blocks
  518. def get_discarded(self, page_no: int) -> list: # 自研模型,只有坐标
  519. blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.ABANDON.value, page_no)
  520. return blocks
  521. def get_text_blocks(self, page_no: int) -> list: # 自研模型搞的,只有坐标,没有字
  522. blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.PLAIN_TEXT.value, page_no)
  523. return blocks
  524. def get_title_blocks(self, page_no: int) -> list: # 自研模型,只有坐标,没字
  525. blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.TITLE.value, page_no)
  526. return blocks
  527. def get_ocr_text(self, page_no: int) -> list: # paddle 搞的,有字也有坐标
  528. text_spans = []
  529. model_page_info = self.__model_list[page_no]
  530. layout_dets = model_page_info['layout_dets']
  531. for layout_det in layout_dets:
  532. if layout_det['category_id'] == '15':
  533. span = {
  534. 'bbox': layout_det['bbox'],
  535. 'content': layout_det['text'],
  536. }
  537. text_spans.append(span)
  538. return text_spans
  539. def get_all_spans(self, page_no: int) -> list:
  540. def remove_duplicate_spans(spans):
  541. new_spans = []
  542. for span in spans:
  543. if not any(span == existing_span for existing_span in new_spans):
  544. new_spans.append(span)
  545. return new_spans
  546. all_spans = []
  547. model_page_info = self.__model_list[page_no]
  548. layout_dets = model_page_info['layout_dets']
  549. allow_category_id_list = [3, 5, 13, 14, 15]
  550. """当成span拼接的"""
  551. # 3: 'image', # 图片
  552. # 5: 'table', # 表格
  553. # 13: 'inline_equation', # 行内公式
  554. # 14: 'interline_equation', # 行间公式
  555. # 15: 'text', # ocr识别文本
  556. for layout_det in layout_dets:
  557. category_id = layout_det['category_id']
  558. if category_id in allow_category_id_list:
  559. span = {'bbox': layout_det['bbox'], 'score': layout_det['score']}
  560. if category_id == 3:
  561. span['type'] = ContentType.Image
  562. elif category_id == 5:
  563. # 获取table模型结果
  564. latex = layout_det.get("latex", None)
  565. html = layout_det.get("html", None)
  566. if latex:
  567. span["latex"] = latex
  568. elif html:
  569. span["html"] = html
  570. span["type"] = ContentType.Table
  571. elif category_id == 13:
  572. span['content'] = layout_det['latex']
  573. span['type'] = ContentType.InlineEquation
  574. elif category_id == 14:
  575. span['content'] = layout_det['latex']
  576. span['type'] = ContentType.InterlineEquation
  577. elif category_id == 15:
  578. span['content'] = layout_det['text']
  579. span['type'] = ContentType.Text
  580. all_spans.append(span)
  581. return remove_duplicate_spans(all_spans)
  582. def get_page_size(self, page_no: int): # 获取页面宽高
  583. # 获取当前页的page对象
  584. page = self.__docs[page_no]
  585. # 获取当前页的宽高
  586. page_w = page.rect.width
  587. page_h = page.rect.height
  588. return page_w, page_h
  589. def __get_blocks_by_type(
  590. self, type: int, page_no: int, extra_col: list[str] = []
  591. ) -> list:
  592. blocks = []
  593. for page_dict in self.__model_list:
  594. layout_dets = page_dict.get('layout_dets', [])
  595. page_info = page_dict.get('page_info', {})
  596. page_number = page_info.get('page_no', -1)
  597. if page_no != page_number:
  598. continue
  599. for item in layout_dets:
  600. category_id = item.get('category_id', -1)
  601. bbox = item.get('bbox', None)
  602. if category_id == type:
  603. block = {
  604. 'bbox': bbox,
  605. 'score': item.get('score'),
  606. }
  607. for col in extra_col:
  608. block[col] = item.get(col, None)
  609. blocks.append(block)
  610. return blocks
  611. def get_model_list(self, page_no):
  612. return self.__model_list[page_no]
  613. if __name__ == '__main__':
  614. drw = DiskReaderWriter(r'D:/project/20231108code-clean')
  615. if 0:
  616. pdf_file_path = r'linshixuqiu\19983-00.pdf'
  617. model_file_path = r'linshixuqiu\19983-00_new.json'
  618. pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
  619. model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
  620. model_list = json.loads(model_json_txt)
  621. write_path = r'D:\project\20231108code-clean\linshixuqiu\19983-00'
  622. img_bucket_path = 'imgs'
  623. img_writer = DiskReaderWriter(join_path(write_path, img_bucket_path))
  624. pdf_docs = fitz.open('pdf', pdf_bytes)
  625. magic_model = MagicModel(model_list, pdf_docs)
  626. if 1:
  627. model_list = json.loads(
  628. drw.read('/opt/data/pdf/20240418/j.chroma.2009.03.042.json')
  629. )
  630. pdf_bytes = drw.read(
  631. '/opt/data/pdf/20240418/j.chroma.2009.03.042.pdf', AbsReaderWriter.MODE_BIN
  632. )
  633. pdf_docs = fitz.open('pdf', pdf_bytes)
  634. magic_model = MagicModel(model_list, pdf_docs)
  635. for i in range(7):
  636. print(magic_model.get_imgs(i))