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    pub fn try_new(
23        origin: Vec3<f32>,
24        voxel_size: f32,
25        dims: [usize; 3],
26        truncation: f32,
27    ) -> SceneResult<Self> {
28        if !(voxel_size.is_finite() && voxel_size > 0.0) {
29            return Err(SceneError::InvalidConfiguration("voxel_size must be > 0".into()));
30        }
31        if !(truncation.is_finite() && truncation > 0.0) {
32            return Err(SceneError::InvalidConfiguration("truncation must be > 0".into()));
33        }
34        if dims.contains(&0) {
35            return Err(SceneError::InvalidConfiguration("dims must be non-zero".into()));
36        }
37        let len = dims[0].saturating_mul(dims[1]).saturating_mul(dims[2]);
38        Ok(Self {
39            origin,
40            voxel_size,
41            dims,
42            distance: vec![truncation; len],
43            weight: vec![0.0; len],
44            truncation,
45        })
46    }
47
48    /// Returns voxel dimensions `[nx, ny, nz]`.
49    #[must_use]
50    pub fn dims(&self) -> [usize; 3] {
51        self.dims
52    }
53
54    /// Integrates one metric depth sample at a world-space point.
55    ///
56    /// Updates every voxel whose center lies within the truncation band, using
57    /// a view-ray signed distance so the field has zero crossings for meshing.
58    pub fn integrate_point(&mut self, point: Vec3<f32>, sensor_origin: Vec3<f32>) {
59        let depth = (point - sensor_origin).length();
60        if !(depth.is_finite() && depth > 1e-5) {
61            return;
62        }
63        let ray = (point - sensor_origin).normalize();
64        let radius = self.truncation;
65        let min = point - Vec3::new(radius, radius, radius);
66        let max = point + Vec3::new(radius, radius, radius);
67        let i0 = self.world_to_index_clamped(min);
68        let i1 = self.world_to_index_clamped(max);
69        for z in i0[2]..=i1[2] {
70            for y in i0[1]..=i1[1] {
71                for x in i0[0]..=i1[0] {
72                    let center = self.index_to_world([x, y, z]);
73                    let sdf = (point - center).dot(ray).clamp(-self.truncation, self.truncation);
74                    let flat = self.flat([x, y, z]);
75                    let w_old = self.weight[flat];
76                    let w_new = w_old + 1.0;
77                    self.distance[flat] = (self.distance[flat] * w_old + sdf) / w_new;
78                    self.weight[flat] = w_new;
79                }
80            }
81        }
82    }
83
84    /// Integrates every XYZ triple from interleaved storage.
85    pub fn integrate_xyz(&mut self, xyz: &[f32], sensor_origin: Vec3<f32>) -> SceneResult<()> {
86        if xyz.len() % 3 != 0 {
87            return Err(SceneError::InvalidConfiguration(
88                "xyz length must be a multiple of 3".into(),
89            ));
90        }
91        for chunk in xyz.chunks_exact(3) {
92            self.integrate_point(Vec3::new(chunk[0], chunk[1], chunk[2]), sensor_origin);
93        }
94        Ok(())
95    }
96
97    /// Integrates XYZ columns directly from a point cloud without interleaving
98    /// or allocating a temporary coordinate vector.
99    ///
100    /// Returns the number of point samples visited. Invalid metric samples are
101    /// ignored by [`TsdfVolume::integrate_point`] with the same semantics as
102    /// [`TsdfVolume::integrate_xyz`].
103    pub fn integrate_cloud(
104        &mut self,
105        cloud: &PointCloud,
106        sensor_origin: Vec3<f32>,
107    ) -> SceneResult<usize> {
108        let (x, y, z) = cloud.positions3()?;
109        for index in 0..cloud.len() {
110            self.integrate_point(Vec3::new(x[index], y[index], z[index]), sensor_origin);
111        }
112        Ok(cloud.len())
113    }
114
115    /// Integrates a point cloud after an explicit sensor-to-volume pose.
116    ///
117    /// Both positions and the supplied sensor origin are transformed by
118    /// `volume_t_sensor`; no transformed point-cloud allocation is created.
119    pub fn integrate_cloud_with_pose(
120        &mut self,
121        cloud: &PointCloud,
122        volume_t_sensor: Isometry3<f32>,
123        sensor_origin: Vec3<f32>,
124    ) -> SceneResult<usize> {
125        let (x, y, z) = cloud.positions3()?;
126        let volume_sensor_origin = volume_t_sensor.transform_point(sensor_origin);
127        for index in 0..cloud.len() {
128            let point = volume_t_sensor.transform_point(Vec3::new(x[index], y[index], z[index]));
129            self.integrate_point(point, volume_sensor_origin);
130        }
131        Ok(cloud.len())
132    }
133
134    /// Extracts the zero isolevel surface with marching tetrahedra.
135    ///
136    /// Voxels with weight `< min_weight` are treated as free space (`+truncation`).
137    pub fn extract_mesh(&self, min_weight: f32) -> TriangleMesh {
138        let mut positions = Vec::new();
139        let mut indices = Vec::new();
140        if self.dims[0] < 2 || self.dims[1] < 2 || self.dims[2] < 2 {
141            return TriangleMesh { positions, indices };
142        }
143
144        for z in 0..self.dims[2] - 1 {
145            for y in 0..self.dims[1] - 1 {
146                for x in 0..self.dims[0] - 1 {
147                    let mut corner_pos = [Vec3::new(0.0, 0.0, 0.0); 8];
148                    let mut corner_val = [0.0f32; 8];
149                    for (corner, offset) in (0..8).zip(CORNER_OFFSETS.iter()) {
150                        let idx = [x + offset[0], y + offset[1], z + offset[2]];
151                        corner_pos[corner] = self.index_to_world(idx);
152                        corner_val[corner] = self.sample(idx, min_weight);
153                    }
154                    for tet in tetrahedra() {
155                        polygonise_tet(
156                            &mut positions,
157                            &mut indices,
158                            [
159                                corner_pos[tet[0]],
160                                corner_pos[tet[1]],
161                                corner_pos[tet[2]],
162                                corner_pos[tet[3]],
163                            ],
164                            [
165                                corner_val[tet[0]],
166                                corner_val[tet[1]],
167                                corner_val[tet[2]],
168                                corner_val[tet[3]],
169                            ],
170                            0.0,
171                        );
172                    }
173                }
174            }
175        }
176        TriangleMesh { positions, indices }
177    }
178
179    fn sample(&self, index: [usize; 3], min_weight: f32) -> f32 {
180        let flat = self.flat(index);
181        if self.weight[flat] < min_weight {
182            self.truncation
183        } else {
184            self.distance[flat]
185        }
186    }
187
188    fn world_to_index_clamped(&self, point: Vec3<f32>) -> [usize; 3] {
189        let ix = ((point.x - self.origin.x) / self.voxel_size).floor() as isize;
190        let iy = ((point.y - self.origin.y) / self.voxel_size).floor() as isize;
191        let iz = ((point.z - self.origin.z) / self.voxel_size).floor() as isize;
192        [
193            ix.clamp(0, self.dims[0] as isize - 1) as usize,
194            iy.clamp(0, self.dims[1] as isize - 1) as usize,
195            iz.clamp(0, self.dims[2] as isize - 1) as usize,
196        ]
197    }
198
199    fn index_to_world(&self, index: [usize; 3]) -> Vec3<f32> {
200        Vec3::new(
201            self.origin.x + (index[0] as f32 + 0.5) * self.voxel_size,
202            self.origin.y + (index[1] as f32 + 0.5) * self.voxel_size,
203            self.origin.z + (index[2] as f32 + 0.5) * self.voxel_size,
204        )
205    }
206
207    fn flat(&self, index: [usize; 3]) -> usize {
208        index[0] + self.dims[0] * (index[1] + self.dims[1] * index[2])
209    }
210}
211
212const CORNER_OFFSETS: [[usize; 3]; 8] =
213    [[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]];
214
215#[cfg(test)]
216mod tests {
217    use super::TsdfVolume;
218    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
219    use spatialrust_math::Vec3;
220
221    #[test]
222    fn integrates_and_extracts_non_empty_mesh() {
223        let mut volume =
224            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
225        volume.integrate_xyz(&[0.0, 0.0, 0.0, 0.2, 0.0, 0.0], Vec3::new(0.0, 0.0, -1.0)).unwrap();
226        let mesh = volume.extract_mesh(0.5);
227        assert!(!mesh.positions.is_empty());
228        assert_eq!(mesh.indices.len() % 3, 0);
229        assert!(mesh.triangle_count() >= 1);
230    }
231
232    #[test]
233    fn empty_weight_yields_empty_mesh() {
234        let volume =
235            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
236        let mesh = volume.extract_mesh(1.0);
237        assert!(mesh.positions.is_empty());
238        assert!(mesh.indices.is_empty());
239    }
240
241    #[test]
242    fn integrates_point_cloud_columns_without_interleave() {
243        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
244        builder.push_point([0.0, 0.0, 0.0]).unwrap();
245        builder.push_point([0.2, 0.0, 0.0]).unwrap();
246        let cloud = builder.build().unwrap();
247        let origin = Vec3::new(0.0, 0.0, -1.0);
248        let mut from_cloud =
249            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
250        let mut from_xyz = from_cloud.clone();
251        assert_eq!(from_cloud.integrate_cloud(&cloud, origin).unwrap(), 2);
252        from_xyz.integrate_xyz(&[0.0, 0.0, 0.0, 0.2, 0.0, 0.0], origin).unwrap();
253        assert_eq!(from_cloud, from_xyz);
254    }
255
256    #[test]
257    fn integrates_point_cloud_with_explicit_sensor_pose() {
258        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
259        builder.push_point([0.0, 0.0, 0.0]).unwrap();
260        builder.push_point([0.2, 0.0, 0.0]).unwrap();
261        let cloud = builder.build().unwrap();
262        let mut volume =
263            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
264        let pose = spatialrust_math::Isometry3::new(
265            spatialrust_math::Quat::<f32>::identity(),
266            Vec3::new(0.5, 0.0, 0.0),
267        );
268        assert_eq!(
269            volume.integrate_cloud_with_pose(&cloud, pose, Vec3::new(0.0, 0.0, -1.0)).unwrap(),
270            2
271        );
272    }
273}
spatialrust_scene/tsdf.rs

spatialrust_scene/
tsdf.rs
