MinerU/mineru/model/table/rec/unet_table/table_recover.py

from typing import Dict, List, Tuple

import numpy as np


class TableRecover:
    def __init__(
        self,
    ):
        pass

    def __call__(
        self, polygons: np.ndarray, rows_thresh=10, col_thresh=15
    ) -> Dict[int, Dict]:
        rows = self.get_rows(polygons, rows_thresh)
        longest_col, each_col_widths, col_nums = self.get_benchmark_cols(
            rows, polygons, col_thresh
        )
        each_row_heights, row_nums = self.get_benchmark_rows(rows, polygons)
        table_res, logic_points_dict = self.get_merge_cells(
            polygons,
            rows,
            row_nums,
            col_nums,
            longest_col,
            each_col_widths,
            each_row_heights,
        )
        logic_points = np.array(
            [logic_points_dict[i] for i in range(len(polygons))]
        ).astype(np.int32)
        return table_res, logic_points

    @staticmethod
    def get_rows(polygons: np.array, rows_thresh=10) -> Dict[int, List[int]]:
        """对每个框进行行分类，框定哪个是一行的"""
        y_axis = polygons[:, 0, 1]
        if y_axis.size == 1:
            return {0: [0]}

        concat_y = np.array(list(zip(y_axis, y_axis[1:])))
        minus_res = concat_y[:, 1] - concat_y[:, 0]

        result = {}
        split_idxs = np.argwhere(abs(minus_res) > rows_thresh).squeeze()
        # 如果都在一行，则将所有下标设置为同一行
        if split_idxs.size == 0:
            return {0: [i for i in range(len(y_axis))]}
        if split_idxs.ndim == 0:
            split_idxs = split_idxs[None, ...]

        if max(split_idxs) != len(minus_res):
            split_idxs = np.append(split_idxs, len(minus_res))

        start_idx = 0
        for row_num, idx in enumerate(split_idxs):
            if row_num != 0:
                start_idx = split_idxs[row_num - 1] + 1
            result.setdefault(row_num, []).extend(range(start_idx, idx + 1))

        # 计算每一行相邻cell的iou，如果大于0.2，则合并为同一个cell
        return result

    def get_benchmark_cols(
        self, rows: Dict[int, List], polygons: np.ndarray, col_thresh=15
    ) -> Tuple[np.ndarray, List[float], int]:
        longest_col = max(rows.values(), key=lambda x: len(x))
        longest_col_points = polygons[longest_col]
        longest_x_start = list(longest_col_points[:, 0, 0])
        longest_x_end = list(longest_col_points[:, 2, 0])
        min_x = longest_x_start[0]
        max_x = longest_x_end[-1]

        # 根据当前col的起始x坐标，更新col的边界
        # 2025.2.22 --- 解决最长列可能漏掉最后一列的问题
        def update_longest_col(col_x_list, cur_v, min_x_, max_x_, insert_last):
            for i, v in enumerate(col_x_list):
                if cur_v - col_thresh <= v <= cur_v + col_thresh:
                    break
                if cur_v < min_x_:
                    col_x_list.insert(0, cur_v)
                    min_x_ = cur_v
                    break
                if cur_v > max_x_:
                    if insert_last:
                        col_x_list.append(cur_v)
                    max_x_ = cur_v
                    break
                if cur_v < v:
                    col_x_list.insert(i, cur_v)
                    break
            return min_x_, max_x_

        for row_value in rows.values():
            cur_row_start = list(polygons[row_value][:, 0, 0])
            cur_row_end = list(polygons[row_value][:, 2, 0])
            for idx, (cur_v_start, cur_v_end) in enumerate(
                zip(cur_row_start, cur_row_end)
            ):
                min_x, max_x = update_longest_col(
                    longest_x_start, cur_v_start, min_x, max_x, True
                )
                min_x, max_x = update_longest_col(
                    longest_x_start, cur_v_end, min_x, max_x, False
                )

        longest_x_start = np.array(longest_x_start)
        each_col_widths = (longest_x_start[1:] - longest_x_start[:-1]).tolist()
        each_col_widths.append(max_x - longest_x_start[-1])
        col_nums = longest_x_start.shape[0]
        return longest_x_start, each_col_widths, col_nums

    def get_benchmark_rows(
        self, rows: Dict[int, List], polygons: np.ndarray
    ) -> Tuple[np.ndarray, List[float], int]:
        leftmost_cell_idxs = [v[0] for v in rows.values()]
        benchmark_x = polygons[leftmost_cell_idxs][:, 0, 1]

        each_row_widths = (benchmark_x[1:] - benchmark_x[:-1]).tolist()

        # 求出最后一行cell中，最大的高度作为最后一行的高度
        bottommost_idxs = list(rows.values())[-1]
        bottommost_boxes = polygons[bottommost_idxs]
        # fix self.compute_L2(v[3, :], v[0, :]), v为逆时针，即v[3]为右上，v[0]为左上,v[1]为左下
        max_height = max([self.compute_L2(v[1, :], v[0, :]) for v in bottommost_boxes])
        each_row_widths.append(max_height)

        row_nums = benchmark_x.shape[0]
        return each_row_widths, row_nums

    @staticmethod
    def compute_L2(a1: np.ndarray, a2: np.ndarray) -> float:
        return np.linalg.norm(a2 - a1)

    def get_merge_cells(
        self,
        polygons: np.ndarray,
        rows: Dict,
        row_nums: int,
        col_nums: int,
        longest_col: np.ndarray,
        each_col_widths: List[float],
        each_row_heights: List[float],
    ) -> Dict[int, Dict[int, int]]:
        col_res_merge, row_res_merge = {}, {}
        logic_points = {}
        merge_thresh = 10
        for cur_row, col_list in rows.items():
            one_col_result, one_row_result = {}, {}
            for one_col in col_list:
                box = polygons[one_col]
                box_width = self.compute_L2(box[3, :], box[0, :])

                # 不一定是从0开始的，应该综合已有值和x坐标位置来确定起始位置
                loc_col_idx = np.argmin(np.abs(longest_col - box[0, 0]))
                col_start = max(sum(one_col_result.values()), loc_col_idx)

                # 计算合并多少个列方向单元格
                for i in range(col_start, col_nums):
                    col_cum_sum = sum(each_col_widths[col_start : i + 1])
                    if i == col_start and col_cum_sum > box_width:
                        one_col_result[one_col] = 1
                        break
                    elif abs(col_cum_sum - box_width) <= merge_thresh:
                        one_col_result[one_col] = i + 1 - col_start
                        break
                    # 这里必须进行修正，不然会出现超越阈值范围后列交错
                    elif col_cum_sum > box_width:
                        idx = (
                            i
                            if abs(col_cum_sum - box_width)
                            < abs(col_cum_sum - each_col_widths[i] - box_width)
                            else i - 1
                        )
                        one_col_result[one_col] = idx + 1 - col_start
                        break
                else:
                    one_col_result[one_col] = col_nums - col_start
                col_end = one_col_result[one_col] + col_start - 1
                box_height = self.compute_L2(box[1, :], box[0, :])
                row_start = cur_row
                for j in range(row_start, row_nums):
                    row_cum_sum = sum(each_row_heights[row_start : j + 1])
                    # box_height 不确定是几行的高度，所以要逐个试验，找一个最近的几行的高
                    # 如果第一次row_cum_sum就比box_height大，那么意味着？丢失了一行
                    if j == row_start and row_cum_sum > box_height:
                        one_row_result[one_col] = 1
                        break
                    elif abs(box_height - row_cum_sum) <= merge_thresh:
                        one_row_result[one_col] = j + 1 - row_start
                        break
                    # 这里必须进行修正，不然会出现超越阈值范围后行交错
                    elif row_cum_sum > box_height:
                        idx = (
                            j
                            if abs(row_cum_sum - box_height)
                            < abs(row_cum_sum - each_row_heights[j] - box_height)
                            else j - 1
                        )
                        one_row_result[one_col] = idx + 1 - row_start
                        break
                else:
                    one_row_result[one_col] = row_nums - row_start
                row_end = one_row_result[one_col] + row_start - 1
                logic_points[one_col] = np.array(
                    [row_start, row_end, col_start, col_end]
                )
            col_res_merge[cur_row] = one_col_result
            row_res_merge[cur_row] = one_row_result

        res = {}
        for i, (c, r) in enumerate(zip(col_res_merge.values(), row_res_merge.values())):
            res[i] = {k: [cc, r[k]] for k, cc in c.items()}
        return res, logic_points