ocr_dict_merge.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205
  1. from loguru import logger
  2. from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold, get_minbox_if_overlap_by_ratio, \
  3. calculate_overlap_area_in_bbox1_area_ratio
  4. from magic_pdf.libs.drop_tag import DropTag
  5. from magic_pdf.libs.ocr_content_type import ContentType, BlockType
  6. from magic_pdf.pre_proc.ocr_fix_block_logic import fix_image_block, fix_table_block, fix_text_block
  7. from magic_pdf.pre_proc.ocr_span_list_modify import modify_y_axis, modify_inline_equation
  8. from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
  9. # 将每一个line中的span从左到右排序
  10. def line_sort_spans_by_left_to_right(lines):
  11. line_objects = []
  12. for line in lines:
  13. # 按照x0坐标排序
  14. line.sort(key=lambda span: span['bbox'][0])
  15. line_bbox = [
  16. min(span['bbox'][0] for span in line), # x0
  17. min(span['bbox'][1] for span in line), # y0
  18. max(span['bbox'][2] for span in line), # x1
  19. max(span['bbox'][3] for span in line), # y1
  20. ]
  21. line_objects.append({
  22. "bbox": line_bbox,
  23. "spans": line,
  24. })
  25. return line_objects
  26. def merge_spans_to_line(spans):
  27. if len(spans) == 0:
  28. return []
  29. else:
  30. # 按照y0坐标排序
  31. spans.sort(key=lambda span: span['bbox'][1])
  32. lines = []
  33. current_line = [spans[0]]
  34. for span in spans[1:]:
  35. # 如果当前的span类型为"interline_equation" 或者 当前行中已经有"interline_equation"
  36. # image和table类型,同上
  37. if span['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] or any(
  38. s['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] for s in
  39. current_line):
  40. # 则开始新行
  41. lines.append(current_line)
  42. current_line = [span]
  43. continue
  44. # 如果当前的span与当前行的最后一个span在y轴上重叠,则添加到当前行
  45. if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
  46. current_line.append(span)
  47. else:
  48. # 否则,开始新行
  49. lines.append(current_line)
  50. current_line = [span]
  51. # 添加最后一行
  52. if current_line:
  53. lines.append(current_line)
  54. return lines
  55. def merge_spans_to_line_by_layout(spans, layout_bboxes):
  56. lines = []
  57. new_spans = []
  58. dropped_spans = []
  59. for item in layout_bboxes:
  60. layout_bbox = item['layout_bbox']
  61. # 遍历spans,将每个span放入对应的layout中
  62. layout_sapns = []
  63. for span in spans:
  64. if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], layout_bbox) > 0.6:
  65. layout_sapns.append(span)
  66. # 如果layout_sapns不为空,则放入new_spans中
  67. if len(layout_sapns) > 0:
  68. new_spans.append(layout_sapns)
  69. # 从spans删除已经放入layout_sapns中的span
  70. for layout_sapn in layout_sapns:
  71. spans.remove(layout_sapn)
  72. if len(new_spans) > 0:
  73. for layout_sapns in new_spans:
  74. layout_lines = merge_spans_to_line(layout_sapns)
  75. lines.extend(layout_lines)
  76. # 对line中的span进行排序
  77. lines = line_sort_spans_by_left_to_right(lines)
  78. for span in spans:
  79. span['tag'] = DropTag.NOT_IN_LAYOUT
  80. dropped_spans.append(span)
  81. return lines, dropped_spans
  82. def merge_lines_to_block(lines):
  83. # 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
  84. blocks = []
  85. for line in lines:
  86. blocks.append(
  87. {
  88. "bbox": line["bbox"],
  89. "lines": [line],
  90. }
  91. )
  92. return blocks
  93. def sort_blocks_by_layout(all_bboxes, layout_bboxes):
  94. new_blocks = []
  95. sort_blocks = []
  96. for item in layout_bboxes:
  97. layout_bbox = item['layout_bbox']
  98. # 遍历blocks,将每个blocks放入对应的layout中
  99. layout_blocks = []
  100. for block in all_bboxes:
  101. # 如果是footnote则跳过
  102. if block[7] == BlockType.Footnote:
  103. continue
  104. block_bbox = [block[0], block[1], block[2], block[3]]
  105. if calculate_overlap_area_in_bbox1_area_ratio(block_bbox, layout_bbox) > 0.8:
  106. layout_blocks.append(block)
  107. # 如果layout_blocks不为空,则放入new_blocks中
  108. if len(layout_blocks) > 0:
  109. new_blocks.append(layout_blocks)
  110. # 从spans删除已经放入layout_sapns中的span
  111. for layout_block in layout_blocks:
  112. all_bboxes.remove(layout_block)
  113. # 如果new_blocks不为空,则对new_blocks中每个block进行排序
  114. if len(new_blocks) > 0:
  115. for bboxes_in_layout_block in new_blocks:
  116. bboxes_in_layout_block.sort(key=lambda x: x[1]) # 一个layout内部的box,按照y0自上而下排序
  117. sort_blocks.extend(bboxes_in_layout_block)
  118. # sort_blocks中已经包含了当前页面所有最终留下的block,且已经排好了顺序
  119. return sort_blocks
  120. def fill_spans_in_blocks(blocks, spans):
  121. '''
  122. 将allspans中的span按位置关系,放入blocks中
  123. '''
  124. block_with_spans = []
  125. for block in blocks:
  126. block_type = block[7]
  127. block_bbox = block[0:4]
  128. block_dict = {
  129. 'block_type': block_type,
  130. 'bbox': block_bbox,
  131. }
  132. block_spans = []
  133. for span in spans:
  134. span_bbox = span['bbox']
  135. if calculate_overlap_area_in_bbox1_area_ratio(span_bbox, block_bbox) > 0.8:
  136. block_spans.append(span)
  137. '''行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)'''
  138. displayed_list = []
  139. text_inline_lines = []
  140. modify_y_axis(block_spans, displayed_list, text_inline_lines)
  141. '''模型识别错误的行间公式, type类型转换成行内公式'''
  142. block_spans = modify_inline_equation(block_spans, displayed_list, text_inline_lines)
  143. '''bbox去除粘连'''
  144. block_spans = remove_overlap_between_bbox(block_spans)
  145. block_dict['spans'] = block_spans
  146. block_with_spans.append(block_dict)
  147. # 从spans删除已经放入block_spans中的span
  148. if len(block_spans) > 0:
  149. for span in block_spans:
  150. spans.remove(span)
  151. return block_with_spans
  152. def fix_block_spans(block_with_spans, img_blocks, table_blocks):
  153. '''
  154. 1、img_block和table_block因为包含caption和footnote的关系,存在block的嵌套关系
  155. 需要将caption和footnote的text_span放入相应img_block和table_block内的
  156. caption_block和footnote_block中
  157. 2、同时需要删除block中的spans字段
  158. '''
  159. fix_blocks = []
  160. for block in block_with_spans:
  161. block_type = block['block_type']
  162. if block_type == BlockType.Image:
  163. block = fix_image_block(block, img_blocks)
  164. elif block_type == BlockType.Table:
  165. block = fix_table_block(block, table_blocks)
  166. elif block_type in [BlockType.Text, BlockType.Title, BlockType.InterlineEquation]:
  167. block = fix_text_block(block)
  168. else:
  169. continue
  170. fix_blocks.append(block)
  171. return fix_blocks