yl-ocr-layout/core/predictor.py

from __future__ import annotations
from typing import Dict, List, Optional

import numpy as np


def predict_img(
    img: np.ndarray, model_name: str, img_size: int, **kwargs
) -> List[LayoutBox]:
    """
    根据模型名称预测布局检测框。模型仅在被调用时加载（以节省GPU资源）。
    """
    if model_name == "ocr-layout":
        from core.detectors.yolov7 import Yolov7Detector

        return Yolov7Detector.predict(img, img_size, **kwargs)

    elif model_name == "ocr-layout-paddle":
        from core.detectors.paddle_yolo import PaddleYoloDetector

        return PaddleYoloDetector.predict(
            img,
            conf_threshold=0.3,
            overlaps_iou_threshold=0.85,
            overlaps_max_count=3,
            **kwargs,
        )
    else:
        raise RuntimeError(f"Invalid model name: {model_name}")


class LayoutBox:
    def __init__(
        self,
        clazz: int,
        clazz_name: Optional[str],
        bbox: List[int],
        conf: float,
    ):
        self.clazz = clazz
        self.clazz_name = clazz_name
        self.bbox = bbox
        self.conf = conf

    @property
    def ltrb(self):
        l, t, r, b = self.bbox
        return [int(x) for x in [l, t, r, b]]

    @property
    def area(self):
        l, t, r, b = self.ltrb
        return (r - l + 1) * (b - t + 1)

    def iou(self, other):
        boxA = self.ltrb
        boxB = other.ltrb
        boxA = [int(x) for x in boxA]
        boxB = [int(x) for x in boxB]

        xA = max(boxA[0], boxB[0])
        yA = max(boxA[1], boxB[1])
        xB = min(boxA[2], boxB[2])
        yB = min(boxA[3], boxB[3])

        interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)

        boxAArea = self.area
        boxBArea = other.area

        iou = interArea / float(boxAArea + boxBArea - interArea)

        return iou

    def to_dict(self) -> Dict:
        return {
            "class": self.clazz,
            "class_name": self.clazz_name,
            "bbox": self.bbox,
            "confidence": self.conf,
        }

    def to_service_dict(self) -> Dict:
        """
        返回中间服务所需的格式
        """
        return {
            "class": self.clazz,
            "class_name": self.clazz_name,
            "bbox": self.ltrb,
            "confidence": self.conf,
        }

    @classmethod
    def from_dict(cls, d) -> LayoutBox:
        return cls(
            clazz=d["class"],
            clazz_name=d["class_name"],
            bbox=d["bbox"],
            conf=d["confidence"],
        )

    def __repr__(self):
        return f"LayoutBox(class={self.clazz}, class_name={self.clazz_name}, bbox={self.bbox}, conf={self.conf})"