From a9e6cad0fbfae199d9aabac7b2c22389a38b312e Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 20 May 2022 15:36:14 -0700 Subject: [PATCH] Refactor field heads --- mattport/nerf/field_modules/field_heads.py | 65 +++++++++++++------ mattport/nerf/fields/base.py | 3 +- mattport/nerf/fields/nerf_field.py | 30 +++++---- mattport/nerf/graph/vanilla_nerf.py | 9 +-- .../nerf/field_modules/test_field_outputs.py | 11 +++- 5 files changed, 80 insertions(+), 38 deletions(-) diff --git a/mattport/nerf/field_modules/field_heads.py b/mattport/nerf/field_modules/field_heads.py index 4416843b4..4d0a0e61a 100644 --- a/mattport/nerf/field_modules/field_heads.py +++ b/mattport/nerf/field_modules/field_heads.py @@ -26,17 +26,35 @@ class FieldHeadNames(Enum): class FieldHead(FieldModule): """Base field output""" - def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None: + def __init__( + self, + out_dim: int, + field_head_name: FieldHeadNames, + in_dim: Optional[int] = None, + activation: Optional[nn.Module] = None, + ) -> None: """ Args: - in_dim (int): input dimension out_dim (int): output dimension for renderer + field_head_name (FieldHeadNames): Field type + in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. activation (Optional[nn.Module]): output head activation """ super().__init__() - self.in_dim = in_dim self.out_dim = out_dim self.activation = activation + self.field_head_name = field_head_name + self.net = None + if in_dim is not None: + self.in_dim = in_dim + self._construct_net() + + def set_in_dim(self, in_dim: int) -> None: + """Set input dimension of Field Head""" + self.in_dim = in_dim + self._construct_net() + + def _construct_net(self): layers = [nn.Linear(self.in_dim, self.out_dim)] if self.activation: layers.append(self.activation) @@ -52,7 +70,7 @@ class FieldHead(FieldModule): TensorType[..., "out_dim"]: Render head output """ if not self.net: - raise SystemError("Render head network not initialized. build_nn_modules() should be called.") + raise SystemError("in_dim not set. Must be provided to construtor, or set_in_dim() should be called.") out_tensor = self.net(in_tensor) return out_tensor @@ -60,23 +78,24 @@ class FieldHead(FieldModule): class DensityFieldHead(FieldHead): """Density output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None: - super().__init__(in_dim, out_dim=1, activation=activation) + def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: + super().__init__(in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.DENSITY, activation=activation) class RGBFieldHead(FieldHead): """RGB output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: - super().__init__(in_dim, out_dim=3, activation=activation) + def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: + super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation) class SHFieldHead(FieldHead): """Spherical harmonics output""" - def __init__(self, in_dim: int, levels: int = 3, channels: int = 3, activation: Optional[nn.Module] = None) -> None: - """_summary_ - + def __init__( + self, in_dim: Optional[int] = None, levels: int = 3, channels: int = 3, activation: Optional[nn.Module] = None + ) -> None: + """ Args: in_dim (int): Input dimension levels (int, optional): Number of spherical harmonics layers. Defaults to 3. @@ -84,37 +103,41 @@ class SHFieldHead(FieldHead): activation (Optional[nn.Module], optional): Output activation. Defaults to None. """ out_dim = channels * levels**2 - super().__init__(in_dim, out_dim=out_dim, activation=activation) + super().__init__(in_dim=in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.SH, activation=activation) class UncertaintyFieldHead(FieldHead): """Uncertainty output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None: - super().__init__(in_dim, out_dim=1, activation=activation) + def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: + super().__init__(in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.UNCERTAINTY, activation=activation) class TransientRGBHead(FieldHead): """Transient RGB output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: - super().__init__(in_dim, out_dim=3, activation=activation) + def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: + super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.TRANSIENT_RGB, activation=activation) class TransientDensityHead(FieldHead): """Transient density output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None: - super().__init__(in_dim, out_dim=1, activation=activation) + def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: + super().__init__( + in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.TRANSIENT_DENSITY, activation=activation + ) class SemanticStuffHead(FieldHead): """Semantic stuff output""" - def __init__(self, in_dim: int, num_classes: int) -> None: + def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None: """ Args: - in_dim (int): Input dimension num_classes (int): Number of semantic classes + in_dim (int): Input dimension """ - super().__init__(in_dim, out_dim=num_classes, activation=None) + super().__init__( + in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_STUFF, activation=None + ) diff --git a/mattport/nerf/fields/base.py b/mattport/nerf/fields/base.py index cbea5dad8..fc765a605 100644 --- a/mattport/nerf/fields/base.py +++ b/mattport/nerf/fields/base.py @@ -7,6 +7,7 @@ from typing import Tuple from torch import nn from torchtyping import TensorType +from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.structures.rays import PointSamples @@ -39,5 +40,5 @@ class Field(nn.Module): """Evaluates the field at points along the ray.""" density, density_embedding = self.get_density(point_samples) field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding) - field_outputs["density"] = density + field_outputs[FieldHeadNames.DENSITY] = density return field_outputs diff --git a/mattport/nerf/fields/nerf_field.py b/mattport/nerf/fields/nerf_field.py index 6aac19caf..79c2949a7 100644 --- a/mattport/nerf/fields/nerf_field.py +++ b/mattport/nerf/fields/nerf_field.py @@ -4,8 +4,10 @@ from typing import Tuple import torch +from torch import nn + from mattport.nerf.field_modules.encoding import Encoding, Identity -from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead +from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP from mattport.nerf.fields.base import Field from mattport.structures.rays import PointSamples @@ -23,7 +25,8 @@ class NeRFField(Field): base_mlp_layer_width: int = 256, head_mlp_num_layers: int = 2, head_mlp_layer_width: int = 128, - skip_connections: Tuple = (4,), + skip_connections: Tuple[int] = (4,), + field_heads: Tuple[FieldHead] = (RGBFieldHead(),), ) -> None: """Create boilerplate NeRF field. @@ -54,7 +57,9 @@ class NeRFField(Field): ) self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) - self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim()) + self.field_heads = nn.ModuleList(field_heads) + for field_head in self.field_heads: + field_head.set_in_dim(self.mlp_head.get_out_dim()) def get_density(self, point_samples: PointSamples): """Computes and returns the densities.""" @@ -79,11 +84,14 @@ class NeRFField(Field): directions = point_samples.directions if not is_not_none(valid_mask): valid_mask = torch.ones_like(directions[..., 0]).bool() - # placeholders for values to return - rgb = torch.zeros(*valid_mask.shape, 3, dtype=torch.float32, device=directions.device) - if not valid_mask.any(): # empty mask - return {"rgb": rgb} - encoded_dir = self.direction_encoding(directions[valid_mask]) - rgb_mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1)) - rgb[valid_mask] = self.field_output_rgb(rgb_mlp_out) - return {"rgb": rgb} + outputs = {} + for field_head in self.field_heads: + # placeholders for values to return + out = torch.zeros(*valid_mask.shape, field_head.out_dim, dtype=torch.float32, device=directions.device) + if not valid_mask.any(): # empty mask + return {field_head.field_head_name: out} + encoded_dir = self.direction_encoding(directions[valid_mask]) + mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1)) + out[valid_mask] = field_head(mlp_out) + outputs[field_head.field_head_name] = out + return outputs diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 4f2275957..875f6c80a 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -10,6 +10,7 @@ from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from mattport.nerf.field_modules.encoding import NeRFEncoding +from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.nerf.fields.nerf_field import NeRFField from mattport.nerf.graph.base import Graph @@ -88,9 +89,9 @@ class NeRFGraph(Graph): # coarse field: field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples()) - weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse["density"]) + weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( - rgb=field_outputs_coarse["rgb"], + rgb=field_outputs_coarse[FieldHeadNames.RGB], weights=weights_coarse, ) accumulation_coarse = self.renderer_accumulation(weights_coarse) @@ -101,9 +102,9 @@ class NeRFGraph(Graph): # fine field: field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) - weights_fine = ray_samples_pdf.get_weights(field_outputs_fine["density"]) + weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( - rgb=field_outputs_fine["rgb"], + rgb=field_outputs_fine[FieldHeadNames.RGB], weights=weights_fine, ) accumulation_fine = self.renderer_accumulation(weights_fine) diff --git a/tests/nerf/field_modules/test_field_outputs.py b/tests/nerf/field_modules/test_field_outputs.py index fcc07fec8..7fd2adf8f 100644 --- a/tests/nerf/field_modules/test_field_outputs.py +++ b/tests/nerf/field_modules/test_field_outputs.py @@ -3,10 +3,12 @@ Field output tests """ import torch from torch import nn +import pytest from mattport.nerf.field_modules.field_heads import ( DensityFieldHead, FieldHead, + FieldHeadNames, RGBFieldHead, SHFieldHead, ) @@ -17,12 +19,19 @@ def test_field_output(): in_dim = 6 out_dim = 4 activation = nn.ReLU() - render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, activation=activation) + render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation) assert render_head.get_out_dim() == out_dim x = torch.ones((9, in_dim)) render_head(x) + # Test in_dim not provided at construction + render_head = FieldHead(out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation) + with pytest.raises(SystemError): + render_head(x) + render_head.set_in_dim(in_dim) + render_head(x) + def test_density_output(): """Test rgb output"""