spatialrust_registration/
kabsch.rs

1use spatialrust_math::{symmetric_eigen3, Isometry3, Mat3, Quat, Vec3};
2
3/// Estimates the rigid transform that best maps `source` onto `target`.
4#[must_use]
5pub fn estimate_rigid_transform(
6    source: &[Vec3<f32>],
7    target: &[Vec3<f32>],
8) -> Option<Isometry3<f32>> {
9    if source.len() != target.len() || source.len() < 3 {
10        return None;
11    }
12
13    let count = source.len() as f64;
14    let mut mean_source = Vec3::new(0.0_f32, 0.0, 0.0);
15    let mut mean_target = Vec3::new(0.0_f32, 0.0, 0.0);
16    for (src, dst) in source.iter().zip(target) {
17        mean_source = mean_source + *src;
18        mean_target = mean_target + *dst;
19    }
20    mean_source = scale_vec3(mean_source, 1.0 / count as f32);
21    mean_target = scale_vec3(mean_target, 1.0 / count as f32);
22
23    let mut h = Mat3::<f64>::from_rows([0.0; 3], [0.0; 3], [0.0; 3]);
24    for (src, dst) in source.iter().zip(target) {
25        let ps = subtract(*src, mean_source);
26        let pt = subtract(*dst, mean_target);
27        h.m[0][0] += f64::from(ps.x) * f64::from(pt.x);
28        h.m[0][1] += f64::from(ps.x) * f64::from(pt.y);
29        h.m[0][2] += f64::from(ps.x) * f64::from(pt.z);
30        h.m[1][0] += f64::from(ps.y) * f64::from(pt.x);
31        h.m[1][1] += f64::from(ps.y) * f64::from(pt.y);
32        h.m[1][2] += f64::from(ps.y) * f64::from(pt.z);
33        h.m[2][0] += f64::from(ps.z) * f64::from(pt.x);
34        h.m[2][1] += f64::from(ps.z) * f64::from(pt.y);
35        h.m[2][2] += f64::from(ps.z) * f64::from(pt.z);
36    }
37
38    let mut rotation_matrix = rotation_from_cross_covariance(h);
39    if mat3_det_f64(rotation_matrix) < 0.0 {
40        let mut v = eigenvectors_from_cross_covariance(h);
41        for row in 0..3 {
42            v.m[row][2] = -v.m[row][2];
43        }
44        rotation_matrix = rotation_from_eigenvectors(h, v);
45    }
46
47    let rotation_f32 = mat3_f64_to_f32(rotation_matrix);
48    let rotated_mean = rotation_f32.mul_vec3(mean_source);
49    let translation = subtract(mean_target, rotated_mean);
50    Some(Isometry3::new(quat_from_mat3(rotation_f32), translation))
51}
52
53fn rotation_from_cross_covariance(h: Mat3<f64>) -> Mat3<f64> {
54    rotation_from_eigenvectors(h, eigenvectors_from_cross_covariance(h))
55}
56
57fn rotation_from_eigenvectors(h: Mat3<f64>, v: Mat3<f64>) -> Mat3<f64> {
58    let eigen = symmetric_eigen3(mul_mat3_f64(h.transpose(), h));
59    let hv = mul_mat3_f64(h, v);
60    let mut u = Mat3::<f64>::identity();
61    for column in 0..3 {
62        let scale = inv_sqrt(eigen.eigenvalues[column]);
63        for row in 0..3 {
64            u.m[row][column] = hv.m[row][column] * scale;
65        }
66    }
67    mul_mat3_f64(v, u.transpose())
68}
69
70fn eigenvectors_from_cross_covariance(h: Mat3<f64>) -> Mat3<f64> {
71    symmetric_eigen3(mul_mat3_f64(h.transpose(), h)).eigenvectors
72}
73
74fn inv_sqrt(value: f64) -> f64 {
75    if value > 1e-12 {
76        1.0 / value.sqrt()
77    } else {
78        0.0
79    }
80}
81
82fn subtract(a: Vec3<f32>, b: Vec3<f32>) -> Vec3<f32> {
83    Vec3::new(a.x - b.x, a.y - b.y, a.z - b.z)
84}
85
86fn scale_vec3(v: Vec3<f32>, scale: f32) -> Vec3<f32> {
87    Vec3::new(v.x * scale, v.y * scale, v.z * scale)
88}
89
90fn mul_mat3_f64(a: Mat3<f64>, b: Mat3<f64>) -> Mat3<f64> {
91    Mat3::from_rows(
92        [
93            a.m[0][0] * b.m[0][0] + a.m[0][1] * b.m[1][0] + a.m[0][2] * b.m[2][0],
94            a.m[0][0] * b.m[0][1] + a.m[0][1] * b.m[1][1] + a.m[0][2] * b.m[2][1],
95            a.m[0][0] * b.m[0][2] + a.m[0][1] * b.m[1][2] + a.m[0][2] * b.m[2][2],
96        ],
97        [
98            a.m[1][0] * b.m[0][0] + a.m[1][1] * b.m[1][0] + a.m[1][2] * b.m[2][0],
99            a.m[1][0] * b.m[0][1] + a.m[1][1] * b.m[1][1] + a.m[1][2] * b.m[2][1],
100            a.m[1][0] * b.m[0][2] + a.m[1][1] * b.m[1][2] + a.m[1][2] * b.m[2][2],
101        ],
102        [
103            a.m[2][0] * b.m[0][0] + a.m[2][1] * b.m[1][0] + a.m[2][2] * b.m[2][0],
104            a.m[2][0] * b.m[0][1] + a.m[2][1] * b.m[1][1] + a.m[2][2] * b.m[2][1],
105            a.m[2][0] * b.m[0][2] + a.m[2][1] * b.m[1][2] + a.m[2][2] * b.m[2][2],
106        ],
107    )
108}
109
110fn mat3_det_f64(m: Mat3<f64>) -> f64 {
111    m.m[0][0] * (m.m[1][1] * m.m[2][2] - m.m[1][2] * m.m[2][1])
112        - m.m[0][1] * (m.m[1][0] * m.m[2][2] - m.m[1][2] * m.m[2][0])
113        + m.m[0][2] * (m.m[1][0] * m.m[2][1] - m.m[1][1] * m.m[2][0])
114}
115
116fn mat3_f64_to_f32(m: Mat3<f64>) -> Mat3<f32> {
117    Mat3::from_rows(
118        [m.m[0][0] as f32, m.m[0][1] as f32, m.m[0][2] as f32],
119        [m.m[1][0] as f32, m.m[1][1] as f32, m.m[1][2] as f32],
120        [m.m[2][0] as f32, m.m[2][1] as f32, m.m[2][2] as f32],
121    )
122}
123
124fn quat_from_mat3(m: Mat3<f32>) -> Quat<f32> {
125    let trace = m.m[0][0] + m.m[1][1] + m.m[2][2];
126    if trace > 0.0 {
127        let s = (trace + 1.0).sqrt() * 2.0;
128        Quat::new(
129            (m.m[2][1] - m.m[1][2]) / s,
130            (m.m[0][2] - m.m[2][0]) / s,
131            (m.m[1][0] - m.m[0][1]) / s,
132            0.25 * s,
133        )
134        .normalize()
135    } else if m.m[0][0] > m.m[1][1] && m.m[0][0] > m.m[2][2] {
136        let s = (1.0 + m.m[0][0] - m.m[1][1] - m.m[2][2]).sqrt() * 2.0;
137        Quat::new(
138            0.25 * s,
139            (m.m[0][1] + m.m[1][0]) / s,
140            (m.m[0][2] + m.m[2][0]) / s,
141            (m.m[2][1] - m.m[1][2]) / s,
142        )
143        .normalize()
144    } else if m.m[1][1] > m.m[2][2] {
145        let s = (1.0 + m.m[1][1] - m.m[0][0] - m.m[2][2]).sqrt() * 2.0;
146        Quat::new(
147            (m.m[0][1] + m.m[1][0]) / s,
148            0.25 * s,
149            (m.m[1][2] + m.m[2][1]) / s,
150            (m.m[0][2] - m.m[2][0]) / s,
151        )
152        .normalize()
153    } else {
154        let s = (1.0 + m.m[2][2] - m.m[0][0] - m.m[1][1]).sqrt() * 2.0;
155        Quat::new(
156            (m.m[0][2] + m.m[2][0]) / s,
157            (m.m[1][2] + m.m[2][1]) / s,
158            0.25 * s,
159            (m.m[1][0] - m.m[0][1]) / s,
160        )
161        .normalize()
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::estimate_rigid_transform;
168    use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
169
170    #[test]
171    fn recovers_pure_translation() {
172        let source: Vec<Vec3<f32>> = (0..20).map(|i| Vec3::new(i as f32 * 0.1, 0.0, 0.0)).collect();
173        let offset = Vec3::new(0.5, -0.2, 0.1);
174        let target: Vec<Vec3<f32>> = source.iter().map(|point| *point + offset).collect();
175
176        let transform = estimate_rigid_transform(&source, &target).unwrap();
177        assert!((transform.translation().x - offset.x).abs() < 1e-4);
178        assert!((transform.translation().y - offset.y).abs() < 1e-4);
179        assert!((transform.translation().z - offset.z).abs() < 1e-4);
180    }
181
182    #[test]
183    fn recovers_known_rigid_transform() {
184        let target: Vec<Vec3<f32>> = (0..4)
185            .flat_map(|x| {
186                (0..4).flat_map(move |y| {
187                    (0..3).map(move |z| Vec3::new(x as f32, y as f32, z as f32 * 0.2))
188                })
189            })
190            .collect();
191        let misalignment = Isometry3::new(
192            Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), 0.2),
193            Vec3::new(0.3, -0.1, 0.05),
194        );
195        let source: Vec<Vec3<f32>> =
196            target.iter().map(|point| misalignment.transform_point(*point)).collect();
197
198        let estimated = estimate_rigid_transform(&source, &target).unwrap();
199        let composed = estimated.compose(misalignment);
200        let probe = Vec3::new(1.0, 2.0, 0.0);
201        let restored = composed.transform_point(probe);
202        assert!((restored.x - probe.x).abs() < 1e-3);
203        assert!((restored.y - probe.y).abs() < 1e-3);
204        assert!((restored.z - probe.z).abs() < 1e-3);
205    }
206}
spatialrust_registration/kabsch.rs

spatialrust_registration/
kabsch.rs
