ocr_span_list_modify.py 3.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. from magic_pdf.config.drop_tag import DropTag
  2. from magic_pdf.config.ocr_content_type import BlockType
  3. from magic_pdf.libs.boxbase import calculate_iou, get_minbox_if_overlap_by_ratio
  4. def remove_overlaps_low_confidence_spans(spans):
  5. dropped_spans = []
  6. # 删除重叠spans中置信度低的的那些
  7. for span1 in spans:
  8. for span2 in spans:
  9. if span1 != span2:
  10. # span1 或 span2 任何一个都不应该在 dropped_spans 中
  11. if span1 in dropped_spans or span2 in dropped_spans:
  12. continue
  13. else:
  14. if calculate_iou(span1['bbox'], span2['bbox']) > 0.9:
  15. if span1['score'] < span2['score']:
  16. span_need_remove = span1
  17. else:
  18. span_need_remove = span2
  19. if (
  20. span_need_remove is not None
  21. and span_need_remove not in dropped_spans
  22. ):
  23. dropped_spans.append(span_need_remove)
  24. if len(dropped_spans) > 0:
  25. for span_need_remove in dropped_spans:
  26. spans.remove(span_need_remove)
  27. span_need_remove['tag'] = DropTag.SPAN_OVERLAP
  28. return spans, dropped_spans
  29. def remove_overlaps_chars(chars):
  30. dropped_chars = []
  31. # 删除重叠的char
  32. for char1 in chars:
  33. for char2 in chars:
  34. if char1 != char2:
  35. # char1 或 char2 任何一个都不应该在 dropped_chars 中
  36. if char1 in dropped_chars or char2 in dropped_chars:
  37. continue
  38. else:
  39. if calculate_iou(char1['bbox'], char2['bbox']) > 0.95:
  40. char_need_remove = char1
  41. if (
  42. char_need_remove is not None
  43. and char_need_remove not in dropped_chars
  44. ):
  45. dropped_chars.append(char_need_remove)
  46. if len(dropped_chars) > 0:
  47. for char_need_remove in dropped_chars:
  48. chars.remove(char_need_remove)
  49. return chars
  50. def remove_overlaps_min_spans(spans):
  51. dropped_spans = []
  52. # 删除重叠spans中较小的那些
  53. for span1 in spans:
  54. for span2 in spans:
  55. if span1 != span2:
  56. # span1 或 span2 任何一个都不应该在 dropped_spans 中
  57. if span1 in dropped_spans or span2 in dropped_spans:
  58. continue
  59. else:
  60. overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.65)
  61. if overlap_box is not None:
  62. span_need_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
  63. if span_need_remove is not None and span_need_remove not in dropped_spans:
  64. dropped_spans.append(span_need_remove)
  65. if len(dropped_spans) > 0:
  66. for span_need_remove in dropped_spans:
  67. spans.remove(span_need_remove)
  68. span_need_remove['tag'] = DropTag.SPAN_OVERLAP
  69. return spans, dropped_spans
  70. def get_qa_need_list_v2(blocks):
  71. # 创建 images, tables, interline_equations, inline_equations 的副本
  72. images = []
  73. tables = []
  74. interline_equations = []
  75. for block in blocks:
  76. if block['type'] == BlockType.Image:
  77. images.append(block)
  78. elif block['type'] == BlockType.Table:
  79. tables.append(block)
  80. elif block['type'] == BlockType.InterlineEquation:
  81. interline_equations.append(block)
  82. return images, tables, interline_equations