Coverage for orchestrant/pipeline/tracking/centroid.py: 98%

1"""Simple centroid-based tracking utilities.""" 

2 

3from __future__ import annotations 

4 

5from collections import deque 

6 

7import numpy as np 

8 

9from orchestrant.pipeline.constants import ( 

10 TRACKER_DEFAULT_MAX_AGE_SECONDS, 

11 TRACKER_DEFAULT_MAX_MATCH_DISTANCE, 

12 TRACKER_DEFAULT_MAX_TRAIL_POINTS, 

13) 

14from orchestrant.pipeline.types import Track 

15 

16 

17class SimpleCentroidTracker: 

18 """Very lightweight centroid tracker for a single class (e.g. persons).""" 

19 

20 def __init__( 

21 self, 

22 *, 

23 max_age_s: float = TRACKER_DEFAULT_MAX_AGE_SECONDS, 

24 max_match_dist_norm: float = TRACKER_DEFAULT_MAX_MATCH_DISTANCE, 

25 max_trail_points: int = TRACKER_DEFAULT_MAX_TRAIL_POINTS, 

26 ) -> None: 

27 """Initialize tracker configuration and internal state.""" 

28 self._max_age_s = float(max_age_s) 

29 self._max_match_dist_norm = float(max_match_dist_norm) 

30 self._max_match_dist_norm_sq = self._max_match_dist_norm**2 

31 self._max_trail_points = int(max_trail_points) 

32 

33 self._next_id = 1 

34 self._tracks: dict[int, Track] = {} 

35 

36 def update( 

37 self, centroids_norm: list[tuple[float, float]], now_ts: float 

38 ) -> dict[int, Track]: 

39 """Update tracks with the latest normalized centroids.""" 

40 self._expire_tracks(now_ts) 

41 

42 if not centroids_norm: 

43 return self._tracks 

44 

45 if not self._tracks: 

46 self._initialize_tracks(centroids_norm, now_ts) 

47 return self._tracks 

48 

49 used_dets = self._associate_tracks(centroids_norm, now_ts) 

50 self._add_unmatched(centroids_norm, used_dets, now_ts) 

51 

52 return self._tracks 

53 

54 def _expire_tracks(self, now_ts: float) -> None: 

55 """Remove tracks that haven't been seen recently.""" 

56 expired_ids = [ 

57 tid 

58 for tid, tr in self._tracks.items() 

59 if (now_ts - tr.last_seen_ts) > self._max_age_s 

60 ] 

61 for tid in expired_ids: 

62 self._tracks.pop(tid, None) 

63 

64 def _initialize_tracks( 

65 self, centroids_norm: list[tuple[float, float]], now_ts: float 

66 ) -> None: 

67 """Create initial tracks from the first set of detections.""" 

68 for centroid in centroids_norm: 

69 self._tracks[self._next_id] = Track( 

70 track_id=self._next_id, 

71 points_norm=deque([centroid], maxlen=self._max_trail_points), 

72 last_seen_ts=now_ts, 

73 ) 

74 self._next_id += 1 

75 

76 def _associate_tracks( 

77 self, centroids_norm: list[tuple[float, float]], now_ts: float 

78 ) -> set[int]: 

79 """Associate detections with existing tracks using greedy matching. 

80 

81 Uses vectorized distance computation for O(n*m) complexity where 

82 n = number of tracks and m = number of detections. 

83 """ 

84 track_ids = list(self._tracks.keys()) 

85 prev_centroids = np.array( 

86 [self._tracks[tid].points_norm[-1] for tid in track_ids] 

87 ) 

88 curr_centroids = np.array(centroids_norm) 

89 

90 if len(track_ids) == 0 or len(curr_centroids) == 0: 

91 return set() 

92 

93 diff = prev_centroids[:, np.newaxis, :] - curr_centroids[np.newaxis, :, :] 

94 dist_sq_matrix = np.sum(diff**2, axis=2) 

95 

96 used_tracks: set[int] = set() 

97 used_dets: set[int] = set() 

98 

99 valid_mask = dist_sq_matrix <= self._max_match_dist_norm_sq 

100 candidates = sorted( 

101 (dist_sq_matrix[ti, di], ti, di) 

102 for ti in range(len(track_ids)) 

103 for di in range(len(centroids_norm)) 

104 if valid_mask[ti, di] 

105 ) 

106 

107 for _dist_sq, ti, di in candidates: 

108 if ti in used_tracks or di in used_dets: 

109 continue 

110 tid = track_ids[ti] 

111 self._tracks[tid].points_norm.append(centroids_norm[di]) 

112 self._tracks[tid].last_seen_ts = now_ts 

113 used_tracks.add(ti) 

114 used_dets.add(di) 

115 

116 return used_dets 

117 

118 def _add_unmatched( 

119 self, 

120 centroids_norm: list[tuple[float, float]], 

121 used_dets: set[int], 

122 now_ts: float, 

123 ) -> None: 

124 """Create new tracks for unmatched detections.""" 

125 for di, centroid in enumerate(centroids_norm): 

126 if di in used_dets: 

127 continue 

128 self._tracks[self._next_id] = Track( 

129 track_id=self._next_id, 

130 points_norm=deque([centroid], maxlen=self._max_trail_points), 

131 last_seen_ts=now_ts, 

132 ) 

133 self._next_id += 1