ocr_dict_merge.py 3.2 KB

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