mule-image/backend/app/services/vision/detect.py

"""
YOLOv8n object detector using raw ONNX Runtime.

Expects {models_dir}/detect/yolov8n.onnx, exported from ultralytics
via bootstrap_models.py. We do NOT ship ultralytics at runtime to
avoid dragging in torch.
"""
import logging
from pathlib import Path

import numpy as np
import onnxruntime as ort

from app.config import VisionSettings
from app.services.vision.base import ObjectDetector, DetectionBox

logger = logging.getLogger(__name__)

_INPUT_SIZE = 640

# COCO class names (80 classes)
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",
]


def _preprocess(image: np.ndarray) -> tuple[np.ndarray, float, float]:
    """Letterbox-resize + normalize to NCHW float32. Returns input tensor
    and scale factors for mapping boxes back to original coords."""
    from PIL import Image

    img = Image.fromarray(image).convert("RGB")
    orig_w, orig_h = img.size

    scale = min(_INPUT_SIZE / orig_w, _INPUT_SIZE / orig_h)
    new_w = int(orig_w * scale)
    new_h = int(orig_h * scale)
    img = img.resize((new_w, new_h), Image.BICUBIC)

    # Paste onto gray canvas
    canvas = np.full((_INPUT_SIZE, _INPUT_SIZE, 3), 114, dtype=np.uint8)
    pad_x = (_INPUT_SIZE - new_w) // 2
    pad_y = (_INPUT_SIZE - new_h) // 2
    canvas[pad_y : pad_y + new_h, pad_x : pad_x + new_w] = np.array(img)

    blob = canvas.astype(np.float32) / 255.0
    blob = blob.transpose(2, 0, 1)[np.newaxis]  # NCHW
    return blob, scale, pad_x, pad_y


def _postprocess(
    outputs: np.ndarray,
    scale: float,
    pad_x: int,
    pad_y: int,
    orig_w: int,
    orig_h: int,
    conf_threshold: float,
    max_detections: int,
) -> list[DetectionBox]:
    """Parse YOLOv8 output (1, 84, N) → list of DetectionBox."""
    # outputs shape: (1, 84, N) where 84 = 4 box coords + 80 class scores
    preds = outputs[0]  # (84, N)
    preds = preds.T  # (N, 84)

    boxes_xywh = preds[:, :4]
    scores = preds[:, 4:]

    class_ids = np.argmax(scores, axis=1)
    confidences = scores[np.arange(len(scores)), class_ids]

    mask = confidences >= conf_threshold
    boxes_xywh = boxes_xywh[mask]
    class_ids = class_ids[mask]
    confidences = confidences[mask]

    if len(confidences) == 0:
        return []

    # Sort by confidence, take top N
    order = np.argsort(-confidences)[:max_detections]
    boxes_xywh = boxes_xywh[order]
    class_ids = class_ids[order]
    confidences = confidences[order]

    results = []
    for i in range(len(confidences)):
        cx, cy, w, h = boxes_xywh[i]
        # Remove letterbox padding and rescale to original image
        x1 = (cx - w / 2 - pad_x) / scale
        y1 = (cy - h / 2 - pad_y) / scale
        x2 = (cx + w / 2 - pad_x) / scale
        y2 = (cy + h / 2 - pad_y) / scale
        # Normalize to 0-1
        bbox = [
            max(0, x1 / orig_w),
            max(0, y1 / orig_h),
            min(1, x2 / orig_w),
            min(1, y2 / orig_h),
        ]
        label = COCO_LABELS[class_ids[i]] if class_ids[i] < len(COCO_LABELS) else f"class_{class_ids[i]}"
        results.append(DetectionBox(label=label, confidence=float(confidences[i]), bbox=bbox))

    return results


class YOLOv8Detector(ObjectDetector):
    def __init__(self, settings: VisionSettings):
        model_path = Path(settings.models_dir) / "detect" / "yolov8n.onnx"

        opts = ort.SessionOptions()
        opts.inter_op_num_threads = 2
        opts.intra_op_num_threads = 2

        logger.info("Loading YOLOv8n from %s", model_path)
        self._session = ort.InferenceSession(str(model_path), opts, providers=["CPUExecutionProvider"])
        self._conf_threshold = settings.detector.min_confidence
        self._max_detections = settings.detector.max_detections

    def detect(self, image: np.ndarray) -> list[DetectionBox]:
        orig_h, orig_w = image.shape[:2]
        blob, scale, pad_x, pad_y = _preprocess(image)
        input_name = self._session.get_inputs()[0].name
        outputs = self._session.run(None, {input_name: blob})[0]
        return _postprocess(
            outputs, scale, pad_x, pad_y, orig_w, orig_h,
            self._conf_threshold, self._max_detections,
        )