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

threecrate_algorithms/
segmentation.rs
