nerfstudio/docs/tutorials/pipelines/creating_pipelines.ipynb

Creating Custom Pipelines

Our goal is to enable Nerfactory users to implement an entire NeRF method by only editing one file.

Here we explain how to create custom pipelines with nerfactory. Pipelines are composed of two components, namely a DataManager and a Model. The pipeline is responsible for routing RayBundle instances returned from the DataManager into the Model.

A NeRF method is implemented as a Pipeline. Users can create a Pipeline with a DataManager and Model. We provide default implementations for these components prefixed with the name Vanilla (e.g., VanillaPipeline, VanillaDataManager which inherit from the Pipeline and DataManager class, respectively) and specific Models (e.g., InstantNGPModel or TensoRFModel which inherit from the Model class). With most NeRF methods, VanillaPipeline and VanillaDataManager can be used and only the model needs to be changed. We provide many already implmented Models for people to get started with 3D reconstruction. However, in this tutorial we will demonstrate how to customize each of the Pipeline, DataManager, and Model with an example.

We'll show how to make an incremental DataManager where we progressively add cameras to train with from multiple scenes and importance sample rays from pixels with a high loss. The extensible features that we show in this tutorial are the following:

• a DataManager that incrementally adds cameras
• a DataManager that uses multiple scenes
• a Model that is conditioned on the scene
• a Pipeline that modifies the DataManager's PixelSampler to sample more rays in regions with high loss

Model Components

The model renders values from 3D line segments defined by a RayBundle.

We provide Models and Model Components (i.e., Fields and Field Modules) which can be strung together to create a custom model.

Models

Vanilla NeRF Model

Mip NeRF Model

Mip NeRF 360 Model

Instant NGP Model

NeRF-W Model

Semantic NeRF Model

Model Components

Fields:

Vanilla NeRF Field

Instant NGP Field

Field Modules:

Encodings

Embeddings

Field Heads

Field Losses

MLP

Spatial Distortions

DataManager Components

The DataManager provides a RayBundle (3D line segments with metadata needed to run the forward pass) and a dictionary of information needed at train-time to supervise the pipeline components.

Similar to the Model components, we provide DataManager components that are strung together to generate RayBundle instances. A RayBundle contains everything needed to generate a ray (e.g., origin and direction, near and far planes) as well as any additional metadata that the Model requires (e.g., camera index, scene index, etc.).

DataManager Components

Datasets Parsers:
Takes a data location and returns a DatsetInputs instance. A DatsetInputs will contain all data needed throughout the pipeline.

Vanilla Image DataParser

Image Datasets:
Takes image filename information and can return images.

Vanilla Image Dataset

Image Dataloader:
Takes an Image Dataset and samples images from it.

Vanilla Image Dataloader

Cache Image Dataloader

Pixel Samplers:
Takes sampled images and returns pixel locations.

Vanilla Pixel Sampler

Ray Generators:
Takes cameras and pixel locations and returns RayBundle objects.

Colliders:
Takes RayBundles and clips them according .

In [1]:

# COLLAPSED
%load_ext autoreload
%autoreload 2

from __future__ import annotations

from typing import Dict, Optional, Tuple, List
import torch

from nerfactory.utils.config import DataManagerConfig, ModelConfig, PipelineConfig, ViewerConfig
from nerfactory.datamanagers.base import DataManager
from nerfactory.datamanagers.structs import DatasetInputs
from nerfactory.models.base import Model
from nerfactory.pipelines.base import Pipeline
from nerfactory.models.instant_ngp import NGPModel
from nerfactory.cameras.rays import RayBundle

In [ ]:

class CustomDataManagerConfig(DataManagerConfig):
    _target = "CustomDataManager"
    scene_names: List[str] = ["chair", "drums", "fern", "lego"]
    steps_per_add_cameras: int = 100  # every X iterations, add more cameras
    num_cameras_per_add: int = 10  # number of cameras to add


class CustomDataManager(DataManager):
    """A custom dataloader that incrementally adds cameras and has multiple scenes."""

    config: CustomDataManagerConfig

    def populate_train_modules(self):
        """Populate the train dataloader modules."""
        self.scene_name_to_image_sampler = {}
        self.scene_name_to_collider = {}  # TODO(ethan): move colliders to the dataloader
        for scene_name in self.config.scene_names:
            dataset_inputs: DatasetInputs = None  # TODO
            image_dataset = ImageDataset(dataset_inputs)
            image_sampler = CacheImageDataloader(image_dataset=image_dataset)  # a torch dataloader
            self.scene_name_to_image_sampler[scene_name] = image_sampler
        self.train_pixel_sampler = PixelSampler
        self.train_ray_generator = RayGenerator  # nn.Module

    def populate_eval_modules(self):
        """Populate the eval dataloader modules."""
        # we'll mostly rely on the train dataloader modules
        self.eval_pixel_sampler = EvalPixelSampler  # this will sample entire images

    def sample_images(self) -> List[Image]:
        images = []
        for scene_name in self.scene_name_to_image_sampler.keys():
            image_sampler = self.scene_name_to_image_sampler[scene_name]
            x = image_sampler.forward()  # note that empty forward calls next()
            images.append(x)
        return images

    def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
        """Get the next batch of training data by stringing together the train modules."""
        # grab some images from each scene
        images = self.sample_images()
        pixels = self.train_pixel_sampler.forward(images)
        ray_bundle = self.train_ray_generator(pixels)
        # shape should be (H, W, :)
        ray_bundle.metadata["scene_indices"] = None  # TODO(ethan): populate this with a tensor of scene indices
        dict_ = {}
        return ray_bundle, dict_

    def next_eval(self, step: int) -> Tuple[RayBundle, Dict]:
        """Get the next batch of eval data by stringing together the eval modules."""
        # grab some images from each scene
        images = self.sample_images()
        pixels = self.train_pixel_sampler.forward(images)
        ray_bundle = self.eval_ray_generator(pixels)  # notice this is eval and not train
        # shape should be (H, W, :)
        ray_bundle.metadata["scene_indices"] = None  # TODO(ethan): populate this with a tensor of scene indices
        dict_ = {}
        return ray_bundle, dict_


class CustomNGPModelConfig(ModelConfig):
    _target = "CustomNGPModel"
    coarse_field: str = "temp"
    fine_field: str = "temp2"


class CustomNGPModel(Model):
    """An instant ngp model modified slightly to output semantics."""

    config: SceneConditionNGPModelConfig

    def populate_modules(self):
        self.ngp_model = NGPModel(coarse_field=self.config.coarse_field, fine_field=self.config.fine_field)

    def get_outputs(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:
        # TODO(ethan): pass in batch from the forward function
        # TODO(ethan): rename batch to something else
        outputs = self.ngp_model.forward(ray_bundle)
        outputs["semantics"] = torch.rand_like(outputs["rgb"])
        return outputs

    def get_loss_dict(self):
        return {}


class CustomPipeline(Pipeline):
    """The Instant NGP pipeline."""

    config: PipelineConfig

    def train_step(self, step: int):
        ray_bundle, batch = self.dataloader.next_train(step=step)
        # TODO: maybe run a CNN on the data before passing into model
        model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch)
        # TODO: update pixel sampler state with loss map to show the flexibilty of our Pipeline
        self.dataloader.pixel_sampler.update_loss_map()
        return model_outputs, loss_dict, metrics_dict

    def eval_step(self):
        ray_bundle, batch = self.dataloader.next_eval(step=step)

In [ ]:

dataloader = DataManager()
model = SceneConditionNGPModel()
pipeline = CustomPipeline.from_dataloader_and_model(dataloader=dataloader, model=model)

Creating pipelines from a config

Now we show how to create a pipeline from a config, which has the following form:

@dataclass
class PipelineConfig:
    """Configuration for pipeline instantiation"""

    _target: ClassVar[Type] = Pipeline
    dataloader: DataloaderConfig = MISSING
    model: ModelConfig = MISSING

See nerfactory/utils/config.py for more details. In this example, we will simply load from an existing configuration from configs/graph_instant_ngp.yaml.

In [ ]:

import pprint
import hydra

hydra.core.global_hydra.GlobalHydra.instance().clear()
from hydra import compose, initialize

initialize(version_base="1.2", config_path="../configs/")
config_name = "graph_instant_ngp.yaml"
config = compose(config_name)
pipeline_config = config.pipeline
pprint.pprint(pipeline_config)
print("----------------------------------------------------")

# from nerfactory.pipelines.base import setup_pipeline
# pipeline = setup_pipeline(pipeline_config, device="cuda")
pipeline = pipeline_config.setup()

In [ ]:

pipeline.get_train_loss_dict()

In [ ]: