New datasets framework (#111)

Datasets have been reformatted in this MR to support easy adding of new datasets. * All datasets should subclass WispDataset, which returns a `Batch` instance from `__get_item__` * Underneath WispDataset, abstract MultiviewDataset and SDFDataset define more specific interfaces. * NeRFSyntheticDataset and RTMVDataset are new concrete classes, which return a `MultiviewBatch`. * Previous SDFDataset separated to concrete classes MeshSampledSDFDataset and OctreeSampledSDFDataset, depending on whether samples are generated directly from mesh or limited by the octree narrowband. Both return a `SDFBatch`. * `load_multiview_dataset` is a convenience method that looks at a given path and tries to determine the correct `MultiviewDataset` subclass to create (out of all available dataset subclasses). This relies on the optional `is_root_of_dataset()` @classmethod each dataset may implement to distinguish itself from other formats. Signed-off-by: operel <[email protected]>
NVIDIAGameWorks · Jan 31, 2023 · 368bd52 · 368bd52
1 parent 170b84e
commit 368bd52
Show file tree

Hide file tree

Showing 26 changed files with 2,060 additions and 1,187 deletions.
diff --git a/README.md b/README.md
@@ -16,6 +16,7 @@ For an overview on neural fields, we recommend you check out the EG STAR report:
 
 ## Latest Updates
 
+* _01/02/23_ `attrdict` dependency added as part of the new datasets framework. If you pull latest, make sure to `pip install attrdict`.
 * _17/01/23_ `pycuda` replaced with `cuda-python`. Wisp can be installed from pip now  (If you pull, run **pip install -r requirements_app.txt**)
 * _05/01/23_ Mains are now introduced as standalone apps, for easier support of new pipelines (**breaking change**) 
 * _21/12/22_ Most modules have been cleaned, reorganized and documented. 

diff --git a/app/nerf/main_nerf.py b/app/nerf/main_nerf.py
@@ -12,11 +12,11 @@
 import logging
 import numpy as np
 import torch
+import wisp
 from wisp.app_utils import default_log_setup, args_to_log_format
 import wisp.config_parser as config_parser
 from wisp.framework import WispState
-from wisp.datasets import MultiviewDataset
-from wisp.datasets.transforms import SampleRays
+from wisp.datasets import MultiviewDataset, SampleRays
 from wisp.models.grids import BLASGrid, OctreeGrid, CodebookOctreeGrid, TriplanarGrid, HashGrid
 from wisp.tracers import BaseTracer, PackedRFTracer
 from wisp.models.nefs import BaseNeuralField, NeuralRadianceField
@@ -240,7 +240,7 @@ def parse_args():
     return args, args_dict
 
 
-def load_dataset(args) -> torch.utils.data.Dataset:
+def load_dataset(args) -> MultiviewDataset:
     """ Loads a multiview dataset comprising of pairs of images and calibrated cameras.
     The types of supported datasets are defined by multiview_dataset_format:
     'standard' - refers to the standard NeRF format popularized by Mildenhall et al. 2020,
@@ -250,16 +250,19 @@ def load_dataset(args) -> torch.utils.data.Dataset:
             This dataset includes depth information which allows for performance improving optimizations in some cases.
     """
     transform = SampleRays(num_samples=args.num_rays_sampled_per_img)
-    train_dataset = MultiviewDataset(dataset_path=args.dataset_path,
-                                     multiview_dataset_format=args.multiview_dataset_format,
-                                     mip=args.mip,
-                                     bg_color=args.bg_color,
-                                     dataset_num_workers=args.dataset_num_workers,
-                                     transform=transform)
-    return train_dataset
-
-
-def load_grid(args, dataset: torch.utils.data.Dataset) -> BLASGrid:
+    train_dataset = wisp.datasets.load_multiview_dataset(dataset_path=args.dataset_path,
+                                                         split='train',
+                                                         mip=args.mip,
+                                                         bg_color=args.bg_color,
+                                                         dataset_num_workers=args.dataset_num_workers,
+                                                         transform=transform)
+    validation_dataset = None
+    if args.valid_every > -1 or args.valid_only:
+        validation_dataset = train_dataset.create_split(split='val', transform=None)
+    return train_dataset, validation_dataset
+
+
+def load_grid(args, dataset: MultiviewDataset) -> BLASGrid:
     """ Wisp's implementation of NeRF uses feature grids to improve the performance and quality (allowing therefore,
     interactivity).
     This function loads the feature grid to use within the neural pipeline.
@@ -269,13 +272,12 @@ def load_grid(args, dataset: torch.utils.data.Dataset) -> BLASGrid:
     See corresponding grid constructors for each of their arg details.
     """
     grid = None
-    # Optimization: For octrees based grids, if dataset contains depth info, initialize only cells known to be occupied
-    has_depth_supervision = getattr(dataset, "coords", None) is not None
 
+    # Optimization: For octrees based grids, if dataset contains depth info, initialize only cells known to be occupied
     if args.grid_type == "OctreeGrid":
-        if has_depth_supervision:
+        if dataset.supports_depth():
             grid = OctreeGrid.from_pointcloud(
-                pointcloud=dataset.coords,
+                pointcloud=dataset.as_pointcloud(),
                 feature_dim=args.feature_dim,
                 base_lod=args.base_lod,
                 num_lods=args.num_lods,
@@ -295,9 +297,9 @@ def load_grid(args, dataset: torch.utils.data.Dataset) -> BLASGrid:
                 feature_bias=args.feature_bias,
             )
     elif args.grid_type == "CodebookOctreeGrid":
-        if has_depth_supervision:
+        if dataset.supports_depth:
             grid = CodebookOctreeGrid.from_pointcloud(
-                pointcloud=dataset.coords,
+                pointcloud=dataset.as_pointcloud(),
                 feature_dim=args.feature_dim,
                 base_lod=args.base_lod,
                 num_lods=args.num_lods,
@@ -359,7 +361,7 @@ def load_grid(args, dataset: torch.utils.data.Dataset) -> BLASGrid:
     return grid
 
 
-def load_neural_field(args, dataset: torch.utils.data.Dataset) -> BaseNeuralField:
+def load_neural_field(args, dataset: MultiviewDataset) -> BaseNeuralField:
     """ Creates a "Neural Field" instance which converts input coordinates to some output signal.
     Here a NeuralRadianceField is created, which maps 3D coordinates (+ 2D view direction) -> RGB + density.
     The NeuralRadianceField uses spatial feature grids internally for faster feature interpolation and raymarching.
@@ -416,7 +418,7 @@ def load_neural_pipeline(args, dataset, device) -> Pipeline:
     return pipeline
 
 
-def load_trainer(pipeline, train_dataset, device, scene_state, args, args_dict) -> BaseTrainer:
+def load_trainer(pipeline, train_dataset, validation_dataset, device, scene_state, args, args_dict) -> BaseTrainer:
     """ Loads the NeRF trainer.
     The trainer is responsible for managing the optimization life-cycles and can be operated in 2 modes:
     - Headless, which will run the train() function until all training steps are exhausted.
@@ -431,7 +433,8 @@ def load_trainer(pipeline, train_dataset, device, scene_state, args, args_dict)
     optimizer_params = config_parser.get_args_for_function(args, optimizer_cls)
 
     trainer = MultiviewTrainer(pipeline=pipeline,
-                               dataset=train_dataset,
+                               train_dataset=train_dataset,
+                               validation_dataset=validation_dataset,
                                num_epochs=args.epochs,
                                batch_size=args.batch_size,
                                optim_cls=optimizer_cls,
@@ -478,10 +481,12 @@ def is_interactive() -> bool:
 default_log_setup(args.log_level)
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-train_dataset = load_dataset(args=args)
+train_dataset, validation_dataset = load_dataset(args=args)
 pipeline = load_neural_pipeline(args=args, dataset=train_dataset, device=device)
 scene_state = WispState()   # Joint trainer / app state
-trainer = load_trainer(pipeline=pipeline, train_dataset=train_dataset, device=device, scene_state=scene_state,
+trainer = load_trainer(pipeline=pipeline,
+                       train_dataset=train_dataset, validation_dataset=validation_dataset,
+                       device=device, scene_state=scene_state,
                        args=args, args_dict=args_dict)
 app = load_app(args=args, scene_state=scene_state, trainer=trainer)
 

diff --git a/app/nglod/main_nglod.py b/app/nglod/main_nglod.py
@@ -14,7 +14,8 @@
 from wisp.app_utils import default_log_setup, args_to_log_format
 import wisp.config_parser as config_parser
 from wisp.framework import WispState
-from wisp.datasets import SDFDataset
+from wisp.datasets import SDFDataset, MeshSampledSDFDataset, OctreeSampledSDFDataset
+from wisp.accelstructs import OctreeAS
 from wisp.models.grids import BLASGrid, OctreeGrid, CodebookOctreeGrid, TriplanarGrid, HashGrid
 from wisp.tracers import BaseTracer, PackedSDFTracer
 from wisp.models.nefs import BaseNeuralField, NeuralSDF
@@ -237,24 +238,31 @@ def parse_args():
     return args, args_dict
 
 
-def load_dataset(args, pipeline: Pipeline) -> torch.utils.data.Dataset:
+def load_dataset(args, pipeline: Pipeline) -> SDFDataset:
     """ Loads a dataset of SDF samples generated over the surface of a mesh. """
-    if isinstance(pipeline.nef.grid, OctreeGrid):
-        train_dataset = SDFDataset.from_grid(
+    if OctreeSampledSDFDataset.supports_blas(pipeline.nef.grid.blas):
+        # The current grid representation uses a bottom-level acceleration structure which can be used for faster
+        # and more precise resampling of sdf values
+        train_dataset = OctreeSampledSDFDataset(
+            occupancy_struct=pipeline.nef.grid.blas,
+            split='train',
+            transform=None,
             sample_mode=args.sample_mode,
             num_samples=args.num_samples,
-            get_normals=args.get_normals,
             sample_tex=args.sample_tex,
-            grid=pipeline.nef.grid,
-            samples_per_voxel=args.samples_per_voxel)
+            samples_per_voxel=args.samples_per_voxel
+        )
     else:
-        train_dataset = SDFDataset.from_mesh(
+        train_dataset = MeshSampledSDFDataset(
+            dataset_path=args.dataset_path,
+            split='train',
+            transform=None,
             sample_mode=args.sample_mode,
             num_samples=args.num_samples,
             get_normals=args.get_normals,
             sample_tex=args.sample_tex,
-            dataset_path=args.dataset_path,
-            mode_norm=args.mode_mesh_norm)
+            mode_norm=args.mode_mesh_norm
+        )
     return train_dataset
 
 
@@ -380,7 +388,7 @@ def load_trainer(pipeline, train_dataset, device, scene_state, args, args_dict)
     optimizer_cls = config_parser.get_module(name=args.optimizer_type)
     optimizer_params = config_parser.get_args_for_function(args, optimizer_cls)
     trainer = SDFTrainer(pipeline=pipeline,
-                         dataset=train_dataset,
+                         train_dataset=train_dataset,
                          num_epochs=args.epochs,
                          batch_size=args.batch_size,
                          optim_cls=optimizer_cls,

diff --git a/examples/latent_nerf/main_demo.py b/examples/latent_nerf/main_demo.py
@@ -13,8 +13,7 @@
 import torch
 from wisp.app_utils import default_log_setup, args_to_log_format
 from wisp.framework import WispState
-from wisp.datasets import MultiviewDataset
-from wisp.datasets.transforms import SampleRays
+from wisp.datasets import NeRFSyntheticDataset, SampleRays
 from wisp.trainers import MultiviewTrainer
 from wisp.models.grids import OctreeGrid
 from wisp.models.pipeline import Pipeline
@@ -37,12 +36,12 @@
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # NeRF is trained with a MultiviewDataset, which knows how to generate RGB rays from a set of images + cameras
-train_dataset = MultiviewDataset(
+train_dataset = NeRFSyntheticDataset(
     dataset_path=args.dataset_path,
-    multiview_dataset_format='standard',
+    split='train',
     mip=0,
     bg_color='black',
-    dataset_num_workers=-1,
+    dataset_num_workers=0,
     transform=SampleRays(
         num_samples=4096
     )
@@ -74,7 +73,7 @@
 weight_decay = 0
 exp_name = 'siggraph_2022_demo'
 trainer = MultiviewTrainer(pipeline=pipeline,
-                           dataset=train_dataset,
+                           train_dataset=train_dataset,
                            num_epochs=args.epochs,
                            batch_size=1,    # 1 image per batch
                            optim_cls=torch.optim.RMSprop,

diff --git a/requirements.txt b/requirements.txt
@@ -24,4 +24,5 @@ PyDispatcher
 pynvml
 setuptools==59.5.0
 wandb>=0.13.5
-pytest>=7.1.0
+pytest>=7.1.0
+attrdict
diff --git a/tests/apps/test_nerf.py b/tests/apps/test_nerf.py
@@ -144,7 +144,7 @@ def test_triplanar_lego(self, lego_path, dataset_num_workers):
         out = run_wisp_script(cmd, cli_args)
         metrics = collect_metrics_from_log(out, ['PSNR'])
 
-        assert float(metrics[100]['PSNR']) > 30.5, 'PSNR is too low.'
+        assert float(metrics[100]['PSNR']) > 30.4, 'PSNR is too low.'
         report_metrics(metrics)  # Prints to log
 
     def test_codebook_V8(self, V8_path, dataset_num_workers):

diff --git a/wisp/config_parser.py b/wisp/config_parser.py
@@ -83,12 +83,14 @@ def get_module(name, module_type=None):
 def get_args_for_function(args, func):
     """ Given a func (for example an __init__(..) function or from_X(..)), and also the parsed args,
     return the subset of args that func expects and args contains. """
+    if isinstance(args, argparse.Namespace):
+        args = vars(args)   # Namespace -> dict
     has_kwargs = inspect.getfullargspec(func).varkw != None
     if has_kwargs:
-        collected_args = vars(args)
+        collected_args = args
     else:
         parameters = dict(inspect.signature(func).parameters)
-        collected_args = {a: getattr(args, a) for a in parameters if hasattr(args, a)}
+        collected_args = {a: args[a] for a in parameters if a in args}
     return collected_args
 
 

diff --git a/wisp/datasets/__init__.py b/wisp/datasets/__init__.py
@@ -6,7 +6,7 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
 
-from .sdf_dataset import SDFDataset
-from .multiview_dataset import MultiviewDataset
-from .random_view_dataset import RandomViewDataset
-from .utils import default_collate
+from .base_datasets import *
+from .utils import *
+from .formats import *
+from .transforms import *