spatialrust_segmentation/
ground.rs1use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
11
12use crate::cloud::{extract_mask, with_labels};
13use crate::segmenter::PointCloudSegmenter;
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub enum UpAxis {
18 X,
20 Y,
22 Z,
24}
25
26#[derive(Clone, Copy, Debug, PartialEq)]
28pub struct GroundConfig {
29 pub cell_size: f32,
31 pub height_threshold: f32,
33 pub erosion_cells: usize,
35 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 #[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#[derive(Clone, Debug, PartialEq)]
55pub struct GroundSegmentation {
56 pub ground: PointCloud,
58 pub non_ground: PointCloud,
60 pub labeled: PointCloud,
62 pub ground_count: usize,
64}
65
66#[derive(Clone, Copy, Debug, PartialEq)]
68pub struct GroundSegmenter {
69 config: GroundConfig,
70}
71
72impl GroundSegmenter {
73 #[must_use]
75 pub const fn new(config: GroundConfig) -> Self {
76 Self { config }
77 }
78
79 #[must_use]
81 pub const fn config(&self) -> GroundConfig {
82 self.config
83 }
84
85 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 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 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 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 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
166fn 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 fn ground_with_building() -> spatialrust_core::PointCloud {
198 let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
199 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 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 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 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 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}