ocr_dict_merge.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold
  2. def merge_spans(spans):
  3. # 按照y0坐标排序
  4. spans.sort(key=lambda span: span['bbox'][1])
  5. lines = []
  6. current_line = [spans[0]]
  7. for span in spans[1:]:
  8. # 如果当前的span类型为"displayed_equation" 或者 当前行中已经有"displayed_equation"
  9. if span['type'] == "displayed_equation" or any(s['type'] == "displayed_equation" for s in current_line):
  10. # 则开始新行
  11. lines.append(current_line)
  12. current_line = [span]
  13. continue
  14. # 如果当前的span与当前行的最后一个span在y轴上重叠,则添加到当前行
  15. if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox']):
  16. current_line.append(span)
  17. else:
  18. # 否则,开始新行
  19. lines.append(current_line)
  20. current_line = [span]
  21. # 添加最后一行
  22. if current_line:
  23. lines.append(current_line)
  24. # 计算每行的边界框,并对每行中的span按照x0进行排序
  25. line_objects = []
  26. for line in lines:
  27. # 按照x0坐标排序
  28. line.sort(key=lambda span: span['bbox'][0])
  29. line_bbox = [
  30. min(span['bbox'][0] for span in line), # x0
  31. min(span['bbox'][1] for span in line), # y0
  32. max(span['bbox'][2] for span in line), # x1
  33. max(span['bbox'][3] for span in line), # y1
  34. ]
  35. line_objects.append({
  36. "bbox": line_bbox,
  37. "spans": line,
  38. })
  39. return line_objects