Working ish T&T (using some manual scaling)

zhangkaiitugithub · Nov 13, 2021 · 71d4730 · 71d4730
1 parent 4b7c4d6
commit 71d4730
Show file tree

Hide file tree

Showing 21 changed files with 757 additions and 452 deletions.
diff --git a/opt/autotune.py b/opt/autotune.py
@@ -148,6 +148,8 @@ def create_prodvars(variables, noise_stds={}):
     def auto_list(x):
         if isinstance(x, list):
             return x
+        elif isinstance(x, dict) or isinstance(x, set):
+            return [x]
         elif isinstance(x, str):
             return eval(x)
         else:
@@ -184,6 +186,7 @@ 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', [])
+    default_config = tasks_file.get('config', '')
 
     if 'eval' in tasks_file:
         args.eval = tasks_file['eval']
@@ -209,7 +212,7 @@ 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['config'] = task.get('config', '')
+            task['config'] = task.get('config', default_config)
             os.makedirs(task['train_dir'], exist_ok=True)
             # santity check
             assert path.exists(task['train_dir']), task['train_dir'] + ' does not exist'

diff --git a/opt/calc_metrics.py b/opt/calc_metrics.py
@@ -76,9 +76,9 @@
 print('SSIM:', avg_ssim)
 print('LPIPS:', avg_lpips)
 postfix = '_cropped' if args.crop != 1.0 else ''
-with open(path.join(args.render_dir, f'psnr{postfix}.txt'), 'w') as f:
-    f.write(str(avg_psnr))
-with open(path.join(args.render_dir, f'ssim{postfix}.txt'), 'w') as f:
-    f.write(str(avg_ssim))
-with open(path.join(args.render_dir, f'lpips{postfix}.txt'), 'w') as f:
-    f.write(str(avg_lpips))
+#  with open(path.join(args.render_dir, f'psnr{postfix}.txt'), 'w') as f:
+#      f.write(str(avg_psnr))
+#  with open(path.join(args.render_dir, f'ssim{postfix}.txt'), 'w') as f:
+#      f.write(str(avg_ssim))
+#  with open(path.join(args.render_dir, f'lpips{postfix}.txt'), 'w') as f:
+#      f.write(str(avg_lpips))
diff --git a/opt/configs/llff.json b/opt/configs/llff.json
@@ -3,14 +3,11 @@
     "upsamp_every": 38400,
     "lr_sigma": 3e1,
     "lr_sh": 1e-2,
-    "lr_sigma_delay_steps": 15000,
     "thresh_type": "sigma",
     "density_thresh": 5,
-    "lambda_tv": 2e-4,
-    "lambda_tv_sh": 2e-3,
-    "enable_random": false,
-    "random_sigma_std": 0.0,
+    "lambda_tv": 5e-4,
+    "lambda_tv_sh": 5e-3,
     "lambda_sparsity": 1e-12,
-    "lambda_beta": 0.0,
-    "background_brightness": 0.5
+    "background_brightness": 0.5,
+    "tv_early_only": 0
 }
diff --git a/opt/configs/tnt.json b/opt/configs/tnt.json
@@ -0,0 +1,20 @@
+{
+    "reso": "[[128, 128, 128], [256, 256, 256], [512, 512, 512], [640, 640, 640]]",
+    "n_iters": 102400,
+    "background_nlayers": 64,
+    "background_reso": 1024,
+    "upsamp_every": 25600,
+    "lr_sigma": 3e1,
+    "lr_sh": 1e-2,
+    "lr_sigma_delay_steps": 15000,
+    "thresh_type": "weight",
+    "weight_thresh": 0.2,
+    "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,
+    "background_brightness": 1.0,
+    "tv_early_only": 0
+}
diff --git a/opt/opt.py b/opt/opt.py
@@ -85,9 +85,7 @@
 
 # TODO: make the lr higher near the end
 group.add_argument('--sigma_optim', choices=['sgd', 'rmsprop'], default='rmsprop', help="Density optimizer")
-group.add_argument('--lr_sigma', type=float, default=
-                        3e1,
-        help='SGD/rmsprop lr for sigma')
+group.add_argument('--lr_sigma', type=float, default=3e1, help='SGD/rmsprop lr for sigma')
 group.add_argument('--lr_sigma_final', type=float, default=5e-2)
 group.add_argument('--lr_sigma_decay_steps', type=int, default=250000)
 group.add_argument('--lr_sigma_delay_steps', type=int, default=15000,
@@ -106,13 +104,25 @@
 group.add_argument('--lr_sh_decay_steps', type=int, default=250000)
 group.add_argument('--lr_sh_delay_steps', type=int, default=0, help="Reverse cosine steps (0 means disable)")
 group.add_argument('--lr_sh_delay_mult', type=float, default=1e-2)
-group.add_argument('--lr_sh_upscale_factor', type=float, default=1.0)
 
-group.add_argument('--bg_optim', choices=['sgd', 'rmsprop'], default='rmsprop', help="Density optimizer")
-group.add_argument('--lr_sigma_bg', type=float, default=3e-1,#3e-2,
+# BG LRs
+group.add_argument('--bg_optim', choices=['sgd', 'rmsprop'], default='rmsprop', help="Background optimizer")
+group.add_argument('--lr_sigma_bg', type=float, default=3e0,
                     help='SGD/rmsprop lr for background')
-group.add_argument('--lr_color_bg', type=float, default=1e-2, #1e-2,
+group.add_argument('--lr_sigma_bg_final', type=float, default=3e-3,
                     help='SGD/rmsprop lr for background')
+group.add_argument('--lr_sigma_bg_decay_steps', type=int, default=250000)
+group.add_argument('--lr_sigma_bg_delay_steps', type=int, default=0, help="Reverse cosine steps (0 means disable)")
+group.add_argument('--lr_sigma_bg_delay_mult', type=float, default=1e-2)
+
+group.add_argument('--lr_color_bg', type=float, default=1e-1,
+                    help='SGD/rmsprop lr for background')
+group.add_argument('--lr_color_bg_final', type=float, default=5e-6,#1e-4,
+                    help='SGD/rmsprop lr for background')
+group.add_argument('--lr_color_bg_decay_steps', type=int, default=250000)
+group.add_argument('--lr_color_bg_delay_steps', type=int, default=0, help="Reverse cosine steps (0 means disable)")
+group.add_argument('--lr_color_bg_delay_mult', type=float, default=1e-2)
+# END BG LRs
 
 group.add_argument('--basis_optim', choices=['sgd', 'rmsprop'], default='rmsprop', help="Learned basis optimizer")
 group.add_argument('--lr_basis', type=float, default=#2e6,
@@ -131,7 +141,7 @@
 group.add_argument('--rms_beta', type=float, default=0.95, help="RMSProp exponential averaging factor")
 
 group.add_argument('--print_every', type=int, default=20, help='print every')
-group.add_argument('--save_every', type=int, default=2, #5,
+group.add_argument('--save_every', type=int, default=5,
                    help='save every x epochs')
 group.add_argument('--eval_every', type=int, default=1,
                    help='evaluate every x epochs')
@@ -148,14 +158,15 @@
                     default="weight",
                    help='Upsample threshold type')
 group.add_argument('--weight_thresh', type=float,
-                    #  default=0.0005,
-                    #  default=0.005 * 256,
                     #  default=0.0005 * 512,
                     default=0.025 * 512,
                    help='Upsample weight threshold; will be divided by resulting z-resolution')
 group.add_argument('--density_thresh', type=float,
-                    default=0.0005,
+                    default=5.0,
                    help='Upsample sigma threshold')
+group.add_argument('--background_density_thresh', type=float,
+                    default=1.0,
+                   help='Background sigma threshold for sparsification')
 group.add_argument('--max_grid_elements', type=int,
                     default=44_000_000,
                    help='Max items to store after upsampling '
@@ -196,8 +207,8 @@
 
 
 # Background TV
-group.add_argument('--lambda_tv_background_sigma', type=float, default=1e-5)
-group.add_argument('--lambda_tv_background_color', type=float, default=1e-4)
+group.add_argument('--lambda_tv_background_sigma', type=float, default=1e-2)
+group.add_argument('--lambda_tv_background_color', type=float, default=1e-2)
 
 group.add_argument('--tv_background_sparsity', type=float, default=0.01)
 # End Background TV
@@ -270,8 +281,8 @@
 grid.density_data.data[:] = args.init_sigma
 
 if grid.use_background:
-    grid.background_cubemap.data[..., -1] = args.init_sigma
-    grid.background_cubemap.data[..., :-1] = 0.0 / svox2.utils.SH_C0
+    grid.background_data.data[..., -1] = args.init_sigma
+    grid.background_data.data[..., :-1] = 0.0 / svox2.utils.SH_C0
 
 #  grid.sh_data.data[:, 0] = 4.0
 #  osh = grid.density_data.data.shape
@@ -321,6 +332,10 @@
                                args.lr_sh_delay_mult, args.lr_sh_decay_steps)
 lr_basis_func = get_expon_lr_func(args.lr_basis, args.lr_basis_final, args.lr_basis_delay_steps,
                                args.lr_basis_delay_mult, args.lr_basis_decay_steps)
+lr_sigma_bg_func = get_expon_lr_func(args.lr_sigma_bg, args.lr_sigma_bg_final, args.lr_sigma_bg_delay_steps,
+                               args.lr_sigma_bg_delay_mult, args.lr_sigma_bg_decay_steps)
+lr_color_bg_func = get_expon_lr_func(args.lr_color_bg, args.lr_color_bg_final, args.lr_color_bg_delay_steps,
+                               args.lr_color_bg_delay_mult, args.lr_color_bg_decay_steps)
 lr_sigma_factor = 1.0
 lr_sh_factor = 1.0
 lr_basis_factor = 1.0
@@ -434,6 +449,8 @@ def train_step():
             lr_sigma = lr_sigma_func(gstep_id) * lr_sigma_factor
             lr_sh = lr_sh_func(gstep_id) * lr_sh_factor
             lr_basis = lr_basis_func(gstep_id - args.lr_basis_begin_step) * lr_basis_factor
+            lr_sigma_bg = lr_sigma_bg_func(gstep_id - args.lr_basis_begin_step) * lr_basis_factor
+            lr_color_bg = lr_color_bg_func(gstep_id - args.lr_basis_begin_step) * lr_basis_factor
             if not args.lr_decay:
                 lr_sigma = args.lr_sigma * lr_sigma_factor
                 lr_sh = args.lr_sh * lr_sh_factor
@@ -484,7 +501,11 @@ def train_step():
                 #  summary_writer.add_scalar("loss_tv_basis", tv_basis, global_step=gstep_id)
                 summary_writer.add_scalar("lr_sh", lr_sh, global_step=gstep_id)
                 summary_writer.add_scalar("lr_sigma", lr_sigma, global_step=gstep_id)
-                summary_writer.add_scalar("lr_basis", lr_basis, global_step=gstep_id)
+                if grid.basis_type == svox2.BASIS_TYPE_3D_TEXTURE:
+                    summary_writer.add_scalar("lr_basis", lr_basis, global_step=gstep_id)
+                if grid.use_background:
+                    summary_writer.add_scalar("lr_sigma_bg", lr_sigma_bg, global_step=gstep_id)
+                    summary_writer.add_scalar("lr_color_bg", lr_color_bg, global_step=gstep_id)
 
                 if args.weight_decay_sh < 1.0:
                     grid.sh_data.data *= args.weight_decay_sigma
@@ -519,7 +540,7 @@ def train_step():
                 grid.inplace_l2_color_grad(grid.sh_data.grad,
                         scaling=args.lambda_l2_sh)
             if grid.use_background and (args.lambda_tv_background_sigma > 0.0 or args.lambda_tv_background_color > 0.0):
-                grid.inplace_tv_background_grad(grid.background_cubemap.grad,
+                grid.inplace_tv_background_grad(grid.background_data.grad,
                         scaling=args.lambda_tv_background_color,
                         scaling_density=args.lambda_tv_background_sigma,
                         sparse_frac=args.tv_background_sparsity)
@@ -534,7 +555,7 @@ def train_step():
             grid.optim_density_step(lr_sigma, beta=args.rms_beta, optim=args.sigma_optim)
             grid.optim_sh_step(lr_sh, beta=args.rms_beta, optim=args.sh_optim)
             if grid.use_background:
-                grid.optim_background_step(args.lr_sigma_bg, args.lr_color_bg, beta=args.rms_beta, optim=args.bg_optim)
+                grid.optim_background_step(lr_sigma_bg, lr_color_bg, beta=args.rms_beta, optim=args.bg_optim)
             if gstep_id >= args.lr_basis_begin_step:
                 if grid.basis_type == svox2.BASIS_TYPE_3D_TEXTURE:
                     grid.optim_basis_step(lr_basis, beta=args.rms_beta, optim=args.basis_optim)
@@ -571,11 +592,11 @@ def train_step():
                     cameras=resample_cameras if args.thresh_type == 'weight' else None,
                     max_elements=args.max_grid_elements)
 
+            if grid.use_background and reso_id <= 1:
+                grid.sparsify_background(args.background_density_thresh)
+
             if args.upsample_density_add:
                 grid.density_data.data[:] += args.upsample_density_add
-            if args.lr_sh_upscale_factor > 1:
-                lr_sh_factor *= args.lr_sh_upscale_factor
-                print('Increased lr to (sigma:)', args.lr_sigma, '(sh:)', args.lr_sh)
 
         if factor > 1 and reso_id < len(reso_list) - 1:
             print('* Using higher resolution images due to large grid; new factor', factor)

diff --git a/opt/render_imgs.py b/opt/render_imgs.py
@@ -108,14 +108,17 @@
 if grid.use_background:
     if args.nobg:
         #  grid.background_cubemap.data = grid.background_cubemap.data.cuda()
-        grid.background_cubemap.data[..., -1] = 0.0
+        grid.background_data.data[..., -1] = 0.0
         render_dir += '_nobg'
     if args.nofg:
         grid.density_data.data[:] = 0.0
         render_dir += '_nofg'
 
 config_util.setup_render_opts(grid.opt, args)
 
+grid.opt.msi_start_layer = 30
+grid.opt.msi_end_layer = 31 #grid.opt.msi_start_layer + 1
+
 if args.blackbg:
     print('Forcing black bg')
     render_dir += '_blackbg'

diff --git a/opt/tasks/eval_ff.json b/opt/tasks/eval_ff.json
@@ -2,7 +2,7 @@
     "eval": true,
     "render": true,
     "data_root": "/home/sxyu/data/nerf_llff_data",
-    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/llff_c2f",
+    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/llff_c2f_10e",
     "variables": {
         "scene": ["fern", "room", "horns", "trex", "flower", "leaves", "orchids", "fortress"]
     },

diff --git a/opt/tasks/eval_tnt.json b/opt/tasks/eval_tnt.json
@@ -0,0 +1,31 @@
+{
+    "eval": true,
+    "data_root": "/home/sxyu/data/TanksAndTempleBG",
+    "train_root": "/home/sxyu/proj/svox2/opt/ckpt_auto/tnt",
+    "config": "configs/tnt.json",
+    "tasks": [{
+        "train_dir": "Train",
+        "data_dir": "Train",
+        "flags": [
+            "--scene_scale", "2.0"
+        ]
+    }, {
+        "train_dir": "M60",
+        "data_dir": "M60",
+        "flags": [
+            "--scene_scale", "2.0"
+        ]
+    }, {
+        "train_dir": "Truck",
+        "data_dir": "Truck",
+        "flags": [
+            "--scene_scale", "1.6"
+        ]
+    }, {
+        "train_dir": "Playground",
+        "data_dir": "Playground",
+        "flags": [
+            "--scene_scale", "1.7"
+        ]
+    }]
+}
diff --git a/opt/util/util.py b/opt/util/util.py
@@ -123,27 +123,39 @@ def save_img(img: np.ndarray, path: str):
     cv2.imwrite(path, img)
 
 
-def equirect2xyz(uv):
+def equirect2xyz(uv, rows, cols):
     """
     Convert equirectangular coordinate to unit vector,
     inverse of xyz2equirect
+    Taken from Vickie Ye
     Args:
         uv: np.ndarray [..., 2] x, y coordinates in image space in [-1.0, 1.0]
     Returns:
         xyz: np.ndarray [..., 3] unit vectors
     """
-    lon = uv[..., 0] * np.pi
-    lat = uv[..., 1] * (np.pi * 0.5)
+    lon = (uv[..., 0] * (1.0 / cols) - 0.5) * (2 * np.pi)
+    lat = -(uv[..., 1] * (1.0 / rows) - 0.5) * np.pi
     coslat = np.cos(lat)
     return np.stack(
         [
             coslat * np.sin(lon),
-            coslat * np.cos(lon),
             np.sin(lat),
+            coslat * np.cos(lon),
         ],
         axis=-1,
     )
 
+def xyz2equirect(bearings, rows, cols):
+    """
+    Convert ray direction vectors into equirectangular pixel coordinates.
+    Inverse of equirect2xyz.
+    Taken from Vickie Ye
+    """
+    lat = np.arcsin(bearings[..., 1])
+    lon = np.arctan2(bearings[..., 0], bearings[..., 2])
+    x = cols * (0.5 + lon / 2 / np.pi)
+    y = rows * (0.5 - lat / np.pi)
+    return np.stack([x, y], axis=-1)
 
 def generate_dirs_equirect(w, h):
     x, y = np.meshgrid(  # pylint: disable=unbalanced-tuple-unpacking

diff --git a/svox2/csrc/include/data_spec.hpp b/svox2/csrc/include/data_spec.hpp
@@ -22,7 +22,8 @@ struct SparseGridSpec {
   Tensor _offset;
   Tensor _scaling;
 
-  Tensor background_cubemap;
+  Tensor background_links;
+  Tensor background_data;
 
   int basis_dim;
   uint8_t basis_type;
@@ -32,20 +33,21 @@ struct SparseGridSpec {
     CHECK_INPUT(density_data);
     CHECK_INPUT(sh_data);
     CHECK_INPUT(links);
-    CHECK_INPUT(background_cubemap);
-    CHECK_INPUT(basis_data);
+    if (background_links.defined()) {
+      CHECK_INPUT(background_links);
+      CHECK_INPUT(background_data);
+      TORCH_CHECK(background_links.ndimension() ==
+                  2);                                 // (H, W) -> [N] \cup {-1}
+      TORCH_CHECK(background_data.ndimension() == 3); // (N, D, C) -> R
+    }
+    if (basis_data.defined()) {
+      CHECK_INPUT(basis_data);
+    }
     CHECK_CPU_INPUT(_offset);
     CHECK_CPU_INPUT(_scaling);
-    TORCH_CHECK(density_data.is_floating_point());
-    TORCH_CHECK(sh_data.is_floating_point());
-    TORCH_CHECK(!links.is_floating_point());
-    TORCH_CHECK(basis_data.is_floating_point());
-    TORCH_CHECK(_offset.is_floating_point());
-    TORCH_CHECK(_scaling.is_floating_point());
     TORCH_CHECK(density_data.ndimension() == 2);
     TORCH_CHECK(sh_data.ndimension() == 2);
     TORCH_CHECK(links.ndimension() == 3);
-    TORCH_CHECK(background_cubemap.ndimension() == 5);
   }
 };
 
@@ -125,4 +127,7 @@ struct RenderOptions {
   float random_sigma_std_background;
   // 32-bit RNG state masks
   uint32_t _m1, _m2, _m3;
+
+  // int msi_start_layer = 0;
+  // int msi_end_layer = 66;
 };