From 22fdcfe235e5dcc98c1a8f12ac396ab42d26c7e4 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 18 May 2022 17:45:31 -0700 Subject: [PATCH 1/3] 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 From 998b98b19a1ecbae7653323bc81b75021e22df80 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 19 May 2022 14:17:03 -0700 Subject: [PATCH 2/3] Clean up sampler --- mattport/nerf/sampler.py | 68 ++++++++++++++++++++++------------------ 1 file changed, 38 insertions(+), 30 deletions(-) diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index e3e4cb091..9269322d9 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -2,6 +2,7 @@ Collection of sampling strategies """ +from abc import abstractmethod from typing import Optional import torch @@ -12,7 +13,34 @@ from mattport.nerf.occupancy_grid import OccupancyGrid from mattport.structures.rays import RayBundle, RaySamples -class UniformSampler(nn.Module): +class Sampler(nn.Module): + """Generate Samples""" + + def __init__( + self, num_samples: int, occupancy_field: Optional[OccupancyGrid] = None, weight_threshold: float = 1e-4 + ) -> None: + super().__init__() + self.num_samples = num_samples + self.occupancy_field = occupancy_field + self.weight_threshold = weight_threshold + + @abstractmethod + def generate_ray_samples(self) -> RaySamples: + """Generate Ray Samples""" + + def forward(self, *args, **kwargs) -> RaySamples: + """Generate ray samples with optional occupancy filtering""" + ray_samples = self.generate_ray_samples(*args, **kwargs) + 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 UniformSampler(Sampler): """Sample uniformly along a ray""" def __init__( @@ -30,14 +58,11 @@ class UniformSampler(nn.Module): 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 + super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified - self.occupancy_field = occupancy_field - self.weight_threshold = weight_threshold @torch.no_grad() - def forward( + def generate_ray_samples( self, ray_bundle: RayBundle = None, num_samples: Optional[int] = None, @@ -51,6 +76,7 @@ class UniformSampler(nn.Module): Returns: RaySamples: Positions and deltas for samples along a ray """ + assert ray_bundle is not None assert ray_bundle.nears is not None assert ray_bundle.fars is not None @@ -73,17 +99,10 @@ 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): +class PDFSampler(Sampler): """Sample based on probability distribution""" def __init__( @@ -103,19 +122,15 @@ class PDFSampler(nn.Module): 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 + super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified - self.occupancy_field = occupancy_field - self.weight_threshold = weight_threshold + self.include_original = include_original @torch.no_grad() - def forward( + def generate_ray_samples( self, - ray_samples: RaySamples, - weights: TensorType[..., "num_samples"], + ray_samples: RaySamples = None, + weights: TensorType[..., "num_samples"] = None, num_samples: Optional[int] = None, eps: float = 1e-5, ) -> RaySamples: @@ -179,11 +194,4 @@ class PDFSampler(nn.Module): 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 From 8a6b7c3e465d272d7aa940c6c60e73f82f3d18c2 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 19 May 2022 18:38:30 -0700 Subject: [PATCH 3/3] Refactor ray samples --- mattport/nerf/graph/vanilla_nerf.py | 2 +- mattport/nerf/sampler.py | 11 +-- mattport/structures/rays.py | 142 +++++++++++++++------------- tests/nerf/test_sampler.py | 2 +- 4 files changed, 81 insertions(+), 76 deletions(-) diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 931e2bec2..5eb4f7fe7 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -179,7 +179,7 @@ class NeRFGraph(Graph): depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.ts) # pdf sampling - ray_samples_pdf = self.sampler_pdf(ray_samples_uniform, weights_coarse) + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field: field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 9269322d9..33a85921d 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -94,10 +94,7 @@ class UniformSampler(Sampler): else: ts = (bins[:, 1:] + bins[:, :-1]) / 2 - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_bundle, - ) + ray_samples = ray_bundle.get_ray_samples(ts) return ray_samples @@ -129,6 +126,7 @@ class PDFSampler(Sampler): @torch.no_grad() def generate_ray_samples( self, + ray_bundle: RayBundle = None, ray_samples: RaySamples = None, weights: TensorType[..., "num_samples"] = None, num_samples: Optional[int] = None, @@ -189,9 +187,6 @@ class PDFSampler(Sampler): # Stop gradients ts = ts.detach() - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_samples.ray_bundle, - ) + ray_samples = ray_bundle.get_ray_samples(ts) return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 56a8134f7..7593041ea 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -21,6 +21,50 @@ class PointSamples: valid_mask: TensorType[..., 1] = None +@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 + + 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, valid_mask=self.valid_mask) + + def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: + """Return weights based on predicted densities + + Args: + densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray + + Returns: + TensorType[..., "num_samples"]: Weights for each sample + """ + + delta_density = self.deltas * densities[..., 0] + alphas = 1 - torch.exp(-delta_density) + + transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) + transmittance = torch.cat( + [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1 + ) + transmittance = torch.exp(-transmittance) # [..., "num_samples"] + + 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 + + @dataclass class RayBundle: """_summary_ @@ -86,6 +130,38 @@ class RayBundle: valid_mask=self.valid_mask[valid_mask] if is_not_none(self.valid_mask) else None, ) + def get_ray_samples(self, ts: TensorType["num_rays", "num_samples"]) -> RaySamples: + """ + Args: + ts (TensorType["num_rays", "num_samples"]): _description_ + + Returns: + RaySamples: _description_ + """ + positions = self.origins[:, None] + ts[:, :, None] * self.directions[:, None] + directions = self.directions.unsqueeze(1).repeat(1, positions.shape[1], 1) + valid_mask = torch.ones_like(ts, dtype=torch.bool) + + dists = ts[..., 1:] - ts[..., :-1] + dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples] + deltas = dists * torch.norm(self.directions[..., None, :], dim=-1) + + if self.camera_indices is not None: + camera_indices = self.camera_indices.unsqueeze(1).repeat(1, positions.shape[1]) + else: + camera_indices = None + + ray_samples = RaySamples( + positions=positions, + directions=directions, + camera_indices=camera_indices, + valid_mask=valid_mask, + ts=ts, + deltas=deltas, + ) + + return ray_samples + @dataclass class CameraRayBundle: @@ -131,69 +207,3 @@ class CameraRayBundle: directions=self.directions.view(-1, 3)[start_idx:end_idx], camera_indices=camera_indices, ) - - -class RaySamples: - """_summary_""" - - def __init__( - self, - ts: TensorType["num_rays", "num_samples"], - ray_bundle: RayBundle, - ) -> None: - self.ray_bundle = ray_bundle - self.ts = ts - 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) -> PointSamples: - """Convert to PointSamples instance and return.""" - # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) - - def get_camera_indices(self): - """Returns camera indices.""" - assert not isinstance(self.ray_bundle.camera_indices, type(None)), "Camera indices cannot be None here." - camera_indices = self.ray_bundle.camera_indices.unsqueeze(1).repeat(1, self.positions.shape[1]) - return camera_indices - - def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]: - """Returns positions.""" - return ray_bundle.origins[:, None] + self.ts[:, :, None] * ray_bundle.directions[:, None] - - def get_deltas(self) -> TensorType[..., "num_samples"]: - """Returns deltas.""" - dists = self.ts[..., 1:] - self.ts[..., :-1] - dist_inf = 1e10 * torch.ones_like(dists[..., -1:]) - dists = torch.cat([dists, dist_inf], -1) # [N_rays, N_samples] - deltas = dists * torch.norm(self.ray_bundle.directions[..., None, :], dim=-1) - return deltas - - def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: - """Return weights based on predicted densities - - Args: - densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray - - Returns: - TensorType[..., "num_samples"]: Weights for each sample - """ - - delta_density = self.deltas * densities[..., 0] - alphas = 1 - torch.exp(-delta_density) - - transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) - transmittance = torch.cat( - [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1 - ) - transmittance = torch.exp(-transmittance) # [..., "num_samples"] - - 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 diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py index 92a6848bb..df88b889c 100644 --- a/tests/nerf/test_sampler.py +++ b/tests/nerf/test_sampler.py @@ -47,7 +47,7 @@ def test_pdf_sampler(): # Just check that it doesn't crash. Need to add some actual tests. pdf_sampler = PDFSampler(num_samples) - pdf_sampler(coarse_ray_samples, weights, num_samples) + pdf_sampler(ray_bundle, coarse_ray_samples, weights, num_samples) # TODO Tancik: Add more precise tests