Fix usage of sdist/tdist in pdf sampler (#552)

Using sdist during sampling

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,
         )
 

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__()
 

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)

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)