| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697 |
- from magic_pdf.config.drop_tag import DropTag
- from magic_pdf.config.ocr_content_type import BlockType
- from magic_pdf.libs.boxbase import calculate_iou, get_minbox_if_overlap_by_ratio
- def remove_overlaps_low_confidence_spans(spans):
- dropped_spans = []
- # 删除重叠spans中置信度低的的那些
- for span1 in spans:
- for span2 in spans:
- if span1 != span2:
- # span1 或 span2 任何一个都不应该在 dropped_spans 中
- if span1 in dropped_spans or span2 in dropped_spans:
- continue
- else:
- if calculate_iou(span1['bbox'], span2['bbox']) > 0.9:
- if span1['score'] < span2['score']:
- span_need_remove = span1
- else:
- span_need_remove = span2
- if (
- span_need_remove is not None
- and span_need_remove not in dropped_spans
- ):
- dropped_spans.append(span_need_remove)
- if len(dropped_spans) > 0:
- for span_need_remove in dropped_spans:
- spans.remove(span_need_remove)
- span_need_remove['tag'] = DropTag.SPAN_OVERLAP
- return spans, dropped_spans
- def remove_overlaps_chars(chars):
- dropped_chars = []
- # 删除重叠的char
- for char1 in chars:
- for char2 in chars:
- if char1 != char2:
- # char1 或 char2 任何一个都不应该在 dropped_chars 中
- if char1 in dropped_chars or char2 in dropped_chars:
- continue
- else:
- if calculate_iou(char1['bbox'], char2['bbox']) > 0.95:
- char_need_remove = char1
- if (
- char_need_remove is not None
- and char_need_remove not in dropped_chars
- ):
- dropped_chars.append(char_need_remove)
- if len(dropped_chars) > 0:
- for char_need_remove in dropped_chars:
- chars.remove(char_need_remove)
- return chars
- def remove_overlaps_min_spans(spans):
- dropped_spans = []
- # 删除重叠spans中较小的那些
- for span1 in spans:
- for span2 in spans:
- if span1 != span2:
- # span1 或 span2 任何一个都不应该在 dropped_spans 中
- if span1 in dropped_spans or span2 in dropped_spans:
- continue
- else:
- overlap_box = get_minbox_if_overlap_by_ratio(span1['bbox'], span2['bbox'], 0.65)
- if overlap_box is not None:
- span_need_remove = next((span for span in spans if span['bbox'] == overlap_box), None)
- if span_need_remove is not None and span_need_remove not in dropped_spans:
- dropped_spans.append(span_need_remove)
- if len(dropped_spans) > 0:
- for span_need_remove in dropped_spans:
- spans.remove(span_need_remove)
- span_need_remove['tag'] = DropTag.SPAN_OVERLAP
- return spans, dropped_spans
- def get_qa_need_list_v2(blocks):
- # 创建 images, tables, interline_equations, inline_equations 的副本
- images = []
- tables = []
- interline_equations = []
- for block in blocks:
- if block['type'] == BlockType.Image:
- images.append(block)
- elif block['type'] == BlockType.Table:
- tables.append(block)
- elif block['type'] == BlockType.InterlineEquation:
- interline_equations.append(block)
- return images, tables, interline_equations
|