1use spatialrust_core::{HasNormals3, HasPositions3, PointCloud, SpatialError, SpatialResult};
10use spatialrust_math::{Isometry3, TransformPoint, Vec3};
11use spatialrust_search::{parallel_worker_count, KdTree, RadiusSearchIndex};
12
13use crate::kabsch::estimate_rigid_transform;
14use crate::registration::{PointCloudRegistration, RegistrationResult};
15
16const BINS: usize = 11;
18const FPFH_DIM: usize = 3 * BINS;
20
21pub const FPFH_DESCRIPTOR_LEN: usize = FPFH_DIM;
23
24pub type FpfhDescriptor = [f32; FPFH_DIM];
26
27fn validate_distance(distance: f32, name: &str) -> SpatialResult<()> {
28 let squared = distance * distance;
29 if !distance.is_finite() || distance <= 0.0 || !squared.is_finite() || squared == 0.0 {
30 return Err(SpatialError::InvalidArgument(format!(
31 "{name} must be positive with finite nonzero f32 square"
32 )));
33 }
34 Ok(())
35}
36
37type Descriptor = FpfhDescriptor;
39
40pub fn fpfh_descriptors(
50 cloud: &PointCloud,
51 feature_radius: f32,
52) -> SpatialResult<Vec<FpfhDescriptor>> {
53 fpfh_descriptors_with_workers(cloud, feature_radius, parallel_worker_count(cloud.len()))
54}
55
56fn fpfh_descriptors_with_workers(
57 cloud: &PointCloud,
58 feature_radius: f32,
59 workers: usize,
60) -> SpatialResult<Vec<FpfhDescriptor>> {
61 validate_distance(feature_radius, "feature_radius")?;
62 let set = PointSet::from_cloud(cloud)?;
63 Ok(compute_fpfh_with_workers(&set, feature_radius, workers))
64}
65
66#[cfg(feature = "register-fpfh-gpu")]
68#[derive(Clone, Debug, PartialEq)]
69pub struct GpuFpfhResult {
70 pub descriptors: Vec<FpfhDescriptor>,
72 pub transfers: spatialrust_core::TransferStats,
74}
75
76#[cfg(feature = "register-fpfh-gpu")]
83pub fn fpfh_descriptors_gpu(
84 runtime: &spatialrust_gpu::WgpuRuntime,
85 cloud: &PointCloud,
86 feature_radius: f32,
87 neighbor_capacity: u32,
88) -> SpatialResult<GpuFpfhResult> {
89 validate_distance(feature_radius, "feature_radius")?;
90 let set = PointSet::from_cloud(cloud)?;
91 let normals: Vec<_> = set.normals.iter().map(|n| [n.x, n.y, n.z]).collect();
92 let input = spatialrust_gpu::GpuFpfhInput::upload(runtime, &set.tree, &normals)?;
93 let neighborhoods =
94 input.radius_neighborhoods_resident(runtime, feature_radius, neighbor_capacity)?;
95 let output = input.descriptors_resident(runtime, &neighborhoods)?;
96 let (descriptors, download) = output.download(runtime)?;
97 let mut transfers = input.upload_transfers();
98 transfers
99 .record(spatialrust_core::TransferDirection::DeviceToHost, download.device_to_host_bytes());
100 Ok(GpuFpfhResult { descriptors, transfers })
101}
102
103#[derive(Clone, Copy, Debug, PartialEq)]
105pub struct FpfhRansacConfig {
106 pub feature_radius: f32,
109 pub max_correspondence_distance: f32,
111 pub ransac_iterations: usize,
113 pub sample_size: usize,
115 pub edge_length_tolerance: f32,
118 pub seed: u64,
120}
121
122impl Default for FpfhRansacConfig {
123 fn default() -> Self {
124 Self {
125 feature_radius: 0.25,
126 max_correspondence_distance: 0.075,
127 ransac_iterations: 4000,
128 sample_size: 3,
129 edge_length_tolerance: 0.9,
130 seed: 0x5eed,
131 }
132 }
133}
134
135impl FpfhRansacConfig {
136 pub fn validate(&self) -> SpatialResult<()> {
138 validate_distance(self.feature_radius, "feature_radius")?;
139 validate_distance(self.max_correspondence_distance, "max_correspondence_distance")?;
140 if self.ransac_iterations == 0 {
141 return Err(SpatialError::InvalidArgument(
142 "ransac_iterations must be positive".to_owned(),
143 ));
144 }
145 if self.sample_size < 3 || self.sample_size.checked_mul(8).is_none() {
146 return Err(SpatialError::InvalidArgument(
147 "sample_size must be at least 3 without overflow".to_owned(),
148 ));
149 }
150 if !self.edge_length_tolerance.is_finite()
151 || self.edge_length_tolerance <= 0.0
152 || self.edge_length_tolerance > 1.0
153 {
154 return Err(SpatialError::InvalidArgument(
155 "edge_length_tolerance must be finite in (0, 1]".to_owned(),
156 ));
157 }
158 Ok(())
159 }
160
161 #[must_use]
163 pub fn with_radius(feature_radius: f32, max_correspondence_distance: f32) -> Self {
164 Self { feature_radius, max_correspondence_distance, ..Self::default() }
165 }
166}
167
168#[derive(Clone, Copy, Debug, PartialEq)]
174pub struct FpfhRansacRegistration {
175 config: FpfhRansacConfig,
176}
177
178impl FpfhRansacRegistration {
179 #[must_use]
181 pub const fn new(config: FpfhRansacConfig) -> Self {
182 Self { config }
183 }
184
185 #[must_use]
187 pub const fn config(&self) -> FpfhRansacConfig {
188 self.config
189 }
190}
191
192impl PointCloudRegistration for FpfhRansacRegistration {
193 fn name(&self) -> &'static str {
194 "FpfhRansacRegistration"
195 }
196
197 fn align(&self, source: &PointCloud, target: &PointCloud) -> SpatialResult<RegistrationResult> {
198 self.config.validate()?;
199 if source.len() < self.config.sample_size || target.len() < self.config.sample_size {
200 return Err(SpatialError::InvalidArgument(
201 "source and target must contain at least sample_size points".to_owned(),
202 ));
203 }
204
205 let src = PointSet::from_cloud(source)?;
206 let tgt = PointSet::from_cloud(target)?;
207
208 let src_features = compute_fpfh(&src, self.config.feature_radius);
209 let tgt_features = compute_fpfh(&tgt, self.config.feature_radius);
210
211 let matches: Vec<usize> =
213 src_features.iter().map(|feature| nearest_feature(feature, &tgt_features)).collect();
214
215 self.ransac(&src, &tgt, &matches)
216 }
217}
218
219impl FpfhRansacRegistration {
220 fn ransac(
221 &self,
222 src: &PointSet,
223 tgt: &PointSet,
224 matches: &[usize],
225 ) -> SpatialResult<RegistrationResult> {
226 let n = src.points.len();
227 let max_sq = self.config.max_correspondence_distance.powi(2);
228 let mut rng = Lcg::new(self.config.seed);
229
230 let mut best_inliers = 0_usize;
231 let mut best_error = f64::INFINITY;
232 let mut best_transform = Isometry3::<f32>::identity();
233
234 for _ in 0..self.config.ransac_iterations {
235 let mut sample = Vec::with_capacity(self.config.sample_size);
237 let mut attempts = 0;
238 while sample.len() < self.config.sample_size && attempts < self.config.sample_size * 8 {
239 let idx = (rng.next_u32() as usize) % n;
240 if !sample.contains(&idx) {
241 sample.push(idx);
242 }
243 attempts += 1;
244 }
245 if sample.len() < self.config.sample_size {
246 continue;
247 }
248
249 let sample_src: Vec<Vec3<f32>> = sample.iter().map(|&i| src.points[i]).collect();
250 let sample_tgt: Vec<Vec3<f32>> =
251 sample.iter().map(|&i| tgt.points[matches[i]]).collect();
252
253 if !edge_lengths_consistent(&sample_src, &sample_tgt, self.config.edge_length_tolerance)
254 {
255 continue;
256 }
257
258 let Some(transform) = estimate_rigid_transform(&sample_src, &sample_tgt) else {
259 continue;
260 };
261
262 let mut inliers = 0_usize;
264 let mut error = 0.0_f64;
265 for (i, &match_idx) in matches.iter().enumerate() {
266 let moved = transform.transform_point(src.points[i]);
267 let dist_sq = (moved - tgt.points[match_idx]).length_squared();
268 if dist_sq <= max_sq {
269 inliers += 1;
270 error += f64::from(dist_sq);
271 }
272 }
273
274 let mean_error = if inliers > 0 { error / inliers as f64 } else { f64::INFINITY };
275 if inliers > best_inliers || (inliers == best_inliers && mean_error < best_error) {
277 best_inliers = inliers;
278 best_error = mean_error;
279 best_transform = transform;
280 }
281 }
282
283 Ok(RegistrationResult {
284 transform: best_transform,
285 fitness: if best_inliers > 0 { best_error } else { f64::INFINITY },
286 iterations: self.config.ransac_iterations,
287 converged: best_inliers >= self.config.sample_size,
288 })
289 }
290}
291
292struct PointSet {
294 points: Vec<Vec3<f32>>,
295 normals: Vec<Vec3<f32>>,
296 tree: KdTree,
297}
298
299impl PointSet {
300 fn from_cloud(cloud: &PointCloud) -> SpatialResult<Self> {
301 let (x, y, z) = cloud.positions3()?;
302 let (nx, ny, nz) = cloud.normals3()?;
303 let points: Vec<Vec3<f32>> = (0..x.len()).map(|i| Vec3::new(x[i], y[i], z[i])).collect();
304 let normals: Vec<Vec3<f32>> =
305 (0..nx.len()).map(|i| Vec3::new(nx[i], ny[i], nz[i])).collect();
306 if points.iter().any(|p| !p.x.is_finite() || !p.y.is_finite() || !p.z.is_finite()) {
307 return Err(SpatialError::InvalidArgument("FPFH positions must be finite".to_owned()));
308 }
309 if !points.is_empty() {
310 let extent_squared: f64 = [x, y, z]
311 .iter()
312 .map(|axis| {
313 let lo = axis.iter().copied().fold(f32::INFINITY, f32::min);
314 let hi = axis.iter().copied().fold(f32::NEG_INFINITY, f32::max);
315 (f64::from(hi) - f64::from(lo)).powi(2)
316 })
317 .sum();
318 if !(extent_squared as f32).is_finite() {
319 return Err(SpatialError::InvalidArgument(
320 "FPFH coordinate extent overflows f32 squared distances".to_owned(),
321 ));
322 }
323 }
324 if normals.iter().any(|n| {
325 !n.x.is_finite()
326 || !n.y.is_finite()
327 || !n.z.is_finite()
328 || !n.length_squared().is_finite()
329 || (n.length_squared() - 1.0).abs() > 1e-3
330 }) {
331 return Err(SpatialError::InvalidArgument(
332 "FPFH normals must be finite unit vectors".to_owned(),
333 ));
334 }
335 let tree = KdTree::from_slices(x, y, z);
336 Ok(Self { points, normals, tree })
337 }
338}
339
340fn darboux(p: Vec3<f32>, np: Vec3<f32>, q: Vec3<f32>, nq: Vec3<f32>) -> Option<(f32, f32, f32)> {
343 let diff = q - p;
344 let dist = diff.length();
345 if dist < 1e-9 {
346 return None;
347 }
348 let u = np;
349 let pq = diff.normalize();
350 let v = pq.cross(u);
351 let v_len = v.length();
352 if v_len < 1e-9 {
353 return None;
354 }
355 let v = v.normalize();
356 let w = u.cross(v);
357
358 let alpha = v.dot(nq); let phi = u.dot(pq); let theta = w.dot(nq).atan2(u.dot(nq)); Some((alpha, phi, theta))
362}
363
364fn bin_index(value: f32, lo: f32, hi: f32) -> usize {
366 let t = ((value - lo) / (hi - lo)).clamp(0.0, 0.999_999);
367 (t * BINS as f32) as usize
368}
369
370fn compute_spfh(set: &PointSet, radius: f32, workers: usize) -> Vec<Descriptor> {
372 let mut spfh = vec![[0.0_f32; FPFH_DIM]; set.points.len()];
373 fill_descriptor_ranges(&mut spfh, workers, |i, spfh_i| {
374 let p = set.points[i];
375 let np = set.normals[i];
376 let neighbors = set.tree.radius_search(p.x, p.y, p.z, radius);
377 let mut count = 0_u32;
378 for neighbor in neighbors {
379 let j = neighbor.index;
380 if j == i {
381 continue;
382 }
383 let Some((alpha, phi, theta)) = darboux(p, np, set.points[j], set.normals[j]) else {
384 continue;
385 };
386 spfh_i[bin_index(alpha, -1.0, 1.0)] += 1.0;
387 spfh_i[BINS + bin_index(phi, -1.0, 1.0)] += 1.0;
388 spfh_i[2 * BINS + bin_index(theta, -std::f32::consts::PI, std::f32::consts::PI)] += 1.0;
389 count += 1;
390 }
391 if count > 0 {
392 for sub in 0..3 {
394 let slice = &mut spfh_i[sub * BINS..(sub + 1) * BINS];
395 let sum: f32 = slice.iter().sum();
396 if sum > 0.0 {
397 for bin in slice {
398 *bin = *bin / sum * 100.0;
399 }
400 }
401 }
402 }
403 });
404 spfh
405}
406
407fn compute_fpfh(set: &PointSet, radius: f32) -> Vec<Descriptor> {
409 compute_fpfh_with_workers(set, radius, parallel_worker_count(set.points.len()))
410}
411
412fn compute_fpfh_with_workers(set: &PointSet, radius: f32, workers: usize) -> Vec<Descriptor> {
413 let spfh = compute_spfh(set, radius, workers);
415 let mut fpfh = vec![[0.0_f32; FPFH_DIM]; set.points.len()];
416 fill_descriptor_ranges(&mut fpfh, workers, |i, fpfh_i| {
417 let p = set.points[i];
418 let neighbors = set.tree.radius_search(p.x, p.y, p.z, radius);
419 let mut sum_weight = 0.0_f32;
420 let mut acc = [0.0_f32; FPFH_DIM];
421 for neighbor in neighbors {
422 let j = neighbor.index;
423 if j == i {
424 continue;
425 }
426 let dist = neighbor.distance_squared.sqrt();
427 if dist < 1e-9 {
428 continue;
429 }
430 let weight = 1.0 / dist;
431 sum_weight += weight;
432 for bin in 0..FPFH_DIM {
433 acc[bin] += weight * spfh[j][bin];
434 }
435 }
436 for bin in 0..FPFH_DIM {
437 fpfh_i[bin] = spfh[i][bin] + if sum_weight > 0.0 { acc[bin] / sum_weight } else { 0.0 };
438 }
439 });
440 fpfh
441}
442
443fn fill_descriptor_ranges<F>(output: &mut [Descriptor], workers: usize, fill: F)
444where
445 F: Fn(usize, &mut Descriptor) + Sync,
446{
447 let workers = workers.max(1).min(output.len().max(1));
448 if workers == 1 {
449 for (i, descriptor) in output.iter_mut().enumerate() {
450 fill(i, descriptor);
451 }
452 return;
453 }
454 let chunk_size = output.len().div_ceil(workers);
455 std::thread::scope(|scope| {
456 for (band, chunk) in output.chunks_mut(chunk_size).enumerate() {
457 let start = band * chunk_size;
458 let fill = &fill;
459 scope.spawn(move || {
460 for (offset, descriptor) in chunk.iter_mut().enumerate() {
461 fill(start + offset, descriptor);
462 }
463 });
464 }
465 });
466}
467
468fn nearest_feature(feature: &Descriptor, candidates: &[Descriptor]) -> usize {
470 let mut best = 0_usize;
471 let mut best_dist = f32::INFINITY;
472 for (idx, candidate) in candidates.iter().enumerate() {
473 let mut dist = 0.0_f32;
474 for bin in 0..FPFH_DIM {
475 let d = feature[bin] - candidate[bin];
476 dist += d * d;
477 if dist >= best_dist {
478 break;
479 }
480 }
481 if dist < best_dist {
482 best_dist = dist;
483 best = idx;
484 }
485 }
486 best
487}
488
489fn edge_lengths_consistent(src: &[Vec3<f32>], tgt: &[Vec3<f32>], tolerance: f32) -> bool {
492 for a in 0..src.len() {
493 for b in (a + 1)..src.len() {
494 let ls = (src[a] - src[b]).length();
495 let lt = (tgt[a] - tgt[b]).length();
496 let lo = ls.min(lt);
497 let hi = ls.max(lt);
498 if hi < 1e-9 {
499 continue;
500 }
501 if lo / hi < tolerance {
502 return false;
503 }
504 }
505 }
506 true
507}
508
509struct Lcg {
511 state: u64,
512}
513
514impl Lcg {
515 fn new(seed: u64) -> Self {
516 let state = seed ^ 0x9e37_79b9_7f4a_7c15;
518 Self { state: if state == 0 { 1 } else { state } }
519 }
520
521 fn next_u32(&mut self) -> u32 {
522 let mut x = self.state;
523 x ^= x << 13;
524 x ^= x >> 7;
525 x ^= x << 17;
526 self.state = x;
527 (x >> 32) as u32
528 }
529}
530
531#[cfg(all(test, not(target_arch = "wasm32")))]
532#[path = "fpfh_parallel_tests.rs"]
533mod parallel_tests;
534
535#[cfg(test)]
536mod tests {
537 use super::{FpfhRansacConfig, FpfhRansacRegistration};
538 use crate::registration::PointCloudRegistration;
539 use spatialrust_core::{DType, FieldSemantic, PointCloudBuilder, PointField, PointSchema};
540 use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
541
542 fn schema_with_normals() -> PointSchema {
543 PointSchema::new()
544 .with_field(PointField::scalar("x", FieldSemantic::PositionX, DType::F32))
545 .with_field(PointField::scalar("y", FieldSemantic::PositionY, DType::F32))
546 .with_field(PointField::scalar("z", FieldSemantic::PositionZ, DType::F32))
547 .with_field(PointField::scalar("normal_x", FieldSemantic::NormalX, DType::F32))
548 .with_field(PointField::scalar("normal_y", FieldSemantic::NormalY, DType::F32))
549 .with_field(PointField::scalar("normal_z", FieldSemantic::NormalZ, DType::F32))
550 }
551
552 fn corner_cloud() -> (Vec<Vec3<f32>>, Vec<Vec3<f32>>) {
555 let mut points = Vec::new();
556 let mut normals = Vec::new();
557 for i in 0..20 {
558 for j in 0..20 {
559 let (a, b) = (i as f32 * 0.05, j as f32 * 0.05);
560 points.push(Vec3::new(a, b, 0.0));
562 normals.push(Vec3::new(0.0, 0.0, 1.0));
563 points.push(Vec3::new(a, 0.0, b + 0.05));
565 normals.push(Vec3::new(0.0, 1.0, 0.0));
566 }
567 }
568 (points, normals)
569 }
570
571 fn build_cloud(points: &[Vec3<f32>], normals: &[Vec3<f32>]) -> spatialrust_core::PointCloud {
572 let mut builder = PointCloudBuilder::new(schema_with_normals());
573 for (p, n) in points.iter().zip(normals) {
574 builder.push_point([p.x, p.y, p.z, n.x, n.y, n.z]).unwrap();
575 }
576 builder.build().unwrap()
577 }
578
579 #[test]
580 fn recovers_coarse_alignment_without_initial_guess() {
581 let (points, normals) = corner_cloud();
582 let target = build_cloud(&points, &normals);
583
584 let misalign = Isometry3::new(
586 Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), 0.6),
587 Vec3::new(0.3, -0.2, 0.1),
588 );
589 let moved_points: Vec<Vec3<f32>> =
590 points.iter().map(|p| misalign.transform_point(*p)).collect();
591 let moved_normals: Vec<Vec3<f32>> =
592 normals.iter().map(|n| misalign.transform_point(*n) - misalign.translation()).collect();
593 let source = build_cloud(&moved_points, &moved_normals);
594
595 let config = FpfhRansacConfig {
596 feature_radius: 0.2,
597 max_correspondence_distance: 0.05,
598 ransac_iterations: 6000,
599 seed: 7,
600 ..FpfhRansacConfig::default()
601 };
602 let result = FpfhRansacRegistration::new(config).align(&source, &target).unwrap();
603 assert!(result.converged);
604
605 let probe = moved_points[123];
609 let restored = result.transform.transform_point(probe);
610 let expected = points[123];
611 let err = (restored - expected).length();
612 assert!(err < 0.15, "coarse alignment residual too large: {err}");
613 }
614
615 #[test]
616 fn rejects_bad_params() {
617 let (points, normals) = corner_cloud();
618 let cloud = build_cloud(&points, &normals);
619 let config = FpfhRansacConfig { feature_radius: 0.0, ..FpfhRansacConfig::default() };
620 assert!(FpfhRansacRegistration::new(config).align(&cloud, &cloud).is_err());
621 }
622
623 #[test]
624 fn public_descriptors_have_expected_shape() {
625 use super::{fpfh_descriptors, FPFH_DESCRIPTOR_LEN};
626 let (points, normals) = corner_cloud();
627 let cloud = build_cloud(&points, &normals);
628 let descriptors = fpfh_descriptors(&cloud, 0.2).unwrap();
629 assert_eq!(descriptors.len(), cloud.len());
630 assert_eq!(FPFH_DESCRIPTOR_LEN, 33);
631 let total: f32 = descriptors[descriptors.len() / 2].iter().sum();
634 assert!(total > 0.0);
635 }
636
637 #[test]
638 fn validates_configuration_before_descriptor_work() {
639 for value in [0.0, -1.0, f32::NAN, f32::INFINITY, f32::MAX, f32::MIN_POSITIVE] {
640 assert!(FpfhRansacConfig { feature_radius: value, ..FpfhRansacConfig::default() }
641 .validate()
642 .is_err());
643 assert!(FpfhRansacConfig {
644 max_correspondence_distance: value,
645 ..FpfhRansacConfig::default()
646 }
647 .validate()
648 .is_err());
649 }
650 for value in [0.0, -1.0, 1.1, f32::NAN, f32::INFINITY] {
651 assert!(FpfhRansacConfig {
652 edge_length_tolerance: value,
653 ..FpfhRansacConfig::default()
654 }
655 .validate()
656 .is_err());
657 }
658 for value in [0, 2, usize::MAX] {
659 assert!(FpfhRansacConfig { sample_size: value, ..FpfhRansacConfig::default() }
660 .validate()
661 .is_err());
662 }
663 assert!(FpfhRansacConfig { ransac_iterations: 0, ..FpfhRansacConfig::default() }
664 .validate()
665 .is_err());
666 assert!(FpfhRansacConfig::default().validate().is_ok());
667 }
668
669 #[test]
670 fn rejects_nonfinite_positions_and_nonunit_normals() {
671 let (mut points, mut normals) = corner_cloud();
672 points[0].x = f32::NAN;
673 assert!(super::fpfh_descriptors(&build_cloud(&points, &normals), 0.2).is_err());
674 points[0].x = 0.0;
675 for normal in
676 [Vec3::new(0.0, 0.0, 0.0), Vec3::new(0.0, 0.0, 2.0), Vec3::new(f32::INFINITY, 0.0, 1.0)]
677 {
678 normals[0] = normal;
679 assert!(super::fpfh_descriptors(&build_cloud(&points, &normals), 0.2).is_err());
680 }
681 assert!(FpfhRansacRegistration::new(FpfhRansacConfig::default())
682 .align(&build_cloud(&points[..2], &normals[..2]), &build_cloud(&points, &normals))
683 .is_err());
684 }
685
686 #[test]
687 fn all_seeds_produce_a_nonzero_rng_state() {
688 let mut rng = super::Lcg::new(0x9e37_79b9_7f4a_7c15);
689 assert!((0..8).any(|_| rng.next_u32() != 0));
690 }
691
692 #[test]
693 fn rejects_overflowing_extent_and_oversized_samples() {
694 let (mut points, normals) = corner_cloud();
695 let ordinary = build_cloud(&points, &normals);
696 let config =
697 FpfhRansacConfig { sample_size: ordinary.len() + 1, ..FpfhRansacConfig::default() };
698 assert!(FpfhRansacRegistration::new(config).align(&ordinary, &ordinary).is_err());
699 points[0].x = f32::MAX;
700 points[1].x = -f32::MAX;
701 let error = super::fpfh_descriptors(&build_cloud(&points, &normals), 0.2).unwrap_err();
702 assert!(error.to_string().contains("extent"));
703 }
704}