diff --git a/docs/developer_guides/pipelines/fields.md b/docs/developer_guides/pipelines/fields.md index e1cb63e0a..8c6c6005f 100644 --- a/docs/developer_guides/pipelines/fields.md +++ b/docs/developer_guides/pipelines/fields.md @@ -11,3 +11,110 @@ :class: only-dark :width: 600 ``` + +## What is a Field? + +A Field is a model component that associates a region of space with some sort of quantity. In the most typical case, the input to a field is a 3D location and viewing direction, and the output is density and color. Let's take a look at the code. + +```python +class Field(nn.Module): + """Base class for fields.""" + + @abstractmethod + def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: + """Computes and returns the densities. Returns a tensor of densities and a tensor of features. + + Args: + ray_samples: Samples locations to compute density. + """ + + @abstractmethod + def get_outputs( + self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None + ) -> Dict[FieldHeadNames, TensorType]: + """Computes and returns the colors. Returns output field values. + + Args: + ray_samples: Samples locations to compute outputs. + density_embedding: Density embeddings to condition on. + """ + + def forward(self, ray_samples: RaySamples): + """Evaluates the field at points along the ray. + + Args: + ray_samples: Samples to evaluate field on. + """ + density, density_embedding = self.get_density(ray_samples) + field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding) + + field_outputs[FieldHeadNames.DENSITY] = density # type: ignore + return field_outputs +``` + +## Separate density and outputs + +The forward function is the main function you'll use, which takes in RaySamples returns quantities for each sample. You'll notice that the get_density function is called for every field, followed by the get_outputs function. + +The get_outputs function is what you need to implement to return custom data. For example, check out of SemanticNerfField where we rely on different FieldHeads to produce correct dimensional outputs for typical quantiites. Our implemented FieldHeads have the following FieldHeadNames names. + +```python +class FieldHeadNames(Enum): + """Possible field outputs""" + + RGB = "rgb" + SH = "sh" + DENSITY = "density" + UNCERTAINTY = "uncertainty" + TRANSIENT_RGB = "transient_rgb" + TRANSIENT_DENSITY = "transient_density" + SEMANTICS = "semantics" + SEMANTICS_STUFF = "semantics_stuff" + SEMANTICS_THING = "semantics_thing" +``` + +```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/field_components/field_heads.py +:color: primary +:outline: +See the code! +``` + +Sometimes all you need is the density from a Field, so we have a helper method called density_fn which takes positions and returns densities. + +## Using Frustums instead of positions + +Let's say you want to query a region of space, rather than a point. Our RaySamples data structure contains Frustums which can be used for exactly this purpose. This enables methods like Mip-NeRF to be implemented in our framework. + +```python +@dataclass +class RaySamples(TensorDataclass): + """Samples along a ray""" + + frustums: Frustums + """Frustums along ray.""" + ... + +@dataclass +class Frustums(TensorDataclass): + """Describes region of space as a frustum.""" + + origins: TensorType["bs":..., 3] + """xyz coordinate for ray origin.""" + directions: TensorType["bs":..., 3] + """Direction of ray.""" + starts: TensorType["bs":..., 1] + """Where the frustum starts along a ray.""" + ends: TensorType["bs":..., 1] + """Where the frustum ends along a ray.""" + pixel_area: TensorType["bs":..., 1] + """Projected area of pixel a distance 1 away from origin.""" + ... +``` + +Take a look at our RaySamples class for more information on the input to our Field classes. + +```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/fields/base_field.py +:color: primary +:outline: +See the code! +``` diff --git a/docs/developer_guides/pipelines/index.rst b/docs/developer_guides/pipelines/index.rst index 527fbff39..9ba4fa3d3 100644 --- a/docs/developer_guides/pipelines/index.rst +++ b/docs/developer_guides/pipelines/index.rst @@ -22,6 +22,7 @@ Here we describe what a Pipeline is and how it works. You can see an overview fi Why Pipelines? +========================== Our goal is for any NeRF paper to be implemented as a Pipeline. @@ -41,6 +42,7 @@ In the following sections, we describe the Pipeline components and look at their pipelines Implementing NeRF Papers +========================== Let's say you want to create a custom Pipeline that has a custom DataManager and a custom Model. Perhaps you care about dynmically adding cameras to the DataManager during training or you want to importance sample and generate rays from pixels where the loss is high. This can be accomplished by mixing and matching components into a Pipeline. The following guide will take you through an example of this. diff --git a/docs/developer_guides/pipelines/pipelines.md b/docs/developer_guides/pipelines/pipelines.md index 410cd21ad..36a05274d 100644 --- a/docs/developer_guides/pipelines/pipelines.md +++ b/docs/developer_guides/pipelines/pipelines.md @@ -11,3 +11,63 @@ :class: only-dark :width: 600 ``` + +## What is a Pipeline? + +The Pipeline contains all the code you need to implement a NeRF method. There are two main functions that you need to implement for the Pipeline. + +```python +class Pipeline(nn.Module): + + datamanager: DataManager + model: Model + + @profiler.time_function + def get_train_loss_dict(self, step: int): + """This function gets your training loss dict. This will be responsible for + getting the next batch of data from the DataManager and interfacing with the + Model class, feeding the data to the model's forward function. + + Args: + step: current iteration step to update sampler if using DDP (distributed) + """ + + @profiler.time_function + def get_eval_loss_dict(self, step: int): + """This function gets your evaluation loss dict. It needs to get the data + from the DataManager and feed it to the model's forward function + + Args: + step: current iteration step + """ +``` + +## Vanilla Implementation + +Here you can see a simple implementation of the get_train_loss_dict from the VanillaPipeline. Essentially, all the pipeline has to do is route data from the DataManager to the Model. + +```python +@profiler.time_function +def get_train_loss_dict(self, step: int): + ray_bundle, batch = self.datamanager.next_train(step) + model_outputs = self.model(ray_bundle) + metrics_dict = self.model.get_metrics_dict(model_outputs, batch) + loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) + return model_outputs, loss_dict, metrics_dict +``` + +## Creating Custom Methods + +:::{admonition} Note +:class: info + +The VanillaPipeline works for most of our methods. +::: + +We also have a DynamicBatchPipeline that is used with InstantNGP to dynamically choose the number of rays to use per training and evaluation iteration. + +```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/pipelines/dynamic_batch.py +:color: primary +:outline: +See the code! +``` \ No newline at end of file