Add NV renderer backend (not really working)

zhangkaiitugithub · Nov 13, 2021 · 3df92a4 · 3df92a4
1 parent d7a784e
commit 3df92a4
Show file tree

Hide file tree

Showing 10 changed files with 529 additions and 41 deletions.
diff --git a/opt/autotune.py b/opt/autotune.py
@@ -29,7 +29,7 @@
 
 PSNR_FILE_NAME = 'test_psnr.txt'
 
-def run_exp(env, eval_mode:bool, enable_render:bool, train_dir, data_dir, config, flags):
+def run_exp(env, eval_mode:bool, enable_render:bool, train_dir, data_dir, config, flags, eval_flags, common_flags):
     opt_base_cmd = [ "python", "-u", "opt.py", "--tune_mode" ]
 
     if not eval_mode:
@@ -55,7 +55,7 @@ def run_exp(env, eval_mode:bool, enable_render:bool, train_dir, data_dir, config
         opt_ret = ""  # Silence
     else:
         print('! RUN opt.py -t', train_dir)
-        opt_cmd = ' '.join(opt_base_cmd + flags)
+        opt_cmd = ' '.join(opt_base_cmd + flags + common_flags)
         print(opt_cmd)
         try:
             opt_ret = subprocess.check_output(opt_cmd, shell=True, env=env).decode(
@@ -76,7 +76,7 @@ def run_exp(env, eval_mode:bool, enable_render:bool, train_dir, data_dir, config
             eval_base_cmd += ['-c', config]
         psnr_file_path = path.join(train_dir, 'test_renders', 'psnr.txt')
         if not path.exists(psnr_file_path):
-            eval_cmd = ' '.join(eval_base_cmd)
+            eval_cmd = ' '.join(eval_base_cmd + eval_flags + common_flags)
             print('! RUN render_imgs.py', ckpt_path)
             print(eval_cmd)
             try:
@@ -90,7 +90,7 @@ def run_exp(env, eval_mode:bool, enable_render:bool, train_dir, data_dir, config
 
         if enable_render:
             eval_base_cmd += ['--render_path']
-            render_cmd = ' '.join(eval_base_cmd)
+            render_cmd = ' '.join(eval_base_cmd + eval_flags + common_flags)
             try:
                 render_ret = subprocess.check_output(render_cmd, shell=True, env=env).decode(
                         sys.stdout.encoding)
@@ -186,6 +186,8 @@ def recursive_replace(data, variables):
     data_root = path.expanduser(tasks_file['data_root'])  # Required
     train_root = path.expanduser(tasks_file['train_root'])  # Required
     base_flags = tasks_file.get('base_flags', [])
+    base_eval_flags = tasks_file.get('base_eval_flags', [])
+    base_common_flags = tasks_file.get('base_common_flags', [])
     default_config = tasks_file.get('config', '')
 
     if 'eval' in tasks_file:
@@ -212,6 +214,8 @@ def recursive_replace(data, variables):
             task['train_dir'] = path.join(train_root, task['train_dir'])  # Required
             task['data_dir'] = path.join(data_root, task.get('data_dir', '')).rstrip('/')
             task['flags'] = task.get('flags', []) + base_flags
+            task['eval_flags'] = task.get('eval_flags', []) + base_eval_flags
+            task['common_flags'] = task.get('common_flags', []) + base_common_flags
             task['config'] = task.get('config', default_config)
             os.makedirs(task['train_dir'], exist_ok=True)
             # santity check

diff --git a/opt/configs/tnt.json b/opt/configs/tnt.json
@@ -2,19 +2,19 @@
     "reso": "[[128, 128, 128], [256, 256, 256], [512, 512, 512], [640, 640, 640]]",
     "n_iters": 102400,
     "background_nlayers": 64,
-    "background_reso": 1024,
+    "background_reso": 768,
     "upsamp_every": 25600,
     "lr_sigma": 3e1,
     "lr_sh": 1e-2,
     "lr_sigma_delay_steps": 15000,
     "thresh_type": "weight",
-    "weight_thresh": 0.2,
+    "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_beta": 1e-5,
-    "lambda_sparsity": 1e-12,
+    "lambda_sparsity": 1e-11,
     "background_brightness": 1.0,
     "tv_early_only": 0
 }
diff --git a/opt/tasks/eval_tnt.json b/opt/tasks/eval_tnt.json
@@ -1,31 +1,23 @@
 {
     "eval": true,
     "data_root": "/home/sxyu/data/TanksAndTempleBG",
-    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/tnt",
+    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/tnt_spars",
     "config": "configs/tnt.json",
     "tasks": [{
         "train_dir": "Train",
         "data_dir": "Train",
-        "flags": [
-            "--scene_scale", "2.0"
-        ]
+        "common_flags": ["--scene_scale", "2.0"]
     }, {
         "train_dir": "M60",
         "data_dir": "M60",
-        "flags": [
-            "--scene_scale", "2.0"
-        ]
+        "common_flags": ["--scene_scale", "2.0"]
     }, {
         "train_dir": "Truck",
         "data_dir": "Truck",
-        "flags": [
-            "--scene_scale", "1.6"
-        ]
+        "common_flags": ["--scene_scale", "1.6"]
     }, {
         "train_dir": "Playground",
         "data_dir": "Playground",
-        "flags": [
-            "--scene_scale", "1.7"
-        ]
+        "common_flags": ["--scene_scale", "1.7"]
     }]
 }
diff --git a/opt/util/config_util.py b/opt/util/config_util.py
@@ -67,16 +67,16 @@ def define_common_args(parser : argparse.ArgumentParser):
                          default=1.0,
                          help="Brightness of the infinite background")
     group.add_argument('--renderer_backend', '-B',
-                         choices=['cuvol', 'svox1'],
+                         choices=['cuvol', 'svox1', 'nvol'],
                          default='cuvol',
                          help="Renderer backend")
     group.add_argument('--random_sigma_std',
                          type=float,
-                         default=1.0,
+                         default=0.0,
                          help="Random Gaussian std to add to density values (only if enable_random)")
     group.add_argument('--random_sigma_std_background',
                          type=float,
-                         default=1.0,
+                         default=0.0,
                          help="Random Gaussian std to add to density values for BG (only if enable_random)")
     group.add_argument('--enable_random',
                          action='store_true',

diff --git a/opt/util/ff_dataset.py b/opt/util/ff_dataset.py
@@ -147,7 +147,6 @@ def _load_images(self):
         self.c2w = torch.stack(all_c2w)
         #  bds_scale = 1.0 / (self.sfm.dmin * 0.75) # 0.9
         bds_scale = 1.0
-        print('scene rescale', bds_scale)
         self.z_bounds = [self.sfm.dmin * bds_scale, self.sfm.dmax * bds_scale]
         if bds_scale != 1.0:
             self.c2w[:, :3, 3] *= bds_scale
@@ -169,24 +168,12 @@ def _load_images(self):
         width = self.sfm.ref_cam['width']
         height = self.sfm.ref_cam['height']
 
-        print('z_bounds from LLFF:', self.z_bounds)
+        print('z_bounds from LLFF:', self.z_bounds, '(not used)')
 
-        # Auto bounds
-        #  [1.496031746031746, 1.6613756613756614, 1.0]
+        # Padded bounds
         radx = 1 + 2 * self.sfm.offset / self.gt.size(2)
         rady = 1 + 2 * self.sfm.offset / self.gt.size(1)
-
-        # Some hardcoded stuff
-        #  radx = 1.4 #1 + 2 * self.sfm.offset / self.gt.size(2)
-        #  rady = 1.75 #+ 2 * self.sfm.offset / self.gt.size(1)
-
-        #  MAX_MPI_LAYERS = 132
-        #  # Our convention is a bit weird so there is some 0.5 voxel
-        #  # empty space on each end, better to deal with this
-        #  radz = MAX_MPI_LAYERS / (MAX_MPI_LAYERS - 1)
         radz = 1.0
-
-
         self.scene_center = [0.0, 0.0, 0.0]
         self.scene_radius = [radx, rady, radz]
         print('scene_radius', self.scene_radius)

diff --git a/setup.py b/setup.py
@@ -39,6 +39,7 @@
             'svox2/csrc/svox2.cpp',
             'svox2/csrc/svox2_kernel.cu',
             'svox2/csrc/render_lerp_kernel_cuvol.cu',
+            'svox2/csrc/render_lerp_kernel_nvol.cu',
             'svox2/csrc/render_svox1_kernel.cu',
             'svox2/csrc/misc_kernel.cu',
             'svox2/csrc/loss_kernel.cu',