update camera structure

.vscode/settings.json

@@ -13,4 +13,7 @@
     "python.linting.pylintEnabled": true,
     "python.linting.flake8Enabled": false,
     "python.linting.enabled": true,
+    "python.linting.pylintArgs": [
+        "--generated-members=numpy.* ,torch.* ,cv2.* , cv.*"
+    ],
 }

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)