Spaces:
Running on Zero
Running on Zero
| 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) | |