ocr_span_list_modify.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. from loguru import logger
  2. from magic_pdf.libs.boxbase import calculate_overlap_area_in_bbox1_area_ratio, get_minbox_if_overlap_by_ratio, \
  3. __is_overlaps_y_exceeds_threshold
  4. from magic_pdf.libs.drop_tag import DropTag
  5. from magic_pdf.libs.ocr_content_type import ContentType
  6. def remove_overlaps_min_spans(spans):
  7. dropped_spans = []
  8. # 删除重叠spans中较小的那些
  9. for span1 in spans.copy():
  10. for span2 in spans.copy():
  11. if span1 != span2:
  12. overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.65)
  13. if overlap_box is not None:
  14. bbox_to_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
  15. if bbox_to_remove is not None:
  16. spans.remove(bbox_to_remove)
  17. bbox_to_remove['tag'] = DropTag.SPAN_OVERLAP
  18. dropped_spans.append(bbox_to_remove)
  19. return spans, dropped_spans
  20. def remove_spans_by_bboxes(spans, need_remove_spans_bboxes):
  21. # 遍历spans, 判断是否在removed_span_block_bboxes中
  22. # 如果是, 则删除该span 否则, 保留该span
  23. need_remove_spans = []
  24. for span in spans:
  25. for removed_bbox in need_remove_spans_bboxes:
  26. if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], removed_bbox) > 0.5:
  27. need_remove_spans.append(span)
  28. break
  29. for span in need_remove_spans:
  30. spans.remove(span)
  31. return spans
  32. def remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict):
  33. dropped_spans = []
  34. for drop_tag, removed_bboxes in need_remove_spans_bboxes_dict.items():
  35. # logger.info(f"remove spans by bbox dict, drop_tag: {drop_tag}, removed_bboxes: {removed_bboxes}")
  36. need_remove_spans = []
  37. for span in spans:
  38. for removed_bbox in removed_bboxes:
  39. if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], removed_bbox) > 0.5:
  40. need_remove_spans.append(span)
  41. break
  42. for span in need_remove_spans:
  43. spans.remove(span)
  44. span['tag'] = drop_tag
  45. dropped_spans.append(span)
  46. return spans, dropped_spans
  47. def adjust_bbox_for_standalone_block(spans):
  48. # 对tpye=["interline_equation", "image", "table"]进行额外处理,如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
  49. for sb_span in spans:
  50. if sb_span['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table]:
  51. for text_span in spans:
  52. if text_span['type'] in [ContentType.Text, ContentType.InlineEquation]:
  53. # 判断span2的纵向高度是否被span所覆盖
  54. if sb_span['bbox'][1] < text_span['bbox'][1] and sb_span['bbox'][3] > text_span['bbox'][3]:
  55. # 判断span2是否在span左边
  56. if text_span['bbox'][0] < sb_span['bbox'][0]:
  57. # 调整span的y0和span2的y0一致
  58. sb_span['bbox'][1] = text_span['bbox'][1]
  59. return spans
  60. def modify_y_axis(spans: list, displayed_list: list, text_inline_lines: list):
  61. # displayed_list = []
  62. # 如果spans为空,则不处理
  63. if len(spans) == 0:
  64. pass
  65. else:
  66. spans.sort(key=lambda span: span['bbox'][1])
  67. lines = []
  68. current_line = [spans[0]]
  69. if spans[0]["type"] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table]:
  70. displayed_list.append(spans[0])
  71. line_first_y0 = spans[0]["bbox"][1]
  72. line_first_y = spans[0]["bbox"][3]
  73. # 用于给行间公式搜索
  74. # text_inline_lines = []
  75. for span in spans[1:]:
  76. # if span.get("content","") == "78.":
  77. # print("debug")
  78. # 如果当前的span类型为"interline_equation" 或者 当前行中已经有"interline_equation"
  79. # image和table类型,同上
  80. if span['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] or any(
  81. s['type'] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table] for s in
  82. current_line):
  83. # 传入
  84. if span["type"] in [ContentType.InterlineEquation, ContentType.Image, ContentType.Table]:
  85. displayed_list.append(span)
  86. # 则开始新行
  87. lines.append(current_line)
  88. if len(current_line) > 1 or current_line[0]["type"] in [ContentType.Text, ContentType.InlineEquation]:
  89. text_inline_lines.append((current_line, (line_first_y0, line_first_y)))
  90. current_line = [span]
  91. line_first_y0 = span["bbox"][1]
  92. line_first_y = span["bbox"][3]
  93. continue
  94. # 如果当前的span与当前行的最后一个span在y轴上重叠,则添加到当前行
  95. if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
  96. if span["type"] == "text":
  97. line_first_y0 = span["bbox"][1]
  98. line_first_y = span["bbox"][3]
  99. current_line.append(span)
  100. else:
  101. # 否则,开始新行
  102. lines.append(current_line)
  103. text_inline_lines.append((current_line, (line_first_y0, line_first_y)))
  104. current_line = [span]
  105. line_first_y0 = span["bbox"][1]
  106. line_first_y = span["bbox"][3]
  107. # 添加最后一行
  108. if current_line:
  109. lines.append(current_line)
  110. if len(current_line) > 1 or current_line[0]["type"] in [ContentType.Text, ContentType.InlineEquation]:
  111. text_inline_lines.append((current_line, (line_first_y0, line_first_y)))
  112. for line in text_inline_lines:
  113. # 按照x0坐标排序
  114. current_line = line[0]
  115. current_line.sort(key=lambda span: span['bbox'][0])
  116. # 调整每一个文字行内bbox统一
  117. for line in text_inline_lines:
  118. current_line, (line_first_y0, line_first_y) = line
  119. for span in current_line:
  120. span["bbox"][1] = line_first_y0
  121. span["bbox"][3] = line_first_y
  122. # return spans, displayed_list, text_inline_lines
  123. def modify_inline_equation(spans: list, displayed_list: list, text_inline_lines: list):
  124. # 错误行间公式转行内公式
  125. j = 0
  126. for i in range(len(displayed_list)):
  127. # if i == 8:
  128. # print("debug")
  129. span = displayed_list[i]
  130. span_y0, span_y = span["bbox"][1], span["bbox"][3]
  131. while j < len(text_inline_lines):
  132. text_line = text_inline_lines[j]
  133. y0, y1 = text_line[1]
  134. if (
  135. span_y0 < y0 and span_y > y0 or span_y0 < y1 and span_y > y1 or span_y0 < y0 and span_y > y1) and __is_overlaps_y_exceeds_threshold(
  136. span['bbox'], (0, y0, 0, y1)):
  137. # 调整公式类型
  138. if span["type"] == ContentType.InterlineEquation:
  139. # 最后一行是行间公式
  140. if j + 1 >= len(text_inline_lines):
  141. span["type"] = ContentType.InlineEquation
  142. span["bbox"][1] = y0
  143. span["bbox"][3] = y1
  144. else:
  145. # 行间公式旁边有多行文字或者行间公式比文字高3倍则不转换
  146. y0_next, y1_next = text_inline_lines[j + 1][1]
  147. if not __is_overlaps_y_exceeds_threshold(span['bbox'], (0, y0_next, 0, y1_next)) and 3 * (
  148. y1 - y0) > span_y - span_y0:
  149. span["type"] = ContentType.InlineEquation
  150. span["bbox"][1] = y0
  151. span["bbox"][3] = y1
  152. break
  153. elif span_y < y0 or span_y0 < y0 and span_y > y0 and not __is_overlaps_y_exceeds_threshold(span['bbox'],
  154. (0, y0, 0, y1)):
  155. break
  156. else:
  157. j += 1
  158. return spans
  159. def get_qa_need_list(blocks):
  160. # 创建 images, tables, interline_equations, inline_equations 的副本
  161. images = []
  162. tables = []
  163. interline_equations = []
  164. inline_equations = []
  165. for block in blocks:
  166. for line in block["lines"]:
  167. for span in line["spans"]:
  168. if span["type"] == ContentType.Image:
  169. images.append(span)
  170. elif span["type"] == ContentType.Table:
  171. tables.append(span)
  172. elif span["type"] == ContentType.InlineEquation:
  173. inline_equations.append(span)
  174. elif span["type"] == ContentType.InterlineEquation:
  175. interline_equations.append(span)
  176. else:
  177. continue
  178. return images, tables, interline_equations, inline_equations