yl-ocr-layout/core/detectors/yolov7.py

from pathlib import Path
from typing import List
from numpy import ndarray
import torch

from core.detectors.base import LayoutDetectorBase
from core.layout import LayoutBox

PROJ_ROOT = Path(__file__).parent.parent.parent
YOLO_DIR = str(PROJ_ROOT / "yolov7")
WEIGHTS = str(PROJ_ROOT / "yiliv7_718.pt")


class Yolov7Detector(LayoutDetectorBase):
    print("======加载 YOLOv7 模型======")
    print(f"是否可使用GPU=======>{torch.cuda.is_available()}")
    _device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    _model = torch.hub.load(YOLO_DIR, "custom", WEIGHTS, source="local").to(
        _device
    )
    print("========>模型加载成功")

    @classmethod
    def predict(cls, img: ndarray, img_size=1824, **kwargs) -> List[LayoutBox]:
        results = cls._model([img], size=img_size)
        return [
            [
                LayoutBox(
                    clazz=int(pred[5]),
                    clazz_name=cls._model.model.names[int(pred[5])],
                    bbox=[
                        int(x) for x in pred[:4].tolist()
                    ],  # convert bbox results to int from float
                    conf=float(pred[4]),
                )
                for pred in result
            ]
            for result in results.xyxy
        ][0]