import torch import torch.nn as nn import torch.nn.functional as F from .backbone import CNNEncoder from .vit_fpn import ViTFeaturePyramid from .mv_transformer import ( MultiViewFeatureTransformer, batch_features_camera_parameters, ) from .matching import warp_with_pose_depth_candidates from .utils import mv_feature_add_position from .dpt_head import DPTHead from .ldm_unet.unet import UNetModel, AttentionBlock from .ldm_unet.util import normalization from .dinov2.dinov2 import DINOv2 def flatten_views(tensor): return tensor.flatten(0, 1) def unflatten_views(tensor, batch_size, num_views): return tensor.reshape(batch_size, num_views, *tensor.shape[1:]) class MultiViewUniMatch(nn.Module): def __init__( self, num_scales=1, feature_channels=128, upsample_factor=8, lowest_feature_resolution=8, num_head=1, ffn_dim_expansion=4, num_transformer_layers=6, num_depth_candidates=128, vit_type="vits", unet_channels=128, unet_channel_mult=[1, 1, 1], unet_num_res_blocks=1, unet_attn_resolutions=[4], grid_sample_disable_cudnn=False, **kwargs, ): super(MultiViewUniMatch, self).__init__() # CNN self.feature_channels = feature_channels self.num_scales = num_scales self.lowest_feature_resolution = lowest_feature_resolution self.upsample_factor = upsample_factor # monocular backbones: final self.vit_type = vit_type # cost volume self.num_depth_candidates = num_depth_candidates # upsampler vit_feature_channel_dict = {"vits": 384, "vitb": 768, "vitl": 1024} vit_feature_channel = vit_feature_channel_dict[vit_type] # CNN self.backbone = CNNEncoder( output_dim=feature_channels, num_output_scales=num_scales, downsample_factor=upsample_factor, lowest_scale=lowest_feature_resolution, return_all_scales=True, ) # Transformer self.transformer = MultiViewFeatureTransformer( num_layers=num_transformer_layers, d_model=feature_channels, nhead=num_head, ffn_dim_expansion=ffn_dim_expansion, ) if self.num_scales > 1: # generate multi-scale features self.mv_pyramid = ViTFeaturePyramid( in_channels=128, scale_factors=[2**i for i in range(self.num_scales)] ) # monodepth encoder = vit_type # can also be 'vitb' or 'vitl' self.pretrained = DINOv2( encoder, use_checkpointing=kwargs.get("use_checkpointing", False), ) del self.pretrained.mask_token # unused if self.num_scales > 1: # generate multi-scale features self.mono_pyramid = ViTFeaturePyramid( in_channels=vit_feature_channel, scale_factors=[2**i for i in range(self.num_scales)], ) # UNet regressor self.regressor = nn.ModuleList() self.regressor_residual = nn.ModuleList() self.depth_head = nn.ModuleList() for i in range(self.num_scales): curr_depth_candidates = num_depth_candidates // (4**i) cnn_feature_channels = 128 - (32 * i) mv_transformer_feature_channels = 128 // (2**i) mono_feature_channels = vit_feature_channel // (2**i) # concat(cost volume, cnn feature, mv feature, mono feature) in_channels = ( curr_depth_candidates + cnn_feature_channels + mv_transformer_feature_channels + mono_feature_channels ) # unet channels channels = unet_channels // (2**i) # unet channel mult & unet_attn_resolutions if i > 0: unet_channel_mult = unet_channel_mult + [1] unet_attn_resolutions = [x * 2 for x in unet_attn_resolutions] # unet modules = [ nn.Conv2d(in_channels, channels, 3, 1, 1), normalization(channels), nn.GELU(), ] modules.append( UNetModel( image_size=None, in_channels=channels, model_channels=channels, out_channels=channels, num_res_blocks=unet_num_res_blocks, attention_resolutions=unet_attn_resolutions, channel_mult=unet_channel_mult, num_head_channels=32, dims=2, postnorm=False, num_frames=2, use_cross_view_self_attn=True, ) ) modules.append(nn.Conv2d(channels, channels, 3, 1, 1)) self.regressor.append(nn.Sequential(*modules)) # regressor residual self.regressor_residual.append(nn.Conv2d(in_channels, channels, 1)) # depth head self.depth_head.append( nn.Sequential( nn.Conv2d( channels, channels * 2, 3, 1, 1, padding_mode="replicate" ), nn.GELU(), nn.Conv2d( channels * 2, curr_depth_candidates, 3, 1, 1, padding_mode="replicate", ), ) ) # upsampler # concat(lowres_depth, cnn feature, mv feature, mono feature) in_channels = ( 1 + cnn_feature_channels + mv_transformer_feature_channels + mono_feature_channels ) model_configs = { "vits": { "in_channels": 384, "features": 32, "out_channels": [48, 96, 192, 384], }, "vitb": { "in_channels": 768, "features": 48, "out_channels": [96, 192, 384, 768], }, "vitl": { "in_channels": 1024, "features": 64, "out_channels": [128, 256, 512, 1024], }, } self.upsampler = DPTHead( **model_configs[vit_type], downsample_factor=upsample_factor, num_scales=num_scales, ) self.grid_sample_disable_cudnn = grid_sample_disable_cudnn def normalize_images(self, images): """Normalize image to match the pretrained UniMatch model. images: (B, V, C, H, W) """ shape = [*[1] * (images.dim() - 3), 3, 1, 1] mean = torch.tensor([0.485, 0.456, 0.406]).reshape(*shape).to(images.device) std = torch.tensor([0.229, 0.224, 0.225]).reshape(*shape).to(images.device) return (images - mean) / std def extract_feature(self, images): # images: [B, V, C, H, W] concat = flatten_views(images) # list of [BV, C, H, W], resolution from high to low features = self.backbone(concat) # reverse: resolution from low to high features = features[::-1] return features def forward( self, images, attn_splits_list=None, intrinsics=None, min_depth=1.0 / 0.5, # inverse depth range max_depth=1.0 / 100, num_depth_candidates=128, extrinsics=None, nn_matrix=None, supervise_intermediate_depth=False, **kwargs, ): results_dict = {} depth_preds = [] match_probs = [] # first normalize images images = self.normalize_images(images) b, v, _, ori_h, ori_w = images.shape # update the num_views in unet attention, useful for random input views set_num_views(self.regressor, num_views=v) # NOTE: in this codebase, intrinsics are normalized by image width and height # in unimatch's codebase: https://github.com/autonomousvision/unimatch, no normalization intrinsics = intrinsics.clone() intrinsics[:, :, 0] *= ori_w intrinsics[:, :, 1] *= ori_h # Batch-major flatten: [B, V] -> [B*V], matching flatten_views(). max_depth = max_depth.view(-1) min_depth = min_depth.view(-1) # Scale list of CNN features, each [B*V, C, H, W], resolution low to high. features_list_cnn = self.extract_feature(images) features_list_cnn_all_scales = features_list_cnn features_list_cnn = features_list_cnn[: self.num_scales] results_dict.update({"features_cnn_all_scales": features_list_cnn_all_scales}) # features_cnn is only used internally; returning it keeps extra feature maps alive. # results_dict.update({"features_cnn": features_list_cnn}) # Add positional encoding on flattened features, then run the transformer # in explicit [B, V, C, H, W] layout. attn_splits = attn_splits_list[0] # [BV, C, H, W] features_cnn_pos = mv_feature_add_position( features_list_cnn[0], attn_splits, self.feature_channels ) features_mv = self.transformer( unflatten_views(features_cnn_pos, b, v), attn_num_splits=attn_splits, nn_matrix=nn_matrix, ) features_mv = flatten_views(features_mv) # [B*V, C, H, W] if self.num_scales > 1: # Scale list of MV features, each [B*V, C, H, W], resolution low to high. features_list_mv = self.mv_pyramid(features_mv) else: features_list_mv = [features_mv] results_dict.update({"features_mv": features_list_mv}) # mono feature ori_h, ori_w = images.shape[-2:] resize_h, resize_w = ori_h // 14 * 14, ori_w // 14 * 14 concat = flatten_views(images) concat = F.interpolate( concat, (resize_h, resize_w), mode="bilinear", align_corners=True ) # get intermediate features intermediate_layer_idx = { "vits": [2, 5, 8, 11], "vitb": [2, 5, 8, 11], "vitl": [4, 11, 17, 23], } mono_intermediate_features = list( self.pretrained.get_intermediate_layers( concat, intermediate_layer_idx[self.vit_type], return_class_token=False ) ) for i in range(len(mono_intermediate_features)): curr_features = ( mono_intermediate_features[i] .reshape(concat.shape[0], resize_h // 14, resize_w // 14, -1) .permute(0, 3, 1, 2) .contiguous() ) # resize to 1/8 resolution curr_features = F.interpolate( curr_features, (ori_h // 8, ori_w // 8), mode="bilinear", align_corners=True, ) mono_intermediate_features[i] = curr_features results_dict.update({"features_mono_intermediate": mono_intermediate_features}) # last mono feature mono_features = mono_intermediate_features[-1] if self.lowest_feature_resolution == 4: mono_features = F.interpolate( mono_features, scale_factor=2, mode="bilinear", align_corners=True ) if self.num_scales > 1: # Scale list of mono features, each [B*V, C, H, W], resolution low to high. features_list_mono = self.mono_pyramid(mono_features) else: features_list_mono = [mono_features] # features_mono is only used internally; returning it keeps extra feature maps alive. # results_dict.update({"features_mono": features_list_mono}) depth = None for scale_idx in range(self.num_scales): downsample_factor = self.upsample_factor * ( 2 ** (self.num_scales - 1 - scale_idx) ) # scale intrinsics intrinsics_curr = intrinsics.clone() # [B, V, 3, 3] intrinsics_curr[:, :, :2] = intrinsics_curr[:, :, :2] / downsample_factor # build cost volume features_mv = features_list_mv[scale_idx] # [BV, C, H, W] # batch_features_camera_parameters returns batch-major query tensors # and M target views per query. # ref: [B*V, C, H, W], [B*V, 3, 3], [B*V, 4, 4] # tgt: [B*V, M, C, H, W], [B*V, M, 3, 3], [B*V, M, 4, 4] ( ref_features, ref_intrinsics, ref_extrinsics, tgt_features, tgt_intrinsics, tgt_extrinsics, ) = batch_features_camera_parameters( unflatten_views(features_mv, b, v), intrinsics_curr, extrinsics, nn_matrix=nn_matrix, ) b_new, _, c, h, w = tgt_features.size() # relative pose # extrinsics: c2w with torch.autocast(device_type=tgt_extrinsics.device.type, enabled=False): pose_curr = torch.matmul( torch.linalg.inv(tgt_extrinsics.float()), ref_extrinsics.float().unsqueeze(1), ) # [B*V, M, 4, 4] if scale_idx > 0: # 2x upsample depth assert depth is not None depth = F.interpolate( depth, scale_factor=2, mode="bilinear", align_corners=True ).detach() num_depth_candidates = self.num_depth_candidates // (4**scale_idx) # generate depth candidates with torch.autocast(device_type=features_mv.device.type, enabled=False): min_depth_fp32 = min_depth.float() max_depth_fp32 = max_depth.float() if scale_idx == 0: # min_depth, max_depth: [B*V] depth_interval = (max_depth_fp32 - min_depth_fp32) / ( self.num_depth_candidates - 1 ) # [B*V] linear_space = torch.linspace( 0, 1, num_depth_candidates, device=features_mv.device, dtype=torch.float32, ).view(1, num_depth_candidates, 1, 1) # [1, D, 1, 1] depth_candidates = ( min_depth_fp32.view(-1, 1, 1, 1) + linear_space * (max_depth_fp32 - min_depth_fp32).view(-1, 1, 1, 1) ) # [B*V, D, 1, 1] else: # half interval each scale depth_interval = ( (max_depth_fp32 - min_depth_fp32) / (self.num_depth_candidates - 1) / (2**scale_idx) ) # [B*V] # [B*V, 1, 1, 1] depth_interval = depth_interval.view(-1, 1, 1, 1) # [B*V, 1, H, W] depth_fp32 = depth.float() depth_range_min = ( depth_fp32 - depth_interval * (num_depth_candidates // 2) ).clamp(min=min_depth_fp32.view(-1, 1, 1, 1)) depth_range_max = ( depth_fp32 + depth_interval * (num_depth_candidates // 2 - 1) ).clamp(max=max_depth_fp32.view(-1, 1, 1, 1)) linear_space = torch.linspace( 0, 1, num_depth_candidates, device=features_mv.device, dtype=torch.float32, ).view(1, num_depth_candidates, 1, 1) # [1, D, 1, 1] depth_candidates = depth_range_min + linear_space * ( depth_range_max - depth_range_min ) # [B*V, D, H, W] num_target_views = tgt_features.size(1) if scale_idx == 0: depth_candidates_full = depth_candidates.expand(-1, -1, h, w) else: depth_candidates_full = depth_candidates depth_candidates_curr = ( depth_candidates_full.unsqueeze(1) .expand(-1, num_target_views, -1, -1, -1) .reshape(-1, num_depth_candidates, h, w) ) # [B*V*M, D, H, W] ref_intrinsics_input = ref_intrinsics.unsqueeze(1).expand( -1, num_target_views, -1, -1 ) # [B*V, M, 3, 3] warped_tgt_features = warp_with_pose_depth_candidates( flatten_views(tgt_features), flatten_views(ref_intrinsics_input), flatten_views(tgt_intrinsics), flatten_views(pose_curr), 1.0 / depth_candidates_curr, # convert inverse depth to depth grid_sample_disable_cudnn=self.grid_sample_disable_cudnn, ) # [B*V*M, C, D, H, W] # ref: [B*V, C, H, W] # warped: [B*V*M, C, D, H, W] -> [B*V, M, C, D, H, W] warped_tgt_features = unflatten_views( warped_tgt_features, b_new, num_target_views ) # [B*V, M, D, H, W] -> [B*V, D, H, W] # Average over target views. cost_volume = ( (ref_features.unsqueeze(-3).unsqueeze(1) * warped_tgt_features).sum(2) / (c**0.5) ).mean(1) # regressor features_cnn = features_list_cnn[scale_idx] # [B*V, C, H, W] features_mono = features_list_mono[scale_idx] # [B*V, C, H, W] concat = torch.cat( (cost_volume, features_cnn, features_mv, features_mono), dim=1 ) out = self.regressor[scale_idx](concat) + self.regressor_residual[ scale_idx ](concat) # Depth distribution and expectation are distance/probability # operations, so keep them in FP32 under mixed precision. with torch.autocast(device_type=out.device.type, enabled=False): depth_logits = self.depth_head[scale_idx](out.float()) match_prob = F.softmax(depth_logits, dim=1) # [B*V, D, H, W] match_probs.append(match_prob) if scale_idx == 0: # [B*V, D, H, W] depth_candidates = depth_candidates.expand(-1, -1, h, w) depth = (match_prob * depth_candidates).sum( dim=1, keepdim=True ) # [B*V, 1, H, W] # Keep intermediate predictions only when they are actually supervised. if ( supervise_intermediate_depth and self.training and scale_idx < self.num_scales - 1 ): depth_bilinear = F.interpolate( depth, scale_factor=downsample_factor, mode="bilinear", align_corners=True, ) depth_preds.append(depth_bilinear) # final output, learned upsampler if scale_idx == self.num_scales - 1: residual_depth = self.upsampler( mono_intermediate_features, # resolution high to low cnn_features=features_list_cnn_all_scales[::-1], mv_features=( features_mv if self.num_scales == 1 else features_list_mv[::-1] ), depth=depth, ) depth_bilinear = F.interpolate( depth, scale_factor=self.upsample_factor, mode="bilinear", align_corners=True, ) depth = (depth_bilinear + residual_depth).clamp( min=min_depth.float().view(-1, 1, 1, 1), max=max_depth.float().view(-1, 1, 1, 1), ) depth_preds.append(depth) # convert inverse depth to depth for i in range(len(depth_preds)): depth_pred = 1.0 / depth_preds[i].squeeze(1) # [B*V, H, W] depth_preds[i] = unflatten_views(depth_pred, b, v) # [B, V, H, W] results_dict.update({"depth_preds": depth_preds}) results_dict.update({"match_probs": match_probs}) return results_dict def set_num_views(module, num_views): if isinstance(module, AttentionBlock): module.attention.n_frames = num_views elif ( isinstance(module, nn.ModuleList) or isinstance(module, nn.Sequential) or isinstance(module, nn.Module) ): for submodule in module.children(): set_num_views(submodule, num_views)