Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,744 Bytes
9b33fca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
"""Projection utilities."""
from __future__ import annotations
import torch
from .transform import inverse_pinhole
def project_points(
points: torch.Tensor, intrinsics: torch.Tensor
) -> torch.Tensor:
"""Project points to pixel coordinates with given intrinsics.
Args:
points: (N, 3) or (B, N, 3) 3D coordinates.
intrinsics: (3, 3) or (B, 3, 3) intrinsic camera matrices.
Returns:
torch.Tensor: (N, 2) or (B, N, 2) 2D pixel coordinates.
Raises:
ValueError: Shape of input points is not valid for computation.
"""
assert points.shape[-1] == 3, "Input coordinates must be 3 dimensional!"
hom_coords = points / points[..., 2:3]
if len(hom_coords.shape) == 2:
assert (
len(intrinsics.shape) == 2
), "Got multiple intrinsics for single point set!"
intrinsics = intrinsics.T
elif len(hom_coords.shape) == 3:
if len(intrinsics.shape) == 2:
intrinsics = intrinsics.unsqueeze(0)
intrinsics = intrinsics.permute(0, 2, 1)
else:
raise ValueError(f"Shape of input points not valid: {points.shape}")
pts_2d = hom_coords @ intrinsics
return pts_2d[..., :2]
def unproject_points(
points: torch.Tensor, depths: torch.Tensor, intrinsics: torch.Tensor
) -> torch.Tensor:
"""Un-projects pixel coordinates to 3D coordinates with given intrinsics.
Args:
points: (N, 2) or (B, N, 2) 2D pixel coordinates.
depths: (N,) / (N, 1) or (B, N,) / (B, N, 1) depth values.
intrinsics: (3, 3) or (B, 3, 3) intrinsic camera matrices.
Returns:
torch.Tensor: (N, 3) or (B, N, 3) 3D coordinates.
Raises:
ValueError: Shape of input points is not valid for computation.
"""
if len(points.shape) == 2:
assert (
len(intrinsics.shape) == 2 or intrinsics.shape[0] == 1
), "Got multiple intrinsics for single point set!"
if len(intrinsics.shape) == 3:
intrinsics = intrinsics.squeeze(0)
inv_intrinsics = inverse_pinhole(intrinsics).transpose(0, 1)
if len(depths.shape) == 1:
depths = depths.unsqueeze(-1)
assert len(depths.shape) == 2, "depths must have same dims as points"
elif len(points.shape) == 3:
inv_intrinsics = inverse_pinhole(intrinsics).transpose(-2, -1)
if len(depths.shape) == 2:
depths = depths.unsqueeze(-1)
assert len(depths.shape) == 3, "depths must have same dims as points"
else:
raise ValueError(f"Shape of input points not valid: {points.shape}")
hom_coords = torch.cat([points, torch.ones_like(points)[..., 0:1]], -1)
pts_3d = hom_coords @ inv_intrinsics
pts_3d *= depths
return pts_3d
def points_inside_image(
points_coord: torch.Tensor,
depths: torch.Tensor,
images_hw: torch.Tensor | tuple[int, int],
) -> torch.Tensor:
"""Generate binary mask.
Creates a mask that is true for all point coordiantes that lie inside the
image,
Args:
points_coord (torch.Tensor): 2D pixel coordinates of shape [..., 2].
depths (torch.Tensor): Associated depth of each 2D pixel coordinate.
images_hw: (torch.Tensor| tuple[int, int]]) Associated tensor of image
dimensions, shape [..., 2] or single height, width pair.
Returns:
torch.Tensor: Binary mask of points inside an image.
"""
mask = torch.ones_like(depths)
h: int | torch.Tensor
w: int | torch.Tensor
if isinstance(images_hw, tuple):
h, w = images_hw
else:
h, w = images_hw[..., 0], images_hw[..., 1]
mask = torch.logical_and(mask, torch.greater(depths, 0))
mask = torch.logical_and(mask, points_coord[..., 0] > 0)
mask = torch.logical_and(mask, points_coord[..., 0] < w - 1)
mask = torch.logical_and(mask, points_coord[..., 1] > 0)
mask = torch.logical_and(mask, points_coord[..., 1] < h - 1)
return mask
def generate_depth_map(
points: torch.Tensor,
intrinsics: torch.Tensor,
image_hw: tuple[int, int],
) -> torch.Tensor:
"""Generate depth map for given pointcloud.
Args:
points: (N, 3) coordinates.
intrinsics: (3, 3) intrinsic camera matrices.
image_hw: (tuple[int,int]) height, width of the image
Returns:
torch.Tensor: Projected depth map of the given pointcloud.
Invalid depth has 0 values
"""
pts_2d = project_points(points, intrinsics).round()
depths = points[:, 2]
depth_map = points.new_zeros(image_hw)
mask = points_inside_image(pts_2d, depths, image_hw)
pts_2d = pts_2d[mask].long()
depth_map[pts_2d[:, 1], pts_2d[:, 0]] = depths[mask]
return depth_map
|