MinerU/magic_pdf/pre_proc/ocr_dict_merge.py

from magic_pdf.config.drop_tag import DropTag
from magic_pdf.config.ocr_content_type import BlockType, ContentType
from magic_pdf.libs.boxbase import (__is_overlaps_y_exceeds_threshold,
                                    _is_in_or_part_overlap_with_area_ratio,
                                    calculate_overlap_area_in_bbox1_area_ratio)


# 将每一个line中的span从左到右排序
def line_sort_spans_by_left_to_right(lines):
    line_objects = []
    for line in lines:
        #  按照x0坐标排序
        line.sort(key=lambda span: span['bbox'][0])
        line_bbox = [
            min(span['bbox'][0] for span in line),  # x0
            min(span['bbox'][1] for span in line),  # y0
            max(span['bbox'][2] for span in line),  # x1
            max(span['bbox'][3] for span in line),  # y1
        ]
        line_objects.append({
            'bbox': line_bbox,
            'spans': line,
        })
    return line_objects


def merge_spans_to_line(spans):
    if len(spans) == 0:
        return []
    else:
        # 按照y0坐标排序
        spans.sort(key=lambda span: span['bbox'][1])

        lines = []
        current_line = [spans[0]]
        for span in spans[1:]:
            # 如果当前的span类型为"interline_equation" 或者 当前行中已经有"interline_equation"
            # image和table类型，同上
            if span['type'] in [
                    ContentType.InterlineEquation, ContentType.Image,
                    ContentType.Table
            ] or any(s['type'] in [
                    ContentType.InterlineEquation, ContentType.Image,
                    ContentType.Table
            ] for s in current_line):
                # 则开始新行
                lines.append(current_line)
                current_line = [span]
                continue

            # 如果当前的span与当前行的最后一个span在y轴上重叠，则添加到当前行
            if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox'], 0.5):
                current_line.append(span)
            else:
                # 否则，开始新行
                lines.append(current_line)
                current_line = [span]

        # 添加最后一行
        if current_line:
            lines.append(current_line)

        return lines


def merge_spans_to_line_by_layout(spans, layout_bboxes):
    lines = []
    new_spans = []
    dropped_spans = []
    for item in layout_bboxes:
        layout_bbox = item['layout_bbox']
        # 遍历spans,将每个span放入对应的layout中
        layout_sapns = []
        for span in spans:
            if calculate_overlap_area_in_bbox1_area_ratio(
                    span['bbox'], layout_bbox) > 0.6:
                layout_sapns.append(span)
        # 如果layout_sapns不为空，则放入new_spans中
        if len(layout_sapns) > 0:
            new_spans.append(layout_sapns)
            # 从spans删除已经放入layout_sapns中的span
            for layout_sapn in layout_sapns:
                spans.remove(layout_sapn)

    if len(new_spans) > 0:
        for layout_sapns in new_spans:
            layout_lines = merge_spans_to_line(layout_sapns)
            lines.extend(layout_lines)

    # 对line中的span进行排序
    lines = line_sort_spans_by_left_to_right(lines)

    for span in spans:
        span['tag'] = DropTag.NOT_IN_LAYOUT
        dropped_spans.append(span)

    return lines, dropped_spans


def merge_lines_to_block(lines):
    # 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
    blocks = []
    for line in lines:
        blocks.append({
            'bbox': line['bbox'],
            'lines': [line],
        })
    return blocks


def sort_blocks_by_layout(all_bboxes, layout_bboxes):
    new_blocks = []
    sort_blocks = []
    for item in layout_bboxes:
        layout_bbox = item['layout_bbox']

        # 遍历blocks,将每个blocks放入对应的layout中
        layout_blocks = []
        for block in all_bboxes:
            # 如果是footnote则跳过
            if block[7] == BlockType.Footnote:
                continue
            block_bbox = block[:4]
            if calculate_overlap_area_in_bbox1_area_ratio(
                    block_bbox, layout_bbox) > 0.8:
                layout_blocks.append(block)

        # 如果layout_blocks不为空，则放入new_blocks中
        if len(layout_blocks) > 0:
            new_blocks.append(layout_blocks)
            # 从all_bboxes删除已经放入layout_blocks中的block
            for layout_block in layout_blocks:
                all_bboxes.remove(layout_block)

    # 如果new_blocks不为空，则对new_blocks中每个block进行排序
    if len(new_blocks) > 0:
        for bboxes_in_layout_block in new_blocks:
            bboxes_in_layout_block.sort(
                key=lambda x: x[1])  # 一个layout内部的box，按照y0自上而下排序
            sort_blocks.extend(bboxes_in_layout_block)

    # sort_blocks中已经包含了当前页面所有最终留下的block，且已经排好了顺序
    return sort_blocks


def fill_spans_in_blocks(blocks, spans, radio):
    """将allspans中的span按位置关系，放入blocks中."""
    block_with_spans = []
    for block in blocks:
        block_type = block[7]
        block_bbox = block[0:4]
        block_dict = {
            'type': block_type,
            'bbox': block_bbox,
        }
        if block_type in [
            BlockType.ImageBody, BlockType.ImageCaption, BlockType.ImageFootnote,
            BlockType.TableBody, BlockType.TableCaption, BlockType.TableFootnote
        ]:
            block_dict['group_id'] = block[-1]
        block_spans = []
        for span in spans:
            span_bbox = span['bbox']
            if calculate_overlap_area_in_bbox1_area_ratio(
                    span_bbox, block_bbox) > radio:
                block_spans.append(span)
        '''行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)'''
        # displayed_list = []
        # text_inline_lines = []
        # modify_y_axis(block_spans, displayed_list, text_inline_lines)
        '''模型识别错误的行间公式, type类型转换成行内公式'''
        # block_spans = modify_inline_equation(block_spans, displayed_list, text_inline_lines)
        '''bbox去除粘连'''  # 去粘连会影响span的bbox，导致后续fill的时候出错
        # block_spans = remove_overlap_between_bbox_for_span(block_spans)

        block_dict['spans'] = block_spans
        block_with_spans.append(block_dict)

        # 从spans删除已经放入block_spans中的span
        if len(block_spans) > 0:
            for span in block_spans:
                spans.remove(span)

    return block_with_spans, spans


def fix_block_spans(block_with_spans, img_blocks, table_blocks):
    """1、img_block和table_block因为包含caption和footnote的关系，存在block的嵌套关系
    需要将caption和footnote的text_span放入相应img_block和table_block内的
    caption_block和footnote_block中 2、同时需要删除block中的spans字段."""
    fix_blocks = []
    for block in block_with_spans:
        block_type = block['type']

        if block_type == BlockType.Image:
            block = fix_image_block(block, img_blocks)
        elif block_type == BlockType.Table:
            block = fix_table_block(block, table_blocks)
        elif block_type in [BlockType.Text, BlockType.Title]:
            block = fix_text_block(block)
        elif block_type == BlockType.InterlineEquation:
            block = fix_interline_block(block)
        else:
            continue
        fix_blocks.append(block)
    return fix_blocks


def fix_block_spans_v2(block_with_spans):
    """1、img_block和table_block因为包含caption和footnote的关系，存在block的嵌套关系
    需要将caption和footnote的text_span放入相应img_block和table_block内的
    caption_block和footnote_block中 2、同时需要删除block中的spans字段."""
    fix_blocks = []
    for block in block_with_spans:
        block_type = block['type']

        if block_type in [BlockType.Text, BlockType.Title,
                          BlockType.ImageCaption, BlockType.ImageFootnote,
                          BlockType.TableCaption, BlockType.TableFootnote
                          ]:
            block = fix_text_block(block)
        elif block_type in [BlockType.InterlineEquation, BlockType.ImageBody, BlockType.TableBody]:
            block = fix_interline_block(block)
        else:
            continue
        fix_blocks.append(block)
    return fix_blocks


def fix_discarded_block(discarded_block_with_spans):
    fix_discarded_blocks = []
    for block in discarded_block_with_spans:
        block = fix_text_block(block)
        fix_discarded_blocks.append(block)
    return fix_discarded_blocks


def merge_spans_to_block(spans: list, block_bbox: list, block_type: str):
    block_spans = []
    # 如果有img_caption，则将img_block中的text_spans放入img_caption_block中
    for span in spans:
        if calculate_overlap_area_in_bbox1_area_ratio(span['bbox'],
                                                      block_bbox) > 0.6:
            block_spans.append(span)
    block_lines = merge_spans_to_line(block_spans)
    # 对line中的span进行排序
    sort_block_lines = line_sort_spans_by_left_to_right(block_lines)
    block = {'bbox': block_bbox, 'type': block_type, 'lines': sort_block_lines}
    return block, block_spans


def make_body_block(span: dict, block_bbox: list, block_type: str):
    # 创建body_block
    body_line = {
        'bbox': block_bbox,
        'spans': [span],
    }
    body_block = {'bbox': block_bbox, 'type': block_type, 'lines': [body_line]}
    return body_block


def fix_image_block(block, img_blocks):
    block['blocks'] = []
    # 遍历img_blocks,找到与当前block匹配的img_block
    for img_block in img_blocks:
        if _is_in_or_part_overlap_with_area_ratio(block['bbox'],
                                                  img_block['bbox'], 0.95):

            # 创建img_body_block
            for span in block['spans']:
                if span['type'] == ContentType.Image and img_block[
                        'img_body_bbox'] == span['bbox']:
                    # 创建img_body_block
                    img_body_block = make_body_block(
                        span, img_block['img_body_bbox'], BlockType.ImageBody)
                    block['blocks'].append(img_body_block)

                    # 从spans中移除img_body_block中已经放入的span
                    block['spans'].remove(span)
                    break

            # 根据list长度，判断img_block中是否有img_caption
            if img_block['img_caption_bbox'] is not None:
                img_caption_block, img_caption_spans = merge_spans_to_block(
                    block['spans'], img_block['img_caption_bbox'],
                    BlockType.ImageCaption)
                block['blocks'].append(img_caption_block)

            if img_block['img_footnote_bbox'] is not None:
                img_footnote_block, img_footnote_spans = merge_spans_to_block(
                    block['spans'], img_block['img_footnote_bbox'],
                    BlockType.ImageFootnote)
                block['blocks'].append(img_footnote_block)
            break
    del block['spans']
    return block


def fix_table_block(block, table_blocks):
    block['blocks'] = []
    # 遍历table_blocks,找到与当前block匹配的table_block
    for table_block in table_blocks:
        if _is_in_or_part_overlap_with_area_ratio(block['bbox'],
                                                  table_block['bbox'], 0.95):

            # 创建table_body_block
            for span in block['spans']:
                if span['type'] == ContentType.Table and table_block[
                        'table_body_bbox'] == span['bbox']:
                    # 创建table_body_block
                    table_body_block = make_body_block(
                        span, table_block['table_body_bbox'],
                        BlockType.TableBody)
                    block['blocks'].append(table_body_block)

                    # 从spans中移除img_body_block中已经放入的span
                    block['spans'].remove(span)
                    break

            # 根据list长度，判断table_block中是否有caption
            if table_block['table_caption_bbox'] is not None:
                table_caption_block, table_caption_spans = merge_spans_to_block(
                    block['spans'], table_block['table_caption_bbox'],
                    BlockType.TableCaption)
                block['blocks'].append(table_caption_block)

                # 如果table_caption_block_spans不为空
                if len(table_caption_spans) > 0:
                    #  一些span已经放入了caption_block中，需要从block['spans']中删除
                    for span in table_caption_spans:
                        block['spans'].remove(span)

            # 根据list长度，判断table_block中是否有table_note
            if table_block['table_footnote_bbox'] is not None:
                table_footnote_block, table_footnote_spans = merge_spans_to_block(
                    block['spans'], table_block['table_footnote_bbox'],
                    BlockType.TableFootnote)
                block['blocks'].append(table_footnote_block)

            break
    del block['spans']
    return block


def fix_text_block(block):
    # 文本block中的公式span都应该转换成行内type
    for span in block['spans']:
        if span['type'] == ContentType.InterlineEquation:
            span['type'] = ContentType.InlineEquation
    block_lines = merge_spans_to_line(block['spans'])
    sort_block_lines = line_sort_spans_by_left_to_right(block_lines)
    block['lines'] = sort_block_lines
    del block['spans']
    return block


def fix_interline_block(block):
    block_lines = merge_spans_to_line(block['spans'])
    sort_block_lines = line_sort_spans_by_left_to_right(block_lines)
    block['lines'] = sort_block_lines
    del block['spans']
    return block