From 0bc278c84880577e616af03ce0e11fcb0b2dd912 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Fri, 8 Apr 2022 18:02:17 -0700
Subject: [PATCH] add structure for renderer

---
 mattport/nerf/render_heads.py |  3 +--
 mattport/nerf/renderer.py     |  0
 mattport/nerf/sampler.py      | 25 +++++++++++++++++++------
 3 files changed, 20 insertions(+), 8 deletions(-)
 delete mode 100644 mattport/nerf/renderer.py

diff --git a/mattport/nerf/render_heads.py b/mattport/nerf/render_heads.py
index 7dabcade8..be509d80b 100644
--- a/mattport/nerf/render_heads.py
+++ b/mattport/nerf/render_heads.py
@@ -1,7 +1,6 @@
 """
 Collection of render heads
 """
-from enum import Enum
 from typing import Optional
 
 from torch import nn
@@ -39,7 +38,7 @@ class RenderHead(nn.Module):
 
 
 class DensityHead(RenderHead):
-    """RGB head"""
+    """Density head"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus) -> None:
         super().__init__(in_dim, 1, activation)
diff --git a/mattport/nerf/renderer.py b/mattport/nerf/renderer.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py
index 52da02629..90428f872 100644
--- a/mattport/nerf/sampler.py
+++ b/mattport/nerf/sampler.py
@@ -1,12 +1,18 @@
 """
 Collection of sampling strategies
 """
+from typing import NamedTuple
 from torchtyping import TensorType
 from torch import nn
 import torch
 import nerf.cameras as cameras
 
 
+RaySamples = NamedTuple(
+    "RaySamples", [("locations", TensorType[..., "num_samples", 3]), ("deltas", TensorType[..., "num_samples"])]
+)
+
+
 class Sampler(nn.Module):
     """Base Sampler. Intended to be subclassed"""
 
@@ -26,17 +32,18 @@ class Sampler(nn.Module):
 
     def generate_samples(
         self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int
-    ) -> TensorType[..., "num_samples", 3]:
+    ) -> RaySamples:
         """Encodes an input tensor.
 
         Args:
             camera_rays (cameras.Rays): Rays to generate samples for
-
-        Returns:
             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
+
+        Returns:
+            RaySamples: Positions and deltas for samples along a ray
         """
         raise NotImplementedError
 
@@ -55,15 +62,21 @@ class UniformSampler(Sampler):
 
         Args:
             camera_rays (cameras.Rays): Rays to generate samples for
-
-        Returns:
             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
+
+        Returns:
+            RaySamples: Positions and deltas for samples along a ray
         """
         near_plane = near_plane or self.near_plane
         far_plane = far_plane or self.far_plane
         num_samples = num_samples or self.num_samples
 
-        return torch.linspace(near_plane, far_plane, num_samples)
+        bins = torch.linspace(near_plane, far_plane, num_samples + 1)
+
+        positions = (bins[..., 1:] + bins[..., :-1]) / 2.0
+        deltas = bins[..., 1:] - bins[..., :-1]
+
+        return RaySamples(locations=positions, deltas=deltas)