magic_pdf/model/magic_model.py

import json
import math

from magic_pdf.libs.commons import fitz
from loguru import logger

from magic_pdf.libs.commons import join_path
from magic_pdf.libs.coordinate_transform import get_scale_ratio
from magic_pdf.libs.ocr_content_type import ContentType
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
from magic_pdf.libs.math import float_gt
from magic_pdf.libs.boxbase import (
    _is_in,
    bbox_relative_pos,
    bbox_distance,
    _is_part_overlap,
    calculate_overlap_area_in_bbox1_area_ratio,
    calculate_iou,
)
from magic_pdf.libs.ModelBlockTypeEnum import ModelBlockTypeEnum

CAPATION_OVERLAP_AREA_RATIO = 0.6


class MagicModel:
    """
    每个函数没有得到元素的时候返回空list

    """

    def __fix_axis(self):
        for model_page_info in self.__model_list:
            need_remove_list = []
            page_no = model_page_info["page_info"]["page_no"]
            horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
                model_page_info, self.__docs[page_no]
            )
            layout_dets = model_page_info["layout_dets"]
            for layout_det in layout_dets:

                if layout_det.get("bbox") is not None:
                    # 兼容直接输出bbox的模型数据,如paddle
                    x0, y0, x1, y1 = layout_det["bbox"]
                else:
                    # 兼容直接输出poly的模型数据，如xxx
                    x0, y0, _, _, x1, y1, _, _ = layout_det["poly"]

                bbox = [
                    int(x0 / horizontal_scale_ratio),
                    int(y0 / vertical_scale_ratio),
                    int(x1 / horizontal_scale_ratio),
                    int(y1 / vertical_scale_ratio),
                ]
                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):
        for model_page_info in self.__model_list:
            need_remove_list = []
            layout_dets = model_page_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):
        for model_page_info in self.__model_list:
            need_remove_list = []
            layout_dets = model_page_info["layout_dets"]
            for layout_det1 in layout_dets:
                for layout_det2 in layout_dets:
                    if layout_det1 == layout_det2:
                        continue
                    if layout_det1["category_id"] in [
                        0,
                        1,
                        2,
                        3,
                        4,
                        5,
                        6,
                        7,
                        8,
                        9,
                    ] and layout_det2["category_id"] in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
                        if (
                            calculate_iou(layout_det1["bbox"], layout_det2["bbox"])
                            > 0.9
                        ):
                            if layout_det1["score"] < layout_det2["score"]:
                                layout_det_need_remove = layout_det1
                            else:
                                layout_det_need_remove = layout_det2

                            if layout_det_need_remove not in need_remove_list:
                                need_remove_list.append(layout_det_need_remove)
                        else:
                            continue
                    else:
                        continue
            for need_remove in need_remove_list:
                layout_dets.remove(need_remove)

    def __init__(self, model_list: list, docs: fitz.Document):
        self.__model_list = model_list
        self.__docs = docs
        """为所有模型数据添加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()

    def __reduct_overlap(self, bboxes):
        N = len(bboxes)
        keep = [True] * N
        for i in range(N):
            for j in range(N):
                if i == j:
                    continue
                if _is_in(bboxes[i]["bbox"], bboxes[j]["bbox"]):
                    keep[i] = False

        return [bboxes[i] for i in range(N) if keep[i]]

    def __tie_up_category_by_distance(
        self, page_no, subject_category_id, object_category_id
    ):
        """
        假定每个 subject 最多有一个 object (可以有多个相邻的 object 合并为单个 object)，每个 object 只能属于一个 subject
        """
        ret = []
        MAX_DIS_OF_POINT = 10**9 + 7

        # subject 和 object 的 bbox 会合并成一个大的 bbox （named: merged bbox）。 筛选出所有和 merged bbox 有 overlap 且 overlap 面积大于 object 的面积的 subjects。
        # 再求出筛选出的 subjects 和 object 的最短距离！
        def may_find_other_nearest_bbox(subject_idx, object_idx):
            ret = float("inf")

            x0 = min(
                all_bboxes[subject_idx]["bbox"][0], all_bboxes[object_idx]["bbox"][0]
            )
            y0 = min(
                all_bboxes[subject_idx]["bbox"][1], all_bboxes[object_idx]["bbox"][1]
            )
            x1 = max(
                all_bboxes[subject_idx]["bbox"][2], all_bboxes[object_idx]["bbox"][2]
            )
            y1 = max(
                all_bboxes[subject_idx]["bbox"][3], all_bboxes[object_idx]["bbox"][3]
            )

            object_area = abs(
                all_bboxes[object_idx]["bbox"][2] - all_bboxes[object_idx]["bbox"][0]
            ) * abs(
                all_bboxes[object_idx]["bbox"][3] - all_bboxes[object_idx]["bbox"][1]
            )

            for i in range(len(all_bboxes)):
                if (
                    i == subject_idx
                    or all_bboxes[i]["category_id"] != subject_category_id
                ):
                    continue
                if _is_part_overlap([x0, y0, x1, y1], all_bboxes[i]["bbox"]) or _is_in(
                    all_bboxes[i]["bbox"], [x0, y0, x1, y1]
                ):

                    i_area = abs(
                        all_bboxes[i]["bbox"][2] - all_bboxes[i]["bbox"][0]
                    ) * abs(all_bboxes[i]["bbox"][3] - all_bboxes[i]["bbox"][1])
                    if i_area >= object_area:
                        ret = min(float("inf"), dis[i][object_idx])

            return ret

        def expand_bbbox(idxes):
            x0s = [all_bboxes[idx]["bbox"][0] for idx in idxes] 
            y0s = [all_bboxes[idx]["bbox"][1] for idx in idxes] 
            x1s = [all_bboxes[idx]["bbox"][2] for idx in idxes] 
            y1s = [all_bboxes[idx]["bbox"][3] for idx in idxes] 
            return min(x0s), min(y0s), max(x1s), max(y1s)

        subjects = self.__reduct_overlap(
            list(
                map(
                    lambda x: {"bbox": x["bbox"], "score": x["score"]},
                    filter(
                        lambda x: x["category_id"] == subject_category_id,
                        self.__model_list[page_no]["layout_dets"],
                    ),
                )
            )
        )

        objects = self.__reduct_overlap(
            list(
                map(
                    lambda x: {"bbox": x["bbox"], "score": x["score"]},
                    filter(
                        lambda x: x["category_id"] == object_category_id,
                        self.__model_list[page_no]["layout_dets"],
                    ),
                )
            )
        )
        subject_object_relation_map = {}

        subjects.sort(
            key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2
        )  # get the distance !

        all_bboxes = []

        for v in subjects:
            all_bboxes.append(
                {
                    "category_id": subject_category_id,
                    "bbox": v["bbox"],
                    "score": v["score"],
                }
            )

        for v in objects:
            all_bboxes.append(
                {
                    "category_id": object_category_id,
                    "bbox": v["bbox"],
                    "score": v["score"],
                }
            )

        N = len(all_bboxes)
        dis = [[MAX_DIS_OF_POINT] * N for _ in range(N)]

        for i in range(N):
            for j in range(i):
                if (
                    all_bboxes[i]["category_id"] == subject_category_id
                    and all_bboxes[j]["category_id"] == subject_category_id
                ):
                    continue

                dis[i][j] = bbox_distance(all_bboxes[i]["bbox"], all_bboxes[j]["bbox"])
                dis[j][i] = dis[i][j]

        used = set()
        for i in range(N):
            # 求第 i 个 subject 所关联的 object
            if all_bboxes[i]["category_id"] != subject_category_id:
                continue
            seen = set()
            candidates = []
            arr = []
            for j in range(N):

                pos_flag_count = sum(
                    list(
                        map(
                            lambda x: 1 if x else 0,
                            bbox_relative_pos(
                                all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]
                            ),
                        )
                    )
                )
                if pos_flag_count > 1:
                    continue
                if (
                    all_bboxes[j]["category_id"] != object_category_id
                    or j in used
                    or dis[i][j] == MAX_DIS_OF_POINT
                ):
                    continue
                left, right, _, _ = bbox_relative_pos(
                    all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]
                )  # 由  pos_flag_count 相关逻辑保证本段逻辑准确性
                if left or right:
                    one_way_dis = all_bboxes[i]["bbox"][2] - all_bboxes[i]["bbox"][0]
                else:
                    one_way_dis = all_bboxes[i]["bbox"][3] - all_bboxes[i]["bbox"][1]
                if dis[i][j] > one_way_dis:
                    continue
                arr.append((dis[i][j], j))

            arr.sort(key=lambda x: x[0])
            if len(arr) > 0:
                # bug: 离该subject 最近的 object 可能跨越了其它的 subject 。比如 [this subect] [some sbuject] [the nearest objec of subject]
                if may_find_other_nearest_bbox(i, arr[0][1]) >= arr[0][0]:

                    candidates.append(arr[0][1])
                    seen.add(arr[0][1])

            # 已经获取初始种子
            for j in set(candidates):
                tmp = []
                for k in range(i + 1, N):
                    pos_flag_count = sum(
                        list(
                            map(
                                lambda x: 1 if x else 0,
                                bbox_relative_pos(
                                    all_bboxes[j]["bbox"], all_bboxes[k]["bbox"]
                                ),
                            )
                        )
                    )

                    if pos_flag_count > 1:
                        continue

                    if (
                        all_bboxes[k]["category_id"] != object_category_id
                        or k in used
                        or k in seen
                        or dis[j][k] == MAX_DIS_OF_POINT
                        or dis[j][k] > dis[i][j]
                    ):
                        continue

                    is_nearest = True
                    for l in range(i + 1, N):
                        if l in (j, k) or l in used or l in seen:
                            continue

                        if not float_gt(dis[l][k], dis[j][k]):
                            is_nearest = False
                            break

                    if is_nearest:
                        nx0, ny0, nx1, ny1 = expand_bbbox(list(seen) + [k])
                        n_dis = bbox_distance(all_bboxes[i]["bbox"], [nx0, ny0, nx1, ny1])
                        if float_gt(dis[i][j], n_dis):
                            continue
                        tmp.append(k)
                        seen.add(k)

                candidates = tmp
                if len(candidates) == 0:
                    break

            # 已经获取到某个 figure 下所有的最靠近的 captions，以及最靠近这些 captions 的 captions 。
            # 先扩一下 bbox，
            ox0, oy0, ox1, oy1 = expand_bbbox(list(seen) + [i])
            ix0, iy0, ix1, iy1 = all_bboxes[i]["bbox"]

            # 分成了 4 个截取空间，需要计算落在每个截取空间下 objects 合并后占据的矩形面积
            caption_poses = [
                [ox0, oy0, ix0, oy1],
                [ox0, oy0, ox1, iy0],
                [ox0, iy1, ox1, oy1],
                [ix1, oy0, ox1, oy1],
            ]

            caption_areas = []
            for bbox in caption_poses:
                embed_arr = []
                for idx in seen:
                    if (
                        calculate_overlap_area_in_bbox1_area_ratio(
                            all_bboxes[idx]["bbox"], bbox
                        )
                        > CAPATION_OVERLAP_AREA_RATIO
                    ):
                        embed_arr.append(idx)

                if len(embed_arr) > 0:
                    embed_x0 = min([all_bboxes[idx]["bbox"][0] for idx in embed_arr])
                    embed_y0 = min([all_bboxes[idx]["bbox"][1] for idx in embed_arr])
                    embed_x1 = max([all_bboxes[idx]["bbox"][2] for idx in embed_arr])
                    embed_y1 = max([all_bboxes[idx]["bbox"][3] for idx in embed_arr])
                    caption_areas.append(
                        int(abs(embed_x1 - embed_x0) * abs(embed_y1 - embed_y0))
                    )
                else:
                    caption_areas.append(0)

            subject_object_relation_map[i] = []
            if max(caption_areas) > 0:
                max_area_idx = caption_areas.index(max(caption_areas))
                caption_bbox = caption_poses[max_area_idx]

                for j in seen:
                    if (
                        calculate_overlap_area_in_bbox1_area_ratio(
                            all_bboxes[j]["bbox"], caption_bbox
                        )
                        > CAPATION_OVERLAP_AREA_RATIO
                    ):
                        used.add(j)
                        subject_object_relation_map[i].append(j)

        for i in sorted(subject_object_relation_map.keys()):
            result = {
                "subject_body": all_bboxes[i]["bbox"],
                "all": all_bboxes[i]["bbox"],
                "score": all_bboxes[i]["score"],
            }

            if len(subject_object_relation_map[i]) > 0:
                x0 = min(
                    [all_bboxes[j]["bbox"][0] for j in subject_object_relation_map[i]]
                )
                y0 = min(
                    [all_bboxes[j]["bbox"][1] for j in subject_object_relation_map[i]]
                )
                x1 = max(
                    [all_bboxes[j]["bbox"][2] for j in subject_object_relation_map[i]]
                )
                y1 = max(
                    [all_bboxes[j]["bbox"][3] for j in subject_object_relation_map[i]]
                )
                result["object_body"] = [x0, y0, x1, y1]
                result["all"] = [
                    min(x0, all_bboxes[i]["bbox"][0]),
                    min(y0, all_bboxes[i]["bbox"][1]),
                    max(x1, all_bboxes[i]["bbox"][2]),
                    max(y1, all_bboxes[i]["bbox"][3]),
                ]
            ret.append(result)

        total_subject_object_dis = 0
        # 计算已经配对的 distance 距离
        for i in subject_object_relation_map.keys():
            for j in subject_object_relation_map[i]:
                total_subject_object_dis += bbox_distance(
                    all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]
                )

        # 计算未匹配的 subject 和 object 的距离（非精确版）
        with_caption_subject = set(
            [
                key
                for key in subject_object_relation_map.keys()
                if len(subject_object_relation_map[i]) > 0
            ]
        )
        for i in range(N):
            if all_bboxes[i]["category_id"] != object_category_id or i in used:
                continue
            candidates = []
            for j in range(N):
                if (
                    all_bboxes[j]["category_id"] != subject_category_id
                    or j in with_caption_subject
                ):
                    continue
                candidates.append((dis[i][j], j))
            if len(candidates) > 0:
                candidates.sort(key=lambda x: x[0])
                total_subject_object_dis += candidates[0][1]
                with_caption_subject.add(j)
        return ret, total_subject_object_dis

    def get_imgs(self, page_no: int):
        figure_captions, _ = self.__tie_up_category_by_distance(
            page_no, 3, 4
        )
        return [
            {
                "bbox": record["all"],
                "img_body_bbox": record["subject_body"],
                "img_caption_bbox": record.get("object_body", None),
                "score": record["score"],
            }
            for record in figure_captions
        ]

    def get_tables(
        self, page_no: int
    ) -> list:  # 3个坐标， caption, table主体，table-note
        with_captions, _ = self.__tie_up_category_by_distance(page_no, 5, 6)
        with_footnotes, _ = self.__tie_up_category_by_distance(page_no, 5, 7)
        ret = []
        N, M = len(with_captions), len(with_footnotes)
        assert N == M
        for i in range(N):
            record = {
                "score": with_captions[i]["score"],
                "table_caption_bbox": with_captions[i].get("object_body", None),
                "table_body_bbox": with_captions[i]["subject_body"],
                "table_footnote_bbox": with_footnotes[i].get("object_body", None),
            }

            x0 = min(with_captions[i]["all"][0], with_footnotes[i]["all"][0])
            y0 = min(with_captions[i]["all"][1], with_footnotes[i]["all"][1])
            x1 = max(with_captions[i]["all"][2], with_footnotes[i]["all"][2])
            y1 = max(with_captions[i]["all"][3], with_footnotes[i]["all"][3])
            record["bbox"] = [x0, y0, x1, y1]
            ret.append(record)
        return ret

    def get_equations(self, page_no: int) -> list:  # 有坐标，也有字
        inline_equations = self.__get_blocks_by_type(
            ModelBlockTypeEnum.EMBEDDING.value, page_no, ["latex"]
        )
        interline_equations = self.__get_blocks_by_type(
            ModelBlockTypeEnum.ISOLATED.value, page_no, ["latex"]
        )
        interline_equations_blocks = self.__get_blocks_by_type(
            ModelBlockTypeEnum.ISOLATE_FORMULA.value, page_no
        )
        return inline_equations, interline_equations, interline_equations_blocks

    def get_discarded(self, page_no: int) -> list:  # 自研模型，只有坐标
        blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.ABANDON.value, page_no)
        return blocks

    def get_text_blocks(self, page_no: int) -> list:  # 自研模型搞的，只有坐标，没有字
        blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.PLAIN_TEXT.value, page_no)
        return blocks

    def get_title_blocks(self, page_no: int) -> list:  # 自研模型，只有坐标，没字
        blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.TITLE.value, page_no)
        return blocks

    def get_ocr_text(self, page_no: int) -> list:  # paddle 搞的，有字也有坐标
        text_spans = []
        model_page_info = self.__model_list[page_no]
        layout_dets = model_page_info["layout_dets"]
        for layout_det in layout_dets:
            if layout_det["category_id"] == "15":
                span = {
                    "bbox": layout_det["bbox"],
                    "content": layout_det["text"],
                }
                text_spans.append(span)
        return text_spans

    def get_all_spans(self, page_no: int) -> 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 = []
        model_page_info = self.__model_list[page_no]
        layout_dets = model_page_info["layout_dets"]
        allow_category_id_list = [3, 5, 13, 14, 15]
        """当成span拼接的"""
        #  3: 'image', # 图片
        #  5: 'table',       # 表格
        #  13: 'inline_equation',     # 行内公式
        #  14: 'interline_equation',      # 行间公式
        #  15: 'text',      # ocr识别文本
        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 == 3:
                    span["type"] = ContentType.Image
                elif category_id == 5:
                    span["type"] = ContentType.Table
                elif category_id == 13:
                    span["content"] = layout_det["latex"]
                    span["type"] = ContentType.InlineEquation
                elif category_id == 14:
                    span["content"] = layout_det["latex"]
                    span["type"] = ContentType.InterlineEquation
                elif category_id == 15:
                    span["content"] = layout_det["text"]
                    span["type"] = ContentType.Text
                all_spans.append(span)
        return remove_duplicate_spans(all_spans)

    def get_page_size(self, page_no: int):  # 获取页面宽高
        # 获取当前页的page对象
        page = self.__docs[page_no]
        # 获取当前页的宽高
        page_w = page.rect.width
        page_h = page.rect.height
        return page_w, page_h

    def __get_blocks_by_type(
        self, type: int, page_no: int, extra_col: list[str] = []
    ) -> list:
        blocks = []
        for page_dict in self.__model_list:
            layout_dets = page_dict.get("layout_dets", [])
            page_info = page_dict.get("page_info", {})
            page_number = page_info.get("page_no", -1)
            if page_no != page_number:
                continue
            for item in layout_dets:
                category_id = item.get("category_id", -1)
                bbox = item.get("bbox", None)

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

    def get_model_list(self, page_no):
        return self.__model_list[page_no]


if __name__ == "__main__":
    drw = DiskReaderWriter(r"D:/project/20231108code-clean")
    if 0:
        pdf_file_path = r"linshixuqiu\19983-00.pdf"
        model_file_path = r"linshixuqiu\19983-00_new.json"
        pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
        model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
        model_list = json.loads(model_json_txt)
        write_path = r"D:\project\20231108code-clean\linshixuqiu\19983-00"
        img_bucket_path = "imgs"
        img_writer = DiskReaderWriter(join_path(write_path, img_bucket_path))
        pdf_docs = fitz.open("pdf", pdf_bytes)
        magic_model = MagicModel(model_list, pdf_docs)

    if 1:
        model_list = json.loads(
            drw.read("/opt/data/pdf/20240418/j.chroma.2009.03.042.json")
        )
        pdf_bytes = drw.read(
            "/opt/data/pdf/20240418/j.chroma.2009.03.042.pdf", AbsReaderWriter.MODE_BIN
        )
        pdf_docs = fitz.open("pdf", pdf_bytes)
        magic_model = MagicModel(model_list, pdf_docs)
        for i in range(7):
            print(magic_model.get_imgs(i))