correct data split for hypernerf & llff datasets

dnerf/provider.py

@@ -206,9 +206,9 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
             self.times = []
 
             # assume frames are already sorted by time!
-            for f in tqdm.tqdm(frames, desc=f'Loading {type} data:'):
+            for f in tqdm.tqdm(frames, desc=f'Loading {type} data'):
                 f_path = os.path.join(self.root_path, f['file_path'])
-                if self.mode == 'blender' and f_path[-4:] != '.png':
+                if self.mode == 'blender' and '.' not in f_path:
                     f_path += '.png' # so silly...
 
                 # there are non-exist paths in fox...

loss.py

@@ -2,6 +2,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+import numpy as np
+
 def mape_loss(pred, target, reduction='mean'):
     # pred, target: [B, 1], torch tenspr
     difference = (pred - target).abs()
@@ -21,4 +23,54 @@ def huber_loss(pred, target, delta=0.1, reduction='mean'):
     if reduction == 'mean':
         loss = loss.mean()
 
-    return loss
+    return loss
+
+
+# ref: https://github.com/sunset1995/torch_efficient_distloss/blob/main/torch_efficient_distloss/eff_distloss.py
+class EffDistLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, w, m, interval):
+        '''
+        Efficient O(N) realization of distortion loss.
+        There are B rays each with N sampled points.
+        w:        Float tensor in shape [B,N]. Volume rendering weights of each point.
+        m:        Float tensor in shape [B,N]. Midpoint distance to camera of each point.
+        interval: Scalar or float tensor in shape [B,N]. The query interval of each point.
+        '''
+        n_rays = np.prod(w.shape[:-1])
+        wm = (w * m)
+        w_cumsum = w.cumsum(dim=-1)
+        wm_cumsum = wm.cumsum(dim=-1)
+
+        w_total = w_cumsum[..., [-1]]
+        wm_total = wm_cumsum[..., [-1]]
+        w_prefix = torch.cat([torch.zeros_like(w_total), w_cumsum[..., :-1]], dim=-1)
+        wm_prefix = torch.cat([torch.zeros_like(wm_total), wm_cumsum[..., :-1]], dim=-1)
+        loss_uni = (1/3) * interval * w.pow(2)
+        loss_bi = 2 * w * (m * w_prefix - wm_prefix)
+        if torch.is_tensor(interval):
+            ctx.save_for_backward(w, m, wm, w_prefix, w_total, wm_prefix, wm_total, interval)
+            ctx.interval = None
+        else:
+            ctx.save_for_backward(w, m, wm, w_prefix, w_total, wm_prefix, wm_total)
+            ctx.interval = interval
+        ctx.n_rays = n_rays
+        return (loss_bi.sum() + loss_uni.sum()) / n_rays
+
+    @staticmethod
+    @torch.autograd.function.once_differentiable
+    def backward(ctx, grad_back):
+        interval = ctx.interval
+        n_rays = ctx.n_rays
+        if interval is None:
+            w, m, wm, w_prefix, w_total, wm_prefix, wm_total, interval = ctx.saved_tensors
+        else:
+            w, m, wm, w_prefix, w_total, wm_prefix, wm_total = ctx.saved_tensors
+        grad_uni = (1/3) * interval * 2 * w
+        w_suffix = w_total - (w_prefix + w)
+        wm_suffix = wm_total - (wm_prefix + wm)
+        grad_bi = 2 * (m * (w_prefix - w_suffix) + (wm_suffix - wm_prefix))
+        grad = grad_back * (grad_bi + grad_uni) / n_rays
+        return grad, None, None, None
+
+eff_distloss = EffDistLoss.apply

main_nerf.py

@@ -20,7 +20,7 @@
     parser.add_argument('--seed', type=int, default=0)
 
     ### training options
-    parser.add_argument('--iters', type=int, default=40000, help="training iters")
+    parser.add_argument('--iters', type=int, default=30000, help="training iters")
     parser.add_argument('--lr', type=float, default=1e-2, help="initial learning rate")
     parser.add_argument('--ckpt', type=str, default='latest')
     parser.add_argument('--num_rays', type=int, default=4096, help="num rays sampled per image for each training step")
@@ -123,7 +123,7 @@
 
         optimizer = lambda model: torch.optim.Adam(model.get_params(opt.lr), betas=(0.9, 0.99), eps=1e-15)
 
-        train_loader = NeRFDataset(opt, device=device, type='trainval').dataloader()
+        train_loader = NeRFDataset(opt, device=device, type='train').dataloader()
 
         # decay to 0.1 * init_lr at last iter step
         scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer, lambda iter: 0.1 ** min(iter / opt.iters, 1))

nerf/provider.py

@@ -192,9 +192,9 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
 
             self.poses = []
             self.images = []
-            for f in tqdm.tqdm(frames, desc=f'Loading {type} data:'):
+            for f in tqdm.tqdm(frames, desc=f'Loading {type} data'):
                 f_path = os.path.join(self.root_path, f['file_path'])
-                if self.mode == 'blender' and f_path[-4:] != '.png':
+                if self.mode == 'blender' and '.' not in f_path:
                     f_path += '.png' # so silly...
 
                 # there are non-exist paths in fox...

readme.md

@@ -63,6 +63,25 @@ For GPUs with lower architecture, `--tcnn` can still be used, but the speed will
 We use the same data format as instant-ngp, e.g., [armadillo](https://github.com/NVlabs/instant-ngp/blob/master/data/sdf/armadillo.obj) and [fox](https://github.com/NVlabs/instant-ngp/tree/master/data/nerf/fox). 
 Please download and put them under `./data`.
 
+We also support self-captured dataset and converting other formats (e.g., LLFF, Tanks&Temples, Mip-NeRF 360) to the nerf-compatible format, with details in the following code block.
+
+<details>
+  <summary> Supported datasets </summary>
+
+  * [nerf_synthetic](https://drive.google.com/drive/folders/128yBriW1IG_3NJ5Rp7APSTZsJqdJdfc1) 
+
+  * [Tanks&Temples](https://dl.fbaipublicfiles.com/nsvf/dataset/TanksAndTemple.zip): [[conversion script]](./scripts/tanks2nerf.py)
+
+  * [LLFF](https://drive.google.com/drive/folders/14boI-o5hGO9srnWaaogTU5_ji7wkX2S7): [[conversion script]](./scripts/llff2nerf.py)
+
+  * [Mip-NeRF 360](http://storage.googleapis.com/gresearch/refraw360/360_v2.zip): [[conversion script]](./scripts/llff2nerf.py)
+
+  * (dynamic) [D-NeRF](https://www.dropbox.com/s/0bf6fl0ye2vz3vr/data.zip?dl=0)
+
+  * (dynamic) [Hyper-NeRF](https://github.com/google/hypernerf/releases/tag/v0.1): [[conversion script]](./scripts/hyper2nerf.py)
+
+</details>
+
 First time running will take some time to compile the CUDA extensions.
 
 ```bash

scripts/hyper2nerf.py

@@ -211,6 +211,15 @@ def rotmat(a, b):
         'frames': train_frames,
     }
     transforms_val = {
+        'w': W,
+        'h': H,
+        'fl_x': fl,
+        'fl_y': fl,
+        'cx': cx,
+        'cy': cy,
+        'frames': val_frames[::10], # only use 1/10 frames for val
+    }
+    transforms_test = {
         'w': W,
         'h': H,
         'fl_x': fl,
@@ -235,5 +244,5 @@ def rotmat(a, b):
     output_path = os.path.join(opt.path, 'transforms_test.json')
     print(f'[INFO] write to {output_path}')
     with open(output_path, 'w') as f:
-        json.dump(transforms_val, f, indent=2)
+        json.dump(transforms_test, f, indent=2)
 

scripts/llff2nerf.py

@@ -65,6 +65,7 @@ def visualize_poses(poses, size=0.1):
     parser.add_argument('path', type=str, help="root directory to the LLFF dataset (contains images/ and pose_bounds.npy)")
     parser.add_argument('--images', type=str, default='images', help="images folder (do not include full path, e.g., just use `images_4`)")
     parser.add_argument('--downscale', type=float, default=1, help="image size down scale, e.g., 4")
+    parser.add_argument('--hold', type=int, default=8, help="hold out for validation every $ images")
 
     opt = parser.parse_args()
     print(f'[INFO] process {opt.path}')
@@ -99,7 +100,7 @@ def visualize_poses(poses, size=0.1):
 
     # visualize_poses(poses)
 
-    # the following stuff are from colmap2nerf... [flower fails, the camera must be in-ward focusing...]
+    # the following stuff are from colmap2nerf... [flower fails, the camera must be in-ward...]
     poses[:, 0:3, 1] *= -1
     poses[:, 0:3, 2] *= -1
     poses = poses[:, [1, 0, 2, 3], :] # swap y and z
@@ -134,28 +135,72 @@ def visualize_poses(poses, size=0.1):
     # visualize_poses(poses)
 
     # construct frames
-    frames = []
-    for i in range(N):
-        frames.append({
+
+    all_ids = np.arange(N)
+    test_ids = all_ids[::opt.hold]
+    train_ids = np.array([i for i in all_ids if i not in test_ids])
+
+    frames_train = []
+    frames_test = []
+    for i in train_ids:
+        frames_train.append({
+            'file_path': images[i],
+            'transform_matrix': poses[i].tolist(),
+        })
+    for i in test_ids:
+        frames_test.append({
             'file_path': images[i],
             'transform_matrix': poses[i].tolist(),
         })
 
+
     # construct a transforms.json
-    transforms = {
+    transforms_train = {
+        'w': W,
+        'h': H,
+        'fl_x': fl,
+        'fl_y': fl,
+        'cx': W // 2,
+        'cy': H // 2,
+        'aabb_scale': 2,
+        'frames': frames_train,
+    }
+
+    transforms_val = {
         'w': W,
         'h': H,
         'fl_x': fl,
         'fl_y': fl,
         'cx': W // 2,
         'cy': H // 2,
         'aabb_scale': 2,
-        'frames': frames,
+        'frames': frames_test[::10],
+    }
+
+    transforms_test = {
+        'w': W,
+        'h': H,
+        'fl_x': fl,
+        'fl_y': fl,
+        'cx': W // 2,
+        'cy': H // 2,
+        'aabb_scale': 2,
+        'frames': frames_test,
     }
 
     # write
-    output_path = os.path.join(opt.path, 'transforms.json')
+    output_path = os.path.join(opt.path, 'transforms_train.json')
+    print(f'[INFO] write to {output_path}')
+    with open(output_path, 'w') as f:
+        json.dump(transforms_train, f, indent=2)
+
+    output_path = os.path.join(opt.path, 'transforms_test.json')
+    print(f'[INFO] write to {output_path}')
+    with open(output_path, 'w') as f:
+        json.dump(transforms_test, f, indent=2)
+
+    output_path = os.path.join(opt.path, 'transforms_val.json')
     print(f'[INFO] write to {output_path}')
     with open(output_path, 'w') as f:
-        json.dump(transforms, f, indent=2)
+        json.dump(transforms_val, f, indent=2)
 

scripts/run_dnerf.sh

@@ -6,4 +6,4 @@
 #OMP_NUM_THREADS=8 CUDA_VISIBLE_DEVICES=1 python main_dnerf.py data/split-cookie/ --workspace trial_dnerf_cookies -O --bound 1 --scale 0.3 --dt_gamma 0
 #OMP_NUM_THREADS=8 CUDA_VISIBLE_DEVICES=1 python main_dnerf.py data/split-cookie/ --workspace trial_dnerf_cookies_ncr --preload --fp16 --bound 1 --scale 0.3 --dt_gamma 0
 
-OMP_NUM_THREADS=8 CUDA_VISIBLE_DEVICES=4 python main_dnerf.py data/vrig-3dprinter/ --workspace trial_dnerf_printer -O --bound 1 --scale 0.3 --dt_gamma 0
+OMP_NUM_THREADS=8 CUDA_VISIBLE_DEVICES=4 python main_dnerf.py data/vrig-3dprinter/ --workspace trial_dnerf_printer -O --bound 2 --scale 0.33 --dt_gamma 0