From f910836b8c39c092e7537fe238b7d19275fe7ae2 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Thu, 1 Sep 2022 20:36:54 +0100 Subject: [PATCH] Cleanup (#325) cleanup --- .vscode/launch.json | 4 +- README.md | 40 +++++++++---------- .../visualize_spatial_distortions.ipynb | 4 +- docs/models/index.md | 4 +- docs/models/nerf.md | 11 +++-- docs/reference/contributing.md | 8 ++-- nerfactory/dataloaders/base.py | 34 +++++++++++----- nerfactory/dataloaders/eval.py | 39 ++++++++++++++---- nerfactory/dataloaders/structs.py | 40 ------------------- nerfactory/engine/trainer.py | 3 ++ nerfactory/models/instant_ngp.py | 2 + nerfactory/models/nerfw.py | 2 + nerfactory/pipelines/base.py | 3 ++ 13 files changed, 99 insertions(+), 95 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 453912d23..e8a656fbd 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -57,9 +57,7 @@ "request": "launch", "program": "scripts/run_train.py", "console": "integratedTerminal", - "args": [ - "tensorf" - ] + "args": ["tensorf"] }, { "name": "Python: run_train nerfw (gpu)", diff --git a/README.md b/README.md index 9efa2a14c..6d6ac9279 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,3 @@ -

- - - - nerfactory - -

- -

The all-in-one repo for NeRFs

-

Test Status @@ -26,9 +16,27 @@ PyPI version

+

+ + + + nerfactory + +

+ +

The all-in-one repo for NeRFs

+ +

+ + documentation + + + viewer + +

+ - [Quickstart](#quickstart) - [Supported Features](#supported-features) -- [Benchmarked Model Architectures](#benchmarked-model-architectures) # Quickstart @@ -177,8 +185,6 @@ We provide the following support strucutures to make life easier for getting sta If you are looking for a feature that is not currently supported, please do not hesitate to contact the Plenoptix team! -#### :metal: Support for [Hydra](https://hydra.cc/) config structure - #### :metal: Support for multiple logging interfaces #### :metal: Built-in support for profiling code @@ -186,11 +192,3 @@ If you are looking for a feature that is not currently supported, please do not #### :metal: Benchmarking scripts #### :metal: Speed up your code with Tiny Cuda NN - -# Benchmarked Model Architectures - -| Method | PSNR | -| --------------------------------------------------------------------------------- | ------------------------ | -| [NeRF](https://arxiv.org/abs/2003.08934) | :hourglass_flowing_sand: | -| [instant NGP](https://nvlabs.github.io/instant-ngp/assets/mueller2022instant.pdf) | :hourglass_flowing_sand: | -| [Mip NeRF](https://arxiv.org/abs/2103.13415) | :hourglass_flowing_sand: | diff --git a/docs/model_components/visualize_spatial_distortions.ipynb b/docs/model_components/visualize_spatial_distortions.ipynb index 6c196f564..08bd68576 100644 --- a/docs/model_components/visualize_spatial_distortions.ipynb +++ b/docs/model_components/visualize_spatial_distortions.ipynb @@ -6,9 +6,9 @@ "source": [ "# Spatial Distortions\n", "\n", - "If you are trying to reconstruct and object floating in an empty void, you can stop reading. However if you are trying to reconstruct a scene or object from images, you may wish to consider adding a spatial distortion.\n", + "If you are trying to reconstruct an object floating in an empty void, you can stop reading. However if you are trying to reconstruct a scene or object from images, you may wish to consider adding a spatial distortion.\n", "\n", - "When rendering a target view of a scene, the camera will emit a camera ray for each pixel and query the scene at points along this ray. We can choose where to query these point using different [samplers](visualize_samplers.ipynb). These samplers have some notion of _bounds_ that define where the ray should start and terminate. If you know that everything in your scenes exists within some predefined bounds (ie. a cube that a room fits in) then the sampler will properly sample the entire space. If however the scene is unbounded (ie. an outdoor scene) defining where to stop sampling is challenging. One option to increase the far sampling distance to a large value (ie. 1km). Alternatively we can warp the space into a fixed volume. Below are supported distortions.\n" + "When rendering a target view of a scene, the camera will emit a camera ray for each pixel and query the scene at points along this ray. We can choose where to query these points using different [samplers](visualize_samplers.ipynb). These samplers have some notion of _bounds_ that define where the ray should start and terminate. If you know that everything in your scenes exists within some predefined bounds (ie. a cube that a room fits in) then the sampler will properly sample the entire space. If however the scene is unbounded (ie. an outdoor scene) defining where to stop sampling is challenging. One option to increase the far sampling distance to a large value (ie. 1km). Alternatively we can warp the space into a fixed volume. Below are supported distortions.\n" ] }, { diff --git a/docs/models/index.md b/docs/models/index.md index 0cf680014..89185ab71 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -1,6 +1,6 @@ # Models -We provide a set of preimplemented nerfactory models. One of the goals of nerfactory is to modularize the various NeRF techniques as much as possible. As a result, many of the techniqes from these preimplemented models can be mixed. +We provide a set of pre implemented nerfactory models. One of the goals of nerfactory is to modularize the various NeRF techniques as much as possible. As a result, many of the techniques from these pre-implemented models can be mixed. ## Running a model @@ -12,7 +12,7 @@ python scripts/run_train.py --config-name MODEL_CONFIG ## Guides -In addition to their implementations, we have provided guides that walk through each of these method. +In addition to their implementations, we have provided guides that walk through each of these method. ```{toctree} :maxdepth: 1 diff --git a/docs/models/nerf.md b/docs/models/nerf.md index 5bf45fdba..3fb426f83 100644 --- a/docs/models/nerf.md +++ b/docs/models/nerf.md @@ -18,7 +18,7 @@ python scripts/run_train.py --config-name=graph_vanilla_nerf.yaml ### Overview -If you have arrived to this site, it is likely that you have atleast heard of NeRFs. This page will discuss the original NeRF paper, _"NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis"_ by Mildenhall, Srinivasan, Tancik et al. (2020). For most tasks, using the original NeRF model is likely not a good choice and hence we provide implementations of various other NeRF related models. It is however useful to understand how NeRF's work as most follow ups follow a similar structure. +If you have arrived to this site, it is likely that you have at least heard of NeRFs. This page will discuss the original NeRF paper, _"NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis"_ by Mildenhall, Srinivasan, Tancik et al. (2020). For most tasks, using the original NeRF model is likely not a good choice and hence we provide implementations of various other NeRF related models. It is however useful to understand how NeRF's work as most follow ups follow a similar structure. The goal is to optimize a volumetric representation of a scene that can be rendered from novel viewpoints. This representation is optimized from a set of images and associated camera poses. @@ -42,7 +42,7 @@ If any of the following assumptions are broken, the reconstructions may fail com :class: only-dark ``` -Here is an overview pipeline for NeRF, we will walk though each component in this guide. +Here is an overview pipeline for NeRF, we will walk through each component in this guide. ### Field Representation @@ -81,7 +81,7 @@ field_fine = NeRFField(position_encoding=pos_enc, direction_encoding=dir_enc) #### Positional Encoding -An extra trick is necessary to making the neural network expressive enough to represent fine details in the scene. The input coordinates $(x,y,z,\theta,\phi)$ need to be encoded to a higher dimensional space prior to being input into the network. You can learn more about encodings [here](../model_components/visualize_encoders.ipynb). +An extra trick is necessary to make the neural network expressive enough to represent fine details in the scene. The input coordinates $(x,y,z,\theta,\phi)$ need to be encoded to a higher dimensional space prior to being input into the network. You can learn more about encodings [here](../model_components/visualize_encoders.ipynb). ```python from nerfactory.fields.modules.encoding import NeRFEncoding @@ -137,7 +137,7 @@ rgb = renderer_rgb( :class: only-dark ``` -How we sample points along rays in space is an important design decision. Various sampling strategies can be used which are dicussed in detail in the [Ray Samplers](../model_components/visualize_samplers.ipynb) guide. In NeRF we take advantage of a hierarchical sampling scheme that first uses a _uniform sampler_ and is followed by a _PDF sampler_. The uniform sampler distributes samples evenly between a predefined distance range from the camera. These are then used to compute an initial render of the scene. The renderer optionally produces _weights_ for each sample that correlate with how important each sample was to the final renderer. The PDF sampler uses these _weights_ to generate a new set of samples that are biased to regions of higher weight. In practice, these regions are near the surface of the object. +How we sample points along rays in space is an important design decision. Various sampling strategies can be used which are discussed in detail in the [Ray Samplers](../model_components/visualize_samplers.ipynb) guide. In NeRF we take advantage of a hierarchical sampling scheme that first uses a _uniform sampler_ and is followed by a _PDF sampler_. The uniform sampler distributes samples evenly between a predefined distance range from the camera. These are then used to compute an initial render of the scene. The renderer optionally produces _weights_ for each sample that correlate with how important each sample was to the final renderer. The PDF sampler uses these _weights_ to generate a new set of samples that are biased to regions of higher weight. In practice, these regions are near the surface of the object. Associated code: @@ -153,9 +153,8 @@ weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadN ray_samples_pdf = sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) ``` - ```{warning} -Described above is specific to scenes that have known bounds (ie. the Blender Sythetic dataset). For unbounded scenes, the original NeRF paper uses Normalized Device Coordinates (NDC) to warp space, along with a _linear in disparity_ sampler. We do not support NDC, for unbounded scenes consider using [Spatial Distortions](../model_components/visualize_spatial_distortions.ipynb). +Described above is specific to scenes that have known bounds (ie. the Blender Synthetic dataset). For unbounded scenes, the original NeRF paper uses Normalized Device Coordinates (NDC) to warp space, along with a _linear in disparity_ sampler. We do not support NDC, for unbounded scenes consider using [Spatial Distortions](../model_components/visualize_spatial_distortions.ipynb). ``` ```{tip} diff --git a/docs/reference/contributing.md b/docs/reference/contributing.md index 8baaa67fa..d2a6eb516 100644 --- a/docs/reference/contributing.md +++ b/docs/reference/contributing.md @@ -1,6 +1,6 @@ # Contributing to nerfactory workflow -The project is set up for development in VSCode, we reccomend using it if you plan to contribute. +The project is set up for development in VSCode, we recommend using it if you plan to contribute. ## TLDR @@ -11,7 +11,7 @@ The project is set up for development in VSCode, we reccomend using it if you pl | Docs | Sphinx | | Docstring style | Google | -## Commiting +## Committing 1. Make your modifications 2. Perform local checks @@ -57,7 +57,7 @@ Run the following to build the documentation: python scripts/docs/build_docs.py ``` -Rerun `make html` when documentation changes are made, `make clean` is necissary if the documentation structure changes. +Rerun `make html` when documentation changes are made, `make clean` is necessary if the documentation structure changes. ### Notebooks @@ -66,7 +66,7 @@ We support jupyter notbooks in our documentation. To improve the readability, th | Tag | Effect | | ------------- | ---------------------------------------------------- | | # HIDDEN | Hide code block and output | -| # COLLAPSED | Collapse the code in a dropdown but show the restuls | +| # COLLAPSED | Collapse the code in a dropdown but show the results | | # OUTPUT_ONLY | Only show the cell's output | ### Auto build diff --git a/nerfactory/dataloaders/base.py b/nerfactory/dataloaders/base.py index 11bca9ee5..7cf2d0857 100644 --- a/nerfactory/dataloaders/base.py +++ b/nerfactory/dataloaders/base.py @@ -19,7 +19,7 @@ from __future__ import annotations import logging from abc import abstractmethod -from typing import Dict, List, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import torch from torch import nn @@ -27,10 +27,11 @@ from torch.nn import Parameter from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg -from nerfactory.dataloaders.eval import FixedIndicesEvalDataloader +from nerfactory.dataloaders.eval import EvalDataloader, FixedIndicesEvalDataloader from nerfactory.dataloaders.image_dataset import ImageDataset, PanopticImageDataset from nerfactory.dataloaders.image_sampler import CacheImageSampler from nerfactory.dataloaders.pixel_sampler import PixelSampler +from nerfactory.dataloaders.structs import DatasetInputs from nerfactory.models.modules.ray_generator import RayGenerator from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfactory.utils.misc import IterableWrapper @@ -71,20 +72,28 @@ class Dataloader(nn.Module): get_eval_iterable: utility that gets a clean pythonic iterator for your eval data Args: - use_train (bool): whether this dataloader is being used for training - use_eval (bool): whether this dataloader is being used for evaluation + use_train: whether this dataloader is being used for training + use_eval: whether this dataloader is being used for evaluation Attributes: use_train (bool): whether or not we are using train use_eval (bool): whether or not we are using eval train_count (int): the step number of our train iteration, needs to be incremented manually eval_count (int): the step number of our eval iteration, needs to be incremented manually + train_datasetinputs (DatasetInputs): the inputs for the train dataset + eval_datasetinputs (DatasetInputs): the inputs for the eval dataset + train_image_dataset (ImageDataset): the image dataset for the train dataset - Additional attributes specific to each subclass are defined in the setuo_train and setup_eval + Additional attributes specific to each subclass are defined in the setup_train and setup_eval functions. """ + train_datasetinputs: Optional[DatasetInputs] = None + eval_datasetinputs: Optional[DatasetInputs] = None + train_image_dataset: Optional[ImageDataset] = None + eval_dataloader: Optional[EvalDataloader] = None + def __init__(self, use_train: bool, use_eval: bool): """Constructor for the Dataloader class. @@ -223,18 +232,21 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method else: logging.info("No eval dataset specified so using train dataset for eval.") dataset_eval = dataset_train - self.config.eval_datasetinputs = dataset_eval.get_dataset_inputs(split="val" if not test_mode else "test") - self.eval_cameras = self.config.eval_datasetinputs.cameras.to(device) + self.eval_datasetinputs = dataset_eval.get_dataset_inputs(split="val" if not test_mode else "test") + self.eval_cameras = self.eval_datasetinputs.cameras.to(device) use_train = self.train_datasetinputs is not None - use_eval = self.config.eval_datasetinputs is not None + use_eval = self.eval_datasetinputs is not None super().__init__(use_train, use_eval) def setup_train(self): """Sets up the dataloader for training""" + assert self.train_datasetinputs is not None if self.config.image_dataset_type == "rgb": self.train_image_dataset = ImageDataset(**self.train_datasetinputs.as_dict()) elif self.config.image_dataset_type == "panoptic": self.train_image_dataset = PanopticImageDataset(**self.train_datasetinputs.as_dict()) + else: + raise ValueError(f"Unknown image dataset type {self.config.image_dataset_type}") self.train_image_sampler = CacheImageSampler( self.train_image_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, @@ -246,10 +258,11 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method def setup_eval(self): """Sets up the dataloader for evaluation""" + assert self.eval_datasetinputs is not None if self.config.image_dataset_type == "rgb": - self.eval_image_dataset = ImageDataset(**self.config.eval_datasetinputs.as_dict()) + self.eval_image_dataset = ImageDataset(**self.eval_datasetinputs.as_dict()) elif self.config.image_dataset_type == "panoptic": - self.eval_image_dataset = PanopticImageDataset(**self.config.eval_datasetinputs.as_dict()) + self.eval_image_dataset = PanopticImageDataset(**self.eval_datasetinputs.as_dict()) self.eval_dataloader = FixedIndicesEvalDataloader( image_dataset=self.eval_image_dataset, cameras=self.eval_cameras, @@ -272,5 +285,6 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method """Returns the next batch of data from the eval dataloader. The RayBundle should be shaped like an image.""" self.eval_count += 1 + assert self.eval_dataloader is not None, "Must setup eval dataloader before calling next_eval" camera_ray_bundle, batch = next(self.eval_dataloader) return camera_ray_bundle, batch diff --git a/nerfactory/dataloaders/eval.py b/nerfactory/dataloaders/eval.py index 9355f4c80..e57b46888 100644 --- a/nerfactory/dataloaders/eval.py +++ b/nerfactory/dataloaders/eval.py @@ -28,13 +28,20 @@ from nerfactory.dataloaders.image_dataset import ImageDataset from nerfactory.utils.misc import get_dict_to_torch -class EvalDataloader: # pylint: disable=too-few-public-methods - """Evaluation dataloader base class""" +class EvalDataloader: + """Evaluation dataloader base class + + Args: + image_dataset: ImageDataset to load data from + cameras: Cameras to use for generating rays + num_rays_per_chunk: Number of camera rays to generate per chunk + device: Device to load data to + """ def __init__( self, image_dataset: ImageDataset, - cameras, + cameras: Cameras, num_rays_per_chunk: int, device: Union[torch.device, str] = "cpu", **kwargs, @@ -55,15 +62,23 @@ class EvalDataloader: # pylint: disable=too-few-public-methods def __next__(self) -> Tuple[RayBundle, Dict]: """Returns the next batch of data""" - def get_camera(self, image_idx) -> Cameras: - """Get camera for the given image index""" + def get_camera(self, image_idx: int = 0) -> Cameras: + """Get camera for the given image index + + Args: + image_idx: Camera image index + """ + distortion_params = None + if self.cameras.distortion_params is not None: + distortion_params = self.cameras.distortion_params[image_idx] + camera = Cameras( fx=self.cameras.fx[image_idx], fy=self.cameras.fy[image_idx], cx=self.cameras.cx, cy=self.cameras.cy, camera_to_worlds=self.cameras.camera_to_worlds[image_idx], - distortion_params=self.cameras.distortion_params[image_idx], + distortion_params=distortion_params, camera_type=self.cameras.camera_type, ) return camera @@ -93,7 +108,10 @@ class FixedIndicesEvalDataloader(EvalDataloader): """ Args: image_dataset: ImageDataset to load data from + cameras: Cameras to use for generating rays + num_rays_per_chunk: Number of camera rays to generate per chunk image_indices: List of image indices to load data from. If None, then use all images. + device: Device to load data to """ super().__init__(image_dataset, cameras, num_rays_per_chunk, device, **kwargs) if image_indices is None: @@ -118,7 +136,14 @@ class FixedIndicesEvalDataloader(EvalDataloader): class RandIndicesEvalDataloader(EvalDataloader): - """Dataloader that returns random images.""" + """Dataloader that returns random images. + + Args: + image_dataset: ImageDataset to load data from + cameras: Cameras to use for generating rays + num_rays_per_chunk: Number of camera rays to generate per chunk + device: Device to load data to + """ def __init__( self, diff --git a/nerfactory/dataloaders/structs.py b/nerfactory/dataloaders/structs.py index 322ce0e12..be6c74b86 100644 --- a/nerfactory/dataloaders/structs.py +++ b/nerfactory/dataloaders/structs.py @@ -24,7 +24,6 @@ import torch from torchtyping import TensorType from nerfactory.cameras.cameras import Cameras -from nerfactory.cameras.rays import RayBundle @dataclass @@ -133,42 +132,3 @@ class DatasetInputs: def as_dict(self) -> dict: """Returns the dataclass as a dictionary.""" return vars(self) - - -@dataclass -class DataloaderOutputs: - """A container for data that is not specific to any dataset. It should be everything - needed by the renderer and loss calculation. Different datasets and models will probably - need to subclass from this differently, since different Field modules will require different - data. - - Args: - rays (RayBundle): The rays for the image. - ground_truth_pixels (TensorType["num_pixels", 3]): The ground truth pixels for the image. - """ - - rays: RayBundle # Raybundle and the cameras will be merged into one thing in a later PR - cameras: torch.Tensor - ground_truth_pixels: Optional[TensorType["num_pixels", 3]] = None - - def check_inputs(self, expected_attr: Dict): - """This will check to make sure that the expected attributes exist. - - This will exist in any subclassed version of this""" - for attr in expected_attr: - assert hasattr(self, attr) - - -@dataclass -class ModelOutputs: - """A container for outputs from the model. - - This will be passed into the loss function along with the base data container. It should be - subclassed for when your model outputs different things. - - Args: - rays (RayBundle): The rays for the image. - ground_truth_pixels (TensorType["num_pixels", 3]): The ground truth pixels for the image. - """ - - rendered_pixels: Optional[TensorType["num_pixels", 3]] = None diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index bd5216c5d..6cb56d236 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -120,6 +120,9 @@ class Trainer: def train(self) -> None: """Train the model.""" + assert self.pipeline.dataloader.train_datasetinputs is not None, "Missing DatasetInputs" + assert self.pipeline.dataloader.train_image_dataset is not None, "Missing ImageDataset" + self.visualizer_state.init_scene( image_dataset=self.pipeline.dataloader.train_image_dataset, dataset_inputs=self.pipeline.dataloader.train_datasetinputs, diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index d1617f056..398a32229 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -50,6 +50,8 @@ class NGPModel(Model): kwargs: additional params to pass up to the parent class model """ + config: cfg.InstantNGPModelConfig + def __init__(self, config: cfg.InstantNGPModelConfig, **kwargs) -> None: assert config.field_implementation in field_implementation_to_class self.field = None diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index 4017f4022..afe0705cf 100644 --- a/nerfactory/models/nerfw.py +++ b/nerfactory/models/nerfw.py @@ -45,6 +45,8 @@ from nerfactory.utils import colors, misc, visualization, writer class NerfWModel(Model): """NeRF-W model""" + config: cfg.NerfWModelConfig + def __init__( self, config: cfg.NerfWModelConfig, diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 774e24d65..6e3c82a30 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -76,6 +76,7 @@ class Pipeline(nn.Module): self.dataloader: Dataloader = config.dataloader.setup(device=device, test_mode=test_mode) self.dataloader.to(device) # TODO(ethan): get rid of scene_bounds from the model + assert self.dataloader.train_datasetinputs is not None, "Missing DatasetInputs" self.model: Model = config.model.setup(scene_bounds=self.dataloader.train_datasetinputs.scene_bounds) self.model.to(device) @@ -101,7 +102,9 @@ class Pipeline(nn.Module): from the dataloader and feed it to the model's forward function""" self.eval() # NOTE(ethan): next_eval() is not being used right now + assert self.dataloader.eval_dataloader is not None for camera_ray_bundle, batch in self.dataloader.eval_dataloader: + assert camera_ray_bundle.camera_indices is not None image_idx = int(camera_ray_bundle.camera_indices[0, 0]) outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) psnr = self.model.log_test_image_outputs(image_idx, step, batch, outputs)