Skip to content

Commit

Permalink
Merge pull request #173 from erizmr/fix_render_test
Browse files Browse the repository at this point in the history 
[taichi] Fix missing volume render test
  • Loading branch information
@ashawkey
ashawkey committed Mar 21, 2023
2 parents ab98abc + c8893b6 commit 3fac34b
Show file tree
Hide file tree
Showing 4 changed files with 86 additions and 38 deletions.
10 changes: 5 additions & 5 deletions taichi_modules/intersection.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,11 @@

@ti.kernel
def simple_ray_aabb_intersec_taichi_forward(
hits_t: ti.types.ndarray(field_dim=2),
rays_o: ti.types.ndarray(field_dim=2),
rays_d: ti.types.ndarray(field_dim=2),
centers: ti.types.ndarray(field_dim=2),
half_sizes: ti.types.ndarray(field_dim=2)):
hits_t: ti.types.ndarray(ndim=2),
rays_o: ti.types.ndarray(ndim=2),
rays_d: ti.types.ndarray(ndim=2),
centers: ti.types.ndarray(ndim=2),
half_sizes: ti.types.ndarray(ndim=2)):

for r in ti.ndrange(hits_t.shape[0]):
ray_o = vec3([rays_o[r, 0], rays_o[r, 1], rays_o[r, 2]])
Expand Down
54 changes: 27 additions & 27 deletions taichi_modules/ray_march.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,17 @@


@ti.kernel
def raymarching_train(rays_o: ti.types.ndarray(field_dim=2),
rays_d: ti.types.ndarray(field_dim=2),
hits_t: ti.types.ndarray(field_dim=2),
density_bitfield: ti.types.ndarray(field_dim=1),
noise: ti.types.ndarray(field_dim=1),
counter: ti.types.ndarray(field_dim=1),
rays_a: ti.types.ndarray(field_dim=2),
xyzs: ti.types.ndarray(field_dim=2),
dirs: ti.types.ndarray(field_dim=2),
deltas: ti.types.ndarray(field_dim=1),
ts: ti.types.ndarray(field_dim=1), cascades: int,
def raymarching_train(rays_o: ti.types.ndarray(ndim=2),
rays_d: ti.types.ndarray(ndim=2),
hits_t: ti.types.ndarray(ndim=2),
density_bitfield: ti.types.ndarray(ndim=1),
noise: ti.types.ndarray(ndim=1),
counter: ti.types.ndarray(ndim=1),
rays_a: ti.types.ndarray(ndim=2),
xyzs: ti.types.ndarray(ndim=2),
dirs: ti.types.ndarray(ndim=2),
deltas: ti.types.ndarray(ndim=1),
ts: ti.types.ndarray(ndim=1), cascades: int,
grid_size: int, scale: float, exp_step_factor: float,
max_samples: float):

Expand Down Expand Up @@ -126,12 +126,12 @@ def raymarching_train(rays_o: ti.types.ndarray(field_dim=2),


@ti.kernel
def raymarching_train_backword(segments: ti.types.ndarray(field_dim=2),
ts: ti.types.ndarray(field_dim=1),
dL_drays_o: ti.types.ndarray(field_dim=2),
dL_drays_d: ti.types.ndarray(field_dim=2),
dL_dxyzs: ti.types.ndarray(field_dim=2),
dL_ddirs: ti.types.ndarray(field_dim=2)):
def raymarching_train_backword(segments: ti.types.ndarray(ndim=2),
ts: ti.types.ndarray(ndim=1),
dL_drays_o: ti.types.ndarray(ndim=2),
dL_drays_d: ti.types.ndarray(ndim=2),
dL_dxyzs: ti.types.ndarray(ndim=2),
dL_ddirs: ti.types.ndarray(ndim=2)):

for s in segments:
index = segments[s]
Expand Down Expand Up @@ -228,22 +228,22 @@ def forward(self, rays_o, rays_d, hits_t, density_bitfield, cascades,

@ti.kernel
def raymarching_test_kernel(
rays_o: ti.types.ndarray(field_dim=2),
rays_d: ti.types.ndarray(field_dim=2),
hits_t: ti.types.ndarray(field_dim=2),
alive_indices: ti.types.ndarray(field_dim=1),
density_bitfield: ti.types.ndarray(field_dim=1),
rays_o: ti.types.ndarray(ndim=2),
rays_d: ti.types.ndarray(ndim=2),
hits_t: ti.types.ndarray(ndim=2),
alive_indices: ti.types.ndarray(ndim=1),
density_bitfield: ti.types.ndarray(ndim=1),
cascades: int,
grid_size: int,
scale: float,
exp_step_factor: float,
N_samples: int,
max_samples: int,
xyzs: ti.types.ndarray(field_dim=2),
dirs: ti.types.ndarray(field_dim=2),
deltas: ti.types.ndarray(field_dim=1),
ts: ti.types.ndarray(field_dim=1),
N_eff_samples: ti.types.ndarray(field_dim=1),
xyzs: ti.types.ndarray(ndim=2),
dirs: ti.types.ndarray(ndim=2),
deltas: ti.types.ndarray(ndim=1),
ts: ti.types.ndarray(ndim=1),
N_eff_samples: ti.types.ndarray(ndim=1),
):

for n in alive_indices:
Expand Down
12 changes: 6 additions & 6 deletions taichi_modules/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,8 +86,8 @@ def __morton3D_invert(x):


@ti.kernel
def morton3D_invert_kernel(indices: ti.types.ndarray(field_dim=1),
coords: ti.types.ndarray(field_dim=2)):
def morton3D_invert_kernel(indices: ti.types.ndarray(ndim=1),
coords: ti.types.ndarray(ndim=2)):
for i in indices:
ind = ti.uint32(indices[i])
coords[i, 0] = __morton3D_invert(ind >> 0)
Expand All @@ -106,8 +106,8 @@ def morton3D_invert(indices):


@ti.kernel
def morton3D_kernel(xyzs: ti.types.ndarray(field_dim=2),
indices: ti.types.ndarray(field_dim=1)):
def morton3D_kernel(xyzs: ti.types.ndarray(ndim=2),
indices: ti.types.ndarray(ndim=1)):
for s in indices:
xyz = uvec3([xyzs[s, 0], xyzs[s, 1], xyzs[s, 2]])
indices[s] = ti.cast(__morton3D(xyz), ti.int32)
Expand All @@ -123,9 +123,9 @@ def morton3D(coords1):


@ti.kernel
def packbits(density_grid: ti.types.ndarray(field_dim=1),
def packbits(density_grid: ti.types.ndarray(ndim=1),
density_threshold: float,
density_bitfield: ti.types.ndarray(field_dim=1)):
density_bitfield: ti.types.ndarray(ndim=1)):

for n in density_bitfield:
bits = ti.uint8(0)
Expand Down
48 changes: 48 additions & 0 deletions taichi_modules/volume_render_test.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import taichi as ti


@ti.kernel
def composite_test(
sigmas: ti.types.ndarray(ndim=2), rgbs: ti.types.ndarray(ndim=3),
deltas: ti.types.ndarray(ndim=2), ts: ti.types.ndarray(ndim=2),
hits_t: ti.types.ndarray(ndim=2),
alive_indices: ti.types.ndarray(ndim=1), T_threshold: float,
N_eff_samples: ti.types.ndarray(ndim=1),
opacity: ti.types.ndarray(ndim=1),
depth: ti.types.ndarray(ndim=1), rgb: ti.types.ndarray(ndim=2)):

for n in alive_indices:
samples = N_eff_samples[n]
if samples == 0:
alive_indices[n] = -1
else:
r = alive_indices[n]

T = 1 - opacity[r]

rgb_temp_0 = 0.0
rgb_temp_1 = 0.0
rgb_temp_2 = 0.0
depth_temp = 0.0
opacity_temp = 0.0

for s in range(samples):
a = 1.0 - ti.exp(-sigmas[n, s] * deltas[n, s])
w = a * T

rgb_temp_0 += w * rgbs[n, s, 0]
rgb_temp_1 += w * rgbs[n, s, 1]
rgb_temp_2 += w * rgbs[n, s, 2]
depth[r] += w * ts[n, s]
opacity[r] += w
T *= 1.0 - a

if T <= T_threshold:
alive_indices[n] = -1
break

rgb[r, 0] += rgb_temp_0
rgb[r, 1] += rgb_temp_1
rgb[r, 2] += rgb_temp_2
depth[r] += depth_temp
opacity[r] += opacity_temp

0 comments on commit 3fac34b

Please sign in to comment.