From 7f32bb91ecf56915f5c5c699dfb4e336132f1479 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 21 Apr 2022 11:31:33 -0700 Subject: [PATCH] Switch to dictionary outputs --- mattport/nerf/field/nerf.py | 5 ++- mattport/nerf/field_modules/field_heads.py | 33 +++++++++---------- mattport/nerf/graph/nerf.py | 6 ++-- mattport/nerf/sampler.py | 6 ++-- .../nerf/field_modules/test_field_outputs.py | 17 ++++++---- tests/nerf/test_sampler.py | 5 +-- 6 files changed, 39 insertions(+), 33 deletions(-) diff --git a/mattport/nerf/field/nerf.py b/mattport/nerf/field/nerf.py index a1746a508..52e9c7880 100644 --- a/mattport/nerf/field/nerf.py +++ b/mattport/nerf/field/nerf.py @@ -6,7 +6,7 @@ import torch from torch import nn from mattport.nerf.field_modules.encoding import NeRFEncoding -from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadOutputs, RGBFieldHead +from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP from mattport.structures.rays import RaySamples @@ -52,6 +52,5 @@ class NeRFField(nn.Module): field_rgb_output = self.field_output_rgb(rgb_mlp_out) field_density_out = self.field_output_density(base_mlp_out) - # TODO(ethan): implement a way to combine the field outputs... - field_outputs = FieldHeadOutputs(rgb=field_rgb_output.rgb, density=field_density_out.density) + field_outputs = field_rgb_output | field_density_out return field_outputs diff --git a/mattport/nerf/field_modules/field_heads.py b/mattport/nerf/field_modules/field_heads.py index 50598117c..5af4ded2a 100644 --- a/mattport/nerf/field_modules/field_heads.py +++ b/mattport/nerf/field_modules/field_heads.py @@ -1,8 +1,8 @@ """ Collection of render heads """ -from dataclasses import dataclass -from typing import Optional +from enum import Enum +from typing import Dict, Optional from torch import nn from torchtyping import TensorType @@ -10,31 +10,32 @@ from torchtyping import TensorType from mattport.nerf.field_modules.base import FieldModule -@dataclass -class FieldHeadOutputs: - """_summary_""" +class FieldHeadNames(Enum): + """Possible field outputs""" - rgb: TensorType["...", 3] = None - density: TensorType["...", 1] = None + RGB = "rgb" + DENSITY = "density" class FieldHead(FieldModule): """Base field output""" def __init__( - self, in_dim: int, out_dim: int, field_quantity_name: str, activation: Optional[nn.Module] = None + self, in_dim: int, out_dim: int, field_head_name: FieldHeadNames, activation: Optional[nn.Module] = None ) -> None: """ Args: in_dim (int): input dimension out_dim (int): output dimension for renderer - field_quantity_name (srt): name of field head + field_head_name (FieldHeadNames): name of field head activation (Optional[nn.Module]): output head activation """ super().__init__() self.in_dim = in_dim self.out_dim = out_dim - self.field_quantity_name = field_quantity_name + # 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() @@ -45,7 +46,7 @@ class FieldHead(FieldModule): layers.append(self.activation) self.net = nn.Sequential(*layers) - def forward(self, in_tensor: TensorType[..., "in_dim"]) -> FieldHeadOutputs: + def forward(self, in_tensor: TensorType[..., "in_dim"]) -> Dict[FieldHeadNames, TensorType]: """Process network output for renderer Args: @@ -54,25 +55,23 @@ class FieldHead(FieldModule): Returns: TensorType[..., "out_dim"]: Render head output """ - if not self.field_quantity_name: + 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) - field_head_outputs = FieldHeadOutputs() - setattr(field_head_outputs, self.field_quantity_name, out_tensor) - return field_head_outputs + return {self.field_head_name: out_tensor} 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, field_quantity_name="density", activation=activation) + super().__init__(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, field_quantity_name="rgb", activation=activation) + super().__init__(in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation) diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py index 4c2598704..e9c5ad98a 100644 --- a/mattport/nerf/graph/nerf.py +++ b/mattport/nerf/graph/nerf.py @@ -49,13 +49,13 @@ class NeRFGraph(Graph): uniform_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs coarse_renderer_outputs = self.renderer_rgb( - rgb=coarse_field_outputs.rgb, density=coarse_field_outputs.density, deltas=uniform_ray_samples.deltas + field_outputs=coarse_field_outputs.rgb, deltas=uniform_ray_samples.deltas ) # RendererOutputs # fine network: - pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs.density) # RaySamples + pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs) # RaySamples fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs fine_renderer_outputs = self.renderer_rgb( - rgb=fine_field_outputs.rgb, density=fine_field_outputs.density, deltas=pdf_ray_samples.deltas + field_outputs=fine_field_outputs, deltas=pdf_ray_samples.deltas ) # RendererOutputs # outputs: outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb} diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index c86dc8458..440af251f 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -2,10 +2,11 @@ Collection of sampling strategies """ -from typing import Optional +from typing import Dict, Optional import torch from torch import nn from torchtyping import TensorType +from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.structures.rays import RayBundle, RaySamples @@ -78,7 +79,7 @@ class PDFSampler(nn.Module): self, ray_bundle: RayBundle, coarse_ray_samples: RaySamples, - density: TensorType[..., "num_samples", 1], + field_outputs: Dict[FieldHeadNames, TensorType], num_samples: Optional[int] = None, randomized: bool = True, eps: float = 1e-5, @@ -102,6 +103,7 @@ class PDFSampler(nn.Module): # Calculate weight contributions along ray # Todo(matt): This computation is duplicated + density = field_outputs[FieldHeadNames.DENSITY] delta_density = coarse_ray_samples.deltas * density[..., 0] alphas = 1 - torch.exp(-delta_density) diff --git a/tests/nerf/field_modules/test_field_outputs.py b/tests/nerf/field_modules/test_field_outputs.py index d6303ffca..68bd590f6 100644 --- a/tests/nerf/field_modules/test_field_outputs.py +++ b/tests/nerf/field_modules/test_field_outputs.py @@ -4,17 +4,16 @@ Field output tests import torch from torch import nn -from mattport.nerf.field_modules import field_heads +from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead def test_field_output(): """Test render output""" in_dim = 6 out_dim = 4 + field_head_name = FieldHeadNames.DENSITY activation = nn.ReLU() - render_head = field_heads.FieldHead( - in_dim=in_dim, out_dim=out_dim, field_quantity_name="rgb", activation=activation - ) + render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, field_head_name=field_head_name, activation=activation) assert render_head.get_out_dim() == out_dim x = torch.ones((9, in_dim)) @@ -24,7 +23,7 @@ def test_field_output(): def test_density_output(): """Test rgb output""" in_dim = 6 - density_head = field_heads.DensityFieldHead(in_dim) + density_head = DensityFieldHead(in_dim) assert density_head.get_out_dim() == 1 x = torch.ones((9, in_dim)) @@ -34,8 +33,14 @@ def test_density_output(): def test_rgb_output(): """Test rgb output""" in_dim = 6 - rgb_head = field_heads.RGBFieldHead(in_dim) + rgb_head = RGBFieldHead(in_dim) assert rgb_head.get_out_dim() == 3 x = torch.ones((9, in_dim)) rgb_head(x) + + +if __name__ == "__main__": + test_field_output() + test_density_output() + test_rgb_output() diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py index bda6d2785..8e884a4be 100644 --- a/tests/nerf/test_sampler.py +++ b/tests/nerf/test_sampler.py @@ -2,6 +2,7 @@ Test samplers """ import torch +from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures.rays import RayBundle @@ -35,14 +36,14 @@ def test_pdf_sampler(): directions = torch.ones_like(origins) ray_bundle = RayBundle(origins=origins, directions=directions) - density = torch.ones((10, num_samples, 1)) + field_outputs = {FieldHeadNames.DENSITY: torch.ones((10, num_samples, 1))} uniform_sampler = UniformSampler(near_plane=near_plane, far_plane=far_plane, num_samples=num_samples) coarse_ray_samples = uniform_sampler(ray_bundle) # Just check that it doesn't crash. Need to add some actual tests. pdf_sampler = PDFSampler(num_samples) - pdf_sampler(ray_bundle, coarse_ray_samples, density, num_samples) + pdf_sampler(ray_bundle, coarse_ray_samples, field_outputs, num_samples) # TODO Tancik: Add more precise tests