1use std::collections::{HashMap, HashSet, VecDeque};
4
5use spatialrust_core::{
6 FieldSemantic, FrameId, HasNormals3, HasPositions3, PointBuffer, PointBufferSet, PointCloud,
7};
8use spatialrust_math::{Isometry3, TransformPoint, Vec3};
9use spatialrust_records::SpatialRecord;
10
11use crate::{SyncError, SyncResult};
12
13#[derive(Clone, Debug, PartialEq)]
15pub struct FrameEdge {
16 pub parent: FrameId,
18 pub child: FrameId,
20 pub child_t_parent: Isometry3<f32>,
22}
23
24#[derive(Clone, Debug, Default, PartialEq)]
26pub struct FrameGraph {
27 edges: HashMap<String, Vec<(String, Isometry3<f32>)>>,
29 reverse: HashMap<String, Vec<(String, Isometry3<f32>)>>,
31}
32
33impl FrameGraph {
34 #[must_use]
36 pub fn new() -> Self {
37 Self::default()
38 }
39
40 pub fn insert_edge(&mut self, edge: FrameEdge) -> SyncResult<()> {
42 if edge.parent.0 == edge.child.0 {
43 return Err(SyncError::InvalidConfiguration(
44 "frame edge parent and child must differ".into(),
45 ));
46 }
47 self.edges
48 .entry(edge.parent.0.clone())
49 .or_default()
50 .retain(|(child, _)| child != &edge.child.0);
51 self.edges
52 .entry(edge.parent.0.clone())
53 .or_default()
54 .push((edge.child.0.clone(), edge.child_t_parent));
55
56 let parent_t_child = edge.child_t_parent.inverse();
57 self.reverse
58 .entry(edge.child.0.clone())
59 .or_default()
60 .retain(|(parent, _)| parent != &edge.parent.0);
61 self.reverse.entry(edge.child.0.clone()).or_default().push((edge.parent.0, parent_t_child));
62 Ok(())
63 }
64
65 pub fn lookup(&self, from: &FrameId, to: &FrameId) -> SyncResult<Isometry3<f32>> {
67 if from == to {
68 return Ok(Isometry3::identity());
69 }
70 let mut queue = VecDeque::from([(from.0.clone(), Isometry3::identity())]);
71 let mut visited = HashSet::from([from.0.clone()]);
72 while let Some((node, acc)) = queue.pop_front() {
73 for (next, edge) in self.neighbors(&node) {
74 if !visited.insert(next.clone()) {
75 continue;
76 }
77 let composed = edge.compose(acc);
78 if next == to.0 {
79 return Ok(composed);
80 }
81 queue.push_back((next, composed));
82 }
83 }
84 Err(SyncError::NoTransformPath { from: from.0.clone(), to: to.0.clone() })
85 }
86
87 pub fn transform_record_to(
94 &self,
95 record: &SpatialRecord,
96 target: &FrameId,
97 ) -> SyncResult<SpatialRecord> {
98 let source = &record.metadata().frame_id;
99 let transform = self.lookup(source, target)?;
100 let cloud = record.cloud();
101 let (x, y, z) = cloud.positions3()?;
102 let mut positions = Vec::with_capacity(cloud.len());
103 for index in 0..cloud.len() {
104 positions.push(transform.transform_point(Vec3::new(x[index], y[index], z[index])));
105 }
106
107 let normal_names = normal_field_names(cloud)?;
108 let normals = if normal_names.is_some() {
109 let (x, y, z) = cloud.normals3()?;
110 Some(
111 (0..cloud.len())
112 .map(|index| {
113 transform
114 .transform_vector(Vec3::new(x[index], y[index], z[index]))
115 .normalize()
116 })
117 .collect::<Vec<_>>(),
118 )
119 } else {
120 None
121 };
122
123 let position_names = position_field_names(cloud);
124 let mut buffers = PointBufferSet::new();
125 for field in cloud.schema().fields() {
126 let buffer = if field.name == position_names.0 {
127 PointBuffer::from_f32(positions.iter().map(|point| point.x).collect())
128 } else if field.name == position_names.1 {
129 PointBuffer::from_f32(positions.iter().map(|point| point.y).collect())
130 } else if field.name == position_names.2 {
131 PointBuffer::from_f32(positions.iter().map(|point| point.z).collect())
132 } else if let Some((names, values)) = normal_names.as_ref().zip(normals.as_ref()) {
133 if field.name == names.0 {
134 PointBuffer::from_f32(values.iter().map(|normal| normal.x).collect())
135 } else if field.name == names.1 {
136 PointBuffer::from_f32(values.iter().map(|normal| normal.y).collect())
137 } else if field.name == names.2 {
138 PointBuffer::from_f32(values.iter().map(|normal| normal.z).collect())
139 } else {
140 clone_buffer(cloud.field(&field.name)?)
141 }
142 } else {
143 clone_buffer(cloud.field(&field.name)?)
144 };
145 buffers.insert(field.name.clone(), buffer);
146 }
147
148 let mut metadata = cloud.metadata().clone();
149 metadata.frame_id = target.clone();
150 metadata.sensor_origin =
151 metadata.sensor_origin.map(|origin| transform.transform_point(origin));
152 let transformed_cloud =
153 PointCloud::try_from_parts(cloud.schema().clone(), buffers, metadata)?;
154 Ok(SpatialRecord::try_new_with_provenance(
155 record.schema().clone(),
156 transformed_cloud,
157 record.provenance().clone(),
158 )?)
159 }
160
161 fn neighbors(&self, node: &str) -> Vec<(String, Isometry3<f32>)> {
162 let mut out = Vec::new();
163 if let Some(forward) = self.edges.get(node) {
164 out.extend(forward.iter().cloned());
165 }
166 if let Some(back) = self.reverse.get(node) {
167 out.extend(back.iter().cloned());
168 }
169 out
170 }
171}
172
173fn position_field_names(cloud: &PointCloud) -> (String, String, String) {
174 (
175 cloud
176 .schema()
177 .find_semantic(FieldSemantic::PositionX)
178 .expect("positions3 validates position fields")
179 .name
180 .clone(),
181 cloud
182 .schema()
183 .find_semantic(FieldSemantic::PositionY)
184 .expect("positions3 validates position fields")
185 .name
186 .clone(),
187 cloud
188 .schema()
189 .find_semantic(FieldSemantic::PositionZ)
190 .expect("positions3 validates position fields")
191 .name
192 .clone(),
193 )
194}
195
196fn normal_field_names(cloud: &PointCloud) -> SyncResult<Option<(String, String, String)>> {
197 let names = [
198 (FieldSemantic::NormalX, "normal x"),
199 (FieldSemantic::NormalY, "normal y"),
200 (FieldSemantic::NormalZ, "normal z"),
201 ]
202 .map(|(semantic, label)| {
203 let fields: Vec<&str> = cloud
204 .schema()
205 .fields()
206 .iter()
207 .filter(|field| field.semantic == semantic)
208 .map(|field| field.name.as_str())
209 .collect();
210 (label, fields)
211 });
212 let present = names.iter().filter(|(_, fields)| !fields.is_empty()).count();
213 if present == 0 {
214 return Ok(None);
215 }
216 if names.iter().any(|(_, fields)| fields.len() != 1) {
217 return Err(SyncError::InvalidConfiguration(
218 "normal fields must contain exactly one x/y/z semantic each".into(),
219 ));
220 }
221 Ok(Some((names[0].1[0].to_owned(), names[1].1[0].to_owned(), names[2].1[0].to_owned())))
222}
223
224fn clone_buffer(buffer: &PointBuffer) -> PointBuffer {
225 match buffer {
226 PointBuffer::F32(values) => PointBuffer::from_f32(values.clone()),
227 PointBuffer::F64(values) => PointBuffer::F64(values.clone()),
228 PointBuffer::U8(values) => PointBuffer::U8(values.clone()),
229 PointBuffer::U16(values) => PointBuffer::U16(values.clone()),
230 PointBuffer::U32(values) => PointBuffer::U32(values.clone()),
231 PointBuffer::I32(values) => PointBuffer::I32(values.clone()),
232 }
233}
234
235#[cfg(test)]
236mod tests {
237 use super::{FrameEdge, FrameGraph};
238 use spatialrust_core::{
239 FrameId, HasNormals3, HasPositions3, PointBuffer, PointBufferSet, PointCloud,
240 SpatialMetadata, StandardSchemas, Timestamp,
241 };
242 use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
243 use spatialrust_records::{RecordProvenance, SchemaVersion, SpatialRecord};
244
245 #[test]
246 fn composes_chain_base_to_lidar() {
247 let mut graph = FrameGraph::new();
248 graph
249 .insert_edge(FrameEdge {
250 parent: FrameId::new("base"),
251 child: FrameId::new("sensor"),
252 child_t_parent: Isometry3::new(
253 Quat::new(0.0, 0.0, 0.0, 1.0),
254 Vec3::new(1.0, 0.0, 0.0),
255 ),
256 })
257 .unwrap();
258 graph
259 .insert_edge(FrameEdge {
260 parent: FrameId::new("sensor"),
261 child: FrameId::new("lidar"),
262 child_t_parent: Isometry3::new(
263 Quat::new(0.0, 0.0, 0.0, 1.0),
264 Vec3::new(0.0, 2.0, 0.0),
265 ),
266 })
267 .unwrap();
268 let t = graph.lookup(&FrameId::new("base"), &FrameId::new("lidar")).unwrap();
269 let p = t.transform_point(Vec3::new(0.0, 0.0, 0.0));
270 assert!((p.x - 1.0).abs() < 1e-5);
271 assert!((p.y - 2.0).abs() < 1e-5);
272 }
273
274 #[test]
275 fn transforms_record_and_preserves_non_geometry_data() {
276 let mut graph = FrameGraph::new();
277 graph
278 .insert_edge(FrameEdge {
279 parent: FrameId::new("base"),
280 child: FrameId::new("sensor"),
281 child_t_parent: Isometry3::new(
282 Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), std::f32::consts::FRAC_PI_2),
283 Vec3::new(1.0, 2.0, 0.0),
284 ),
285 })
286 .unwrap();
287
288 let mut buffers = PointBufferSet::new();
289 buffers.insert("x", PointBuffer::from_f32(vec![1.0]));
290 buffers.insert("y", PointBuffer::from_f32(vec![0.0]));
291 buffers.insert("z", PointBuffer::from_f32(vec![0.0]));
292 buffers.insert("intensity", PointBuffer::from_f32(vec![7.0]));
293 buffers.insert("normal_x", PointBuffer::from_f32(vec![1.0]));
294 buffers.insert("normal_y", PointBuffer::from_f32(vec![0.0]));
295 buffers.insert("normal_z", PointBuffer::from_f32(vec![0.0]));
296 let mut metadata = SpatialMetadata::new("base", Timestamp::from_nanos(42));
297 metadata.sensor_origin = Some(Vec3::new(0.0, 0.0, 0.0));
298 let cloud =
299 PointCloud::try_from_parts(StandardSchemas::point_xyzinormal(), buffers, metadata)
300 .unwrap();
301 let provenance = RecordProvenance::try_new("bag")
302 .unwrap()
303 .with_stream_id("/lidar")
304 .with_sequence(Some(9));
305 let record = SpatialRecord::try_from_cloud_with_provenance(
306 "point",
307 SchemaVersion::new(1, 0),
308 cloud,
309 provenance.clone(),
310 )
311 .unwrap();
312
313 let transformed = graph.transform_record_to(&record, &FrameId::new("sensor")).unwrap();
314 let (x, y, z) = transformed.cloud().positions3().unwrap();
315 assert!((x[0] - 1.0).abs() < 1e-5);
316 assert!((y[0] - 3.0).abs() < 1e-5);
317 assert!((z[0]).abs() < 1e-5);
318 let (nx, ny, nz) = transformed.cloud().normals3().unwrap();
319 assert!(nx[0].abs() < 1e-5);
320 assert!((ny[0] - 1.0).abs() < 1e-5);
321 assert!(nz[0].abs() < 1e-5);
322 assert_eq!(transformed.metadata().frame_id, FrameId::new("sensor"));
323 assert_eq!(transformed.metadata().sensor_origin, Some(Vec3::new(1.0, 2.0, 0.0)));
324 assert_eq!(transformed.cloud().field("intensity").unwrap().as_f32().unwrap(), &[7.0]);
325 assert_eq!(transformed.provenance(), &provenance);
326 }
327
328 #[test]
329 fn transform_record_rejects_missing_path() {
330 let mut buffers = PointBufferSet::new();
331 buffers.insert("x", PointBuffer::from_f32(vec![0.0]));
332 buffers.insert("y", PointBuffer::from_f32(vec![0.0]));
333 buffers.insert("z", PointBuffer::from_f32(vec![0.0]));
334 let cloud = PointCloud::try_from_parts(
335 StandardSchemas::point_xyz(),
336 buffers,
337 SpatialMetadata::new("source", Timestamp::from_nanos(0)),
338 )
339 .unwrap();
340 let record =
341 SpatialRecord::try_from_cloud("point", SchemaVersion::new(1, 0), cloud).unwrap();
342 let error =
343 FrameGraph::new().transform_record_to(&record, &FrameId::new("target")).unwrap_err();
344 assert!(matches!(error, crate::SyncError::NoTransformPath { .. }));
345 }
346}