diff --git a/mattport/nerf/fields/base.py b/mattport/nerf/fields/base.py index e930b2174..cbea5dad8 100644 --- a/mattport/nerf/fields/base.py +++ b/mattport/nerf/fields/base.py @@ -3,21 +3,17 @@ Base class for the graphs. """ from abc import abstractmethod +from typing import Tuple from torch import nn +from torchtyping import TensorType -from mattport.structures.rays import PointSamples, RaySamples +from mattport.structures.rays import PointSamples class Field(nn.Module): """Base class for fields.""" - def __init__(self) -> None: - super().__init__() - # minimum weight for it to be valid. used to skip querying the field - # TODO(ethan): avoid the hardcoding - self.weights_threshold = 1e-4 - def density_fn(self, positions): """Returns only the density. Used primarily with the occupancy grid.""" point_samples = PointSamples(positions=positions) @@ -25,8 +21,15 @@ class Field(nn.Module): return density @abstractmethod - def get_density(self, point_samples: PointSamples): - """Computes and returns the densities.""" + def get_density(self, point_samples: PointSamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: + """Computes and returns the densities. + + Args: + point_samples (PointSamples): Samples locations to compute density. + + Returns: + Tuple[TensorType[...,1], TensorType[...,"num_features"]]: A tensor of densities and a tensor of features. + """ @abstractmethod def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): @@ -38,14 +41,3 @@ class Field(nn.Module): field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding) field_outputs["density"] = density return field_outputs - - def forward_with_weight_pruning(self, ray_samples: RaySamples): - """Forward but with pruning based on weights, from the densities.""" - density, density_embedding = self.get_density(ray_samples) - weights = ray_samples.get_weights(density) - valid_mask_outputs = weights >= self.weights_threshold - field_outputs = self.get_outputs( - ray_samples, density_embedding=density_embedding, valid_mask=valid_mask_outputs - ) - field_outputs["density"] = density - return field_outputs diff --git a/mattport/nerf/fields/density_field.py b/mattport/nerf/fields/density_field.py deleted file mode 100644 index a31a0a67e..000000000 --- a/mattport/nerf/fields/density_field.py +++ /dev/null @@ -1,40 +0,0 @@ -""" -Density Field. -""" - -from torch import nn - -from torchtyping import TensorType -from mattport.nerf.field_modules.field_heads import DensityFieldHead -from mattport.structures.rays import RaySamples - - -class DensityField(nn.Module): - """Base class for fields.""" - - def __init__(self, model: nn.Sequential, density_activation: nn.Module = nn.Softplus()) -> None: - super().__init__() - self.model = model - self.density_head = DensityFieldHead(in_dim=self.model[-1].get_out_dim(), activation=density_activation) - - def process_point(self, in_tensor: TensorType): - """_summary_""" - features = self.model(in_tensor) - density = self.density_head(features) - return density, features - - def proccess_ray(self, rays: RaySamples): - """_summary_""" - densities, features = self.process_point(in_tensor=rays.positions) - - weights = rays.get_weights(densities) - return densities, weights, features - - def forward(self, in_tensor): - """summary""" - return self.process_point(in_tensor) - - def density_fn(self, positions): - """Returns only the density.""" - density, _ = self.process_point(in_tensor=positions) - return density diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index f1ba247ba..763d67747 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -15,7 +15,6 @@ from mattport.nerf.field_modules.encoding import NeRFEncoding from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP from mattport.nerf.fields.base import Field -from mattport.nerf.fields.density_field import DensityField from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer @@ -43,9 +42,6 @@ class NeRFField(Field): self.build_mlp_rgb() self.build_heads() - sequence = nn.Sequential(self.encoding_xyz, self.mlp_base) - self.density_field = DensityField(sequence) - def build_encodings(self): """Build the encodings.""" self.encoding_xyz = NeRFEncoding( @@ -95,7 +91,9 @@ class NeRFField(Field): if not valid_mask.any(): # empty mask return density, base_mlp_out - density[valid_mask], base_mlp_out[valid_mask] = self.density_field(positions[valid_mask]) + encoded_xyz = self.encoding_xyz(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]) return density, base_mlp_out def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 7593041ea..616c00850 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -15,22 +15,22 @@ from mattport.utils.misc import is_not_none class PointSamples: """Samples in space.""" - positions: TensorType[..., 3] = None - directions: TensorType[..., 3] = None - camera_indices: TensorType[..., 1] = None - valid_mask: TensorType[..., 1] = None + positions: TensorType[..., 3] = None # XYZ locations + directions: TensorType[..., 3] = None # Unit direction vector + camera_indices: TensorType[..., 1] = None # Camera index + valid_mask: TensorType[..., 1] = None # Rays that are valid @dataclass class RaySamples: """Samples along a ray""" - positions: TensorType[..., 3] = None - directions: TensorType[..., 3] = None - camera_indices: TensorType[..., 1] = None - valid_mask: TensorType[..., 1] = None - ts: TensorType[..., 1] = None - deltas: TensorType[..., 1] = None + positions: TensorType[..., 3] = None # XYZ locations + directions: TensorType[..., 3] = None # Unit direction vector + camera_indices: TensorType[..., 1] = None # Camera index + valid_mask: TensorType[..., 1] = None # Rays that are valid + ts: TensorType[..., 1] = None # "time steps", distances along ray + deltas: TensorType[..., 1] = None # "width" of each sample def to_point_samples(self) -> PointSamples: """Convert to PointSamples instance and return.""" @@ -67,14 +67,10 @@ class RaySamples: @dataclass class RayBundle: - """_summary_ + """A bundle of ray parameters.""" - Returns: - _type_: _description_ - """ - - origins: TensorType["num_rays", 3] - directions: TensorType["num_rays", 3] + origins: TensorType["num_rays", 3] # Ray origins + directions: TensorType["num_rays", 3] # camera_indices: Optional[TensorType["num_rays"]] = None nears: Optional[TensorType["num_rays"]] = None fars: Optional[TensorType["num_rays"]] = None