Spaces:
Running on Zero
Running on Zero
| # Modified for SplatAnything: inference-only release. | |
| import torch | |
| import numpy as np | |
| from typing import Tuple | |
| def depth_to_camera_coords(depthmap, camera_intrinsics): | |
| """ | |
| Convert depth map to 3D camera coordinates. | |
| Args: | |
| depthmap (BxHxW tensor): Batch of depth maps | |
| camera_intrinsics (Bx3x3 tensor): Camera intrinsics matrix for each camera | |
| Returns: | |
| X_cam (BxHxWx3 tensor): 3D points in camera coordinates | |
| valid_mask (BxHxW tensor): Mask indicating valid depth pixels | |
| """ | |
| (B, H, W) = depthmap.shape | |
| device = depthmap.device | |
| dtype = depthmap.dtype | |
| camera_intrinsics = camera_intrinsics.float() | |
| fx = camera_intrinsics[:, 0, 0] | |
| fy = camera_intrinsics[:, 1, 1] | |
| cx = camera_intrinsics[:, 0, 2] | |
| cy = camera_intrinsics[:, 1, 2] | |
| (v_grid, u_grid) = torch.meshgrid(torch.arange(H, dtype=dtype, device=device), torch.arange(W, dtype=dtype, device=device), indexing='ij') | |
| u_grid = u_grid.unsqueeze(0) | |
| v_grid = v_grid.unsqueeze(0) | |
| z_cam = depthmap | |
| x_cam = (u_grid - cx.view(B, 1, 1)) * z_cam / fx.view(B, 1, 1) | |
| y_cam = (v_grid - cy.view(B, 1, 1)) * z_cam / fy.view(B, 1, 1) | |
| X_cam = torch.stack([x_cam, y_cam, z_cam], dim=-1) | |
| valid_mask = depthmap > 0.0 | |
| return (X_cam, valid_mask) | |
| def depth_to_world_coords_points(depth_map: torch.Tensor, extrinsic: torch.Tensor, intrinsic: torch.Tensor, eps=1e-08) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """ | |
| Convert a batch of depth maps to world coordinates. | |
| Args: | |
| depth_map (torch.Tensor): (B, H, W) Depth map | |
| extrinsic (torch.Tensor): (B, 4, 4) Camera extrinsic matrix (camera-to-world transformation) | |
| intrinsic (torch.Tensor): (B, 3, 3) Camera intrinsic matrix | |
| Returns: | |
| world_coords_points (torch.Tensor): (B, H, W, 3) World coordinates | |
| camera_points (torch.Tensor): (B, H, W, 3) Camera coordinates | |
| point_mask (torch.Tensor): (B, H, W) Valid depth mask | |
| """ | |
| if depth_map is None: | |
| return (None, None, None) | |
| point_mask = depth_map > eps | |
| (camera_points, _) = depth_to_camera_coords(depth_map, intrinsic) | |
| R_cam_to_world = extrinsic[:, :3, :3] | |
| t_cam_to_world = extrinsic[:, :3, 3] | |
| world_coords_points = torch.einsum('bhwi,bji->bhwj', camera_points, R_cam_to_world) + t_cam_to_world[:, None, None, :] | |
| return (world_coords_points, camera_points, point_mask) | |