tiny_trigger/detector.py

from __future__ import annotations

import re
from typing import Any, Protocol

from .models import Detection


def parse_class_prompt(class_prompt: str | list[str] | tuple[str, ...]) -> list[str]:
    if isinstance(class_prompt, (list, tuple)):
        parts = [str(item) for item in class_prompt]
    else:
        parts = re.split(r"[,;\n]+", class_prompt)

    seen: set[str] = set()
    classes: list[str] = []
    for part in parts:
        label = " ".join(part.strip().split())
        key = label.lower()
        if label and key not in seen:
            seen.add(key)
            classes.append(label)
    return classes


class Detector(Protocol):
    class_names: list[str]

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        ...


class GreedyIoUTracker:
    """Assign stable track IDs to detections by greedy IoU matching across frames.

    Detection stays the source of truth: every detection is returned, matched to an
    existing track when their same-label boxes overlap enough, or given a fresh ID
    otherwise. Nothing is dropped for failing to match — unlike a confirm-before-emit
    tracker (e.g. ByteTrack), which withholds unconfirmed detections and lowers recall.
    A track that goes unmatched for more than ``max_age`` frames is forgotten.
    """

    def __init__(self, *, iou_threshold: float = 0.3, max_age: int = 2) -> None:
        self.iou_threshold = iou_threshold
        self.max_age = max_age
        self._tracks: dict[int, dict[str, Any]] = {}
        self._next_id = 1

    def assign(self, detections: list[Detection]) -> list[Detection]:
        for track in self._tracks.values():
            track["age"] += 1
        assigned: list[Detection] = []
        used: set[int] = set()
        for detection in sorted(detections, key=lambda item: item.confidence, reverse=True):
            label = detection.label.strip().lower()
            best_id: int | None = None
            best_iou = self.iou_threshold
            for track_id, track in self._tracks.items():
                if track_id in used or track["label"] != label:
                    continue
                iou = _box_iou(detection.bbox_xyxy_norm, track["bbox"])
                if iou >= best_iou:
                    best_iou = iou
                    best_id = track_id
            if best_id is None:
                best_id = self._next_id
                self._next_id += 1
            used.add(best_id)
            self._tracks[best_id] = {"bbox": detection.bbox_xyxy_norm, "label": label, "age": 0}
            assigned.append(detection.model_copy(update={"track_id": best_id}))
        self._tracks = {track_id: track for track_id, track in self._tracks.items() if track["age"] <= self.max_age}
        return assigned


class UltralyticsYOLOEDetector:
    def __init__(
        self,
        *,
        class_names: list[str],
        model_name: str = "yoloe-26s-seg.pt",
        device: str | None = None,
        tracking_enabled: bool = False,
    ) -> None:
        if not class_names:
            raise ValueError("YOLOE needs at least one open-vocabulary class.")

        try:
            from ultralytics import YOLOE
        except ImportError as exc:  # pragma: no cover - optional heavy dependency
            raise RuntimeError("Install ultralytics to use the YOLOE detector.") from exc

        self.class_names = class_names
        self.model_name = model_name
        self.device = device
        self.tracking_enabled = tracking_enabled
        self._tracker = GreedyIoUTracker() if tracking_enabled else None
        self.model = YOLOE(model_name)
        self.model.set_classes(class_names)

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        kwargs: dict[str, Any] = {"conf": confidence, "verbose": False}
        if self.device:
            kwargs["device"] = self.device
        if image_size:
            kwargs["imgsz"] = image_size
        if max_detections:
            kwargs["max_det"] = max_detections

        results = self.model.predict(frame, **kwargs)
        if not results:
            return []
        detections = detections_from_ultralytics_result(
            results[0],
            frame_shape=frame.shape,
            frame_index=frame_index,
            timestamp_sec=timestamp_sec,
            fallback_names=self.class_names,
        )
        if self._tracker is not None:
            detections = self._tracker.assign(detections)
        return detections


def detections_from_ultralytics_result(
    result: Any,
    *,
    frame_shape: tuple[int, int, int],
    frame_index: int,
    timestamp_sec: float,
    fallback_names: list[str],
) -> list[Detection]:
    boxes = getattr(result, "boxes", None)
    if boxes is None or boxes.xyxy is None:
        return []

    height, width = frame_shape[:2]
    xyxy_values = boxes.xyxy.cpu().tolist()
    confidences = boxes.conf.cpu().tolist()
    class_ids = boxes.cls.cpu().tolist()
    raw_track_ids = getattr(boxes, "id", None)
    track_ids = raw_track_ids.cpu().tolist() if raw_track_ids is not None else [None] * len(xyxy_values)
    names = getattr(result, "names", {}) or {}

    detections: list[Detection] = []
    for bbox, score, class_id, track_id in zip(xyxy_values, confidences, class_ids, track_ids):
        label = _label_from_names(names, int(class_id), fallback_names)
        detections.append(
            Detection(
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
                label=label,
                confidence=float(score),
                bbox_xyxy=tuple(float(value) for value in bbox),
                bbox_xyxy_norm=_normalize_box(bbox, width, height),
                track_id=int(track_id) if track_id is not None else None,
            )
        )
    return suppress_duplicate_detections(detections)


def suppress_duplicate_detections(
    detections: list[Detection],
    *,
    iou_threshold: float = 0.8,
) -> list[Detection]:
    """Collapse heavily overlapping same-label boxes into one.

    Keeps the highest-confidence box's geometry/label/score, but carries the
    *oldest* track_id in the overlap group. Ultralytics assigns track IDs from
    an incrementing counter, so the smallest ID is the longest-lived track;
    preferring it keeps identity stable across frames instead of letting it flip
    with per-frame confidence (which otherwise caused cooldown/count misfires).
    """
    kept: list[Detection] = []
    for detection in sorted(detections, key=lambda item: item.confidence, reverse=True):
        match_index = next(
            (
                index
                for index, existing in enumerate(kept)
                if _same_label(detection, existing)
                and _box_iou(detection.bbox_xyxy_norm, existing.bbox_xyxy_norm) >= iou_threshold
            ),
            None,
        )
        if match_index is None:
            kept.append(detection)
            continue
        existing = kept[match_index]
        oldest_id = _oldest_track_id(existing.track_id, detection.track_id)
        if oldest_id != existing.track_id:
            kept[match_index] = existing.model_copy(update={"track_id": oldest_id})
    return sorted(kept, key=lambda item: (item.frame_index, item.label, item.bbox_xyxy_norm))


def _label_from_names(names: Any, class_id: int, fallback_names: list[str]) -> str:
    if isinstance(names, dict) and class_id in names:
        return str(names[class_id])
    if isinstance(names, list) and 0 <= class_id < len(names):
        return str(names[class_id])
    if 0 <= class_id < len(fallback_names):
        return fallback_names[class_id]
    return f"class_{class_id}"


def _normalize_box(bbox: list[float], width: int, height: int) -> tuple[float, float, float, float]:
    x1, y1, x2, y2 = bbox
    if width <= 0 or height <= 0:
        return (0.0, 0.0, 0.0, 0.0)
    return (
        _clamp01(float(x1) / width),
        _clamp01(float(y1) / height),
        _clamp01(float(x2) / width),
        _clamp01(float(y2) / height),
    )


def _clamp01(value: float) -> float:
    return max(0.0, min(1.0, value))


def _same_label(left: Detection, right: Detection) -> bool:
    return left.label.strip().lower() == right.label.strip().lower()


def _oldest_track_id(left: int | None, right: int | None) -> int | None:
    ids = [value for value in (left, right) if value is not None]
    if not ids:
        return None
    return min(ids)


def _box_iou(
    left: tuple[float, float, float, float],
    right: tuple[float, float, float, float],
) -> float:
    ax1, ay1, ax2, ay2 = left
    bx1, by1, bx2, by2 = right
    intersection_width = max(0.0, min(ax2, bx2) - max(ax1, bx1))
    intersection_height = max(0.0, min(ay2, by2) - max(ay1, by1))
    intersection = intersection_width * intersection_height
    if intersection <= 0:
        return 0.0
    left_area = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
    right_area = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
    union = left_area + right_area - intersection
    if union <= 0:
        return 0.0
    return intersection / union