robot0321 / svox_custom

customizing svox library, especially for the rendering process

SVOX custom: customize the rendering method following NeRF and NeuS

Folked from the repository https://github.com/sxyu/svox

• I customize the trace_ray function in svox/csrc/rt_kernel.cu file.

  svox_custom/svox/csrc/rt_kernel.cu

  Lines 222 to 321 in 4d923d5

  	template <typename scalar_t>
  	__device__ __inline__ void trace_ray(
  	PackedTreeSpec<scalar_t>& __restrict__ tree,
  	SingleRaySpec<scalar_t> ray,
  	RenderOptions& __restrict__ opt,
  	torch::TensorAccessor<scalar_t, 1, torch::RestrictPtrTraits, int32_t> out) {
  	const scalar_t delta_scale = _get_delta_scale(tree.scaling, ray.dir);
  	scalar_t tmin, tmax;
  	scalar_t invdir[3];
  	const int tree_N = tree.child.size(1);
  	const int data_dim = tree.data.size(4);
  	const int out_data_dim = out.size(0);
  	#pragma unroll
  	for (int i = 0; i < 3; ++i) {
  	invdir[i] = 1.0 / (ray.dir[i] + 1e-9);
  	}
  	_dda_unit(ray.origin, invdir, &tmin, &tmax);
  	if (tmax < 0 || tmin > tmax) {
  	// Ray doesn't hit box
  	for (int j = 0; j < out_data_dim; ++j) {
  	out[j] = opt.background_brightness;
  	}
  	return;
  	} else {
  	for (int j = 0; j < out_data_dim; ++j) {
  	out[j] = 0.f;
  	}
  	scalar_t pos[3];
  	scalar_t basis_fn[25];
  	maybe_precalc_basis<scalar_t>(opt.format, opt.basis_dim,
  	tree.extra_data, ray.vdir, basis_fn);
  	scalar_t light_intensity = 1.f;
  	scalar_t t = tmin;
  	scalar_t cube_sz;
  	const scalar_t d_rgb_pad = 1 + 2 * opt.rgb_padding;
  	while (t < tmax) {
  	for (int j = 0; j < 3; ++j) {
  	pos[j] = ray.origin[j] + t * ray.dir[j];
  	}
  	int64_t node_id;
  	scalar_t* tree_val = query_single_from_root<scalar_t>(tree.data, tree.child,
  	pos, &cube_sz, tree.weight_accum != nullptr ? &node_id : nullptr);
  	scalar_t att;
  	scalar_t subcube_tmin, subcube_tmax;
  	_dda_unit(pos, invdir, &subcube_tmin, &subcube_tmax);
  	const scalar_t t_subcube = (subcube_tmax - subcube_tmin) / cube_sz;
  	const scalar_t delta_t = t_subcube + opt.step_size;
  	scalar_t sigma = tree_val[data_dim - 1];
  	// if (opt.density_softplus) sigma = _SOFTPLUS_M1(sigma);
  	if (sigma > opt.sigma_thresh) {
  	att = expf(-delta_t * delta_scale * sigma);
  	const scalar_t weight = light_intensity * (1.f - att);
  	if (opt.format != FORMAT_RGBA) {
  	for (int t = 0; t < out_data_dim; ++ t) {
  	int off = t * opt.basis_dim;
  	scalar_t tmp = 0.0;
  	for (int i = opt.min_comp; i <= opt.max_comp; ++i) {
  	tmp += basis_fn[i] * tree_val[off + i];
  	}
  	out[t] += weight * (_SIGMOID(tmp) * d_rgb_pad - opt.rgb_padding);
  	}
  	} else {
  	for (int j = 0; j < out_data_dim; ++j) {
  	// out[j] += weight * (_SIGMOID(tree_val[j]) * d_rgb_pad - opt.rgb_padding);
  	out[j] += weight * tree_val[j];
  	}
  	}
  	light_intensity *= att;
  	if (tree.weight_accum != nullptr) {
  	if (tree.weight_accum_max) {
  	atomicMax(&tree.weight_accum[node_id], weight);
  	} else {
  	atomicAdd(&tree.weight_accum[node_id], weight);
  	}
  	}
  	if (light_intensity <= opt.stop_thresh) {
  	// Full opacity, stop
  	scalar_t scale = 1.0 / (1.0 - light_intensity);
  	for (int j = 0; j < out_data_dim; ++j) {
  	out[j] *= scale;
  	}
  	return;
  	}
  	}
  	t += delta_t;
  	}
  	for (int j = 0; j < out_data_dim; ++j) {
  	out[j] += light_intensity * opt.background_brightness;
  	}
  	}
  	}

• This code is not complete since I only fixed the forward function (only for customized rendering, not training)

• Whenever you change the library code, do uninstall and re-install. (There is dev-mode but I did not use it)

• svox/csrc/rt_kernel_neus.cu is for NeuS and svox/csrc/rt_kernel.cu is for NeRF

• Note that the file named svox/csrc/rt_kernel.cu is used in this library. If you want to change, swap the name of files.

Install

python setup.py install

Uninstall

pip uninstall svox

Documentation

Please see https://svox.readthedocs.io