1#![allow(clippy::needless_range_loop)]
6#[allow(unused_imports)]
7use super::functions::*;
8#[allow(unused_imports)]
9use super::functions_2::*;
10use std::collections::HashMap;
11
12pub struct PointCloudFilter;
14impl PointCloudFilter {
15 pub fn voxel_downsample(cloud: &PointCloud, voxel_size: f64) -> PointCloud {
17 if cloud.points.is_empty() || voxel_size <= 0.0 {
18 return PointCloud::new();
19 }
20 let inv = 1.0 / voxel_size;
21 let mut voxels: HashMap<(i64, i64, i64), Vec<usize>> = HashMap::new();
22 for (i, &p) in cloud.points.iter().enumerate() {
23 let key = (
24 (p[0] * inv).floor() as i64,
25 (p[1] * inv).floor() as i64,
26 (p[2] * inv).floor() as i64,
27 );
28 voxels.entry(key).or_default().push(i);
29 }
30 let mut out = PointCloud::new();
31 for indices in voxels.values() {
32 let n = indices.len() as f64;
33 let mut sum = [0.0f64; 3];
34 for &i in indices {
35 let p = cloud.points[i];
36 for k in 0..3 {
37 sum[k] += p[k];
38 }
39 }
40 out.add_point([sum[0] / n, sum[1] / n, sum[2] / n]);
41 }
42 out
43 }
44 pub fn statistical_outlier_removal(
47 cloud: &PointCloud,
48 k: usize,
49 std_factor: f64,
50 ) -> PointCloud {
51 if cloud.points.is_empty() {
52 return PointCloud::new();
53 }
54 let tree = KdTree3D::build(&cloud.points);
55 let mean_dists: Vec<f64> = cloud
56 .points
57 .iter()
58 .map(|&p| {
59 let neighbors = tree.k_nearest(p, k + 1);
60 let dists: Vec<f64> = neighbors
61 .iter()
62 .filter(|&&(idx, _)| cloud.points[idx] != p || idx != 0)
63 .map(|&(_, d)| d.sqrt())
64 .collect();
65 if dists.is_empty() {
66 0.0
67 } else {
68 dists.iter().sum::<f64>() / dists.len() as f64
69 }
70 })
71 .collect();
72 let mean = mean_dists.iter().sum::<f64>() / mean_dists.len() as f64;
73 let variance = mean_dists
74 .iter()
75 .map(|&d| (d - mean) * (d - mean))
76 .sum::<f64>()
77 / mean_dists.len() as f64;
78 let std = variance.sqrt();
79 let threshold = mean + std_factor * std;
80 let mut out = PointCloud::new();
81 for (i, &md) in mean_dists.iter().enumerate() {
82 if md <= threshold {
83 out.add_point(cloud.points[i]);
84 if !cloud.normals.is_empty() {
85 out.normals.push(cloud.normals[i]);
86 }
87 if !cloud.colors.is_empty() {
88 out.colors.push(cloud.colors[i]);
89 }
90 }
91 }
92 out
93 }
94 pub fn radius_outlier_removal(
96 cloud: &PointCloud,
97 radius: f64,
98 min_neighbors: usize,
99 ) -> PointCloud {
100 if cloud.points.is_empty() {
101 return PointCloud::new();
102 }
103 let tree = KdTree3D::build(&cloud.points);
104 let mut out = PointCloud::new();
105 for (i, &p) in cloud.points.iter().enumerate() {
106 let count = tree.range_search(p, radius).len();
107 if count.saturating_sub(1) >= min_neighbors {
108 out.add_point(p);
109 if !cloud.normals.is_empty() {
110 out.normals.push(cloud.normals[i]);
111 }
112 if !cloud.colors.is_empty() {
113 out.colors.push(cloud.colors[i]);
114 }
115 }
116 }
117 out
118 }
119}
120pub struct KdTree3D {
122 pub(super) nodes: Vec<KdNode3D>,
123 pub(super) points: Vec<[f64; 3]>,
124 pub(super) root: Option<Box<KdNode3D>>,
125}
126impl KdTree3D {
127 pub fn build(points: &[[f64; 3]]) -> Self {
129 let mut indices: Vec<usize> = (0..points.len()).collect();
130 let root = Self::build_recursive(points, &mut indices, 0);
131 Self {
132 nodes: Vec::new(),
133 points: points.to_vec(),
134 root,
135 }
136 }
137 fn build_recursive(
138 points: &[[f64; 3]],
139 indices: &mut [usize],
140 depth: usize,
141 ) -> Option<Box<KdNode3D>> {
142 if indices.is_empty() {
143 return None;
144 }
145 let axis = depth % 3;
146 indices.sort_unstable_by(|&a, &b| {
147 points[a][axis]
148 .partial_cmp(&points[b][axis])
149 .unwrap_or(std::cmp::Ordering::Equal)
150 });
151 let mid = indices.len() / 2;
152 let idx = indices[mid];
153 let left = Self::build_recursive(points, &mut indices[..mid], depth + 1);
154 let right = Self::build_recursive(points, &mut indices[mid + 1..], depth + 1);
155 Some(Box::new(KdNode3D {
156 point_idx: idx,
157 split_axis: axis,
158 left,
159 right,
160 }))
161 }
162 pub fn nearest_neighbor(&self, query: [f64; 3]) -> (usize, f64) {
164 let mut best = (0, f64::MAX);
165 if let Some(ref root) = self.root {
166 Self::nn_search(&self.points, root, query, &mut best);
167 }
168 best
169 }
170 fn nn_search(points: &[[f64; 3]], node: &KdNode3D, query: [f64; 3], best: &mut (usize, f64)) {
171 let p = points[node.point_idx];
172 let d2 = dist2(p, query);
173 if d2 < best.1 {
174 *best = (node.point_idx, d2);
175 }
176 let axis = node.split_axis;
177 let diff = query[axis] - p[axis];
178 let (near, far) = if diff <= 0.0 {
179 (node.left.as_deref(), node.right.as_deref())
180 } else {
181 (node.right.as_deref(), node.left.as_deref())
182 };
183 if let Some(near_node) = near {
184 Self::nn_search(points, near_node, query, best);
185 }
186 if diff * diff < best.1
187 && let Some(far_node) = far
188 {
189 Self::nn_search(points, far_node, query, best);
190 }
191 }
192 pub fn k_nearest(&self, query: [f64; 3], k: usize) -> Vec<(usize, f64)> {
194 if k == 0 || self.points.is_empty() {
195 return Vec::new();
196 }
197 let mut heap: Vec<(usize, f64)> = Vec::with_capacity(k + 1);
198 if let Some(ref root) = self.root {
199 Self::knn_search(&self.points, root, query, k, &mut heap);
200 }
201 heap.sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
202 heap
203 }
204 fn knn_search(
205 points: &[[f64; 3]],
206 node: &KdNode3D,
207 query: [f64; 3],
208 k: usize,
209 heap: &mut Vec<(usize, f64)>,
210 ) {
211 let p = points[node.point_idx];
212 let d2 = dist2(p, query);
213 heap.push((node.point_idx, d2));
214 if heap.len() > k {
215 let worst_pos = heap
216 .iter()
217 .enumerate()
218 .max_by(|a, b| {
219 a.1.1
220 .partial_cmp(&b.1.1)
221 .unwrap_or(std::cmp::Ordering::Equal)
222 })
223 .map(|(i, _)| i)
224 .unwrap_or(0);
225 heap.swap_remove(worst_pos);
226 }
227 let worst_dist2 = heap.iter().map(|e| e.1).fold(0.0_f64, f64::max);
228 let axis = node.split_axis;
229 let diff = query[axis] - p[axis];
230 let (near, far) = if diff <= 0.0 {
231 (node.left.as_deref(), node.right.as_deref())
232 } else {
233 (node.right.as_deref(), node.left.as_deref())
234 };
235 if let Some(near_node) = near {
236 Self::knn_search(points, near_node, query, k, heap);
237 }
238 let current_worst = heap.iter().map(|e| e.1).fold(0.0_f64, f64::max);
239 if (heap.len() < k || diff * diff < current_worst)
240 && let Some(far_node) = far
241 {
242 Self::knn_search(points, far_node, query, k, heap);
243 }
244 let _ = worst_dist2;
245 }
246 pub fn range_search(&self, center: [f64; 3], radius: f64) -> Vec<usize> {
248 let mut result = Vec::new();
249 let r2 = radius * radius;
250 if let Some(ref root) = self.root {
251 Self::range_recursive(&self.points, root, center, r2, &mut result);
252 }
253 result
254 }
255 fn range_recursive(
256 points: &[[f64; 3]],
257 node: &KdNode3D,
258 center: [f64; 3],
259 r2: f64,
260 result: &mut Vec<usize>,
261 ) {
262 let p = points[node.point_idx];
263 if dist2(p, center) <= r2 {
264 result.push(node.point_idx);
265 }
266 let axis = node.split_axis;
267 let diff = center[axis] - p[axis];
268 if (diff * diff <= r2 || diff <= 0.0)
269 && let Some(ref left) = node.left
270 {
271 Self::range_recursive(points, left, center, r2, result);
272 }
273 if (diff * diff <= r2 || diff > 0.0)
274 && let Some(ref right) = node.right
275 {
276 Self::range_recursive(points, right, center, r2, result);
277 }
278 }
279}
280pub struct KdNode3D {
282 pub point_idx: usize,
284 pub split_axis: usize,
286 pub left: Option<Box<KdNode3D>>,
288 pub right: Option<Box<KdNode3D>>,
290}
291pub struct NormalEstimation {
293 pub k_neighbors: usize,
295}
296impl NormalEstimation {
297 pub fn new(k_neighbors: usize) -> Self {
299 Self { k_neighbors }
300 }
301 pub fn estimate(&self, cloud: &PointCloud) -> Vec<[f64; 3]> {
303 let tree = KdTree3D::build(&cloud.points);
304 cloud
305 .points
306 .iter()
307 .map(|&p| {
308 let neighbors = tree.k_nearest(p, self.k_neighbors);
309 let pts: Vec<[f64; 3]> = neighbors
310 .iter()
311 .map(|&(idx, _)| cloud.points[idx])
312 .collect();
313 if pts.len() < 3 {
314 return [0.0, 0.0, 1.0];
315 }
316 let cov = Self::covariance_matrix(&pts);
317 Self::smallest_eigenvector_3x3_sym(&cov)
318 })
319 .collect()
320 }
321 pub fn covariance_matrix(pts: &[[f64; 3]]) -> [[f64; 3]; 3] {
323 if pts.is_empty() {
324 return [[0.0; 3]; 3];
325 }
326 let n = pts.len() as f64;
327 let mut centroid = [0.0f64; 3];
328 for &p in pts {
329 for i in 0..3 {
330 centroid[i] += p[i];
331 }
332 }
333 for i in 0..3 {
334 centroid[i] /= n;
335 }
336 let mut cov = [[0.0f64; 3]; 3];
337 for &p in pts {
338 let d = sub(p, centroid);
339 for i in 0..3 {
340 for j in 0..3 {
341 cov[i][j] += d[i] * d[j];
342 }
343 }
344 }
345 cov
346 }
347 pub fn smallest_eigenvector_3x3_sym(m: &[[f64; 3]; 3]) -> [f64; 3] {
350 let mut v = [1.0f64, 1.0, 1.0];
351 for _ in 0..20 {
352 v = mat_vec_3(m, v);
353 let len = length(v);
354 if len < 1e-15 {
355 break;
356 }
357 v = scale(v, 1.0 / len);
358 }
359 let lambda_max = dot(mat_vec_3(m, v), v);
360 let shifted = [
361 [lambda_max - m[0][0], -m[0][1], -m[0][2]],
362 [-m[1][0], lambda_max - m[1][1], -m[1][2]],
363 [-m[2][0], -m[2][1], lambda_max - m[2][2]],
364 ];
365 let mut u = [1.0f64, 0.0, 0.0];
366 for _ in 0..50 {
367 u = mat_vec_3(&shifted, u);
368 let len = length(u);
369 if len < 1e-15 {
370 break;
371 }
372 u = scale(u, 1.0 / len);
373 }
374 normalize(u)
375 }
376}
377pub struct RansacPlaneResult {
379 pub normal: [f64; 3],
381 pub point_on_plane: [f64; 3],
383 pub inliers: Vec<usize>,
385 pub n_inliers: usize,
387}
388pub struct IcpResult {
390 pub rotation: [[f64; 3]; 3],
392 pub translation: [f64; 3],
394 pub final_error: f64,
396 pub iterations: usize,
398}
399impl IcpResult {
400 pub fn apply_to(&self, cloud: &PointCloud) -> PointCloud {
402 let mut out = PointCloud::new();
403 for &p in &cloud.points {
404 let rp = mat_vec_3(&self.rotation, p);
405 out.add_point(add(rp, self.translation));
406 }
407 if !cloud.normals.is_empty() {
408 for &n in &cloud.normals {
409 out.normals.push(mat_vec_3(&self.rotation, n));
410 }
411 }
412 out.colors = cloud.colors.clone();
413 out
414 }
415}
416pub struct PointCloud {
418 pub points: Vec<[f64; 3]>,
420 pub normals: Vec<[f64; 3]>,
422 pub colors: Vec<[f32; 3]>,
424}
425impl PointCloud {
426 pub fn new() -> Self {
428 Self {
429 points: Vec::new(),
430 normals: Vec::new(),
431 colors: Vec::new(),
432 }
433 }
434 pub fn add_point(&mut self, p: [f64; 3]) {
436 self.points.push(p);
437 }
438 pub fn len(&self) -> usize {
440 self.points.len()
441 }
442 pub fn is_empty(&self) -> bool {
444 self.points.is_empty()
445 }
446 pub fn bounding_box(&self) -> ([f64; 3], [f64; 3]) {
448 if self.points.is_empty() {
449 return ([0.0; 3], [0.0; 3]);
450 }
451 let mut mn = self.points[0];
452 let mut mx = self.points[0];
453 for &p in &self.points {
454 for i in 0..3 {
455 if p[i] < mn[i] {
456 mn[i] = p[i];
457 }
458 if p[i] > mx[i] {
459 mx[i] = p[i];
460 }
461 }
462 }
463 (mn, mx)
464 }
465 pub fn centroid(&self) -> [f64; 3] {
467 if self.points.is_empty() {
468 return [0.0; 3];
469 }
470 let n = self.points.len() as f64;
471 let mut sum = [0.0f64; 3];
472 for &p in &self.points {
473 for i in 0..3 {
474 sum[i] += p[i];
475 }
476 }
477 [sum[0] / n, sum[1] / n, sum[2] / n]
478 }
479 pub fn translate(&mut self, offset: [f64; 3]) {
481 for p in &mut self.points {
482 *p = add(*p, offset);
483 }
484 }
485 pub fn scale_uniform(&mut self, s: f64) {
487 for p in &mut self.points {
488 *p = scale(*p, s);
489 }
490 }
491 pub fn from_grid(nx: usize, ny: usize, height_fn: impl Fn(f64, f64) -> f64, dx: f64) -> Self {
495 let mut cloud = Self::new();
496 for i in 0..nx {
497 for j in 0..ny {
498 let x = i as f64 * dx;
499 let y = j as f64 * dx;
500 let z = height_fn(x, y);
501 cloud.add_point([x, y, z]);
502 }
503 }
504 cloud
505 }
506}
507impl PointCloud {
508 pub fn from_points(points: Vec<[f64; 3]>) -> Self {
510 Self {
511 points,
512 normals: Vec::new(),
513 colors: Vec::new(),
514 }
515 }
516 pub fn estimate_normals_pca(&mut self, k: usize) {
519 let tree = KdTree3D::build(&self.points);
520 self.normals = self
521 .points
522 .iter()
523 .map(|&p| {
524 let neighbors = tree.k_nearest(p, k);
525 let pts: Vec<[f64; 3]> =
526 neighbors.iter().map(|&(idx, _)| self.points[idx]).collect();
527 if pts.len() < 3 {
528 return [0.0, 0.0, 1.0];
529 }
530 let cov = NormalEstimation::covariance_matrix(&pts);
531 NormalEstimation::smallest_eigenvector_3x3_sym(&cov)
532 })
533 .collect();
534 }
535 pub fn voxel_downsample(&self, voxel_size: f64) -> Self {
537 PointCloudFilter::voxel_downsample(self, voxel_size)
538 }
539 pub fn statistical_outlier_removal(&self, k: usize, std_ratio: f64) -> Self {
541 PointCloudFilter::statistical_outlier_removal(self, k, std_ratio)
542 }
543}
544impl PointCloud {
545 pub fn farthest_point_sample(&self, k: usize) -> Self {
550 let indices = farthest_point_sampling(&self.points, k, 0);
551 let mut out = PointCloud::new();
552 for &i in &indices {
553 out.add_point(self.points[i]);
554 if !self.normals.is_empty() {
555 out.normals.push(self.normals[i]);
556 }
557 if !self.colors.is_empty() {
558 out.colors.push(self.colors[i]);
559 }
560 }
561 out
562 }
563 pub fn fit_plane_ransac(&self, n_iter: usize, threshold: f64) -> Option<RansacPlaneResult> {
567 ransac_fit_plane(&self.points, n_iter, threshold)
568 }
569 pub fn pca_obb(&self) -> ([[f64; 3]; 3], [f64; 3], [f64; 3]) {
573 pca_obb(&self.points)
574 }
575 pub fn compute_principal_curvatures(&self, k: usize) -> Vec<(f64, f64)> {
583 if self.points.is_empty() {
584 return Vec::new();
585 }
586 let tree = KdTree3D::build(&self.points);
587 self.points
588 .iter()
589 .map(|&p| {
590 let neighbors = tree.k_nearest(p, k + 1);
591 let pts: Vec<[f64; 3]> =
592 neighbors.iter().map(|&(idx, _)| self.points[idx]).collect();
593 if pts.len() < 3 {
594 return (0.0, 0.0);
595 }
596 let n = pts.len() as f64;
597 let mut centroid = [0.0f64; 3];
598 for &q in &pts {
599 for i in 0..3 {
600 centroid[i] += q[i];
601 }
602 }
603 for i in 0..3 {
604 centroid[i] /= n;
605 }
606 let mut cov = [[0.0f64; 3]; 3];
607 for &q in &pts {
608 let d = sub(q, centroid);
609 for i in 0..3 {
610 for j in 0..3 {
611 cov[i][j] += d[i] * d[j];
612 }
613 }
614 }
615 for i in 0..3 {
616 for j in 0..3 {
617 cov[i][j] /= n;
618 }
619 }
620 let eigenvalues = jacobi_eigenvalues_3x3(cov);
621 let mut ev = eigenvalues;
622 ev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
623 (ev[0], ev[1])
624 })
625 .collect()
626 }
627 pub fn icp_register(&self, target: &PointCloud) -> IcpResult {
632 IcpRegistration::align(self, target)
633 }
634}
635pub struct IcpRegistration {
637 pub max_iterations: usize,
639 pub convergence_tol: f64,
641}
642impl IcpRegistration {
643 pub fn new(max_iterations: usize, convergence_tol: f64) -> Self {
645 Self {
646 max_iterations,
647 convergence_tol,
648 }
649 }
650 pub fn align(source: &PointCloud, target: &PointCloud) -> IcpResult {
652 let icp = Self::new(50, 1e-6);
653 icp.run(source, target)
654 }
655 pub fn run(&self, source: &PointCloud, target: &PointCloud) -> IcpResult {
657 let mut rot = identity_3x3();
658 let mut trans = [0.0f64; 3];
659 let mut current: Vec<[f64; 3]> = source.points.clone();
660 let target_tree = KdTree3D::build(&target.points);
661 let mut prev_error = f64::MAX;
662 let mut iters = 0;
663 for iter in 0..self.max_iterations {
664 iters = iter + 1;
665 let pairs: Vec<(usize, usize)> = current
666 .iter()
667 .enumerate()
668 .map(|(i, &p)| {
669 let (ti, _) = target_tree.nearest_neighbor(p);
670 (i, ti)
671 })
672 .collect();
673 let error: f64 = pairs
674 .iter()
675 .map(|&(si, ti)| dist2(current[si], target.points[ti]))
676 .sum::<f64>()
677 / pairs.len() as f64;
678 if (prev_error - error).abs() < self.convergence_tol {
679 break;
680 }
681 prev_error = error;
682 let src_centroid =
683 mean_points(&pairs.iter().map(|&(si, _)| current[si]).collect::<Vec<_>>());
684 let tgt_centroid = mean_points(
685 &pairs
686 .iter()
687 .map(|&(_, ti)| target.points[ti])
688 .collect::<Vec<_>>(),
689 );
690 let mut h = [[0.0f64; 3]; 3];
691 for &(si, ti) in &pairs {
692 let s = sub(current[si], src_centroid);
693 let t = sub(target.points[ti], tgt_centroid);
694 for i in 0..3 {
695 for j in 0..3 {
696 h[i][j] += s[i] * t[j];
697 }
698 }
699 }
700 let (u, _s_vals, v) = svd_3x3_sym_jacobi(h);
701 let det = mat_det_3(&mat_mul_3(&v, &mat_transpose_3(&u)));
702 let mut d = identity_3x3();
703 if det < 0.0 {
704 d[2][2] = -1.0;
705 }
706 let r_step = mat_mul_3(&v, &mat_mul_3(&d, &mat_transpose_3(&u)));
707 let t_step = sub(tgt_centroid, mat_vec_3(&r_step, src_centroid));
708 for p in &mut current {
709 *p = add(mat_vec_3(&r_step, *p), t_step);
710 }
711 rot = mat_mul_3(&r_step, &rot);
712 trans = add(mat_vec_3(&r_step, trans), t_step);
713 }
714 let target_tree2 = KdTree3D::build(&target.points);
715 let final_error = if current.is_empty() {
716 0.0
717 } else {
718 current
719 .iter()
720 .map(|&p| {
721 let (ti, _) = target_tree2.nearest_neighbor(p);
722 dist2(p, target.points[ti])
723 })
724 .sum::<f64>()
725 / current.len() as f64
726 };
727 IcpResult {
728 rotation: rot,
729 translation: trans,
730 final_error,
731 iterations: iters,
732 }
733 }
734 pub fn closest_point_pairs(source: &PointCloud, target: &PointCloud) -> Vec<(usize, usize)> {
736 let tree = KdTree3D::build(&target.points);
737 source
738 .points
739 .iter()
740 .enumerate()
741 .map(|(i, &p)| {
742 let (ti, _) = tree.nearest_neighbor(p);
743 (i, ti)
744 })
745 .collect()
746 }
747}