diff --git a/.vscode/settings.json b/.vscode/settings.json index f476de807..2f20bab22 100644 --- a/.vscode/settings.json +++ b/.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" + ], } \ 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) 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)