spatialrust_scene/
tsdf.rs

1//! Dense TSDF volume integration and mesh extraction.
2
3use spatialrust_core::{HasPositions3, PointCloud};
4use spatialrust_math::{Isometry3, TransformPoint, Vec3};
5
6use crate::marching_cubes::{polygonise_tet, tetrahedra};
7use crate::{SceneError, SceneResult, TriangleMesh};
8
9/// Axis-aligned truncated signed distance volume.
10#[derive(Clone, Debug, PartialEq)]
11pub struct TsdfVolume {
12    origin: Vec3<f32>,
13    voxel_size: f32,
14    dims: [usize; 3],
15    distance: Vec<f32>,
16    weight: Vec<f32>,
17    truncation: f32,
18}
19
20impl TsdfVolume {
21    /// Creates an empty TSDF volume initialized to truncation distance.
22    ///
23    /// Rejects non-finite origins or voxel centers, invalid scales/dimensions
24    /// and field capacities that cannot be represented. Both fields are
25    /// reserved fallibly before initialization; reservation failures return
26    /// [`SceneError::InvalidConfiguration`]. This does not impose a memory
27    /// budget or guarantee recovery from operating-system overcommit.
28    pub fn try_new(
29        origin: Vec3<f32>,
30        voxel_size: f32,
31        dims: [usize; 3],
32        truncation: f32,
33    ) -> SceneResult<Self> {
34        if !(voxel_size.is_finite() && voxel_size > 0.0) {
35            return Err(SceneError::InvalidConfiguration("voxel_size must be > 0".into()));
36        }
37        if !(truncation.is_finite() && truncation > 0.0) {
38            return Err(SceneError::InvalidConfiguration("truncation must be > 0".into()));
39        }
40        if dims.contains(&0) {
41            return Err(SceneError::InvalidConfiguration("dims must be non-zero".into()));
42        }
43        if ![origin.x, origin.y, origin.z].into_iter().all(f32::is_finite) {
44            return Err(SceneError::InvalidConfiguration("origin must be finite".into()));
45        }
46        let len = dims[0]
47            .checked_mul(dims[1])
48            .and_then(|count| count.checked_mul(dims[2]))
49            .filter(|&count| count <= isize::MAX as usize / std::mem::size_of::<f32>())
50            .ok_or_else(|| {
51                SceneError::InvalidConfiguration("TSDF field capacity is not representable".into())
52            })?;
53        for (coordinate, dim) in [origin.x, origin.y, origin.z].into_iter().zip(dims) {
54            // Match index_to_world's f32 evaluation, including intermediate
55            // offset overflow. Positive offsets are monotonic on each axis.
56            let last_center = coordinate + ((dim - 1) as f32 + 0.5) * voxel_size;
57            if !last_center.is_finite() {
58                return Err(SceneError::InvalidConfiguration(
59                    "voxel centers must be finite".into(),
60                ));
61            }
62        }
63        let mut distance = Vec::new();
64        let mut weight = Vec::new();
65        distance.try_reserve_exact(len).map_err(|_| {
66            SceneError::InvalidConfiguration("unable to reserve TSDF distance storage".into())
67        })?;
68        weight.try_reserve_exact(len).map_err(|_| {
69            SceneError::InvalidConfiguration("unable to reserve TSDF weight storage".into())
70        })?;
71        distance.resize(len, truncation);
72        weight.resize(len, 0.0);
73        Ok(Self { origin, voxel_size, dims, distance, weight, truncation })
74    }
75
76    /// Returns voxel dimensions `[nx, ny, nz]`.
77    #[must_use]
78    pub fn dims(&self) -> [usize; 3] {
79        self.dims
80    }
81
82    /// Integrates one metric depth sample at a world-space point.
83    ///
84    /// Updates valid cells in the inclusive floor-index range of the cube
85    /// `point +/- truncation`, using a view-ray signed distance for meshing.
86    /// Fully disjoint support is ignored; partial overlap is clipped to the
87    /// grid. This retains the discretized cube rather than a spherical band.
88    pub fn integrate_point(&mut self, point: Vec3<f32>, sensor_origin: Vec3<f32>) {
89        let depth = (point - sensor_origin).length();
90        if !(depth.is_finite() && depth > 1e-5) {
91            return;
92        }
93        let ray = (point - sensor_origin).normalize();
94        let radius = self.truncation;
95        let min = point - Vec3::new(radius, radius, radius);
96        let max = point + Vec3::new(radius, radius, radius);
97        let i0 = self.world_to_index(min);
98        let i1 = self.world_to_index(max);
99        // Reject disjoint raw ranges before either endpoint is clamped onto
100        // a boundary cell. Keep the original quantization for every overlap.
101        if (0..3).any(|axis| i1[axis] < 0 || i0[axis] >= self.dims[axis] as isize) {
102            return;
103        }
104        let i0: [usize; 3] =
105            std::array::from_fn(|axis| i0[axis].clamp(0, self.dims[axis] as isize - 1) as usize);
106        let i1: [usize; 3] =
107            std::array::from_fn(|axis| i1[axis].clamp(0, self.dims[axis] as isize - 1) as usize);
108        for z in i0[2]..=i1[2] {
109            for y in i0[1]..=i1[1] {
110                for x in i0[0]..=i1[0] {
111                    let center = self.index_to_world([x, y, z]);
112                    let sdf = (point - center).dot(ray).clamp(-self.truncation, self.truncation);
113                    let flat = self.flat([x, y, z]);
114                    let w_old = self.weight[flat];
115                    let w_new = w_old + 1.0;
116                    self.distance[flat] = (self.distance[flat] * w_old + sdf) / w_new;
117                    self.weight[flat] = w_new;
118                }
119            }
120        }
121    }
122
123    /// Integrates every XYZ triple from interleaved storage.
124    pub fn integrate_xyz(&mut self, xyz: &[f32], sensor_origin: Vec3<f32>) -> SceneResult<()> {
125        if xyz.len() % 3 != 0 {
126            return Err(SceneError::InvalidConfiguration(
127                "xyz length must be a multiple of 3".into(),
128            ));
129        }
130        for chunk in xyz.chunks_exact(3) {
131            self.integrate_point(Vec3::new(chunk[0], chunk[1], chunk[2]), sensor_origin);
132        }
133        Ok(())
134    }
135
136    /// Integrates XYZ columns directly from a point cloud without interleaving
137    /// or allocating a temporary coordinate vector.
138    ///
139    /// Returns the number of point samples visited. Invalid metric samples are
140    /// ignored by [`TsdfVolume::integrate_point`] with the same semantics as
141    /// [`TsdfVolume::integrate_xyz`].
142    pub fn integrate_cloud(
143        &mut self,
144        cloud: &PointCloud,
145        sensor_origin: Vec3<f32>,
146    ) -> SceneResult<usize> {
147        let (x, y, z) = cloud.positions3()?;
148        for index in 0..cloud.len() {
149            self.integrate_point(Vec3::new(x[index], y[index], z[index]), sensor_origin);
150        }
151        Ok(cloud.len())
152    }
153
154    /// Integrates a point cloud after an explicit sensor-to-volume pose.
155    ///
156    /// Both positions and the supplied sensor origin are transformed by
157    /// `volume_t_sensor`; no transformed point-cloud allocation is created.
158    pub fn integrate_cloud_with_pose(
159        &mut self,
160        cloud: &PointCloud,
161        volume_t_sensor: Isometry3<f32>,
162        sensor_origin: Vec3<f32>,
163    ) -> SceneResult<usize> {
164        let (x, y, z) = cloud.positions3()?;
165        let volume_sensor_origin = volume_t_sensor.transform_point(sensor_origin);
166        for index in 0..cloud.len() {
167            let point = volume_t_sensor.transform_point(Vec3::new(x[index], y[index], z[index]));
168            self.integrate_point(point, volume_sensor_origin);
169        }
170        Ok(cloud.len())
171    }
172
173    /// Extracts the zero isolevel surface with marching tetrahedra.
174    ///
175    /// Voxels with weight `< min_weight` are treated as free space (`+truncation`).
176    pub fn extract_mesh(&self, min_weight: f32) -> TriangleMesh {
177        let mut positions = Vec::new();
178        let mut indices = Vec::new();
179        if self.dims[0] < 2 || self.dims[1] < 2 || self.dims[2] < 2 {
180            return TriangleMesh { positions, indices };
181        }
182
183        for z in 0..self.dims[2] - 1 {
184            for y in 0..self.dims[1] - 1 {
185                for x in 0..self.dims[0] - 1 {
186                    let mut corner_pos = [Vec3::new(0.0, 0.0, 0.0); 8];
187                    let mut corner_val = [0.0f32; 8];
188                    for (corner, offset) in (0..8).zip(CORNER_OFFSETS.iter()) {
189                        let idx = [x + offset[0], y + offset[1], z + offset[2]];
190                        corner_pos[corner] = self.index_to_world(idx);
191                        corner_val[corner] = self.sample(idx, min_weight);
192                    }
193                    for tet in tetrahedra() {
194                        polygonise_tet(
195                            &mut positions,
196                            &mut indices,
197                            [
198                                corner_pos[tet[0]],
199                                corner_pos[tet[1]],
200                                corner_pos[tet[2]],
201                                corner_pos[tet[3]],
202                            ],
203                            [
204                                corner_val[tet[0]],
205                                corner_val[tet[1]],
206                                corner_val[tet[2]],
207                                corner_val[tet[3]],
208                            ],
209                            0.0,
210                        );
211                    }
212                }
213            }
214        }
215        TriangleMesh { positions, indices }
216    }
217
218    fn sample(&self, index: [usize; 3], min_weight: f32) -> f32 {
219        let flat = self.flat(index);
220        if self.weight[flat] < min_weight {
221            self.truncation
222        } else {
223            self.distance[flat]
224        }
225    }
226
227    fn world_to_index(&self, point: Vec3<f32>) -> [isize; 3] {
228        let ix = ((point.x - self.origin.x) / self.voxel_size).floor() as isize;
229        let iy = ((point.y - self.origin.y) / self.voxel_size).floor() as isize;
230        let iz = ((point.z - self.origin.z) / self.voxel_size).floor() as isize;
231        [ix, iy, iz]
232    }
233
234    fn index_to_world(&self, index: [usize; 3]) -> Vec3<f32> {
235        Vec3::new(
236            self.origin.x + (index[0] as f32 + 0.5) * self.voxel_size,
237            self.origin.y + (index[1] as f32 + 0.5) * self.voxel_size,
238            self.origin.z + (index[2] as f32 + 0.5) * self.voxel_size,
239        )
240    }
241
242    fn flat(&self, index: [usize; 3]) -> usize {
243        index[0] + self.dims[0] * (index[1] + self.dims[1] * index[2])
244    }
245}
246
247const CORNER_OFFSETS: [[usize; 3]; 8] =
248    [[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]];
249
250#[cfg(test)]
251#[path = "tsdf_support_tests.rs"]
252mod support_tests;
253
254#[cfg(test)]
255#[path = "tsdf_grid_tests.rs"]
256mod grid_tests;
257
258#[cfg(test)]
259#[path = "tsdf_mesh_tests.rs"]
260mod mesh_tests;
261
262#[cfg(test)]
263mod tests {
264    use super::TsdfVolume;
265    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
266    use spatialrust_math::Vec3;
267
268    #[test]
269    fn integrates_and_extracts_non_empty_mesh() {
270        let mut volume =
271            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
272        volume.integrate_xyz(&[0.0, 0.0, 0.0, 0.2, 0.0, 0.0], Vec3::new(0.0, 0.0, -1.0)).unwrap();
273        let mesh = volume.extract_mesh(0.5);
274        assert!(!mesh.positions.is_empty());
275        assert_eq!(mesh.indices.len() % 3, 0);
276        assert!(mesh.triangle_count() >= 1);
277    }
278
279    #[test]
280    fn empty_weight_yields_empty_mesh() {
281        let volume =
282            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
283        let mesh = volume.extract_mesh(1.0);
284        assert!(mesh.positions.is_empty());
285        assert!(mesh.indices.is_empty());
286    }
287
288    #[test]
289    fn integrates_point_cloud_columns_without_interleave() {
290        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
291        builder.push_point([0.0, 0.0, 0.0]).unwrap();
292        builder.push_point([0.2, 0.0, 0.0]).unwrap();
293        let cloud = builder.build().unwrap();
294        let origin = Vec3::new(0.0, 0.0, -1.0);
295        let mut from_cloud =
296            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
297        let mut from_xyz = from_cloud.clone();
298        assert_eq!(from_cloud.integrate_cloud(&cloud, origin).unwrap(), 2);
299        from_xyz.integrate_xyz(&[0.0, 0.0, 0.0, 0.2, 0.0, 0.0], origin).unwrap();
300        assert_eq!(from_cloud, from_xyz);
301    }
302
303    #[test]
304    fn integrates_point_cloud_with_explicit_sensor_pose() {
305        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
306        builder.push_point([0.0, 0.0, 0.0]).unwrap();
307        builder.push_point([0.2, 0.0, 0.0]).unwrap();
308        let cloud = builder.build().unwrap();
309        let mut volume =
310            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
311        let pose = spatialrust_math::Isometry3::new(
312            spatialrust_math::Quat::<f32>::identity(),
313            Vec3::new(0.5, 0.0, 0.0),
314        );
315        assert_eq!(
316            volume.integrate_cloud_with_pose(&cloud, pose, Vec3::new(0.0, 0.0, -1.0)).unwrap(),
317            2
318        );
319    }
320}
spatialrust_scene/tsdf.rs

spatialrust_scene/
tsdf.rs
