diff --git a/.vscode/settings.json b/.vscode/settings.json index 8c50b4b35..51b829b7f 100644 --- a/.vscode/settings.json +++ b/.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.*" + ], } \ No newline at end of file diff --git a/mattport/nerf/cameras.py b/mattport/nerf/cameras.py index 51dd72a62..9f9afe388 100644 --- a/mattport/nerf/cameras.py +++ b/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)