spatialrust_segmentation/
cluster_gpu.rs

1use spatialrust_core::{HasPositions3, PointCloud, SpatialResult, TransferStats};
2use spatialrust_gpu::{euclidean_cluster_roots_gpu_with_receipt, WgpuRuntime};
3
4use crate::cluster::{finalize_euclidean_clusters, EuclideanClusterConfig, EuclideanClusterResult};
5use crate::segmenter::PointCloudSegmenter;
6
7/// GPU-grid-backed Euclidean cluster extractor.
8///
9/// Sparse grid key generation, sorting, and compaction run on wgpu. Connected
10/// component labeling and cluster-size filtering remain deterministic host-side
11/// stages, matching [`crate::EuclideanClusterExtractor`] semantics.
12#[derive(Clone, Copy, Debug, PartialEq)]
13pub struct GpuEuclideanClusterExtractor {
14    config: EuclideanClusterConfig,
15}
16
17impl GpuEuclideanClusterExtractor {
18    /// Creates a GPU extractor from config.
19    #[must_use]
20    pub const fn new(config: EuclideanClusterConfig) -> Self {
21        Self { config }
22    }
23
24    /// Returns the extractor config.
25    #[must_use]
26    pub const fn config(&self) -> EuclideanClusterConfig {
27        self.config
28    }
29
30    /// Clusters the input cloud using the grid union-find backend.
31    pub fn extract(&self, input: &PointCloud) -> SpatialResult<EuclideanClusterResult> {
32        self.extract_with_receipt(input).map(|(result, _)| result)
33    }
34
35    /// Clusters the input cloud and returns GPU grid transfer accounting.
36    pub fn extract_with_receipt(
37        &self,
38        input: &PointCloud,
39    ) -> SpatialResult<(EuclideanClusterResult, TransferStats)> {
40        if input.is_empty() {
41            return Ok((
42                EuclideanClusterResult {
43                    cloud: input.clone(),
44                    cluster_count: 0,
45                    cluster_sizes: Vec::new(),
46                },
47                TransferStats::default(),
48            ));
49        }
50
51        let (x, y, z) = input.positions3()?;
52        let runtime = WgpuRuntime::shared()?;
53        let (roots, transfers) = euclidean_cluster_roots_gpu_with_receipt(
54            &runtime,
55            x,
56            y,
57            z,
58            self.config.cluster_tolerance,
59        )?;
60        Ok((finalize_euclidean_clusters(input, &roots, self.config)?, transfers))
61    }
62}
63
64impl PointCloudSegmenter for GpuEuclideanClusterExtractor {
65    fn name(&self) -> &'static str {
66        "GpuEuclideanClusterExtractor"
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::GpuEuclideanClusterExtractor;
73    use crate::cluster::EuclideanClusterConfig;
74    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
75
76    fn three_clusters() -> spatialrust_core::PointCloud {
77        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
78        for center in [(0.0, 0.0, 0.0), (10.0, 0.0, 0.0), (0.0, 10.0, 0.0)] {
79            for dx in 0..3 {
80                for dy in 0..3 {
81                    builder
82                        .push_point([center.0 + dx as f32, center.1 + dy as f32, center.2])
83                        .unwrap();
84                }
85            }
86        }
87        builder.build().unwrap()
88    }
89
90    #[test]
91    fn gpu_finds_three_separated_clusters() {
92        let input = three_clusters();
93        let extractor = GpuEuclideanClusterExtractor::new(EuclideanClusterConfig {
94            cluster_tolerance: 1.5,
95            min_cluster_size: 3,
96            max_cluster_size: usize::MAX,
97            gpu_min_points: None,
98        });
99        let result = extractor.extract(&input).unwrap();
100        assert_eq!(result.cluster_count, 3);
101        assert!(result.cluster_sizes.iter().all(|&size| size == 9));
102    }
103
104    fn long_chain(len: usize, spacing: f32) -> spatialrust_core::PointCloud {
105        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
106        for index in 0..len {
107            builder.push_point([index as f32 * spacing, 0.0, 0.0]).unwrap();
108        }
109        builder.build().unwrap()
110    }
111
112    #[test]
113    fn gpu_matches_cpu_on_long_chain() {
114        use crate::cluster::EuclideanClusterExtractor;
115
116        let input = long_chain(300, 1.0);
117        let config = EuclideanClusterConfig {
118            cluster_tolerance: 1.5,
119            min_cluster_size: 1,
120            max_cluster_size: usize::MAX,
121            gpu_min_points: None,
122        };
123        let cpu = EuclideanClusterExtractor::new(config).extract(&input).unwrap();
124        let gpu = GpuEuclideanClusterExtractor::new(config).extract(&input).unwrap();
125        assert_eq!(cpu.cluster_count, 1);
126        assert_eq!(gpu.cluster_count, cpu.cluster_count);
127        assert_eq!(gpu.cluster_sizes, cpu.cluster_sizes);
128    }
129
130    #[test]
131    fn gpu_matches_cpu_on_three_clusters() {
132        use crate::cluster::EuclideanClusterExtractor;
133
134        let input = three_clusters();
135        let config = EuclideanClusterConfig {
136            cluster_tolerance: 1.5,
137            min_cluster_size: 3,
138            max_cluster_size: usize::MAX,
139            gpu_min_points: None,
140        };
141        let cpu = EuclideanClusterExtractor::new(config).extract(&input).unwrap();
142        let gpu = GpuEuclideanClusterExtractor::new(config).extract(&input).unwrap();
143        assert_eq!(gpu.cluster_count, cpu.cluster_count);
144        assert_eq!(gpu.cluster_sizes, cpu.cluster_sizes);
145    }
146}
spatialrust_segmentation/cluster_gpu.rs

spatialrust_segmentation/
cluster_gpu.rs
