Implemented Plenoxel. (#1)

* add codes to run the blender dataset

* change the default batch size

* add white_bkgd option

* corrct the ndc coordiinate

* change the sampling strategy

* add dataloader for tanks and temples

* add tanks and temples

* change the batch size for a fair comparison

* hotfix: minor bug

* hotfix: change the upsampe frequency

* add torch.no_grad

* add codes to load model

* modify configs to run sweep

* change the config to run sweep

* hotfix: fix loading error

* hotfix: fix typo

* hotfix: minor fix

* hotfix: change the expname

* hotfix: change the expname

Author: jeongyw12382

.gitignore

@@ -6,4 +6,8 @@ logs
 *pycache*
 *swap-pane*
 *.idea*
-logs_collection
+logs_collection
+*build*
+*.egg-info*
+*.eggs*
+*_debug*

README.md

@@ -13,14 +13,18 @@ This project is created and maintained by Yoonwoo Jeong, Jinoh Cho, and Seungjoo
     - Several options are different from the original NeRF (TensorFlow)
     - NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis: [[Code](https://github.com/google-research/google-research/tree/master/jaxnerf) | [Paper](https://arxiv.org/pdf/2003.08934) | [Arxiv](https://arxiv.org/abs/2003.08934)]
     - Document: [Link](./docs/jaxnerf.md)
-- SNeRG (Torch Version)
+- SNeRG (Torch)
     - Coming Soon
-- MipNeRF (Torch Version)
+- MipNeRF (Torch)
     - Coming Soon
-- MipNeRF360 (Torch Version)
+- MipNeRF360 (Torch)
     - Coming Soon
-- NeRF++ (Torch Version)
+- NeRF++ (Torch)
     - Coming Soon
+- Plenoxel (Torch)
+    - Reorganized the code structure for easier modification. 
+    - Plenoxels: Radiance Fields without Neural Networks [[Code](https://github.com/sxyu/svox2) | [Paper]()]
+    - Document: [Link](./docs/plenoxel.md)
 
 ## Dataset
 ```

config.py

@@ -1,5 +1,15 @@
 import argparse
 
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise Exception("Boolean value expected.")
+        
 def config_parser():
 
     parser = argparse.ArgumentParser()
@@ -41,7 +51,11 @@ def config_parser():
 
     ray = parser.add_argument_group("rays") 
     ray.add_argument(
-        "--use_pixel_centers", action="store_true", default=False,
+        "--use_pixel_centers", 
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False,
         help="add a half pixel while generating rays"
     )
 
@@ -55,19 +69,28 @@ def config_parser():
         "--testskip", type=int, default=8,
         help="will load 1/N images from test/val sets, useful for large datasets like deepvoxels",
     )
+    dataset.add_argument(
+        "--scene_scale", type=float, default=1,
+        help="resize the scale of scenes"
+    )
+    dataset.add_argument(
+        "--shuffle_train", 
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False,
+        help="shuffle the train set"
+    )
 
     ## blender flags
     blender = parser.add_argument_group("blender")
     blender.add_argument(
         "--white_bkgd",
-        action="store_true",
+        type=str2bool,
+        nargs="?",
+        const=True,
         help="set to render synthetic data on a white bkgd (always use for dvoxels)",
     )
-    blender.add_argument(
-        "--half_res",
-        action="store_true",
-        help="load blender synthetic data at 400x400 instead of 800x800",
-    )
 
     ## llff flags
     llff = parser.add_argument_group("llff flags")
@@ -76,16 +99,23 @@ def config_parser():
     )
     llff.add_argument(
         "--no_ndc",
-        action="store_true",
+        type=str2bool,
+        nargs="?",
+        const=True,
         help="do not use normalized device coordinates (set for non-forward facing scenes)",
     )
     llff.add_argument(
         "--lindisp",
-        action="store_true",
+        type=str2bool,
+        nargs="?",
+        const=True,
         help="sampling linearly in disparity rather than depth",
     )
     llff.add_argument(
-        "--spherify", action="store_true", help="set for spherical 360 scenes"
+        "--spherify",   
+        type=str2bool,
+        nargs="?",
+        const=True,help="set for spherical 360 scenes"
     )
     llff.add_argument(
         "--llffhold",
@@ -100,20 +130,20 @@ def config_parser():
         "--i_print", type=int, default=1000,
         help="frequency of console printout and metric logging",
     )
-    metadata.add_argument(
-        "--i_weights", type=int, default=50000,
-        help="frequency of storing weights"
-    )
     metadata.add_argument(
         "--i_validation", type=int, default=50000,
         help="frequency of validation"
     )
     metadata.add_argument(
-        "--debug", action="store_true", default=False,
+        "--debug", 
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False,
         help="run with debug mode"
     )
     metadata.add_argument(
-        "--expname", type=str, required=True, help="experiment name"
+        "--expname", type=str, default=None, help="experiment name"
     )
     metadata.add_argument(
         "--basedir", type=str, default="./logs/",
@@ -126,22 +156,49 @@ def config_parser():
 
     runmode = parser.add_argument_group("running mode")
     runmode.add_argument(
-        "--train", action="store_true", default=False, help="run with train mode"
+        "--train",  
+        type=str2bool,
+        nargs="?",
+        const=True, 
+        default=False, 
+        help="run with train mode"
     )
     runmode.add_argument(
-        "--eval", action="store_true", default=False, help="run with eval mode"
+        "--eval",    
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False, 
+        help="run with eval mode"
     )
     runmode.add_argument(
-        "--bake", action="store_true", default=False, help="bake the trained model"
+        "--bake",         
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False, 
+        help="bake the trained model"
     )
     runmode.add_argument(
-        "--render", action="store_true", default=False, help="render to generate video"
+        "--render", 
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False, help="render to generate video"
     )
     runmode.add_argument(
-        "--skip_validation", action="store_true", default=False, 
+        "--skip_validation",        
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False, 
     )
     runmode.add_argument(
-        "--tpu", action="store_true", default=False, help="run with tpus"
+        "--tpu", 
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False, help="run with tpus"
     )
     runmode.add_argument(
         "--tpu_num", type=int, default=0, help="number of tpu"
@@ -152,10 +209,26 @@ def config_parser():
     runmode.add_argument(
         "--seed", type=int, default=0, help="seed to fix"
     )
+    runmode.add_argument(
+        "--use_custom_optim",         
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False, 
+        help="Run with a custom optimization step"
+    )
+    runmode.add_argument(
+        "--run_large_model",         
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False,
+        help="For wandb sweep: run large nerf model"
+    )
 
     config = parser.add_argument_group("config")
     config.add_argument(
         "--config", type=str, default=None, help="path to config file"
     )
 
-    return parser.parse_args(), parser
+    return parser.parse_known_args()[0], parser

configs/jaxnerf_torch/blender.yaml

@@ -17,7 +17,6 @@ lr_delay_mult: 0.2
 max_steps: 200000
 
 i_print: 100
-i_weights: 50000
 i_validation: 50000
 use_pixel_centers: true
 

configs/jaxnerf_torch/blender_large.yaml

@@ -17,7 +17,6 @@ lr_delay_mult: 0.01
 max_steps: 200000
 
 i_print: 100
-i_weights: 50000
 i_validation: 50000
 use_pixel_centers: true
 

configs/jaxnerf_torch/llff.yaml

@@ -19,7 +19,6 @@ lr_delay_mult: 1.
 max_steps: 200000
 
 i_print: 100
-i_weights: 50000
 i_validation: 50000
 use_pixel_centers: true
 

configs/jaxnerf_torch/llff_large.yaml

@@ -19,7 +19,6 @@ lr_delay_mult: 1.
 max_steps: 200000
 
 i_print: 100
-i_weights: 50000
 i_validation: 50000
 use_pixel_centers: true
 

configs/plenoxel_torch/blender.yaml

@@ -0,0 +1,23 @@
+dataset_type: blender
+model: plenoxel_torch
+batching: all_images_wo_replace
+
+N_rand: 4096
+
+max_steps: 200000
+
+i_print: 100
+i_validation: 50000
+
+no_ndc: true
+
+reso: "[[256, 256, 256], [512, 512, 512]]"
+upsamp_every: 50000
+lr_sigma: 3.0e+1
+lr_sh: 1.0e-2
+lambda_tv: 1.0e-5
+lambda_tv_sh: 1.0e-3
+
+use_pixel_centers: true
+white_bkgd: true
+scene_scale: 0.6666

configs/plenoxel_torch/llff.yaml

@@ -0,0 +1,28 @@
+dataset_type: llff
+model: plenoxel_torch
+batching: all_images_wo_replace
+
+N_rand: 4096
+max_steps: 200000
+
+llffhold: 8
+factor: 4
+
+i_print: 100
+i_validation: 50000
+use_pixel_centers: true
+
+no_ndc: false
+
+reso: "[[256, 256, 128], [512, 512, 128], [1408, 1156, 128]]"
+upsamp_every: 50000
+lr_sigma: 3.0e+1
+lr_sh: 1.0e-2
+thresh_type: "sigma"
+density_thresh: 5.0
+lambda_tv: 5.0e-4
+lambda_tv_sh: 5.0e-3
+lambda_sparsity: 1.0e-12
+background_brightness: 0.5
+last_sample_opaque: false,
+tv_early_only: 0

configs/plenoxel_torch/tanks_and_temples.yaml

@@ -0,0 +1,30 @@
+dataset_type: tanks_and_temples
+model: plenoxel_torch
+batching: all_images
+
+N_rand: 4096
+max_steps: 200000
+
+i_print: 100
+i_validation: 50000
+use_pixel_centers: true
+
+no_ndc: true
+chunk: 16384
+
+upsamp_every: 50000
+reso: "[[128, 128, 128], [256, 256, 256], [512, 512, 512], [640, 640, 640]]"
+background_nlayers: 64
+background_reso: 1024
+lr_sigma: 3.0e+1
+lr_sh: 1.0e-2
+lr_sigma_delay_steps: 15000
+thresh_type: "weight"
+lambda_tv": 5.0e-5
+lambda_tv_sh": 5.0e-3
+lambda_tv_background_sigma": 1.0e-3
+lambda_tv_background_color: 1.0e-3
+lambda_beta: 1.0e-5
+lambda_sparsity: 1.0e-11
+background_brightness: 1.0
+tv_early_only: 0