Added timing to autotune; (probably) fixed NV renderer kernels; auto …

…scene scale for T&T

opt/autotune.py

@@ -69,8 +69,8 @@ def run_exp(env, eval_mode:bool, enable_render:bool, train_dir, data_dir, config
         with open(log_file_path, 'w') as f:
             f.write(opt_ret)
         secs = (end_time - start_time).total_seconds()
-        timings_file = open(os.path.join(train_dir, 'timings.txt'), 'a')
-        timings_file.write(f"{secs} s = {secs / 60} m = {secs / (60 * 60)} h\n")
+        timings_file = open(os.path.join(train_dir, 'time_mins.txt'), 'a')
+        timings_file.write(f"{secs / 60}\n")
 
     if eval_mode:
         eval_base_cmd = [
@@ -255,16 +255,18 @@ def recursive_replace(data, variables):
     if args.eval:
         print('Done')
         with open(leaderboard_path, 'w') as leaderboard_file:
-            lines = [f'dir\tPSNR\tSSIM\tLPIPS\n']
+            lines = [f'dir\tPSNR\tSSIM\tLPIPS\nminutes\n']
             all_tasks = sorted(all_tasks, key=lambda task:task['train_dir'])
             all_psnr = []
             all_ssim = []
             all_lpips = []
+            all_times = []
             for task in all_tasks:
                 train_dir = task['train_dir']
                 psnr_file_path = path.join(train_dir, 'test_renders', 'psnr.txt')
                 ssim_file_path = path.join(train_dir, 'test_renders', 'ssim.txt')
                 lpips_file_path = path.join(train_dir, 'test_renders', 'lpips.txt')
+                time_file_path = path.join(train_dir, 'time_mins.txt')
 
                 if path.isfile(psnr_file_path):
                     with open(psnr_file_path, 'r') as f:
@@ -287,7 +289,14 @@ def recursive_replace(data, variables):
                     lpips_txt = f'{lpips:.10f}'
                 else:
                     lpips_txt = 'ERR'
-                line = f'{path.basename(train_dir.rstrip("/"))}\t{psnr_txt}\t{ssim_txt}\t{lpips_txt}\n'
+                if path.isfile(time_file_path):
+                    with open(time_file_path, 'r') as f:
+                        time_mins = float(f.read())
+                    all_times.append(time_mins)
+                    time_txt = f'{time_mins:.10f}'
+                else:
+                    time_txt = 'ERR'
+                line = f'{path.basename(train_dir.rstrip("/"))}\t{psnr_txt}\t{ssim_txt}\t{lpips_txt}\t{time_txt}\n'
                 lines.append(line)
             lines.append('---------\n')
             if len(all_psnr):
@@ -296,6 +305,8 @@ def recursive_replace(data, variables):
                 lines.append('Average SSIM: ' + str(sum(all_ssim) / len(all_ssim)) + '\n')
             if len(all_lpips):
                 lines.append('Average LPIPS: ' + str(sum(all_lpips) / len(all_lpips)) + '\n')
+            if len(all_times):
+                lines.append('Average Time (mins): ' + str(sum(all_times) / len(all_times)) + '\n')
             leaderboard_file.writelines(lines)
 
     else:

opt/configs/syn_nv.json

@@ -0,0 +1,11 @@
+{
+    "reso": "[[256, 256, 256]]",
+    "upsamp_every": 38400,
+    "renderer_backend": "nvol",
+    "lr_sigma": 3e-1,
+    "lr_sigma_final": 3e-5,
+    "lr_sigma_delay_steps": 0,
+    "lr_sh": 1e-2,
+    "lambda_tv": 0.0,
+    "lambda_tv_sh": 0.0
+}

opt/configs/tnt.json

@@ -2,7 +2,7 @@
     "reso": "[[128, 128, 128], [256, 256, 256], [512, 512, 512], [640, 640, 640]]",
     "n_iters": 102400,
     "background_nlayers": 64,
-    "background_reso": 768,
+    "background_reso": 1024,
     "upsamp_every": 25600,
     "lr_sigma": 3e1,
     "lr_sh": 1e-2,
@@ -11,8 +11,8 @@
     "weight_thresh": 1.28,
     "lambda_tv": 5e-5,
     "lambda_tv_sh": 5e-3,
-    "lambda_tv_background_sigma": 0.01,
-    "lambda_tv_background_color": 0.01,
+    "lambda_tv_background_sigma": 1e-3,
+    "lambda_tv_background_color": 1e-3,
     "lambda_beta": 1e-5,
     "lambda_sparsity": 1e-11,
     "background_brightness": 1.0,

opt/opt.py

@@ -165,7 +165,7 @@
                     default=5.0,
                    help='Upsample sigma threshold')
 group.add_argument('--background_density_thresh', type=float,
-                    default=1.0,
+                    default=0.25,
                    help='Background sigma threshold for sparsification')
 group.add_argument('--max_grid_elements', type=int,
                     default=44_000_000,

opt/render_imgs.py

@@ -3,6 +3,7 @@
 
 import torch
 import svox2
+import svox2.utils
 import math
 import argparse
 import numpy as np
@@ -121,6 +122,9 @@
         render_dir += '_nobg'
     if args.nofg:
         grid.density_data.data[:] = 0.0
+        #  grid.sh_data.data[..., 0] = 1.0 / svox2.utils.SH_C0
+        #  grid.sh_data.data[..., 9] = 1.0 / svox2.utils.SH_C0
+        #  grid.sh_data.data[..., 18] = 1.0 / svox2.utils.SH_C0
         render_dir += '_nofg'
 
 config_util.setup_render_opts(grid.opt, args)

opt/tasks/eval_tnt.json

@@ -1,23 +1,13 @@
 {
     "eval": true,
     "data_root": "/home/sxyu/data/TanksAndTempleBG",
-    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/tnt_spars_fasttv",
-    "config": "configs/tnt.json",
+    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/tnt_equirectlin_fasttv_autoscale",
+    "variables": {
+        "scene": ["Train", "M60", "Truck", "Playground"]
+    },
     "tasks": [{
-        "train_dir": "Train",
-        "data_dir": "Train",
-        "common_flags": ["--scene_scale", "2.0"]
-    }, {
-        "train_dir": "M60",
-        "data_dir": "M60",
-        "common_flags": ["--scene_scale", "2.0"]
-    }, {
-        "train_dir": "Truck",
-        "data_dir": "Truck",
-        "common_flags": ["--scene_scale", "1.6"]
-    }, {
-        "train_dir": "Playground",
-        "data_dir": "Playground",
-        "common_flags": ["--scene_scale", "1.7"]
+        "train_dir": "{scene}",
+        "data_dir": "{scene}",
+        "config": "configs/tnt.json"
     }]
 }

opt/util/config_util.py

@@ -39,12 +39,16 @@ def define_common_args(parser : argparse.ArgumentParser):
                          help="LLFF holdout every")
     group.add_argument('--normalize_by_bbox',
                          type=bool,
-                         default=True,
-                         help="Normalize by bounding box in bbox.txt, if available (NSVF dataset only)")
+                         default=False,
+                         help="Normalize by bounding box in bbox.txt, if available (NSVF dataset only); precedes normalize_by_camera")
     group.add_argument('--data_bbox_scale',
                          type=float,
                          default=1.2,
                          help="Data bbox scaling (NSVF dataset only)")
+    group.add_argument('--normalize_by_camera',
+                         type=bool,
+                         default=True,
+                         help="Normalize using cameras, assuming a 360 capture (NSVF dataset only); only used if not normalize_by_bbox")
     group.add_argument('--perm', action='store_true', default=False,
                          help='sample by permutation of rays (true epoch) instead of '
                               'uniformly random rays')
@@ -89,6 +93,9 @@ def define_common_args(parser : argparse.ArgumentParser):
 
 
 def build_data_options(args):
+    """
+    Arguments to pass as kwargs to the dataset constructor
+    """
     return {
         'dataset_type': args.dataset_type,
         'epoch_size': args.epoch_size * args.__dict__.get('batch_size', 5000),
@@ -98,6 +105,7 @@ def build_data_options(args):
         'hold_every': args.llffhold,
         'normalize_by_bbox': args.normalize_by_bbox,
         'data_bbox_scale': args.data_bbox_scale,
+        'normalize_by_camera': args.normalize_by_camera,
         'permutation': args.perm
     }
 
@@ -114,6 +122,9 @@ def maybe_merge_config_file(args, allow_invalid=False):
         args.__dict__.update(configs)
 
 def setup_render_opts(opt, args):
+    """
+    Pass render arguments to the SparseGrid renderer options
+    """
     opt.step_size = args.step_size
     opt.sigma_thresh = args.sigma_thresh
     opt.stop_thresh = args.stop_thresh

opt/util/nsvf_dataset.py

@@ -39,6 +39,7 @@ def __init__(
         white_bkgd: bool = True,
         normalize_by_bbox: bool = False,
         data_bbox_scale : float = 1.1,
+        normalize_by_camera: bool = True,
         **kwargs
     ):
         assert path.isdir(root), f"'{root}' is not a directory"
@@ -113,7 +114,10 @@ def look_for_dir(cands, required=True):
             pose_path = path.join(root, pose_dir_name, pose_fname)
             #  intrin_path = path.join(root, intrin_dir_name, pose_fname)
 
-            cam_mtx = np.loadtxt(pose_path).reshape(4, 4)
+            cam_mtx = np.loadtxt(pose_path).reshape(-1, 4)
+            if len(cam_mtx == 3):
+                bottom = np.array([[0.0, 0.0, 0.0, 1.0]])
+                cam_mtx = np.concatenate([cam_mtx, bottom], axis=0)
             all_c2w.append(torch.from_numpy(cam_mtx))  # C2W (4, 4) OpenCV
             full_size = list(image.shape[:2])
             rsz_h, rsz_w = [round(hw * scale) for hw in full_size]
@@ -125,6 +129,7 @@ def look_for_dir(cands, required=True):
 
         self.c2w_f64 = torch.stack(all_c2w)
 
+        print('NORMALIZE BY?', 'bbox' if normalize_by_bbox else 'camera' if normalize_by_camera else 'manual')
         if normalize_by_bbox:
             # Not used, but could be helpful
             bbox_path = path.join(root, "bbox.txt")
@@ -137,10 +142,21 @@ def look_for_dir(cands, required=True):
                 self.c2w_f64[:, :3, 3] -= center
                 # Rescale
                 scene_scale = 1.0 / radius.max()
-                print(' Overriding scene_scale by ', scene_scale)
             else:
                 warn('normalize_by_bbox=True but bbox.txt was not available')
-
+        elif normalize_by_camera:
+            norm_pose_files = sorted(os.listdir(path.join(root, pose_dir_name)), key=sort_key)
+            norm_pose_files = [x for x in norm_pose_files if not x.startswith('2_')]
+            norm_poses = np.stack([np.loadtxt(path.join(root, pose_dir_name, x)).reshape(-1, 4)
+                                    for x in norm_pose_files], axis=0)
+
+            # Select subset of files
+            center = np.median(norm_poses[:, :3, 3], axis=0)
+            radius = np.median(np.linalg.norm(norm_poses[:, :3, 3] - center, axis=-1))
+            self.c2w_f64[:, :3, 3] -= center
+            scene_scale = 1.0 / radius
+
+        print('scene_scale', scene_scale)
         self.c2w_f64[:, :3, 3] *= scene_scale
         self.c2w = self.c2w_f64.float()
 

svox2/csrc/render_lerp_kernel_cuvol.cu

@@ -288,21 +288,20 @@ __device__ __inline__ void render_background_forward(
     ConcentricSpheresIntersector csi(ray.origin, ray.dir);
 
     const float inner_radius = fmaxf(_dist_ray_to_origin(ray.origin, ray.dir) + 1e-3f, 1.f);
-    float t, t_last;
+    float t, invr_last = 1.f / inner_radius;
     const int n_steps = int(grid.background_nlayers / opt.step_size) + 2;
 
-    csi.intersect(inner_radius, &t_last);
+    // csi.intersect(inner_radius, &t_last);
 
     float outv[3] = {0.f, 0.f, 0.f};
     for (int i = 0; i < n_steps; ++i) {
         // Between 1 and infty
         float r = n_steps / (n_steps - i - 0.5);
         if (r < inner_radius || !csi.intersect(r, &t)) continue;
-        const float t_mid = (t + t_last) * 0.5f;
 
 #pragma unroll 3
         for (int j = 0; j < 3; ++j) {
-            ray.pos[j] = fmaf(t_mid, ray.dir[j], ray.origin[j]);
+            ray.pos[j] = fmaf(t, ray.dir[j], ray.origin[j]);
         }
         const float invr_mid = _rnorm(ray.pos);
 #pragma unroll 3
@@ -338,7 +337,7 @@ __device__ __inline__ void render_background_forward(
         // if (opt.randomize && opt.random_sigma_std_background > 0.0)
         //     sigma += ray.rng.randn() * opt.random_sigma_std_background;
         if (sigma > 0.f) {
-            const float pcnt = (t - t_last) * ray.world_step * sigma;
+            const float pcnt = (invr_last - invr_mid) * ray.world_step * sigma;
             const float weight = _EXP(log_transmit) * (1.f - _EXP(-pcnt));
             log_transmit -= pcnt;
 #pragma unroll 3
@@ -359,7 +358,7 @@ __device__ __inline__ void render_background_forward(
                 break;
             }
         }
-        t_last = t;
+        invr_last = invr_mid;
     }
 #pragma unroll 3
     for (int i = 0; i < 3; ++i) {
@@ -380,20 +379,19 @@ __device__ __inline__ void render_background_backward(
     // printf("log_transmit_init=%f\n", log_transmit);
     ConcentricSpheresIntersector csi(ray.origin, ray.dir);
 
-    float t, t_last;
     const int n_steps = int(grid.background_nlayers / opt.step_size) + 2;
 
     const float inner_radius = fmaxf(_dist_ray_to_origin(ray.origin, ray.dir) + 1e-3f, 1.f);
-    csi.intersect(inner_radius, &t_last);
+    float t, invr_last = 1.f / inner_radius;
+    // csi.intersect(inner_radius, &t_last);
     for (int i = 0; i < n_steps; ++i) {
         float r = n_steps / (n_steps - i - 0.5);
 
         if (r < inner_radius || !csi.intersect(r, &t)) continue;
 
-        const float t_mid = (t + t_last) * 0.5f;
 #pragma unroll 3
         for (int j = 0; j < 3; ++j) {
-            ray.pos[j] = fmaf(t_mid, ray.dir[j], ray.origin[j]);
+            ray.pos[j] = fmaf(t, ray.dir[j], ray.origin[j]);
         }
 
         const float invr_mid = _rnorm(ray.pos);
@@ -431,7 +429,7 @@ __device__ __inline__ void render_background_backward(
         //     sigma += ray.rng.randn() * opt.random_sigma_std_background;
         if (sigma > 0.f) {
             float total_color = 0.f;
-            const float pcnt = ray.world_step * (t - t_last) * sigma;
+            const float pcnt = ray.world_step * (invr_last - invr_mid) * sigma;
             const float weight = _EXP(log_transmit) * (1.f - _EXP(-pcnt));
             log_transmit -= pcnt;
 
@@ -464,7 +462,7 @@ __device__ __inline__ void render_background_backward(
             }
 
             accum -= weight * total_color;
-            float curr_grad_sigma = ray.world_step * (t - t_last) * (
+            float curr_grad_sigma = ray.world_step * (invr_last - invr_mid) * (
                     total_color * _EXP(log_transmit) - accum);
 
             trilerp_backward_bg_one(
@@ -483,7 +481,7 @@ __device__ __inline__ void render_background_backward(
                 break;
             }
         }
-        t_last = t;
+        invr_last = invr_mid;
     }
 }