From 800461e7ae809409ebade023c58cacea56ae360e Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 23 May 2022 17:20:43 -0700 Subject: [PATCH] Move masking abstraction --- mattport/nerf/field_modules/encoding.py | 2 +- mattport/nerf/fields/base.py | 48 ++++++++++++++++++++--- mattport/nerf/fields/instant_ngp_field.py | 2 +- mattport/nerf/fields/nerf_field.py | 43 ++++++-------------- mattport/nerf/occupancy_grid.py | 2 +- mattport/structures/rays.py | 23 +++++++++++ 6 files changed, 82 insertions(+), 38 deletions(-) diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index 748664792..0155fce94 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -215,7 +215,7 @@ class HashEncoding(Encoding): assert in_tensor.shape[-1] == 3 in_tensor = in_tensor[..., None, :] # [..., 1, 3] - scaled = in_tensor * self.scalings.view(1, 1, -1, 1).to(in_tensor.device) # [..., L, 3] + scaled = in_tensor * self.scalings.view(-1, 1).to(in_tensor.device) # [..., L, 3] scaled_c = torch.ceil(scaled).type(torch.int32) scaled_f = torch.floor(scaled).type(torch.int32) diff --git a/mattport/nerf/fields/base.py b/mattport/nerf/fields/base.py index fc765a605..1e03db539 100644 --- a/mattport/nerf/fields/base.py +++ b/mattport/nerf/fields/base.py @@ -3,13 +3,15 @@ Base class for the graphs. """ from abc import abstractmethod -from typing import Tuple +from typing import Optional, Tuple from torch import nn +import torch from torchtyping import TensorType from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.structures.rays import PointSamples +from mattport.utils.misc import is_not_none class Field(nn.Module): @@ -33,12 +35,48 @@ class Field(nn.Module): """ @abstractmethod - def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): - """Computes and returns the colors.""" + def get_outputs( + self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + ) -> dict[FieldHeadNames, TensorType]: + """Computes and returns the colors. + + Args: + point_samples (PointSamples): Samples locations to compute outputs. + density_embedding (TensorType, optional): Density embeddings to condition on. + + Returns: + Dict[FieldHeadNames, TensorType]: Output field values. + """ def forward(self, point_samples: PointSamples): """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) + valid_mask = point_samples.valid_mask + + if is_not_none(valid_mask): + # Hacky handling of empty masks. Tests on a single ray but doesn't use results + if not valid_mask.any(): + point_samples = PointSamples( + positions=point_samples.positions[0, :], directions=point_samples.directions[0, :] + ) + else: + point_samples = point_samples.apply_masks() + density_masked, density_embedding_masked = self.get_density(point_samples) + field_outputs_masked = self.get_outputs(point_samples, density_embedding=density_embedding_masked) + + field_outputs = {} + for k, value in field_outputs_masked.items(): + zeros = torch.zeros(*valid_mask.shape, value.shape[-1], dtype=torch.float32, device=valid_mask.device) + if valid_mask.any(): + zeros[valid_mask] = value + else: + zeros[0, :] = value + field_outputs[k] = zeros + density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=valid_mask.device) + if valid_mask.any(): + density[valid_mask] = density_masked + else: + density, density_embedding = self.get_density(point_samples) + field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding) + field_outputs[FieldHeadNames.DENSITY] = density return field_outputs diff --git a/mattport/nerf/fields/instant_ngp_field.py b/mattport/nerf/fields/instant_ngp_field.py index deae732b2..092635136 100644 --- a/mattport/nerf/fields/instant_ngp_field.py +++ b/mattport/nerf/fields/instant_ngp_field.py @@ -102,7 +102,7 @@ class TCNNInstantNGPField(Field): density = F.softplus(density_before_activation) return density, base_mlp_out - def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): + def get_outputs(self, point_samples: PointSamples, density_embedding=None): # TODO: add valid_mask masking! assert is_not_none(density_embedding) directions = point_samples.directions diff --git a/mattport/nerf/fields/nerf_field.py b/mattport/nerf/fields/nerf_field.py index 79c2949a7..99c9267f8 100644 --- a/mattport/nerf/fields/nerf_field.py +++ b/mattport/nerf/fields/nerf_field.py @@ -1,17 +1,17 @@ """ Classic NeRF field""" -from typing import Tuple +from typing import Optional, Tuple import torch from torch import nn +from torchtyping import TensorType from mattport.nerf.field_modules.encoding import Encoding, Identity -from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, RGBFieldHead +from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP from mattport.nerf.fields.base import Field from mattport.structures.rays import PointSamples -from mattport.utils.misc import is_not_none class NeRFField(Field): @@ -63,35 +63,18 @@ class NeRFField(Field): def get_density(self, point_samples: PointSamples): """Computes and returns the densities.""" - positions = point_samples.positions - valid_mask = point_samples.valid_mask - if not is_not_none(valid_mask): - valid_mask = torch.ones_like(positions[..., 0]).bool() - # placeholders for values to return - density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=positions.device) - base_mlp_out = torch.zeros( - *valid_mask.shape, self.mlp_base.out_dim, dtype=torch.float32, device=positions.device - ) - if not valid_mask.any(): # empty mask - return density, base_mlp_out - - encoded_xyz = self.position_encoding(positions[valid_mask]) - base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz) - density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask]) + encoded_xyz = self.position_encoding(point_samples.positions) + base_mlp_out = self.mlp_base(encoded_xyz) + density = self.field_output_density(base_mlp_out) return density, base_mlp_out - def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): - directions = point_samples.directions - if not is_not_none(valid_mask): - valid_mask = torch.ones_like(directions[..., 0]).bool() + def get_outputs( + self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + ) -> dict[FieldHeadNames, TensorType]: + """Computes and returns the outputs.""" 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 + encoded_dir = self.direction_encoding(point_samples.directions) + mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) + outputs[field_head.field_head_name] = field_head(mlp_out) return outputs diff --git a/mattport/nerf/occupancy_grid.py b/mattport/nerf/occupancy_grid.py index dfe4a9b8f..9b72a29d7 100644 --- a/mattport/nerf/occupancy_grid.py +++ b/mattport/nerf/occupancy_grid.py @@ -21,7 +21,7 @@ class OccupancyGrid(nn.Module): self.num_cascades = num_cascades # the number of levels (i.e, cascades) self.resolution = resolution self.register_buffer("aabb", aabb) # axis-aligned bounding box - occupancy_grid = torch.zeros( + occupancy_grid = torch.ones( [self.num_cascades] + [self.resolution] * 3 ) # (num_cascades, x_resolution, y_resolution, z_resolution) # not a module parameter, but it still should be saved and part of `state_dict` diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 780208a22..9b2589b2f 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -27,6 +27,29 @@ class PointSamples: camera_indices: TensorType[..., 1] = None valid_mask: TensorType[...] = None + def apply_masks(self) -> "PointSamples": + """Use valid_mask to mask samples. + + Returns: + PointSamples: New set of masked samples. + """ + if is_not_none(self.valid_mask): + positions = self.positions[self.valid_mask] + directions = self.directions[self.valid_mask] if is_not_none(self.directions) else self.directions + camera_indices = ( + self.camera_indices[self.valid_mask] if is_not_none(self.camera_indices) else self.camera_indices + ) + return PointSamples( + positions=positions, + directions=directions, + camera_indices=camera_indices, + ) + return PointSamples( + positions=self.positions, + directions=self.directions, + camera_indices=self.camera_indices, + ) + @dataclass class RaySamples: