From 22fdcfe235e5dcc98c1a8f12ac396ab42d26c7e4 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 18 May 2022 17:45:31 -0700 Subject: [PATCH] Add occupancy to samplers --- mattport/nerf/fields/base.py | 10 +-- mattport/nerf/graph/vanilla_nerf.py | 3 +- mattport/nerf/sampler.py | 121 +++++++++++----------------- mattport/structures/rays.py | 12 ++- 4 files changed, 62 insertions(+), 84 deletions(-) diff --git a/mattport/nerf/fields/base.py b/mattport/nerf/fields/base.py index 5f3b01f37..e930b2174 100644 --- a/mattport/nerf/fields/base.py +++ b/mattport/nerf/fields/base.py @@ -25,23 +25,23 @@ class Field(nn.Module): return density @abstractmethod - def get_density(self, point_samples: PointSamples, valid_mask=None): + def get_density(self, point_samples: PointSamples): """Computes and returns the densities.""" @abstractmethod def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): """Computes and returns the colors.""" - def forward(self, point_samples: PointSamples, valid_mask_density=None): + def forward(self, point_samples: PointSamples): """Evaluates the field at points along the ray.""" - density, density_embedding = self.get_density(point_samples, valid_mask=valid_mask_density) + density, density_embedding = self.get_density(point_samples) 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, valid_mask_density=None): + 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, valid_mask=valid_mask_density) + 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( diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index ffb90bea2..931e2bec2 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -81,9 +81,10 @@ class NeRFField(Field): self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim()) self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) - def get_density(self, point_samples: PointSamples, valid_mask=None): + 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 diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 0ddea9091..e3e4cb091 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -3,98 +3,43 @@ Collection of sampling strategies """ from typing import Optional + import torch from torch import nn from torchtyping import TensorType + from mattport.nerf.occupancy_grid import OccupancyGrid - from mattport.structures.rays import RayBundle, RaySamples -class OccupancyGridSampler(nn.Module): - """Sample based on occupancy grid""" - - def __init__( - self, - occupancy_grid: OccupancyGrid, - num_samples: int = 256, - weight_threshold: float = 1e-4, - train_stratified=True, - ) -> None: - """ - Args: - """ - super().__init__() - self.num_samples = num_samples - self.train_stratified = train_stratified - self.weight_threshold = weight_threshold - self.occupancy_grid = occupancy_grid - - @torch.no_grad() - def forward( - self, - ray_bundle: RayBundle, - num_samples: Optional[int] = None, - ) -> RaySamples: - """Generates position samples uniformly. - - Args: - ray_bundle (RayBundle): Rays to generate samples for - num_samples (Optional[int]): Number of samples per ray - - Returns: - RaySamples: Positions and deltas for samples along a ray - """ - assert ray_bundle.nears is not None - assert ray_bundle.fars is not None - - num_samples = num_samples or self.num_samples - num_rays = ray_bundle.origins.shape[0] - - bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device) # shape (num_samples+1,) - bins = ray_bundle.nears[:, None] + bins[None, :] * ( - ray_bundle.fars[:, None] - ray_bundle.nears[:, None] - ) # shape (num_rays, num_samples+1) - - if self.train_stratified and self.training: - t_rand = torch.rand((num_rays, num_samples), dtype=bins.dtype, device=bins.device) - ts = bins[:, :-1] + t_rand * (bins[:, 1:] - bins[:, :-1]) # shape (num_rays, num_samples) - else: - ts = (bins[:, 1:] + bins[:, :-1]) / 2 - - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_bundle, - ) - - densities = self.occupancy_grid.get_densities(ray_samples.positions, update_iter_count=True) - weights = ray_samples.get_weights(densities) - - valid_mask = weights >= self.weight_threshold - - return ray_samples, weights, valid_mask - - class UniformSampler(nn.Module): """Sample uniformly along a ray""" - def __init__(self, num_samples: int, train_stratified=True) -> None: + def __init__( + self, + num_samples: int, + train_stratified=True, + occupancy_field: Optional[OccupancyGrid] = None, + weight_threshold: float = 1e-4, + ) -> None: """ Args: - near_plane (float): Minimum distance along ray to sample - far_plane (float): Maximum distance along ray to sample num_samples (int): Number of samples per ray train_stratified (bool): Use stratified sampling during training. Defults to True + occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + samples below weight_threshold as set as invalid. + weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. """ super().__init__() - self.num_samples = num_samples self.train_stratified = train_stratified + self.occupancy_field = occupancy_field + self.weight_threshold = weight_threshold @torch.no_grad() def forward( self, - ray_bundle: RayBundle, + ray_bundle: RayBundle = None, num_samples: Optional[int] = None, ) -> RaySamples: """Generates position samples uniformly. @@ -128,29 +73,48 @@ class UniformSampler(nn.Module): ray_bundle=ray_bundle, ) + if self.occupancy_field is not None: + densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) + weights = ray_samples.get_weights(densities) + + valid_mask = weights >= self.weight_threshold + ray_samples.set_valid_mask(valid_mask) + return ray_samples class PDFSampler(nn.Module): """Sample based on probability distribution""" - def __init__(self, num_samples: int, train_stratified: bool = True, include_original: bool = True) -> None: + def __init__( + self, + num_samples: int, + train_stratified: bool = True, + include_original: bool = True, + occupancy_field: OccupancyGrid = None, + weight_threshold: float = 1e-4, + ) -> None: """ Args: num_samples (int): Number of samples per ray train_stratified: boolean: Randomize location within each bin during training. Defaults to True include_original: Add original samples to ray. Defaults to True + occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + samples below weight_threshold as set as invalid. + weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. """ super().__init__() self.num_samples = num_samples self.include_original = include_original self.train_stratified = train_stratified + self.occupancy_field = occupancy_field + self.weight_threshold = weight_threshold @torch.no_grad() def forward( self, - coarse_ray_samples: RaySamples, + ray_samples: RaySamples, weights: TensorType[..., "num_samples"], num_samples: Optional[int] = None, eps: float = 1e-5, @@ -195,7 +159,7 @@ class PDFSampler(nn.Module): indicies_g = torch.stack([below, above], -1) matched_shape = (indicies_g.shape[0], indicies_g.shape[1], cdf.shape[-1]) cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) - bins_g = torch.gather(coarse_ray_samples.ts.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) + bins_g = torch.gather(ray_samples.ts.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) denom = cdf_g[..., 1] - cdf_g[..., 0] denom = torch.where(denom < eps, torch.ones_like(denom), denom) @@ -203,7 +167,7 @@ class PDFSampler(nn.Module): ts = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0]) if self.include_original: - ts, _ = torch.sort(torch.cat([coarse_ray_samples.ts, ts], -1), -1) + ts, _ = torch.sort(torch.cat([ray_samples.ts, ts], -1), -1) else: ts, _ = torch.sort(ts, -1) @@ -212,7 +176,14 @@ class PDFSampler(nn.Module): ray_samples = RaySamples( ts=ts, - ray_bundle=coarse_ray_samples.ray_bundle, + ray_bundle=ray_samples.ray_bundle, ) + if self.occupancy_field is not None: + densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) + weights = ray_samples.get_weights(densities) + + valid_mask = weights >= self.weight_threshold + ray_samples.set_valid_mask(valid_mask) + return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index d1df55b97..56a8134f7 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -4,8 +4,8 @@ Some ray datastructures. import random from dataclasses import dataclass from typing import Optional -import torch +import torch from torchtyping import TensorType from mattport.utils.misc import is_not_none @@ -18,6 +18,7 @@ class PointSamples: positions: TensorType[..., 3] = None directions: TensorType[..., 3] = None camera_indices: TensorType[..., 1] = None + valid_mask: TensorType[..., 1] = None @dataclass @@ -145,11 +146,12 @@ class RaySamples: self.positions = self.get_positions(ray_bundle) self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1) self.deltas = self.get_deltas() + self.valid_mask = torch.ones_like(self.ts, dtype=torch.bool) - def to_point_samples(self): + def to_point_samples(self) -> PointSamples: """Convert to PointSamples instance and return.""" # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions) + return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) def get_camera_indices(self): """Returns camera indices.""" @@ -191,3 +193,7 @@ class RaySamples: weights = alphas * transmittance # [..., "num_samples"] return weights + + def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: + """Sets valid mask""" + self.valid_mask = valid_mask