1use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub enum VoxelFill {
8 Occupancy,
10 Count,
12}
13
14#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct VoxelGridConfig {
17 pub voxel_size: f32,
19 pub origin: Option<[f32; 3]>,
21 pub dims: Option<[usize; 3]>,
23 pub fill: VoxelFill,
25}
26
27impl Default for VoxelGridConfig {
28 fn default() -> Self {
29 Self { voxel_size: 0.1, origin: None, dims: None, fill: VoxelFill::Occupancy }
30 }
31}
32
33impl VoxelGridConfig {
34 #[must_use]
36 pub fn new(voxel_size: f32) -> Self {
37 Self { voxel_size, ..Self::default() }
38 }
39}
40
41#[derive(Clone, Debug, PartialEq)]
43pub struct OccupancyGrid {
44 pub dims: [usize; 3],
46 pub origin: [f32; 3],
48 pub voxel_size: f32,
50 pub data: Vec<f32>,
52}
53
54impl OccupancyGrid {
55 #[must_use]
57 pub fn len(&self) -> usize {
58 self.data.len()
59 }
60
61 #[must_use]
63 pub fn is_empty(&self) -> bool {
64 self.data.is_empty()
65 }
66
67 #[must_use]
69 pub fn occupied_count(&self) -> usize {
70 self.data.iter().filter(|&&v| v != 0.0).count()
71 }
72
73 #[must_use]
75 pub fn get(&self, x: usize, y: usize, z: usize) -> Option<f32> {
76 let [nx, ny, nz] = self.dims;
77 if x >= nx || y >= ny || z >= nz {
78 return None;
79 }
80 Some(self.data[z * ny * nx + y * nx + x])
81 }
82}
83
84pub fn voxelize(cloud: &PointCloud, config: VoxelGridConfig) -> SpatialResult<OccupancyGrid> {
86 if config.voxel_size <= 0.0 || config.voxel_size.is_nan() {
87 return Err(SpatialError::InvalidArgument("voxel_size must be positive".to_owned()));
88 }
89 let (x, y, z) = cloud.positions3()?;
90 let len = cloud.len();
91
92 let origin = match config.origin {
93 Some(o) => o,
94 None => {
95 if len == 0 {
96 [0.0, 0.0, 0.0]
97 } else {
98 [
99 x.iter().copied().fold(f32::INFINITY, f32::min),
100 y.iter().copied().fold(f32::INFINITY, f32::min),
101 z.iter().copied().fold(f32::INFINITY, f32::min),
102 ]
103 }
104 }
105 };
106
107 let inv = 1.0 / config.voxel_size;
108 let dims = match config.dims {
109 Some(d) => d,
110 None => {
111 if len == 0 {
112 [1, 1, 1]
113 } else {
114 let span = |vals: &[f32], o: f32| {
115 let max = vals.iter().copied().fold(f32::NEG_INFINITY, f32::max);
116 (((max - o) * inv).floor() as usize) + 1
117 };
118 [span(x, origin[0]), span(y, origin[1]), span(z, origin[2])]
119 }
120 }
121 };
122 let [nx, ny, nz] = dims;
123 let total = nx.checked_mul(ny).and_then(|v| v.checked_mul(nz)).ok_or_else(|| {
124 SpatialError::InvalidArgument("voxel grid dimensions overflow".to_owned())
125 })?;
126
127 let mut data = vec![0.0_f32; total];
128 for i in 0..len {
129 let vx = ((x[i] - origin[0]) * inv).floor();
130 let vy = ((y[i] - origin[1]) * inv).floor();
131 let vz = ((z[i] - origin[2]) * inv).floor();
132 if vx < 0.0 || vy < 0.0 || vz < 0.0 {
133 continue;
134 }
135 let (vx, vy, vz) = (vx as usize, vy as usize, vz as usize);
136 if vx >= nx || vy >= ny || vz >= nz {
137 continue;
138 }
139 let idx = vz * ny * nx + vy * nx + vx;
140 match config.fill {
141 VoxelFill::Occupancy => data[idx] = 1.0,
142 VoxelFill::Count => data[idx] += 1.0,
143 }
144 }
145
146 Ok(OccupancyGrid { dims, origin, voxel_size: config.voxel_size, data })
147}
148
149#[cfg(test)]
150mod tests {
151 use super::*;
152 use spatialrust_core::{PointCloudBuilder, StandardSchemas};
153
154 fn cloud(points: &[[f32; 3]]) -> PointCloud {
155 let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
156 for p in points {
157 builder.push_point(*p).unwrap();
158 }
159 builder.build().unwrap()
160 }
161
162 #[test]
163 fn occupancy_marks_filled_voxels() {
164 let c = cloud(&[[0.05, 0.05, 0.05], [0.15, 0.05, 0.05]]);
166 let grid = voxelize(&c, VoxelGridConfig::new(0.1)).unwrap();
167 assert_eq!(grid.dims, [2, 1, 1]);
168 assert_eq!(grid.occupied_count(), 2);
169 assert_eq!(grid.get(0, 0, 0), Some(1.0));
170 assert_eq!(grid.get(1, 0, 0), Some(1.0));
171 }
172
173 #[test]
174 fn count_accumulates_points_per_voxel() {
175 let c = cloud(&[
177 [0.01, 0.01, 0.01],
178 [0.02, 0.02, 0.02],
179 [0.03, 0.03, 0.03],
180 [0.55, 0.01, 0.01],
181 ]);
182 let config = VoxelGridConfig { fill: VoxelFill::Count, ..VoxelGridConfig::new(0.1) };
183 let grid = voxelize(&c, config).unwrap();
184 assert_eq!(grid.get(0, 0, 0), Some(3.0));
185 assert_eq!(grid.occupied_count(), 2);
186 let total: f32 = grid.data.iter().sum();
188 assert_eq!(total, 4.0);
189 }
190
191 #[test]
192 fn fixed_dims_drop_out_of_bounds_points() {
193 let c = cloud(&[[0.05, 0.05, 0.05], [100.0, 0.0, 0.0]]);
194 let config = VoxelGridConfig {
195 origin: Some([0.0, 0.0, 0.0]),
196 dims: Some([2, 2, 2]),
197 ..VoxelGridConfig::new(0.1)
198 };
199 let grid = voxelize(&c, config).unwrap();
200 assert_eq!(grid.dims, [2, 2, 2]);
201 assert_eq!(grid.len(), 8);
202 assert_eq!(grid.occupied_count(), 1);
204 }
205
206 #[test]
207 fn rejects_bad_voxel_size() {
208 let c = cloud(&[[0.0, 0.0, 0.0]]);
209 assert!(voxelize(&c, VoxelGridConfig::new(0.0)).is_err());
210 }
211}