cleanup
This commit is contained in:
Matthew Tancik authored and GitHub committed 2022-09-01 12:36:54 -07:00
1 parent 9b0cf2ab16
commit f910836b8c
13 files changed
+99 -95

No files matched your search

+1 -3
View File
@@ -57,9 +57,7 @@
"request": "launch",
"program": "scripts/run_train.py",
"console": "integratedTerminal",
"args": [
"tensorf"
]
"args": ["tensorf"]
},
{
"name": "Python: run_train nerfw (gpu)",
+19 -21
View File
@@ -1,13 +1,3 @@
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="docs/_static/imgs/logo-dark.png">
<source media="(prefers-color-scheme: light)" srcset="docs/_static/imgs/logo.png">
<img alt="nerfactory" src="docs/_static/imgs/logo.png" width="400">
</picture>
</p>
<p align="center"> The all-in-one repo for NeRFs </p>
<p align="center">
<a href='https://github.com/plenoptix/nerfactory/actions/workflows/core_code_checks.yml'>
<img src='https://github.com/plenoptix/nerfactory/actions/workflows/core_code_checks.yml/badge.svg' alt='Test Status' />
@@ -26,9 +16,27 @@
<a href="https://badge.fury.io/py/nerfactory"><img src="https://badge.fury.io/py/nerfactory.svg" alt="PyPI version" height="18"></a>
</p>
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="docs/_static/imgs/logo-dark.png">
<source media="(prefers-color-scheme: light)" srcset="docs/_static/imgs/logo.png">
<img alt="nerfactory" src="docs/_static/imgs/logo.png" width="400">
</picture>
</p>
<p align="center"> The all-in-one repo for NeRFs </p>
<p align="center">
<picture>
<a herf="http://www.nerfactory.com/"><img alt="documentation" src="docs/_static/imgs/readme_documentation.png" width="150">
</picture>
<picture>
<a herf="https://viewer.nerfactory.com/"><img alt="viewer" src="docs/_static/imgs/readme_viewer.png" width="150">
</picture>
</p>
- [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: |
@@ -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"
]
},
{
+2 -2
View File
@@ -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
+5 -6
View File
@@ -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}
+4 -4
View File
@@ -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
+24 -10
View File
@@ -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
+32 -7
View File
@@ -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,
-40
View File
@@ -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
+3
View File
@@ -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,
+2
View File
@@ -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
+2
View File
@@ -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,
+3
View File
@@ -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)