MinerU/mineru/backend/pipeline/pipeline_magic_model.py

from mineru.utils.boxbase import bbox_relative_pos, calculate_iou, bbox_distance, is_in, get_minbox_if_overlap_by_ratio
from mineru.utils.enum_class import CategoryId, ContentType
from mineru.utils.magic_model_utils import tie_up_category_by_distance_v3, reduct_overlap


class MagicModel:
    """每个函数没有得到元素的时候返回空list."""
    def __init__(self, page_model_info: dict, scale: float):
        self.__page_model_info = page_model_info
        self.__scale = scale
        """为所有模型数据添加bbox信息(缩放，poly->bbox)"""
        self.__fix_axis()
        """删除置信度特别低的模型数据(<0.05),提高质量"""
        self.__fix_by_remove_low_confidence()
        """删除高iou(>0.9)数据中置信度较低的那个"""
        self.__fix_by_remove_high_iou_and_low_confidence()
        """将部分tbale_footnote修正为image_footnote"""
        self.__fix_footnote()
        """处理重叠的image_body和table_body"""
        self.__fix_by_remove_overlap_image_table_body()

    def __fix_by_remove_overlap_image_table_body(self):
        need_remove_list = []
        layout_dets = self.__page_model_info['layout_dets']
        image_blocks = list(filter(
            lambda x: x['category_id'] == CategoryId.ImageBody, layout_dets
        ))
        table_blocks = list(filter(
            lambda x: x['category_id'] == CategoryId.TableBody, layout_dets
        ))

        def add_need_remove_block(blocks):
            for i in range(len(blocks)):
                for j in range(i + 1, len(blocks)):
                    block1 = blocks[i]
                    block2 = blocks[j]
                    overlap_box = get_minbox_if_overlap_by_ratio(
                        block1['bbox'], block2['bbox'], 0.8
                    )
                    if overlap_box is not None:
                        # 判断哪个区块的面积更小，移除较小的区块
                        area1 = (block1['bbox'][2] - block1['bbox'][0]) * (block1['bbox'][3] - block1['bbox'][1])
                        area2 = (block2['bbox'][2] - block2['bbox'][0]) * (block2['bbox'][3] - block2['bbox'][1])

                        if area1 <= area2:
                            block_to_remove = block1
                            large_block = block2
                        else:
                            block_to_remove = block2
                            large_block = block1

                        if block_to_remove not in need_remove_list:
                            # 扩展大区块的边界框
                            x1, y1, x2, y2 = large_block['bbox']
                            sx1, sy1, sx2, sy2 = block_to_remove['bbox']
                            x1 = min(x1, sx1)
                            y1 = min(y1, sy1)
                            x2 = max(x2, sx2)
                            y2 = max(y2, sy2)
                            large_block['bbox'] = [x1, y1, x2, y2]
                            need_remove_list.append(block_to_remove)

        # 处理图像-图像重叠
        add_need_remove_block(image_blocks)
        # 处理表格-表格重叠
        add_need_remove_block(table_blocks)

        # 从布局中移除标记的区块
        for need_remove in need_remove_list:
            if need_remove in layout_dets:
                layout_dets.remove(need_remove)


    def __fix_axis(self):
        need_remove_list = []
        layout_dets = self.__page_model_info['layout_dets']
        for layout_det in layout_dets:
            x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
            bbox = [
                int(x0 / self.__scale),
                int(y0 / self.__scale),
                int(x1 / self.__scale),
                int(y1 / self.__scale),
            ]
            layout_det['bbox'] = bbox
            # 删除高度或者宽度小于等于0的spans
            if bbox[2] - bbox[0] <= 0 or bbox[3] - bbox[1] <= 0:
                need_remove_list.append(layout_det)
        for need_remove in need_remove_list:
            layout_dets.remove(need_remove)

    def __fix_by_remove_low_confidence(self):
        need_remove_list = []
        layout_dets = self.__page_model_info['layout_dets']
        for layout_det in layout_dets:
            if layout_det['score'] <= 0.05:
                need_remove_list.append(layout_det)
            else:
                continue
        for need_remove in need_remove_list:
            layout_dets.remove(need_remove)

    def __fix_by_remove_high_iou_and_low_confidence(self):
        need_remove_list = []
        layout_dets = list(filter(
            lambda x: x['category_id'] in [
                    CategoryId.Title,
                    CategoryId.Text,
                    CategoryId.ImageBody,
                    CategoryId.ImageCaption,
                    CategoryId.TableBody,
                    CategoryId.TableCaption,
                    CategoryId.TableFootnote,
                    CategoryId.InterlineEquation_Layout,
                    CategoryId.InterlineEquationNumber_Layout,
                ], self.__page_model_info['layout_dets']
            )
        )
        for i in range(len(layout_dets)):
            for j in range(i + 1, len(layout_dets)):
                layout_det1 = layout_dets[i]
                layout_det2 = layout_dets[j]

                if calculate_iou(layout_det1['bbox'], layout_det2['bbox']) > 0.9:

                    layout_det_need_remove = layout_det1 if layout_det1['score'] < layout_det2['score'] else layout_det2

                    if layout_det_need_remove not in need_remove_list:
                        need_remove_list.append(layout_det_need_remove)

        for need_remove in need_remove_list:
            self.__page_model_info['layout_dets'].remove(need_remove)

    def __fix_footnote(self):
        footnotes = []
        figures = []
        tables = []

        for obj in self.__page_model_info['layout_dets']:
            if obj['category_id'] == CategoryId.TableFootnote:
                footnotes.append(obj)
            elif obj['category_id'] == CategoryId.ImageBody:
                figures.append(obj)
            elif obj['category_id'] == CategoryId.TableBody:
                tables.append(obj)
            if len(footnotes) * len(figures) == 0:
                continue
        dis_figure_footnote = {}
        dis_table_footnote = {}

        for i in range(len(footnotes)):
            for j in range(len(figures)):
                pos_flag_count = sum(
                    list(
                        map(
                            lambda x: 1 if x else 0,
                            bbox_relative_pos(
                                footnotes[i]['bbox'], figures[j]['bbox']
                            ),
                        )
                    )
                )
                if pos_flag_count > 1:
                    continue
                dis_figure_footnote[i] = min(
                    self._bbox_distance(figures[j]['bbox'], footnotes[i]['bbox']),
                    dis_figure_footnote.get(i, float('inf')),
                )
        for i in range(len(footnotes)):
            for j in range(len(tables)):
                pos_flag_count = sum(
                    list(
                        map(
                            lambda x: 1 if x else 0,
                            bbox_relative_pos(
                                footnotes[i]['bbox'], tables[j]['bbox']
                            ),
                        )
                    )
                )
                if pos_flag_count > 1:
                    continue

                dis_table_footnote[i] = min(
                    self._bbox_distance(tables[j]['bbox'], footnotes[i]['bbox']),
                    dis_table_footnote.get(i, float('inf')),
                )
        for i in range(len(footnotes)):
            if i not in dis_figure_footnote:
                continue
            if dis_table_footnote.get(i, float('inf')) > dis_figure_footnote[i]:
                footnotes[i]['category_id'] = CategoryId.ImageFootnote

    def _bbox_distance(self, bbox1, bbox2):
        left, right, bottom, top = bbox_relative_pos(bbox1, bbox2)
        flags = [left, right, bottom, top]
        count = sum([1 if v else 0 for v in flags])
        if count > 1:
            return float('inf')
        if left or right:
            l1 = bbox1[3] - bbox1[1]
            l2 = bbox2[3] - bbox2[1]
        else:
            l1 = bbox1[2] - bbox1[0]
            l2 = bbox2[2] - bbox2[0]

        if l2 > l1 and (l2 - l1) / l1 > 0.3:
            return float('inf')

        return bbox_distance(bbox1, bbox2)

    def __tie_up_category_by_distance_v3(self, subject_category_id, object_category_id):
        # 定义获取主体和客体对象的函数
        def get_subjects():
            return reduct_overlap(
                list(
                    map(
                        lambda x: {'bbox': x['bbox'], 'score': x['score']},
                        filter(
                            lambda x: x['category_id'] == subject_category_id,
                            self.__page_model_info['layout_dets'],
                        ),
                    )
                )
            )

        def get_objects():
            return reduct_overlap(
                list(
                    map(
                        lambda x: {'bbox': x['bbox'], 'score': x['score']},
                        filter(
                            lambda x: x['category_id'] == object_category_id,
                            self.__page_model_info['layout_dets'],
                        ),
                    )
                )
            )

        # 调用通用方法
        return tie_up_category_by_distance_v3(
            get_subjects,
            get_objects
        )

    def get_imgs(self):
        with_captions = self.__tie_up_category_by_distance_v3(
            CategoryId.ImageBody, CategoryId.ImageCaption
        )
        with_footnotes = self.__tie_up_category_by_distance_v3(
            CategoryId.ImageBody, CategoryId.ImageFootnote
        )
        ret = []
        for v in with_captions:
            record = {
                'image_body': v['sub_bbox'],
                'image_caption_list': v['obj_bboxes'],
            }
            filter_idx = v['sub_idx']
            d = next(filter(lambda x: x['sub_idx'] == filter_idx, with_footnotes))
            record['image_footnote_list'] = d['obj_bboxes']
            ret.append(record)
        return ret

    def get_tables(self) -> list:
        with_captions = self.__tie_up_category_by_distance_v3(
            CategoryId.TableBody, CategoryId.TableCaption
        )
        with_footnotes = self.__tie_up_category_by_distance_v3(
            CategoryId.TableBody, CategoryId.TableFootnote
        )
        ret = []
        for v in with_captions:
            record = {
                'table_body': v['sub_bbox'],
                'table_caption_list': v['obj_bboxes'],
            }
            filter_idx = v['sub_idx']
            d = next(filter(lambda x: x['sub_idx'] == filter_idx, with_footnotes))
            record['table_footnote_list'] = d['obj_bboxes']
            ret.append(record)
        return ret

    def get_equations(self) -> tuple[list, list, list]:  # 有坐标，也有字
        inline_equations = self.__get_blocks_by_type(
            CategoryId.InlineEquation, ['latex']
        )
        interline_equations = self.__get_blocks_by_type(
            CategoryId.InterlineEquation_YOLO, ['latex']
        )
        interline_equations_blocks = self.__get_blocks_by_type(
            CategoryId.InterlineEquation_Layout
        )
        return inline_equations, interline_equations, interline_equations_blocks

    def get_discarded(self) -> list:  # 自研模型，只有坐标
        blocks = self.__get_blocks_by_type(CategoryId.Abandon)
        return blocks

    def get_text_blocks(self) -> list:  # 自研模型搞的，只有坐标，没有字
        blocks = self.__get_blocks_by_type(CategoryId.Text)
        return blocks

    def get_title_blocks(self) -> list:  # 自研模型，只有坐标，没字
        blocks = self.__get_blocks_by_type(CategoryId.Title)
        return blocks

    def get_all_spans(self) -> list:

        def remove_duplicate_spans(spans):
            new_spans = []
            for span in spans:
                if not any(span == existing_span for existing_span in new_spans):
                    new_spans.append(span)
            return new_spans

        all_spans = []
        layout_dets = self.__page_model_info['layout_dets']
        allow_category_id_list = [
            CategoryId.ImageBody,
            CategoryId.TableBody,
            CategoryId.InlineEquation,
            CategoryId.InterlineEquation_YOLO,
            CategoryId.OcrText,
        ]
        """当成span拼接的"""
        for layout_det in layout_dets:
            category_id = layout_det['category_id']
            if category_id in allow_category_id_list:
                span = {'bbox': layout_det['bbox'], 'score': layout_det['score']}
                if category_id == CategoryId.ImageBody:
                    span['type'] = ContentType.IMAGE
                elif category_id == CategoryId.TableBody:
                    # 获取table模型结果
                    latex = layout_det.get('latex', None)
                    html = layout_det.get('html', None)
                    if latex:
                        span['latex'] = latex
                    elif html:
                        span['html'] = html
                    span['type'] = ContentType.TABLE
                elif category_id == CategoryId.InlineEquation:
                    span['content'] = layout_det['latex']
                    span['type'] = ContentType.INLINE_EQUATION
                elif category_id == CategoryId.InterlineEquation_YOLO:
                    span['content'] = layout_det['latex']
                    span['type'] = ContentType.INTERLINE_EQUATION
                elif category_id == CategoryId.OcrText:
                    span['content'] = layout_det['text']
                    span['type'] = ContentType.TEXT
                all_spans.append(span)
        return remove_duplicate_spans(all_spans)

    def __get_blocks_by_type(
        self, category_type: int, extra_col=None
    ) -> list:
        if extra_col is None:
            extra_col = []
        blocks = []
        layout_dets = self.__page_model_info.get('layout_dets', [])
        for item in layout_dets:
            category_id = item.get('category_id', -1)
            bbox = item.get('bbox', None)

            if category_id == category_type:
                block = {
                    'bbox': bbox,
                    'score': item.get('score'),
                }
                for col in extra_col:
                    block[col] = item.get(col, None)
                blocks.append(block)
        return blocks