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                ..Default::default()
109            };
110            let Ok(result) = RansacPlaneSegmenter::new(config).segment(&sub) else {
111                // No plane with enough inliers remains.
112                break;
113            };
114
115            // Re-classify the *remaining original* points against the fitted
116            // model so labels map back to the input cloud.
117            let model = result.model;
118            let mut next_remaining = Vec::with_capacity(remaining.len());
119            let mut assigned = 0_usize;
120            for &orig in &remaining {
121                if model.distance_xyz(x[orig], y[orig], z[orig]) <= self.config.distance_threshold {
122                    labels[orig] = plane_index as i32;
123                    assigned += 1;
124                } else {
125                    next_remaining.push(orig);
126                }
127            }
128
129            if assigned < self.config.min_inliers {
130                // The model did not actually cover enough original points; undo.
131                for &orig in &remaining {
132                    if labels[orig] == plane_index as i32 {
133                        labels[orig] = -1;
134                    }
135                }
136                break;
137            }
138
139            planes.push(model);
140            plane_sizes.push(assigned);
141            remaining = next_remaining;
142        }
143
144        Ok(MultiPlaneSegmentation {
145            labeled: with_labels(input, "label", labels)?,
146            planes,
147            plane_sizes,
148        })
149    }
150}
151
152impl PointCloudSegmenter for MultiPlaneSegmenter {
153    fn name(&self) -> &'static str {
154        "MultiPlaneSegmenter"
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::{MultiPlaneConfig, MultiPlaneSegmenter};
161    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
162
163    /// A floor (z=0), a wall (y=0), and a ceiling (z=2): three planes.
164    fn room() -> spatialrust_core::PointCloud {
165        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
166        for i in 0..20 {
167            for j in 0..20 {
168                let (a, b) = (i as f32 * 0.1, j as f32 * 0.1);
169                builder.push_point([a, b, 0.0]).unwrap(); // floor
170                builder.push_point([a, 0.0, b]).unwrap(); // wall
171                builder.push_point([a, b, 2.0]).unwrap(); // ceiling
172            }
173        }
174        builder.build().unwrap()
175    }
176
177    #[test]
178    fn extracts_three_room_planes() {
179        let cloud = room();
180        let seg = MultiPlaneSegmenter::new(MultiPlaneConfig {
181            max_planes: 4,
182            distance_threshold: 0.02,
183            min_inliers: 100,
184            max_iterations: 500,
185            seed: 7,
186        })
187        .segment(&cloud)
188        .unwrap();
189
190        assert_eq!(seg.planes.len(), 3, "expected floor, wall, ceiling");
191        assert!(seg.labeled.field("label").is_ok());
192        // The three planes together cover almost every point.
193        let covered: usize = seg.plane_sizes.iter().sum();
194        assert!(covered as f32 > cloud.len() as f32 * 0.95);
195    }
196
197    #[test]
198    fn stops_when_no_dominant_plane_remains() {
199        // Pure noise has no plane with enough inliers.
200        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
201        let mut seed = 99_u64;
202        let mut rng = || {
203            seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
204            (seed >> 40) as f32 / (1u64 << 24) as f32
205        };
206        for _ in 0..400 {
207            builder.push_point([rng() * 5.0, rng() * 5.0, rng() * 5.0]).unwrap();
208        }
209        let cloud = builder.build().unwrap();
210        let seg = MultiPlaneSegmenter::new(MultiPlaneConfig::new(3, 0.01)).segment(&cloud).unwrap();
211        assert!(seg.planes.is_empty(), "noise should yield no planes");
212    }
213
214    #[test]
215    fn rejects_bad_threshold() {
216        let cloud = room();
217        let config = MultiPlaneConfig { distance_threshold: -1.0, ..MultiPlaneConfig::default() };
218        assert!(MultiPlaneSegmenter::new(config).segment(&cloud).is_err());
219    }
220}
spatialrust_segmentation/multi_plane.rs

spatialrust_segmentation/
multi_plane.rs
