threecrate_algorithms/
segmentation.rs

1//! Segmentation algorithms
2
3use crate::KdTree;
4use nalgebra::Vector4;
5use rand::prelude::*;
6use rand::rng;
7use rayon::prelude::*;
8use std::collections::{HashSet, VecDeque};
9use threecrate_core::NearestNeighborSearch;
10use threecrate_core::{Error, Point3f, PointCloud, Result, Vector3f};
11
12/// A 3D plane model defined by the equation ax + by + cz + d = 0
13#[derive(Debug, Clone, PartialEq)]
14pub struct PlaneModel {
15    /// Plane coefficients [a, b, c, d] where ax + by + cz + d = 0
16    pub coefficients: Vector4<f32>,
17}
18
19impl PlaneModel {
20    /// Create a new plane model from coefficients
21    pub fn new(a: f32, b: f32, c: f32, d: f32) -> Self {
22        Self {
23            coefficients: Vector4::new(a, b, c, d),
24        }
25    }
26
27    /// Create a plane model from three points
28    pub fn from_points(p1: &Point3f, p2: &Point3f, p3: &Point3f) -> Option<Self> {
29        // Calculate two vectors in the plane
30        let v1 = p2 - p1;
31        let v2 = p3 - p1;
32
33        // Calculate normal vector using cross product
34        let normal = v1.cross(&v2);
35
36        // Check if points are collinear
37        if normal.magnitude() < 1e-8 {
38            return None;
39        }
40
41        let normal = normal.normalize();
42
43        // Calculate d coefficient using point p1
44        let d = -normal.dot(&p1.coords);
45
46        Some(PlaneModel::new(normal.x, normal.y, normal.z, d))
47    }
48
49    /// Get the normal vector of the plane
50    pub fn normal(&self) -> Vector3f {
51        Vector3f::new(
52            self.coefficients.x,
53            self.coefficients.y,
54            self.coefficients.z,
55        )
56    }
57
58    /// Calculate the distance from a point to the plane
59    pub fn distance_to_point(&self, point: &Point3f) -> f32 {
60        let normal = self.normal();
61        let normal_magnitude = normal.magnitude();
62
63        if normal_magnitude < 1e-8 {
64            return f32::INFINITY;
65        }
66
67        (self.coefficients.x * point.x
68            + self.coefficients.y * point.y
69            + self.coefficients.z * point.z
70            + self.coefficients.w)
71            .abs()
72            / normal_magnitude
73    }
74
75    /// Count inliers within a distance threshold
76    pub fn count_inliers(&self, points: &[Point3f], threshold: f32) -> usize {
77        points
78            .iter()
79            .filter(|point| self.distance_to_point(point) <= threshold)
80            .count()
81    }
82
83    /// Get indices of inlier points within a distance threshold
84    pub fn get_inliers(&self, points: &[Point3f], threshold: f32) -> Vec<usize> {
85        points
86            .iter()
87            .enumerate()
88            .filter(|(_, point)| self.distance_to_point(point) <= threshold)
89            .map(|(i, _)| i)
90            .collect()
91    }
92}
93
94/// RANSAC plane segmentation result
95#[derive(Debug, Clone)]
96pub struct PlaneSegmentationResult {
97    /// The best plane model found
98    pub model: PlaneModel,
99    /// Indices of inlier points
100    pub inliers: Vec<usize>,
101    /// Number of RANSAC iterations performed
102    pub iterations: usize,
103}
104
105/// Plane segmentation using RANSAC algorithm
106///
107/// This function finds the best plane that fits the most points in the cloud
108/// using the RANSAC (Random Sample Consensus) algorithm.
109///
110/// # Arguments
111/// * `cloud` - Input point cloud
112/// * `threshold` - Maximum distance for a point to be considered an inlier
113/// * `max_iters` - Maximum number of RANSAC iterations
114///
115/// # Returns
116/// * `Result<PlaneSegmentationResult>` - The best plane model and inlier indices
117pub fn segment_plane(
118    cloud: &PointCloud<Point3f>,
119    threshold: f32,
120    max_iters: usize,
121) -> Result<PlaneSegmentationResult> {
122    if cloud.len() < 3 {
123        return Err(Error::InvalidData(
124            "Need at least 3 points for plane segmentation".to_string(),
125        ));
126    }
127
128    if threshold <= 0.0 {
129        return Err(Error::InvalidData("Threshold must be positive".to_string()));
130    }
131
132    if max_iters == 0 {
133        return Err(Error::InvalidData(
134            "Max iterations must be positive".to_string(),
135        ));
136    }
137
138    let points = &cloud.points;
139    let mut rng = rng();
140    let mut best_model: Option<PlaneModel> = None;
141    let mut best_inliers = Vec::new();
142    let mut best_score = 0;
143
144    for _iteration in 0..max_iters {
145        // Randomly sample 3 points
146        let mut indices = HashSet::new();
147        while indices.len() < 3 {
148            indices.insert(rng.random_range(0..points.len()));
149        }
150        let indices: Vec<usize> = indices.into_iter().collect();
151
152        let p1 = &points[indices[0]];
153        let p2 = &points[indices[1]];
154        let p3 = &points[indices[2]];
155
156        // Try to create a plane model from these points
157        if let Some(model) = PlaneModel::from_points(p1, p2, p3) {
158            // Count inliers
159            let inlier_count = model.count_inliers(points, threshold);
160
161            // Update best model if this one is better
162            if inlier_count > best_score {
163                best_score = inlier_count;
164                best_inliers = model.get_inliers(points, threshold);
165                best_model = Some(model);
166            }
167        }
168    }
169
170    match best_model {
171        Some(model) => Ok(PlaneSegmentationResult {
172            model,
173            inliers: best_inliers,
174            iterations: max_iters,
175        }),
176        None => Err(Error::Algorithm(
177            "Failed to find valid plane model".to_string(),
178        )),
179    }
180}
181
182/// Parallel RANSAC plane segmentation for better performance on large point clouds
183///
184/// This version uses parallel processing to speed up the RANSAC algorithm
185/// by running multiple iterations in parallel.
186///
187/// # Arguments
188/// * `cloud` - Input point cloud
189/// * `threshold` - Maximum distance for a point to be considered an inlier
190/// * `max_iters` - Maximum number of RANSAC iterations
191///
192/// # Returns
193/// * `Result<PlaneSegmentationResult>` - The best plane model and inlier indices
194pub fn segment_plane_parallel(
195    cloud: &PointCloud<Point3f>,
196    threshold: f32,
197    max_iters: usize,
198) -> Result<PlaneSegmentationResult> {
199    if cloud.len() < 3 {
200        return Err(Error::InvalidData(
201            "Need at least 3 points for plane segmentation".to_string(),
202        ));
203    }
204
205    if threshold <= 0.0 {
206        return Err(Error::InvalidData("Threshold must be positive".to_string()));
207    }
208
209    if max_iters == 0 {
210        return Err(Error::InvalidData(
211            "Max iterations must be positive".to_string(),
212        ));
213    }
214
215    let points = &cloud.points;
216
217    // Run RANSAC iterations in parallel. Each try only counts its inliers; the
218    // list is built once for the winner (keeping a list per try used
219    // iterations x points memory: 1.5 GB for 1M points and 200 tries).
220    let best = (0..max_iters)
221        .into_par_iter()
222        .filter_map(|_| {
223            let mut rng = rng();
224
225            // Randomly sample 3 points
226            let mut indices = HashSet::new();
227            while indices.len() < 3 {
228                indices.insert(rng.random_range(0..points.len()));
229            }
230            let indices: Vec<usize> = indices.into_iter().collect();
231
232            let p1 = &points[indices[0]];
233            let p2 = &points[indices[1]];
234            let p3 = &points[indices[2]];
235
236            // Try to create a plane model from these points
237            PlaneModel::from_points(p1, p2, p3).map(|model| {
238                let score = model.count_inliers(points, threshold);
239                (model, score)
240            })
241        })
242        .max_by_key(|(_, score)| *score);
243
244    match best {
245        Some((model, _)) => Ok(PlaneSegmentationResult {
246            inliers: model.get_inliers(points, threshold),
247            model,
248            iterations: max_iters,
249        }),
250        None => Err(Error::Algorithm(
251            "Failed to find valid plane model".to_string(),
252        )),
253    }
254}
255
256/// Legacy function for backward compatibility
257#[deprecated(note = "Use segment_plane instead which returns a complete result")]
258pub fn segment_plane_legacy(cloud: &PointCloud<Point3f>, threshold: f32) -> Result<Vec<usize>> {
259    let result = segment_plane(cloud, threshold, 1000)?;
260    Ok(result.inliers)
261}
262
263/// RANSAC plane segmentation with simplified interface
264///
265/// This function provides a simplified interface for RANSAC plane segmentation
266/// that returns plane coefficients and inlier indices directly.
267///
268/// # Arguments
269/// * `cloud` - Input point cloud
270/// * `max_iters` - Maximum number of RANSAC iterations
271/// * `threshold` - Maximum distance for a point to be considered an inlier
272///
273/// # Returns
274/// * `Result<(Vector4<f32>, Vec<usize>)>` - Plane coefficients and inlier indices
275///
276/// # Example
277/// ```rust
278/// use threecrate_algorithms::segment_plane_ransac;
279/// use threecrate_core::{PointCloud, Point3f};
280/// use nalgebra::Vector4;
281///
282/// fn main() -> Result<(), Box<dyn std::error::Error>> {
283///     let cloud = PointCloud::from_points(vec![
284///         Point3f::new(0.0, 0.0, 0.0),
285///         Point3f::new(1.0, 0.0, 0.0),
286///         Point3f::new(0.0, 1.0, 0.0),
287///     ]);
288///
289///     let (coefficients, inliers) = segment_plane_ransac(&cloud, 1000, 0.01)?;
290///     println!("Plane coefficients: {:?}", coefficients);
291///     println!("Found {} inliers", inliers.len());
292///     Ok(())
293/// }
294/// ```
295pub fn segment_plane_ransac(
296    cloud: &PointCloud<Point3f>,
297    max_iters: usize,
298    threshold: f32,
299) -> Result<(Vector4<f32>, Vec<usize>)> {
300    let result = segment_plane(cloud, threshold, max_iters)?;
301    Ok((result.model.coefficients, result.inliers))
302}
303
304/// RANSAC plane segmentation (alias for segment_plane_ransac)
305///
306/// This function is an alias for `segment_plane_ransac` to maintain compatibility
307/// with the README documentation.
308///
309/// # Arguments
310/// * `cloud` - Input point cloud
311/// * `max_iters` - Maximum number of RANSAC iterations
312/// * `threshold` - Maximum distance for a point to be considered an inlier
313///
314/// # Returns
315/// * `Result<(Vector4<f32>, Vec<usize>)>` - Plane coefficients and inlier indices
316pub fn plane_segmentation_ransac(
317    cloud: &PointCloud<Point3f>,
318    max_iters: usize,
319    threshold: f32,
320) -> Result<(Vector4<f32>, Vec<usize>)> {
321    segment_plane_ransac(cloud, max_iters, threshold)
322}
323
324/// Configuration for Euclidean cluster extraction
325#[derive(Debug, Clone)]
326pub struct EuclideanClusterConfig {
327    /// Maximum distance between points to be considered part of the same cluster
328    pub tolerance: f32,
329    /// Minimum number of points for a valid cluster
330    pub min_cluster_size: usize,
331    /// Maximum number of points allowed in a cluster
332    pub max_cluster_size: usize,
333}
334
335impl Default for EuclideanClusterConfig {
336    fn default() -> Self {
337        Self {
338            tolerance: 0.02,
339            min_cluster_size: 100,
340            max_cluster_size: 25000,
341        }
342    }
343}
344
345impl EuclideanClusterConfig {
346    /// Create a new config with given parameters
347    pub fn new(tolerance: f32, min_cluster_size: usize, max_cluster_size: usize) -> Self {
348        Self {
349            tolerance,
350            min_cluster_size,
351            max_cluster_size,
352        }
353    }
354}
355
356/// Result of Euclidean cluster extraction
357#[derive(Debug, Clone)]
358pub struct ClusterExtractionResult {
359    /// Each inner Vec contains the point indices belonging to one cluster,
360    /// ordered from largest to smallest cluster.
361    pub clusters: Vec<Vec<usize>>,
362}
363
364impl ClusterExtractionResult {
365    /// Number of clusters found
366    pub fn num_clusters(&self) -> usize {
367        self.clusters.len()
368    }
369
370    /// Extract a sub-cloud for the cluster at `index`
371    pub fn get_cluster_cloud(
372        &self,
373        cloud: &PointCloud<Point3f>,
374        index: usize,
375    ) -> Option<PointCloud<Point3f>> {
376        self.clusters.get(index).map(|indices| {
377            PointCloud::from_points(indices.iter().map(|&i| cloud.points[i]).collect())
378        })
379    }
380}
381
382/// Euclidean cluster extraction using a region-growing BFS approach.
383///
384/// Builds a KD-tree over the cloud, then grows clusters by expanding into
385/// all unvisited points within `config.tolerance` of any seed point.
386/// Clusters outside the `[min_cluster_size, max_cluster_size]` range are discarded.
387///
388/// # Arguments
389/// * `cloud` - Input point cloud
390/// * `config` - Cluster extraction parameters
391///
392/// # Returns
393/// * `Result<ClusterExtractionResult>` - The extracted clusters (largest first)
394pub fn extract_euclidean_clusters(
395    cloud: &PointCloud<Point3f>,
396    config: &EuclideanClusterConfig,
397) -> Result<ClusterExtractionResult> {
398    if cloud.is_empty() {
399        return Err(Error::InvalidData("Point cloud is empty".to_string()));
400    }
401    if config.tolerance <= 0.0 {
402        return Err(Error::InvalidData("Tolerance must be positive".to_string()));
403    }
404    if config.min_cluster_size == 0 {
405        return Err(Error::InvalidData(
406            "min_cluster_size must be at least 1".to_string(),
407        ));
408    }
409    if config.min_cluster_size > config.max_cluster_size {
410        return Err(Error::InvalidData(
411            "min_cluster_size must not exceed max_cluster_size".to_string(),
412        ));
413    }
414
415    let points = &cloud.points;
416    let kdtree = KdTree::new(points)?;
417
418    let mut visited = vec![false; points.len()];
419    let mut clusters: Vec<Vec<usize>> = Vec::new();
420
421    for seed_idx in 0..points.len() {
422        if visited[seed_idx] {
423            continue;
424        }
425
426        // BFS from seed_idx
427        let mut cluster = Vec::new();
428        let mut queue = VecDeque::new();
429
430        visited[seed_idx] = true;
431        queue.push_back(seed_idx);
432
433        while let Some(current) = queue.pop_front() {
434            cluster.push(current);
435
436            let neighbors = kdtree.find_radius_neighbors(&points[current], config.tolerance);
437            for (neighbor_idx, _dist) in neighbors {
438                if !visited[neighbor_idx] {
439                    visited[neighbor_idx] = true;
440                    queue.push_back(neighbor_idx);
441                }
442            }
443        }
444
445        if cluster.len() >= config.min_cluster_size && cluster.len() <= config.max_cluster_size {
446            clusters.push(cluster);
447        }
448    }
449
450    // Sort largest cluster first
451    clusters.sort_by(|a, b| b.len().cmp(&a.len()));
452
453    Ok(ClusterExtractionResult { clusters })
454}
455
456/// Parallel Euclidean cluster extraction.
457///
458/// Uses parallel processing to speed up the radius neighbor search phase.
459/// The BFS is still sequential per cluster but neighbor lookups are parallelised
460/// by pre-computing all neighbor lists up front.
461///
462/// # Arguments
463/// * `cloud` - Input point cloud
464/// * `config` - Cluster extraction parameters
465///
466/// # Returns
467/// * `Result<ClusterExtractionResult>` - The extracted clusters (largest first)
468pub fn extract_euclidean_clusters_parallel(
469    cloud: &PointCloud<Point3f>,
470    config: &EuclideanClusterConfig,
471) -> Result<ClusterExtractionResult> {
472    if cloud.is_empty() {
473        return Err(Error::InvalidData("Point cloud is empty".to_string()));
474    }
475    if config.tolerance <= 0.0 {
476        return Err(Error::InvalidData("Tolerance must be positive".to_string()));
477    }
478    if config.min_cluster_size == 0 {
479        return Err(Error::InvalidData(
480            "min_cluster_size must be at least 1".to_string(),
481        ));
482    }
483    if config.min_cluster_size > config.max_cluster_size {
484        return Err(Error::InvalidData(
485            "min_cluster_size must not exceed max_cluster_size".to_string(),
486        ));
487    }
488
489    let points = &cloud.points;
490    let kdtree = KdTree::new(points)?;
491    let tolerance = config.tolerance;
492
493    // Pre-compute neighbor lists in parallel
494    let neighbor_lists: Vec<Vec<usize>> = (0..points.len())
495        .into_par_iter()
496        .map(|i| {
497            kdtree
498                .find_radius_neighbors(&points[i], tolerance)
499                .into_iter()
500                .map(|(idx, _)| idx)
501                .collect()
502        })
503        .collect();
504
505    let mut visited = vec![false; points.len()];
506    let mut clusters: Vec<Vec<usize>> = Vec::new();
507
508    for seed_idx in 0..points.len() {
509        if visited[seed_idx] {
510            continue;
511        }
512
513        let mut cluster = Vec::new();
514        let mut queue = VecDeque::new();
515
516        visited[seed_idx] = true;
517        queue.push_back(seed_idx);
518
519        while let Some(current) = queue.pop_front() {
520            cluster.push(current);
521            for &neighbor_idx in &neighbor_lists[current] {
522                if !visited[neighbor_idx] {
523                    visited[neighbor_idx] = true;
524                    queue.push_back(neighbor_idx);
525                }
526            }
527        }
528
529        if cluster.len() >= config.min_cluster_size && cluster.len() <= config.max_cluster_size {
530            clusters.push(cluster);
531        }
532    }
533
534    clusters.sort_by(|a, b| b.len().cmp(&a.len()));
535
536    Ok(ClusterExtractionResult { clusters })
537}
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542    use approx::assert_relative_eq;
543
544    #[test]
545    fn test_plane_model_from_points() {
546        // Create a plane in XY plane (z=0)
547        let p1 = Point3f::new(0.0, 0.0, 0.0);
548        let p2 = Point3f::new(1.0, 0.0, 0.0);
549        let p3 = Point3f::new(0.0, 1.0, 0.0);
550
551        let model = PlaneModel::from_points(&p1, &p2, &p3).unwrap();
552
553        // Normal should be close to (0, 0, 1) or (0, 0, -1)
554        let normal = model.normal();
555        assert!(
556            normal.z.abs() > 0.9,
557            "Normal should be primarily in Z direction: {:?}",
558            normal
559        );
560
561        // Distance to points on the plane should be ~0
562        assert!(model.distance_to_point(&p1) < 1e-6);
563        assert!(model.distance_to_point(&p2) < 1e-6);
564        assert!(model.distance_to_point(&p3) < 1e-6);
565    }
566
567    #[test]
568    fn test_plane_model_collinear_points() {
569        // Create collinear points
570        let p1 = Point3f::new(0.0, 0.0, 0.0);
571        let p2 = Point3f::new(1.0, 0.0, 0.0);
572        let p3 = Point3f::new(2.0, 0.0, 0.0);
573
574        let model = PlaneModel::from_points(&p1, &p2, &p3);
575        assert!(model.is_none(), "Should return None for collinear points");
576    }
577
578    #[test]
579    fn test_plane_distance_calculation() {
580        // Create a plane at z=1
581        let model = PlaneModel::new(0.0, 0.0, 1.0, -1.0);
582
583        let point_on_plane = Point3f::new(0.0, 0.0, 1.0);
584        let point_above_plane = Point3f::new(0.0, 0.0, 2.0);
585        let point_below_plane = Point3f::new(0.0, 0.0, 0.0);
586
587        assert_relative_eq!(
588            model.distance_to_point(&point_on_plane),
589            0.0,
590            epsilon = 1e-6
591        );
592        assert_relative_eq!(
593            model.distance_to_point(&point_above_plane),
594            1.0,
595            epsilon = 1e-6
596        );
597        assert_relative_eq!(
598            model.distance_to_point(&point_below_plane),
599            1.0,
600            epsilon = 1e-6
601        );
602    }
603
604    #[test]
605    fn test_segment_plane_simple() {
606        // Create a point cloud with most points on a plane
607        let mut cloud = PointCloud::new();
608
609        // Add points on XY plane (z=0)
610        for i in 0..10 {
611            for j in 0..10 {
612                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
613            }
614        }
615
616        // Add a few outliers
617        cloud.push(Point3f::new(5.0, 5.0, 10.0));
618        cloud.push(Point3f::new(5.0, 5.0, -10.0));
619
620        let result = segment_plane(&cloud, 0.1, 100).unwrap();
621
622        // Should find most of the points as inliers
623        assert!(
624            result.inliers.len() >= 95,
625            "Should find most points as inliers"
626        );
627
628        // Normal should be close to (0, 0, 1) or (0, 0, -1)
629        let normal = result.model.normal();
630        assert!(
631            normal.z.abs() > 0.9,
632            "Normal should be primarily in Z direction"
633        );
634    }
635
636    #[test]
637    fn test_segment_plane_insufficient_points() {
638        let mut cloud = PointCloud::new();
639        cloud.push(Point3f::new(0.0, 0.0, 0.0));
640        cloud.push(Point3f::new(1.0, 0.0, 0.0));
641
642        let result = segment_plane(&cloud, 0.1, 100);
643        assert!(result.is_err(), "Should fail with insufficient points");
644    }
645
646    #[test]
647    fn test_segment_plane_invalid_threshold() {
648        let mut cloud = PointCloud::new();
649        cloud.push(Point3f::new(0.0, 0.0, 0.0));
650        cloud.push(Point3f::new(1.0, 0.0, 0.0));
651        cloud.push(Point3f::new(0.0, 1.0, 0.0));
652
653        let result = segment_plane(&cloud, -0.1, 100);
654        assert!(result.is_err(), "Should fail with negative threshold");
655    }
656
657    #[test]
658    fn test_segment_plane_parallel() {
659        // Create a point cloud with most points on a plane
660        let mut cloud = PointCloud::new();
661
662        // Add points on XY plane (z=0)
663        for i in 0..10 {
664            for j in 0..10 {
665                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
666            }
667        }
668
669        let result = segment_plane_parallel(&cloud, 0.1, 100).unwrap();
670
671        // Should find most of the points as inliers
672        assert!(
673            result.inliers.len() >= 95,
674            "Should find most points as inliers"
675        );
676    }
677
678    #[test]
679    fn test_segment_plane_ransac_simple() {
680        // Create a point cloud with most points on a plane
681        let mut cloud = PointCloud::new();
682
683        // Add points on XY plane (z=0)
684        for i in 0..10 {
685            for j in 0..10 {
686                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
687            }
688        }
689
690        // Add a few outliers
691        cloud.push(Point3f::new(5.0, 5.0, 10.0));
692        cloud.push(Point3f::new(5.0, 5.0, -10.0));
693
694        let (coefficients, inliers) = segment_plane_ransac(&cloud, 100, 0.1).unwrap();
695
696        // Should find most of the points as inliers
697        assert!(inliers.len() >= 95, "Should find most points as inliers");
698
699        // Normal should be close to (0, 0, 1) or (0, 0, -1)
700        let normal = Vector3f::new(coefficients.x, coefficients.y, coefficients.z);
701        assert!(
702            normal.z.abs() > 0.9,
703            "Normal should be primarily in Z direction: {:?}",
704            normal
705        );
706    }
707
708    #[test]
709    fn test_segment_plane_ransac_noisy() {
710        // Create a point cloud with noisy planar points
711        let mut cloud = PointCloud::new();
712        let mut rng = rng();
713
714        // Add points on XY plane (z=0) with noise
715        for i in 0..20 {
716            for j in 0..20 {
717                let x = i as f32;
718                let y = j as f32;
719                let z = rng.random_range(-0.05..0.05); // Add noise to z coordinate
720                cloud.push(Point3f::new(x, y, z));
721            }
722        }
723
724        // Add some outliers
725        for _ in 0..20 {
726            let x = rng.random_range(0.0..20.0);
727            let y = rng.random_range(0.0..20.0);
728            let z = rng.random_range(1.0..5.0); // Outliers above the plane
729            cloud.push(Point3f::new(x, y, z));
730        }
731
732        let (coefficients, inliers) = segment_plane_ransac(&cloud, 1000, 0.1).unwrap();
733
734        // Should find most of the planar points as inliers
735        assert!(
736            inliers.len() >= 350,
737            "Should find most planar points as inliers"
738        );
739
740        // Normal should be close to (0, 0, 1) or (0, 0, -1)
741        let normal = Vector3f::new(coefficients.x, coefficients.y, coefficients.z);
742        assert!(
743            normal.z.abs() > 0.8,
744            "Normal should be primarily in Z direction: {:?}",
745            normal
746        );
747
748        // Test that outliers are not included in inliers
749        let outlier_indices: Vec<usize> = (400..420).collect();
750        let outlier_inliers: Vec<usize> = inliers
751            .iter()
752            .filter(|&&idx| outlier_indices.contains(&idx))
753            .cloned()
754            .collect();
755        assert!(
756            outlier_inliers.len() <= 2,
757            "Should not include many outliers in inliers"
758        );
759    }
760
761    #[test]
762    fn test_segment_plane_ransac_tilted_plane() {
763        // Create a tilted plane (not aligned with coordinate axes)
764        let mut cloud = PointCloud::new();
765        let mut rng = rng();
766
767        // Create a tilted plane: x + y + z = 0
768        for i in 0..15 {
769            for j in 0..15 {
770                let x = i as f32;
771                let y = j as f32;
772                let z = -(x + y); // Points on the plane x + y + z = 0
773
774                // Add some noise
775                let noise_x = rng.random_range(-0.02..0.02);
776                let noise_y = rng.random_range(-0.02..0.02);
777                let noise_z = rng.random_range(-0.02..0.02);
778
779                cloud.push(Point3f::new(x + noise_x, y + noise_y, z + noise_z));
780            }
781        }
782
783        // Add outliers
784        for _ in 0..30 {
785            let x = rng.random_range(0.0..15.0);
786            let y = rng.random_range(0.0..15.0);
787            let z = rng.random_range(5.0..10.0); // Outliers above the plane
788            cloud.push(Point3f::new(x, y, z));
789        }
790
791        let (coefficients, inliers) = segment_plane_ransac(&cloud, 1000, 0.1).unwrap();
792
793        // Should find most of the planar points as inliers
794        assert!(
795            inliers.len() >= 200,
796            "Should find most planar points as inliers"
797        );
798
799        // Normal should be close to (1, 1, 1) normalized
800        let normal = Vector3f::new(coefficients.x, coefficients.y, coefficients.z);
801        let expected_normal = Vector3f::new(1.0, 1.0, 1.0).normalize();
802        let dot_product = normal.dot(&expected_normal).abs();
803        assert!(
804            dot_product > 0.8,
805            "Normal should be close to expected direction: {:?}",
806            normal
807        );
808    }
809
810    #[test]
811    fn test_plane_segmentation_ransac_alias() {
812        // Test that plane_segmentation_ransac is an alias for segment_plane_ransac
813        let mut cloud = PointCloud::new();
814
815        // Add points on XY plane (z=0)
816        for i in 0..5 {
817            for j in 0..5 {
818                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
819            }
820        }
821
822        let result1 = segment_plane_ransac(&cloud, 100, 0.1).unwrap();
823        let result2 = plane_segmentation_ransac(&cloud, 100, 0.1).unwrap();
824
825        // Both should return valid results (RANSAC is stochastic, so exact values may differ)
826        assert!(result1.1.len() >= 20, "Should find most points as inliers");
827        assert!(result2.1.len() >= 20, "Should find most points as inliers");
828
829        // Both should have similar inlier counts (within reasonable bounds)
830        let diff = (result1.1.len() as i32 - result2.1.len() as i32).abs();
831        assert!(
832            diff <= 5,
833            "Inlier counts should be similar: {} vs {}",
834            result1.1.len(),
835            result2.1.len()
836        );
837    }
838
839    #[test]
840    fn test_segment_plane_ransac_insufficient_points() {
841        let mut cloud = PointCloud::new();
842        cloud.push(Point3f::new(0.0, 0.0, 0.0));
843        cloud.push(Point3f::new(1.0, 0.0, 0.0));
844
845        let result = segment_plane_ransac(&cloud, 100, 0.1);
846        assert!(result.is_err(), "Should fail with insufficient points");
847    }
848
849    #[test]
850    fn test_segment_plane_ransac_invalid_threshold() {
851        let mut cloud = PointCloud::new();
852        cloud.push(Point3f::new(0.0, 0.0, 0.0));
853        cloud.push(Point3f::new(1.0, 0.0, 0.0));
854        cloud.push(Point3f::new(0.0, 1.0, 0.0));
855
856        let result = segment_plane_ransac(&cloud, 100, -0.1);
857        assert!(result.is_err(), "Should fail with negative threshold");
858    }
859
860    #[test]
861    fn test_segment_plane_ransac_zero_iterations() {
862        let mut cloud = PointCloud::new();
863        cloud.push(Point3f::new(0.0, 0.0, 0.0));
864        cloud.push(Point3f::new(1.0, 0.0, 0.0));
865        cloud.push(Point3f::new(0.0, 1.0, 0.0));
866
867        let result = segment_plane_ransac(&cloud, 0, 0.1);
868        assert!(result.is_err(), "Should fail with zero iterations");
869    }
870
871    // ---- Euclidean cluster extraction tests ----
872
873    fn make_sphere_cloud(center: Point3f, radius: f32, count: usize) -> Vec<Point3f> {
874        let mut rng = rng();
875        let mut pts = Vec::with_capacity(count);
876        while pts.len() < count {
877            let x: f32 = rng.random_range(-radius..radius);
878            let y: f32 = rng.random_range(-radius..radius);
879            let z: f32 = rng.random_range(-radius..radius);
880            if x * x + y * y + z * z <= radius * radius {
881                pts.push(Point3f::new(center.x + x, center.y + y, center.z + z));
882            }
883        }
884        pts
885    }
886
887    #[test]
888    fn test_euclidean_cluster_two_blobs() {
889        let mut cloud = PointCloud::new();
890        // Blob 1 near origin
891        for p in make_sphere_cloud(Point3f::new(0.0, 0.0, 0.0), 0.3, 200) {
892            cloud.push(p);
893        }
894        // Blob 2 far away
895        for p in make_sphere_cloud(Point3f::new(10.0, 0.0, 0.0), 0.3, 150) {
896            cloud.push(p);
897        }
898
899        let config = EuclideanClusterConfig::new(0.5, 50, 10000);
900        let result = extract_euclidean_clusters(&cloud, &config).unwrap();
901
902        assert_eq!(result.num_clusters(), 2, "Should detect exactly 2 clusters");
903        // Largest cluster first
904        assert!(result.clusters[0].len() >= result.clusters[1].len());
905    }
906
907    #[test]
908    fn test_euclidean_cluster_three_blobs() {
909        let mut cloud = PointCloud::new();
910        for p in make_sphere_cloud(Point3f::new(0.0, 0.0, 0.0), 0.4, 300) {
911            cloud.push(p);
912        }
913        for p in make_sphere_cloud(Point3f::new(5.0, 0.0, 0.0), 0.4, 200) {
914            cloud.push(p);
915        }
916        for p in make_sphere_cloud(Point3f::new(0.0, 5.0, 0.0), 0.4, 100) {
917            cloud.push(p);
918        }
919
920        let config = EuclideanClusterConfig::new(0.6, 50, 10000);
921        let result = extract_euclidean_clusters(&cloud, &config).unwrap();
922
923        assert_eq!(result.num_clusters(), 3);
924    }
925
926    #[test]
927    fn test_euclidean_cluster_min_size_filter() {
928        let mut cloud = PointCloud::new();
929        // Large cluster
930        for p in make_sphere_cloud(Point3f::new(0.0, 0.0, 0.0), 0.4, 300) {
931            cloud.push(p);
932        }
933        // Small cluster (should be filtered out)
934        for p in make_sphere_cloud(Point3f::new(10.0, 0.0, 0.0), 0.2, 5) {
935            cloud.push(p);
936        }
937
938        let config = EuclideanClusterConfig::new(0.5, 50, 10000);
939        let result = extract_euclidean_clusters(&cloud, &config).unwrap();
940
941        assert_eq!(result.num_clusters(), 1, "Small cluster should be filtered");
942    }
943
944    #[test]
945    fn test_euclidean_cluster_max_size_filter() {
946        let mut cloud = PointCloud::new();
947        for p in make_sphere_cloud(Point3f::new(0.0, 0.0, 0.0), 0.5, 500) {
948            cloud.push(p);
949        }
950
951        // max_cluster_size smaller than the blob
952        let config = EuclideanClusterConfig::new(0.6, 1, 100);
953        let result = extract_euclidean_clusters(&cloud, &config).unwrap();
954
955        // The single large blob exceeds max_cluster_size, so no clusters returned
956        assert_eq!(
957            result.num_clusters(),
958            0,
959            "Oversized cluster should be filtered"
960        );
961    }
962
963    #[test]
964    fn test_euclidean_cluster_get_cloud() {
965        let mut cloud = PointCloud::new();
966        for p in make_sphere_cloud(Point3f::new(0.0, 0.0, 0.0), 0.3, 200) {
967            cloud.push(p);
968        }
969        for p in make_sphere_cloud(Point3f::new(10.0, 0.0, 0.0), 0.3, 100) {
970            cloud.push(p);
971        }
972
973        let config = EuclideanClusterConfig::new(0.5, 50, 10000);
974        let result = extract_euclidean_clusters(&cloud, &config).unwrap();
975
976        let sub = result.get_cluster_cloud(&cloud, 0).unwrap();
977        assert_eq!(sub.len(), result.clusters[0].len());
978    }
979
980    #[test]
981    fn test_euclidean_cluster_parallel_matches_sequential() {
982        let mut cloud = PointCloud::new();
983        for p in make_sphere_cloud(Point3f::new(0.0, 0.0, 0.0), 0.3, 200) {
984            cloud.push(p);
985        }
986        for p in make_sphere_cloud(Point3f::new(5.0, 0.0, 0.0), 0.3, 150) {
987            cloud.push(p);
988        }
989
990        let config = EuclideanClusterConfig::new(0.5, 50, 10000);
991        let seq = extract_euclidean_clusters(&cloud, &config).unwrap();
992        let par = extract_euclidean_clusters_parallel(&cloud, &config).unwrap();
993
994        assert_eq!(seq.num_clusters(), par.num_clusters());
995        for (s, p) in seq.clusters.iter().zip(par.clusters.iter()) {
996            assert_eq!(s.len(), p.len());
997        }
998    }
999
1000    #[test]
1001    fn test_euclidean_cluster_invalid_config() {
1002        let mut cloud = PointCloud::new();
1003        cloud.push(Point3f::new(0.0, 0.0, 0.0));
1004
1005        assert!(
1006            extract_euclidean_clusters(&cloud, &EuclideanClusterConfig::new(-1.0, 1, 100)).is_err()
1007        );
1008        assert!(
1009            extract_euclidean_clusters(&cloud, &EuclideanClusterConfig::new(0.1, 0, 100)).is_err()
1010        );
1011        assert!(
1012            extract_euclidean_clusters(&cloud, &EuclideanClusterConfig::new(0.1, 10, 5)).is_err()
1013        );
1014    }
1015
1016    #[test]
1017    fn test_euclidean_cluster_empty_cloud() {
1018        let cloud: PointCloud<Point3f> = PointCloud::new();
1019        let config = EuclideanClusterConfig::default();
1020        assert!(extract_euclidean_clusters(&cloud, &config).is_err());
1021    }
1022}
threecrate_algorithms/segmentation.rs

threecrate_algorithms/
segmentation.rs
