providers/object_detection/yolov8.py

"""
YOLOv8 object detection provider (ONNX).

Uses YOLOv8n (nano) — 6MB model, 80 COCO classes, designed for
real-time CPU inference.  ~150ms per image on CPU.

Design:
  - Model loaded LAZILY via cores.onnx.get_session()
  - Preprocessing: letterbox resize to 640x640 (cores.vision.letterbox)
  - Postprocessing: NMS + box scaling (cores.vision.nms + scale_boxes)
  - If onnxruntime is not installed, is_available() returns False

License: AGPL-3.0 (model weights freely usable; commercial license available)
"""

from __future__ import annotations

import numpy as np

from config.settings import Settings, settings as _default_settings
from cores.onnx import is_onnx_available, get_session, ensure_model
from cores.vision import letterbox, nms, xywh2xyxy, scale_boxes
from pipeline.feature_extraction import PipelineOutput
from providers.base import BaseProvider, ProviderCapability


# COCO 80-class labels
COCO_LABELS = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
    "truck", "boat", "traffic light", "fire hydrant", "stop sign",
    "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep",
    "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
    "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
    "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
    "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork",
    "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
    "couch", "potted plant", "bed", "dining table", "toilet", "tv",
    "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
    "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
    "scissors", "teddy bear", "hair drier", "toothbrush",
]


class YOLOv8Provider(BaseProvider):
    name = "yolov8"
    capability = ProviderCapability.OBJECT_DETECTION

    MODEL_FILE = "yolov8n.onnx"
    INPUT_SIZE = (640, 640)
    CONFIDENCE_THRESHOLD = 0.25
    IOU_THRESHOLD = 0.45

    def __init__(self, settings: Settings | None = None) -> None:
        super().__init__(settings=settings or _default_settings)
        self._available = is_onnx_available()

    def is_available(self) -> bool:
        return self._available

    def _run(self, pipeline_output: PipelineOutput) -> tuple[dict, dict]:
        if not self._available:
            raise RuntimeError("onnxruntime not installed")

        model_file = ensure_model(self.MODEL_FILE, settings=self._settings)
        session = get_session(model_file, self._settings)

        img: np.ndarray = pipeline_output.image
        h, w = img.shape[:2]

        # Preprocess: letterbox to 640x640
        padded, scale, pad = letterbox(img, self.INPUT_SIZE)
        import cv2
        rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
        normalized = rgb.astype(np.float32) / 255.0
        nchw = normalized.transpose(2, 0, 1)[None]

        # Inference
        output = session.run_single(nchw)
        # YOLOv8 output shape: (1, 84, 8400) → transpose to (8400, 84)
        predictions = output[0].T  # (N, 84) = [cx, cy, w, h, 80 class scores]

        # Filter by confidence
        scores = predictions[:, 4:].max(axis=1)
        class_ids = predictions[:, 4:].argmax(axis=1)
        mask = scores >= self.CONFIDENCE_THRESHOLD
        if not mask.any():
            raw = {"num_objects": 0, "model": "yolov8n"}
            normalized = {"objects": [], "model": "yolov8n"}
            return raw, normalized

        filtered = predictions[mask]
        filtered_scores = scores[mask]
        filtered_classes = class_ids[mask]

        # Convert cx,cy,w,h → x1,y1,x2,y2
        boxes = xywh2xyxy(filtered[:, :4])

        # NMS
        keep = nms(boxes, filtered_scores, self.IOU_THRESHOLD)
        boxes = boxes[keep]
        filtered_scores = filtered_scores[keep]
        filtered_classes = filtered_classes[keep]

        # Scale back to original image
        boxes = scale_boxes(boxes, scale, pad, (h, w))

        # Build output
        objects: list[dict] = []
        for box, score, cls_id in zip(boxes, filtered_scores, filtered_classes):
            label = COCO_LABELS[int(cls_id)] if int(cls_id) < len(COCO_LABELS) else f"class_{int(cls_id)}"
            objects.append({
                "label": label,
                "confidence": round(float(score), 4),
                "box": {
                    "x": int(box[0]), "y": int(box[1]),
                    "w": int(box[2] - box[0]), "h": int(box[3] - box[1]),
                },
            })

        raw = {
            "num_objects": len(objects),
            "model": "yolov8n",
            "input_size": self.INPUT_SIZE,
        }
        normalized = {
            "objects": objects,
            "model": "yolov8n",
        }
        return raw, normalized