src/inference.py

"""Inference core shared by the Streamlit app, scripts, and tests.

Keeps every consumer on one code path: one weights default, one device-selection
rule, one filtering behavior. Deliberately imports no Streamlit — this module
must stay importable headless (tests, batch scripts, CI).
"""

from __future__ import annotations

import logging
from collections.abc import Sequence
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from ultralytics import YOLO
from ultralytics.engine.results import Results

logger = logging.getLogger(__name__)

DEFAULT_WEIGHTS = "runs/detect/yolov8n_v1_train/weights/best.pt"

# Anything Ultralytics accepts as a single-image source.
ImageSource = str | Path | np.ndarray | Image.Image


def available_devices() -> list[str]:
    """Inference devices usable on this machine, preferred-first after cpu."""
    devices = ["cpu"]
    if torch.cuda.is_available():
        devices.append("cuda")
    if torch.backends.mps.is_available():
        devices.append("mps")
    return devices


def pick_device() -> str:
    """Fastest available device: cuda > mps > cpu."""
    devices = available_devices()
    return devices[-1] if len(devices) > 1 else "cpu"


def load_model(weights: str | Path = DEFAULT_WEIGHTS, device: str | None = None) -> YOLO:
    """Load YOLO weights onto `device` (auto-picks the fastest when None)."""
    device = device or pick_device()
    model = YOLO(weights)
    model.to(device)
    logger.info("Loaded %s on %s", weights, device)
    return model


def detect(
    model: YOLO,
    image: ImageSource,
    conf: float = 0.25,
    classes: Sequence[int] | None = None,
    imgsz: int | None = None,
) -> Results:
    """Run detection on one image, filtered to `conf` and (optionally) `classes`.

    Filtering uses Ultralytics' native predict-time `conf`/`classes` arguments,
    which apply before NMS — unlike post-hoc box masking, low-confidence boxes
    can't suppress real ones first.

    `imgsz` overrides the model's inference resolution (default 640). Lowering it
    to 480/320 trades accuracy for a large CPU speedup — the difference between
    frozen and watchable video on weak hardware like HF Spaces cpu-basic.
    """
    kwargs = {"imgsz": imgsz} if imgsz else {}
    results = model.predict(
        image, conf=conf, classes=list(classes) if classes else None, verbose=False, **kwargs
    )
    return results[0]


def track(
    model: YOLO,
    frame: ImageSource,
    conf: float = 0.25,
    classes: Sequence[int] | None = None,
    imgsz: int | None = None,
) -> Results:
    """Like detect(), but with ByteTrack identity persistence across calls.

    Feed consecutive frames of ONE stream to ONE model instance: `persist=True`
    carries tracker state on the model object, so interleaving streams (or
    switching videos without a fresh model) corrupts identities. Detections
    carry `.boxes.id` (may be None for not-yet-confirmed tracks).
    """
    kwargs = {"imgsz": imgsz} if imgsz else {}
    results = model.track(
        frame,
        persist=True,
        tracker="bytetrack.yaml",
        conf=conf,
        classes=list(classes) if classes else None,
        verbose=False,
        **kwargs,
    )
    return results[0]


def reset_tracker(model: YOLO) -> None:
    """Clear ByteTrack state so a new video starts with fresh identities.

    A cached/reused model carries tracker state on its predictor; without a
    reset, track ids from the previous stream bleed into the next one. No-op
    if the model has never tracked (predictor/trackers don't exist yet).
    """
    trackers = getattr(getattr(model, "predictor", None), "trackers", None) or []
    for t in trackers:
        if hasattr(t, "reset"):
            t.reset()


def sightings(result: Results) -> list[tuple[int | None, str, float]]:
    """(track_id, class_name, confidence) triples for one tracked frame."""
    out: list[tuple[int | None, str, float]] = []
    boxes = result.boxes
    if boxes is None:
        return out
    ids = boxes.id.int().tolist() if boxes.id is not None else [None] * len(boxes)
    for track_id, cls_idx, conf in zip(ids, boxes.cls.int().tolist(), boxes.conf.tolist(), strict=True):
        out.append((track_id, result.names[cls_idx], float(conf)))
    return out


def annotate(result: Results) -> np.ndarray:
    """Draw the detections onto the image and return it as an ndarray.

    Channel order follows the input: ndarray sources come back as given (the app
    passes RGB and displays RGB); path sources come back BGR per Ultralytics.
    """
    return result.plot()


def detect_and_annotate(
    model: YOLO,
    image: ImageSource,
    conf: float = 0.25,
    classes: Sequence[int] | None = None,
    imgsz: int | None = None,
) -> np.ndarray:
    """Convenience wrapper: detect() then annotate()."""
    return annotate(detect(model, image, conf=conf, classes=classes, imgsz=imgsz))