Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[taichi] Fix missing volume render test #173

Merged
merged 1 commit into from
Mar 21, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
[taichi] Fix missing volume render test
  • Loading branch information
erizmr committed Mar 21, 2023
commit c8893b6a988eec3df3989e8606cd0c85bd57892f
10 changes: 5 additions & 5 deletions taichi_modules/intersection.py
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
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
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
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