ocr_dict_merge.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  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. # 将每一个line中的span从左到右排序
  5. def line_sort_spans_by_left_to_right(lines):
  6. line_objects = []
  7. for line in lines:
  8. # 按照x0坐标排序
  9. line.sort(key=lambda span: span['bbox'][0])
  10. line_bbox = [
  11. min(span['bbox'][0] for span in line), # x0
  12. min(span['bbox'][1] for span in line), # y0
  13. max(span['bbox'][2] for span in line), # x1
  14. max(span['bbox'][3] for span in line), # y1
  15. ]
  16. line_objects.append({
  17. "bbox": line_bbox,
  18. "spans": line,
  19. })
  20. return line_objects
  21. def merge_spans_to_line(spans):
  22. # 按照y0坐标排序
  23. spans.sort(key=lambda span: span['bbox'][1])
  24. lines = []
  25. current_line = [spans[0]]
  26. for span in spans[1:]:
  27. # 如果当前的span类型为"displayed_equation" 或者 当前行中已经有"displayed_equation"
  28. # image和table类型,同上
  29. if span['type'] in ["displayed_equation", "image", "table"] or any(
  30. s['type'] in ["displayed_equation", "image", "table"] for s in current_line):
  31. # 则开始新行
  32. lines.append(current_line)
  33. current_line = [span]
  34. continue
  35. # 如果当前的span与当前行的最后一个span在y轴上重叠,则添加到当前行
  36. if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
  37. current_line.append(span)
  38. else:
  39. # 否则,开始新行
  40. lines.append(current_line)
  41. current_line = [span]
  42. # 添加最后一行
  43. if current_line:
  44. lines.append(current_line)
  45. return lines
  46. def merge_spans_to_line_by_layout(spans, layout_bboxes):
  47. lines = []
  48. new_spans = []
  49. for item in layout_bboxes:
  50. layout_bbox = item['layout_bbox']
  51. # 遍历spans,将每个span放入对应的layout中
  52. layout_sapns = []
  53. for span in spans:
  54. if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], layout_bbox) > 0.8:
  55. layout_sapns.append(span)
  56. # 如果layout_sapns不为空,则放入new_spans中
  57. if len(layout_sapns) > 0:
  58. new_spans.append(layout_sapns)
  59. # 从spans删除已经放入layout_sapns中的span
  60. for layout_sapn in layout_sapns:
  61. spans.remove(layout_sapn)
  62. if len(new_spans) > 0:
  63. for layout_sapns in new_spans:
  64. layout_lines = merge_spans_to_line(layout_sapns)
  65. lines.extend(layout_lines)
  66. #对line中的span进行排序
  67. lines = line_sort_spans_by_left_to_right(lines)
  68. return lines