app/services/object_detector.py

import io
import numpy as np
import onnxruntime as ort
from PIL import Image
from huggingface_hub import hf_hub_download
from ..config import (
    YOLO_MODEL_ID,
    YOLO_MODEL_FILENAME,
    INTRA_OP_NUM_THREADS,
    INTER_OP_NUM_THREADS,
)

_yolo_session = None


def _get_yolo_session() -> ort.InferenceSession:
    """Lazy-load the YOLO ONNX session on first use instead of at import time."""
    global _yolo_session
    if _yolo_session is not None:
        return _yolo_session

    print(f"[object_detector] Downloading YOLO model '{YOLO_MODEL_ID}/{YOLO_MODEL_FILENAME}'...")
    model_file = hf_hub_download(repo_id=YOLO_MODEL_ID, filename=YOLO_MODEL_FILENAME)

    ort_options = ort.SessionOptions()
    ort_options.intra_op_num_threads = INTRA_OP_NUM_THREADS
    ort_options.inter_op_num_threads = INTER_OP_NUM_THREADS
    _yolo_session = ort.InferenceSession(
        model_file, sess_options=ort_options, providers=["CPUExecutionProvider"]
    )
    print("[object_detector] YOLO model loaded successfully!")
    return _yolo_session


def _letterbox(
    img: Image.Image, new_shape: tuple[int, int] = (640, 640)
) -> tuple[Image.Image, float, tuple[int, int]]:
    """Resize image with letterboxing (preserve aspect ratio, pad with gray)."""
    w, h = img.size
    r = min(new_shape[0] / h, new_shape[1] / w)
    new_unpad = (int(round(w * r)), int(round(h * r)))
    dw = (new_shape[1] - new_unpad[0]) / 2
    dh = (new_shape[0] - new_unpad[1]) / 2

    resized = img.resize(new_unpad, Image.Resampling.BILINEAR)

    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    padded = Image.new("RGB", (new_shape[1], new_shape[0]), (114, 114, 114))
    padded.paste(resized, (left, top))
    return padded, r, (dw, dh)


def _nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> np.ndarray:
    """Non-Maximum Suppression. Returns indices of kept boxes."""
    if len(boxes) == 0:
        return np.array([], dtype=int)

    order = scores.argsort()[::-1]
    keep = []

    while len(order) > 0:
        i = order[0]
        keep.append(i)
        if len(order) == 1:
            break

        xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
        yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
        xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
        yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])

        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
        area_rest = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (
            boxes[order[1:], 3] - boxes[order[1:], 1]
        )
        iou = inter / (area_i + area_rest - inter + 1e-7)

        inds = np.where(iou <= iou_threshold)[0]
        order = order[inds + 1]

    return np.array(keep, dtype=int)


def _preprocess(contents: bytes) -> tuple[np.ndarray, float, tuple[int, int], tuple[int, int]]:
    """Decode image, letterbox to 640x640, return NCHW float32 tensor + metadata."""
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    orig_w, orig_h = image.size

    img, ratio, (dw, dh) = _letterbox(image)
    arr = np.array(img, dtype=np.float32) / 255.0
    arr = arr.transpose(2, 0, 1)  # HWC -> CHW
    arr = np.expand_dims(arr, 0)  # add batch dim -> NCHW
    return arr, ratio, (dw, dh), (orig_w, orig_h)


def _postprocess(
    output: np.ndarray,
    conf_threshold: float,
    iou_threshold: float,
    ratio: float,
    pad: tuple[int, int],
    orig_size: tuple[int, int],
    max_detections: int,
) -> list[dict]:
    """Parse YOLO output tensor, apply NMS, return list of detections."""
    # output shape: (1, 84, 8400) -> transpose to (8400, 84)
    preds = output[0].T  # (8400, 84)

    boxes_xywh = preds[:, :4]
    class_scores = preds[:, 4:]
    max_scores = class_scores.max(axis=1)
    class_ids = class_scores.argmax(axis=1)

    # Filter by confidence
    mask = max_scores > conf_threshold
    boxes_xywh = boxes_xywh[mask]
    max_scores = max_scores[mask]
    class_ids = class_ids[mask]

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

    # Convert xywh -> xyxy
    x1 = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2
    y1 = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2
    x2 = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2
    y2 = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2
    boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)

    # NMS
    keep = _nms(boxes_xyxy, max_scores, iou_threshold)

    # Take top max_detections
    if len(keep) > max_detections:
        keep = keep[:max_detections]

    # Scale back to original image coordinates (undo letterbox padding + ratio)
    dw, dh = pad
    orig_w, orig_h = orig_size
    detections = []
    for idx in keep:
        bx1, by1, bx2, by2 = boxes_xyxy[idx]
        # Undo letterbox padding
        bx1 = (bx1 - dw) / ratio
        by1 = (by1 - dh) / ratio
        bx2 = (bx2 - dw) / ratio
        by2 = (by2 - dh) / ratio
        # Clamp to original image bounds
        bx1 = max(0, min(bx1, orig_w))
        by1 = max(0, min(by1, orig_h))
        bx2 = max(0, min(bx2, orig_w))
        by2 = max(0, min(by2, orig_h))
        detections.append(
            {
                "bbox": [float(bx1), float(by1), float(bx2), float(by2)],
                "confidence": float(max_scores[idx]),
                "class_id": int(class_ids[idx]),
            }
        )

    return detections


def detect_objects(
    contents: bytes,
    conf_threshold: float = 0.25,
    iou_threshold: float = 0.45,
    max_detections: int = 3,
) -> list[dict]:
    """Run YOLO detection on image bytes. Returns list of detections sorted by confidence."""
    session = _get_yolo_session()
    pixel_values, ratio, pad, orig_size = _preprocess(contents)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    raw_output = session.run([output_name], {input_name: pixel_values})[0]
    detections = _postprocess(
        raw_output, conf_threshold, iou_threshold, ratio, pad, orig_size, max_detections
    )
    return detections


def crop_objects(
    contents: bytes, detections: list[dict], padding_ratio: float = 0.10
) -> list[bytes]:
    """Crop detected bounding box regions from the original image. Returns list of JPEG bytes."""
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    orig_w, orig_h = image.size
    crops = []

    for det in detections:
        x1, y1, x2, y2 = det["bbox"]
        bw, bh = x2 - x1, y2 - y1
        pad_x, pad_y = bw * padding_ratio, bh * padding_ratio

        cx1 = max(0, int(x1 - pad_x))
        cy1 = max(0, int(y1 - pad_y))
        cx2 = min(orig_w, int(x2 + pad_x))
        cy2 = min(orig_h, int(y2 + pad_y))

        if cx2 <= cx1 or cy2 <= cy1:
            continue

        crop = image.crop((cx1, cy1, cx2, cy2))
        buf = io.BytesIO()
        crop.save(buf, format="JPEG", quality=95)
        crops.append(buf.getvalue())

    return crops