Skip to main content

spatialrust_segmentation/
ground.rs

1//! Grid-based ground segmentation for outdoor (LiDAR) scans.
2//!
3//! The cloud is binned into a 2D grid over the ground plane; each cell's minimum
4//! height seeds a ground-elevation estimate, which is then eroded against its
5//! neighbors (a morphological opening) so that isolated high cells — rooftops,
6//! vehicle tops — are pulled down to the surrounding ground level rather than
7//! seeding their own false ground. A point is ground if it sits within
8//! `height_threshold` of that estimate.
9
10use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
11
12use crate::cloud::{extract_mask, with_labels};
13use crate::segmenter::PointCloudSegmenter;
14
15/// Which axis points "up" (its minimum defines local ground height).
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub enum UpAxis {
18    /// +X is up.
19    X,
20    /// +Y is up.
21    Y,
22    /// +Z is up (the common case).
23    Z,
24}
25
26/// Configuration for [`GroundSegmenter`].
27#[derive(Clone, Copy, Debug, PartialEq)]
28pub struct GroundConfig {
29    /// Side length of each grid cell in the ground plane.
30    pub cell_size: f32,
31    /// Max height above the local ground estimate for a point to be ground.
32    pub height_threshold: f32,
33    /// Erosion window radius in cells (0 disables erosion).
34    pub erosion_cells: usize,
35    /// Which axis is "up".
36    pub up_axis: UpAxis,
37}
38
39impl Default for GroundConfig {
40    fn default() -> Self {
41        Self { cell_size: 0.5, height_threshold: 0.2, erosion_cells: 1, up_axis: UpAxis::Z }
42    }
43}
44
45impl GroundConfig {
46    /// Creates a config from the cell size and height threshold (Z up).
47    #[must_use]
48    pub const fn new(cell_size: f32, height_threshold: f32) -> Self {
49        Self { cell_size, height_threshold, erosion_cells: 1, up_axis: UpAxis::Z }
50    }
51}
52
53/// Result of ground segmentation.
54#[derive(Clone, Debug, PartialEq)]
55pub struct GroundSegmentation {
56    /// Points classified as ground.
57    pub ground: PointCloud,
58    /// Points classified as non-ground (objects, vegetation, structures).
59    pub non_ground: PointCloud,
60    /// Input cloud with a `label` field (1 = ground, 0 = non-ground).
61    pub labeled: PointCloud,
62    /// Number of ground points.
63    pub ground_count: usize,
64}
65
66/// Grid-based ground segmenter.
67#[derive(Clone, Copy, Debug, PartialEq)]
68pub struct GroundSegmenter {
69    config: GroundConfig,
70}
71
72impl GroundSegmenter {
73    /// Creates a segmenter from config.
74    #[must_use]
75    pub const fn new(config: GroundConfig) -> Self {
76        Self { config }
77    }
78
79    /// Returns the segmenter config.
80    #[must_use]
81    pub const fn config(&self) -> GroundConfig {
82        self.config
83    }
84
85    /// Computes the per-point ground mask (`true` = ground).
86    pub fn ground_mask(&self, input: &PointCloud) -> SpatialResult<Vec<bool>> {
87        if self.config.cell_size <= 0.0 || self.config.cell_size.is_nan() {
88            return Err(SpatialError::InvalidArgument("cell_size must be positive".to_owned()));
89        }
90        let len = input.len();
91        if len == 0 {
92            return Ok(Vec::new());
93        }
94
95        let (x, y, z) = input.positions3()?;
96        // (plane_a, plane_b) span the ground; `up` is the height axis.
97        let (pa, pb, up): (&[f32], &[f32], &[f32]) = match self.config.up_axis {
98            UpAxis::X => (y, z, x),
99            UpAxis::Y => (x, z, y),
100            UpAxis::Z => (x, y, z),
101        };
102
103        let inv_cell = 1.0 / self.config.cell_size;
104        let min_a = pa.iter().copied().fold(f32::INFINITY, f32::min);
105        let min_b = pb.iter().copied().fold(f32::INFINITY, f32::min);
106        let max_a = pa.iter().copied().fold(f32::NEG_INFINITY, f32::max);
107        let max_b = pb.iter().copied().fold(f32::NEG_INFINITY, f32::max);
108        let cols = (((max_a - min_a) * inv_cell) as usize) + 1;
109        let rows = (((max_b - min_b) * inv_cell) as usize) + 1;
110
111        let cell_of = |i: usize| {
112            let c = (((pa[i] - min_a) * inv_cell) as usize).min(cols - 1);
113            let r = (((pb[i] - min_b) * inv_cell) as usize).min(rows - 1);
114            r * cols + c
115        };
116
117        // Per-cell minimum height.
118        let mut cell_min = vec![f32::INFINITY; cols * rows];
119        for (i, &up_i) in up.iter().enumerate() {
120            let cell = cell_of(i);
121            if up_i < cell_min[cell] {
122                cell_min[cell] = up_i;
123            }
124        }
125
126        // Morphological erosion: each cell's ground reference is the minimum of
127        // its neighborhood, so a lone high cell adopts the surrounding ground.
128        let ground_ref = if self.config.erosion_cells == 0 {
129            cell_min.clone()
130        } else {
131            erode(&cell_min, cols, rows, self.config.erosion_cells)
132        };
133
134        let mut mask = vec![false; len];
135        for (i, &up_i) in up.iter().enumerate() {
136            let reference = ground_ref[cell_of(i)];
137            if reference.is_finite() && up_i - reference <= self.config.height_threshold {
138                mask[i] = true;
139            }
140        }
141        Ok(mask)
142    }
143
144    /// Segments the cloud into ground and non-ground partitions.
145    pub fn segment(&self, input: &PointCloud) -> SpatialResult<GroundSegmentation> {
146        let mask = self.ground_mask(input)?;
147        let non_ground_mask: Vec<bool> = mask.iter().map(|&g| !g).collect();
148        let labels: Vec<i32> = mask.iter().map(|&g| i32::from(g)).collect();
149        let ground_count = mask.iter().filter(|&&g| g).count();
150
151        Ok(GroundSegmentation {
152            ground: extract_mask(input, &mask)?,
153            non_ground: extract_mask(input, &non_ground_mask)?,
154            labeled: with_labels(input, "label", labels)?,
155            ground_count,
156        })
157    }
158}
159
160impl PointCloudSegmenter for GroundSegmenter {
161    fn name(&self) -> &'static str {
162        "GroundSegmenter"
163    }
164}
165
166/// Greyscale morphological erosion of a grid of heights by `radius` cells.
167fn erode(cells: &[f32], cols: usize, rows: usize, radius: usize) -> Vec<f32> {
168    let mut out = vec![f32::INFINITY; cells.len()];
169    let r = radius as isize;
170    for row in 0..rows {
171        for col in 0..cols {
172            let mut m = f32::INFINITY;
173            for dr in -r..=r {
174                for dc in -r..=r {
175                    let nr = row as isize + dr;
176                    let nc = col as isize + dc;
177                    if nr >= 0 && nr < rows as isize && nc >= 0 && nc < cols as isize {
178                        let value = cells[nr as usize * cols + nc as usize];
179                        if value < m {
180                            m = value;
181                        }
182                    }
183                }
184            }
185            out[row * cols + col] = m;
186        }
187    }
188    out
189}
190
191#[cfg(test)]
192mod tests {
193    use super::{GroundConfig, GroundSegmenter};
194    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
195
196    /// A flat ground grid plus a raised "building" block above part of it.
197    fn ground_with_building() -> spatialrust_core::PointCloud {
198        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
199        // Ground: 20x20 at z ~ 0.
200        for i in 0..20 {
201            for j in 0..20 {
202                builder.push_point([i as f32 * 0.5, j as f32 * 0.5, 0.0]).unwrap();
203            }
204        }
205        // Building roof: a 5x5 block at z = 3 over one corner.
206        for i in 0..5 {
207            for j in 0..5 {
208                builder.push_point([i as f32 * 0.5, j as f32 * 0.5, 3.0]).unwrap();
209            }
210        }
211        builder.build().unwrap()
212    }
213
214    #[test]
215    fn separates_ground_from_building() {
216        let cloud = ground_with_building();
217        let seg = GroundSegmenter::new(GroundConfig::new(0.6, 0.3)).segment(&cloud).unwrap();
218        // 400 ground points, 25 roof points.
219        assert_eq!(seg.ground_count, 400);
220        assert_eq!(seg.ground.len(), 400);
221        assert_eq!(seg.non_ground.len(), 25);
222        assert!(seg.labeled.field("label").is_ok());
223    }
224
225    #[test]
226    fn sloped_ground_stays_ground() {
227        // A ramp rising along x: each cell's local min tracks the slope, so the
228        // whole ramp should be classified as ground.
229        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
230        for i in 0..40 {
231            for j in 0..10 {
232                let x = i as f32 * 0.5;
233                builder.push_point([x, j as f32 * 0.5, x * 0.2]).unwrap();
234            }
235        }
236        let cloud = builder.build().unwrap();
237        let seg = GroundSegmenter::new(GroundConfig::new(0.6, 0.3)).segment(&cloud).unwrap();
238        // Almost everything is ground (the gentle slope fits the cell threshold).
239        assert!(seg.ground_count as f32 > cloud.len() as f32 * 0.95);
240    }
241
242    #[test]
243    fn rejects_bad_params() {
244        let cloud = ground_with_building();
245        assert!(GroundSegmenter::new(GroundConfig::new(0.0, 0.3)).ground_mask(&cloud).is_err());
246    }
247}