Merge branch 'master' into ethan/initial_trainer_with_config

.vscode/settings.json

@@ -11,7 +11,6 @@
         "--line-length=120"
     ],
     "python.linting.pylintEnabled": true,
-    "python.linting.pylintArgs": [],
     "python.linting.flake8Enabled": false,
     "python.linting.enabled": true,
     "python.testing.pytestArgs": [
@@ -19,4 +18,8 @@
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,
+    "python.linting.pylintArgs": [
+        "--generated-members=numpy.* ,torch.* ,cv2.* , cv.*",
+        "--max-line-length=120"
+    ],
 }

mattport/nerf/cameras.py

@@ -3,6 +3,7 @@ Camera Models
 """
 from dataclasses import dataclass
 from torchtyping import TensorType
+import torch
 
 
 @dataclass
@@ -16,26 +17,44 @@ class Rays:
 class Camera:
     """Base Camera. Intended to be subclassed"""
 
-    def __init__(self, height, width, camera_to_world) -> None:
+    def __init__(self, height: int, width: int, camera_to_world: TensorType[4, 3]) -> None:
+        """Initialize camera
+
+        Args:
+            height (int): Number of camera sensor rows
+            width (int): Number of camera sensor columns
+            camera_to_world (TensorType[4, 3]): Pose transorm for camera
+        """
         self.height = height
         self.width = width
         self.camera_to_world = camera_to_world
 
-    def get_rays(self, coords: TensorType[..., 2]) -> Rays:
-        """Encodes an input tensor.
+    def generate_rays(self, coords: TensorType[..., 2]) -> Rays:
+        """Generates camera rays associated to grid of pixel coordinates.
 
         Args:
-            coords (TensorType[..., 2]): Input tensor to be encoded
+            coords (TensorType[..., 2]): Pixel coordiantes
 
         Returns:
-            TensorType[..., "output_dim"]: A encoded input tensor
+            Rays: Camera rays associated to pixel coordinates.
         """
         raise NotImplementedError
 
-    def get_all_rays(self) -> Rays:
+    def get_image_coords(self) -> TensorType["height", "width", 2]:
+        """Get grid of image coordinates
+
+        Returns:
+            TensorType["height", "width", 2]: Image coordinates
+        """
+        coords = torch.meshgrid(torch.arange(self.height), torch.arange(self.width))
+        coords = torch.stack(coords, dim=-1)
+        return coords
+
+    def generate_all_rays(self) -> Rays:
         """Generate rays for entire camera
 
         Returns:
             Rays: Camera rays of shape [height, width]
         """
-        raise NotImplementedError
+        coords = self.get_image_coords()
+        return self.generate_rays(coords)

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)