Refactor weights

mattport/nerf/graph/vanilla_nerf.py

@@ -137,19 +137,21 @@ class NeRFGraph(Graph):
         coarse_field_outputs = self.field_coarse(uniform_ray_samples)  # FieldOutputs
         # time end coarse network
 
+        coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY])
+
         coarse_renderer_outputs = self.renderer_rgb(
             rgb=coarse_field_outputs[FieldHeadNames.RGB],
-            density=coarse_field_outputs[FieldHeadNames.DENSITY],
-            deltas=uniform_ray_samples.deltas,
+            weights=coarse_weights,
         )  # RendererOutputs
         # fine network:
-        pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs)  # RaySamples
+        pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights)  # RaySamples
         fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
 
+        fine_weights = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY])
+
         fine_renderer_outputs = self.renderer_rgb(
             rgb=fine_field_outputs[FieldHeadNames.RGB],
-            density=fine_field_outputs[FieldHeadNames.DENSITY],
-            deltas=pdf_ray_samples.deltas,
+            weights=fine_weights,
         )  # RendererOutputs
         # outputs:
         outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb}

mattport/nerf/renderers.py

@@ -2,14 +2,11 @@
 Collection of renderers
 """
 from dataclasses import dataclass
-from typing import Set
 
 import torch
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf.field_modules import field_heads
-
 
 @dataclass
 class RendererOutputs:
@@ -19,84 +16,25 @@ class RendererOutputs:
     density: TensorType["num_rays", 1] = None
 
 
-class Renderer(nn.Module):
-    """Base Renderer. Intended to be subclassed"""
-
-    def required_field_outputs(self) -> Set[field_heads.FieldHead]:
-        """
-        Returns:
-            Set[RenderHead]: RenderHeads required for this renderer.
-        """
-        raise NotImplementedError
-
-    def forward(
-        self,
-        rgb: TensorType[..., "num_samples", 3],
-        density: TensorType[..., "num_samples", 1],
-        deltas: TensorType[..., "num_samples"],
-    ) -> RendererOutputs:
-        """Composite samples along ray and render image
-
-        Args:
-            in_tensor (TensorType[..., "num_samples", -1]): feature for each sample
-            deltas (TensorType[..., "num_samples"]): depth of each sample
-
-        Returns:
-            TensorType[..., "out_dim"]: Rendered ray
-        """
-        raise NotImplementedError
-
-
-class RGBRenderer(Renderer):
+class RGBRenderer(nn.Module):
     """Standard volumetic rendering."""
 
-    def required_field_outputs(self) -> Set[field_heads.FieldHead]:
-        return set(field_heads.RGBFieldHead, field_heads.DensityFieldHead)
-
     def forward(
         self,
         rgb: TensorType[..., "num_samples", 3],
-        density: TensorType[..., "num_samples", 1],
-        deltas: TensorType[..., "num_samples"],
+        weights: TensorType[..., "num_samples"],
     ) -> RendererOutputs:
         """Composite samples along ray and render color image
 
         Args:
-            in_tensor (TensorType[..., "num_samples", -1]): RGB for each sample
-            deltas (TensorType[..., "num_samples"]): depth influence for each sample
+            rgb (TensorType[..., "num_samples", -1]): RGB for each sample
+            weights (TensorType[..., "num_samples"]): Weights for each sample
 
         Returns:
             TensorType[..., "out_dim"]: Composited RGB ray
         """
-        delta_density = 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(rgb.device), transmittance], axis=-1)
-        transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
-
-        weights = alphas * transmittance  # [..., "num_samples"]
 
         rgb = torch.sum(weights[..., None] * rgb, dim=-2)
 
         renderer_outputs = RendererOutputs(rgb=rgb)
         return renderer_outputs
-
-
-class DepthRenderer(Renderer):
-    """_summary_
-
-    Args:
-        Renderer (_type_): _description_
-    """
-
-    def required_field_outputs(self) -> Set[field_heads.FieldHead]:
-        return set(field_heads.DensityFieldHead)
-
-    def forward(
-        self,
-        rgb: TensorType[..., "num_samples", 3],
-        density: TensorType[..., "num_samples", 1],
-        deltas: TensorType[..., "num_samples"],
-    ) -> RendererOutputs:
-        raise NotImplementedError

mattport/nerf/sampler.py

@@ -2,11 +2,10 @@
 Collection of sampling strategies
 """
 
-from typing import Dict, Optional
+from typing import Optional
 import torch
 from torch import nn
 from torchtyping import TensorType
-from mattport.nerf.field_modules.field_heads import FieldHeadNames
 
 from mattport.structures.rays import RayBundle, RaySamples
 
@@ -89,9 +88,8 @@ class PDFSampler(nn.Module):
 
     def forward(
         self,
-        ray_bundle: RayBundle,
         coarse_ray_samples: RaySamples,
-        field_outputs: Dict[FieldHeadNames, TensorType],
+        weights: TensorType[..., "num_samples"],
         num_samples: Optional[int] = None,
         eps: float = 1e-5,
     ) -> RaySamples:
@@ -108,27 +106,15 @@ class PDFSampler(nn.Module):
         Returns:
             RaySamples: Positions and deltas for samples along a ray
         """
+        # TODO (matt) Look into torch.no_grad(), or torch.inference_mode
 
         num_samples = num_samples or self.num_samples
-
-        # Calculate weight contributions along ray
-        # Todo(matt): This computation is duplicated
-        density = field_outputs[FieldHeadNames.DENSITY]
-        delta_density = coarse_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"]
-        weights = weights[..., 1:]
+        weights = weights[..., :-1]
 
         # 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)
-        weights += padding / weights.shape[-1]
+        weights = weights + padding / weights.shape[-1]
         weights_sum += padding
 
         pdf = weights / weights_sum
@@ -158,12 +144,15 @@ class PDFSampler(nn.Module):
 
         if self.include_original:
             ts, _ = torch.sort(torch.cat([coarse_ray_samples.ts, ts], -1), -1)
+        else:
+            ts, _ = torch.sort(ts, -1)
 
+        # Stop gradients
         ts = ts.detach()
 
         ray_samples = RaySamples(
             ts=ts,
-            ray_bundle=ray_bundle,
+            ray_bundle=coarse_ray_samples.ray_bundle,
         )
 
         return ray_samples

mattport/structures/rays.py

@@ -92,3 +92,26 @@ class RaySamples:
         dists = torch.cat([dists, dists[..., -1:]], -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

tests/nerf/test_renderers.py

@@ -12,20 +12,15 @@ def test_rgb_renderer():
     num_samples = 10
 
     rgb = torch.ones((3, num_samples, 3))
-    densities = torch.ones((3, num_samples, 1))
-    deltas = torch.ones((3, num_samples))
+    weights = torch.ones((3, num_samples))
+    weights /= torch.sum(weights, axis=-1, keepdim=True)
 
     rgb_renderer = renderers.RGBRenderer()
 
-    assert rgb_renderer.required_field_outputs is not None
-
-    out = rgb_renderer(rgb=rgb, density=densities, deltas=deltas)
+    out = rgb_renderer(rgb=rgb, weights=weights)
     assert torch.max(out.rgb) > 0.9
 
-    out = rgb_renderer(rgb=rgb * 0, density=densities, deltas=deltas)
-    assert torch.max(out.rgb) == pytest.approx(0)
-
-    out = rgb_renderer(rgb=rgb, density=densities * 0, deltas=deltas)
+    out = rgb_renderer(rgb=rgb * 0, weights=weights)
     assert torch.max(out.rgb) == pytest.approx(0)
 
     # TODO Tancik: Add more precise tests

tests/nerf/test_sampler.py

@@ -36,14 +36,14 @@ def test_pdf_sampler():
     directions = torch.ones_like(origins)
     ray_bundle = RayBundle(origins=origins, directions=directions)
 
-    field_outputs = {FieldHeadNames.DENSITY: torch.ones((10, num_samples, 1))}
-
     uniform_sampler = UniformSampler(near_plane=near_plane, far_plane=far_plane, num_samples=num_samples)
     coarse_ray_samples = uniform_sampler(ray_bundle)
 
+    weights = torch.ones((10, num_samples))
+
     # Just check that it doesn't crash. Need to add some actual tests.
     pdf_sampler = PDFSampler(num_samples)
-    pdf_sampler(ray_bundle, coarse_ray_samples, field_outputs, num_samples)
+    pdf_sampler(coarse_ray_samples, weights, num_samples)
 
     # TODO Tancik: Add more precise tests