jk-rag-python/mineru/model/mfr/unimernet/Unimernet.py

import torch
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from mineru.utils.boxbase import calculate_iou


class MathDataset(Dataset):
    def __init__(self, image_paths, transform=None):
        self.image_paths = image_paths
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        raw_image = self.image_paths[idx]
        if self.transform:
            image = self.transform(raw_image)
            return image


class UnimernetModel(object):
    def __init__(self, weight_dir, _device_="cpu"):
        from .unimernet_hf import UnimernetModel
        if _device_.startswith("mps") or _device_.startswith("npu") or _device_.startswith("musa"):
            self.model = UnimernetModel.from_pretrained(weight_dir, attn_implementation="eager")
        else:
            self.model = UnimernetModel.from_pretrained(weight_dir)
        self.device = torch.device(_device_)
        self.model.to(self.device)
        if not _device_.startswith("cpu"):
            self.model = self.model.to(dtype=torch.float16)
        self.model.eval()

    @staticmethod
    def _filter_boxes_by_iou(xyxy, conf, cla, iou_threshold=0.8):
        """过滤IOU超过阈值的重叠框，保留置信度较高的框。

        Args:
            xyxy: 框坐标张量，shape为(N, 4)
            conf: 置信度张量，shape为(N,)
            cla: 类别张量，shape为(N,)
            iou_threshold: IOU阈值，默认0.9

        Returns:
            过滤后的xyxy, conf, cla张量
        """
        if len(xyxy) == 0:
            return xyxy, conf, cla

        # 转换为CPU进行处理
        xyxy_cpu = xyxy.cpu()
        conf_cpu = conf.cpu()

        n = len(xyxy_cpu)
        keep = [True] * n

        for i in range(n):
            if not keep[i]:
                continue
            bbox1 = xyxy_cpu[i].tolist()
            for j in range(i + 1, n):
                if not keep[j]:
                    continue
                bbox2 = xyxy_cpu[j].tolist()
                iou = calculate_iou(bbox1, bbox2)
                if iou > iou_threshold:
                    # 保留置信度较高的框
                    if conf_cpu[i] >= conf_cpu[j]:
                        keep[j] = False
                    else:
                        keep[i] = False
                        break  # i被删除，跳出内循环

        keep_indices = [i for i in range(n) if keep[i]]
        if len(keep_indices) == n:
            return xyxy, conf, cla

        keep_indices = torch.tensor(keep_indices, dtype=torch.long)
        return xyxy[keep_indices], conf[keep_indices], cla[keep_indices]

    def predict(self, mfd_res, image):
        formula_list = []
        mf_image_list = []

        # 对检测框进行IOU去重，保留置信度较高的框
        xyxy_filtered, conf_filtered, cla_filtered = self._filter_boxes_by_iou(
            mfd_res.boxes.xyxy, mfd_res.boxes.conf, mfd_res.boxes.cls
        )

        for xyxy, conf, cla in zip(
            xyxy_filtered.cpu(), conf_filtered.cpu(), cla_filtered.cpu()
        ):
            xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
            new_item = {
                "category_id": 13 + int(cla.item()),
                "poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
                "score": round(float(conf.item()), 2),
                "latex": "",
            }
            formula_list.append(new_item)
            bbox_img = image[ymin:ymax, xmin:xmax]
            mf_image_list.append(bbox_img)

        dataset = MathDataset(mf_image_list, transform=self.model.transform)
        dataloader = DataLoader(dataset, batch_size=32, num_workers=0)
        mfr_res = []
        for mf_img in dataloader:
            mf_img = mf_img.to(dtype=self.model.dtype)
            mf_img = mf_img.to(self.device)
            with torch.no_grad():
                output = self.model.generate({"image": mf_img})
            mfr_res.extend(output["fixed_str"])
        for res, latex in zip(formula_list, mfr_res):
            res["latex"] = latex
        return formula_list

    def batch_predict(
            self,
            images_mfd_res: list,
            images: list,
            batch_size: int = 64,
            interline_enable: bool = True,
    ) -> list:
        images_formula_list = []
        mf_image_list = []
        backfill_list = []
        image_info = []  # Store (area, original_index, image) tuples

        # Collect images with their original indices
        for image_index in range(len(images_mfd_res)):
            mfd_res = images_mfd_res[image_index]
            image = images[image_index]
            formula_list = []

            # 对检测框进行IOU去重，保留置信度较高的框
            xyxy_filtered, conf_filtered, cla_filtered = self._filter_boxes_by_iou(
                mfd_res.boxes.xyxy, mfd_res.boxes.conf, mfd_res.boxes.cls
            )

            for idx, (xyxy, conf, cla) in enumerate(zip(
                    xyxy_filtered, conf_filtered, cla_filtered
            )):
                if not interline_enable and cla.item() == 1:
                    continue  # Skip interline regions if not enabled
                xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
                new_item = {
                    "category_id": 13 + int(cla.item()),
                    "poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
                    "score": round(float(conf.item()), 2),
                    "latex": "",
                }
                formula_list.append(new_item)
                bbox_img = image[ymin:ymax, xmin:xmax]
                area = (xmax - xmin) * (ymax - ymin)

                curr_idx = len(mf_image_list)
                image_info.append((area, curr_idx, bbox_img))
                mf_image_list.append(bbox_img)

            images_formula_list.append(formula_list)
            backfill_list += formula_list

        # Stable sort by area
        image_info.sort(key=lambda x: x[0])  # sort by area
        sorted_indices = [x[1] for x in image_info]
        sorted_images = [x[2] for x in image_info]

        # Create mapping for results
        index_mapping = {new_idx: old_idx for new_idx, old_idx in enumerate(sorted_indices)}

        # Create dataset with sorted images
        dataset = MathDataset(sorted_images, transform=self.model.transform)

        # 如果batch_size > len(sorted_images)，则设置为不超过len(sorted_images)的2的幂
        batch_size = min(batch_size, max(1, 2 ** (len(sorted_images).bit_length() - 1))) if sorted_images else 1

        dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=0)

        # Process batches and store results
        mfr_res = []
        # for mf_img in dataloader:

        with tqdm(total=len(sorted_images), desc="MFR Predict") as pbar:
            for index, mf_img in enumerate(dataloader):
                mf_img = mf_img.to(dtype=self.model.dtype)
                mf_img = mf_img.to(self.device)
                with torch.no_grad():
                    output = self.model.generate({"image": mf_img}, batch_size=batch_size)
                mfr_res.extend(output["fixed_str"])

                # 更新进度条，每次增加batch_size，但要注意最后一个batch可能不足batch_size
                current_batch_size = min(batch_size, len(sorted_images) - index * batch_size)
                pbar.update(current_batch_size)

        # Restore original order
        unsorted_results = [""] * len(mfr_res)
        for new_idx, latex in enumerate(mfr_res):
            original_idx = index_mapping[new_idx]
            unsorted_results[original_idx] = latex

        # Fill results back
        for res, latex in zip(backfill_list, unsorted_results):
            res["latex"] = latex

        return images_formula_list