Spaces:
Running on Zero
Running on Zero
File size: 3,477 Bytes
7405e85 | 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 | import torch
import torch.nn.functional as F
def coords_grid(b, h, w, homogeneous=False, device=None):
y, x = torch.meshgrid(
torch.arange(h), torch.arange(w), indexing="ij"
) # [H, W]
stacks = [x, y]
if homogeneous:
ones = torch.ones_like(x) # [H, W]
stacks.append(ones)
grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W]
grid = grid[None].expand(b, -1, -1, -1) # [B, 2, H, W] or [B, 3, H, W]
if device is not None:
grid = grid.to(device)
return grid
def warp_with_pose_depth_candidates(
feature1,
ref_intrinsics,
tgt_intrinsics,
pose,
depth,
clamp_min_depth=1e-3,
grid_sample_disable_cudnn=False,
):
"""
feature1: [B, C, H, W]
ref_intrinsics: [B, 3, 3], used to backproject reference pixels
tgt_intrinsics: [B, 3, 3], used to reproject into target-view pixels
pose: [B, 4, 4]
depth: [B, D, H, W]
"""
assert ref_intrinsics.size(1) == ref_intrinsics.size(2) == 3
assert tgt_intrinsics.size(1) == tgt_intrinsics.size(2) == 3
assert pose.size(1) == pose.size(2) == 4
assert depth.dim() == 4
b, d, h, w = depth.size()
c = feature1.size(1)
feature_dtype = feature1.dtype
with torch.no_grad(), torch.autocast(
device_type=depth.device.type, enabled=False
):
ref_intrinsics = ref_intrinsics.float()
tgt_intrinsics = tgt_intrinsics.float()
pose = pose.float()
depth = depth.float()
# pixel coordinates
grid = coords_grid(
b, h, w, homogeneous=True, device=depth.device
) # [B, 3, H, W]
# back project to 3D and transform viewpoint
points = torch.linalg.inv(ref_intrinsics).bmm(
grid.reshape(b, 3, -1)
) # [B, 3, H*W]
points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2) * depth.reshape(
b, 1, d, h * w
) # [B, 3, D, H*W]
points = points + pose[:, :3, -1:].unsqueeze(-1) # [B, 3, D, H*W]
# reproject to 2D image plane
points = torch.bmm(tgt_intrinsics, points.view(b, 3, -1)).view(
b, 3, d, h * w
) # [B, 3, D, H*W]
pixel_coords = points[:, :2] / points[:, -1:].clamp(
min=clamp_min_depth
) # [B, 2, D, H*W]
# normalize to [-1, 1]
x_grid = 2 * pixel_coords[:, 0] / (w - 1) - 1
y_grid = 2 * pixel_coords[:, 1] / (h - 1) - 1
grid = torch.stack([x_grid, y_grid], dim=-1) # [B, D, H*W, 2]
# Disable cuDNN for large grid_sample inputs; this avoids a known PyTorch
# issue while keeping the default path for smaller inputs.
# Ref: https://github.com/pytorch/pytorch/issues/88380
if feature1.numel() > 1000000:
grid_sample_disable_cudnn = True
with torch.backends.cudnn.flags(enabled=not grid_sample_disable_cudnn):
# grid_sample requires the feature and grid dtypes to match. Compute the
# warp in FP32 for sub-pixel grid accuracy, then return to the feature dtype.
with torch.autocast(device_type=feature1.device.type, enabled=False):
warped_feature = F.grid_sample(
feature1.float(),
grid.view(b, d * h, w, 2),
mode="bilinear",
padding_mode="zeros",
align_corners=True,
).view(
b, c, d, h, w
) # [B, C, D, H, W]
return warped_feature.to(feature_dtype)
|