From cc7bc66afc2312f1aa8d1b9e850dcf04ca0a0ff2 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 26 Sep 2022 19:12:41 -0700 Subject: [PATCH] Fix usage of sdist/tdist in pdf sampler (#552) Using sdist during sampling --- nerfactory/cameras/rays.py | 28 ++++++--- nerfactory/fields/compound_field.py | 2 +- nerfactory/models/modules/ray_sampler.py | 72 ++++++++++++++++++++---- nerfactory/models/proposal.py | 10 ++-- 4 files changed, 88 insertions(+), 24 deletions(-) diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py index bb10b6e7c..26201023a 100644 --- a/nerfactory/cameras/rays.py +++ b/nerfactory/cameras/rays.py @@ -17,7 +17,7 @@ Some ray datastructures. """ import random from dataclasses import dataclass -from typing import Dict, Optional +from typing import Callable, Dict, Optional import torch from torchtyping import TensorType @@ -93,6 +93,12 @@ class RaySamples(TensorDataclass): """Rays that are valid.""" deltas: Optional[TensorType["bs":..., 1]] = None """"width" of each sample.""" + spacing_starts: Optional[TensorType["bs":..., "num_samples", 1]] = None + """Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling.""" + spacing_ends: Optional[TensorType["bs":..., "num_samples", 1]] = None + """Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling.""" + spacing_to_euclidean_fn: Optional[Callable] = None + """Function to convert bins to euclidean distance.""" metadata: Optional[Dict[str, TensorType["bs":..., "latent_dims"]]] = None """addtional information relevant to generating ray samples""" @@ -207,7 +213,12 @@ class RayBundle(TensorDataclass): return self.flatten()[start_idx:end_idx] def get_ray_samples( - self, bin_starts: TensorType["bs":..., "num_samples", 1], bin_ends: TensorType["bs":..., "num_samples", 1] + self, + bin_starts: TensorType["bs":..., "num_samples", 1], + bin_ends: TensorType["bs":..., "num_samples", 1], + spacing_starts: Optional[TensorType["bs":..., "num_samples", 1]], + spacing_ends: Optional[TensorType["bs":..., "num_samples", 1]], + spacing_to_euclidean_fn: Optional[Callable], ) -> RaySamples: """Produces samples for each ray by projection points along the ray direction. Currently samples uniformly. @@ -222,7 +233,7 @@ class RayBundle(TensorDataclass): valid_mask = torch.ones((bin_starts.shape), dtype=torch.bool, device=device) - dists = bin_ends - bin_starts # [..., N_samples, 1] + dists = bin_ends - bin_starts # [..., num_samples, 1] deltas = dists * torch.norm(self.directions[:, :], dim=-1)[..., None, None] if self.camera_indices is not None: @@ -235,16 +246,19 @@ class RayBundle(TensorDataclass): frustums = Frustums( origins=shaped_raybundle_fields.origins, # [..., 1, 3] directions=shaped_raybundle_fields.directions, # [..., 1, 3] - starts=bin_starts, # [..., N_samples, 1] - ends=bin_ends, # [..., N_samples, 1] + starts=bin_starts, # [..., num_samples, 1] + ends=bin_ends, # [..., num_samples, 1] pixel_area=shaped_raybundle_fields.pixel_area, # [..., 1, 1] ).to(device) ray_samples = RaySamples( frustums=frustums, camera_indices=camera_indices, # [..., 1, 1] - valid_mask=valid_mask, # [..., N_samples, 1] - deltas=deltas, # [..., N_samples, 1] + valid_mask=valid_mask, # [..., num_samples, 1] + deltas=deltas, # [..., num_samples, 1] + spacing_starts=spacing_starts, # [..., num_samples, 1] + spacing_ends=spacing_ends, # [..., num_samples, 1] + spacing_to_euclidean_fn=spacing_to_euclidean_fn, metadata=shaped_raybundle_fields.metadata, ) diff --git a/nerfactory/fields/compound_field.py b/nerfactory/fields/compound_field.py index a8071ab6f..49f1ac944 100644 --- a/nerfactory/fields/compound_field.py +++ b/nerfactory/fields/compound_field.py @@ -83,7 +83,7 @@ class TCNNCompoundField(Field): num_layers_color: int = 3, hidden_dim_color: int = 64, appearance_embedding_dim: int = 32, - spatial_distortion: SpatialDistortion = SceneContraction(), + spatial_distortion: Optional[SpatialDistortion] = None, ) -> None: super().__init__() diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index 8edc2dd63..d3cadbf7c 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/nerfactory/models/modules/ray_sampler.py @@ -110,13 +110,12 @@ class SpacedSampler(Sampler): assert ray_bundle.fars is not None num_samples = num_samples or self.num_samples + assert num_samples is not None num_rays = ray_bundle.origins.shape[0] bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)[None, ...] # [1, num_samples+1] - s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars)) - bins = self.spacing_fn_inv(bins * s_far + (1 - bins) * s_near) # [num_rays, num_samples+1] - + # TODO More complicated than it needs to be. if self.train_stratified and self.training: t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device) bin_centers = (bins[..., 1:] + bins[..., :-1]) / 2.0 @@ -124,7 +123,17 @@ class SpacedSampler(Sampler): bin_lower = torch.cat([bins[..., :1], bin_centers], -1) bins = bin_lower + (bin_upper - bin_lower) * t_rand - ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None]) + s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars)) + spacing_to_euclidean_fn = lambda x: self.spacing_fn_inv(x * s_far + (1 - x) * s_near) + euclidean_bins = spacing_to_euclidean_fn(bins) # [num_rays, num_samples+1] + + ray_samples = ray_bundle.get_ray_samples( + bin_starts=euclidean_bins[..., :-1, None], + bin_ends=euclidean_bins[..., 1:, None], + spacing_starts=bins[..., :-1, None], + spacing_ends=bins[..., 1:, None], + spacing_to_euclidean_fn=spacing_to_euclidean_fn, + ) return ray_samples @@ -237,6 +246,35 @@ class LogSampler(SpacedSampler): ) +class UniformLinDispPiecewiseSampler(SpacedSampler): + """Piecewise sampler along a ray that allocates the first half of the samples uniformly and the second half + using linearly in disparity spacing. + + + Args: + num_samples: Number of samples per ray + train_stratified: Use stratified sampling during training. Defults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + """ + + def __init__( + self, + num_samples: Optional[int] = None, + train_stratified=True, + density_field: Optional[DensityGrid] = None, + weight_threshold: float = 1e-2, + ) -> None: + super().__init__( + num_samples=num_samples, + spacing_fn=lambda x: torch.where(x < 1, x / 2, 1 - 1 / (2 * x)), + spacing_fn_inv=lambda x: torch.where(x < 0.5, 2 * x, 1 / (2 - 2 * x)), + train_stratified=train_stratified, + density_field=density_field, + weight_threshold=weight_threshold, + ) + + class PDFSampler(Sampler): """Sample based on probability distribution @@ -289,6 +327,7 @@ class PDFSampler(Sampler): raise ValueError("ray_samples and ray_bundle must be provided") num_samples = num_samples or self.num_samples + assert num_samples is not None num_bins = num_samples + 1 weights = weights[..., 0] + self.histogram_padding @@ -314,12 +353,14 @@ class PDFSampler(Sampler): u = u.expand(size=(*cdf.shape[:-1], num_bins)) u = u.contiguous() - # Force bins to not have a gap between them. Kinda hacky, should reconsider. + assert ( + ray_samples.spacing_starts is not None and ray_samples.spacing_ends is not None + ), "ray_sample spacing_starts and spacing_ends must be provided" + assert ray_samples.spacing_to_euclidean_fn is not None, "ray_samples.spacing_to_euclidean_fn must be provided" existing_bins = torch.cat( [ - ray_samples.frustums.starts[..., :1, 0], - (ray_samples.frustums.starts[..., 1:, 0] + ray_samples.frustums.ends[..., :-1, 0]) / 2.0, - ray_samples.frustums.ends[..., -1:, 0], + ray_samples.spacing_starts[..., 0], + ray_samples.spacing_ends[..., -1:, 0], ], dim=-1, ) @@ -341,7 +382,15 @@ class PDFSampler(Sampler): # Stop gradients bins = bins.detach() - ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None]) + euclidean_bins = ray_samples.spacing_to_euclidean_fn(bins) + + ray_samples = ray_bundle.get_ray_samples( + bin_starts=euclidean_bins[..., :-1, None], + bin_ends=euclidean_bins[..., 1:, None], + spacing_starts=bins[..., :-1, None], + spacing_ends=bins[..., 1:, None], + spacing_to_euclidean_fn=ray_samples.spacing_to_euclidean_fn, + ) return ray_samples @@ -502,8 +551,7 @@ class ProposalNetworkSampler(Sampler): self.num_nerf_samples_per_ray = num_nerf_samples_per_ray self.num_proposal_network_iterations = num_proposal_network_iterations - # NOTE: a linear disparity sampler won't work until we sample in s space instead of t space - self.uniform_sampler = UniformSampler() + self.initial_sampler = UniformLinDispPiecewiseSampler() self.pdf_sampler = PDFSampler(include_original=False) def generate_ray_samples( @@ -521,7 +569,7 @@ class ProposalNetworkSampler(Sampler): num_samples = self.num_proposal_samples_per_ray if is_prop else self.num_nerf_samples_per_ray if i_level == 0: # Uniform sampling because we need to start with some samples - ray_samples = self.uniform_sampler(ray_bundle, num_samples=num_samples) + ray_samples = self.initial_sampler(ray_bundle, num_samples=num_samples) else: # PDF sampling based on the last samples and their weights ray_samples = self.pdf_sampler(ray_bundle, ray_samples, weights, num_samples=num_samples) diff --git a/nerfactory/models/proposal.py b/nerfactory/models/proposal.py index 875f7e9e1..2555b5d34 100644 --- a/nerfactory/models/proposal.py +++ b/nerfactory/models/proposal.py @@ -53,7 +53,7 @@ class ProposalModelConfig(ModelConfig): _target: Type = field(default_factory=lambda: ProposalModel) near_plane: float = 0.05 """How far along the ray to start sampling.""" - far_plane: float = 20.0 + far_plane: float = 1000.0 """How far along the ray to stop sampling.""" randomize_background: bool = True """Whether to randomize the background color.""" @@ -84,14 +84,16 @@ class ProposalModel(Model): """Set the fields and modules.""" super().populate_modules() + scene_contraction = SceneContraction() + # Fields if self.config.use_appearance_conditioning: self.field = TCNNCompoundField( - self.scene_bounds.aabb, spatial_distortion=SceneContraction(), num_images=self.num_train_data + self.scene_bounds.aabb, spatial_distortion=scene_contraction, num_images=self.num_train_data ) else: - self.field = TCNNInstantNGPField(self.scene_bounds.aabb, spatial_distortion=SceneContraction()) - self.proposal_network = DensityField(self.scene_bounds.aabb, spatial_distortion=SceneContraction()) + self.field = TCNNInstantNGPField(self.scene_bounds.aabb, spatial_distortion=scene_contraction) + self.proposal_network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction) # Collider self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)