mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 10:34:29 +08:00
6879a8a9bd
* starting docs for pipelines * working on an example * clearing output * starting to implement config changes and add documentation * removing diff * format * templated pipeline notebook a bit * black
277 lines
9.9 KiB
Plaintext
277 lines
9.9 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Creating Custom Pipelines\n",
|
|
"\n",
|
|
"Here we explain how to create custom pipelines in nerfactory. Pipelines are composed of two components, namely a Dataloader and a Model. We'll show how to make an incremental dataloader with multiple scenes, a model with scene conditioning, and importance sampling at pixels with high loss.\n",
|
|
"\n",
|
|
"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 Dataloader PixelSampler that samples more densely in regions with high loss\n",
|
|
"\n",
|
|
"Our goal is to enable Nerfactory users to create an entire custom Pipeline by only editing one file.\n",
|
|
"\n",
|
|
"```text\n",
|
|
"#TODO(ethan): show a figure depicting the different modules that go into a pipeline\n",
|
|
"```"
|
|
]
|
|
},
|
|
{
|
|
"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",
|
|
"\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",
|
|
"\n",
|
|
"from nerfactory.models.instant_ngp import NGPModel\n",
|
|
"from nerfactory.cameras.rays import RayBundle"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### RayBundle class\n",
|
|
"\n",
|
|
"In the RayBundle class, we allow uses to add non-default attributes for conditioning.\n"
|
|
]
|
|
},
|
|
{
|
|
"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
|
|
}
|