cecile_supercool_tracker/examples/
iou.rs

1use crate::track::{
2    MetricOutput, MetricQuery, NoopLookup, Observation, ObservationAttributes, ObservationMetric,
3    ObservationsDb, TrackAttributes, TrackAttributesUpdate, TrackStatus,
4};
5use crate::utils::bbox::BoundingBox;
6use anyhow::Result;
7
8#[derive(Debug, Clone, Default)]
9pub struct BBoxAttributes {
10    pub bboxes: Vec<BoundingBox>,
11}
12
13#[derive(Clone, Debug)]
14pub struct BBoxAttributesUpdate;
15
16impl TrackAttributesUpdate<BBoxAttributes> for BBoxAttributesUpdate {
17    fn apply(&self, _attrs: &mut BBoxAttributes) -> Result<()> {
18        Ok(())
19    }
20}
21
22impl TrackAttributes<BBoxAttributes, BoundingBox> for BBoxAttributes {
23    type Update = BBoxAttributesUpdate;
24    type Lookup = NoopLookup<BBoxAttributes, BoundingBox>;
25
26    fn compatible(&self, _other: &BBoxAttributes) -> bool {
27        true
28    }
29
30    fn merge(&mut self, other: &BBoxAttributes) -> Result<()> {
31        self.bboxes.extend_from_slice(&other.bboxes);
32        Ok(())
33    }
34
35    fn baked(&self, _observations: &ObservationsDb<BoundingBox>) -> Result<TrackStatus> {
36        Ok(TrackStatus::Ready)
37    }
38}
39
40#[derive(Clone)]
41pub struct IOUMetric {
42    history: usize,
43}
44
45impl Default for IOUMetric {
46    fn default() -> Self {
47        Self { history: 3 }
48    }
49}
50
51impl ObservationMetric<BBoxAttributes, BoundingBox> for IOUMetric {
52    fn metric(&self, mq: &MetricQuery<'_, BBoxAttributes, BoundingBox>) -> MetricOutput<f32> {
53        let (e1, e2) = (mq.candidate_observation, mq.track_observation);
54        let box_m_opt =
55            BoundingBox::calculate_metric_object(&e1.attr().as_ref(), &e2.attr().as_ref());
56        if let Some(box_m) = &box_m_opt {
57            if *box_m < 0.01 {
58                None
59            } else {
60                Some((box_m_opt, None))
61            }
62        } else {
63            None
64        }
65    }
66
67    fn optimize(
68        &mut self,
69        _feature_class: u64,
70        _merge_history: &[u64],
71        attrs: &mut BBoxAttributes,
72        features: &mut Vec<Observation<BoundingBox>>,
73        prev_length: usize,
74        is_merge: bool,
75    ) -> Result<()> {
76        if !is_merge {
77            if let Some(bb) = &features[prev_length].attr() {
78                attrs.bboxes.push(*bb);
79            }
80        }
81        // Kalman filter should be used here to generate better prediction for next
82        // comparison
83        features.reverse();
84        features.truncate(self.history);
85        features.reverse();
86        Ok(())
87    }
88}