Added depth rendering function

README.md

@@ -125,8 +125,8 @@ python autotune.py -g '<space delimited GPU ids>' tasks/eval_tnt.json
 
 First make sure you have colmap installed. Then
 
-(in opt/)
-`bash scripts/proc_colmap.sh <img_dir>`
+(in opt/scripts)
+`bash proc_colmap.sh <img_dir>`
 
 Where `<img_dir>` should be a directory directly containing png/jpg images from a 
 normal perspective camera.
@@ -136,14 +136,15 @@ For custom datasets we adopt a data format similar to that in NSVF
 You should be able to use this dataset directly afterwards. The format will be auto-detected.
 
 To view the data use:
-`python scripts/view_data.py <img_dir>`
+`python view_data.py <img_dir>`
 
 This should launch a server at localhost:8889
 
 
-You may need to tune the TV. For forward-facing scenes, often making the TV weights 10x
+You may need to tune the TV.
+ For forward-facing scenes, often making the TV weights 10x
 higher is helpful (`configs/llff_hitv.json`).
-For the real lego scene I used the config `configs/custom.json`.
+For the real lego scene I used the config `configs/custom.json`, or you can also try `configs/custom_alt.json` which has some minor differences.
 
 ## Random tip: how to make pip install faster for native extensions
 

environment.yml

@@ -23,6 +23,7 @@ dependencies:
       - matplotlib
       - scipy>=1.6.0
   - pytorch
+  - torchvision
   - cudatoolkit
   - tqdm
 

opt/opt.py

@@ -156,6 +156,11 @@
                    default=0.1,
                    help='initialization sigma (for BG)')
 
+# Extra logging
+group.add_argument('--log_mse_image', action='store_true', default=False)
+group.add_argument('--log_depth_map', action='store_true', default=False)
+group.add_argument('--log_depth_map_use_thresh', type=float, default=None,
+        help="If specified, uses the Dex-neRF version of depth with given thresh; else returns expected term")
 
 
 group = parser.add_argument_group("misc experiments")
@@ -238,6 +243,7 @@
 
 group.add_argument('--nosphereinit', action='store_true', default=False,
                      help='do not start with sphere bounds (please do not use for 360)')
+
 args = parser.parse_args()
 config_util.maybe_merge_config_file(args)
 
@@ -405,9 +411,18 @@ def eval_step():
                     img_pred.clamp_max_(1.0)
                     summary_writer.add_image(f'test/image_{img_id:04d}',
                             img_pred, global_step=gstep_id_base, dataformats='HWC')
-                    mse_img = all_mses / all_mses.max()
-                    summary_writer.add_image(f'test/mse_map_{img_id:04d}',
-                            mse_img, global_step=gstep_id_base, dataformats='HWC')
+                    if args.log_mse_image:
+                        mse_img = all_mses / all_mses.max()
+                        summary_writer.add_image(f'test/mse_map_{img_id:04d}',
+                                mse_img, global_step=gstep_id_base, dataformats='HWC')
+                    if args.log_depth_map:
+                        depth_img = grid.volume_render_depth_image(cam,
+                                    15.0 if args.log_depth_map_use_thresh else None
+                                )
+                        depth_img = viridis_cmap(depth_img.cpu())
+                        summary_writer.add_image(f'test/depth_map_{img_id:04d}',
+                                depth_img,
+                                global_step=gstep_id_base, dataformats='HWC')
 
                 rgb_pred_test = rgb_gt_test = None
                 mse_num : float = all_mses.mean().item()

svox2/csrc/render_lerp_kernel_cuvol.cu

@@ -119,6 +119,111 @@ __device__ __inline__ void trace_ray_cuvol(
     }
 }
 
+__device__ __inline__ void trace_ray_expected_term(
+        const PackedSparseGridSpec& __restrict__ grid,
+        SingleRaySpec& __restrict__ ray,
+        const RenderOptions& __restrict__ opt,
+        float* __restrict__ out) {
+    if (ray.tmin > ray.tmax) {
+        *out = 0.f;
+        return;
+    }
+
+    float t = ray.tmin;
+    float outv = 0.f;
+
+    float log_transmit = 0.f;
+    // printf("tmin %f, tmax %f \n", ray.tmin, ray.tmax);
+
+    while (t <= ray.tmax) {
+#pragma unroll 3
+        for (int j = 0; j < 3; ++j) {
+            ray.pos[j] = fmaf(t, ray.dir[j], ray.origin[j]);
+            ray.pos[j] = min(max(ray.pos[j], 0.f), grid.size[j] - 1.f);
+            ray.l[j] = min(static_cast<int32_t>(ray.pos[j]), grid.size[j] - 2);
+            ray.pos[j] -= static_cast<float>(ray.l[j]);
+        }
+
+        const float skip = compute_skip_dist(ray,
+                       grid.links, grid.stride_x,
+                       grid.size[2], 0);
+
+        if (skip >= opt.step_size) {
+            // For consistency, we skip the by step size
+            t += ceilf(skip / opt.step_size) * opt.step_size;
+            continue;
+        }
+        float sigma = trilerp_cuvol_one(
+                grid.links, grid.density_data,
+                grid.stride_x,
+                grid.size[2],
+                1,
+                ray.l, ray.pos,
+                0);
+        if (sigma > opt.sigma_thresh) {
+            const float pcnt = ray.world_step * sigma;
+            const float weight = _EXP(log_transmit) * (1.f - _EXP(-pcnt));
+            log_transmit -= pcnt;
+
+            outv += weight * (t / opt.step_size) * ray.world_step;
+            if (_EXP(log_transmit) < opt.stop_thresh) {
+                log_transmit = -1e3f;
+                break;
+            }
+        }
+        t += opt.step_size;
+    }
+    *out = outv;
+}
+
+// From Dex-NeRF
+__device__ __inline__ void trace_ray_sigma_thresh(
+        const PackedSparseGridSpec& __restrict__ grid,
+        SingleRaySpec& __restrict__ ray,
+        const RenderOptions& __restrict__ opt,
+        float sigma_thresh,
+        float* __restrict__ out) {
+    if (ray.tmin > ray.tmax) {
+        *out = 0.f;
+        return;
+    }
+
+    float t = ray.tmin;
+    *out = 0.f;
+
+    while (t <= ray.tmax) {
+#pragma unroll 3
+        for (int j = 0; j < 3; ++j) {
+            ray.pos[j] = fmaf(t, ray.dir[j], ray.origin[j]);
+            ray.pos[j] = min(max(ray.pos[j], 0.f), grid.size[j] - 1.f);
+            ray.l[j] = min(static_cast<int32_t>(ray.pos[j]), grid.size[j] - 2);
+            ray.pos[j] -= static_cast<float>(ray.l[j]);
+        }
+
+        const float skip = compute_skip_dist(ray,
+                       grid.links, grid.stride_x,
+                       grid.size[2], 0);
+
+        if (skip >= opt.step_size) {
+            // For consistency, we skip the by step size
+            t += ceilf(skip / opt.step_size) * opt.step_size;
+            continue;
+        }
+        float sigma = trilerp_cuvol_one(
+                grid.links, grid.density_data,
+                grid.stride_x,
+                grid.size[2],
+                1,
+                ray.l, ray.pos,
+                0);
+        if (sigma > sigma_thresh) {
+            *out = (t / opt.step_size) * ray.world_step;
+            break;
+        }
+        t += opt.step_size;
+    }
+}
+
 __device__ __inline__ void trace_ray_cuvol_backward(
         const PackedSparseGridSpec& __restrict__ grid,
         const float* __restrict__ grad_output,
@@ -753,6 +858,48 @@ __global__ void render_background_backward_kernel(
         grads);
 }
 
+__launch_bounds__(TRACE_RAY_CUDA_THREADS, MIN_BLOCKS_PER_SM)
+__global__ void render_ray_expected_term_kernel(
+        PackedSparseGridSpec grid,
+        PackedRaysSpec rays,
+        RenderOptions opt,
+        float* __restrict__ out) {
+        // const PackedSparseGridSpec& __restrict__ grid,
+        // SingleRaySpec& __restrict__ ray,
+        // const RenderOptions& __restrict__ opt,
+        // float* __restrict__ out) {
+    CUDA_GET_THREAD_ID(ray_id, rays.origins.size(0));
+    SingleRaySpec ray_spec(rays.origins[ray_id].data(), rays.dirs[ray_id].data());
+    ray_find_bounds(ray_spec, grid, opt, ray_id);
+    trace_ray_expected_term(
+        grid,
+        ray_spec,
+        opt,
+        out + ray_id);
+}
+
+__launch_bounds__(TRACE_RAY_CUDA_THREADS, MIN_BLOCKS_PER_SM)
+__global__ void render_ray_sigma_thresh_kernel(
+        PackedSparseGridSpec grid,
+        PackedRaysSpec rays,
+        RenderOptions opt,
+        float sigma_thresh,
+        float* __restrict__ out) {
+        // const PackedSparseGridSpec& __restrict__ grid,
+        // SingleRaySpec& __restrict__ ray,
+        // const RenderOptions& __restrict__ opt,
+        // float* __restrict__ out) {
+    CUDA_GET_THREAD_ID(ray_id, rays.origins.size(0));
+    SingleRaySpec ray_spec(rays.origins[ray_id].data(), rays.dirs[ray_id].data());
+    ray_find_bounds(ray_spec, grid, opt, ray_id);
+    trace_ray_sigma_thresh(
+        grid,
+        ray_spec,
+        opt,
+        sigma_thresh,
+        out + ray_id);
+}
+
 }  // namespace device
 
 torch::Tensor _get_empty_1d(const torch::Tensor& origins) {
@@ -1002,3 +1149,46 @@ void volume_render_cuvol_fused(
 
     CUDA_CHECK_ERRORS;
 }
+
+torch::Tensor volume_render_expected_term(SparseGridSpec& grid,
+        RaysSpec& rays, RenderOptions& opt) {
+    auto options =
+        torch::TensorOptions()
+        .dtype(rays.origins.dtype())
+        .layout(torch::kStrided)
+        .device(rays.origins.device())
+        .requires_grad(false);
+    torch::Tensor results = torch::empty({rays.origins.size(0)}, options);
+    const auto Q = rays.origins.size(0);
+    const int blocks = CUDA_N_BLOCKS_NEEDED(Q, TRACE_RAY_CUDA_THREADS);
+    device::render_ray_expected_term_kernel<<<blocks, TRACE_RAY_CUDA_THREADS>>>(
+            grid,
+            rays,
+            opt,
+            results.data_ptr<float>()
+        );
+    return results;
+}
+
+torch::Tensor volume_render_sigma_thresh(SparseGridSpec& grid,
+        RaysSpec& rays,
+        RenderOptions& opt,
+        float sigma_thresh) {
+    auto options =
+        torch::TensorOptions()
+        .dtype(rays.origins.dtype())
+        .layout(torch::kStrided)
+        .device(rays.origins.device())
+        .requires_grad(false);
+    torch::Tensor results = torch::empty({rays.origins.size(0)}, options);
+    const auto Q = rays.origins.size(0);
+    const int blocks = CUDA_N_BLOCKS_NEEDED(Q, TRACE_RAY_CUDA_THREADS);
+    device::render_ray_sigma_thresh_kernel<<<blocks, TRACE_RAY_CUDA_THREADS>>>(
+            grid,
+            rays,
+            opt,
+            sigma_thresh,
+            results.data_ptr<float>()
+        );
+    return results;
+}

svox2/csrc/svox2.cpp

@@ -13,20 +13,29 @@ std::tuple<torch::Tensor, torch::Tensor> sample_grid(SparseGridSpec &, Tensor,
 void sample_grid_backward(SparseGridSpec &, Tensor, Tensor, Tensor, Tensor,
                           Tensor, bool);
 
+// ** NeRF rendering formula (trilerp)
 Tensor volume_render_cuvol(SparseGridSpec &, RaysSpec &, RenderOptions &);
 Tensor volume_render_cuvol_image(SparseGridSpec &, CameraSpec &,
                                  RenderOptions &);
 void volume_render_cuvol_backward(SparseGridSpec &, RaysSpec &, RenderOptions &,
                                   Tensor, Tensor, GridOutputGrads &);
 void volume_render_cuvol_fused(SparseGridSpec &, RaysSpec &, RenderOptions &,
                                Tensor, float, float, Tensor, GridOutputGrads &);
-
+// Expected termination (depth) rendering
+torch::Tensor volume_render_expected_term(SparseGridSpec &, RaysSpec &,
+                                          RenderOptions &);
+// Depth rendering based on sigma-threshold as in Dex-NeRF
+torch::Tensor volume_render_sigma_thresh(SparseGridSpec &, RaysSpec &,
+                                         RenderOptions &, float);
+
+// ** NV rendering formula (trilerp)
 Tensor volume_render_nvol(SparseGridSpec &, RaysSpec &, RenderOptions &);
 void volume_render_nvol_backward(SparseGridSpec &, RaysSpec &, RenderOptions &,
                                  Tensor, Tensor, GridOutputGrads &);
 void volume_render_nvol_fused(SparseGridSpec &, RaysSpec &, RenderOptions &,
                               Tensor, float, float, Tensor, GridOutputGrads &);
 
+// ** NeRF rendering formula (nearest-neighbor, infinitely many steps)
 Tensor volume_render_svox1(SparseGridSpec &, RaysSpec &, RenderOptions &);
 void volume_render_svox1_backward(SparseGridSpec &, RaysSpec &, RenderOptions &,
                                   Tensor, Tensor, GridOutputGrads &);
@@ -70,6 +79,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   _REG_FUNC(volume_render_cuvol_image);
   _REG_FUNC(volume_render_cuvol_backward);
   _REG_FUNC(volume_render_cuvol_fused);
+  _REG_FUNC(volume_render_expected_term);
+  _REG_FUNC(volume_render_sigma_thresh);
 
   _REG_FUNC(volume_render_nvol);
   _REG_FUNC(volume_render_nvol_backward);