84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""Session-local single-stream IoU tracker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from .models import BoundingBox, Detection, TrackedObject
|
|
|
|
|
|
def _iou(first: BoundingBox, second: BoundingBox) -> float:
|
|
intersection = BoundingBox(
|
|
max(first.left, second.left),
|
|
max(first.top, second.top),
|
|
min(first.right, second.right),
|
|
min(first.bottom, second.bottom),
|
|
).area
|
|
union = first.area + second.area - intersection
|
|
return intersection / union if union else 0.0
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _Track:
|
|
track_id: str
|
|
detection: Detection
|
|
missed: int = 0
|
|
|
|
|
|
class SingleStreamTracker:
|
|
"""Tracks anonymous boxes only within one process session and one stream."""
|
|
|
|
def __init__(self, *, iou_threshold: float = 0.2, max_missed: int = 2) -> None:
|
|
if not 0.0 <= iou_threshold <= 1.0 or max_missed < 0:
|
|
raise ValueError("invalid tracker threshold or missed-frame limit")
|
|
self._iou_threshold = iou_threshold
|
|
self._max_missed = max_missed
|
|
self._tracks: dict[str, _Track] = {}
|
|
self._next_id = 1
|
|
|
|
def update(
|
|
self,
|
|
detections: tuple[Detection, ...],
|
|
*,
|
|
frame_sequence: int,
|
|
timestamp_ns: int,
|
|
) -> tuple[TrackedObject, ...]:
|
|
unmatched_tracks = set(self._tracks)
|
|
results: list[TrackedObject] = []
|
|
for detection in detections:
|
|
candidates = [
|
|
(track_id, _iou(self._tracks[track_id].detection.box, detection.box))
|
|
for track_id in unmatched_tracks
|
|
if self._tracks[track_id].detection.category == detection.category
|
|
]
|
|
track_id, score = max(candidates, key=lambda item: item[1], default=("", -1.0))
|
|
if score < self._iou_threshold:
|
|
track_id = f"track-{self._next_id:06d}"
|
|
self._next_id += 1
|
|
self._tracks[track_id] = _Track(track_id, detection)
|
|
else:
|
|
unmatched_tracks.remove(track_id)
|
|
self._tracks[track_id].detection = detection
|
|
self._tracks[track_id].missed = 0
|
|
results.append(
|
|
TrackedObject(
|
|
track_id=track_id,
|
|
box=detection.box,
|
|
category=detection.category,
|
|
confidence=detection.confidence,
|
|
frame_sequence=frame_sequence,
|
|
timestamp_ns=timestamp_ns,
|
|
)
|
|
)
|
|
for track_id in unmatched_tracks:
|
|
track = self._tracks[track_id]
|
|
track.missed += 1
|
|
if track.missed > self._max_missed:
|
|
del self._tracks[track_id]
|
|
return tuple(results)
|
|
|
|
def finish(self) -> tuple[str, ...]:
|
|
ended = tuple(sorted(self._tracks))
|
|
self._tracks.clear()
|
|
return ended
|