spatialrust_segmentation/
multi_plane.rs

1//! Sequential multi-plane extraction (floor, walls, ceiling, …).
2//!
3//! Indoor and structured scenes contain several dominant planes, not one. This
4//! repeatedly fits the dominant plane with RANSAC, labels and removes its
5//! inliers, and continues on the remainder — the standard way to decompose a
6//! room into its planar surfaces.
7
8use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
9
10use crate::cloud::{extract_indices, with_labels};
11use crate::plane::{PlaneModel, RansacPlaneConfig, RansacPlaneSegmenter};
12use crate::segmenter::PointCloudSegmenter;
13
14/// Configuration for [`MultiPlaneSegmenter`].
15#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct MultiPlaneConfig {
17    /// Maximum number of planes to extract.
18    pub max_planes: usize,
19    /// Maximum distance from a plane for inlier classification.
20    pub distance_threshold: f32,
21    /// Minimum inliers to accept a plane; extraction stops below this.
22    pub min_inliers: usize,
23    /// RANSAC iterations per plane.
24    pub max_iterations: usize,
25    /// Seed for deterministic sampling.
26    pub seed: u64,
27}
28
29impl Default for MultiPlaneConfig {
30    fn default() -> Self {
31        Self {
32            max_planes: 4,
33            distance_threshold: 0.02,
34            min_inliers: 100,
35            max_iterations: 1_000,
36            seed: 42,
37        }
38    }
39}
40
41impl MultiPlaneConfig {
42    /// Creates a config from the plane count and distance threshold.
43    #[must_use]
44    pub const fn new(max_planes: usize, distance_threshold: f32) -> Self {
45        Self { max_planes, distance_threshold, min_inliers: 100, max_iterations: 1_000, seed: 42 }
46    }
47}
48
49/// Result of multi-plane segmentation.
50#[derive(Clone, Debug, PartialEq)]
51pub struct MultiPlaneSegmentation {
52    /// Fitted plane models, in extraction order (most dominant first).
53    pub planes: Vec<PlaneModel>,
54    /// Input cloud with a `label` field: plane index `0..planes.len()`, or `-1`
55    /// for points not assigned to any plane.
56    pub labeled: PointCloud,
57    /// Number of points assigned to each plane, in plane order.
58    pub plane_sizes: Vec<usize>,
59}
60
61/// Sequential RANSAC multi-plane segmenter.
62#[derive(Clone, Copy, Debug, PartialEq)]
63pub struct MultiPlaneSegmenter {
64    config: MultiPlaneConfig,
65}
66
67impl MultiPlaneSegmenter {
68    /// Creates a segmenter from config.
69    #[must_use]
70    pub const fn new(config: MultiPlaneConfig) -> Self {
71        Self { config }
72    }
73
74    /// Returns the segmenter config.
75    #[must_use]
76    pub const fn config(&self) -> MultiPlaneConfig {
77        self.config
78    }
79
80    /// Extracts up to `max_planes` dominant planes and labels each point.
81    pub fn segment(&self, input: &PointCloud) -> SpatialResult<MultiPlaneSegmentation> {
82        if self.config.distance_threshold < 0.0 {
83            return Err(SpatialError::InvalidArgument(
84                "distance_threshold must be non-negative".to_owned(),
85            ));
86        }
87        let len = input.len();
88        let (x, y, z) = input.positions3()?;
89
90        let mut labels = vec![-1_i32; len];
91        let mut remaining: Vec<usize> = (0..len).collect();
92        let mut planes = Vec::new();
93        let mut plane_sizes = Vec::new();
94
95        for plane_index in 0..self.config.max_planes {
96            if remaining.len() < 3 || remaining.len() < self.config.min_inliers {
97                break;
98            }
99
100            // Fit the dominant plane of whatever points are left.
101            let sub = extract_indices(input, &remaining)?;
102            let config = RansacPlaneConfig {
103                distance_threshold: self.config.distance_threshold,
104                max_iterations: self.config.max_iterations,
105                min_inliers: self.config.min_inliers,
106                // Vary the seed per plane so successive fits explore differently.
107                seed: self.config.seed.wrapping_add(plane_index as u64),
108            };
109            let Ok(result) = RansacPlaneSegmenter::new(config).segment(&sub) else {
110                // No plane with enough inliers remains.
111                break;
112            };
113
114            // Re-classify the *remaining original* points against the fitted
115            // model so labels map back to the input cloud.
116            let model = result.model;
117            let mut next_remaining = Vec::with_capacity(remaining.len());
118            let mut assigned = 0_usize;
119            for &orig in &remaining {
120                if model.distance_xyz(x[orig], y[orig], z[orig]) <= self.config.distance_threshold {
121                    labels[orig] = plane_index as i32;
122                    assigned += 1;
123                } else {
124                    next_remaining.push(orig);
125                }
126            }
127
128            if assigned < self.config.min_inliers {
129                // The model did not actually cover enough original points; undo.
130                for &orig in &remaining {
131                    if labels[orig] == plane_index as i32 {
132                        labels[orig] = -1;
133                    }
134                }
135                break;
136            }
137
138            planes.push(model);
139            plane_sizes.push(assigned);
140            remaining = next_remaining;
141        }
142
143        Ok(MultiPlaneSegmentation {
144            labeled: with_labels(input, "label", labels)?,
145            planes,
146            plane_sizes,
147        })
148    }
149}
150
151impl PointCloudSegmenter for MultiPlaneSegmenter {
152    fn name(&self) -> &'static str {
153        "MultiPlaneSegmenter"
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::{MultiPlaneConfig, MultiPlaneSegmenter};
160    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
161
162    /// A floor (z=0), a wall (y=0), and a ceiling (z=2): three planes.
163    fn room() -> spatialrust_core::PointCloud {
164        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
165        for i in 0..20 {
166            for j in 0..20 {
167                let (a, b) = (i as f32 * 0.1, j as f32 * 0.1);
168                builder.push_point([a, b, 0.0]).unwrap(); // floor
169                builder.push_point([a, 0.0, b]).unwrap(); // wall
170                builder.push_point([a, b, 2.0]).unwrap(); // ceiling
171            }
172        }
173        builder.build().unwrap()
174    }
175
176    #[test]
177    fn extracts_three_room_planes() {
178        let cloud = room();
179        let seg = MultiPlaneSegmenter::new(MultiPlaneConfig {
180            max_planes: 4,
181            distance_threshold: 0.02,
182            min_inliers: 100,
183            max_iterations: 500,
184            seed: 7,
185        })
186        .segment(&cloud)
187        .unwrap();
188
189        assert_eq!(seg.planes.len(), 3, "expected floor, wall, ceiling");
190        assert!(seg.labeled.field("label").is_ok());
191        // The three planes together cover almost every point.
192        let covered: usize = seg.plane_sizes.iter().sum();
193        assert!(covered as f32 > cloud.len() as f32 * 0.95);
194    }
195
196    #[test]
197    fn stops_when_no_dominant_plane_remains() {
198        // Pure noise has no plane with enough inliers.
199        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
200        let mut seed = 99_u64;
201        let mut rng = || {
202            seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
203            (seed >> 40) as f32 / (1u64 << 24) as f32
204        };
205        for _ in 0..400 {
206            builder.push_point([rng() * 5.0, rng() * 5.0, rng() * 5.0]).unwrap();
207        }
208        let cloud = builder.build().unwrap();
209        let seg = MultiPlaneSegmenter::new(MultiPlaneConfig::new(3, 0.01)).segment(&cloud).unwrap();
210        assert!(seg.planes.is_empty(), "noise should yield no planes");
211    }
212
213    #[test]
214    fn rejects_bad_threshold() {
215        let cloud = room();
216        let config = MultiPlaneConfig { distance_threshold: -1.0, ..MultiPlaneConfig::default() };
217        assert!(MultiPlaneSegmenter::new(config).segment(&cloud).is_err());
218    }
219}
spatialrust_segmentation/multi_plane.rs

spatialrust_segmentation/
multi_plane.rs
