Spaces:
Sleeping
Sleeping
| from typing import List, Optional, Tuple | |
| import numpy as np | |
| import torch | |
| from torch.nn import functional as F | |
| def get_position_map_from_depth(depth, mask, intrinsics, extrinsics, image_wh=None): | |
| """Compute the position map from the depth map and the camera parameters for a batch of views. | |
| Args: | |
| depth (torch.Tensor): The depth maps with the shape (B, H, W, 1). | |
| mask (torch.Tensor): The masks with the shape (B, H, W, 1). | |
| intrinsics (torch.Tensor): The camera intrinsics matrices with the shape (B, 3, 3). | |
| extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4). | |
| image_wh (Tuple[int, int]): The image width and height. | |
| Returns: | |
| torch.Tensor: The position maps with the shape (B, H, W, 3). | |
| """ | |
| if image_wh is None: | |
| image_wh = depth.shape[2], depth.shape[1] | |
| B, H, W, _ = depth.shape | |
| depth = depth.squeeze(-1) | |
| u_coord, v_coord = torch.meshgrid( | |
| torch.arange(image_wh[0]), torch.arange(image_wh[1]), indexing="xy" | |
| ) | |
| u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1) | |
| v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1) | |
| # Compute the position map by back-projecting depth pixels to 3D space | |
| x = ( | |
| (u_coord - intrinsics[:, 0, 2].unsqueeze(-1).unsqueeze(-1)) | |
| * depth | |
| / intrinsics[:, 0, 0].unsqueeze(-1).unsqueeze(-1) | |
| ) | |
| y = ( | |
| (v_coord - intrinsics[:, 1, 2].unsqueeze(-1).unsqueeze(-1)) | |
| * depth | |
| / intrinsics[:, 1, 1].unsqueeze(-1).unsqueeze(-1) | |
| ) | |
| z = depth | |
| # Concatenate to form the 3D coordinates in the camera frame | |
| camera_coords = torch.stack([x, y, z], dim=-1) | |
| # Apply the extrinsic matrix to get coordinates in the world frame | |
| coords_homogeneous = torch.nn.functional.pad( | |
| camera_coords, (0, 1), "constant", 1.0 | |
| ) # Add a homogeneous coordinate | |
| world_coords = torch.matmul( | |
| coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2) | |
| ).view(B, H, W, 4) | |
| # Apply the mask to the position map | |
| position_map = world_coords[..., :3] * mask | |
| return position_map | |
| def get_position_map_from_depth_ortho( | |
| depth, mask, extrinsics, ortho_scale, image_wh=None | |
| ): | |
| """Compute the position map from the depth map and the camera parameters for a batch of views | |
| using orthographic projection with a given ortho_scale. | |
| Args: | |
| depth (torch.Tensor): The depth maps with the shape (B, H, W, 1). | |
| mask (torch.Tensor): The masks with the shape (B, H, W, 1). | |
| extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4). | |
| ortho_scale (torch.Tensor): The scaling factor for the orthographic projection with the shape (B, 1, 1, 1). | |
| image_wh (Tuple[int, int]): Optional. The image width and height. | |
| Returns: | |
| torch.Tensor: The position maps with the shape (B, H, W, 3). | |
| """ | |
| if image_wh is None: | |
| image_wh = depth.shape[2], depth.shape[1] | |
| B, H, W, _ = depth.shape | |
| depth = depth.squeeze(-1) | |
| # Generating grid of coordinates in the image space | |
| u_coord, v_coord = torch.meshgrid( | |
| torch.arange(0, image_wh[0]), torch.arange(0, image_wh[1]), indexing="xy" | |
| ) | |
| u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1) | |
| v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1) | |
| # Compute the position map using orthographic projection with ortho_scale | |
| x = (u_coord - image_wh[0] / 2) / ortho_scale / image_wh[0] | |
| y = (v_coord - image_wh[1] / 2) / ortho_scale / image_wh[1] | |
| z = depth | |
| # Concatenate to form the 3D coordinates in the camera frame | |
| camera_coords = torch.stack([x, y, z], dim=-1) | |
| # Apply the extrinsic matrix to get coordinates in the world frame | |
| coords_homogeneous = torch.nn.functional.pad( | |
| camera_coords, (0, 1), "constant", 1.0 | |
| ) # Add a homogeneous coordinate | |
| world_coords = torch.matmul( | |
| coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2) | |
| ).view(B, H, W, 4) | |
| # Apply the mask to the position map | |
| position_map = world_coords[..., :3] * mask | |
| return position_map | |
| def get_opencv_from_blender(matrix_world, fov=None, image_size=None): | |
| # convert matrix_world to opencv format extrinsics | |
| opencv_world_to_cam = matrix_world.inverse() | |
| opencv_world_to_cam[1, :] *= -1 | |
| opencv_world_to_cam[2, :] *= -1 | |
| R, T = opencv_world_to_cam[:3, :3], opencv_world_to_cam[:3, 3] | |
| if fov is None: # orthographic camera | |
| return R, T | |
| R, T = R.unsqueeze(0), T.unsqueeze(0) | |
| # convert fov to opencv format intrinsics | |
| focal = 1 / np.tan(fov / 2) | |
| intrinsics = np.diag(np.array([focal, focal, 1])).astype(np.float32) | |
| opencv_cam_matrix = ( | |
| torch.from_numpy(intrinsics).unsqueeze(0).float().to(matrix_world.device) | |
| ) | |
| opencv_cam_matrix[:, :2, -1] += torch.tensor([image_size / 2, image_size / 2]).to( | |
| matrix_world.device | |
| ) | |
| opencv_cam_matrix[:, [0, 1], [0, 1]] *= image_size / 2 | |
| return R, T, opencv_cam_matrix | |
| def get_ray_directions( | |
| H: int, | |
| W: int, | |
| focal: float, | |
| principal: Optional[Tuple[float, float]] = None, | |
| use_pixel_centers: bool = True, | |
| ) -> torch.Tensor: | |
| """ | |
| Get ray directions for all pixels in camera coordinate. | |
| Args: | |
| H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers | |
| Outputs: | |
| directions: (H, W, 3), the direction of the rays in camera coordinate | |
| """ | |
| pixel_center = 0.5 if use_pixel_centers else 0 | |
| cx, cy = W / 2, H / 2 if principal is None else principal | |
| i, j = torch.meshgrid( | |
| torch.arange(W, dtype=torch.float32) + pixel_center, | |
| torch.arange(H, dtype=torch.float32) + pixel_center, | |
| indexing="xy", | |
| ) | |
| directions = torch.stack( | |
| [(i - cx) / focal, -(j - cy) / focal, -torch.ones_like(i)], -1 | |
| ) | |
| return F.normalize(directions, dim=-1) | |
| def get_rays( | |
| directions: torch.Tensor, c2w: torch.Tensor | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Get ray origins and directions from camera coordinates to world coordinates | |
| Args: | |
| directions: (H, W, 3) ray directions in camera coordinates | |
| c2w: (4, 4) camera-to-world transformation matrix | |
| Outputs: | |
| rays_o, rays_d: (H, W, 3) ray origins and directions in world coordinates | |
| """ | |
| # Rotate ray directions from camera coordinate to the world coordinate | |
| rays_d = directions @ c2w[:3, :3].T | |
| rays_o = c2w[:3, 3].expand(rays_d.shape) | |
| return rays_o, rays_d | |
| def compute_plucker_embed( | |
| c2w: torch.Tensor, image_width: int, image_height: int, focal: float | |
| ) -> torch.Tensor: | |
| """ | |
| Computes Plucker coordinates for a camera. | |
| Args: | |
| c2w: (4, 4) camera-to-world transformation matrix | |
| image_width: Image width | |
| image_height: Image height | |
| focal: Focal length of the camera | |
| Returns: | |
| plucker: (6, H, W) Plucker embedding | |
| """ | |
| directions = get_ray_directions(image_height, image_width, focal) | |
| rays_o, rays_d = get_rays(directions, c2w) | |
| # Cross product to get Plucker coordinates | |
| cross = torch.cross(rays_o, rays_d, dim=-1) | |
| plucker = torch.cat((rays_d, cross), dim=-1) | |
| return plucker.permute(2, 0, 1) | |
| def get_plucker_embeds_from_cameras( | |
| c2w: List[torch.Tensor], fov: List[float], image_size: int | |
| ) -> torch.Tensor: | |
| """ | |
| Given lists of camera transformations and fov, returns the batched plucker embeddings. | |
| Args: | |
| c2w: list of camera-to-world transformation matrices | |
| fov: list of field of view values | |
| image_size: size of the image | |
| Returns: | |
| plucker_embeds: (B, 6, H, W) batched plucker embeddings | |
| """ | |
| plucker_embeds = [] | |
| for cam_matrix, cam_fov in zip(c2w, fov): | |
| focal = 0.5 * image_size / np.tan(0.5 * cam_fov) | |
| plucker = compute_plucker_embed(cam_matrix, image_size, image_size, focal) | |
| plucker_embeds.append(plucker) | |
| return torch.stack(plucker_embeds) | |
| def get_plucker_embeds_from_cameras_ortho( | |
| c2w: List[torch.Tensor], ortho_scale: List[float], image_size: int | |
| ): | |
| """ | |
| Given lists of camera transformations and fov, returns the batched plucker embeddings. | |
| Parameters: | |
| c2w: list of camera-to-world transformation matrices | |
| fov: list of field of view values | |
| image_size: size of the image | |
| Returns: | |
| plucker_embeds: plucker embeddings (B, 6, H, W) | |
| """ | |
| plucker_embeds = [] | |
| # compute pairwise mask and plucker embeddings | |
| for cam_matrix, scale in zip(c2w, ortho_scale): | |
| # blender to opencv to pytorch3d | |
| R, T = get_opencv_from_blender(cam_matrix) | |
| cam_pos = -R.T @ T | |
| view_dir = R.T @ torch.tensor([0, 0, 1]).float().to(cam_matrix.device) | |
| # normalize camera position | |
| cam_pos = F.normalize(cam_pos, dim=0) | |
| plucker = torch.concat([view_dir, cam_pos]) | |
| plucker = plucker.unsqueeze(-1).unsqueeze(-1).repeat(1, image_size, image_size) | |
| plucker_embeds.append(plucker) | |
| plucker_embeds = torch.stack(plucker_embeds) | |
| return plucker_embeds | |