spatialrust_segmentation/
cluster.rs

1use std::collections::{HashMap, VecDeque};
2
3#[cfg(feature = "segment-euclidean-gpu")]
4use spatialrust_core::DeviceKind;
5use spatialrust_core::{
6    ExecutionOutput, ExecutionPolicy, ExecutionReceipt, HasPositions3, PointCloud, SpatialError,
7    SpatialResult, TransferDirection,
8};
9use spatialrust_search::{KdTree, RadiusSearchIndex};
10
11use crate::cloud::with_labels;
12use crate::segmenter::PointCloudSegmenter;
13
14/// Minimum point count before GPU Euclidean clustering uses KD-tree BFS instead of grid UF.
15pub const DEFAULT_GPU_KDTREE_MIN_POINTS: usize = 50_000;
16
17/// Minimum point count before GPU Euclidean clustering is selected under `Auto`.
18pub const DEFAULT_GPU_MIN_POINTS_EUCLIDEAN: usize = 2_000;
19
20/// Configuration for Euclidean clustering.
21#[derive(Clone, Copy, Debug, PartialEq)]
22pub struct EuclideanClusterConfig {
23    /// Maximum distance between points in the same cluster.
24    pub cluster_tolerance: f32,
25    /// Minimum number of points required to form a cluster.
26    pub min_cluster_size: usize,
27    /// Maximum number of points allowed in a cluster.
28    pub max_cluster_size: usize,
29    /// Minimum input point count before GPU execution is considered under `Auto`.
30    ///
31    /// `None` always uses GPU when requested.
32    pub gpu_min_points: Option<usize>,
33}
34
35impl Default for EuclideanClusterConfig {
36    fn default() -> Self {
37        Self {
38            cluster_tolerance: 0.02,
39            min_cluster_size: 1,
40            max_cluster_size: usize::MAX,
41            gpu_min_points: Some(DEFAULT_GPU_MIN_POINTS_EUCLIDEAN),
42        }
43    }
44}
45
46impl EuclideanClusterConfig {
47    /// Creates a config with the given tolerance and minimum cluster size.
48    #[must_use]
49    pub const fn with_tolerance(cluster_tolerance: f32, min_cluster_size: usize) -> Self {
50        Self {
51            cluster_tolerance,
52            min_cluster_size,
53            max_cluster_size: usize::MAX,
54            gpu_min_points: Some(DEFAULT_GPU_MIN_POINTS_EUCLIDEAN),
55        }
56    }
57
58    /// Disables the GPU point-count heuristic so GPU is always used when requested.
59    #[must_use]
60    pub const fn without_gpu_min_points(mut self) -> Self {
61        self.gpu_min_points = None;
62        self
63    }
64
65    /// Returns the point-count threshold used by [`ExecutionPolicy::Auto`].
66    #[must_use]
67    pub const fn effective_gpu_min_points(&self) -> Option<usize> {
68        self.gpu_min_points
69    }
70}
71
72/// Result of Euclidean clustering.
73#[derive(Clone, Debug, PartialEq)]
74pub struct EuclideanClusterResult {
75    /// Input points annotated with cluster labels.
76    pub cloud: PointCloud,
77    /// Number of valid clusters found.
78    pub cluster_count: usize,
79    /// Size of each cluster in label order.
80    pub cluster_sizes: Vec<usize>,
81}
82
83/// Euclidean region-growing cluster extractor.
84#[derive(Clone, Copy, Debug, PartialEq)]
85pub struct EuclideanClusterExtractor {
86    config: EuclideanClusterConfig,
87}
88
89impl EuclideanClusterExtractor {
90    /// Creates an extractor from config.
91    #[must_use]
92    pub const fn new(config: EuclideanClusterConfig) -> Self {
93        Self { config }
94    }
95
96    /// Returns the extractor config.
97    #[must_use]
98    pub const fn config(&self) -> EuclideanClusterConfig {
99        self.config
100    }
101
102    /// Clusters the input cloud and adds a `label` field.
103    pub fn extract(&self, input: &PointCloud) -> SpatialResult<EuclideanClusterResult> {
104        validate_cluster_config(self.config)?;
105        if input.is_empty() {
106            return Ok(EuclideanClusterResult {
107                cloud: input.clone(),
108                cluster_count: 0,
109                cluster_sizes: Vec::new(),
110            });
111        }
112        let roots = extract_cpu_roots(input, self.config)?;
113        finalize_euclidean_clusters(input, &roots, self.config)
114    }
115
116    /// Clusters the input cloud using the given execution policy.
117    ///
118    /// With the `segment-euclidean-gpu` feature, [`ExecutionPolicy::Auto`]
119    /// runs grid union-find clustering when the input meets
120    /// [`EuclideanClusterConfig::effective_gpu_min_points`] and the uniform
121    /// grid fits within the cell cap. An explicit GPU policy is strict and
122    /// returns an error when the grid cannot fit; use `Auto` for fallback.
123    pub fn extract_with_policy(
124        &self,
125        input: &PointCloud,
126        policy: ExecutionPolicy,
127    ) -> SpatialResult<EuclideanClusterResult> {
128        self.extract_with_policy_and_receipt(input, policy).map(ExecutionOutput::into_output)
129    }
130
131    /// Clusters the input cloud and returns execution/transfer accounting.
132    pub fn extract_with_policy_and_receipt(
133        &self,
134        input: &PointCloud,
135        policy: ExecutionPolicy,
136    ) -> SpatialResult<ExecutionOutput<EuclideanClusterResult>> {
137        policy.validate()?;
138        let (output, resolved_policy, transfers) = self.extract_policy_output(input, policy)?;
139        let mut receipt = ExecutionReceipt::new(policy, resolved_policy);
140        receipt.record_stage("euclidean-clustering");
141        if let Some(transfers) = transfers {
142            receipt
143                .record_transfer(TransferDirection::HostToDevice, transfers.host_to_device_bytes());
144            receipt.record_transfer(
145                TransferDirection::DeviceToDevice,
146                transfers.device_to_device_bytes(),
147            );
148            receipt
149                .record_transfer(TransferDirection::DeviceToHost, transfers.device_to_host_bytes());
150            receipt.record_stage("gpu-sparse-grid");
151            receipt.record_stage("cpu-component-labeling");
152        }
153        Ok(ExecutionOutput::new(output, receipt))
154    }
155
156    fn extract_policy_output(
157        &self,
158        input: &PointCloud,
159        policy: ExecutionPolicy,
160    ) -> SpatialResult<(
161        EuclideanClusterResult,
162        ExecutionPolicy,
163        Option<spatialrust_core::TransferStats>,
164    )> {
165        #[cfg(feature = "segment-euclidean-gpu")]
166        {
167            let resolved = self.resolve_policy(input, policy)?;
168            if matches!(resolved, ExecutionPolicy::Gpu(DeviceKind::Wgpu)) {
169                if !self.gpu_grid_fits(input) {
170                    if policy.allows_fallback() {
171                        return Ok((self.extract(input)?, ExecutionPolicy::CpuSingle, None));
172                    }
173                    return Err(SpatialError::InvalidArgument(
174                        "point cloud does not fit the GPU Euclidean clustering grid".to_owned(),
175                    ));
176                }
177                let (output, transfers) =
178                    crate::cluster_gpu::GpuEuclideanClusterExtractor::new(self.config)
179                        .extract_with_receipt(input)?;
180                return Ok((output, resolved, Some(transfers)));
181            }
182        }
183
184        #[cfg(not(feature = "segment-euclidean-gpu"))]
185        if policy.requests_gpu() {
186            return Err(SpatialError::InvalidArgument(
187                "GPU Euclidean clustering requires the segment-euclidean-gpu feature".to_owned(),
188            ));
189        }
190
191        let resolved_policy = match policy {
192            ExecutionPolicy::Auto => ExecutionPolicy::CpuSingle,
193            other => other,
194        };
195        Ok((self.extract(input)?, resolved_policy, None))
196    }
197
198    #[cfg(feature = "segment-euclidean-gpu")]
199    fn gpu_grid_fits(&self, input: &PointCloud) -> bool {
200        let Ok((x, y, z)) = input.positions3() else {
201            return false;
202        };
203        spatialrust_search::uniform_grid_fits(x, y, z, self.config.cluster_tolerance)
204    }
205
206    #[cfg(feature = "segment-euclidean-gpu")]
207    fn should_use_gpu(&self, input: &PointCloud) -> bool {
208        self.config.effective_gpu_min_points().map_or(true, |min_points| input.len() >= min_points)
209            && spatialrust_gpu::WgpuRuntime::shared().is_ok()
210    }
211
212    #[cfg(feature = "segment-euclidean-gpu")]
213    fn resolve_policy(
214        &self,
215        input: &PointCloud,
216        policy: ExecutionPolicy,
217    ) -> SpatialResult<ExecutionPolicy> {
218        match policy {
219            ExecutionPolicy::Auto => {
220                if self.should_use_gpu(input) {
221                    Ok(ExecutionPolicy::Gpu(DeviceKind::Wgpu))
222                } else {
223                    Ok(ExecutionPolicy::CpuSingle)
224                }
225            }
226            ExecutionPolicy::Gpu(DeviceKind::Cpu) => Err(SpatialError::InvalidArgument(
227                "GPU execution policy cannot target the CPU device".to_owned(),
228            )),
229            ExecutionPolicy::Gpu(DeviceKind::Cuda) => Err(SpatialError::InvalidArgument(
230                "CUDA Euclidean clustering is not available".to_owned(),
231            )),
232            other => Ok(other),
233        }
234    }
235}
236
237impl PointCloudSegmenter for EuclideanClusterExtractor {
238    fn name(&self) -> &'static str {
239        "EuclideanClusterExtractor"
240    }
241}
242
243/// Builds labeled output from per-point connected-component roots.
244pub(crate) fn finalize_euclidean_clusters(
245    input: &PointCloud,
246    component_roots: &[u32],
247    config: EuclideanClusterConfig,
248) -> SpatialResult<EuclideanClusterResult> {
249    if component_roots.len() != input.len() {
250        return Err(SpatialError::InvalidArgument(
251            "component root labels must match point count".to_owned(),
252        ));
253    }
254
255    let mut sizes: HashMap<u32, usize> = HashMap::new();
256    for &root in component_roots {
257        *sizes.entry(root).or_insert(0) += 1;
258    }
259
260    let mut valid_roots: Vec<u32> = sizes
261        .iter()
262        .filter(|(_, &size)| size >= config.min_cluster_size && size <= config.max_cluster_size)
263        .map(|(&root, _)| root)
264        .collect();
265    valid_roots.sort_unstable();
266
267    let mut remap: HashMap<u32, i32> = HashMap::new();
268    for (cluster_id, root) in valid_roots.iter().enumerate() {
269        remap.insert(*root, cluster_id as i32);
270    }
271
272    let mut labels = vec![-1_i32; input.len()];
273    let mut cluster_sizes = Vec::with_capacity(valid_roots.len());
274    for root in &valid_roots {
275        cluster_sizes.push(sizes[root]);
276        for (index, &point_root) in component_roots.iter().enumerate() {
277            if point_root == *root {
278                labels[index] = remap[root];
279            }
280        }
281    }
282
283    Ok(EuclideanClusterResult {
284        cloud: with_labels(input, "label", labels)?,
285        cluster_count: cluster_sizes.len(),
286        cluster_sizes,
287    })
288}
289
290pub(crate) fn extract_cpu_roots(
291    input: &PointCloud,
292    config: EuclideanClusterConfig,
293) -> SpatialResult<Vec<u32>> {
294    let (x, y, z) = input.positions3()?;
295    let tree = KdTree::from_slices(x, y, z);
296    let len = input.len();
297    let mut processed = vec![false; len];
298    let mut roots = vec![0u32; len];
299
300    for seed in 0..len {
301        if processed[seed] {
302            continue;
303        }
304
305        let mut queue = VecDeque::from([seed]);
306        let mut cluster_indices = Vec::new();
307        processed[seed] = true;
308
309        while let Some(index) = queue.pop_front() {
310            cluster_indices.push(index);
311            let neighbors =
312                tree.radius_search(x[index], y[index], z[index], config.cluster_tolerance);
313            for neighbor in neighbors {
314                let candidate = neighbor.index;
315                if processed[candidate] {
316                    continue;
317                }
318                processed[candidate] = true;
319                queue.push_back(candidate);
320            }
321        }
322
323        let root = *cluster_indices.iter().min().unwrap_or(&seed) as u32;
324        for index in cluster_indices {
325            roots[index] = root;
326        }
327    }
328
329    Ok(roots)
330}
331
332fn validate_cluster_config(config: EuclideanClusterConfig) -> SpatialResult<()> {
333    if config.cluster_tolerance < 0.0 {
334        return Err(SpatialError::InvalidArgument(
335            "cluster_tolerance must be non-negative".to_owned(),
336        ));
337    }
338    if config.min_cluster_size == 0 {
339        return Err(SpatialError::InvalidArgument(
340            "min_cluster_size must be greater than zero".to_owned(),
341        ));
342    }
343    if config.max_cluster_size < config.min_cluster_size {
344        return Err(SpatialError::InvalidArgument(
345            "max_cluster_size must be >= min_cluster_size".to_owned(),
346        ));
347    }
348    Ok(())
349}
350
351#[cfg(test)]
352mod tests {
353    use super::{EuclideanClusterConfig, EuclideanClusterExtractor};
354    use spatialrust_core::{DeviceKind, ExecutionPolicy, PointCloudBuilder, StandardSchemas};
355
356    fn three_clusters() -> spatialrust_core::PointCloud {
357        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
358        for center in [(0.0, 0.0, 0.0), (10.0, 0.0, 0.0), (0.0, 10.0, 0.0)] {
359            for dx in 0..3 {
360                for dy in 0..3 {
361                    builder
362                        .push_point([center.0 + dx as f32, center.1 + dy as f32, center.2])
363                        .unwrap();
364                }
365            }
366        }
367        builder.build().unwrap()
368    }
369
370    #[test]
371    fn finds_three_separated_clusters() {
372        let input = three_clusters();
373        let extractor = EuclideanClusterExtractor::new(EuclideanClusterConfig {
374            cluster_tolerance: 1.5,
375            min_cluster_size: 3,
376            max_cluster_size: usize::MAX,
377            ..Default::default()
378        });
379        let result = extractor.extract(&input).unwrap();
380        assert_eq!(result.cluster_count, 3);
381        assert!(result.cluster_sizes.iter().all(|&size| size == 9));
382        assert!(result.cloud.field("label").is_ok());
383    }
384
385    #[test]
386    fn rejects_clusters_smaller_than_minimum() {
387        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
388        builder.push_point([0.0, 0.0, 0.0]).unwrap();
389        builder.push_point([0.1, 0.0, 0.0]).unwrap();
390        let input = builder.build().unwrap();
391
392        let extractor = EuclideanClusterExtractor::new(EuclideanClusterConfig {
393            cluster_tolerance: 0.5,
394            min_cluster_size: 3,
395            max_cluster_size: usize::MAX,
396            ..Default::default()
397        });
398        let result = extractor.extract(&input).unwrap();
399        assert_eq!(result.cluster_count, 0);
400    }
401
402    #[test]
403    fn rejects_unsupported_explicit_gpu_policy() {
404        let extractor = EuclideanClusterExtractor::new(EuclideanClusterConfig::default());
405        let error = extractor
406            .extract_with_policy(&three_clusters(), ExecutionPolicy::Gpu(DeviceKind::Cuda))
407            .unwrap_err();
408        assert!(matches!(error, spatialrust_core::SpatialError::InvalidArgument(_)));
409    }
410}
spatialrust_segmentation/cluster.rs

spatialrust_segmentation/
cluster.rs
