From d512391df7c44189bfd85272eeec0e043e0b10f0 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 23 May 2022 11:09:28 -0700 Subject: [PATCH] Remove CameraRayBundle --- mattport/nerf/graph/base.py | 2 +- mattport/structures/cameras.py | 46 ++++---- mattport/structures/rays.py | 177 ++++++++++++++++--------------- tests/structures/test_cameras.py | 4 +- 4 files changed, 118 insertions(+), 111 deletions(-) diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index af309e976..fa0ee1cec 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -138,7 +138,7 @@ class Graph(AbstractGraph): image_height, image_width = camera_ray_bundle.origins.shape[:2] device = self.get_device() - num_rays = camera_ray_bundle.get_num_rays() + num_rays = len(camera_ray_bundle) outputs = {} outputs_lists = defaultdict(list) diff --git a/mattport/structures/cameras.py b/mattport/structures/cameras.py index 1d27fe408..60c721246 100644 --- a/mattport/structures/cameras.py +++ b/mattport/structures/cameras.py @@ -9,7 +9,7 @@ import torch from torch.nn.functional import normalize from torchtyping import TensorType -from mattport.structures.rays import CameraRayBundle, RayBundle +from mattport.structures.rays import RayBundle @dataclass @@ -98,21 +98,19 @@ class Camera: """ return - def generate_camera_rays(self) -> CameraRayBundle: + def generate_camera_rays(self) -> RayBundle: """Generate rays for the camera. Returns: Rays: Camera rays of shape [height, width] """ - image_height = self.get_image_height() - image_width = self.get_image_width() - num_rays = image_height * image_width - intrinsics = self.get_intrinsics().unsqueeze(0).repeat(num_rays, 1) # (num_rays, num_intrinsics_params) - camera_to_world = self.camera_to_world.unsqueeze(0).repeat(num_rays, 1, 1) # (num_rays, 3, 4) - coords = self.get_image_coords().view(num_rays, 2) # (num_rays, 2) + height = self.get_image_height() + width = self.get_image_width() + intrinsics = self.get_intrinsics().unsqueeze(0).repeat(height, height, 1) # (num_rays, num_intrinsics_params) + camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1) # (num_rays, 3, 4) + coords = self.get_image_coords() ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords) - camera_ray_bundle = ray_bundle.to_camera_ray_bundle(image_height, image_width) - return camera_ray_bundle + return ray_bundle class PinholeCamera(Camera): @@ -173,24 +171,24 @@ class PinholeCamera(Camera): @classmethod def generate_rays( cls, - intrinsics: TensorType["num_rays", "num_intrinsics_params"], - camera_to_world: TensorType["num_rays", 3, 4], - coords: TensorType["num_rays", 2], + intrinsics: TensorType[..., "num_intrinsics_params"], + camera_to_world: TensorType[..., 3, 4], + coords: TensorType[..., 2], ) -> RayBundle: - cx = intrinsics[:, 0:1] - cy = intrinsics[:, 1:2] - fx = intrinsics[:, cls.fx_index() : cls.fx_index() + 1] - fy = intrinsics[:, cls.fy_index() : cls.fy_index() + 1] - y = coords[:, 0:1] # (num_rays, 1) - x = coords[:, 1:2] # (num_rays, 1) - original_directions = torch.cat([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1) # (num_rays, 3) - rotation = camera_to_world[:, :3, :3] # (num_rays, 3, 3) + cx = intrinsics[..., 0:1] + cy = intrinsics[..., 1:2] + fx = intrinsics[..., cls.fx_index() : cls.fx_index() + 1] + fy = intrinsics[..., cls.fy_index() : cls.fy_index() + 1] + y = coords[..., 0:1] # (..., 1) + x = coords[..., 1:2] # (..., 1) + original_directions = torch.cat([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1) # (..., 3) + rotation = camera_to_world[..., :3, :3] # (..., 3, 3) directions = torch.sum( - original_directions[:, None, :] * rotation, dim=-1 - ) # (num_rays, 1, 3) * (num_rays, 3, 3) -> (num_rays, 3) + original_directions[..., None, :] * rotation, dim=-1 + ) # (..., 1, 3) * (..., 3, 3) -> (..., 3) directions = normalize(directions, dim=-1) - origins = camera_to_world[:, :3, 3] # (num_rays, 3) + origins = camera_to_world[..., :3, 3] # (..., 3) return RayBundle(origins=origins, directions=directions) diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 616c00850..780208a22 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -13,24 +13,40 @@ from mattport.utils.misc import is_not_none @dataclass class PointSamples: - """Samples in space.""" + """Samples in space. - positions: TensorType[..., 3] = None # XYZ locations - directions: TensorType[..., 3] = None # Unit direction vector - camera_indices: TensorType[..., 1] = None # Camera index - valid_mask: TensorType[..., 1] = None # Rays that are valid + Args: + positions (TensorType[..., 3]): XYZ locations + directions (TensorType[..., 3]): Unit direction vector + camera_indices (TensorType[..., 1]): Camera index + valid_mask (TensorType[...]): Rays that are valid + """ + + positions: TensorType[..., 3] + directions: TensorType[..., 3] = None + camera_indices: TensorType[..., 1] = None + valid_mask: TensorType[...] = None @dataclass class RaySamples: - """Samples along a ray""" + """Samples along a ray - positions: TensorType[..., 3] = None # XYZ locations - directions: TensorType[..., 3] = None # Unit direction vector - camera_indices: TensorType[..., 1] = None # Camera index - valid_mask: TensorType[..., 1] = None # Rays that are valid - ts: TensorType[..., 1] = None # "time steps", distances along ray - deltas: TensorType[..., 1] = None # "width" of each sample + Args: + positions (TensorType[..., 3]): XYZ locations + directions (TensorType[..., 3]): Unit direction vector + camera_indices (TensorType[..., 1]): Camera index + valid_mask (TensorType[...]): Rays that are valid + ts (TensorType[..., 1]): "time steps", distances along ray + deltas )TensorType[..., 1]): "width" of each sample + """ + + positions: TensorType[..., 3] + directions: TensorType[..., 3] = None + camera_indices: TensorType[..., 1] = None + valid_mask: TensorType[...] = None + ts: TensorType[..., 1] = None + deltas: TensorType[..., 1] = None def to_point_samples(self) -> PointSamples: """Convert to PointSamples instance and return.""" @@ -67,45 +83,55 @@ class RaySamples: @dataclass class RayBundle: - """A bundle of ray parameters.""" + """A bundle of ray parameters. - origins: TensorType["num_rays", 3] # Ray origins - directions: TensorType["num_rays", 3] # - camera_indices: Optional[TensorType["num_rays"]] = None + Args: + origins (TensorType[..., 3]): Ray origins (XYZ) + directions (TensorType[..., 3]): Unit ray direction vector + camera_indices (TensorType[..., 1]): Camera indices + nears (TensorType[..., 1]): Distance along ray to start sampling + fars (TensorType[..., 1]): Rays Distance along ray to stop sampling + valid_mask (TensorType[...]): Rays that are valid + """ + + origins: TensorType["num_rays", 3] + directions: TensorType["num_rays", 3] + camera_indices: Optional[TensorType["num_rays", 1]] = None nears: Optional[TensorType["num_rays"]] = None fars: Optional[TensorType["num_rays"]] = None valid_mask: Optional[TensorType["num_rays"]] = None - def to_camera_ray_bundle(self, image_height, image_width) -> "CameraRayBundle": - """Returns a CameraRayBundle from this object.""" - camera_ray_bundle = CameraRayBundle( - origins=self.origins.view(image_height, image_width, 3), - directions=self.directions.view(image_height, image_width, 3), - camera_indices=self.camera_indices.view(image_height, image_width) - if not isinstance(self.camera_indices, type(None)) - else None, - ) - return camera_ray_bundle + def move_to_device(self, device: torch.device) -> None: + """Move bundle data to a device. - def move_to_device(self, device): - """Move to a device.""" + Args: + device (torch.device): Device to move tensors to. + """ self.origins = self.origins.to(device) self.directions = self.directions.to(device) if not isinstance(self.camera_indices, type(None)): self.camera_indices = self.camera_indices.to(device) + def set_camera_indices(self, camera_index: int) -> None: + """Sets all of the the camera indices to a specific camera index. + + Args: + camera_index (int): Camera index. + """ + self.camera_indices = torch.ones_like(self.origins[..., 0]).long() * camera_index + def __len__(self): - num_rays = self.origins.shape[0] + num_rays = torch.numel(self.origins) // self.origins.shape[-1] return num_rays - def sample(self, num_rays: int): + def sample(self, num_rays: int) -> "RayBundle": """Returns a RayBundle as a subset of rays. Args: - num_rays (int): + num_rays (int): Number of rays in output RayBundle Returns: - RayBundle: _description_ + RayBundle: RayBundle with subset of rays. """ assert num_rays <= len(self) indices = random.sample(range(len(self)), k=num_rays) @@ -115,8 +141,15 @@ class RayBundle: camera_indices=self.camera_indices[indices], ) - def get_masked_ray_bundle(self, valid_mask): - """Return a masked instance of the ray bundle.""" + def get_masked_ray_bundle(self, valid_mask: TensorType) -> "RayBundle": + """Return a masked instance of the ray bundle. + + Args: + valid_mask (TensorType): Mask of rays to keep. + + Returns: + RayBundle: RayBundle of masked rays. + """ return RayBundle( origins=self.origins[valid_mask], directions=self.directions[valid_mask], @@ -126,13 +159,35 @@ class RayBundle: valid_mask=self.valid_mask[valid_mask] if is_not_none(self.valid_mask) else None, ) - def get_ray_samples(self, ts: TensorType["num_rays", "num_samples"]) -> RaySamples: - """ + def get_row_major_sliced_ray_bundle(self, start_idx: int, end_idx: int) -> "RayBundle": + """Flattens RayBundle and extracts chunk given start and end indicies. + Args: - ts (TensorType["num_rays", "num_samples"]): _description_ + start_idx (int): Start index of RayBundle chunk. + end_idx (int): End index of RayBundle chunk. Returns: - RaySamples: _description_ + RayBundle: Flattened RayBundle with end_idx-start_idx rays. + + """ + camera_indices = ( + self.camera_indices.view(-1)[start_idx:end_idx] if not isinstance(self.camera_indices, type(None)) else None + ) + return RayBundle( + origins=self.origins.view(-1, 3)[start_idx:end_idx], + directions=self.directions.view(-1, 3)[start_idx:end_idx], + camera_indices=camera_indices, + ) + + def get_ray_samples(self, ts: TensorType[..., "num_samples"]) -> RaySamples: + """Produces samples for each ray by projection points along the ray direction. + + + Args: + ts (TensorType[..., "num_samples"]): Distance from origin of sample points. + + Returns: + RaySamples: Samples projected along ray. """ positions = self.origins[:, None] + ts[:, :, None] * self.directions[:, None] directions = self.directions.unsqueeze(1).repeat(1, positions.shape[1], 1) @@ -157,49 +212,3 @@ class RayBundle: ) return ray_samples - - -@dataclass -class CameraRayBundle: - """_summary_""" - - origins: TensorType["image_height", "image_width", 3] - directions: TensorType["image_height", "image_width", 3] - camera_indices: Optional[TensorType["image_height", "image_width", 3]] = None - camera_index: int = None - - def __post_init__(self): - if not isinstance(self.camera_index, type(None)): - self.set_camera_indices(self.camera_index) - - def set_camera_indices(self, camera_index: int): - """Sets the camera indices for a specific camera index.""" - self.camera_index = camera_index - self.camera_indices = torch.ones_like(self.origins[..., 0]).long() * camera_index - - def get_num_rays(self): - """Return the number of rays in this bundle.""" - image_height, image_width = self.origins.shape[:2] - num_rays = image_height * image_width - return num_rays - - def to_ray_bundle(self) -> RayBundle: - """_summary_ - - Returns: - RayBundle: _description_ - """ - # TODO(ethan): handle camera_index - ray_bundle = RayBundle(origins=self.origins.view(-1, 3), directions=self.directions.view(-1, 3)) - return ray_bundle - - def get_row_major_sliced_ray_bundle(self, start_idx, end_idx): - """Return a RayBundle""" - camera_indices = ( - self.camera_indices.view(-1)[start_idx:end_idx] if not isinstance(self.camera_indices, type(None)) else None - ) - return RayBundle( - origins=self.origins.view(-1, 3)[start_idx:end_idx], - directions=self.directions.view(-1, 3)[start_idx:end_idx], - camera_indices=camera_indices, - ) diff --git a/tests/structures/test_cameras.py b/tests/structures/test_cameras.py index e0d841a6b..936116806 100644 --- a/tests/structures/test_cameras.py +++ b/tests/structures/test_cameras.py @@ -5,12 +5,12 @@ Test the camera classes. import torch from mattport.structures.cameras import PinholeCamera -from mattport.structures.rays import CameraRayBundle +from mattport.structures.rays import RayBundle def test_pinhole_camera(): """Test that the pinhole camera model works.""" pinhole_camera = PinholeCamera(cx=400.0, cy=400.0, fx=10.0, fy=10.0) camera_ray_bundle = pinhole_camera.generate_camera_rays() - assert isinstance(camera_ray_bundle, CameraRayBundle) + assert isinstance(camera_ray_bundle, RayBundle) assert torch.allclose(camera_ray_bundle.origins[0], torch.tensor([0.0, 0.0, 0.0]))