File size: 6,968 Bytes
d725335 114ea19 8536e24 d725335 114ea19 d725335 8536e24 114ea19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | 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
|