1use uuid::Uuid;
2
3use crate::projection::project_to_ground;
4use crate::proto::{Detection, Frame, MapObject, SceneMap};
5
6#[derive(Debug, Clone, Copy)]
8pub struct EngineConfig {
9 pub association_radius_meters: f32,
27 pub stale_timeout_seconds: f64,
41}
42
43impl Default for EngineConfig {
44 fn default() -> Self {
45 Self { association_radius_meters: 0.3, stale_timeout_seconds: 2.0 }
46 }
47}
48
49pub struct Engine {
50 scene_map: SceneMap,
51 config: EngineConfig,
52}
53
54impl Engine {
55 pub fn new() -> Self {
56 Self::with_config(EngineConfig::default())
57 }
58
59 pub fn with_config(config: EngineConfig) -> Self {
60 Self { scene_map: SceneMap::default(), config }
61 }
62
63 pub fn ingest_frame(&mut self, frame: Frame) -> &SceneMap {
64 let (Some(pose), Some(intrinsics)) = (frame.pose.as_ref(), frame.intrinsics.as_ref())
65 else {
66 return &self.scene_map;
67 };
68
69 for detection in &frame.detections {
70 if let Some((x, y)) = project_to_ground(pose, intrinsics, detection) {
71 self.associate_and_fuse(detection, x, y, pose.timestamp);
72 }
73 }
74
75 self.prune_stale(pose.timestamp);
76 self.scene_map.updated_at = pose.timestamp;
77 &self.scene_map
78 }
79
80 fn associate_and_fuse(&mut self, detection: &Detection, x: f32, y: f32, timestamp: f64) {
81 let radius = self.config.association_radius_meters;
82 let nearest = self
83 .scene_map
84 .objects
85 .iter_mut()
86 .filter(|obj| obj.label == detection.label)
87 .filter_map(|obj| {
88 let dist = ((obj.x - x).powi(2) + (obj.y - y).powi(2)).sqrt();
89 (dist <= radius).then_some((dist, obj))
90 })
91 .min_by(|a, b| a.0.total_cmp(&b.0))
92 .map(|(_, obj)| obj);
93
94 match nearest {
95 Some(obj) => fuse(obj, detection, x, y, timestamp),
96 None => self.scene_map.objects.push(new_object(detection, x, y, timestamp)),
97 }
98 }
99
100 fn prune_stale(&mut self, now: f64) {
101 let stale_timeout = self.config.stale_timeout_seconds;
102 self.scene_map
103 .objects
104 .retain(|obj| now - obj.last_seen <= stale_timeout);
105 }
106}
107
108impl Default for Engine {
109 fn default() -> Self {
110 Self::new()
111 }
112}
113
114fn new_object(detection: &Detection, x: f32, y: f32, timestamp: f64) -> MapObject {
115 MapObject {
116 id: Uuid::new_v4().to_string(),
117 label: detection.label.clone(),
118 x,
119 y,
120 confidence: detection.confidence,
121 observation_count: 1,
122 first_seen: timestamp,
123 last_seen: timestamp,
124 }
125}
126
127fn fuse(obj: &mut MapObject, detection: &Detection, x: f32, y: f32, timestamp: f64) {
128 let n = obj.observation_count as f32;
129 obj.x = (obj.x * n + x) / (n + 1.0);
130 obj.y = (obj.y * n + y) / (n + 1.0);
131 obj.confidence = 1.0 - (1.0 - obj.confidence) * (1.0 - detection.confidence);
134 obj.observation_count += 1;
135 obj.last_seen = timestamp;
136}
137
138#[cfg(test)]
139mod tests {
140 use super::*;
141 use crate::proto::{CameraIntrinsics, CameraPose};
142
143 fn frame_at(timestamp: f64, x: f32, y: f32, z: f32, detections: Vec<Detection>) -> Frame {
144 Frame {
145 pose: Some(CameraPose {
146 timestamp,
147 x,
148 y,
149 z,
150 qx: 1.0,
151 qy: 0.0,
152 qz: 0.0,
153 qw: 0.0,
154 }),
155 intrinsics: Some(CameraIntrinsics { fx: 500.0, fy: 500.0, cx: 320.0, cy: 240.0 }),
156 detections,
157 }
158 }
159
160 fn detection(label: &str, confidence: f32) -> Detection {
161 Detection {
162 label: label.to_string(),
163 confidence,
164 bbox_x: 0.5,
165 bbox_y: 0.5,
166 bbox_w: 0.0,
167 bbox_h: 0.0,
168 }
169 }
170
171 #[test]
172 fn repeated_observations_fuse_into_one_object_with_growing_confidence() {
173 let mut engine = Engine::new();
174 engine.ingest_frame(frame_at(0.0, 0.0, 0.0, 2.0, vec![detection("chair", 0.5)]));
175 let map = engine.ingest_frame(frame_at(0.5, 0.0, 0.0, 2.0, vec![detection("chair", 0.5)]));
176
177 assert_eq!(map.objects.len(), 1);
178 let obj = &map.objects[0];
179 assert_eq!(obj.observation_count, 2);
180 assert!(obj.confidence > 0.5, "confidence should grow past a single observation");
181 }
182
183 #[test]
184 fn different_labels_at_same_spot_stay_separate() {
185 let mut engine = Engine::new();
186 let map = engine.ingest_frame(frame_at(
187 0.0,
188 0.0,
189 0.0,
190 2.0,
191 vec![detection("chair", 0.5), detection("door", 0.5)],
192 ));
193
194 assert_eq!(map.objects.len(), 2);
195 }
196
197 #[test]
198 fn objects_not_reinforced_go_stale_and_get_pruned() {
199 let mut engine = Engine::new();
200 engine.ingest_frame(frame_at(0.0, 0.0, 0.0, 2.0, vec![detection("chair", 0.5)]));
201 let map = engine.ingest_frame(frame_at(10.0, 0.0, 0.0, 2.0, vec![]));
202
203 assert!(map.objects.is_empty());
204 }
205}