From 69977da8e641fda3f0bf866b7e17fe1f6e6f350d Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 5 Apr 2022 14:33:25 -0700 Subject: [PATCH 1/2] update camera structure --- .vscode/settings.json | 3 +++ mattport/nerf/cameras.py | 33 ++++++++++++++++++++++++++------- 2 files changed, 29 insertions(+), 7 deletions(-) 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) From 219f51a5acd592cf136c0af4ac9b4fc04eea3626 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 5 Apr 2022 15:06:36 -0700 Subject: [PATCH 2/2] 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)