Files
Magic123/taichi_modules/volume_render_test.py
Guocheng Qian 13e18567fa first commit
2023-08-02 19:51:43 -07:00

49 lines
1.4 KiB
Python

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