Skip to main content

del_msh_cpu/
pix2depth.rs

1pub struct Depth;
2
3impl<T> crate::trimesh3_raycast::ScalarRender<T> for Depth
4where
5    T: num_traits::Float,
6{
7    fn fwd(
8        &self,
9        bc: &[T; 3],
10        i_tri: u32,
11        tri2vtx: &[u32],
12        vtx2xyz: &[T],
13        transform_world2ndc: &[T; 16],
14    ) -> T {
15        if i_tri == u32::MAX {
16            return T::zero();
17        };
18        let q = crate::trimesh3::to_tri3(tri2vtx, vtx2xyz, i_tri as usize)
19            .position_from_barycentric_coordinates(bc[0], bc[1]);
20        let ndc =
21            del_geo_core::mat4_col_major::transform_homogeneous(transform_world2ndc, &q).unwrap();
22        let one = T::one();
23        let half = one / (one + one);
24        (ndc[2] + one) * half
25    }
26
27    fn bwd(
28        &self,
29        dldw_depth: T,
30        p0: &[T; 3],
31        p1: &[T; 3],
32        p2: &[T; 3],
33        ray_org: &[T; 3],
34        ray_dir: &[T; 3],
35        transform_world2ndc: &[T; 16],
36    ) -> ([T; 3], [T; 3], [T; 3]) {
37        let zero = T::zero();
38        let one = T::one();
39        let half = one / (one + one);
40        let dldw_ndc = [zero, zero, half * dldw_depth];
41        let (_t, bc) =
42            del_geo_core::tri3::intersection_against_line(p0, p1, p2, ray_org, ray_dir).unwrap();
43        let q = del_geo_core::tri3::position_from_barycentric_coords(p0, p1, p2, &bc);
44        let dndcdq = del_geo_core::mat4_col_major::jacobian_transform(transform_world2ndc, &q);
45        let dndcdq_t = del_geo_core::mat3_col_major::transpose(&dndcdq);
46        let dldw_q = del_geo_core::mat3_col_major::mult_vec(&dndcdq_t, &dldw_ndc);
47        let dldw_t = del_geo_core::vec3::dot(ray_dir, &dldw_q);
48        let (_t, _u, _v, dldw_p0, dldw_p1, dldw_p2) =
49            del_geo_core::tri3::intersection_against_line_bwd_wrt_tri(
50                p0, p1, p2, ray_org, ray_dir, dldw_t, zero, zero,
51            );
52        (dldw_p0, dldw_p1, dldw_p2)
53    }
54}
55
56#[test]
57fn test_hoge() {
58    use del_geo_core::vec3::Vec3;
59    let p0: [[f64; 3]; 3] = [[-13., -5., 8.], [14., -5., 8.], [1., 3., -3.]];
60    let ray_org = [8., 11., 10.];
61    let ray_dir = [1., 0., 2.].sub(&ray_org);
62    let transform_world2ndc = [
63        4., -1., 5., 1., 1., 3., 9., 3., 1., 4., 2., 2., -1., -2., 3., 3.,
64    ];
65    use crate::trimesh3_raycast::ScalarRender;
66    let depth_layer = Depth;
67    let (_t0, bc0) =
68        del_geo_core::tri3::intersection_against_line(&p0[0], &p0[1], &p0[2], &ray_org, &ray_dir)
69            .unwrap();
70    let depth0 = depth_layer.fwd(&bc0, 0, &[0, 1, 2], p0.as_flattened(), &transform_world2ndc);
71    let dldw_depth = 1.3;
72    let l0 = depth0 * dldw_depth;
73    let (dldw_p0, dldw_p1, dldw_p2) = depth_layer.bwd(
74        dldw_depth,
75        &p0[0],
76        &p0[1],
77        &p0[2],
78        &ray_org,
79        &ray_dir,
80        &transform_world2ndc,
81    );
82    let eps = 1.0e-5;
83    for (i_node, i_dim) in itertools::iproduct!(0..3, 0..3) {
84        let p1 = {
85            let mut p1 = p0;
86            p1[i_node][i_dim] += eps;
87            p1
88        };
89        let (_t, bc1) = del_geo_core::tri3::intersection_against_line(
90            &p1[0], &p1[1], &p1[2], &ray_org, &ray_dir,
91        )
92        .unwrap();
93        let depth1 = depth_layer.fwd(&bc1, 0, &[0, 1, 2], p1.as_flattened(), &transform_world2ndc);
94        let l1 = depth1 * dldw_depth;
95        let num_diff = (l1 - l0) / eps;
96        let ana_diff = match i_node {
97            0 => dldw_p0[i_dim],
98            1 => dldw_p1[i_dim],
99            2 => dldw_p2[i_dim],
100            _ => unreachable!(),
101        };
102        println!("{i_node} {i_dim} {num_diff}, {ana_diff}");
103        assert!((num_diff - ana_diff).abs() < 1.0e-5);
104    }
105}
106
107pub fn pix2depth_from_pix2tri(
108    pix2depth: &mut [f32],
109    pix2tri: &[u32],
110    tri2vtx: &[u32],
111    vtx2xyz: &[f32],
112    img_shape: (usize, usize), // (width, height)
113    transform_ndc2world: &[f32; 16],
114) {
115    let transform_world2ndc =
116        del_geo_core::mat4_col_major::try_inverse_with_pivot(transform_ndc2world).unwrap();
117    let fn_pix2depth = |i_pix: usize| -> f32 {
118        let i_tri = pix2tri[i_pix];
119        if i_tri == u32::MAX {
120            return 0f32;
121        }
122        let i_w = i_pix % img_shape.0;
123        let i_h = i_pix / img_shape.0;
124        let (ray_org, ray_dir) =
125            del_geo_core::mat4_col_major::ray_from_transform_ndc2world_and_pixel_coordinates(
126                (i_w as f32 + 0.5, i_h as f32 + 0.5),
127                &(img_shape.0 as f32, img_shape.1 as f32),
128                transform_ndc2world,
129            );
130        let tri = crate::trimesh3::to_tri3(tri2vtx, vtx2xyz, i_tri as usize);
131        let (coeff, _bc) = del_geo_core::tri3::intersection_against_line(
132            tri.p0, tri.p1, tri.p2, &ray_org, &ray_dir,
133        )
134        .unwrap();
135        let pos_world = del_geo_core::vec3::axpy(coeff, &ray_dir, &ray_org);
136        let pos_ndc =
137            del_geo_core::mat4_col_major::transform_homogeneous(&transform_world2ndc, &pos_world)
138                .unwrap();
139        (pos_ndc[2] + 1f32) * 0.5f32
140    };
141    use rayon::prelude::*;
142    pix2depth
143        .par_iter_mut()
144        .enumerate()
145        .for_each(|(i_pix, depth)| *depth = fn_pix2depth(i_pix));
146}
147
148pub fn render_depth_bvh(
149    image_size: (usize, usize),
150    pix2depth: &mut [f32],
151    transform_ndc2world: &[f32; 16],
152    tri2vtx: &[usize],
153    vtx2xyz: &[f32],
154    bvhnodes: &[usize],
155    bvhnode2aabb: &[f32],
156) {
157    let transform_world2ndc: [f32; 16] =
158        del_geo_core::mat4_col_major::try_inverse(transform_ndc2world).unwrap();
159    let (width, height) = image_size;
160    for ih in 0..height {
161        for iw in 0..width {
162            let (ray_org, ray_dir) =
163                del_geo_core::mat4_col_major::ray_from_transform_ndc2world_and_pixel_coordinates(
164                    (iw as f32 + 0.5, ih as f32 + 0.5),
165                    &(image_size.0 as f32, image_size.1 as f32),
166                    transform_ndc2world,
167                );
168            let mut hits = vec![];
169            crate::search_bvh3::intersections_ray(
170                &mut hits,
171                &ray_org,
172                &ray_dir,
173                &crate::search_bvh3::TriMeshWithBvh {
174                    tri2vtx,
175                    vtx2xyz,
176                    bvhnodes,
177                    bvhnode2aabb,
178                },
179                0,
180            );
181            hits.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
182            let Some(&(depth, _i_tri)) = hits.first() else {
183                continue;
184            };
185            let pos = del_geo_core::vec3::axpy(depth, &ray_dir, &ray_org);
186            let ndc =
187                del_geo_core::mat4_col_major::transform_homogeneous(&transform_world2ndc, &pos)
188                    .unwrap();
189            let depth_ndc = (ndc[2] + 1f32) * 0.5f32;
190            pix2depth[ih * width + iw] = depth_ndc;
191        }
192    }
193}
194
195#[test]
196fn test_depthmap() {
197    for i_case in 0..3 {
198        let (tri2vtx, vtx2xyz) = match i_case {
199            0 => {
200                let vtx2xyz_polyline = crate::polyline3::helix(100, 0.05, 0.7, 0.4);
201                use std::f32::consts::PI;
202                let rot_y = del_geo_core::mat4_col_major::from_rot_y(PI * 0.51);
203                let rot_x = del_geo_core::mat4_col_major::from_rot_x(-PI * 0.25);
204                let transform = del_geo_core::mat4_col_major::mult_mat_col_major(&rot_x, &rot_y);
205                let vtx2xyz_polyline =
206                    crate::vtx2xyz::transform_homogeneous(&vtx2xyz_polyline, &transform);
207                crate::polyline3::to_trimesh3_capsule(&vtx2xyz_polyline, 32, 32, 0.05)
208            }
209            1 => {
210                let vtx2xyz_polyline = del_geo_core::bezier_cubic::sample_uniform_param(
211                    100,
212                    &[0.9, 0.0, -0.2],
213                    &[-3.0, 0.9, -0.2],
214                    &[0.9, -3.0, 0.2],
215                    &[0.0, 0.9, 0.2],
216                    true,
217                    true,
218                );
219                use slice_of_array::SliceFlatExt;
220                let vtx2xyz_polyline = vtx2xyz_polyline.flat().to_owned();
221                crate::polyline3::to_trimesh3_capsule(&vtx2xyz_polyline, 32, 32, 0.05)
222            }
223            2 => {
224                let (tri2vtx, vtx2xyz) = crate::trimesh3_primitive::torus_zup(0.8, 0.05, 32, 32);
225                let transform =
226                    del_geo_core::mat4_col_major::from_rot_x(std::f32::consts::PI / 12.0);
227                let vtx2xyz = crate::vtx2xyz::transform_homogeneous(&vtx2xyz, &transform);
228                (tri2vtx, vtx2xyz)
229            }
230            _ => unreachable!(),
231        };
232        crate::io_wavefront_obj::save_tri2vtx_vtx2xyz(
233            format!("../target/trimesh3_raycast_mesh{i_case}.obj"),
234            &tri2vtx,
235            &vtx2xyz,
236            3,
237        )
238        .unwrap();
239        let aabb3 = crate::vtx2xyz::aabb3(&vtx2xyz, 0.);
240        dbg!(aabb3);
241        let bvhnodes = crate::bvhnodes_morton::from_triangle_mesh(&tri2vtx, &vtx2xyz, 3);
242        let bvhnode2aabb = crate::bvhnode2aabb3::from_uniform_mesh_with_bvh(
243            0, &bvhnodes, &tri2vtx, 3, &vtx2xyz, None,
244        );
245        let img_shape = (300, 300);
246        let mut pix2depth = vec![0f32; img_shape.0 * img_shape.1];
247        let transform_ndc2world = del_geo_core::mat4_col_major::from_identity();
248        dbg!(&transform_ndc2world);
249        render_depth_bvh(
250            img_shape,
251            &mut pix2depth,
252            &transform_ndc2world,
253            &tri2vtx,
254            &vtx2xyz,
255            &bvhnodes,
256            &bvhnode2aabb,
257        );
258        pix2depth.iter_mut().for_each(|v| *v = (*v) + 0.0);
259        del_canvas::write_png_from_float_image(
260            format!("../target/trimesh3_raycast_depth_{i_case}.png"),
261            img_shape,
262            1,
263            &pix2depth,
264        )
265        .unwrap();
266        let (quad2vtx, vtx2xyz) =
267            crate::grid2::to_quadmesh3_hightmap(img_shape, &pix2depth, 1.0 / img_shape.0 as f32);
268        crate::io_wavefront_obj::save_quad2vtx_vtx2xyz(
269            format!("../target/trimesh3_raycast_hightmap_{i_case}.obj"),
270            &quad2vtx,
271            &vtx2xyz,
272            3,
273        )
274        .unwrap();
275    }
276}