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 5b30d940a..f1ba247ba 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 @@ -178,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 0ddea9091..33a85921d 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -2,99 +2,69 @@ Collection of sampling strategies """ +from abc import abstractmethod from typing import Optional + import torch from torch import nn from torchtyping import TensorType -from mattport.nerf.occupancy_grid import OccupancyGrid +from mattport.nerf.occupancy_grid import OccupancyGrid from mattport.structures.rays import RayBundle, RaySamples -class OccupancyGridSampler(nn.Module): - """Sample based on occupancy grid""" +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__( self, - occupancy_grid: OccupancyGrid, - num_samples: int = 256, - weight_threshold: float = 1e-4, + num_samples: int, train_stratified=True, + occupancy_field: Optional[OccupancyGrid] = None, + weight_threshold: float = 1e-4, ) -> 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: - """ - 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 + super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified @torch.no_grad() - def forward( + def generate_ray_samples( self, - ray_bundle: RayBundle, + ray_bundle: RayBundle = None, num_samples: Optional[int] = None, ) -> RaySamples: """Generates position samples uniformly. @@ -106,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 @@ -123,35 +94,41 @@ class UniformSampler(nn.Module): 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 -class PDFSampler(nn.Module): +class PDFSampler(Sampler): """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 + super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified + self.include_original = include_original @torch.no_grad() - def forward( + def generate_ray_samples( self, - coarse_ray_samples: RaySamples, - weights: TensorType[..., "num_samples"], + ray_bundle: RayBundle = None, + ray_samples: RaySamples = None, + weights: TensorType[..., "num_samples"] = None, num_samples: Optional[int] = None, eps: float = 1e-5, ) -> RaySamples: @@ -195,7 +172,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,16 +180,13 @@ 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) # Stop gradients ts = ts.detach() - ray_samples = RaySamples( - ts=ts, - ray_bundle=coarse_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 d1df55b97..7593041ea 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,51 @@ class PointSamples: positions: TensorType[..., 3] = None directions: TensorType[..., 3] = None camera_indices: TensorType[..., 1] = None + 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 @@ -85,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: @@ -130,64 +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() - - def to_point_samples(self): - """Convert to PointSamples instance and return.""" - # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions) - - 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 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