yl-ocr-layout/core/layout.py

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

import numpy as np


class LayoutBox:
    def __init__(
        self,
        clazz: int,
        bbox: List[int],
        conf: float,
        clazz_name: Optional[str] = None,
    ):
        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})"