threecrate_algorithms/
segmentation.rs

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

threecrate_algorithms/
segmentation.rs
