1use std::fmt;
35
36use crate::cluster::ClusterError;
37use crate::representation::{EventFrame, EventFrameData};
38
39#[derive(Clone, Copy, Debug, PartialEq)]
41pub struct Blob {
42 pub label: u64,
44 pub centroid: (f64, f64),
46 pub bounds: (usize, usize, usize, usize),
48 pub area: usize,
50}
51
52impl Blob {
53 pub fn extent(&self) -> usize {
55 let (x0, y0, x1, y1) = self.bounds;
56 (x1 - x0).max(y1 - y0) + 1
57 }
58}
59
60#[derive(Clone, Copy, Debug, PartialEq)]
62pub struct Track {
63 pub id: u64,
65 pub centroid: (f64, f64),
66 pub velocity: (f64, f64),
68 pub area: usize,
69 pub age: usize,
71 pub missed: usize,
73}
74
75impl Track {
76 fn predicted(&self) -> (f64, f64) {
82 (
83 self.centroid.0 + self.velocity.0,
84 self.centroid.1 + self.velocity.1,
85 )
86 }
87}
88
89#[derive(Clone, Copy, Debug, PartialEq)]
91pub struct TrackerConfig {
92 pub connectivity: u8,
94 pub min_area: usize,
97 pub max_distance: f64,
100 pub max_missed: usize,
103}
104
105impl Default for TrackerConfig {
106 fn default() -> Self {
107 Self {
108 connectivity: 8,
109 min_area: 4,
110 max_distance: 20.0,
111 max_missed: 3,
112 }
113 }
114}
115
116pub struct Tracker {
118 config: TrackerConfig,
119 tracks: Vec<Track>,
120 next_id: u64,
121}
122
123impl Tracker {
124 pub fn new(config: TrackerConfig) -> Self {
125 Self {
126 config,
127 tracks: Vec::new(),
128 next_id: 1,
129 }
130 }
131
132 pub fn tracks(&self) -> &[Track] {
134 &self.tracks
135 }
136
137 pub fn reset(&mut self) {
140 self.tracks.clear();
141 }
142
143 pub fn update(&mut self, frame: &EventFrame) -> Result<&[Track], ClusterError> {
145 let blobs = blobs_of(frame, self.config.connectivity, self.config.min_area)?;
146 self.associate(&blobs);
147 Ok(&self.tracks)
148 }
149
150 fn associate(&mut self, blobs: &[Blob]) {
152 let mut pairs: Vec<(f64, usize, usize)> = Vec::new();
155 for (track_index, track) in self.tracks.iter().enumerate() {
156 let (px, py) = track.predicted();
157 for (blob_index, blob) in blobs.iter().enumerate() {
158 let distance =
159 ((blob.centroid.0 - px).powi(2) + (blob.centroid.1 - py).powi(2)).sqrt();
160 if distance <= self.config.max_distance {
161 pairs.push((distance, track_index, blob_index));
162 }
163 }
164 }
165 pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
166
167 let mut track_taken = vec![false; self.tracks.len()];
168 let mut blob_taken = vec![false; blobs.len()];
169 for (_, track_index, blob_index) in pairs {
170 if track_taken[track_index] || blob_taken[blob_index] {
171 continue;
172 }
173 track_taken[track_index] = true;
174 blob_taken[blob_index] = true;
175
176 let blob = blobs[blob_index];
177 let track = &mut self.tracks[track_index];
178 track.velocity = (
179 blob.centroid.0 - track.centroid.0,
180 blob.centroid.1 - track.centroid.1,
181 );
182 track.centroid = blob.centroid;
183 track.area = blob.area;
184 track.missed = 0;
185 }
186
187 for (index, track) in self.tracks.iter_mut().enumerate() {
191 track.age += 1;
192 if !track_taken[index] {
193 track.missed += 1;
194 track.centroid = track.predicted();
195 }
196 }
197 self.tracks
198 .retain(|track| track.missed <= self.config.max_missed);
199
200 for (index, blob) in blobs.iter().enumerate() {
201 if !blob_taken[index] {
202 self.tracks.push(Track {
203 id: self.next_id,
204 centroid: blob.centroid,
205 velocity: (0.0, 0.0),
206 area: blob.area,
207 age: 0,
208 missed: 0,
209 });
210 self.next_id += 1;
211 }
212 }
213 }
214}
215
216pub fn blobs_of(
218 frame: &EventFrame,
219 connectivity: u8,
220 min_area: usize,
221) -> Result<Vec<Blob>, ClusterError> {
222 let labels = frame.connected_components(connectivity)?;
223 let (_, height, width) = labels.shape();
224 let EventFrameData::U64(values) = labels.data() else {
225 return Ok(Vec::new());
228 };
229
230 let count = values.iter().copied().max().unwrap_or(0) as usize;
232 let mut sums = vec![
233 (
234 0.0_f64,
235 0.0_f64,
236 0_usize,
237 usize::MAX,
238 usize::MAX,
239 0_usize,
240 0_usize
241 );
242 count + 1
243 ];
244 for y in 0..height {
245 for x in 0..width {
246 let label = values[y * width + x] as usize;
247 if label == 0 {
248 continue;
249 }
250 let entry = &mut sums[label];
251 entry.0 += x as f64;
252 entry.1 += y as f64;
253 entry.2 += 1;
254 entry.3 = entry.3.min(x);
255 entry.4 = entry.4.min(y);
256 entry.5 = entry.5.max(x);
257 entry.6 = entry.6.max(y);
258 }
259 }
260
261 Ok((1..=count)
262 .filter_map(|label| {
263 let (sx, sy, area, x0, y0, x1, y1) = sums[label];
264 if area < min_area.max(1) {
265 return None;
266 }
267 Some(Blob {
268 label: label as u64,
269 centroid: (sx / area as f64, sy / area as f64),
270 bounds: (x0, y0, x1, y1),
271 area,
272 })
273 })
274 .collect())
275}
276
277#[derive(Debug, PartialEq, Eq)]
279pub enum TrackError {
280 Cluster(ClusterError),
281}
282
283impl fmt::Display for TrackError {
284 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
285 match self {
286 Self::Cluster(error) => write!(formatter, "{error}"),
287 }
288 }
289}
290
291impl std::error::Error for TrackError {}
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296 use crate::representation::EventFrame;
297
298 fn frame_with(
300 width: usize,
301 height: usize,
302 centres: &[(usize, usize)],
303 size: usize,
304 ) -> EventFrame {
305 let mut data = vec![0_u8; width * height];
306 for &(cx, cy) in centres {
307 for dy in 0..size {
308 for dx in 0..size {
309 let (x, y) = (cx + dx, cy + dy);
310 if x < width && y < height {
311 data[y * width + x] = 1;
312 }
313 }
314 }
315 }
316 EventFrame::intensity(EventFrameData::U8(data), width, height).unwrap()
317 }
318
319 #[test]
320 fn blobs_measure_centroid_area_and_bounds() {
321 let frame = frame_with(32, 32, &[(10, 6)], 4);
323 let blobs = blobs_of(&frame, 8, 1).unwrap();
324 assert_eq!(blobs.len(), 1);
325 assert_eq!(blobs[0].area, 16);
326 assert!((blobs[0].centroid.0 - 11.5).abs() < 1e-9);
327 assert!((blobs[0].centroid.1 - 7.5).abs() < 1e-9);
328 assert_eq!(blobs[0].bounds, (10, 6, 13, 9));
329 assert_eq!(blobs[0].extent(), 4);
330 }
331
332 #[test]
333 fn separate_objects_are_separate_blobs() {
334 let frame = frame_with(64, 64, &[(5, 5), (40, 40)], 4);
335 assert_eq!(blobs_of(&frame, 8, 1).unwrap().len(), 2);
336 }
337
338 #[test]
339 fn min_area_filters_noise() {
340 let mut data = vec![0_u8; 32 * 32];
342 for dy in 0..5 {
343 for dx in 0..5 {
344 data[(10 + dy) * 32 + 10 + dx] = 1;
345 }
346 }
347 data[2 * 32 + 30] = 1;
348 let frame = EventFrame::intensity(EventFrameData::U8(data), 32, 32).unwrap();
349 assert_eq!(blobs_of(&frame, 8, 1).unwrap().len(), 2);
350 assert_eq!(blobs_of(&frame, 8, 4).unwrap().len(), 1);
351 }
352
353 #[test]
354 fn an_empty_frame_has_no_blobs() {
355 let frame = frame_with(16, 16, &[], 0);
356 assert!(blobs_of(&frame, 8, 1).unwrap().is_empty());
357 }
358
359 #[test]
360 fn a_track_keeps_its_id_while_the_object_moves() {
361 let mut tracker = Tracker::new(TrackerConfig::default());
364 let mut ids = Vec::new();
365 for step in 0..10 {
366 let frame = frame_with(96, 64, &[(5 + step * 6, 20)], 5);
367 let tracks = tracker.update(&frame).unwrap();
368 assert_eq!(tracks.len(), 1, "step {step}");
369 ids.push(tracks[0].id);
370 }
371 assert!(ids.windows(2).all(|w| w[0] == w[1]), "id changed: {ids:?}");
372 }
373
374 #[test]
375 fn velocity_matches_the_motion() {
376 let mut tracker = Tracker::new(TrackerConfig::default());
377 for step in 0..5 {
378 let frame = frame_with(96, 64, &[(5 + step * 6, 20)], 5);
379 tracker.update(&frame).unwrap();
380 }
381 let track = tracker.tracks()[0];
382 assert!(
383 (track.velocity.0 - 6.0).abs() < 1e-6,
384 "vx {}",
385 track.velocity.0
386 );
387 assert!(track.velocity.1.abs() < 1e-6, "vy {}", track.velocity.1);
388 }
389
390 #[test]
391 fn two_objects_get_two_ids() {
392 let mut tracker = Tracker::new(TrackerConfig::default());
393 for step in 0..6 {
394 let frame = frame_with(96, 64, &[(5 + step * 4, 10), (5 + step * 4, 45)], 5);
395 let tracks = tracker.update(&frame).unwrap();
396 assert_eq!(tracks.len(), 2, "step {step}");
397 }
398 let ids: Vec<u64> = tracker.tracks().iter().map(|t| t.id).collect();
399 assert_ne!(ids[0], ids[1]);
400 }
401
402 #[test]
403 fn a_track_survives_a_brief_disappearance() {
404 let config = TrackerConfig {
405 max_missed: 3,
406 ..TrackerConfig::default()
407 };
408 let mut tracker = Tracker::new(config);
409 for step in 0..4 {
410 tracker
411 .update(&frame_with(96, 64, &[(5 + step * 5, 20)], 5))
412 .unwrap();
413 }
414 let id = tracker.tracks()[0].id;
415
416 for _ in 0..2 {
418 tracker.update(&frame_with(96, 64, &[], 0)).unwrap();
419 }
420 assert_eq!(tracker.tracks().len(), 1);
421 assert_eq!(tracker.tracks()[0].id, id, "the id must survive the gap");
422
423 let tracks = tracker.update(&frame_with(96, 64, &[(35, 20)], 5)).unwrap();
425 assert_eq!(tracks.len(), 1);
426 assert_eq!(tracks[0].id, id, "and be recognised as the same object");
427 }
428
429 #[test]
430 fn a_track_dies_after_max_missed() {
431 let config = TrackerConfig {
432 max_missed: 2,
433 ..TrackerConfig::default()
434 };
435 let mut tracker = Tracker::new(config);
436 tracker.update(&frame_with(64, 64, &[(20, 20)], 5)).unwrap();
437 assert_eq!(tracker.tracks().len(), 1);
438 for _ in 0..3 {
439 tracker.update(&frame_with(64, 64, &[], 0)).unwrap();
440 }
441 assert!(tracker.tracks().is_empty(), "the track should have expired");
442 }
443
444 #[test]
445 fn a_distant_jump_starts_a_new_track_rather_than_teleporting() {
446 let config = TrackerConfig {
448 max_distance: 10.0,
449 max_missed: 0,
450 ..TrackerConfig::default()
451 };
452 let mut tracker = Tracker::new(config);
453 let first = tracker.update(&frame_with(96, 96, &[(5, 5)], 5)).unwrap()[0].id;
454 let tracks = tracker.update(&frame_with(96, 96, &[(80, 80)], 5)).unwrap();
455 assert_eq!(tracks.len(), 1);
456 assert_ne!(
457 tracks[0].id, first,
458 "a jump beyond the gate is a new object"
459 );
460 }
461
462 #[test]
463 fn ids_are_never_reused() {
464 let mut tracker = Tracker::new(TrackerConfig {
465 max_missed: 0,
466 ..TrackerConfig::default()
467 });
468 let first = tracker.update(&frame_with(64, 64, &[(10, 10)], 5)).unwrap()[0].id;
469 tracker.update(&frame_with(64, 64, &[], 0)).unwrap();
470 let second = tracker.update(&frame_with(64, 64, &[(10, 10)], 5)).unwrap()[0].id;
471 assert_ne!(first, second, "a new object must not inherit a retired id");
472
473 tracker.reset();
474 let third = tracker.update(&frame_with(64, 64, &[(10, 10)], 5)).unwrap()[0].id;
475 assert!(third > second, "ids continue past a reset");
476 }
477
478 #[test]
479 fn fast_motion_is_followed_by_predicting_ahead() {
480 let mut tracker = Tracker::new(TrackerConfig::default());
484 let mut ids = Vec::new();
485 for step in 0..6 {
486 let frame = frame_with(160, 64, &[(5 + step * 15, 20)], 6);
487 let tracks = tracker.update(&frame).unwrap();
488 assert_eq!(tracks.len(), 1, "lost the object at step {step}");
489 ids.push(tracks[0].id);
490 }
491 assert!(ids.windows(2).all(|w| w[0] == w[1]));
492 }
493}