jk-rag-python/mineru/utils/magic_model_utils.py

"""
包含两个MagicModel类中重复使用的方法和逻辑
"""
from typing import List, Dict, Any, Callable

from loguru import logger
from mineru.utils.boxbase import bbox_distance, bbox_center_distance, is_in


def reduct_overlap(bboxes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    去除重叠的bbox，保留不被其他bbox包含的bbox

    Args:
        bboxes: 包含bbox信息的字典列表

    Returns:
        去重后的bbox列表
    """
    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_v3(
        get_subjects_func: Callable,
        get_objects_func: Callable,
        extract_subject_func: Callable = None,
        extract_object_func: Callable = None
):
    """
    通用的类别关联方法，用于将主体对象与客体对象进行关联

    参数:
        get_subjects_func: 函数，提取主体对象
        get_objects_func: 函数，提取客体对象
        extract_subject_func: 函数，自定义提取主体属性（默认使用bbox和其他属性）
        extract_object_func: 函数，自定义提取客体属性（默认使用bbox和其他属性）

    返回:
        关联后的对象列表
    """
    subjects = get_subjects_func()
    objects = get_objects_func()

    # 如果没有提供自定义提取函数，使用默认函数
    if extract_subject_func is None:
        extract_subject_func = lambda x: x
    if extract_object_func is None:
        extract_object_func = lambda x: x

    ret = []
    N, M = len(subjects), len(objects)
    subjects.sort(key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2)
    objects.sort(key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2)

    OBJ_IDX_OFFSET = 10000
    SUB_BIT_KIND, OBJ_BIT_KIND = 0, 1

    all_boxes_with_idx = [(i, SUB_BIT_KIND, sub["bbox"][0], sub["bbox"][1]) for i, sub in enumerate(subjects)] + [
        (i + OBJ_IDX_OFFSET, OBJ_BIT_KIND, obj["bbox"][0], obj["bbox"][1]) for i, obj in enumerate(objects)
    ]
    seen_idx = set()
    seen_sub_idx = set()

    while N > len(seen_sub_idx):
        candidates = []
        for idx, kind, x0, y0 in all_boxes_with_idx:
            if idx in seen_idx:
                continue
            candidates.append((idx, kind, x0, y0))

        if len(candidates) == 0:
            break
        left_x = min([v[2] for v in candidates])
        top_y = min([v[3] for v in candidates])

        candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y) ** 2)

        fst_idx, fst_kind, left_x, top_y = candidates[0]
        fst_bbox = subjects[fst_idx]['bbox'] if fst_kind == SUB_BIT_KIND else objects[fst_idx - OBJ_IDX_OFFSET]['bbox']
        candidates.sort(
            key=lambda x: bbox_distance(fst_bbox, subjects[x[0]]['bbox']) if x[1] == SUB_BIT_KIND else bbox_distance(
                fst_bbox, objects[x[0] - OBJ_IDX_OFFSET]['bbox']))
        nxt = None

        for i in range(1, len(candidates)):
            if candidates[i][1] ^ fst_kind == 1:
                nxt = candidates[i]
                break
        if nxt is None:
            break

        if fst_kind == SUB_BIT_KIND:
            sub_idx, obj_idx = fst_idx, nxt[0] - OBJ_IDX_OFFSET
        else:
            sub_idx, obj_idx = nxt[0], fst_idx - OBJ_IDX_OFFSET

        pair_dis = bbox_distance(subjects[sub_idx]["bbox"], objects[obj_idx]["bbox"])
        nearest_dis = float("inf")
        for i in range(N):
            # 取消原先算法中 1对1 匹配的偏置
            # if i in seen_idx or i == sub_idx:continue
            nearest_dis = min(nearest_dis, bbox_distance(subjects[i]["bbox"], objects[obj_idx]["bbox"]))

        if pair_dis >= 3 * nearest_dis:
            seen_idx.add(sub_idx)
            continue

        seen_idx.add(sub_idx)
        seen_idx.add(obj_idx + OBJ_IDX_OFFSET)
        seen_sub_idx.add(sub_idx)

        ret.append(
            {
                "sub_bbox": extract_subject_func(subjects[sub_idx]),
                "obj_bboxes": [extract_object_func(objects[obj_idx])],
                "sub_idx": sub_idx,
            }
        )

    for i in range(len(objects)):
        j = i + OBJ_IDX_OFFSET
        if j in seen_idx:
            continue
        seen_idx.add(j)
        nearest_dis, nearest_sub_idx = float("inf"), -1
        for k in range(len(subjects)):
            dis = bbox_distance(objects[i]["bbox"], subjects[k]["bbox"])
            if dis < nearest_dis:
                nearest_dis = dis
                nearest_sub_idx = k

        for k in range(len(subjects)):
            if k != nearest_sub_idx:
                continue
            if k in seen_sub_idx:
                for kk in range(len(ret)):
                    if ret[kk]["sub_idx"] == k:
                        ret[kk]["obj_bboxes"].append(extract_object_func(objects[i]))
                        break
            else:
                ret.append(
                    {
                        "sub_bbox": extract_subject_func(subjects[k]),
                        "obj_bboxes": [extract_object_func(objects[i])],
                        "sub_idx": k,
                    }
                )
            seen_sub_idx.add(k)
            seen_idx.add(k)

    for i in range(len(subjects)):
        if i in seen_sub_idx:
            continue
        ret.append(
            {
                "sub_bbox": extract_subject_func(subjects[i]),
                "obj_bboxes": [],
                "sub_idx": i,
            }
        )

    return ret


def tie_up_category_by_index(
        get_subjects_func: Callable,
        get_objects_func: Callable,
        extract_subject_func: Callable = None,
        extract_object_func: Callable = None,
        object_block_type: str = "object",
):
    """
    基于index的类别关联方法，用于将主体对象与客体对象进行关联
    客体优先匹配给index最接近的主体，匹配优先级为：
    1. index差值（最高优先级）
    2. bbox边缘距离（相邻边距离）
    3. bbox中心点距离（最低优先级，作为最终tiebreaker）

    参数:
        get_subjects_func: 函数，提取主体对象
        get_objects_func: 函数，提取客体对象
        extract_subject_func: 函数，自定义提取主体属性（默认使用bbox和其他属性）
        extract_object_func: 函数，自定义提取客体属性（默认使用bbox和其他属性）

    返回:
        关联后的对象列表，按主体index升序排列
    """
    subjects = get_subjects_func()
    objects = get_objects_func()

    # 如果没有提供自定义提取函数，使用默认函数
    if extract_subject_func is None:
        extract_subject_func = lambda x: x
    if extract_object_func is None:
        extract_object_func = lambda x: x

    # 初始化结果字典，key为主体索引，value为关联信息
    result_dict = {}

    # 初始化所有主体
    for i, subject in enumerate(subjects):
        result_dict[i] = {
            "sub_bbox": extract_subject_func(subject),
            "obj_bboxes": [],
            "sub_idx": i,
        }

    # 提取所有客体的index集合，用于计算有效index差值
    object_indices = set(obj["index"] for obj in objects)

    def calc_effective_index_diff(obj_index: int, sub_index: int) -> int:
        """
        计算有效的index差值
        有效差值 = 绝对差值 - 区间内其他客体的数量
        即：如果obj_index和sub_index之间的差值是由其他客体造成的，则应该扣除这部分差值
        """
        if obj_index == sub_index:
            return 0

        start, end = min(obj_index, sub_index), max(obj_index, sub_index)
        abs_diff = end - start

        # 计算区间(start, end)内有多少个其他客体的index
        other_objects_count = 0
        for idx in range(start + 1, end):
            if idx in object_indices:
                other_objects_count += 1

        return abs_diff - other_objects_count

    # 为每个客体找到最匹配的主体
    for obj in objects:
        if len(subjects) == 0:
            # 如果没有主体，跳过客体
            continue

        obj_index = obj["index"]
        min_index_diff = float("inf")
        best_subject_indices = []

        # 找出有效index差值最小的所有主体
        for i, subject in enumerate(subjects):
            sub_index = subject["index"]
            index_diff = calc_effective_index_diff(obj_index, sub_index)

            if index_diff < min_index_diff:
                min_index_diff = index_diff
                best_subject_indices = [i]
            elif index_diff == min_index_diff:
                best_subject_indices.append(i)

        if len(best_subject_indices) == 1:
            best_subject_idx = best_subject_indices[0]
        # 如果有多个主体的index差值相同（最多两个），根据边缘距离进行筛选
        elif len(best_subject_indices) == 2:
            # 计算所有候选主体的边缘距离
            edge_distances = [(idx, bbox_distance(obj["bbox"], subjects[idx]["bbox"])) for idx in best_subject_indices]
            edge_dist_diff = abs(edge_distances[0][1] - edge_distances[1][1])

            for idx, edge_dist in edge_distances:
                logger.debug(f"Obj index: {obj_index}, Sub index: {subjects[idx]['index']}, Edge distance: {edge_dist}")

            if edge_dist_diff > 2:
                # 边缘距离差值大于2，匹配边缘距离更小的主体
                best_subject_idx = min(edge_distances, key=lambda x: x[1])[0]
                logger.debug(f"Obj index: {obj_index}, edge_dist_diff > 2, matching to subject with min edge distance, index: {subjects[best_subject_idx]['index']}")
            elif object_block_type == "table_caption":
                # 边缘距离差值<=2且为table_caption，匹配index更大的主体
                best_subject_idx = max(best_subject_indices, key=lambda idx: subjects[idx]["index"])
                logger.debug(f"Obj index: {obj_index}, edge_dist_diff <= 2 and table_caption, matching to later subject with index: {subjects[best_subject_idx]['index']}")
            elif object_block_type.endswith("footnote"):
                # 边缘距离差值<=2且为footnote，匹配index更小的主体
                best_subject_idx = min(best_subject_indices, key=lambda idx: subjects[idx]["index"])
                logger.debug(f"Obj index: {obj_index}, edge_dist_diff <= 2 and footnote, matching to earlier subject with index: {subjects[best_subject_idx]['index']}")
            else:
                # 边缘距离差值<=2 且不适用特殊匹配规则，使用中心点距离匹配
                center_distances = [(idx, bbox_center_distance(obj["bbox"], subjects[idx]["bbox"])) for idx in best_subject_indices]
                for idx, center_dist in center_distances:
                    logger.debug(f"Obj index: {obj_index}, Sub index: {subjects[idx]['index']}, Center distance: {center_dist}")
                best_subject_idx = min(center_distances, key=lambda x: x[1])[0]
        else:
            raise ValueError("More than two subjects have the same minimal index difference, which is unexpected.")

        # 将客体添加到最佳主体的obj_bboxes中
        result_dict[best_subject_idx]["obj_bboxes"].append(extract_object_func(obj))

    # 转换为列表并按主体index排序
    ret = list(result_dict.values())
    ret.sort(key=lambda x: x["sub_idx"])

    return ret