From 8ccc98a3892e0f32a4249ab75b388294b1ea2107 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 16 May 2022 18:07:22 -0700 Subject: [PATCH] refactor field heads --- mattport/nerf/field_modules/base.py | 10 ++-- mattport/nerf/field_modules/field_heads.py | 18 +++---- mattport/nerf/renderers.py | 52 +++++++++++-------- .../nerf/field_modules/test_field_outputs.py | 21 +++++++- 4 files changed, 59 insertions(+), 42 deletions(-) diff --git a/mattport/nerf/field_modules/base.py b/mattport/nerf/field_modules/base.py index b3c0a8e52..1b74d2d2a 100644 --- a/mattport/nerf/field_modules/base.py +++ b/mattport/nerf/field_modules/base.py @@ -44,15 +44,11 @@ class FieldModule(nn.Module): @abstractmethod def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """_summary_ - + """ Args: - x (TensorType[..., "input_dim"]): Input tensor to be encoded - - Raises: - NotImplementedError: _description_ + in_tensor (TensorType[..., "input_dim"]): Input tensor to process Returns: - TensorType[..., "output_dim"]: A encoded input tensor + TensorType[..., "output_dim"]: Processed tensor """ raise NotImplementedError diff --git a/mattport/nerf/field_modules/field_heads.py b/mattport/nerf/field_modules/field_heads.py index f5e52930b..3d61d22d0 100644 --- a/mattport/nerf/field_modules/field_heads.py +++ b/mattport/nerf/field_modules/field_heads.py @@ -14,6 +14,7 @@ class FieldHeadNames(Enum): """Possible field outputs""" RGB = "rgb" + SH = "sh" DENSITY = "density" UNCERTAINTY = "uncertainty" TRANSIENT_RGB = "transient_rgb" @@ -38,14 +39,8 @@ class FieldHead(FieldModule): super().__init__() self.in_dim = in_dim self.out_dim = out_dim - # if field_head_name not in FieldHeadNames: - # raise ValueError("Incorrect field output name. Should be one of FieldHeadNames.") self.field_head_name = field_head_name self.activation = activation - self.net = None - self.build_nn_modules() - - def build_nn_modules(self) -> None: layers = [nn.Linear(self.in_dim, self.out_dim)] if self.activation: layers.append(self.activation) @@ -60,10 +55,6 @@ class FieldHead(FieldModule): Returns: TensorType[..., "out_dim"]: Render head output """ - if not self.field_head_name: - raise ValueError("field_quantity_name should be set in the child class. E.g., as 'rgb' or 'density'.") - if not self.net: - raise SystemError("Render head network not initialized. build_nn_modules() should be called.") out_tensor = self.net(in_tensor) return {self.field_head_name: out_tensor} @@ -95,7 +86,7 @@ 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, field_head_name=FieldHeadNames.RGB, activation=activation) + super().__init__(in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.SH, activation=activation) class UncertaintyFieldHead(FieldHead): @@ -123,4 +114,9 @@ class SemanticStuffHead(FieldHead): """Semantic stuff output""" def __init__(self, in_dim: int, num_classes: int) -> None: + """ + Args: + in_dim (int): Input dimension + num_classes (int): Number of semantic classes + """ super().__init__(in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_STUFF, activation=None) diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index ffad86620..eb9174b71 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -35,6 +35,29 @@ class RGBRenderer(nn.Module): super().__init__() self.background_color = background_color + @classmethod + def combine_rgb( + cls, + rgb: TensorType[..., "num_samples", 3], + weights: TensorType[..., "num_samples"], + background_color: Optional[TensorType[3]] = None, + ) -> TensorType[..., 3]: + """Composite samples along ray and render color image + + Args: + rgb (TensorType[..., "num_samples", -1]): RGB for each sample + weights (TensorType[..., "num_samples"]): Weights for each sample + + Returns: + TensorType[..., 3]: Outputs rgb values. + """ + rgb = torch.sum(weights[..., None] * rgb, dim=-2) + + if background_color is not None: + rgb = rgb + background_color.to(weights.device)[None, ...] * (1.0 - torch.sum(weights, dim=-1)[..., None]) + + return rgb + def forward( self, rgb: TensorType[..., "num_samples", 3], @@ -50,17 +73,9 @@ class RGBRenderer(nn.Module): RendererOutputs: Outputs with rgb values. """ - rgb = torch.sum(weights[..., None] * rgb, dim=-2) - - if self.background_color is not None: - rgb = rgb + self.background_color.to(weights.device)[None, ...] * ( - 1.0 - torch.sum(weights, dim=-1)[..., None] - ) - torch.clamp_(rgb, min=0.0, max=1.0) - - assert torch.max(rgb) <= 1.0 - assert torch.min(rgb) >= 0.0 - + rgb = self.combine_rgb(rgb, weights) + if not self.training: + torch.clamp_(rgb, min=0.0, max=1.0) renderer_outputs = RendererOutputs(rgb=rgb) return renderer_outputs @@ -69,7 +84,7 @@ class SHRenderer(nn.Module): """Render RGB value from spherical harmonics.""" def __init__( - self, activation: Optional[nn.Module] = nn.Sigmoid(), background_color: Optional[TensorType[3]] = None + self, background_color: Optional[TensorType[3]] = None, activation: Optional[nn.Module] = nn.Sigmoid() ) -> None: """ Args: @@ -103,21 +118,12 @@ class SHRenderer(nn.Module): components = components_from_spherical_harmonics(levels=levels, directions=directions) rgb = sh * components[..., None, :] # [..., num_samples, 3, sh_components] - rgb = torch.sum(sh, dim=-1) # [..., num_samples, 3] - - rgb = torch.sum(weights[..., None] * rgb, dim=-2) + 0.5 + rgb = torch.sum(sh, dim=-1) + 0.5 # [..., num_samples, 3] if self.activation is not None: self.activation(rgb) - if self.background_color is not None: - rgb = rgb + self.background_color.to(weights.device)[None, ...] * ( - 1.0 - torch.sum(weights, dim=-1)[..., None] - ) - torch.clamp_(rgb, min=0.0, max=1.0) - - assert torch.max(rgb) <= 1.0 - assert torch.min(rgb) >= 0.0 + rgb = RGBRenderer.combine_rgb(rgb, weights, background_color=self.background_color) renderer_outputs = RendererOutputs(rgb=rgb) return renderer_outputs diff --git a/tests/nerf/field_modules/test_field_outputs.py b/tests/nerf/field_modules/test_field_outputs.py index 68bd590f6..3175ea673 100644 --- a/tests/nerf/field_modules/test_field_outputs.py +++ b/tests/nerf/field_modules/test_field_outputs.py @@ -4,7 +4,13 @@ Field output tests import torch from torch import nn -from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead +from mattport.nerf.field_modules.field_heads import ( + DensityFieldHead, + FieldHead, + FieldHeadNames, + RGBFieldHead, + SHFieldHead, +) def test_field_output(): @@ -40,7 +46,20 @@ def test_rgb_output(): rgb_head(x) +def test_sh_output(): + """Test sh output""" + in_dim = 6 + levels = 4 + channels = 3 + rgb_head = SHFieldHead(in_dim, levels=levels, channels=channels) + assert rgb_head.get_out_dim() == channels * levels**2 + + x = torch.ones((9, in_dim)) + rgb_head(x) + + if __name__ == "__main__": test_field_output() test_density_output() test_rgb_output() + test_sh_output()