spatialrust_scene/
tsdf.rs

1//! Dense TSDF volume integration and mesh extraction.
2
3use spatialrust_math::Vec3;
4
5use crate::marching_cubes::{polygonise_tet, tetrahedra};
6use crate::{SceneError, SceneResult, TriangleMesh};
7
8/// Axis-aligned truncated signed distance volume.
9#[derive(Clone, Debug, PartialEq)]
10pub struct TsdfVolume {
11    origin: Vec3<f32>,
12    voxel_size: f32,
13    dims: [usize; 3],
14    distance: Vec<f32>,
15    weight: Vec<f32>,
16    truncation: f32,
17}
18
19impl TsdfVolume {
20    /// Creates an empty TSDF volume initialized to truncation distance.
21    pub fn try_new(
22        origin: Vec3<f32>,
23        voxel_size: f32,
24        dims: [usize; 3],
25        truncation: f32,
26    ) -> SceneResult<Self> {
27        if !(voxel_size.is_finite() && voxel_size > 0.0) {
28            return Err(SceneError::InvalidConfiguration("voxel_size must be > 0".into()));
29        }
30        if !(truncation.is_finite() && truncation > 0.0) {
31            return Err(SceneError::InvalidConfiguration("truncation must be > 0".into()));
32        }
33        if dims.contains(&0) {
34            return Err(SceneError::InvalidConfiguration("dims must be non-zero".into()));
35        }
36        let len = dims[0].saturating_mul(dims[1]).saturating_mul(dims[2]);
37        Ok(Self {
38            origin,
39            voxel_size,
40            dims,
41            distance: vec![truncation; len],
42            weight: vec![0.0; len],
43            truncation,
44        })
45    }
46
47    /// Returns voxel dimensions `[nx, ny, nz]`.
48    #[must_use]
49    pub fn dims(&self) -> [usize; 3] {
50        self.dims
51    }
52
53    /// Integrates one metric depth sample at a world-space point.
54    ///
55    /// Updates every voxel whose center lies within the truncation band, using
56    /// a view-ray signed distance so the field has zero crossings for meshing.
57    pub fn integrate_point(&mut self, point: Vec3<f32>, sensor_origin: Vec3<f32>) {
58        let depth = (point - sensor_origin).length();
59        if !(depth.is_finite() && depth > 1e-5) {
60            return;
61        }
62        let ray = (point - sensor_origin).normalize();
63        let radius = self.truncation;
64        let min = point - Vec3::new(radius, radius, radius);
65        let max = point + Vec3::new(radius, radius, radius);
66        let i0 = self.world_to_index_clamped(min);
67        let i1 = self.world_to_index_clamped(max);
68        for z in i0[2]..=i1[2] {
69            for y in i0[1]..=i1[1] {
70                for x in i0[0]..=i1[0] {
71                    let center = self.index_to_world([x, y, z]);
72                    let sdf = (point - center).dot(ray).clamp(-self.truncation, self.truncation);
73                    let flat = self.flat([x, y, z]);
74                    let w_old = self.weight[flat];
75                    let w_new = w_old + 1.0;
76                    self.distance[flat] = (self.distance[flat] * w_old + sdf) / w_new;
77                    self.weight[flat] = w_new;
78                }
79            }
80        }
81    }
82
83    /// Integrates every XYZ triple from interleaved storage.
84    pub fn integrate_xyz(&mut self, xyz: &[f32], sensor_origin: Vec3<f32>) -> SceneResult<()> {
85        if xyz.len() % 3 != 0 {
86            return Err(SceneError::InvalidConfiguration(
87                "xyz length must be a multiple of 3".into(),
88            ));
89        }
90        for chunk in xyz.chunks_exact(3) {
91            self.integrate_point(Vec3::new(chunk[0], chunk[1], chunk[2]), sensor_origin);
92        }
93        Ok(())
94    }
95
96    /// Extracts the zero isolevel surface with marching tetrahedra.
97    ///
98    /// Voxels with weight `< min_weight` are treated as free space (`+truncation`).
99    pub fn extract_mesh(&self, min_weight: f32) -> TriangleMesh {
100        let mut positions = Vec::new();
101        let mut indices = Vec::new();
102        if self.dims[0] < 2 || self.dims[1] < 2 || self.dims[2] < 2 {
103            return TriangleMesh { positions, indices };
104        }
105
106        for z in 0..self.dims[2] - 1 {
107            for y in 0..self.dims[1] - 1 {
108                for x in 0..self.dims[0] - 1 {
109                    let mut corner_pos = [Vec3::new(0.0, 0.0, 0.0); 8];
110                    let mut corner_val = [0.0f32; 8];
111                    for (corner, offset) in (0..8).zip(CORNER_OFFSETS.iter()) {
112                        let idx = [x + offset[0], y + offset[1], z + offset[2]];
113                        corner_pos[corner] = self.index_to_world(idx);
114                        corner_val[corner] = self.sample(idx, min_weight);
115                    }
116                    for tet in tetrahedra() {
117                        polygonise_tet(
118                            &mut positions,
119                            &mut indices,
120                            [
121                                corner_pos[tet[0]],
122                                corner_pos[tet[1]],
123                                corner_pos[tet[2]],
124                                corner_pos[tet[3]],
125                            ],
126                            [
127                                corner_val[tet[0]],
128                                corner_val[tet[1]],
129                                corner_val[tet[2]],
130                                corner_val[tet[3]],
131                            ],
132                            0.0,
133                        );
134                    }
135                }
136            }
137        }
138        TriangleMesh { positions, indices }
139    }
140
141    fn sample(&self, index: [usize; 3], min_weight: f32) -> f32 {
142        let flat = self.flat(index);
143        if self.weight[flat] < min_weight {
144            self.truncation
145        } else {
146            self.distance[flat]
147        }
148    }
149
150    fn world_to_index_clamped(&self, point: Vec3<f32>) -> [usize; 3] {
151        let ix = ((point.x - self.origin.x) / self.voxel_size).floor() as isize;
152        let iy = ((point.y - self.origin.y) / self.voxel_size).floor() as isize;
153        let iz = ((point.z - self.origin.z) / self.voxel_size).floor() as isize;
154        [
155            ix.clamp(0, self.dims[0] as isize - 1) as usize,
156            iy.clamp(0, self.dims[1] as isize - 1) as usize,
157            iz.clamp(0, self.dims[2] as isize - 1) as usize,
158        ]
159    }
160
161    fn index_to_world(&self, index: [usize; 3]) -> Vec3<f32> {
162        Vec3::new(
163            self.origin.x + (index[0] as f32 + 0.5) * self.voxel_size,
164            self.origin.y + (index[1] as f32 + 0.5) * self.voxel_size,
165            self.origin.z + (index[2] as f32 + 0.5) * self.voxel_size,
166        )
167    }
168
169    fn flat(&self, index: [usize; 3]) -> usize {
170        index[0] + self.dims[0] * (index[1] + self.dims[1] * index[2])
171    }
172}
173
174const CORNER_OFFSETS: [[usize; 3]; 8] =
175    [[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]];
176
177#[cfg(test)]
178mod tests {
179    use super::TsdfVolume;
180    use spatialrust_math::Vec3;
181
182    #[test]
183    fn integrates_and_extracts_non_empty_mesh() {
184        let mut volume =
185            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
186        volume.integrate_xyz(&[0.0, 0.0, 0.0, 0.2, 0.0, 0.0], Vec3::new(0.0, 0.0, -1.0)).unwrap();
187        let mesh = volume.extract_mesh(0.5);
188        assert!(!mesh.positions.is_empty());
189        assert_eq!(mesh.indices.len() % 3, 0);
190        assert!(mesh.triangle_count() >= 1);
191    }
192
193    #[test]
194    fn empty_weight_yields_empty_mesh() {
195        let volume =
196            TsdfVolume::try_new(Vec3::new(-1.0, -1.0, -1.0), 0.25, [8, 8, 8], 0.5).unwrap();
197        let mesh = volume.extract_mesh(1.0);
198        assert!(mesh.positions.is_empty());
199        assert!(mesh.indices.is_empty());
200    }
201}
spatialrust_scene/tsdf.rs

spatialrust_scene/
tsdf.rs
