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) @staticmethod 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)