auto-detect dataset mode

dnerf/provider.py

@@ -94,7 +94,6 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
         self.type = type # train, val, test
         self.downscale = downscale
         self.root_path = opt.path
-        self.mode = opt.mode # colmap, blender
         self.preload = opt.preload # preload data into GPU
         self.scale = opt.scale # camera radius scale to make sure camera are inside the bounding box.
         self.bound = opt.bound # bounding box half length, also used as the radius to random sample poses.
@@ -105,6 +104,12 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
 
         self.rand_pose = opt.rand_pose
 
+        # auto-detect transforms.json and split mode.
+        if os.path.exists(os.path.join(self.root_path, 'transforms.json')):
+            self.mode = 'colmap' # manually split, use view-interpolation for test.
+        elif os.path.exists(os.path.join(self.root_path, 'transforms_train.json')):
+            self.mode = 'blender' # provided split
+
         # load nerf-compatible format data.
         if self.mode == 'colmap':
             with open(os.path.join(self.root_path, 'transforms.json'), 'r') as f:
@@ -235,7 +240,7 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
             self.error_map = None
 
         # [debug] uncomment to view all training poses.
-        # visualize_poses(self.poses.numpy())
+        visualize_poses(self.poses.numpy())
 
         # [debug] uncomment to view examples of randomly generated poses.
         # visualize_poses(rand_poses(100, self.device, radius=self.radius).cpu().numpy())
@@ -328,4 +333,5 @@ def dataloader(self):
             size += size // self.rand_pose # index >= size means we use random pose.
         loader = DataLoader(list(range(size)), batch_size=1, collate_fn=self.collate, shuffle=self.training, num_workers=0)
         loader._data = self # an ugly fix... we need to access error_map & poses in trainer.
+        loader.has_gt = self.images is not None
         return loader

main_CCNeRF.py

@@ -39,7 +39,6 @@
     parser.add_argument("--upsample_model_steps", type=int, action="append", default=[2000, 3000, 4000, 5500, 7000])
 
     ### dataset options
-    parser.add_argument('--mode', type=str, default='blender', help="dataset mode, supports (colmap, blender)")
     parser.add_argument('--color_space', type=str, default='linear', help="Color space, supports (linear, srgb)")
     parser.add_argument('--preload', action='store_true', help="preload all data into GPU, accelerate training but use more GPU memory")
     parser.add_argument('--bound', type=float, default=1, help="assume the scene is bounded in box[-bound, bound]^3, if > 1, will invoke adaptive ray marching.")
@@ -161,7 +160,7 @@ def load_model(path):
             # compose mode have no gt, do not evaulate
             if opt.compose:
                 trainer.test(test_loader, save_path=os.path.join(opt.workspace, 'compose'))
-            elif opt.mode == 'blender':
+            elif test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.
@@ -215,7 +214,7 @@ def load_model(path):
                 print(model)
                 trainer.save_checkpoint(name, full=False, remove_old=False)
 
-                if opt.mode == 'blender':
+                if test_loader.has_gt:
                     trainer.evaluate(test_loader, name=name) # blender has gt, so evaluate it.
                 else:
                     trainer.test(test_loader, name=name) # colmap doesn't have gt, so just test.

main_dnerf.py

@@ -38,7 +38,6 @@
     # parser.add_argument('--tcnn', action='store_true', help="use TCNN backend")
 
     ### dataset options
-    parser.add_argument('--mode', type=str, default='colmap', help="dataset mode, supports (colmap, blender)")
     parser.add_argument('--color_space', type=str, default='srgb', help="Color space, supports (linear, srgb)")
     parser.add_argument('--preload', action='store_true', help="preload all data into GPU, accelerate training but use more GPU memory")
     # (the default value is for the fox dataset)
@@ -103,7 +102,7 @@
         else:
             test_loader = NeRFDataset(opt, device=device, type='test').dataloader()
 
-            if opt.mode == 'blender':
+            if test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.
@@ -134,7 +133,7 @@
             # also test
             test_loader = NeRFDataset(opt, device=device, type='test').dataloader()
 
-            if opt.mode == 'blender':
+            if test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.

main_nerf.py

@@ -37,7 +37,6 @@
     parser.add_argument('--tcnn', action='store_true', help="use TCNN backend")
 
     ### dataset options
-    parser.add_argument('--mode', type=str, default='colmap', help="dataset mode, supports (colmap, blender)")
     parser.add_argument('--color_space', type=str, default='srgb', help="Color space, supports (linear, srgb)")
     parser.add_argument('--preload', action='store_true', help="preload all data into GPU, accelerate training but use more GPU memory")
     # (the default value is for the fox dataset)
@@ -112,7 +111,7 @@
         else:
             test_loader = NeRFDataset(opt, device=device, type='test').dataloader()
 
-            if opt.mode == 'blender':
+            if test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.
@@ -143,7 +142,7 @@
             # also test
             test_loader = NeRFDataset(opt, device=device, type='test').dataloader()
 
-            if opt.mode == 'blender':
+            if test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.

main_tensoRF.py

@@ -37,7 +37,6 @@
     parser.add_argument("--upsample_model_steps", type=int, action="append", default=[2000, 3000, 4000, 5500, 7000])
 
     ### dataset options
-    parser.add_argument('--mode', type=str, default='colmap', help="dataset mode, supports (colmap, blender)")
     parser.add_argument('--color_space', type=str, default='srgb', help="Color space, supports (linear, srgb)")
     parser.add_argument('--preload', action='store_true', help="preload all data into GPU, accelerate training but use more GPU memory")
     # (the default value is for the fox dataset)
@@ -104,7 +103,7 @@
         else:
             test_loader = NeRFDataset(opt, device=device, type='test').dataloader()
 
-            if opt.mode == 'blender':
+            if test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.
@@ -140,7 +139,7 @@
             # also test
             test_loader = NeRFDataset(opt, device=device, type='test').dataloader()
 
-            if opt.mode == 'blender':
+            if test_loader.has_gt:
                 trainer.evaluate(test_loader) # blender has gt, so evaluate it.
             else:
                 trainer.test(test_loader) # colmap doesn't have gt, so just test.

nerf/provider.py

@@ -94,7 +94,6 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
         self.type = type # train, val, test
         self.downscale = downscale
         self.root_path = opt.path
-        self.mode = opt.mode # colmap, blender
         self.preload = opt.preload # preload data into GPU
         self.scale = opt.scale # camera radius scale to make sure camera are inside the bounding box.
         self.bound = opt.bound # bounding box half length, also used as the radius to random sample poses.
@@ -105,6 +104,12 @@ def __init__(self, opt, device, type='train', downscale=1, n_test=10):
 
         self.rand_pose = opt.rand_pose
 
+        # auto-detect transforms.json and split mode.
+        if os.path.exists(os.path.join(self.root_path, 'transforms.json')):
+            self.mode = 'colmap' # manually split, use view-interpolation for test.
+        elif os.path.exists(os.path.join(self.root_path, 'transforms_train.json')):
+            self.mode = 'blender' # provided split
+
         # load nerf-compatible format data.
         if self.mode == 'colmap':
             with open(os.path.join(self.root_path, 'transforms.json'), 'r') as f:
@@ -314,4 +319,5 @@ def dataloader(self):
             size += size // self.rand_pose # index >= size means we use random pose.
         loader = DataLoader(list(range(size)), batch_size=1, collate_fn=self.collate, shuffle=self.training, num_workers=0)
         loader._data = self # an ugly fix... we need to access error_map & poses in trainer.
+        loader.has_gt = self.images is not None
         return loader

readme.md

@@ -66,7 +66,7 @@ First time running will take some time to compile the CUDA extensions.
 ```bash
 ### Instant-ngp NeRF
 # train with different backbones (with slower pytorch ray marching)
-# for the colmap dataset, the default dataset setting `--mode colmap --bound 2 --scale 0.33` is used.
+# for the colmap dataset, the default dataset setting `--bound 2 --scale 0.33` is used.
 python main_nerf.py data/fox --workspace trial_nerf # fp32 mode
 python main_nerf.py data/fox --workspace trial_nerf --fp16 # fp16 mode (pytorch amp)
 python main_nerf.py data/fox --workspace trial_nerf --fp16 --ff # fp16 mode + FFMLP (this repo's implementation)
@@ -97,13 +97,12 @@ python main_nerf.py data/fox --workspace trial_nerf -O --gui
 # test mode for GUI
 python main_nerf.py data/fox --workspace trial_nerf -O --gui --test
 
-# for the blender dataset, you should add `--mode blender --bound 1.0 --scale 0.8 --dt_gamma 0 --color_space linear`
-# --mode specifies dataset type ('blender' or 'colmap')
+# for the blender dataset, you should add `--bound 1.0 --scale 0.8 --dt_gamma 0 --color_space linear`
 # --bound means the scene is assumed to be inside box[-bound, bound]
 # --scale adjusts the camera locaction to make sure it falls inside the above bounding box. 
 # --dt_gamma controls the adaptive ray marching speed, set to 0 turns it off.
-python main_nerf.py data/nerf_synthetic/lego --workspace trial_nerf -O --mode blender --bound 1.0 --scale 0.8 --dt_gamma 0 --color_space linear
-python main_nerf.py data/nerf_synthetic/lego --workspace trial_nerf -O --mode blender --bound 1.0 --scale 0.8 --dt_gamma 0 --color_space linear --gui
+python main_nerf.py data/nerf_synthetic/lego --workspace trial_nerf -O --bound 1.0 --scale 0.8 --dt_gamma 0 --color_space linear
+python main_nerf.py data/nerf_synthetic/lego --workspace trial_nerf -O --bound 1.0 --scale 0.8 --dt_gamma 0 --color_space linear --gui
 
 # for the LLFF dataset, you should first convert it to nerf-compatible format:
 python scripts/llff2nerf.py data/nerf_llff_data/fern # by default it use full-resolution images, and write `transforms.json` to the folder
@@ -115,8 +114,8 @@ python main_nerf.py data/nerf_llff_data/fern --workspace trial_nerf -O --gui
 # for the Tanks&Temples dataset, you should first convert it to nerf-compatible format:
 python scripts/tanks2nerf.py data/TanksAndTemple/Family # write `trainsforms_{split}.json` for [train, val, test]
 # then you can train as a blender dataset (you'll need to tune the scale & bound if necessary)
-python main_nerf.py data/TanksAndTemple/Family --workspace trial_nerf_family -O --mode blender --bound 1.0 --scale 0.33 --dt_gamma 0
-python main_nerf.py data/TanksAndTemple/Family --workspace trial_nerf_family -O --mode blender --bound 1.0 --scale 0.33 --dt_gamma 0 --gui
+python main_nerf.py data/TanksAndTemple/Family --workspace trial_nerf_family -O --bound 1.0 --scale 0.33 --dt_gamma 0
+python main_nerf.py data/TanksAndTemple/Family --workspace trial_nerf_family -O --bound 1.0 --scale 0.33 --dt_gamma 0 --gui
 
 # for custom dataset, you should:
 # 1. take a video / many photos from different views 
@@ -138,22 +137,22 @@ python main_sdf.py data/armadillo.obj --workspace trial_sdf --fp16 --test
 ### TensoRF
 # almost the same as Instant-ngp NeRF, just replace the main script.
 python main_tensoRF.py data/fox --workspace trial_tensoRF -O
-python main_tensoRF.py data/nerf_synthetic/lego --workspace trial_tensoRF -O --mode blender --bound 1.0 --scale 0.8 --dt_gamma 0 
+python main_tensoRF.py data/nerf_synthetic/lego --workspace trial_tensoRF -O --bound 1.0 --scale 0.8 --dt_gamma 0 
 
 ### CCNeRF
 # training on single objects, turn on --error_map for better quality.
-python main_CCNeRF.py data/nerf_synthetic/chair --workspace trial_cc_chair -O --mode blender --bound 1.0 --scale 0.67 --dt_gamma 0 --error_map
-python main_CCNeRF.py data/nerf_synthetic/ficus --workspace trial_cc_ficus -O --mode blender --bound 1.0 --scale 0.67 --dt_gamma 0 --error_map
-python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O --mode blender --bound 1.0 --scale 0.67 --dt_gamma 0 --error_map
+python main_CCNeRF.py data/nerf_synthetic/chair --workspace trial_cc_chair -O --bound 1.0 --scale 0.67 --dt_gamma 0 --error_map
+python main_CCNeRF.py data/nerf_synthetic/ficus --workspace trial_cc_ficus -O --bound 1.0 --scale 0.67 --dt_gamma 0 --error_map
+python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O --bound 1.0 --scale 0.67 --dt_gamma 0 --error_map
 # compose, use a larger bound and more samples per ray for better quality.
-python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O --mode blender --bound 2.0 --scale 0.67 --dt_gamma 0 --max_steps 2048 --test --compose
+python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O --bound 2.0 --scale 0.67 --dt_gamma 0 --max_steps 2048 --test --compose
 # compose + gui, only about 1 FPS without dynamic resolution... just for quick verification of composition results.
-python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O --mode blender --bound 2.0 --scale 0.67 --dt_gamma 0 --test --compose --gui
+python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O --bound 2.0 --scale 0.67 --dt_gamma 0 --test --compose --gui
 
 ### D-NeRF
 # almost the same as Instant-ngp NeRF, just replace the main script.
-python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0 --mode blender
-python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0 --mode blender --gui
+python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0
+python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0 --gui
 ```
 
 check the `scripts` directory for more provided examples.

scripts/colmap2nerf.py

@@ -10,7 +10,7 @@
 
 import argparse
 import os
-from pathlib import Path, PurePosixPath
+from pathlib import Path
 
 import numpy as np
 import json
@@ -128,12 +128,15 @@ def qvec2rotmat(qvec):
     ])
 
 def rotmat(a, b):
-    a, b = a / np.linalg.norm(a), b / np.linalg.norm(b)
-    v = np.cross(a, b)
-    c = np.dot(a, b)
-    s = np.linalg.norm(v)
-    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
-    return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2 + 1e-10))
+	a, b = a / np.linalg.norm(a), b / np.linalg.norm(b)
+	v = np.cross(a, b)
+	c = np.dot(a, b)
+	# handle exception for the opposite direction input
+	if c < -1 + 1e-10:
+		return rotmat(a + np.random.uniform(-1e-2, 1e-2, 3), b)
+	s = np.linalg.norm(v)
+	kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
+	return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2 + 1e-10))
 
 def closest_point_2_lines(oa, da, ob, db): # returns point closest to both rays of form o+t*d, and a weight factor that goes to 0 if the lines are parallel
     da = da / np.linalg.norm(da)
@@ -271,6 +274,7 @@ def closest_point_2_lines(oa, da, ob, db): # returns point closest to both rays
                 t = tvec.reshape([3,1])
                 m = np.concatenate([np.concatenate([R, t], 1), bottom], 0)
                 c2w = np.linalg.inv(m)
+
                 c2w[0:3, 2] *= -1 # flip the y and z axis
                 c2w[0:3, 1] *= -1
                 c2w = c2w[[1, 0, 2, 3],:] # swap y and z