ocr_dict_merge.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold, get_minbox_if_overlap_by_ratio
  2. # 删除重叠spans中较小的那些
  3. def remove_overlaps_min_spans(spans):
  4. for span1 in spans.copy():
  5. for span2 in spans.copy():
  6. if span1 != span2:
  7. overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.8)
  8. if overlap_box is not None:
  9. bbox_to_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
  10. if bbox_to_remove is not None:
  11. spans.remove(bbox_to_remove)
  12. return spans
  13. def merge_spans_to_line(spans):
  14. # 按照y0坐标排序
  15. spans.sort(key=lambda span: span['bbox'][1])
  16. lines = []
  17. current_line = [spans[0]]
  18. for span in spans[1:]:
  19. # 如果当前的span类型为"displayed_equation" 或者 当前行中已经有"displayed_equation"
  20. # image和table类型,同上
  21. if span['type'] in ["displayed_equation", "image", "table"] or any(s['type'] in ["displayed_equation", "image", "table"] for s in current_line):
  22. # 则开始新行
  23. lines.append(current_line)
  24. current_line = [span]
  25. continue
  26. # 如果当前的span与当前行的最后一个span在y轴上重叠,则添加到当前行
  27. if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
  28. current_line.append(span)
  29. else:
  30. # 否则,开始新行
  31. lines.append(current_line)
  32. current_line = [span]
  33. # 添加最后一行
  34. if current_line:
  35. lines.append(current_line)
  36. # 计算每行的边界框,并对每行中的span按照x0进行排序
  37. line_objects = []
  38. for line in lines:
  39. # 按照x0坐标排序
  40. line.sort(key=lambda span: span['bbox'][0])
  41. line_bbox = [
  42. min(span['bbox'][0] for span in line), # x0
  43. min(span['bbox'][1] for span in line), # y0
  44. max(span['bbox'][2] for span in line), # x1
  45. max(span['bbox'][3] for span in line), # y1
  46. ]
  47. line_objects.append({
  48. "bbox": line_bbox,
  49. "spans": line,
  50. })
  51. return line_objects