spatialrust_segmentation/
plane.rs

1#[cfg(feature = "segment-ransac-plane-gpu")]
2use spatialrust_core::TransferDirection;
3use spatialrust_core::{
4    DeviceKind, ExecutionOutput, ExecutionPolicy, ExecutionReceipt, HasPositions3, PointCloud,
5    SpatialError, SpatialResult,
6};
7use spatialrust_math::Vec3;
8
9use crate::cloud::extract_mask;
10use crate::plane_ransac::{
11    collect_inliers, plane_from_indices, refine_plane_from_inliers, sample_indices, Rng,
12};
13use crate::segmenter::PointCloudSegmenter;
14
15/// Plane model in Hessian form: `normal · p + d = 0` with unit normal.
16#[derive(Clone, Copy, Debug, PartialEq)]
17pub struct PlaneModel {
18    /// Unit-length plane normal.
19    pub normal: Vec3<f32>,
20    /// Plane offset term.
21    pub d: f32,
22}
23
24impl PlaneModel {
25    /// Returns the signed distance from a point to the plane.
26    #[must_use]
27    pub fn signed_distance(&self, point: Vec3<f32>) -> f32 {
28        self.normal.dot(point) + self.d
29    }
30
31    /// Returns the absolute distance from a point to the plane.
32    #[must_use]
33    pub fn distance(&self, point: Vec3<f32>) -> f32 {
34        self.signed_distance(point).abs()
35    }
36
37    /// Returns the absolute distance from XYZ coordinates to the plane.
38    #[must_use]
39    pub fn distance_xyz(&self, x: f32, y: f32, z: f32) -> f32 {
40        (self.normal.x * x + self.normal.y * y + self.normal.z * z + self.d).abs()
41    }
42}
43
44/// Minimum point count before GPU RANSAC plane scoring is selected under `Auto`.
45///
46/// Full-cloud bench on the public PCL `table_scene_lms400` sample (460k points,
47/// 1000 iterations) showed ~11× GPU speedup. After MVP-style voxel downsampling
48/// (leaf=0.05) + normals the same scene is ~2k points and GPU remains ~2.7×
49/// faster in local release measurements.
50pub const DEFAULT_GPU_MIN_POINTS_PLANE: usize = 2_000;
51
52/// Configuration for RANSAC plane segmentation.
53#[derive(Clone, Copy, Debug, PartialEq)]
54pub struct RansacPlaneConfig {
55    /// Maximum distance from the plane for inlier classification.
56    pub distance_threshold: f32,
57    /// Maximum number of RANSAC iterations.
58    pub max_iterations: usize,
59    /// Minimum number of inliers required to accept a model.
60    pub min_inliers: usize,
61    /// Seed for deterministic sampling in tests.
62    pub seed: u64,
63    /// Minimum input point count before GPU execution is considered under `Auto`.
64    ///
65    /// `None` always uses GPU when requested.
66    pub gpu_min_points: Option<usize>,
67}
68
69impl Default for RansacPlaneConfig {
70    fn default() -> Self {
71        Self {
72            distance_threshold: 0.01,
73            max_iterations: 1_000,
74            min_inliers: 3,
75            seed: 42,
76            gpu_min_points: Some(DEFAULT_GPU_MIN_POINTS_PLANE),
77        }
78    }
79}
80
81impl RansacPlaneConfig {
82    /// Creates a config with the given distance threshold.
83    #[must_use]
84    pub const fn with_distance_threshold(distance_threshold: f32) -> Self {
85        Self {
86            distance_threshold,
87            max_iterations: 1_000,
88            min_inliers: 3,
89            seed: 42,
90            gpu_min_points: Some(DEFAULT_GPU_MIN_POINTS_PLANE),
91        }
92    }
93
94    /// Disables the GPU point-count heuristic so GPU is always used when requested.
95    #[must_use]
96    pub const fn without_gpu_min_points(mut self) -> Self {
97        self.gpu_min_points = None;
98        self
99    }
100
101    /// Returns the point-count threshold used by [`ExecutionPolicy::Auto`].
102    #[must_use]
103    pub const fn effective_gpu_min_points(&self) -> Option<usize> {
104        self.gpu_min_points
105    }
106}
107
108/// Result of RANSAC plane segmentation.
109#[derive(Clone, Debug, PartialEq)]
110pub struct RansacPlaneSegmentation {
111    /// Fitted plane model refined from inliers.
112    pub model: PlaneModel,
113    /// Points classified as inliers.
114    pub inliers: PointCloud,
115    /// Points classified as outliers.
116    pub outliers: PointCloud,
117    /// Number of inlier points.
118    pub inlier_count: usize,
119}
120
121/// RANSAC-based dominant plane segmenter.
122#[derive(Clone, Copy, Debug, PartialEq)]
123pub struct RansacPlaneSegmenter {
124    config: RansacPlaneConfig,
125}
126
127impl RansacPlaneSegmenter {
128    /// Creates a segmenter from config.
129    #[must_use]
130    pub const fn new(config: RansacPlaneConfig) -> Self {
131        Self { config }
132    }
133
134    /// Returns the segmenter config.
135    #[must_use]
136    pub const fn config(&self) -> RansacPlaneConfig {
137        self.config
138    }
139
140    /// Segments the dominant plane and returns inlier/outlier clouds.
141    pub fn segment(&self, input: &PointCloud) -> SpatialResult<RansacPlaneSegmentation> {
142        self.segment_with_policy(input, ExecutionPolicy::CpuSingle)
143    }
144
145    /// Segments the dominant plane using the given execution policy.
146    ///
147    /// With the `segment-ransac-plane-gpu` feature, [`ExecutionPolicy::Auto`]
148    /// runs hypothesis scoring on wgpu when the input meets
149    /// [`RansacPlaneConfig::effective_gpu_min_points`]. An explicit GPU policy
150    /// is strict and does not use the threshold; use `Auto` for fallback.
151    pub fn segment_with_policy(
152        &self,
153        input: &PointCloud,
154        policy: ExecutionPolicy,
155    ) -> SpatialResult<RansacPlaneSegmentation> {
156        self.segment_with_policy_and_receipt(input, policy).map(ExecutionOutput::into_output)
157    }
158
159    /// Segments a plane and returns execution/transfer accounting.
160    pub fn segment_with_policy_and_receipt(
161        &self,
162        input: &PointCloud,
163        policy: ExecutionPolicy,
164    ) -> SpatialResult<ExecutionOutput<RansacPlaneSegmentation>> {
165        policy.validate()?;
166        let resolved_policy = self.receipt_policy(input, policy)?;
167        let output = self.segment_policy_output(input, policy)?;
168        let mut receipt = ExecutionReceipt::new(policy, resolved_policy);
169        receipt.record_stage("plane-segmentation");
170        if matches!(resolved_policy, ExecutionPolicy::Gpu(DeviceKind::Wgpu)) {
171            #[cfg(feature = "segment-ransac-plane-gpu")]
172            {
173                let hypothesis_count = crate::plane_ransac::generate_hypotheses(
174                    input.len(),
175                    self.config.max_iterations,
176                    self.config.seed,
177                )
178                .len();
179                let input_bytes = (input.len() * 3 * std::mem::size_of::<f32>()) as u64;
180                let hypothesis_bytes = (hypothesis_count * 4 * std::mem::size_of::<u32>()) as u64;
181                let score_bytes = (hypothesis_count
182                    * std::mem::size_of::<spatialrust_gpu::GpuPlaneScore>())
183                    as u64;
184                receipt.record_transfer(TransferDirection::HostToDevice, input_bytes);
185                receipt.record_transfer(TransferDirection::HostToDevice, hypothesis_bytes);
186                receipt.record_transfer(TransferDirection::DeviceToHost, score_bytes);
187                receipt.record_stage("gpu-score-readback");
188            }
189        }
190        Ok(ExecutionOutput::new(output, receipt))
191    }
192
193    fn segment_policy_output(
194        &self,
195        input: &PointCloud,
196        policy: ExecutionPolicy,
197    ) -> SpatialResult<RansacPlaneSegmentation> {
198        #[cfg(feature = "segment-ransac-plane-gpu")]
199        {
200            let resolved = self.resolve_policy(input, policy)?;
201            if matches!(resolved, ExecutionPolicy::Gpu(DeviceKind::Wgpu)) {
202                return crate::plane_gpu::GpuRansacPlaneSegmenter::new(self.config).segment(input);
203            }
204        }
205
206        #[cfg(not(feature = "segment-ransac-plane-gpu"))]
207        if policy.requests_gpu() {
208            return Err(SpatialError::InvalidArgument(
209                "GPU plane segmentation requires the segment-ransac-plane-gpu feature".to_owned(),
210            ));
211        }
212
213        let _ = policy;
214        self.segment_cpu(input)
215    }
216
217    fn receipt_policy(
218        &self,
219        input: &PointCloud,
220        policy: ExecutionPolicy,
221    ) -> SpatialResult<ExecutionPolicy> {
222        #[cfg(feature = "segment-ransac-plane-gpu")]
223        {
224            self.resolve_policy(input, policy)
225        }
226        #[cfg(not(feature = "segment-ransac-plane-gpu"))]
227        {
228            let _ = input;
229            Ok(match policy {
230                ExecutionPolicy::Auto => ExecutionPolicy::CpuSingle,
231                other => other,
232            })
233        }
234    }
235
236    #[cfg(feature = "segment-ransac-plane-gpu")]
237    fn should_use_gpu(&self, input: &PointCloud) -> bool {
238        self.config.effective_gpu_min_points().map_or(true, |min_points| input.len() >= min_points)
239            && spatialrust_gpu::WgpuRuntime::shared().is_ok()
240    }
241
242    #[cfg(feature = "segment-ransac-plane-gpu")]
243    fn resolve_policy(
244        &self,
245        input: &PointCloud,
246        policy: ExecutionPolicy,
247    ) -> SpatialResult<ExecutionPolicy> {
248        match policy {
249            ExecutionPolicy::Auto => {
250                if self.should_use_gpu(input) {
251                    Ok(ExecutionPolicy::Gpu(DeviceKind::Wgpu))
252                } else {
253                    Ok(ExecutionPolicy::CpuSingle)
254                }
255            }
256            ExecutionPolicy::Gpu(DeviceKind::Cpu) => Err(SpatialError::InvalidArgument(
257                "GPU execution policy cannot target the CPU device".to_owned(),
258            )),
259            ExecutionPolicy::Gpu(DeviceKind::Cuda) => Err(SpatialError::InvalidArgument(
260                "CUDA plane segmentation is not available".to_owned(),
261            )),
262            other => Ok(other),
263        }
264    }
265
266    fn segment_cpu(&self, input: &PointCloud) -> SpatialResult<RansacPlaneSegmentation> {
267        if input.is_empty() {
268            return Err(SpatialError::InvalidArgument(
269                "cannot segment plane from empty point cloud".to_owned(),
270            ));
271        }
272
273        let (x, y, z) = input.positions3()?;
274        let len = input.len();
275        if len < 3 {
276            return Err(SpatialError::InvalidArgument(
277                "plane segmentation requires at least three points".to_owned(),
278            ));
279        }
280
281        let mut rng = Rng::new(self.config.seed);
282        let mut best_inliers = Vec::new();
283        let mut best_model = None;
284
285        for _ in 0..self.config.max_iterations {
286            let Some(sample) = sample_indices(&mut rng, len) else {
287                continue;
288            };
289            let Some(candidate) = plane_from_indices(x, y, z, sample) else {
290                continue;
291            };
292
293            let inliers = collect_inliers(x, y, z, &candidate, self.config.distance_threshold);
294            if inliers.len() > best_inliers.len() {
295                best_inliers = inliers;
296                best_model = Some(candidate);
297            }
298        }
299
300        finalize_plane_segmentation(input, x, y, z, &self.config, best_inliers, best_model)
301    }
302
303    /// Returns only the outlier cloud after removing the dominant plane.
304    pub fn extract_outliers(&self, input: &PointCloud) -> SpatialResult<PointCloud> {
305        self.segment(input).map(|result| result.outliers)
306    }
307}
308
309pub(crate) fn finalize_plane_segmentation(
310    input: &PointCloud,
311    x: &[f32],
312    y: &[f32],
313    z: &[f32],
314    config: &RansacPlaneConfig,
315    best_inliers: Vec<usize>,
316    best_model: Option<PlaneModel>,
317) -> SpatialResult<RansacPlaneSegmentation> {
318    if best_inliers.len() < config.min_inliers {
319        return Err(SpatialError::InvalidArgument(format!(
320            "RANSAC found only {} inliers, minimum is {}",
321            best_inliers.len(),
322            config.min_inliers
323        )));
324    }
325
326    let model = refine_plane_from_inliers(x, y, z, &best_inliers)
327        .or(best_model)
328        .ok_or_else(|| SpatialError::InvalidArgument("failed to refine plane model".to_owned()))?;
329
330    let len = input.len();
331    let mut inlier_mask = vec![false; len];
332    for index in &best_inliers {
333        inlier_mask[*index] = true;
334    }
335    let mut outlier_mask = inlier_mask.clone();
336    for selected in &mut outlier_mask {
337        *selected = !*selected;
338    }
339
340    let inliers = extract_mask(input, &inlier_mask)?;
341    let outliers = extract_mask(input, &outlier_mask)?;
342
343    Ok(RansacPlaneSegmentation { inlier_count: best_inliers.len(), model, inliers, outliers })
344}
345
346impl PointCloudSegmenter for RansacPlaneSegmenter {
347    fn name(&self) -> &'static str {
348        "RansacPlaneSegmenter"
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use super::{PlaneModel, RansacPlaneConfig, RansacPlaneSegmenter};
355    use spatialrust_core::{
356        DeviceKind, ExecutionPolicy, HasPositions3, PointCloudBuilder, StandardSchemas,
357    };
358    use spatialrust_math::Vec3;
359
360    fn plane_with_outliers() -> spatialrust_core::PointCloud {
361        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
362        for x in 0..10 {
363            for y in 0..10 {
364                builder.push_point([x as f32, y as f32, 0.0]).unwrap();
365            }
366        }
367        builder.push_point([0.0, 0.0, 5.0]).unwrap();
368        builder.push_point([1.0, 1.0, 5.0]).unwrap();
369        builder.build().unwrap()
370    }
371
372    #[test]
373    fn segments_dominant_plane() {
374        let input = plane_with_outliers();
375        let segmenter = RansacPlaneSegmenter::new(RansacPlaneConfig {
376            distance_threshold: 0.05,
377            max_iterations: 500,
378            min_inliers: 50,
379            seed: 7,
380            ..Default::default()
381        });
382        let result = segmenter.segment(&input).unwrap();
383        assert_eq!(result.inlier_count, 100);
384        assert_eq!(result.outliers.len(), 2);
385        assert!(result.model.normal.z.abs() > 0.9);
386    }
387
388    #[test]
389    fn plane_distance_matches_point() {
390        let model = PlaneModel { normal: Vec3::new(0.0, 0.0, 1.0), d: 0.0 };
391        assert!((model.distance(Vec3::new(0.0, 0.0, 1.0)) - 1.0).abs() < 1e-6);
392    }
393
394    #[test]
395    fn extract_outliers_removes_plane() {
396        let input = plane_with_outliers();
397        let segmenter = RansacPlaneSegmenter::new(RansacPlaneConfig {
398            distance_threshold: 0.05,
399            max_iterations: 500,
400            min_inliers: 50,
401            seed: 7,
402            ..Default::default()
403        });
404        let outliers = segmenter.extract_outliers(&input).unwrap();
405        let (_, _, z) = outliers.positions3().unwrap();
406        assert!(z.iter().all(|value| *value > 1.0));
407    }
408
409    #[test]
410    fn rejects_unsupported_explicit_gpu_policy() {
411        let segmenter = RansacPlaneSegmenter::new(RansacPlaneConfig::default());
412        let error = segmenter
413            .segment_with_policy(&plane_with_outliers(), ExecutionPolicy::Gpu(DeviceKind::Cuda))
414            .unwrap_err();
415        assert!(matches!(error, spatialrust_core::SpatialError::InvalidArgument(_)));
416    }
417}
spatialrust_segmentation/plane.rs

spatialrust_segmentation/
plane.rs
