Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn.functional as F | |
| from einops import rearrange, repeat | |
| from torch import nn | |
| from ....geometry.projection import ( | |
| homogenize_points, | |
| project, | |
| sample_image_grid, | |
| transform_cam2world, | |
| transform_world2cam, | |
| unproject, | |
| ) | |
| class DepthGuidedMVRefiner(nn.Module): | |
| def __init__( | |
| self, | |
| feature_channels: int, | |
| hidden_channels: int | None = None, | |
| concat_ref_depth: bool = False, | |
| ) -> None: | |
| super().__init__() | |
| hidden_channels = feature_channels if hidden_channels is None else hidden_channels | |
| depth_channels = max(hidden_channels // 4, 16) | |
| self.concat_ref_depth = concat_ref_depth | |
| self.depth_encoder = nn.Sequential( | |
| nn.Conv2d( | |
| 1, | |
| depth_channels, | |
| 1, | |
| 1, | |
| 0, | |
| ), | |
| nn.GELU(), | |
| nn.Conv2d( | |
| depth_channels, | |
| depth_channels, | |
| 1, | |
| 1, | |
| 0, | |
| ), | |
| nn.GELU(), | |
| ) | |
| self.feature_encoder = nn.Sequential( | |
| nn.Conv2d( | |
| feature_channels * 2 | |
| + depth_channels * (2 if self.concat_ref_depth else 1), | |
| hidden_channels, | |
| 3, | |
| 1, | |
| 1, | |
| padding_mode="replicate", | |
| ), | |
| nn.GELU(), | |
| nn.Conv2d( | |
| hidden_channels, | |
| hidden_channels, | |
| 3, | |
| 1, | |
| 1, | |
| padding_mode="replicate", | |
| ), | |
| nn.GELU(), | |
| ) | |
| self.post_fuse = nn.Sequential( | |
| nn.Conv2d( | |
| feature_channels + hidden_channels + depth_channels * int(self.concat_ref_depth), | |
| hidden_channels, | |
| 3, | |
| 1, | |
| 1, | |
| padding_mode="replicate", | |
| ), | |
| nn.GELU(), | |
| nn.Conv2d( | |
| hidden_channels, | |
| feature_channels, | |
| 3, | |
| 1, | |
| 1, | |
| padding_mode="replicate", | |
| ), | |
| ) | |
| nn.init.zeros_(self.post_fuse[-1].weight) | |
| nn.init.zeros_(self.post_fuse[-1].bias) | |
| def _normalized_coords_to_grid( | |
| xy: torch.Tensor, | |
| h: int, | |
| w: int, | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| x_pixel = xy[..., 0] * w - 0.5 | |
| y_pixel = xy[..., 1] * h - 0.5 | |
| if w > 1: | |
| x_grid = 2 * x_pixel / (w - 1) - 1 | |
| else: | |
| x_grid = torch.zeros_like(x_pixel) | |
| if h > 1: | |
| y_grid = 2 * y_pixel / (h - 1) - 1 | |
| else: | |
| y_grid = torch.zeros_like(y_pixel) | |
| grid = torch.stack((x_grid, y_grid), dim=-1) | |
| in_bounds = ( | |
| (x_pixel >= 0) | |
| & (x_pixel <= (w - 1)) | |
| & (y_pixel >= 0) | |
| & (y_pixel <= (h - 1)) | |
| ) | |
| return grid, in_bounds | |
| def forward( | |
| self, | |
| features: torch.Tensor, | |
| depth: torch.Tensor, | |
| intrinsics: torch.Tensor, | |
| extrinsics: torch.Tensor, | |
| source_view_indices: torch.Tensor | None = None, | |
| ) -> torch.Tensor: | |
| b, v, c, h, w = features.shape | |
| if v <= 1: | |
| return features | |
| depth = depth.detach() | |
| intrinsics = intrinsics.detach() | |
| extrinsics = extrinsics.detach() | |
| feature_dtype = features.dtype | |
| xy_ref, _ = sample_image_grid((h, w), device=features.device) | |
| xy_ref = repeat(xy_ref, "h w xy -> b h w xy", b=b) | |
| refined_features = [] | |
| geom_eps = torch.finfo(torch.float32).eps | |
| count_eps = torch.finfo(features.dtype).eps | |
| batch_indices = torch.arange(b, device=features.device) | |
| if source_view_indices is not None: | |
| if source_view_indices.shape[:2] != (b, v): | |
| raise ValueError( | |
| "source_view_indices must have shape [B, V, M], got " | |
| f"{tuple(source_view_indices.shape)} for features " | |
| f"{tuple(features.shape)}" | |
| ) | |
| source_view_indices = source_view_indices.to( | |
| device=features.device, | |
| dtype=torch.long, | |
| ) | |
| for ref_idx in range(v): | |
| ref_features = features[:, ref_idx] | |
| ref_depth = depth[:, ref_idx].float() | |
| ref_intrinsics = intrinsics[:, ref_idx].float() | |
| ref_extrinsics = extrinsics[:, ref_idx].float() | |
| ref_depth_repr = None | |
| if self.concat_ref_depth: | |
| ref_depth_input = ref_depth.unsqueeze(1).clamp_min(geom_eps).reciprocal() | |
| ref_depth_repr = self.depth_encoder(ref_depth_input.to(feature_dtype)) | |
| with torch.no_grad(), torch.autocast( | |
| device_type=features.device.type, enabled=False | |
| ): | |
| ref_points_cam = unproject( | |
| xy_ref, | |
| ref_depth, | |
| ref_intrinsics[:, None, None], | |
| ) | |
| ref_points_world = transform_cam2world( | |
| homogenize_points(ref_points_cam), | |
| ref_extrinsics[:, None, None], | |
| )[..., :-1] | |
| fused_sum = None | |
| valid_count = None | |
| if source_view_indices is None: | |
| selected_src_indices = [ | |
| src_idx for src_idx in range(v) if src_idx != ref_idx | |
| ] | |
| else: | |
| selected_src_indices = source_view_indices[:, ref_idx].unbind(dim=-1) | |
| for src_idx in selected_src_indices: | |
| if isinstance(src_idx, int): | |
| src_features = features[:, src_idx] | |
| src_extrinsics = extrinsics[:, src_idx].float()[:, None, None] | |
| src_intrinsics = intrinsics[:, src_idx].float()[:, None, None] | |
| else: | |
| src_features = features[batch_indices, src_idx] | |
| src_extrinsics = ( | |
| extrinsics[batch_indices, src_idx].float()[:, None, None] | |
| ) | |
| src_intrinsics = ( | |
| intrinsics[batch_indices, src_idx].float()[:, None, None] | |
| ) | |
| with torch.no_grad(), torch.autocast( | |
| device_type=features.device.type, enabled=False | |
| ): | |
| src_xy, src_in_front = project( | |
| ref_points_world, | |
| src_extrinsics, | |
| src_intrinsics, | |
| ) | |
| src_grid, src_in_bounds = self._normalized_coords_to_grid( | |
| src_xy, h, w | |
| ) | |
| valid_mask_fp32 = (src_in_front & src_in_bounds).unsqueeze(1).float() | |
| with torch.autocast(device_type=features.device.type, enabled=False): | |
| warped_features = F.grid_sample( | |
| src_features.float(), | |
| src_grid, | |
| mode="bilinear", | |
| padding_mode="zeros", | |
| align_corners=True, | |
| ) | |
| valid_mask = valid_mask_fp32.to(feature_dtype) | |
| warped_features = warped_features.to(feature_dtype) | |
| warped_features = warped_features * valid_mask | |
| with torch.no_grad(), torch.autocast( | |
| device_type=features.device.type, enabled=False | |
| ): | |
| src_points_cam = transform_world2cam( | |
| homogenize_points(ref_points_world), | |
| src_extrinsics, | |
| )[..., :-1] | |
| warped_depth = rearrange( | |
| src_points_cam[..., -1], "b h w -> b 1 h w" | |
| ) | |
| warped_depth_input = torch.where( | |
| valid_mask_fp32 > 0, | |
| warped_depth.clamp_min(geom_eps).reciprocal(), | |
| torch.zeros_like(warped_depth), | |
| ) | |
| depth_repr = self.depth_encoder(warped_depth_input.to(feature_dtype)) | |
| feature_inputs = [ref_features, warped_features] | |
| if ref_depth_repr is not None: | |
| feature_inputs.append(ref_depth_repr) | |
| feature_inputs.append(depth_repr) | |
| value = self.feature_encoder(torch.cat(feature_inputs, dim=1)) * valid_mask | |
| if fused_sum is None: | |
| fused_sum = value | |
| valid_count = valid_mask | |
| else: | |
| fused_sum = fused_sum + value | |
| valid_count = valid_count + valid_mask | |
| if fused_sum is None or valid_count is None: | |
| refined_features.append(ref_features) | |
| continue | |
| fused = fused_sum / valid_count.clamp_min(count_eps) | |
| post_fuse_inputs = [ref_features, fused] | |
| if ref_depth_repr is not None: | |
| post_fuse_inputs.append(ref_depth_repr) | |
| residual = self.post_fuse(torch.cat(post_fuse_inputs, dim=1)) | |
| refined_features.append(ref_features + residual) | |
| return torch.stack(refined_features, dim=1) | |