103 lines
3.4 KiB
Python
103 lines
3.4 KiB
Python
|
"""Lightweight deterministic person tracking for the single-camera V1 flow."""
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from math import hypot
|
||
|
|
from typing import Dict, Sequence, Tuple
|
||
|
|
|
||
|
|
from v1.pose import PersonPose
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class TrackedPersonPose:
|
||
|
|
track_id: str
|
||
|
|
detected_at_monotonic: float
|
||
|
|
pose: PersonPose
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class _Track:
|
||
|
|
center: Tuple[float, float]
|
||
|
|
last_seen_at: float
|
||
|
|
|
||
|
|
|
||
|
|
class PersonTracker:
|
||
|
|
"""Assign stable IDs by nearest normalized box center across adjacent frames."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self, max_match_distance_ratio: float = 0.2, max_age_seconds: float = 2.0
|
||
|
|
) -> None:
|
||
|
|
if not 0.0 < max_match_distance_ratio <= 1.0:
|
||
|
|
raise ValueError("max_match_distance_ratio must be in (0, 1]")
|
||
|
|
if max_age_seconds <= 0:
|
||
|
|
raise ValueError("max_age_seconds must be positive")
|
||
|
|
self._max_match_distance_ratio = max_match_distance_ratio
|
||
|
|
self._max_age_seconds = max_age_seconds
|
||
|
|
self._tracks: Dict[str, _Track] = {}
|
||
|
|
self._next_track_number = 1
|
||
|
|
|
||
|
|
def update(
|
||
|
|
self,
|
||
|
|
poses: Sequence[PersonPose],
|
||
|
|
detected_at_monotonic: float,
|
||
|
|
frame_size: Tuple[int, int],
|
||
|
|
) -> Sequence[TrackedPersonPose]:
|
||
|
|
width, height = frame_size
|
||
|
|
if width <= 0 or height <= 0:
|
||
|
|
raise ValueError("frame_size must contain positive width and height")
|
||
|
|
self._expire_tracks(detected_at_monotonic)
|
||
|
|
available_ids = set(self._tracks)
|
||
|
|
tracked = []
|
||
|
|
for pose in poses:
|
||
|
|
center = self._box_center(pose)
|
||
|
|
track_id = self._nearest_available_track(center, available_ids, width, height)
|
||
|
|
if track_id is None:
|
||
|
|
track_id = "P-{0:04d}".format(self._next_track_number)
|
||
|
|
self._next_track_number += 1
|
||
|
|
else:
|
||
|
|
available_ids.remove(track_id)
|
||
|
|
self._tracks[track_id] = _Track(center=center, last_seen_at=detected_at_monotonic)
|
||
|
|
tracked.append(
|
||
|
|
TrackedPersonPose(
|
||
|
|
track_id=track_id,
|
||
|
|
detected_at_monotonic=float(detected_at_monotonic),
|
||
|
|
pose=pose,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return tuple(tracked)
|
||
|
|
|
||
|
|
def _nearest_available_track(
|
||
|
|
self,
|
||
|
|
center: Tuple[float, float],
|
||
|
|
available_ids: set,
|
||
|
|
width: int,
|
||
|
|
height: int,
|
||
|
|
):
|
||
|
|
closest_id = None
|
||
|
|
closest_distance = None
|
||
|
|
for track_id in available_ids:
|
||
|
|
previous = self._tracks[track_id].center
|
||
|
|
distance = hypot(
|
||
|
|
(center[0] - previous[0]) / float(width),
|
||
|
|
(center[1] - previous[1]) / float(height),
|
||
|
|
)
|
||
|
|
if distance <= self._max_match_distance_ratio and (
|
||
|
|
closest_distance is None or distance < closest_distance
|
||
|
|
):
|
||
|
|
closest_id = track_id
|
||
|
|
closest_distance = distance
|
||
|
|
return closest_id
|
||
|
|
|
||
|
|
def _expire_tracks(self, now: float) -> None:
|
||
|
|
expired_ids = [
|
||
|
|
track_id
|
||
|
|
for track_id, track in self._tracks.items()
|
||
|
|
if now - track.last_seen_at > self._max_age_seconds
|
||
|
|
]
|
||
|
|
for track_id in expired_ids:
|
||
|
|
del self._tracks[track_id]
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _box_center(pose: PersonPose) -> Tuple[float, float]:
|
||
|
|
left, top, right, bottom = pose.box_xyxy
|
||
|
|
return ((left + right) / 2.0, (top + bottom) / 2.0)
|