mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-14 11:54:01 +08:00
1 parent
9b0cf2ab16
commit
f910836b8c
13 files changed
+99
-95
No files matched your search
Vendored
+1
-3
@@ -57,9 +57,7 @@
|
||||
"request": "launch",
|
||||
"program": "scripts/run_train.py",
|
||||
"console": "integratedTerminal",
|
||||
"args": [
|
||||
"tensorf"
|
||||
]
|
||||
"args": ["tensorf"]
|
||||
},
|
||||
{
|
||||
"name": "Python: run_train nerfw (gpu)",
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -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
@@ -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}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user