spatialrust_segmentation/
plane_gpu.rs

1use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
2use spatialrust_gpu::{score_ransac_plane_hypotheses_gpu, WgpuRuntime};
3use spatialrust_math::Vec3;
4
5use crate::plane::{
6    finalize_plane_segmentation, PlaneModel, RansacPlaneConfig, RansacPlaneSegmentation,
7};
8use crate::plane_ransac::{collect_inliers, generate_hypotheses};
9use crate::segmenter::PointCloudSegmenter;
10
11/// GPU-accelerated RANSAC plane segmenter.
12///
13/// Hypothesis scoring (inlier counting) runs on wgpu; refinement and cloud
14/// extraction reuse the CPU helpers so results match [`crate::RansacPlaneSegmenter`].
15#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct GpuRansacPlaneSegmenter {
17    config: RansacPlaneConfig,
18}
19
20impl GpuRansacPlaneSegmenter {
21    /// Creates a GPU segmenter from config.
22    #[must_use]
23    pub const fn new(config: RansacPlaneConfig) -> Self {
24        Self { config }
25    }
26
27    /// Returns the segmenter config.
28    #[must_use]
29    pub const fn config(&self) -> RansacPlaneConfig {
30        self.config
31    }
32
33    /// Segments the dominant plane using GPU hypothesis scoring.
34    pub fn segment(&self, input: &PointCloud) -> SpatialResult<RansacPlaneSegmentation> {
35        if input.is_empty() {
36            return Err(SpatialError::InvalidArgument(
37                "cannot segment plane from empty point cloud".to_owned(),
38            ));
39        }
40
41        let (x, y, z) = input.positions3()?;
42        let len = input.len();
43        if len < 3 {
44            return Err(SpatialError::InvalidArgument(
45                "plane segmentation requires at least three points".to_owned(),
46            ));
47        }
48
49        let runtime = WgpuRuntime::shared()?;
50        let hypotheses_usize =
51            generate_hypotheses(len, self.config.max_iterations, self.config.seed);
52        let hypotheses_u32: Vec<[u32; 3]> = hypotheses_usize
53            .iter()
54            .map(|indices| [indices[0] as u32, indices[1] as u32, indices[2] as u32])
55            .collect();
56
57        let scores = score_ransac_plane_hypotheses_gpu(
58            &runtime,
59            x,
60            y,
61            z,
62            &hypotheses_u32,
63            self.config.distance_threshold,
64        )?;
65
66        let (best_index, best_count) = scores
67            .iter()
68            .enumerate()
69            .map(|(index, score)| (index, score.inlier_count as usize))
70            .max_by_key(|(_, count)| *count)
71            .unwrap_or((0, 0));
72
73        let best_score = &scores[best_index];
74        let best_model = PlaneModel {
75            normal: Vec3::new(best_score.normal[0], best_score.normal[1], best_score.normal[2]),
76            d: best_score.d,
77        };
78        let best_inliers = if best_count > 0 {
79            collect_inliers(x, y, z, &best_model, self.config.distance_threshold)
80        } else {
81            Vec::new()
82        };
83
84        finalize_plane_segmentation(input, x, y, z, &self.config, best_inliers, Some(best_model))
85    }
86}
87
88impl PointCloudSegmenter for GpuRansacPlaneSegmenter {
89    fn name(&self) -> &'static str {
90        "GpuRansacPlaneSegmenter"
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::GpuRansacPlaneSegmenter;
97    use crate::plane::{RansacPlaneConfig, RansacPlaneSegmenter};
98    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
99    use spatialrust_gpu::WgpuRuntime;
100
101    fn plane_with_outliers() -> spatialrust_core::PointCloud {
102        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
103        for x in 0..10 {
104            for y in 0..10 {
105                builder.push_point([x as f32, y as f32, 0.0]).unwrap();
106            }
107        }
108        builder.push_point([0.0, 0.0, 5.0]).unwrap();
109        builder.push_point([1.0, 1.0, 5.0]).unwrap();
110        builder.build().unwrap()
111    }
112
113    #[test]
114    fn gpu_matches_cpu_on_planar_patch() {
115        if WgpuRuntime::shared().is_err() {
116            return;
117        }
118
119        let input = plane_with_outliers();
120        let config = RansacPlaneConfig {
121            distance_threshold: 0.05,
122            max_iterations: 500,
123            min_inliers: 50,
124            seed: 7,
125            ..Default::default()
126        };
127        let cpu = RansacPlaneSegmenter::new(config).segment(&input).unwrap();
128        let gpu = GpuRansacPlaneSegmenter::new(config).segment(&input).unwrap();
129        assert_eq!(cpu.inlier_count, gpu.inlier_count);
130        assert_eq!(cpu.outliers.len(), gpu.outliers.len());
131        assert!((cpu.model.normal.z - gpu.model.normal.z).abs() < 1e-3);
132    }
133}
spatialrust_segmentation/plane_gpu.rs

spatialrust_segmentation/
plane_gpu.rs
