Get code running

.vscode/settings.json

@@ -31,6 +31,7 @@
         "--max-line-length=120",
         "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u",
         "--disable=too-many-arguments",
-        "--disable=fixme"
+        "--disable=fixme",
+        "--disable-msg=too-many-locals"
     ],
 }

mattport/nerf/graph/nerf.py

@@ -10,7 +10,7 @@ from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
 from mattport.nerf.renderers import RGBRenderer
-from mattport.nerf.sampler import UniformSampler
+from mattport.nerf.sampler import PDFSampler, UniformSampler
 
 
 class NeRFGraph(Graph):
@@ -29,7 +29,7 @@ class NeRFGraph(Graph):
 
         # samplers
         self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
-        self.sampler_pdf = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
+        self.sampler_pdf = PDFSampler(num_samples=64)
 
         # field
         self.field_coarse = NeRFField()
@@ -52,7 +52,7 @@ class NeRFGraph(Graph):
             rgb=coarse_field_outputs.rgb, density=coarse_field_outputs.density, deltas=uniform_ray_samples.deltas
         )  # RendererOutputs
         # fine network:
-        pdf_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
+        pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs.density)  # RaySamples
         fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
         fine_renderer_outputs = self.renderer_rgb(
             rgb=fine_field_outputs.rgb, density=fine_field_outputs.density, deltas=pdf_ray_samples.deltas

mattport/nerf/sampler.py

@@ -2,6 +2,7 @@
 Collection of sampling strategies
 """
 
+from typing import Optional
 import torch
 from torch import nn
 from torchtyping import TensorType
@@ -12,17 +13,34 @@ from mattport.structures.rays import RayBundle, RaySamples
 class UniformSampler(nn.Module):
     """Sample uniformly along a ray"""
 
+    def __init__(self, near_plane: float, far_plane: float, num_samples: int) -> 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
+        """
+        super().__init__()
+
+        self.near_plane = near_plane
+        self.far_plane = far_plane
+        self.num_samples = num_samples
+
     def forward(
-        self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int
+        self,
+        ray_bundle: RayBundle,
+        near_plane: Optional[float] = None,
+        far_plane: Optional[float] = None,
+        num_samples: Optional[int] = None,
     ) -> TensorType[..., "num_samples", 3]:
         """Generates position samples uniformly.
 
         Args:
-            camera_rays (cameras.Rays): Rays to generate samples for
+            ray_bundle (RayBundle): Rays to generate samples for
             TensorType[..., "output_dim"]: A encoded input tensor
-            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
+            near_plane (Optional[float]): Minimum distance along ray to sample
+            far_plane (Optional[float]): Maximum distance along ray to sample
+            num_samples (Optional[int]): Number of samples per ray
 
         Returns:
             RaySamples: Positions and deltas for samples along a ray
@@ -31,36 +49,48 @@ class UniformSampler(nn.Module):
         far_plane = far_plane or self.far_plane
         num_samples = num_samples or self.num_samples
 
-        bins = torch.linspace(near_plane, far_plane, num_samples + 1)
+        num_rays = ray_bundle.origins.shape[0]
 
-        t = (bins[..., 1:] + bins[..., :-1]) / 2.0
-        deltas = bins[..., 1:] - bins[..., :-1]
+        bins = torch.linspace(near_plane, far_plane, num_samples + 1).to(ray_bundle.origins.device)  # (num_samples+1,)
+        bins = bins.unsqueeze(0).repeat(num_rays, 1)  # (num_rays, num_samples+1)
 
-        positions = camera_rays.origin + camera_rays.direction * t
+        ray_samples = RaySamples(
+            bins=bins,
+            ray_bundle=ray_bundle,
+        )
 
-        return RaySamples(locations=positions, deltas=deltas)
+        return ray_samples
 
 
 class PDFSampler(nn.Module):
     """Sample based on probability distribution"""
 
+    def __init__(self, num_samples: int) -> None:
+        """
+        Args:
+            num_samples (int): Number of samples per ray
+        """
+        super().__init__()
+
+        self.num_samples = num_samples
+
     def forward(
         self,
-        camera_rays: cameras.Rays,
-        num_samples: int,
-        bins: TensorType[..., "num_samples"],
-        weights: TensorType[..., "num_samples"],
+        ray_bundle: RayBundle,
+        coase_ray_samples: RaySamples,
+        density: TensorType[..., "num_samples", 1],
+        num_samples: Optional[int] = None,
         randomized: bool = True,
         eps: float = 1e-5,
     ) -> TensorType[..., "num_samples", 3]:
         """Generates position samples given a distribution.
 
         Args:
-            camera_rays (cameras.Rays): Rays to generate samples for
+            ray_bundle (RayBundle): Rays to generate samples for
             TensorType[..., "output_dim"]: A encoded input tensor
-            num_samples (int): Number of samples per ray
             bins (TensorType[..., "num_samples"]): Ray bins
             weights: (TensorType[..., "num_samples"]): Weights for each bin
+            num_samples (Optional[int]): Number of samples per ray
             randomized: boolean: Randomize location within each bin. Defaults to True
             eps: float: Small value to prevent numerical issues. Defaults to 1e-5
 
@@ -68,6 +98,20 @@ class PDFSampler(nn.Module):
             RaySamples: Positions and deltas for samples along a ray
         """
 
+        num_samples = num_samples or self.num_samples
+
+        # Calculate weight contributions along ray
+        # Todo(matt): This computation is duplicated
+        delta_density = coase_ray_samples.deltas * density[..., 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(density.device), transmittance], axis=-1
+        )
+        transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
+        weights = alphas * transmittance  # [..., "num_samples"]
+
         # Add small offset to rays with zero weight to prevent NaNs
         weights_sum = torch.sum(weights, dim=-1, keepdim=True)
         padding = torch.relu(eps - weights_sum)
@@ -75,38 +119,33 @@ class PDFSampler(nn.Module):
         weights_sum += padding
 
         pdf = weights / weights_sum
-        cdf = torch.min(1, torch.cumsum(pdf, dim=-1))
+        cdf = torch.min(torch.ones_like(pdf), torch.cumsum(pdf, dim=-1))
         cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)
 
         if randomized:
-            u = torch.rand(size=(*cdf.shape[:-1], num_samples))
+            u = torch.rand(size=(*cdf.shape[:-1], num_samples), device=cdf.device)
         else:
-            u = torch.linspace(0.0, 1.0, steps=num_samples)
+            u = torch.linspace(0.0, 1.0, steps=num_samples, device=cdf.device)
             u = torch.expand(size=(*cdf.shape[:-1], num_samples))
 
-        # u = u.contigous()
+        # u = u.contiguous()
         indicies = torch.searchsorted(cdf, u, right=True)
         below = torch.max(torch.zeros_like(indicies), indicies - 1)
-        above = torch.min(cdf.shape[-1] - 1, indicies)
+        above = torch.min(cdf.shape[-1] - torch.ones_like(indicies), indicies)
 
         indicies_g = torch.stack([below, above], -1)
-        cdf_g = torch.gather(cdf, indicies_g, axis=-1, batch_dims=len(indicies_g.shape) - 2)
-        bins_g = torch.gather(bins, indicies_g, axis=-1, batch_dims=len(indicies_g.shape) - 2)
+        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(coase_ray_samples.bins.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)
         t = (u - cdf_g[..., 0]) / denom
         samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
 
-        deltas = samples[..., 1:] - samples[..., :-1]
-
-        positions = camera_rays.origin + camera_rays.direction * t
-
-        num_rays = ray_bundle.origins.shape[0]
         ray_samples = RaySamples(
-            t_min=torch.ones_like(ray_bundle.origins[:, 0]) * near_plane,
-            t_max=torch.ones_like(ray_bundle.origins[:, 0]) * far_plane,
-            ts=bins.unsqueeze(0).repeat(num_rays, 1),  # (num_rays, num_samples)
+            bins=samples,
             ray_bundle=ray_bundle,
         )
+
         return ray_samples

mattport/structures/rays.py

@@ -72,14 +72,10 @@ class RaySamples:
 
     def __init__(
         self,
-        t_min: TensorType["num_rays"],
-        t_max: TensorType["num_rays"],
-        ts: TensorType["num_rays", "num_samples+1"],
+        bins: TensorType["num_rays", "num_samples+1"],
         ray_bundle: RayBundle,
     ) -> None:
-        self.t_min = t_min
-        self.t_max = t_max
-        self.ts = ts
+        self.bins = bins
         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()
@@ -87,11 +83,11 @@ class RaySamples:
     def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]:
         """Returns positions."""
         # TODO(ethan): check with Matt about this
-        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2  # (num_rays, num_samples)
+        t_mid = (self.bins[:, 1:] + self.bins[:, :-1]) / 2  # (num_rays, num_samples)
         return ray_bundle.origins[:, None] + t_mid[:, :, None] * ray_bundle.directions[:, None]
 
     def get_deltas(self) -> TensorType["num_samples"]:
         """Returns deltas."""
-        deltas = self.ts[..., 1:] - self.ts[..., :-1]
+        deltas = self.bins[..., 1:] - self.bins[..., :-1]
         # TODO(ethan): check this with Matt
         return deltas