nerfstudio/docs/tutorials/pipelines/creating_pipelines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating Custom Pipelines\n",
    "\n",
    "> Our goal is to enable Nerfactory users to implement an entire NeRF method by only editing one file.\n",
    "\n",
    "Here we explain how to create custom pipelines with nerfactory. Pipelines are composed of two components, namely a Dataloader and a Model. The pipeline is responsible for routing RayBundle instances returned from the Dataloader into the Model.\n",
    "\n",
    "A NeRF method is implemented as a Pipeline. Users can create a Pipeline with a Dataloader and Model. We provide default implementations for these components prefixed with the name Vanilla (e.g., `VanillaPipeline`, `VanillaDataloader` which inherit from the `Pipeline` and `Dataloader` class, respectively) and specific Models (e.g., `InstantNGPModel` or `TensoRFModel` which inherit from the `Model` class). With most NeRF methods, `VanillaPipeline` and `VanillaDataloader` 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, Dataloader, and Model with an example.\n",
    "\n",
    "We'll show how to make an incremental dataloader 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:\n",
    "- a Dataloader that incrementally adds cameras\n",
    "- a Dataloader that uses multiple scenes\n",
    "- a Model that is conditioned on the scene\n",
    "- a Pipeline that modifies the Dataloader's PixelSampler to sample more rays in regions with high loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model Components\n",
    "\n",
    "> The model renders values from 3D line segments defined by a RayBundle.\n",
    "\n",
    "We provide Models and Model Components (i.e., Fields and Field Modules) which can be strung together to create a custom model.\n",
    "\n",
    "<details closed>\n",
    "<summary>Models</summary>\n",
    "<br>\n",
    "<li>Vanilla NeRF Model</li>\n",
    "<li>Mip NeRF Model</li>\n",
    "<li>Mip NeRF 360 Model</li>\n",
    "<li>Instant NGP Model</li>\n",
    "<li>NeRF-W Model</li>\n",
    "<li>Semantic NeRF Model</li>\n",
    "</details>\n",
    "\n",
    "<br>\n",
    "\n",
    "<details closed>\n",
    "<summary>Model Components</summary>\n",
    "<br>\n",
    "Fields:\n",
    "<li>Vanilla NeRF Field</li>\n",
    "<li>Instant NGP Field</li>\n",
    "<br>\n",
    "Field Modules:\n",
    "<li>Encodings</li>\n",
    "<li>Embeddings</li>\n",
    "<li>Field Heads</li>\n",
    "<li>Field Losses</li>\n",
    "<li>MLP</li>\n",
    "<li>Spatial Distortions</li>\n",
    "</details>\n",
    "\n",
    "#### Dataloader Components\n",
    "\n",
    "> The Dataloader 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.\n",
    "\n",
    "Similar to the Model components, we provide Dataloader 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.).\n",
    "\n",
    "<details closed>\n",
    "<summary>Dataloader Components</summary>\n",
    "<br>\n",
    "Scene Datasets:<br>\n",
    "    <em style=\"padding-left:2em\">Takes a data location and returns a DatasetInputs instance. A DatasetInputs will contain all data needed throughout the pipeline.</em>\n",
    "<li>Vanilla Image Dataset</li>\n",
    "<li>Panoptic Image Dataset</li>\n",
    "<br>\n",
    "Image Datasets:<br>\n",
    "    <em style=\"padding-left:2em\">Takes image filename information and can return images.</em>\n",
    "<li>Vanilla Image Dataset</li>\n",
    "<li>Panoptic Image Dataset</li>\n",
    "<br>\n",
    "Image Samplers:<br>\n",
    "    <em style=\"padding-left:2em\">Takes an Image Dataset and samples images from it.</em>\n",
    "<li>Vanilla Image Sampler</li>\n",
    "<li>Cache Image Sampler</li>\n",
    "<br>\n",
    "Pixel Samplers:<br>\n",
    "    <em style=\"padding-left:2em\">Takes sampled images and returns pixel locations.</em>\n",
    "<li>Vanilla Image Sampler</li>\n",
    "<li>Cache Image Sampler</li>\n",
    "<br>\n",
    "Pixel Samplers:<br>\n",
    "    <em style=\"padding-left:2em\">Takes sampled images and returns pixel locations.</em>\n",
    "<li>Vanilla Pixel Sampler</li>\n",
    "<br>\n",
    "Ray Generators:<br>\n",
    "    <em style=\"padding-left:2em\">Takes cameras and pixel locations and returns RayBundle objects.</em>\n",
    "<br>\n",
    "\n",
    "Colliders:<br>\n",
    "    <em style=\"padding-left:2em\">Takes RayBundles and clips them according .</em>\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "remove-cell",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# COLLAPSED\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from __future__ import annotations\n",
    "\n",
    "from typing import Dict, Optional, Tuple, List\n",
    "import torch\n",
    "\n",
    "from nerfactory.utils.config import DataloaderConfig, ModelConfig, PipelineConfig, ViewerConfig\n",
    "from nerfactory.dataloaders.base import Dataloader\n",
    "from nerfactory.dataloaders.structs import DatasetInputs\n",
    "from nerfactory.models.base import Model\n",
    "from nerfactory.pipelines.base import Pipeline\n",
    "from nerfactory.models.instant_ngp import NGPModel\n",
    "from nerfactory.cameras.rays import RayBundle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataloaderConfig(DataloaderConfig):\n",
    "    _target = \"CustomDataloader\"\n",
    "    scene_names: List[str] = [\"chair\", \"drums\", \"fern\", \"lego\"]\n",
    "    steps_per_add_cameras: int = 100  # every X iterations, add more cameras\n",
    "    num_cameras_per_add: int = 10  # number of cameras to add\n",
    "\n",
    "\n",
    "class CustomDataloader(Dataloader):\n",
    "    \"\"\"A custom dataloader that incrementally adds cameras and has multiple scenes.\"\"\"\n",
    "\n",
    "    config: CustomDataloaderConfig\n",
    "\n",
    "    def populate_train_modules(self):\n",
    "        \"\"\"Populate the train dataloader modules.\"\"\"\n",
    "        self.scene_name_to_image_sampler = {}\n",
    "        self.scene_name_to_collider = {}  # TODO(ethan): move colliders to the dataloader\n",
    "        for scene_name in self.config.scene_names:\n",
    "            dataset_inputs: DatasetInputs = None  # TODO\n",
    "            image_dataset = ImageDataset(dataset_inputs)\n",
    "            image_sampler = CacheImageSampler(image_dataset=image_dataset)  # a torch dataloader\n",
    "            self.scene_name_to_image_sampler[scene_name] = image_sampler\n",
    "        self.train_pixel_sampler = PixelSampler\n",
    "        self.train_ray_generator = RayGenerator  # nn.Module\n",
    "\n",
    "    def populate_eval_modules(self):\n",
    "        \"\"\"Populate the eval dataloader modules.\"\"\"\n",
    "        # we'll mostly rely on the train dataloader modules\n",
    "        self.eval_pixel_sampler = EvalPixelSampler  # this will sample entire images\n",
    "\n",
    "    def sample_images(self) -> List[Image]:\n",
    "        images = []\n",
    "        for scene_name in self.scene_name_to_image_sampler.keys():\n",
    "            image_sampler = self.scene_name_to_image_sampler[scene_name]\n",
    "            x = image_sampler.forward()  # note that empty forward calls next()\n",
    "            images.append(x)\n",
    "        return images\n",
    "\n",
    "    def next_train(self, step: int) -> Tuple[RayBundle, Dict]:\n",
    "        \"\"\"Get the next batch of training data by stringing together the train modules.\"\"\"\n",
    "        # grab some images from each scene\n",
    "        images = self.sample_images()\n",
    "        pixels = self.train_pixel_sampler.forward(images)\n",
    "        ray_bundle = self.train_ray_generator(pixels)\n",
    "        # shape should be (H, W, :)\n",
    "        ray_bundle.metadata[\"scene_indices\"] = None  # TODO(ethan): populate this with a tensor of scene indices\n",
    "        dict_ = {}\n",
    "        return ray_bundle, dict_\n",
    "\n",
    "    def next_eval(self, step: int) -> Tuple[RayBundle, Dict]:\n",
    "        \"\"\"Get the next batch of eval data by stringing together the eval modules.\"\"\"\n",
    "        # grab some images from each scene\n",
    "        images = self.sample_images()\n",
    "        pixels = self.train_pixel_sampler.forward(images)\n",
    "        ray_bundle = self.eval_ray_generator(pixels)  # notice this is eval and not train\n",
    "        # shape should be (H, W, :)\n",
    "        ray_bundle.metadata[\"scene_indices\"] = None  # TODO(ethan): populate this with a tensor of scene indices\n",
    "        dict_ = {}\n",
    "        return ray_bundle, dict_\n",
    "\n",
    "\n",
    "class CustomNGPModelConfig(ModelConfig):\n",
    "    _target = \"CustomNGPModel\"\n",
    "    coarse_field: str = \"temp\"\n",
    "    fine_field: str = \"temp2\"\n",
    "\n",
    "\n",
    "class CustomNGPModel(Model):\n",
    "    \"\"\"An instant ngp model modified slightly to output semantics.\"\"\"\n",
    "\n",
    "    config: SceneConditionNGPModelConfig\n",
    "\n",
    "    def populate_modules(self):\n",
    "        self.ngp_model = NGPModel(coarse_field=self.config.coarse_field, fine_field=self.config.fine_field)\n",
    "\n",
    "    def get_outputs(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:\n",
    "        # TODO(ethan): pass in batch from the forward function\n",
    "        # TODO(ethan): rename batch to something else\n",
    "        outputs = self.ngp_model.forward(ray_bundle)\n",
    "        outputs[\"semantics\"] = torch.rand_like(outputs[\"rgb\"])\n",
    "        return outputs\n",
    "\n",
    "    def get_loss_dict(self):\n",
    "        return {}\n",
    "\n",
    "\n",
    "class CustomPipeline(Pipeline):\n",
    "    \"\"\"The Instant NGP pipeline.\"\"\"\n",
    "\n",
    "    config: PipelineConfig\n",
    "\n",
    "    def train_step(self, step: int):\n",
    "        ray_bundle, batch = self.dataloader.next_train(step=step)\n",
    "        # TODO: maybe run a CNN on the data before passing into model\n",
    "        model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch)\n",
    "        # TODO: update pixel sampler state with loss map to show the flexibilty of our Pipeline\n",
    "        self.dataloader.pixel_sampler.update_loss_map()\n",
    "        return model_outputs, loss_dict, metrics_dict\n",
    "\n",
    "    def eval_step(self):\n",
    "        ray_bundle, batch = self.dataloader.next_eval(step=step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = Dataloader()\n",
    "model = SceneConditionNGPModel()\n",
    "pipeline = CustomPipeline.from_dataloader_and_model(dataloader=dataloader, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating pipelines from a config\n",
    "\n",
    "Now we show how to create a pipeline from a config, which has the following form:\n",
    "\n",
    "```python\n",
    "@dataclass\n",
    "class PipelineConfig:\n",
    "    \"\"\"Configuration for pipeline instantiation\"\"\"\n",
    "\n",
    "    _target: ClassVar[Type] = Pipeline\n",
    "    dataloader: DataloaderConfig = MISSING\n",
    "    model: ModelConfig = MISSING\n",
    "```\n",
    "\n",
    "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`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "import hydra\n",
    "\n",
    "hydra.core.global_hydra.GlobalHydra.instance().clear()\n",
    "from hydra import compose, initialize\n",
    "\n",
    "initialize(version_base=\"1.2\", config_path=\"../configs/\")\n",
    "config_name = \"graph_instant_ngp.yaml\"\n",
    "config = compose(config_name)\n",
    "pipeline_config = config.pipeline\n",
    "pprint.pprint(pipeline_config)\n",
    "print(\"----------------------------------------------------\")\n",
    "\n",
    "# from nerfactory.pipelines.base import setup_pipeline\n",
    "# pipeline = setup_pipeline(pipeline_config, device=\"cuda\")\n",
    "pipeline = pipeline_config.setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.get_train_loss_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('nerfactory')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "e691f033e0f2f1b9c0a11e7e81375a1e27aec87a71fd1b1eaa545f7a5e61b3f6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}