spatialrust_registration/
icp.rs

1use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
2use spatialrust_math::{Isometry3, TransformPoint, Vec3};
3use spatialrust_search::{KdTree, NearestNeighborIndex};
4
5use crate::kabsch::estimate_rigid_transform;
6use crate::registration::{PointCloudRegistration, RegistrationResult};
7
8/// Configuration for point-to-point ICP.
9#[derive(Clone, Copy, Debug, PartialEq)]
10pub struct IcpConfig {
11    /// Maximum number of ICP iterations.
12    pub max_iterations: usize,
13    /// Maximum correspondence distance.
14    pub max_correspondence_distance: f32,
15    /// Stop when the transform update is smaller than this threshold.
16    pub transformation_epsilon: f64,
17    /// Stop when fitness improvement is smaller than this threshold.
18    pub fitness_epsilon: f64,
19    /// Minimum number of correspondences required per iteration.
20    pub min_correspondences: usize,
21    /// Initial transform guess mapping source into target frame.
22    pub initial_guess: Isometry3<f32>,
23}
24
25impl Default for IcpConfig {
26    fn default() -> Self {
27        Self {
28            max_iterations: 50,
29            max_correspondence_distance: 1.0,
30            transformation_epsilon: 1e-8,
31            fitness_epsilon: 1e-6,
32            min_correspondences: 3,
33            initial_guess: Isometry3::identity(),
34        }
35    }
36}
37
38impl IcpConfig {
39    /// Creates a config with the given correspondence distance.
40    #[must_use]
41    pub fn with_correspondence_distance(max_correspondence_distance: f32) -> Self {
42        Self { max_correspondence_distance, ..Self::default() }
43    }
44}
45
46/// Point-to-point ICP registration.
47#[derive(Clone, Copy, Debug, PartialEq)]
48pub struct IcpRegistration {
49    config: IcpConfig,
50}
51
52impl IcpRegistration {
53    /// Creates an ICP registration algorithm from config.
54    #[must_use]
55    pub const fn new(config: IcpConfig) -> Self {
56        Self { config }
57    }
58
59    /// Returns the ICP config.
60    #[must_use]
61    pub const fn config(&self) -> IcpConfig {
62        self.config
63    }
64
65    /// Aligns `source` to `target` using iterative closest point.
66    pub fn align_with_diagnostics(
67        &self,
68        source: &PointCloud,
69        target: &PointCloud,
70    ) -> SpatialResult<RegistrationResult> {
71        if source.is_empty() || target.is_empty() {
72            return Err(SpatialError::InvalidArgument(
73                "ICP requires non-empty source and target point clouds".to_owned(),
74            ));
75        }
76
77        let (source_x, source_y, source_z) = source.positions3()?;
78        let (target_x, target_y, target_z) = target.positions3()?;
79        let tree = KdTree::from_slices(target_x, target_y, target_z);
80        let max_distance_squared =
81            self.config.max_correspondence_distance * self.config.max_correspondence_distance;
82
83        let mut transform = self.config.initial_guess;
84        let mut transformed = Vec::with_capacity(source.len());
85        for index in 0..source.len() {
86            transformed.push(Vec3::new(source_x[index], source_y[index], source_z[index]));
87        }
88        apply_transform_in_place(&mut transformed, source_x, source_y, source_z, transform);
89
90        let mut iterations = 0usize;
91        let mut converged = false;
92
93        for _ in 0..self.config.max_iterations {
94            iterations += 1;
95            let mut pairs_source = Vec::new();
96            let mut pairs_target = Vec::new();
97
98            for point in &transformed {
99                let Some(neighbor) = tree.nearest_one(point.x, point.y, point.z) else {
100                    continue;
101                };
102                if neighbor.distance_squared <= max_distance_squared {
103                    pairs_source.push(*point);
104                    pairs_target.push(Vec3::new(
105                        target_x[neighbor.index],
106                        target_y[neighbor.index],
107                        target_z[neighbor.index],
108                    ));
109                }
110            }
111
112            if pairs_source.len() < self.config.min_correspondences {
113                return Err(SpatialError::InvalidArgument(format!(
114                    "ICP found only {} correspondences, minimum is {}",
115                    pairs_source.len(),
116                    self.config.min_correspondences
117                )));
118            }
119
120            let Some(delta) = estimate_rigid_transform(&pairs_source, &pairs_target) else {
121                return Err(SpatialError::InvalidArgument(
122                    "ICP failed to estimate a rigid transform".to_owned(),
123                ));
124            };
125
126            transform = delta.compose(transform);
127            apply_transform_in_place(&mut transformed, source_x, source_y, source_z, transform);
128
129            let fitness = final_fitness(
130                &transformed,
131                target_x,
132                target_y,
133                target_z,
134                &tree,
135                max_distance_squared,
136            );
137            if fitness < self.config.fitness_epsilon {
138                converged = true;
139                break;
140            }
141            if transform_delta_below_epsilon(delta, self.config.transformation_epsilon) {
142                converged = true;
143                break;
144            }
145        }
146
147        Ok(RegistrationResult {
148            transform,
149            fitness: final_fitness(
150                &transformed,
151                target_x,
152                target_y,
153                target_z,
154                &tree,
155                max_distance_squared,
156            ),
157            iterations,
158            converged,
159        })
160    }
161}
162
163impl PointCloudRegistration for IcpRegistration {
164    fn name(&self) -> &'static str {
165        "IcpRegistration"
166    }
167
168    fn align(&self, source: &PointCloud, target: &PointCloud) -> SpatialResult<RegistrationResult> {
169        self.align_with_diagnostics(source, target)
170    }
171}
172
173fn apply_transform_in_place(
174    transformed: &mut [Vec3<f32>],
175    source_x: &[f32],
176    source_y: &[f32],
177    source_z: &[f32],
178    transform: Isometry3<f32>,
179) {
180    for (index, point) in transformed.iter_mut().enumerate() {
181        *point =
182            transform.transform_point(Vec3::new(source_x[index], source_y[index], source_z[index]));
183    }
184}
185
186fn mean_squared_error(source: &[Vec3<f32>], target: &[Vec3<f32>]) -> f64 {
187    let mut sum = 0.0_f64;
188    for (src, dst) in source.iter().zip(target) {
189        let dx = f64::from(src.x - dst.x);
190        let dy = f64::from(src.y - dst.y);
191        let dz = f64::from(src.z - dst.z);
192        sum += dx * dx + dy * dy + dz * dz;
193    }
194    sum / source.len() as f64
195}
196
197fn final_fitness(
198    transformed: &[Vec3<f32>],
199    target_x: &[f32],
200    target_y: &[f32],
201    target_z: &[f32],
202    tree: &KdTree,
203    max_distance_squared: f32,
204) -> f64 {
205    let mut pairs_source = Vec::new();
206    let mut pairs_target = Vec::new();
207    for point in transformed {
208        let Some(neighbor) = tree.nearest_one(point.x, point.y, point.z) else {
209            continue;
210        };
211        if neighbor.distance_squared <= max_distance_squared {
212            pairs_source.push(*point);
213            pairs_target.push(Vec3::new(
214                target_x[neighbor.index],
215                target_y[neighbor.index],
216                target_z[neighbor.index],
217            ));
218        }
219    }
220    if pairs_source.is_empty() {
221        return f64::MAX;
222    }
223    mean_squared_error(&pairs_source, &pairs_target)
224}
225
226fn transform_delta_below_epsilon(delta: Isometry3<f32>, epsilon: f64) -> bool {
227    let translation = delta.translation();
228    let translation_norm = f64::from(
229        (translation.x * translation.x
230            + translation.y * translation.y
231            + translation.z * translation.z)
232            .sqrt(),
233    );
234    translation_norm < epsilon
235}
236
237#[cfg(test)]
238mod tests {
239    use super::{IcpConfig, IcpRegistration};
240    use crate::registration::PointCloudRegistration;
241    use crate::transform::transform_point_cloud;
242    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
243    use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
244
245    fn plane_cloud() -> spatialrust_core::PointCloud {
246        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
247        for x in 0..6 {
248            for y in 0..6 {
249                for z in 0..3 {
250                    builder
251                        .push_point([x as f32 * 0.05, y as f32 * 0.05, z as f32 * 0.05])
252                        .unwrap();
253                }
254            }
255        }
256        builder.build().unwrap()
257    }
258
259    #[test]
260    fn aligns_translated_source_to_target() {
261        let target = plane_cloud();
262        let shift = Isometry3::new(Quat::<f32>::identity(), Vec3::new(0.02, -0.01, 0.0));
263        let source = transform_point_cloud(&target, shift).unwrap();
264
265        let registration = IcpRegistration::new(IcpConfig {
266            max_correspondence_distance: 0.1,
267            max_iterations: 30,
268            ..IcpConfig::default()
269        });
270        let result = registration.align(&source, &target).unwrap();
271        assert!(result.fitness < 1e-4);
272        assert!(result.converged);
273
274        let composed = result.transform.compose(shift);
275        let probe = Vec3::new(0.2, 0.3, 0.0);
276        let restored = composed.transform_point(probe);
277        assert!((restored.x - probe.x).abs() < 5e-3);
278        assert!((restored.y - probe.y).abs() < 5e-3);
279    }
280
281    #[test]
282    fn aligns_rotated_and_translated_source() {
283        let target = plane_cloud();
284        let misalignment = Isometry3::new(
285            Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), 0.15),
286            Vec3::new(0.02, 0.01, 0.0),
287        );
288        let source = transform_point_cloud(&target, misalignment).unwrap();
289
290        let registration = IcpRegistration::new(IcpConfig {
291            max_correspondence_distance: 0.1,
292            max_iterations: 40,
293            ..IcpConfig::default()
294        });
295        let result = registration.align(&source, &target).unwrap();
296        assert!(result.fitness < 1e-3);
297
298        let composed = result.transform.compose(misalignment);
299        let probe = Vec3::new(0.15, 0.25, 0.0);
300        let restored = composed.transform_point(probe);
301        assert!((restored.x - probe.x).abs() < 1e-2);
302        assert!((restored.y - probe.y).abs() < 1e-2);
303    }
304}
spatialrust_registration/icp.rs

spatialrust_registration/
icp.rs
