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 std::collections::HashSet;
8
9/// A 3D plane model defined by the equation ax + by + cz + d = 0
10#[derive(Debug, Clone, PartialEq)]
11pub struct PlaneModel {
12    /// Plane coefficients [a, b, c, d] where ax + by + cz + d = 0
13    pub coefficients: Vector4<f32>,
14}
15
16impl PlaneModel {
17    /// Create a new plane model from coefficients
18    pub fn new(a: f32, b: f32, c: f32, d: f32) -> Self {
19        Self {
20            coefficients: Vector4::new(a, b, c, d),
21        }
22    }
23
24    /// Create a plane model from three points
25    pub fn from_points(p1: &Point3f, p2: &Point3f, p3: &Point3f) -> Option<Self> {
26        // Calculate two vectors in the plane
27        let v1 = p2 - p1;
28        let v2 = p3 - p1;
29        
30        // Calculate normal vector using cross product
31        let normal = v1.cross(&v2);
32        
33        // Check if points are collinear
34        if normal.magnitude() < 1e-8 {
35            return None;
36        }
37        
38        let normal = normal.normalize();
39        
40        // Calculate d coefficient using point p1
41        let d = -normal.dot(&p1.coords);
42        
43        Some(PlaneModel::new(normal.x, normal.y, normal.z, d))
44    }
45
46    /// Get the normal vector of the plane
47    pub fn normal(&self) -> Vector3f {
48        Vector3f::new(
49            self.coefficients.x,
50            self.coefficients.y,
51            self.coefficients.z,
52        )
53    }
54
55    /// Calculate the distance from a point to the plane
56    pub fn distance_to_point(&self, point: &Point3f) -> f32 {
57        let normal = self.normal();
58        let normal_magnitude = normal.magnitude();
59        
60        if normal_magnitude < 1e-8 {
61            return f32::INFINITY;
62        }
63        
64        (self.coefficients.x * point.x + 
65         self.coefficients.y * point.y + 
66         self.coefficients.z * point.z + 
67         self.coefficients.w).abs() / normal_magnitude
68    }
69
70    /// Count inliers within a distance threshold
71    pub fn count_inliers(&self, points: &[Point3f], threshold: f32) -> usize {
72        points.iter()
73            .filter(|point| self.distance_to_point(point) <= threshold)
74            .count()
75    }
76
77    /// Get indices of inlier points within a distance threshold
78    pub fn get_inliers(&self, points: &[Point3f], threshold: f32) -> Vec<usize> {
79        points.iter()
80            .enumerate()
81            .filter(|(_, point)| self.distance_to_point(point) <= threshold)
82            .map(|(i, _)| i)
83            .collect()
84    }
85}
86
87/// RANSAC plane segmentation result
88#[derive(Debug, Clone)]
89pub struct PlaneSegmentationResult {
90    /// The best plane model found
91    pub model: PlaneModel,
92    /// Indices of inlier points
93    pub inliers: Vec<usize>,
94    /// Number of RANSAC iterations performed
95    pub iterations: usize,
96}
97
98/// Plane segmentation using RANSAC algorithm
99/// 
100/// This function finds the best plane that fits the most points in the cloud
101/// using the RANSAC (Random Sample Consensus) algorithm.
102/// 
103/// # Arguments
104/// * `cloud` - Input point cloud
105/// * `threshold` - Maximum distance for a point to be considered an inlier
106/// * `max_iters` - Maximum number of RANSAC iterations
107/// 
108/// # Returns
109/// * `Result<PlaneSegmentationResult>` - The best plane model and inlier indices
110pub fn segment_plane(
111    cloud: &PointCloud<Point3f>, 
112    threshold: f32, 
113    max_iters: usize
114) -> Result<PlaneSegmentationResult> {
115    if cloud.len() < 3 {
116        return Err(Error::InvalidData("Need at least 3 points for plane segmentation".to_string()));
117    }
118
119    if threshold <= 0.0 {
120        return Err(Error::InvalidData("Threshold must be positive".to_string()));
121    }
122
123    if max_iters == 0 {
124        return Err(Error::InvalidData("Max iterations must be positive".to_string()));
125    }
126
127    let points = &cloud.points;
128    let mut rng = thread_rng();
129    let mut best_model: Option<PlaneModel> = None;
130    let mut best_inliers = Vec::new();
131    let mut best_score = 0;
132
133    for _iteration in 0..max_iters {
134        // Randomly sample 3 points
135        let mut indices = HashSet::new();
136        while indices.len() < 3 {
137            indices.insert(rng.gen_range(0..points.len()));
138        }
139        let indices: Vec<usize> = indices.into_iter().collect();
140
141        let p1 = &points[indices[0]];
142        let p2 = &points[indices[1]];
143        let p3 = &points[indices[2]];
144
145        // Try to create a plane model from these points
146        if let Some(model) = PlaneModel::from_points(p1, p2, p3) {
147            // Count inliers
148            let inlier_count = model.count_inliers(points, threshold);
149            
150            // Update best model if this one is better
151            if inlier_count > best_score {
152                best_score = inlier_count;
153                best_inliers = model.get_inliers(points, threshold);
154                best_model = Some(model);
155            }
156        }
157    }
158
159    match best_model {
160        Some(model) => Ok(PlaneSegmentationResult {
161            model,
162            inliers: best_inliers,
163            iterations: max_iters,
164        }),
165        None => Err(Error::Algorithm("Failed to find valid plane model".to_string())),
166    }
167}
168
169/// Parallel RANSAC plane segmentation for better performance on large point clouds
170/// 
171/// This version uses parallel processing to speed up the RANSAC algorithm
172/// by running multiple iterations in parallel.
173/// 
174/// # Arguments
175/// * `cloud` - Input point cloud
176/// * `threshold` - Maximum distance for a point to be considered an inlier
177/// * `max_iters` - Maximum number of RANSAC iterations
178/// 
179/// # Returns
180/// * `Result<PlaneSegmentationResult>` - The best plane model and inlier indices
181pub fn segment_plane_parallel(
182    cloud: &PointCloud<Point3f>, 
183    threshold: f32, 
184    max_iters: usize
185) -> Result<PlaneSegmentationResult> {
186    if cloud.len() < 3 {
187        return Err(Error::InvalidData("Need at least 3 points for plane segmentation".to_string()));
188    }
189
190    if threshold <= 0.0 {
191        return Err(Error::InvalidData("Threshold must be positive".to_string()));
192    }
193
194    if max_iters == 0 {
195        return Err(Error::InvalidData("Max iterations must be positive".to_string()));
196    }
197
198    let points = &cloud.points;
199    
200    // Run RANSAC iterations in parallel
201    let results: Vec<_> = (0..max_iters)
202        .into_par_iter()
203        .filter_map(|_| {
204            let mut rng = thread_rng();
205            
206            // Randomly sample 3 points
207            let mut indices = HashSet::new();
208            while indices.len() < 3 {
209                indices.insert(rng.gen_range(0..points.len()));
210            }
211            let indices: Vec<usize> = indices.into_iter().collect();
212
213            let p1 = &points[indices[0]];
214            let p2 = &points[indices[1]];
215            let p3 = &points[indices[2]];
216
217            // Try to create a plane model from these points
218            PlaneModel::from_points(p1, p2, p3).map(|model| {
219                let inliers = model.get_inliers(points, threshold);
220                let score = inliers.len();
221                (model, inliers, score)
222            })
223        })
224        .collect();
225
226    // Find the best result
227    let best = results.into_iter()
228        .max_by_key(|(_, _, score)| *score);
229
230    match best {
231        Some((model, inliers, _)) => Ok(PlaneSegmentationResult {
232            model,
233            inliers,
234            iterations: max_iters,
235        }),
236        None => Err(Error::Algorithm("Failed to find valid plane model".to_string())),
237    }
238}
239
240/// Legacy function for backward compatibility
241#[deprecated(note = "Use segment_plane instead which returns a complete result")]
242pub fn segment_plane_legacy(cloud: &PointCloud<Point3f>, threshold: f32) -> Result<Vec<usize>> {
243    let result = segment_plane(cloud, threshold, 1000)?;
244    Ok(result.inliers)
245}
246
247/// RANSAC plane segmentation with simplified interface
248/// 
249/// This function provides a simplified interface for RANSAC plane segmentation
250/// that returns plane coefficients and inlier indices directly.
251/// 
252/// # Arguments
253/// * `cloud` - Input point cloud
254/// * `max_iters` - Maximum number of RANSAC iterations
255/// * `threshold` - Maximum distance for a point to be considered an inlier
256/// 
257/// # Returns
258/// * `Result<(Vector4<f32>, Vec<usize>)>` - Plane coefficients and inlier indices
259/// 
260/// # Example
261/// ```rust
262/// use threecrate_algorithms::segment_plane_ransac;
263/// use threecrate_core::{PointCloud, Point3f};
264/// use nalgebra::Vector4;
265/// 
266/// fn main() -> Result<(), Box<dyn std::error::Error>> {
267///     let cloud = PointCloud::from_points(vec![
268///         Point3f::new(0.0, 0.0, 0.0),
269///         Point3f::new(1.0, 0.0, 0.0),
270///         Point3f::new(0.0, 1.0, 0.0),
271///     ]);
272/// 
273///     let (coefficients, inliers) = segment_plane_ransac(&cloud, 1000, 0.01)?;
274///     println!("Plane coefficients: {:?}", coefficients);
275///     println!("Found {} inliers", inliers.len());
276///     Ok(())
277/// }
278/// ```
279pub fn segment_plane_ransac(
280    cloud: &PointCloud<Point3f>,
281    max_iters: usize,
282    threshold: f32,
283) -> Result<(Vector4<f32>, Vec<usize>)> {
284    let result = segment_plane(cloud, threshold, max_iters)?;
285    Ok((result.model.coefficients, result.inliers))
286}
287
288/// RANSAC plane segmentation (alias for segment_plane_ransac)
289/// 
290/// This function is an alias for `segment_plane_ransac` to maintain compatibility
291/// with the README documentation.
292/// 
293/// # Arguments
294/// * `cloud` - Input point cloud
295/// * `max_iters` - Maximum number of RANSAC iterations
296/// * `threshold` - Maximum distance for a point to be considered an inlier
297/// 
298/// # Returns
299/// * `Result<(Vector4<f32>, Vec<usize>)>` - Plane coefficients and inlier indices
300pub fn plane_segmentation_ransac(
301    cloud: &PointCloud<Point3f>,
302    max_iters: usize,
303    threshold: f32,
304) -> Result<(Vector4<f32>, Vec<usize>)> {
305    segment_plane_ransac(cloud, max_iters, threshold)
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use approx::assert_relative_eq;
312
313    #[test]
314    fn test_plane_model_from_points() {
315        // Create a plane in XY plane (z=0)
316        let p1 = Point3f::new(0.0, 0.0, 0.0);
317        let p2 = Point3f::new(1.0, 0.0, 0.0);
318        let p3 = Point3f::new(0.0, 1.0, 0.0);
319
320        let model = PlaneModel::from_points(&p1, &p2, &p3).unwrap();
321        
322        // Normal should be close to (0, 0, 1) or (0, 0, -1)
323        let normal = model.normal();
324        assert!(normal.z.abs() > 0.9, "Normal should be primarily in Z direction: {:?}", normal);
325        
326        // Distance to points on the plane should be ~0
327        assert!(model.distance_to_point(&p1) < 1e-6);
328        assert!(model.distance_to_point(&p2) < 1e-6);
329        assert!(model.distance_to_point(&p3) < 1e-6);
330    }
331
332    #[test]
333    fn test_plane_model_collinear_points() {
334        // Create collinear points
335        let p1 = Point3f::new(0.0, 0.0, 0.0);
336        let p2 = Point3f::new(1.0, 0.0, 0.0);
337        let p3 = Point3f::new(2.0, 0.0, 0.0);
338
339        let model = PlaneModel::from_points(&p1, &p2, &p3);
340        assert!(model.is_none(), "Should return None for collinear points");
341    }
342
343    #[test]
344    fn test_plane_distance_calculation() {
345        // Create a plane at z=1
346        let model = PlaneModel::new(0.0, 0.0, 1.0, -1.0);
347        
348        let point_on_plane = Point3f::new(0.0, 0.0, 1.0);
349        let point_above_plane = Point3f::new(0.0, 0.0, 2.0);
350        let point_below_plane = Point3f::new(0.0, 0.0, 0.0);
351        
352        assert_relative_eq!(model.distance_to_point(&point_on_plane), 0.0, epsilon = 1e-6);
353        assert_relative_eq!(model.distance_to_point(&point_above_plane), 1.0, epsilon = 1e-6);
354        assert_relative_eq!(model.distance_to_point(&point_below_plane), 1.0, epsilon = 1e-6);
355    }
356
357    #[test]
358    fn test_segment_plane_simple() {
359        // Create a point cloud with most points on a plane
360        let mut cloud = PointCloud::new();
361        
362        // Add points on XY plane (z=0)
363        for i in 0..10 {
364            for j in 0..10 {
365                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
366            }
367        }
368        
369        // Add a few outliers
370        cloud.push(Point3f::new(5.0, 5.0, 10.0));
371        cloud.push(Point3f::new(5.0, 5.0, -10.0));
372        
373        let result = segment_plane(&cloud, 0.1, 100).unwrap();
374        
375        // Should find most of the points as inliers
376        assert!(result.inliers.len() >= 95, "Should find most points as inliers");
377        
378        // Normal should be close to (0, 0, 1) or (0, 0, -1)
379        let normal = result.model.normal();
380        assert!(normal.z.abs() > 0.9, "Normal should be primarily in Z direction");
381    }
382
383    #[test]
384    fn test_segment_plane_insufficient_points() {
385        let mut cloud = PointCloud::new();
386        cloud.push(Point3f::new(0.0, 0.0, 0.0));
387        cloud.push(Point3f::new(1.0, 0.0, 0.0));
388        
389        let result = segment_plane(&cloud, 0.1, 100);
390        assert!(result.is_err(), "Should fail with insufficient points");
391    }
392
393    #[test]
394    fn test_segment_plane_invalid_threshold() {
395        let mut cloud = PointCloud::new();
396        cloud.push(Point3f::new(0.0, 0.0, 0.0));
397        cloud.push(Point3f::new(1.0, 0.0, 0.0));
398        cloud.push(Point3f::new(0.0, 1.0, 0.0));
399        
400        let result = segment_plane(&cloud, -0.1, 100);
401        assert!(result.is_err(), "Should fail with negative threshold");
402    }
403
404    #[test]
405    fn test_segment_plane_parallel() {
406        // Create a point cloud with most points on a plane
407        let mut cloud = PointCloud::new();
408        
409        // Add points on XY plane (z=0)
410        for i in 0..10 {
411            for j in 0..10 {
412                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
413            }
414        }
415        
416        let result = segment_plane_parallel(&cloud, 0.1, 100).unwrap();
417        
418        // Should find most of the points as inliers
419        assert!(result.inliers.len() >= 95, "Should find most points as inliers");
420    }
421
422    #[test]
423    fn test_segment_plane_ransac_simple() {
424        // Create a point cloud with most points on a plane
425        let mut cloud = PointCloud::new();
426        
427        // Add points on XY plane (z=0)
428        for i in 0..10 {
429            for j in 0..10 {
430                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
431            }
432        }
433        
434        // Add a few outliers
435        cloud.push(Point3f::new(5.0, 5.0, 10.0));
436        cloud.push(Point3f::new(5.0, 5.0, -10.0));
437        
438        let (coefficients, inliers) = segment_plane_ransac(&cloud, 100, 0.1).unwrap();
439        
440        // Should find most of the points as inliers
441        assert!(inliers.len() >= 95, "Should find most points as inliers");
442        
443        // Normal should be close to (0, 0, 1) or (0, 0, -1)
444        let normal = Vector3f::new(coefficients.x, coefficients.y, coefficients.z);
445        assert!(normal.z.abs() > 0.9, "Normal should be primarily in Z direction: {:?}", normal);
446    }
447
448    #[test]
449    fn test_segment_plane_ransac_noisy() {
450        // Create a point cloud with noisy planar points
451        let mut cloud = PointCloud::new();
452        let mut rng = thread_rng();
453        
454        // Add points on XY plane (z=0) with noise
455        for i in 0..20 {
456            for j in 0..20 {
457                let x = i as f32;
458                let y = j as f32;
459                let z = rng.gen_range(-0.05..0.05); // Add noise to z coordinate
460                cloud.push(Point3f::new(x, y, z));
461            }
462        }
463        
464        // Add some outliers
465        for _ in 0..20 {
466            let x = rng.gen_range(0.0..20.0);
467            let y = rng.gen_range(0.0..20.0);
468            let z = rng.gen_range(1.0..5.0); // Outliers above the plane
469            cloud.push(Point3f::new(x, y, z));
470        }
471        
472        let (coefficients, inliers) = segment_plane_ransac(&cloud, 1000, 0.1).unwrap();
473        
474        // Should find most of the planar points as inliers
475        assert!(inliers.len() >= 350, "Should find most planar points as inliers");
476        
477        // Normal should be close to (0, 0, 1) or (0, 0, -1)
478        let normal = Vector3f::new(coefficients.x, coefficients.y, coefficients.z);
479        assert!(normal.z.abs() > 0.8, "Normal should be primarily in Z direction: {:?}", normal);
480        
481        // Test that outliers are not included in inliers
482        let outlier_indices: Vec<usize> = (400..420).collect();
483        let outlier_inliers: Vec<usize> = inliers.iter()
484            .filter(|&&idx| outlier_indices.contains(&idx))
485            .cloned()
486            .collect();
487        assert!(outlier_inliers.len() <= 2, "Should not include many outliers in inliers");
488    }
489
490    #[test]
491    fn test_segment_plane_ransac_tilted_plane() {
492        // Create a tilted plane (not aligned with coordinate axes)
493        let mut cloud = PointCloud::new();
494        let mut rng = thread_rng();
495        
496        // Create a tilted plane: x + y + z = 0
497        for i in 0..15 {
498            for j in 0..15 {
499                let x = i as f32;
500                let y = j as f32;
501                let z = -(x + y); // Points on the plane x + y + z = 0
502                
503                // Add some noise
504                let noise_x = rng.gen_range(-0.02..0.02);
505                let noise_y = rng.gen_range(-0.02..0.02);
506                let noise_z = rng.gen_range(-0.02..0.02);
507                
508                cloud.push(Point3f::new(x + noise_x, y + noise_y, z + noise_z));
509            }
510        }
511        
512        // Add outliers
513        for _ in 0..30 {
514            let x = rng.gen_range(0.0..15.0);
515            let y = rng.gen_range(0.0..15.0);
516            let z = rng.gen_range(5.0..10.0); // Outliers above the plane
517            cloud.push(Point3f::new(x, y, z));
518        }
519        
520        let (coefficients, inliers) = segment_plane_ransac(&cloud, 1000, 0.1).unwrap();
521        
522        // Should find most of the planar points as inliers
523        assert!(inliers.len() >= 200, "Should find most planar points as inliers");
524        
525        // Normal should be close to (1, 1, 1) normalized
526        let normal = Vector3f::new(coefficients.x, coefficients.y, coefficients.z);
527        let expected_normal = Vector3f::new(1.0, 1.0, 1.0).normalize();
528        let dot_product = normal.dot(&expected_normal).abs();
529        assert!(dot_product > 0.8, "Normal should be close to expected direction: {:?}", normal);
530    }
531
532    #[test]
533    fn test_plane_segmentation_ransac_alias() {
534        // Test that plane_segmentation_ransac is an alias for segment_plane_ransac
535        let mut cloud = PointCloud::new();
536        
537        // Add points on XY plane (z=0)
538        for i in 0..5 {
539            for j in 0..5 {
540                cloud.push(Point3f::new(i as f32, j as f32, 0.0));
541            }
542        }
543        
544        let result1 = segment_plane_ransac(&cloud, 100, 0.1).unwrap();
545        let result2 = plane_segmentation_ransac(&cloud, 100, 0.1).unwrap();
546        
547        // Both should return valid results (RANSAC is stochastic, so exact values may differ)
548        assert!(result1.1.len() >= 20, "Should find most points as inliers");
549        assert!(result2.1.len() >= 20, "Should find most points as inliers");
550        
551        // Both should have similar inlier counts (within reasonable bounds)
552        let diff = (result1.1.len() as i32 - result2.1.len() as i32).abs();
553        assert!(diff <= 5, "Inlier counts should be similar: {} vs {}", result1.1.len(), result2.1.len());
554    }
555
556    #[test]
557    fn test_segment_plane_ransac_insufficient_points() {
558        let mut cloud = PointCloud::new();
559        cloud.push(Point3f::new(0.0, 0.0, 0.0));
560        cloud.push(Point3f::new(1.0, 0.0, 0.0));
561        
562        let result = segment_plane_ransac(&cloud, 100, 0.1);
563        assert!(result.is_err(), "Should fail with insufficient points");
564    }
565
566    #[test]
567    fn test_segment_plane_ransac_invalid_threshold() {
568        let mut cloud = PointCloud::new();
569        cloud.push(Point3f::new(0.0, 0.0, 0.0));
570        cloud.push(Point3f::new(1.0, 0.0, 0.0));
571        cloud.push(Point3f::new(0.0, 1.0, 0.0));
572        
573        let result = segment_plane_ransac(&cloud, 100, -0.1);
574        assert!(result.is_err(), "Should fail with negative threshold");
575    }
576
577    #[test]
578    fn test_segment_plane_ransac_zero_iterations() {
579        let mut cloud = PointCloud::new();
580        cloud.push(Point3f::new(0.0, 0.0, 0.0));
581        cloud.push(Point3f::new(1.0, 0.0, 0.0));
582        cloud.push(Point3f::new(0.0, 1.0, 0.0));
583        
584        let result = segment_plane_ransac(&cloud, 0, 0.1);
585        assert!(result.is_err(), "Should fail with zero iterations");
586    }
587}
threecrate_algorithms/segmentation.rs

threecrate_algorithms/
segmentation.rs
