mingxi_platform/mingxi-vision/engine/ultralytics_adapter.py

import time
from pathlib import Path

import numpy as np

from .base import BaseInferenceEngine, DefectBox, DetectResult
from .defect_classes import get_class_info


class UltralyticsAdapter(BaseInferenceEngine):

    def __init__(self, model_path: str):
        from ultralytics import YOLO
        self._model_path = str(model_path)
        self._model = YOLO(self._model_path)
        self._version = Path(model_path).stem

    def detect(self, image: np.ndarray, conf: float = 0.45) -> DetectResult:
        h, w = image.shape[:2]
        t0 = time.perf_counter()
        results = self._model(image, conf=conf, verbose=False)
        duration_ms = (time.perf_counter() - t0) * 1000

        defects = []
        for r in results:
            if r.boxes is None:
                continue
            for box in r.boxes:
                class_id = int(box.cls[0])
                confidence = float(box.conf[0])
                xyxy = box.xyxy[0].tolist()
                info = get_class_info(class_id)
                defects.append(DefectBox(
                    class_id=class_id,
                    class_name=info["name"],
                    class_name_zh=info["zh"],
                    confidence=round(confidence, 4),
                    severity=info["severity"],
                    box_xyxy=[round(v, 1) for v in xyxy],
                ))

        return DetectResult(
            defects=defects,
            duration_ms=round(duration_ms, 1),
            image_width=w,
            image_height=h,
            model_version=self._version,
        )

    def warmup(self) -> None:
        dummy = np.zeros((640, 640, 3), dtype=np.uint8)
        self.detect(dummy, conf=0.45)

    @property
    def model_version(self) -> str:
        return self._version