From 219f51a5acd592cf136c0af4ac9b4fc04eea3626 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Tue, 5 Apr 2022 15:06:36 -0700
Subject: [PATCH] add sampler boilerplate

---
 .vscode/settings.json    |  3 +-
 mattport/nerf/sampler.py | 69 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 51b829b7f..5c0e9d025 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -14,6 +14,7 @@
     "python.linting.flake8Enabled": false,
     "python.linting.enabled": true,
     "python.linting.pylintArgs": [
-        "--generated-members=numpy.* ,torch.* ,cv2.* , cv.*"
+        "--generated-members=numpy.* ,torch.* ,cv2.* , cv.*",
+        "--max-line-length=120"
     ],
 }
\ No newline at end of file
diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py
index e69de29bb..52da02629 100644
--- a/mattport/nerf/sampler.py
+++ b/mattport/nerf/sampler.py
@@ -0,0 +1,69 @@
+"""
+Collection of sampling strategies
+"""
+from torchtyping import TensorType
+from torch import nn
+import torch
+import nerf.cameras as cameras
+
+
+class Sampler(nn.Module):
+    """Base Sampler. Intended to be subclassed"""
+
+    def __init__(self, near_plane: float, far_plane: float, num_samples: int) -> None:
+        """Inialize sampler
+
+        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 generate_samples(
+        self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int
+    ) -> TensorType[..., "num_samples", 3]:
+        """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
+        """
+        raise NotImplementedError
+
+    def forward(self, camera_rays: cameras.Rays, num_samples: int) -> TensorType[..., "num_samples", 3]:
+        """Call forward"""
+        return self.Sampler(camera_rays, num_samples)
+
+
+class UniformSampler(Sampler):
+    """Base Sampler. Intended to be subclassed"""
+
+    def generate_samples(
+        self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int
+    ) -> TensorType[..., "num_samples", 3]:
+        """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
+        """
+        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)