tests/test_detector.py

from __future__ import annotations

from tiny_trigger.detector import (
    GreedyIoUTracker,
    UltralyticsYOLOEDetector,
    detections_from_ultralytics_result,
    parse_class_prompt,
    suppress_duplicate_detections,
)
from tiny_trigger.models import Detection


class TensorLike:
    def __init__(self, values):
        self.values = values

    def cpu(self):
        return self

    def tolist(self):
        return self.values


class BoxesLike:
    def __init__(self, *, ids=None):
        self.xyxy = TensorLike([[10.0, 20.0, 30.0, 40.0], [40.0, 50.0, 70.0, 90.0]])
        self.conf = TensorLike([0.91, 0.82])
        self.cls = TensorLike([0, 1])
        self.id = TensorLike(ids) if ids is not None else None


class ResultLike:
    def __init__(self, *, ids=None):
        self.boxes = BoxesLike(ids=ids)
        self.names = {0: "person", 1: "car"}


class FrameLike:
    shape = (100, 100, 3)


def test_parse_class_prompt_splits_and_dedupes() -> None:
    assert parse_class_prompt(" cat, feeder robot\npackage; cat ") == ["cat", "feeder robot", "package"]


def detection(label: str, confidence: float, box: tuple[float, float, float, float]) -> Detection:
    return Detection(
        frame_index=0,
        timestamp_sec=0.0,
        label=label,
        confidence=confidence,
        bbox_xyxy=(0.0, 0.0, 10.0, 10.0),
        bbox_xyxy_norm=box,
    )


def test_suppress_duplicate_detections_keeps_highest_confidence_same_label_box() -> None:
    detections = [
        detection("person", 0.62, (0.10, 0.10, 0.50, 0.50)),
        detection("person", 0.91, (0.11, 0.11, 0.51, 0.51)),
        detection("person", 0.80, (0.60, 0.10, 0.80, 0.30)),
        detection("bag", 0.70, (0.11, 0.11, 0.51, 0.51)),
    ]

    filtered = suppress_duplicate_detections(detections, iou_threshold=0.8)

    assert len(filtered) == 3
    assert ("person", 0.91) in [(item.label, item.confidence) for item in filtered]
    assert ("person", 0.62) not in [(item.label, item.confidence) for item in filtered]
    assert ("person", 0.80) in [(item.label, item.confidence) for item in filtered]
    assert ("bag", 0.70) in [(item.label, item.confidence) for item in filtered]


def test_suppress_duplicate_detections_keeps_oldest_track_id() -> None:
    def tracked(confidence: float, box: tuple[float, float, float, float], track_id: int) -> Detection:
        return Detection(
            frame_index=0,
            timestamp_sec=0.0,
            label="person",
            confidence=confidence,
            bbox_xyxy=(0.0, 0.0, 10.0, 10.0),
            bbox_xyxy_norm=box,
            track_id=track_id,
        )

    # Higher-confidence box is the newer track (id 7); the suppressed
    # lower-confidence box is the established track (id 5).
    detections = [
        tracked(0.91, (0.11, 0.11, 0.51, 0.51), 7),
        tracked(0.62, (0.10, 0.10, 0.50, 0.50), 5),
    ]

    filtered = suppress_duplicate_detections(detections, iou_threshold=0.8)

    assert len(filtered) == 1
    assert filtered[0].confidence == 0.91  # highest-confidence geometry kept
    assert filtered[0].track_id == 5  # identity follows the oldest track, not confidence


def test_greedy_tracker_matches_overlap_and_keeps_unmatched() -> None:
    tracker = GreedyIoUTracker(iou_threshold=0.3)

    def det(label: str, box: tuple[float, float, float, float]) -> Detection:
        return Detection(
            frame_index=0,
            timestamp_sec=0.0,
            label=label,
            confidence=0.9,
            bbox_xyxy=(0.0, 0.0, 10.0, 10.0),
            bbox_xyxy_norm=box,
        )

    first = tracker.assign([det("person", (0.10, 0.10, 0.20, 0.40))])
    person_id = first[0].track_id
    assert person_id is not None

    # Next frame: the same person (overlapping box) keeps its id; a second person
    # that matches no existing track is still returned with a fresh id, never dropped.
    second = tracker.assign(
        [
            det("person", (0.11, 0.10, 0.21, 0.40)),
            det("person", (0.70, 0.10, 0.80, 0.40)),
        ]
    )
    ids = [item.track_id for item in second]
    assert len(second) == 2  # nothing dropped for failing to match
    assert person_id in ids
    assert len(set(ids)) == 2  # the unmatched one got its own id


def test_detections_from_ultralytics_result_reads_track_ids() -> None:
    detections = detections_from_ultralytics_result(
        ResultLike(ids=[7, 8]),
        frame_shape=(100, 100, 3),
        frame_index=3,
        timestamp_sec=1.5,
        fallback_names=["person", "car"],
    )

    assert [(item.label, item.track_id) for item in detections] == [("car", 8), ("person", 7)]


def test_detections_from_ultralytics_result_allows_missing_track_ids() -> None:
    detections = detections_from_ultralytics_result(
        ResultLike(),
        frame_shape=(100, 100, 3),
        frame_index=3,
        timestamp_sec=1.5,
        fallback_names=["person", "car"],
    )

    assert [item.track_id for item in detections] == [None, None]


def test_ultralytics_detector_uses_predict_by_default() -> None:
    calls = []

    class ModelLike:
        def predict(self, frame, **kwargs):
            calls.append(("predict", frame, kwargs))
            return [ResultLike()]

    detector = UltralyticsYOLOEDetector.__new__(UltralyticsYOLOEDetector)
    detector.class_names = ["person", "car"]
    detector.model_name = "fake.pt"
    detector.device = None
    detector.tracking_enabled = False
    detector._tracker = None
    detector.model = ModelLike()
    frame = FrameLike()

    detections = detector.detect(frame, frame_index=1, timestamp_sec=0.5, confidence=0.25)

    assert [item.track_id for item in detections] == [None, None]
    assert calls == [("predict", frame, {"conf": 0.25, "verbose": False})]


def test_ultralytics_detector_predicts_and_assigns_track_ids_when_tracking_enabled() -> None:
    calls = []

    class ModelLike:
        def predict(self, frame, **kwargs):
            calls.append("predict")
            return [ResultLike()]

    detector = UltralyticsYOLOEDetector.__new__(UltralyticsYOLOEDetector)
    detector.class_names = ["person", "car"]
    detector.model_name = "fake.pt"
    detector.device = None
    detector.tracking_enabled = True
    detector._tracker = GreedyIoUTracker()
    detector.model = ModelLike()
    frame = FrameLike()

    first = detector.detect(frame, frame_index=1, timestamp_sec=0.5, confidence=0.25)
    second = detector.detect(frame, frame_index=2, timestamp_sec=1.0, confidence=0.25)

    # Detection always runs through predict (never track); IDs come from our tracker.
    assert calls == ["predict", "predict"]
    first_ids = {item.label: item.track_id for item in first}
    assert all(track_id is not None for track_id in first_ids.values())
    # Same boxes the next frame -> same IDs (stable identity, nothing dropped).
    second_ids = {item.label: item.track_id for item in second}
    assert second_ids == first_ids