"""MapAnything model. Copyright (c) Meta Platforms, Inc. and affiliates. This source code is licensed under the Apache License, Version 2.0 found in the LICENSE file in the root directory of this source tree. """ from __future__ import annotations import warnings from functools import partial from typing import Any, Callable, Dict, List, Tuple, Type, Union import torch from huggingface_hub import PyTorchModelHubMixin from torch import Tensor, nn from uniception.models.encoders import ( EncoderGlobalRepInput, ViTEncoderInput, ViTEncoderNonImageInput, encoder_factory, ) from uniception.models.info_sharing.alternating_attention_transformer import ( MultiViewAlternatingAttentionTransformer, MultiViewAlternatingAttentionTransformerIFR, ) from uniception.models.info_sharing.base import MultiViewTransformerInput from uniception.models.info_sharing.cross_attention_transformer import ( MultiViewCrossAttentionTransformer, MultiViewCrossAttentionTransformerIFR, ) from uniception.models.info_sharing.global_attention_transformer import ( MultiViewGlobalAttentionTransformer, MultiViewGlobalAttentionTransformerIFR, ) from uniception.models.prediction_heads.adaptors import ( CamTranslationPlusQuatsAdaptor, PointMapAdaptor, PointMapPlusRayDirectionsPlusDepthAdaptor, PointMapPlusRayDirectionsPlusDepthWithConfidenceAdaptor, PointMapPlusRayDirectionsPlusDepthWithConfidenceAndMaskAdaptor, PointMapPlusRayDirectionsPlusDepthWithMaskAdaptor, PointMapWithConfidenceAdaptor, PointMapWithConfidenceAndMaskAdaptor, PointMapWithMaskAdaptor, RayDirectionsPlusDepthAdaptor, RayDirectionsPlusDepthWithConfidenceAdaptor, RayDirectionsPlusDepthWithConfidenceAndMaskAdaptor, RayDirectionsPlusDepthWithMaskAdaptor, RayMapPlusDepthAdaptor, RayMapPlusDepthWithConfidenceAdaptor, RayMapPlusDepthWithConfidenceAndMaskAdaptor, RayMapPlusDepthWithMaskAdaptor, ScaleAdaptor, ) from uniception.models.prediction_heads.base import ( AdaptorInput, PredictionHeadInput, PredictionHeadLayeredInput, PredictionHeadTokenInput, ) from uniception.models.prediction_heads.dpt import ( DPTFeature, DPTRegressionProcessor, ) from uniception.models.prediction_heads.linear import LinearFeature from uniception.models.prediction_heads.mlp_head import MLPHead from uniception.models.prediction_heads.pose_head import PoseHead from uniception.models.utils.transformer_blocks import Mlp, SwiGLUFFNFused from mapdet3d.op.mapanything.geometry import ( apply_log_to_norm, convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap, normalize_depth_using_non_zero_pixels, normalize_pose_translations, transform_pose_using_quats_and_trans_2_to_1, ) from mapdet3d.op.mapanything.inference import ( postprocess_model_outputs_for_inference, preprocess_input_views_for_inference, validate_input_views_for_inference, ) class MapAnything(PyTorchModelHubMixin, nn.Module): "Modular MapAnything model class that supports input of images & optional geometric modalities (multiple reconstruction tasks)." def __init__( self, name: str, encoder_config: Dict, info_sharing_config: Dict, pred_head_config: Dict, geometric_input_config: Dict, fusion_norm_layer: Union[ Type[nn.Module], Callable[..., nn.Module] ] = partial(nn.LayerNorm, eps=1e-6), pretrained_checkpoint_path: str = None, load_specific_pretrained_submodules: bool = False, specific_pretrained_submodules: list = None, torch_hub_force_reload: bool = False, use_register_tokens_from_encoder: bool = False, info_sharing_mlp_layer_str: str = "mlp", ): """ Multi-view model containing an image encoder fused with optional geometric modalities followed by a multi-view attention transformer and respective downstream heads. The goal is to output scene representation. The multi-view attention transformer also takes as input a scale token to predict the metric scaling factor for the predicted scene representation. Args: name (str): Name of the model. encoder_config (Dict): Configuration for the encoder. info_sharing_config (Dict): Configuration for the multi-view attention transformer. pred_head_config (Dict): Configuration for the prediction heads. geometric_input_config (Dict): Configuration for the input of optional geometric modalities. fusion_norm_layer (Union[Type[nn.Module], Callable[..., nn.Module]]): Normalization layer to use after fusion (addition) of encoder and geometric modalities. (default: partial(nn.LayerNorm, eps=1e-6)) pretrained_checkpoint_path (str): Path to pretrained checkpoint. (default: None) load_specific_pretrained_submodules (bool): Whether to load specific pretrained submodules. (default: False) specific_pretrained_submodules (list): List of specific pretrained submodules to load. Must be provided when load_specific_pretrained_submodules is True. (default: None) torch_hub_force_reload (bool): Whether to force reload the encoder from torch hub. (default: False) use_register_tokens_from_encoder (bool): Whether to use register tokens from encoder. (default: False) info_sharing_mlp_layer_str (str): Type of MLP layer to use in the multi-view transformer. Useful for DINO init of the multi-view transformer. Options: "mlp" or "swiglufused". (default: "mlp") """ super().__init__() # Initialize the attributes self.name = name self.encoder_config = encoder_config self.info_sharing_config = info_sharing_config self.pred_head_config = pred_head_config self.geometric_input_config = geometric_input_config self.pretrained_checkpoint_path = pretrained_checkpoint_path self.load_specific_pretrained_submodules = ( load_specific_pretrained_submodules ) self.specific_pretrained_submodules = specific_pretrained_submodules self.torch_hub_force_reload = torch_hub_force_reload self.use_register_tokens_from_encoder = ( use_register_tokens_from_encoder ) self.info_sharing_mlp_layer_str = info_sharing_mlp_layer_str self.class_init_args = { "name": self.name, "encoder_config": self.encoder_config, "info_sharing_config": self.info_sharing_config, "pred_head_config": self.pred_head_config, "geometric_input_config": self.geometric_input_config, "pretrained_checkpoint_path": self.pretrained_checkpoint_path, "load_specific_pretrained_submodules": self.load_specific_pretrained_submodules, "specific_pretrained_submodules": self.specific_pretrained_submodules, "torch_hub_force_reload": self.torch_hub_force_reload, "use_register_tokens_from_encoder": self.use_register_tokens_from_encoder, "info_sharing_mlp_layer_str": self.info_sharing_mlp_layer_str, } # Get relevant parameters from the configs self.info_sharing_type = info_sharing_config["model_type"] self.info_sharing_return_type = info_sharing_config[ "model_return_type" ] self.pred_head_type = pred_head_config["type"] # Initialize image encoder if self.encoder_config["uses_torch_hub"]: self.encoder_config["torch_hub_force_reload"] = ( torch_hub_force_reload ) # Create a copy of the config before deleting the key to preserve it for serialization encoder_config_copy = self.encoder_config.copy() del encoder_config_copy["uses_torch_hub"] self.encoder = encoder_factory(**encoder_config_copy) # Initialize the encoder for ray directions ray_dirs_encoder_config = self.geometric_input_config[ "ray_dirs_encoder_config" ] ray_dirs_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim ray_dirs_encoder_config["patch_size"] = self.encoder.patch_size self.ray_dirs_encoder = encoder_factory(**ray_dirs_encoder_config) # Initialize the encoder for depth (normalized per view and values after normalization are scaled logarithmically) depth_encoder_config = self.geometric_input_config[ "depth_encoder_config" ] depth_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim depth_encoder_config["patch_size"] = self.encoder.patch_size self.depth_encoder = encoder_factory(**depth_encoder_config) # Initialize the encoder for log scale factor of depth depth_scale_encoder_config = self.geometric_input_config[ "scale_encoder_config" ] depth_scale_encoder_config["enc_embed_dim"] = ( self.encoder.enc_embed_dim ) self.depth_scale_encoder = encoder_factory( **depth_scale_encoder_config ) # Initialize the encoder for camera rotation cam_rot_encoder_config = self.geometric_input_config[ "cam_rot_encoder_config" ] cam_rot_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim self.cam_rot_encoder = encoder_factory(**cam_rot_encoder_config) # Initialize the encoder for camera translation (normalized across all provided camera translations) cam_trans_encoder_config = self.geometric_input_config[ "cam_trans_encoder_config" ] cam_trans_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim self.cam_trans_encoder = encoder_factory(**cam_trans_encoder_config) # Initialize the encoder for log scale factor of camera translation cam_trans_scale_encoder_config = self.geometric_input_config[ "scale_encoder_config" ] cam_trans_scale_encoder_config["enc_embed_dim"] = ( self.encoder.enc_embed_dim ) self.cam_trans_scale_encoder = encoder_factory( **cam_trans_scale_encoder_config ) # Initialize the fusion norm layer self.fusion_norm_layer = fusion_norm_layer(self.encoder.enc_embed_dim) # Initialize the Scale Token # Used to scale the final scene predictions to metric scale # During inference extended to (B, C, T), where T is the number of tokens (i.e., 1) self.scale_token = nn.Parameter( torch.zeros(self.encoder.enc_embed_dim) ) torch.nn.init.trunc_normal_(self.scale_token, std=0.02) # Set the MLP layer config for the info sharing transformer if info_sharing_mlp_layer_str == "mlp": info_sharing_config["module_args"]["mlp_layer"] = Mlp elif info_sharing_mlp_layer_str == "swiglufused": info_sharing_config["module_args"]["mlp_layer"] = SwiGLUFFNFused else: raise ValueError( f"Invalid info_sharing_mlp_layer_str: {info_sharing_mlp_layer_str}. Valid options: ['mlp', 'swiglufused']" ) # Initialize the info sharing module (multi-view transformer) self._initialize_info_sharing(info_sharing_config) # Initialize the prediction heads self._initialize_prediction_heads(pred_head_config) # Initialize the final adaptors self._initialize_adaptors(pred_head_config) # Load pretrained weights self._load_pretrained_weights() @property def device(self) -> torch.device: return next(self.parameters()).device @property def dtype(self) -> torch.dtype: return next(self.parameters()).dtype def _initialize_info_sharing(self, info_sharing_config): """ Initialize the information sharing module based on the configuration. This method sets up the custom positional encoding if specified and initializes the appropriate multi-view transformer based on the configuration type. Args: info_sharing_config (Dict): Configuration for the multi-view attention transformer. Should contain 'custom_positional_encoding', 'model_type', and 'model_return_type'. Returns: None Raises: ValueError: If invalid configuration options are provided. """ # Initialize Custom Positional Encoding if required custom_positional_encoding = info_sharing_config[ "custom_positional_encoding" ] if custom_positional_encoding is not None: if isinstance(custom_positional_encoding, str): print( f"Using custom positional encoding for multi-view attention transformer: {custom_positional_encoding}" ) raise ValueError( f"Invalid custom_positional_encoding: {custom_positional_encoding}. None implemented." ) elif isinstance(custom_positional_encoding, Callable): print( "Using callable function as custom positional encoding for multi-view attention transformer." ) self.custom_positional_encoding = custom_positional_encoding else: self.custom_positional_encoding = None # Add dependecies to info_sharing_config info_sharing_config["module_args"][ "input_embed_dim" ] = self.encoder.enc_embed_dim info_sharing_config["module_args"][ "custom_positional_encoding" ] = self.custom_positional_encoding # Initialize Multi-View Transformer if self.info_sharing_return_type == "no_intermediate_features": # Returns only normalized last layer features # Intialize multi-view transformer based on type if self.info_sharing_type == "cross_attention": self.info_sharing = MultiViewCrossAttentionTransformer( **info_sharing_config["module_args"] ) elif self.info_sharing_type == "global_attention": self.info_sharing = MultiViewGlobalAttentionTransformer( **info_sharing_config["module_args"] ) elif self.info_sharing_type == "alternating_attention": self.info_sharing = MultiViewAlternatingAttentionTransformer( **info_sharing_config["module_args"] ) else: raise ValueError( f"Invalid info_sharing_type: {self.info_sharing_type}. Valid options: ['cross_attention', 'global_attention', 'alternating_attention']" ) elif self.info_sharing_return_type == "intermediate_features": # Returns intermediate features and normalized last layer features # Initialize mulit-view transformer based on type if self.info_sharing_type == "cross_attention": self.info_sharing = MultiViewCrossAttentionTransformerIFR( **info_sharing_config["module_args"] ) elif self.info_sharing_type == "global_attention": self.info_sharing = MultiViewGlobalAttentionTransformerIFR( **info_sharing_config["module_args"] ) elif self.info_sharing_type == "alternating_attention": self.info_sharing = ( MultiViewAlternatingAttentionTransformerIFR( **info_sharing_config["module_args"] ) ) else: raise ValueError( f"Invalid info_sharing_type: {self.info_sharing_type}. Valid options: ['cross_attention', 'global_attention', 'alternating_attention']" ) # Assess if the DPT needs to use encoder features if len(self.info_sharing.indices) == 2: self.use_encoder_features_for_dpt = True elif len(self.info_sharing.indices) == 3: self.use_encoder_features_for_dpt = False else: raise ValueError( "Invalid number of indices provided for info sharing feature returner. Please provide 2 or 3 indices." ) else: raise ValueError( f"Invalid info_sharing_return_type: {self.info_sharing_return_type}. Valid options: ['no_intermediate_features', 'intermediate_features']" ) def _initialize_prediction_heads(self, pred_head_config): """ Initialize the prediction heads based on the prediction head configuration. This method configures and initializes the appropriate prediction heads based on the specified prediction head type (linear, DPT, or DPT+pose). It sets up the necessary dependencies and creates the required model components. Args: pred_head_config (Dict): Configuration for the prediction heads. Returns: None Raises: ValueError: If an invalid pred_head_type is provided. """ # Add dependencies to prediction head config pred_head_config["feature_head"][ "patch_size" ] = self.encoder.patch_size if self.pred_head_type == "linear": pred_head_config["feature_head"][ "input_feature_dim" ] = self.info_sharing.dim elif "dpt" in self.pred_head_type: # Add dependencies for DPT & Regressor head if self.use_encoder_features_for_dpt: pred_head_config["feature_head"]["input_feature_dims"] = [ self.encoder.enc_embed_dim ] + [self.info_sharing.dim] * 3 else: pred_head_config["feature_head"]["input_feature_dims"] = [ self.info_sharing.dim ] * 4 pred_head_config["regressor_head"]["input_feature_dim"] = ( pred_head_config["feature_head"]["feature_dim"] ) # Add dependencies for Pose head if required if "pose" in self.pred_head_type: pred_head_config["pose_head"][ "patch_size" ] = self.encoder.patch_size pred_head_config["pose_head"][ "input_feature_dim" ] = self.info_sharing.dim else: raise ValueError( f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" ) pred_head_config["scale_head"][ "input_feature_dim" ] = self.info_sharing.dim # Initialize Prediction Heads if self.pred_head_type == "linear": # Initialize Dense Prediction Head for all views self.dense_head = LinearFeature(**pred_head_config["feature_head"]) elif "dpt" in self.pred_head_type: # Initialze Dense Predction Head for all views self.dpt_feature_head = DPTFeature( **pred_head_config["feature_head"] ) self.dpt_regressor_head = DPTRegressionProcessor( **pred_head_config["regressor_head"] ) self.dense_head = nn.Sequential( self.dpt_feature_head, self.dpt_regressor_head ) # Initialize Pose Head for all views if required if "pose" in self.pred_head_type: self.pose_head = PoseHead(**pred_head_config["pose_head"]) else: raise ValueError( f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" ) self.scale_head = MLPHead(**pred_head_config["scale_head"]) def _initialize_adaptors(self, pred_head_config): """ Initialize the adaptors based on the prediction head configuration. This method sets up the appropriate adaptors for different scene representation types, such as pointmaps, ray maps with depth, or ray directions with depth and pose. Args: pred_head_config (Dict): Configuration for the prediction heads including adaptor type. Returns: None Raises: ValueError: If an invalid adaptor_type is provided. AssertionError: If ray directions + depth + pose is used with an incompatible head type. """ if pred_head_config["adaptor_type"] == "pointmap": self.dense_adaptor = PointMapAdaptor(**pred_head_config["adaptor"]) self.scene_rep_type = "pointmap" elif pred_head_config["adaptor_type"] == "pointmap+confidence": self.dense_adaptor = PointMapWithConfidenceAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "pointmap+confidence" elif pred_head_config["adaptor_type"] == "pointmap+mask": self.dense_adaptor = PointMapWithMaskAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "pointmap+mask" elif pred_head_config["adaptor_type"] == "pointmap+confidence+mask": self.dense_adaptor = PointMapWithConfidenceAndMaskAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "pointmap+confidence+mask" elif pred_head_config["adaptor_type"] == "raymap+depth": self.dense_adaptor = RayMapPlusDepthAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "raymap+depth" elif pred_head_config["adaptor_type"] == "raymap+depth+confidence": self.dense_adaptor = RayMapPlusDepthWithConfidenceAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "raymap+depth+confidence" elif pred_head_config["adaptor_type"] == "raymap+depth+mask": self.dense_adaptor = RayMapPlusDepthWithMaskAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "raymap+depth+mask" elif ( pred_head_config["adaptor_type"] == "raymap+depth+confidence+mask" ): self.dense_adaptor = RayMapPlusDepthWithConfidenceAndMaskAdaptor( **pred_head_config["adaptor"] ) self.scene_rep_type = "raymap+depth+confidence+mask" elif pred_head_config["adaptor_type"] == "raydirs+depth+pose": assert ( self.pred_head_type == "dpt+pose" ), "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = RayDirectionsPlusDepthAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "raydirs+depth+pose" elif ( pred_head_config["adaptor_type"] == "raydirs+depth+pose+confidence" ): assert ( self.pred_head_type == "dpt+pose" ), "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = RayDirectionsPlusDepthWithConfidenceAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "raydirs+depth+pose+confidence" elif pred_head_config["adaptor_type"] == "raydirs+depth+pose+mask": assert ( self.pred_head_type == "dpt+pose" ), "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = RayDirectionsPlusDepthWithMaskAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "raydirs+depth+pose+mask" elif ( pred_head_config["adaptor_type"] == "raydirs+depth+pose+confidence+mask" ): assert ( self.pred_head_type == "dpt+pose" ), "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = ( RayDirectionsPlusDepthWithConfidenceAndMaskAdaptor( **pred_head_config["dpt_adaptor"] ) ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "raydirs+depth+pose+confidence+mask" elif pred_head_config["adaptor_type"] == "campointmap+pose": assert ( self.pred_head_type == "dpt+pose" ), "Camera pointmap + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = PointMapAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "campointmap+pose" elif pred_head_config["adaptor_type"] == "campointmap+pose+confidence": assert ( self.pred_head_type == "dpt+pose" ), "Camera pointmap + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = PointMapWithConfidenceAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "campointmap+pose+confidence" elif pred_head_config["adaptor_type"] == "campointmap+pose+mask": assert ( self.pred_head_type == "dpt+pose" ), "Camera pointmap + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = PointMapWithMaskAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "campointmap+pose+mask" elif ( pred_head_config["adaptor_type"] == "campointmap+pose+confidence+mask" ): assert ( self.pred_head_type == "dpt+pose" ), "Camera pointmap + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = PointMapWithConfidenceAndMaskAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "campointmap+pose+confidence+mask" elif pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose": assert ( self.pred_head_type == "dpt+pose" ), "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = PointMapPlusRayDirectionsPlusDepthAdaptor( **pred_head_config["dpt_adaptor"] ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "pointmap+raydirs+depth+pose" elif ( pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose+confidence" ): assert ( self.pred_head_type == "dpt+pose" ), "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = ( PointMapPlusRayDirectionsPlusDepthWithConfidenceAdaptor( **pred_head_config["dpt_adaptor"] ) ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "pointmap+raydirs+depth+pose+confidence" elif ( pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose+mask" ): assert ( self.pred_head_type == "dpt+pose" ), "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = ( PointMapPlusRayDirectionsPlusDepthWithMaskAdaptor( **pred_head_config["dpt_adaptor"] ) ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "pointmap+raydirs+depth+pose+mask" elif ( pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose+confidence+mask" ): assert ( self.pred_head_type == "dpt+pose" ), "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." self.dense_adaptor = ( PointMapPlusRayDirectionsPlusDepthWithConfidenceAndMaskAdaptor( **pred_head_config["dpt_adaptor"] ) ) self.pose_adaptor = CamTranslationPlusQuatsAdaptor( **pred_head_config["pose_adaptor"] ) self.scene_rep_type = "pointmap+raydirs+depth+pose+confidence+mask" else: raise ValueError( f"Invalid adaptor_type: {pred_head_config['adaptor_type']}. \ Valid options: ['pointmap', 'raymap+depth', 'raydirs+depth+pose', 'campointmap+pose', 'pointmap+raydirs+depth+pose' \ 'pointmap+confidence', 'raymap+depth+confidence', 'raydirs+depth+pose+confidence', 'campointmap+pose+confidence', 'pointmap+raydirs+depth+pose+confidence' \ 'pointmap+mask', 'raymap+depth+mask', 'raydirs+depth+pose+mask', 'campointmap+pose+mask', 'pointmap+raydirs+depth+pose+mask' \ 'pointmap+confidence+mask', 'raymap+depth+confidence+mask', 'raydirs+depth+pose+confidence+mask', 'campointmap+pose+confidence+mask', 'pointmap+raydirs+depth+pose+confidence+mask']" ) self.scale_adaptor = ScaleAdaptor(**pred_head_config["scale_adaptor"]) def _load_pretrained_weights(self): """ Load pretrained weights from a checkpoint file. If load_specific_pretrained_submodules is True, only loads weights for the specified submodules. Otherwise, loads all weights from the checkpoint. Returns: None """ if self.pretrained_checkpoint_path is not None: if not self.load_specific_pretrained_submodules: print( f"Loading pretrained MapAnything weights from {self.pretrained_checkpoint_path} ..." ) ckpt = torch.load( self.pretrained_checkpoint_path, weights_only=False ) print(self.load_state_dict(ckpt["model"])) else: print( f"Loading pretrained MapAnything weights from {self.pretrained_checkpoint_path} for specific submodules: {self.specific_pretrained_submodules} ..." ) assert ( self.pred_head_type is not None ), "Specific submodules to load cannot be None." ckpt = torch.load( self.pretrained_checkpoint_path, weights_only=False ) filtered_ckpt = {} for ckpt_key, ckpt_value in ckpt["model"].items(): for submodule in self.specific_pretrained_submodules: if ckpt_key.startswith(submodule): filtered_ckpt[ckpt_key] = ckpt_value print(self.load_state_dict(filtered_ckpt, strict=False)) def _encode_n_views(self, views) -> list[Tensor]: """ Encode all the input views (batch of images) in a single forward pass. Assumes all the input views have the same image shape, batch size, and data normalization type. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. Returns: A tuple containing: List[torch.Tensor]: A list containing the encoded features for all N views. List[torch.Tensor]: A list containing the encoded per-view registers for all N views. """ num_views = len(views) data_norm_type = views[0]["data_norm_type"][0] imgs_list = [view["img"] for view in views] all_imgs_across_views = torch.cat(imgs_list, dim=0) encoder_input = ViTEncoderInput( image=all_imgs_across_views, data_norm_type=data_norm_type ) encoder_output = self.encoder(encoder_input) all_encoder_features_across_views = encoder_output.features.chunk( num_views, dim=0 ) all_encoder_registers_across_views = None if ( self.use_register_tokens_from_encoder and encoder_output.registers is not None ): all_encoder_registers_across_views = ( encoder_output.registers.chunk(num_views, dim=0) ) return ( all_encoder_features_across_views, all_encoder_registers_across_views, ) def _compute_pose_quats_and_trans_for_across_views_in_ref_view( self, views, num_views, device, dtype, batch_size_per_view, per_sample_cam_input_mask, ): """ Compute the pose quats and trans for all the views in the frame of the reference view 0. Returns identity pose for views where the camera input mask is False or the pose is not provided. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. num_views (int): Number of views. device (torch.device): Device to use for the computation. dtype (torch.dtype): Data type to use for the computation. per_sample_cam_input_mask (Tensor): Tensor containing the per sample camera input mask. Returns: Tensor: A tensor containing the pose quats for all the views in the frame of the reference view 0. (batch_size_per_view * view, 4) Tensor: A tensor containing the pose trans for all the views in the frame of the reference view 0. (batch_size_per_view * view, 3) Tensor: A tensor containing the per sample camera input mask. """ # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 pose_quats_non_ref_views = [] pose_trans_non_ref_views = [] pose_quats_ref_view_0 = [] pose_trans_ref_view_0 = [] for view_idx in range(num_views): per_sample_cam_input_mask_for_curr_view = ( per_sample_cam_input_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] ) if ( "camera_pose_quats" in views[view_idx] and "camera_pose_trans" in views[view_idx] and per_sample_cam_input_mask_for_curr_view.any() ): # Get the camera pose quats and trans for the current view cam_pose_quats = views[view_idx]["camera_pose_quats"][ per_sample_cam_input_mask_for_curr_view ] cam_pose_trans = views[view_idx]["camera_pose_trans"][ per_sample_cam_input_mask_for_curr_view ] # Append to the list pose_quats_non_ref_views.append(cam_pose_quats) pose_trans_non_ref_views.append(cam_pose_trans) # Get the camera pose quats and trans for the reference view 0 cam_pose_quats = views[0]["camera_pose_quats"][ per_sample_cam_input_mask_for_curr_view ] cam_pose_trans = views[0]["camera_pose_trans"][ per_sample_cam_input_mask_for_curr_view ] # Append to the list pose_quats_ref_view_0.append(cam_pose_quats) pose_trans_ref_view_0.append(cam_pose_trans) else: per_sample_cam_input_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] = False # Initialize the pose quats and trans for all views as identity pose_quats_across_views = torch.tensor( [0.0, 0.0, 0.0, 1.0], dtype=dtype, device=device ).repeat( batch_size_per_view * num_views, 1 ) # (q_x, q_y, q_z, q_w) pose_trans_across_views = torch.zeros( (batch_size_per_view * num_views, 3), dtype=dtype, device=device ) # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 if len(pose_quats_non_ref_views) > 0: # Stack the pose quats and trans for all the non-reference views and reference view 0 pose_quats_non_ref_views = torch.cat( pose_quats_non_ref_views, dim=0 ) pose_trans_non_ref_views = torch.cat( pose_trans_non_ref_views, dim=0 ) pose_quats_ref_view_0 = torch.cat(pose_quats_ref_view_0, dim=0) pose_trans_ref_view_0 = torch.cat(pose_trans_ref_view_0, dim=0) # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 ( pose_quats_non_ref_views_in_ref_view_0, pose_trans_non_ref_views_in_ref_view_0, ) = transform_pose_using_quats_and_trans_2_to_1( pose_quats_ref_view_0, pose_trans_ref_view_0, pose_quats_non_ref_views, pose_trans_non_ref_views, ) # Update the pose quats and trans for all the non-reference views pose_quats_across_views[per_sample_cam_input_mask] = ( pose_quats_non_ref_views_in_ref_view_0.to(dtype=dtype) ) pose_trans_across_views[per_sample_cam_input_mask] = ( pose_trans_non_ref_views_in_ref_view_0.to(dtype=dtype) ) return ( pose_quats_across_views, pose_trans_across_views, per_sample_cam_input_mask, ) def _encode_and_fuse_ray_dirs( self, views, num_views, batch_size_per_view, all_encoder_features_across_views, per_sample_ray_dirs_input_mask, ): """ Encode the ray directions for all the views and fuse it with the other encoder features in a single forward pass. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. num_views (int): Number of views. batch_size_per_view (int): Batch size per view. all_encoder_features_across_views (Tensor): Tensor containing the encoded features for all N views. per_sample_ray_dirs_input_mask (Tensor): Tensor containing the per sample ray direction input mask. Returns: Tensor: A tensor containing the encoded features for all the views. """ # Get the height and width of the images _, _, height, width = views[0]["img"].shape # Get the ray directions for all the views where info is provided and the ray direction input mask is True ray_dirs_list = [] for view_idx in range(num_views): per_sample_ray_dirs_input_mask_for_curr_view = ( per_sample_ray_dirs_input_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] ) ray_dirs_for_curr_view = torch.zeros( (batch_size_per_view, height, width, 3), dtype=all_encoder_features_across_views.dtype, device=all_encoder_features_across_views.device, ) if ( "ray_directions_cam" in views[view_idx] and per_sample_ray_dirs_input_mask_for_curr_view.any() ): ray_dirs_for_curr_view[ per_sample_ray_dirs_input_mask_for_curr_view ] = views[view_idx]["ray_directions_cam"][ per_sample_ray_dirs_input_mask_for_curr_view ] else: per_sample_ray_dirs_input_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] = False ray_dirs_list.append(ray_dirs_for_curr_view) # Stack the ray directions for all the views and permute to (B * V, C, H, W) ray_dirs = torch.cat(ray_dirs_list, dim=0) # (B * V, H, W, 3) ray_dirs = ray_dirs.permute( 0, 3, 1, 2 ).contiguous() # (B * V, 3, H, W) # Encode the ray directions ray_dirs_features_across_views = self.ray_dirs_encoder( ViTEncoderNonImageInput(data=ray_dirs) ).features # Fuse the ray direction features with the other encoder features (zero out the features where the ray direction input mask is False) ray_dirs_features_across_views = ( ray_dirs_features_across_views * per_sample_ray_dirs_input_mask.unsqueeze(-1) .unsqueeze(-1) .unsqueeze(-1) ) all_encoder_features_across_views = ( all_encoder_features_across_views + ray_dirs_features_across_views ) return all_encoder_features_across_views def _encode_and_fuse_depths( self, views, num_views, batch_size_per_view, all_encoder_features_across_views, per_sample_depth_input_mask, ): """ Encode the z depths for all the views and fuse it with the other encoder features in a single forward pass. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. num_views (int): Number of views. batch_size_per_view (int): Batch size per view. all_encoder_features_across_views (Tensor): Tensor containing the encoded features for all N views. per_sample_depth_input_mask (Tensor): Tensor containing the per sample depth input mask. Returns: Tensor: A tensor containing the encoded features for all the views. """ # Get the device and height and width of the images device = all_encoder_features_across_views.device _, _, height, width = views[0]["img"].shape # Decide to use randomly sampled sparse depth or dense depth if torch.rand(1) < self.geometric_input_config["sparse_depth_prob"]: use_sparse_depth = True else: use_sparse_depth = False # Get the depths for all the views depth_list = [] depth_norm_factors_list = [] metric_scale_depth_mask_list = [] for view_idx in range(num_views): # Get the input mask for current view per_sample_depth_input_mask_for_curr_view = ( per_sample_depth_input_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] ) depth_for_curr_view = torch.zeros( (batch_size_per_view, height, width, 1), dtype=all_encoder_features_across_views.dtype, device=device, ) depth_norm_factor_for_curr_view = torch.zeros( (batch_size_per_view), dtype=all_encoder_features_across_views.dtype, device=device, ) metric_scale_mask_for_curr_view = torch.zeros( (batch_size_per_view), dtype=torch.bool, device=device, ) if ( "depth_along_ray" in views[view_idx] ) and per_sample_depth_input_mask_for_curr_view.any(): # Get depth for current view depth_for_curr_view_input = views[view_idx]["depth_along_ray"][ per_sample_depth_input_mask_for_curr_view ] # Get the metric scale mask if "is_metric_scale" in views[view_idx]: metric_scale_mask = views[view_idx]["is_metric_scale"][ per_sample_depth_input_mask_for_curr_view ] else: metric_scale_mask = torch.zeros( depth_for_curr_view_input.shape[0], dtype=torch.bool, device=device, ) # Turn off indication of metric scale samples based on the depth_scale_norm_all_prob depth_scale_norm_all_mask = ( torch.rand(metric_scale_mask.shape[0]) < self.geometric_input_config["depth_scale_norm_all_prob"] ) if depth_scale_norm_all_mask.any(): metric_scale_mask[depth_scale_norm_all_mask] = False # Assign the metric scale mask to the respective indices metric_scale_mask_for_curr_view[ per_sample_depth_input_mask_for_curr_view ] = metric_scale_mask # Sparsely sample the depth if required if use_sparse_depth: # Create a mask of ones sparsification_mask = torch.ones_like( depth_for_curr_view_input, device=device ) # Create a mask for valid pixels (depth > 0) valid_pixel_mask = depth_for_curr_view_input > 0 # Calculate the number of valid pixels num_valid_pixels = valid_pixel_mask.sum().item() # Calculate the number of valid pixels to set to zero num_to_zero = int( num_valid_pixels * self.geometric_input_config[ "sparsification_removal_percent" ] ) if num_to_zero > 0: # Get the indices of valid pixels valid_indices = valid_pixel_mask.nonzero(as_tuple=True) # Randomly select indices to zero out indices_to_zero = torch.randperm(num_valid_pixels)[ :num_to_zero ] # Set selected valid indices to zero in the mask sparsification_mask[ valid_indices[0][indices_to_zero], valid_indices[1][indices_to_zero], valid_indices[2][indices_to_zero], valid_indices[3][indices_to_zero], ] = 0 # Apply the mask on the depth depth_for_curr_view_input = ( depth_for_curr_view_input * sparsification_mask ) # Normalize the depth scaled_depth_for_curr_view_input, depth_norm_factor = ( normalize_depth_using_non_zero_pixels( depth_for_curr_view_input, return_norm_factor=True ) ) # Assign the depth and depth norm factor to the respective indices depth_for_curr_view[ per_sample_depth_input_mask_for_curr_view ] = scaled_depth_for_curr_view_input depth_norm_factor_for_curr_view[ per_sample_depth_input_mask_for_curr_view ] = depth_norm_factor else: per_sample_depth_input_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] = False # Append the depths, depth norm factor and metric scale mask for the current view depth_list.append(depth_for_curr_view) depth_norm_factors_list.append(depth_norm_factor_for_curr_view) metric_scale_depth_mask_list.append( metric_scale_mask_for_curr_view ) # Stack the depths for all the views and permute to (B * V, C, H, W) depths = torch.cat(depth_list, dim=0) # (B * V, H, W, 1) depths = apply_log_to_norm( depths ) # Scale logarithimically (norm is computed along last dim) depths = depths.permute(0, 3, 1, 2).contiguous() # (B * V, 1, H, W) # Encode the depths using the depth encoder depth_features_across_views = self.depth_encoder( ViTEncoderNonImageInput(data=depths) ).features # Zero out the depth features where the depth input mask is False depth_features_across_views = ( depth_features_across_views * per_sample_depth_input_mask.unsqueeze(-1) .unsqueeze(-1) .unsqueeze(-1) ) # Stack the depth norm factors for all the views depth_norm_factors = torch.cat( depth_norm_factors_list, dim=0 ) # (B * V, ) # Encode the depth norm factors using the log scale encoder for depth log_depth_norm_factors = torch.log( depth_norm_factors + 1e-8 ) # (B * V, ) depth_scale_features_across_views = self.depth_scale_encoder( EncoderGlobalRepInput(data=log_depth_norm_factors.unsqueeze(-1)) ).features # Zero out the depth scale features where the depth input mask is False depth_scale_features_across_views = ( depth_scale_features_across_views * per_sample_depth_input_mask.unsqueeze(-1) ) # Stack the metric scale mask for all the views metric_scale_depth_mask = torch.cat( metric_scale_depth_mask_list, dim=0 ) # (B * V, ) # Zero out the depth scale features where the metric scale mask is False # Scale encoding is only provided for metric scale samples depth_scale_features_across_views = ( depth_scale_features_across_views * metric_scale_depth_mask.unsqueeze(-1) ) # Fuse the depth features & depth scale features with the other encoder features all_encoder_features_across_views = ( all_encoder_features_across_views + depth_features_across_views + depth_scale_features_across_views.unsqueeze(-1).unsqueeze(-1) ) return all_encoder_features_across_views def _encode_and_fuse_cam_quats_and_trans( self, views, num_views, batch_size_per_view, all_encoder_features_across_views, pose_quats_across_views, pose_trans_across_views, per_sample_cam_input_mask, ): """ Encode the camera quats and trans for all the views and fuse it with the other encoder features in a single forward pass. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. num_views (int): Number of views. batch_size_per_view (int): Batch size per view. all_encoder_features_across_views (Tensor): Tensor containing the encoded features for all N views. pose_quats_across_views (Tensor): Tensor containing the pose quats for all the views in the frame of the reference view 0. (batch_size_per_view * view, 4) pose_trans_across_views (Tensor): Tensor containing the pose trans for all the views in the frame of the reference view 0. (batch_size_per_view * view, 3) per_sample_cam_input_mask (Tensor): Tensor containing the per sample camera input mask. Returns: Tensor: A tensor containing the encoded features for all the views. """ # Encode the pose quats pose_quats_features_across_views = self.cam_rot_encoder( EncoderGlobalRepInput(data=pose_quats_across_views) ).features # Zero out the pose quat features where the camera input mask is False pose_quats_features_across_views = ( pose_quats_features_across_views * per_sample_cam_input_mask.unsqueeze(-1) ) # Get the metric scale mask for all samples device = all_encoder_features_across_views.device metric_scale_pose_trans_mask = torch.zeros( (batch_size_per_view * num_views), dtype=torch.bool, device=device ) for view_idx in range(num_views): if "is_metric_scale" in views[view_idx]: # Get the metric scale mask for the input pose priors metric_scale_mask = views[view_idx]["is_metric_scale"] else: metric_scale_mask = torch.zeros( batch_size_per_view, dtype=torch.bool, device=device ) metric_scale_pose_trans_mask[ view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view ] = metric_scale_mask # Turn off indication of metric scale samples based on the pose_scale_norm_all_prob pose_norm_all_mask = ( torch.rand(batch_size_per_view * num_views) < self.geometric_input_config["pose_scale_norm_all_prob"] ) if pose_norm_all_mask.any(): metric_scale_pose_trans_mask[pose_norm_all_mask] = False # Get the scale norm factor for all the samples and scale the pose translations pose_trans_across_views = torch.split( pose_trans_across_views, batch_size_per_view, dim=0 ) # Split into num_views chunks pose_trans_across_views = torch.stack( pose_trans_across_views, dim=1 ) # Stack the views along a new dimension (batch_size_per_view, num_views, 3) scaled_pose_trans_across_views, pose_trans_norm_factors = ( normalize_pose_translations( pose_trans_across_views, return_norm_factor=True ) ) # Resize the pose translation back to (batch_size_per_view * num_views, 3) and extend the norm factor to (batch_size_per_view * num_views, 1) scaled_pose_trans_across_views = scaled_pose_trans_across_views.unbind( dim=1 ) # Convert back to list of views, where each view has batch_size_per_view tensor scaled_pose_trans_across_views = torch.cat( scaled_pose_trans_across_views, dim=0 ) # Concatenate back to (batch_size_per_view * num_views, 3) pose_trans_norm_factors_across_views = ( pose_trans_norm_factors.unsqueeze(-1).repeat(num_views, 1) ) # (B, ) -> (B * V, 1) # Encode the pose trans pose_trans_features_across_views = self.cam_trans_encoder( EncoderGlobalRepInput(data=scaled_pose_trans_across_views) ).features # Zero out the pose trans features where the camera input mask is False pose_trans_features_across_views = ( pose_trans_features_across_views * per_sample_cam_input_mask.unsqueeze(-1) ) # Encode the pose translation norm factors using the log scale encoder for pose trans log_pose_trans_norm_factors_across_views = torch.log( pose_trans_norm_factors_across_views + 1e-8 ) pose_trans_scale_features_across_views = self.cam_trans_scale_encoder( EncoderGlobalRepInput( data=log_pose_trans_norm_factors_across_views ) ).features # Zero out the pose trans scale features where the camera input mask is False pose_trans_scale_features_across_views = ( pose_trans_scale_features_across_views * per_sample_cam_input_mask.unsqueeze(-1) ) # Zero out the pose trans scale features where the metric scale mask is False # Scale encoding is only provided for metric scale samples pose_trans_scale_features_across_views = ( pose_trans_scale_features_across_views * metric_scale_pose_trans_mask.unsqueeze(-1) ) # Fuse the pose quat features, pose trans features, pose trans scale features and pose trans type PE features with the other encoder features all_encoder_features_across_views = ( all_encoder_features_across_views + pose_quats_features_across_views.unsqueeze(-1).unsqueeze(-1) + pose_trans_features_across_views.unsqueeze(-1).unsqueeze(-1) + pose_trans_scale_features_across_views.unsqueeze(-1).unsqueeze( -1 ) ) return all_encoder_features_across_views def _encode_and_fuse_optional_geometric_inputs( self, views, all_encoder_features_across_views_list ): """ Encode all the input optional geometric modalities and fuses it with the image encoder features in a single forward pass. Assumes all the input views have the same shape and batch size. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. all_encoder_features_across_views (List[Tensor]): List of tensors containing the encoded image features for all N views. Returns: List[Tensor]: A list containing the encoded features for all N views. """ num_views = len(views) batch_size_per_view, _, _, _ = views[0]["img"].shape device = all_encoder_features_across_views_list[0].device dtype = all_encoder_features_across_views_list[0].dtype all_encoder_features_across_views = torch.cat( all_encoder_features_across_views_list, dim=0 ) # Get the overall input mask for all the views overall_geometric_input_mask = ( torch.rand(batch_size_per_view, device=device) < self.geometric_input_config["overall_prob"] ) overall_geometric_input_mask = overall_geometric_input_mask.repeat( num_views ) # Get the per sample input mask after dropout # Per sample input mask is in view-major order so that index v*B + b in each mask corresponds to sample b of view v: (B * V) per_sample_geometric_input_mask = torch.rand( batch_size_per_view * num_views, device=device ) < (1 - self.geometric_input_config["dropout_prob"]) per_sample_geometric_input_mask = ( per_sample_geometric_input_mask & overall_geometric_input_mask ) # Get the ray direction input mask per_sample_ray_dirs_input_mask = ( torch.rand(batch_size_per_view, device=device) < self.geometric_input_config["ray_dirs_prob"] ) per_sample_ray_dirs_input_mask = per_sample_ray_dirs_input_mask.repeat( num_views ) per_sample_ray_dirs_input_mask = ( per_sample_ray_dirs_input_mask & per_sample_geometric_input_mask ) # Get the depth input mask per_sample_depth_input_mask = ( torch.rand(batch_size_per_view, device=device) < self.geometric_input_config["depth_prob"] ) per_sample_depth_input_mask = per_sample_depth_input_mask.repeat( num_views ) per_sample_depth_input_mask = ( per_sample_depth_input_mask & per_sample_geometric_input_mask ) # Get the camera input mask per_sample_cam_input_mask = ( torch.rand(batch_size_per_view, device=device) < self.geometric_input_config["cam_prob"] ) per_sample_cam_input_mask = per_sample_cam_input_mask.repeat(num_views) per_sample_cam_input_mask = ( per_sample_cam_input_mask & per_sample_geometric_input_mask ) # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 # Returned pose quats and trans represent identity pose for views/samples where the camera input mask is False ( pose_quats_across_views, pose_trans_across_views, per_sample_cam_input_mask, ) = self._compute_pose_quats_and_trans_for_across_views_in_ref_view( views, num_views, device, dtype, batch_size_per_view, per_sample_cam_input_mask, ) # Encode the ray directions and fuse with the image encoder features all_encoder_features_across_views = self._encode_and_fuse_ray_dirs( views, num_views, batch_size_per_view, all_encoder_features_across_views, per_sample_ray_dirs_input_mask, ) # Encode the depths and fuse with the image encoder features all_encoder_features_across_views = self._encode_and_fuse_depths( views, num_views, batch_size_per_view, all_encoder_features_across_views, per_sample_depth_input_mask, ) # Encode the cam quat and trans and fuse with the image encoder features all_encoder_features_across_views = ( self._encode_and_fuse_cam_quats_and_trans( views, num_views, batch_size_per_view, all_encoder_features_across_views, pose_quats_across_views, pose_trans_across_views, per_sample_cam_input_mask, ) ) # Normalize the fused features (permute -> normalize -> permute) all_encoder_features_across_views = ( all_encoder_features_across_views.permute(0, 2, 3, 1).contiguous() ) all_encoder_features_across_views = self.fusion_norm_layer( all_encoder_features_across_views ) all_encoder_features_across_views = ( all_encoder_features_across_views.permute(0, 3, 1, 2).contiguous() ) # Split the batched views into individual views fused_all_encoder_features_across_views = ( all_encoder_features_across_views.chunk(num_views, dim=0) ) return fused_all_encoder_features_across_views def _compute_adaptive_minibatch_size( self, memory_safety_factor: float = 0.95, ) -> int: """ Compute adaptive minibatch size based on available PyTorch memory. Args: memory_safety_factor: Safety factor to avoid OOM (0.95 = use 95% of available memory) Returns: Computed minibatch size """ device = self.device if device.type == "cuda": # Get available GPU memory torch.cuda.empty_cache() available_memory = torch.cuda.mem_get_info()[ 0 ] # Free memory in bytes usable_memory = ( available_memory * memory_safety_factor ) # Use safety factor to avoid OOM else: # For non-CUDA devices, use conservative default print( "Non-CUDA device detected. Using conservative default minibatch size of 1 for memory efficient dense prediction head inference." ) return 1 # Determine minibatch size based on available memory max_estimated_memory_per_sample = ( 680 * 1024 * 1024 ) # 680 MB per sample (upper bound profiling using a 518 x 518 input) computed_minibatch_size = int( usable_memory / max_estimated_memory_per_sample ) if computed_minibatch_size < 1: computed_minibatch_size = 1 return computed_minibatch_size def downstream_dense_head( self, dense_head_inputs: Tensor | list[Tensor], img_shape: Tuple[int, int], ): """ Run the downstream dense prediction head """ if self.pred_head_type == "linear": dense_head_outputs = self.dense_head( PredictionHeadInput(last_feature=dense_head_inputs) ) dense_final_outputs = self.dense_adaptor( AdaptorInput( adaptor_feature=dense_head_outputs.decoded_channels, output_shape_hw=img_shape, ) ) elif self.pred_head_type in ["dpt", "dpt+pose"]: dense_head_outputs = self.dense_head( PredictionHeadLayeredInput( list_features=dense_head_inputs, target_output_shape=img_shape, ) ) dense_final_outputs = self.dense_adaptor( AdaptorInput( adaptor_feature=dense_head_outputs.decoded_channels, output_shape_hw=img_shape, ) ) else: raise ValueError( f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" ) return dense_final_outputs def downstream_head( self, dense_head_inputs: Tensor | list[Tensor], scale_head_inputs: Tensor, img_shape: Tuple[int, int], memory_efficient_inference: bool = False, ): """Run Prediction Heads & Post-Process Outputs.""" # Get device device = self.device # Use mini-batch inference to run the dense prediction head (the memory bottleneck) # This saves memory and is slower than running the dense prediction head in one go if memory_efficient_inference: # Obtain the batch size of the dense head inputs if self.pred_head_type == "linear": batch_size = dense_head_inputs.shape[0] elif self.pred_head_type in ["dpt", "dpt+pose"]: batch_size = dense_head_inputs[0].shape[0] else: raise ValueError( f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" ) # Compute the mini batch size and number of mini batches adaptively based on available memory minibatch = self._compute_adaptive_minibatch_size() num_batches = (batch_size + minibatch - 1) // minibatch # Run prediction for each mini-batch dense_final_outputs_list = [] pose_final_outputs_list = ( [] if self.pred_head_type == "dpt+pose" else None ) for batch_idx in range(num_batches): start_idx = batch_idx * minibatch end_idx = min((batch_idx + 1) * minibatch, batch_size) # Get the inputs for the current mini-batch if self.pred_head_type == "linear": dense_head_inputs_batch = dense_head_inputs[ start_idx:end_idx ] elif self.pred_head_type in ["dpt", "dpt+pose"]: dense_head_inputs_batch = [ x[start_idx:end_idx] for x in dense_head_inputs ] else: raise ValueError( f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" ) # Dense prediction (mini-batched) dense_final_outputs_batch = self.downstream_dense_head( dense_head_inputs_batch, img_shape ) dense_final_outputs_list.append(dense_final_outputs_batch) # Pose prediction (mini-batched) if self.pred_head_type == "dpt+pose": pose_head_inputs_batch = dense_head_inputs[-1][ start_idx:end_idx ] pose_head_outputs_batch = self.pose_head( PredictionHeadInput( last_feature=pose_head_inputs_batch ) ) pose_final_outputs_batch = self.pose_adaptor( AdaptorInput( adaptor_feature=pose_head_outputs_batch.decoded_channels, output_shape_hw=img_shape, ) ) pose_final_outputs_list.append(pose_final_outputs_batch) # Concatenate the dense prediction head outputs from all mini-batches available_keys = dense_final_outputs_batch.__dict__.keys() dense_pred_data_dict = { key: torch.cat( [ getattr(output, key) for output in dense_final_outputs_list ], dim=0, ) for key in available_keys } dense_final_outputs = dense_final_outputs_batch.__class__( **dense_pred_data_dict ) # Concatenate the pose prediction head outputs from all mini-batches pose_final_outputs = None if self.pred_head_type == "dpt+pose": available_keys = pose_final_outputs_batch.__dict__.keys() pose_pred_data_dict = { key: torch.cat( [ getattr(output, key) for output in pose_final_outputs_list ], dim=0, ) for key in available_keys } pose_final_outputs = pose_final_outputs_batch.__class__( **pose_pred_data_dict ) # Clear CUDA cache for better memory efficiency if device.type == "cuda": torch.cuda.empty_cache() else: # Run prediction for all (batch_size * num_views) in one go # Dense prediction dense_final_outputs = self.downstream_dense_head( dense_head_inputs, img_shape ) # Pose prediction pose_final_outputs = None if self.pred_head_type == "dpt+pose": pose_head_outputs = self.pose_head( PredictionHeadInput(last_feature=dense_head_inputs[-1]) ) pose_final_outputs = self.pose_adaptor( AdaptorInput( adaptor_feature=pose_head_outputs.decoded_channels, output_shape_hw=img_shape, ) ) # Scale prediction is lightweight, so we can run it in one go scale_head_output = self.scale_head( PredictionHeadTokenInput(last_feature=scale_head_inputs) ) scale_final_output = self.scale_adaptor( AdaptorInput( adaptor_feature=scale_head_output.decoded_channels, output_shape_hw=img_shape, ) ) scale_final_output = scale_final_output.value.squeeze( -1 ) # (B, 1, 1) -> (B, 1) # Clear CUDA cache for better memory efficiency if memory_efficient_inference and device.type == "cuda": torch.cuda.empty_cache() return dense_final_outputs, pose_final_outputs, scale_final_output def forward(self, views, memory_efficient_inference: bool = False): """Forward pass. Process: 1. Encodes the N input views (images). 2. Encodes the optional geometric inputs (ray directions, depths, camera rotations, camera translations). 3. Fuses the encoded features from the N input views and the optional geometric inputs using addition and normalization. 4. Information sharing across the encoded features and a scale token using a multi-view attention transformer. 5. Passes the final features from transformer through the prediction heads. 6. Returns the processed final outputs for N views. Assumption: - All the input views and dense geometric inputs have the same image shape. Args: views (List[dict]): List of dictionaries containing the input views' images and instance information. Each dictionary should contain the following keys: "img" (tensor): Image tensor of shape (B, C, H, W). Input images must be normalized based on the data norm type of image encoder. "data_norm_type" (list): [model.encoder.data_norm_type] Optionally, each dictionary can also contain the following keys for the respective optional geometric inputs: "ray_directions_cam" (tensor): Ray directions in the local camera frame. Tensor of shape (B, H, W, 3). "depth_along_ray" (tensor): Depth along the ray. Tensor of shape (B, H, W, 1). "camera_pose_quats" (tensor): Camera pose quaternions. Tensor of shape (B, 4). Camera pose is opencv (RDF) cam2world transformation. "camera_pose_trans" (tensor): Camera pose translations. Tensor of shape (B, 3). Camera pose is opencv (RDF) cam2world transformation. "is_metric_scale" (tensor): Boolean tensor indicating whether the geometric inputs are in metric scale or not. Tensor of shape (B, 1). memory_efficient_inference (bool): Whether to use memory efficient inference or not. This runs the dense prediction head (the memory bottleneck) in a memory efficient manner. Default is False. Returns: List[dict]: A list containing the final outputs for all N views. """ # Get input shape of the images, number of views, and batch size per view batch_size_per_view, _, height, width = views[0]["img"].shape img_shape = (int(height), int(width)) num_views = len(views) # Run the image encoder on all the input views ( all_encoder_features_across_views, all_encoder_registers_across_views, ) = self._encode_n_views(views) # Encode the optional geometric inputs and fuse with the encoded features from the N input views # Use high precision to prevent NaN values after layer norm in dense representation encoder (due to high variance in last dim of features) with torch.autocast("cuda", enabled=False): all_encoder_features_across_views = ( self._encode_and_fuse_optional_geometric_inputs( views, all_encoder_features_across_views ) ) # Expand the scale token to match the batch size input_scale_token = ( self.scale_token.unsqueeze(0) .unsqueeze(-1) .repeat(batch_size_per_view, 1, 1) ) # (B, C, 1) # Combine all images into view-centric representation # Output is a list containing the encoded features for all N views after information sharing. info_sharing_input = MultiViewTransformerInput( features=all_encoder_features_across_views, additional_input_tokens_per_view=all_encoder_registers_across_views, additional_input_tokens=input_scale_token, ) if self.info_sharing_return_type == "no_intermediate_features": final_info_sharing_multi_view_feat = self.info_sharing( info_sharing_input ) elif self.info_sharing_return_type == "intermediate_features": ( final_info_sharing_multi_view_feat, intermediate_info_sharing_multi_view_feat, ) = self.info_sharing(info_sharing_input) if self.pred_head_type == "linear": # Stack the features for all views dense_head_inputs = torch.cat( final_info_sharing_multi_view_feat.features, dim=0 ) elif self.pred_head_type in ["dpt", "dpt+pose"]: # Get the list of features for all views dense_head_inputs_list = [] if self.use_encoder_features_for_dpt: # Stack all the image encoder features for all views stacked_encoder_features = torch.cat( all_encoder_features_across_views, dim=0 ) dense_head_inputs_list.append(stacked_encoder_features) # Stack the first intermediate features for all views stacked_intermediate_features_1 = torch.cat( intermediate_info_sharing_multi_view_feat[0].features, dim=0, ) dense_head_inputs_list.append(stacked_intermediate_features_1) # Stack the second intermediate features for all views stacked_intermediate_features_2 = torch.cat( intermediate_info_sharing_multi_view_feat[1].features, dim=0, ) dense_head_inputs_list.append(stacked_intermediate_features_2) # Stack the last layer features for all views stacked_final_features = torch.cat( final_info_sharing_multi_view_feat.features, dim=0 ) dense_head_inputs_list.append(stacked_final_features) else: # Stack the first intermediate features for all views stacked_intermediate_features_1 = torch.cat( intermediate_info_sharing_multi_view_feat[0].features, dim=0, ) dense_head_inputs_list.append(stacked_intermediate_features_1) # Stack the second intermediate features for all views stacked_intermediate_features_2 = torch.cat( intermediate_info_sharing_multi_view_feat[1].features, dim=0, ) dense_head_inputs_list.append(stacked_intermediate_features_2) # Stack the third intermediate features for all views stacked_intermediate_features_3 = torch.cat( intermediate_info_sharing_multi_view_feat[2].features, dim=0, ) dense_head_inputs_list.append(stacked_intermediate_features_3) # Stack the last layer stacked_final_features = torch.cat( final_info_sharing_multi_view_feat.features, dim=0 ) dense_head_inputs_list.append(stacked_final_features) else: raise ValueError( f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" ) with torch.autocast("cuda", enabled=False): # Prepare inputs for the downstream heads if self.pred_head_type == "linear": dense_head_inputs = dense_head_inputs elif self.pred_head_type in ["dpt", "dpt+pose"]: dense_head_inputs = dense_head_inputs_list scale_head_inputs = ( final_info_sharing_multi_view_feat.additional_token_features ) # Run the downstream heads dense_final_outputs, pose_final_outputs, scale_final_output = ( self.downstream_head( dense_head_inputs=dense_head_inputs, scale_head_inputs=scale_head_inputs, img_shape=img_shape, memory_efficient_inference=memory_efficient_inference, ) ) # Prepare the final scene representation for all views if self.scene_rep_type in [ "pointmap", "pointmap+confidence", "pointmap+mask", "pointmap+confidence+mask", ]: output_pts3d = dense_final_outputs.value # Reshape final scene representation to (B * V, H, W, C) output_pts3d = output_pts3d.permute(0, 2, 3, 1).contiguous() # Split the predicted pointmaps back to their respective views output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) # Pack the output as a list of dictionaries res = [] for i in range(num_views): res.append( { "pts3d": output_pts3d_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "metric_scaling_factor": scale_final_output, } ) elif self.scene_rep_type in [ "raymap+depth", "raymap+depth+confidence", "raymap+depth+mask", "raymap+depth+confidence+mask", ]: # Reshape final scene representation to (B * V, H, W, C) output_scene_rep = dense_final_outputs.value.permute( 0, 2, 3, 1 ).contiguous() # Get the predicted ray origins, directions, and depths along rays ( output_ray_origins, output_ray_directions, output_depth_along_ray, ) = output_scene_rep.split([3, 3, 1], dim=-1) # Get the predicted pointmaps output_pts3d = ( output_ray_origins + output_ray_directions * output_depth_along_ray ) # Split the predicted quantities back to their respective views output_ray_origins_per_view = output_ray_origins.chunk( num_views, dim=0 ) output_ray_directions_per_view = output_ray_directions.chunk( num_views, dim=0 ) output_depth_along_ray_per_view = output_depth_along_ray.chunk( num_views, dim=0 ) output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) # Pack the output as a list of dictionaries res = [] for i in range(num_views): res.append( { "pts3d": output_pts3d_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "ray_origins": output_ray_origins_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "ray_directions": output_ray_directions_per_view[ i ], "depth_along_ray": output_depth_along_ray_per_view[ i ] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "metric_scaling_factor": scale_final_output, } ) elif self.scene_rep_type in [ "raydirs+depth+pose", "raydirs+depth+pose+confidence", "raydirs+depth+pose+mask", "raydirs+depth+pose+confidence+mask", ]: # Reshape output dense rep to (B * V, H, W, C) output_dense_rep = dense_final_outputs.value.permute( 0, 2, 3, 1 ).contiguous() # Get the predicted ray directions and depths along rays output_ray_directions, output_depth_along_ray = ( output_dense_rep.split([3, 1], dim=-1) ) # Get the predicted camera translations and quaternions output_cam_translations, output_cam_quats = ( pose_final_outputs.value.split([3, 4], dim=-1) ) # Get the predicted pointmaps in world frame and camera frame output_pts3d = convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( output_ray_directions, output_depth_along_ray, output_cam_translations, output_cam_quats, ) output_pts3d_cam = ( output_ray_directions * output_depth_along_ray ) # Split the predicted quantities back to their respective views output_ray_directions_per_view = output_ray_directions.chunk( num_views, dim=0 ) output_depth_along_ray_per_view = output_depth_along_ray.chunk( num_views, dim=0 ) output_cam_translations_per_view = ( output_cam_translations.chunk(num_views, dim=0) ) output_cam_quats_per_view = output_cam_quats.chunk( num_views, dim=0 ) output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) output_pts3d_cam_per_view = output_pts3d_cam.chunk( num_views, dim=0 ) # Pack the output as a list of dictionaries res = [] for i in range(num_views): res.append( { "pts3d": output_pts3d_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "pts3d_cam": output_pts3d_cam_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "ray_directions": output_ray_directions_per_view[ i ], "depth_along_ray": output_depth_along_ray_per_view[ i ] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "cam_trans": output_cam_translations_per_view[i] * scale_final_output, "cam_quats": output_cam_quats_per_view[i], "metric_scaling_factor": scale_final_output, } ) elif self.scene_rep_type in [ "campointmap+pose", "campointmap+pose+confidence", "campointmap+pose+mask", "campointmap+pose+confidence+mask", ]: # Get the predicted camera frame pointmaps output_pts3d_cam = dense_final_outputs.value # Reshape final scene representation to (B * V, H, W, C) output_pts3d_cam = output_pts3d_cam.permute( 0, 2, 3, 1 ).contiguous() # Get the predicted camera translations and quaternions output_cam_translations, output_cam_quats = ( pose_final_outputs.value.split([3, 4], dim=-1) ) # Get the ray directions and depths along rays output_depth_along_ray = torch.norm( output_pts3d_cam, dim=-1, keepdim=True ) output_ray_directions = ( output_pts3d_cam / output_depth_along_ray ) # Get the predicted pointmaps in world frame output_pts3d = convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( output_ray_directions, output_depth_along_ray, output_cam_translations, output_cam_quats, ) # Split the predicted quantities back to their respective views output_ray_directions_per_view = output_ray_directions.chunk( num_views, dim=0 ) output_depth_along_ray_per_view = output_depth_along_ray.chunk( num_views, dim=0 ) output_cam_translations_per_view = ( output_cam_translations.chunk(num_views, dim=0) ) output_cam_quats_per_view = output_cam_quats.chunk( num_views, dim=0 ) output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) output_pts3d_cam_per_view = output_pts3d_cam.chunk( num_views, dim=0 ) # Pack the output as a list of dictionaries res = [] for i in range(num_views): res.append( { "pts3d": output_pts3d_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "pts3d_cam": output_pts3d_cam_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "ray_directions": output_ray_directions_per_view[ i ], "depth_along_ray": output_depth_along_ray_per_view[ i ] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "cam_trans": output_cam_translations_per_view[i] * scale_final_output, "cam_quats": output_cam_quats_per_view[i], "metric_scaling_factor": scale_final_output, } ) elif self.scene_rep_type in [ "pointmap+raydirs+depth+pose", "pointmap+raydirs+depth+pose+confidence", "pointmap+raydirs+depth+pose+mask", "pointmap+raydirs+depth+pose+confidence+mask", ]: # Reshape final scene representation to (B * V, H, W, C) output_dense_rep = dense_final_outputs.value.permute( 0, 2, 3, 1 ).contiguous() # Get the predicted pointmaps, ray directions and depths along rays output_pts3d, output_ray_directions, output_depth_along_ray = ( output_dense_rep.split([3, 3, 1], dim=-1) ) # Get the predicted camera translations and quaternions output_cam_translations, output_cam_quats = ( pose_final_outputs.value.split([3, 4], dim=-1) ) # Get the predicted pointmaps in camera frame output_pts3d_cam = ( output_ray_directions * output_depth_along_ray ) # Replace the predicted world-frame pointmaps if required if self.pred_head_config["adaptor_config"][ "use_factored_predictions_for_global_pointmaps" ]: output_pts3d = convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( output_ray_directions, output_depth_along_ray, output_cam_translations, output_cam_quats, ) # Split the predicted quantities back to their respective views output_ray_directions_per_view = output_ray_directions.chunk( num_views, dim=0 ) output_depth_along_ray_per_view = output_depth_along_ray.chunk( num_views, dim=0 ) output_cam_translations_per_view = ( output_cam_translations.chunk(num_views, dim=0) ) output_cam_quats_per_view = output_cam_quats.chunk( num_views, dim=0 ) output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) output_pts3d_cam_per_view = output_pts3d_cam.chunk( num_views, dim=0 ) # Pack the output as a list of dictionaries res = [] for i in range(num_views): res.append( { "pts3d": output_pts3d_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "pts3d_cam": output_pts3d_cam_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "ray_directions": output_ray_directions_per_view[ i ], "depth_along_ray": output_depth_along_ray_per_view[ i ] * scale_final_output.unsqueeze(-1).unsqueeze(-1), "cam_trans": output_cam_translations_per_view[i] * scale_final_output, "cam_quats": output_cam_quats_per_view[i], "metric_scaling_factor": scale_final_output, } ) else: raise ValueError( f"Invalid scene_rep_type: {self.scene_rep_type}. \ Valid options: ['pointmap', 'raymap+depth', 'raydirs+depth+pose', 'campointmap+pose', 'pointmap+raydirs+depth+pose' \ 'pointmap+confidence', 'raymap+depth+confidence', 'raydirs+depth+pose+confidence', 'campointmap+pose+confidence', 'pointmap+raydirs+depth+pose+confidence' \ 'pointmap+mask', 'raymap+depth+mask', 'raydirs+depth+pose+mask', 'campointmap+pose+mask', 'pointmap+raydirs+depth+pose+mask' \ 'pointmap+confidence+mask', 'raymap+depth+confidence+mask', 'raydirs+depth+pose+confidence+mask', 'campointmap+pose+confidence+mask', 'pointmap+raydirs+depth+pose+confidence+mask']" ) # Get the output confidences for all views (if available) and add them to the result if "confidence" in self.scene_rep_type: output_confidences = dense_final_outputs.confidence # Reshape confidences to (B * V, H, W) output_confidences = ( output_confidences.permute(0, 2, 3, 1) .squeeze(-1) .contiguous() ) # Split the predicted confidences back to their respective views output_confidences_per_view = output_confidences.chunk( num_views, dim=0 ) # Add the confidences to the result for i in range(num_views): res[i]["conf"] = output_confidences_per_view[i] # Get the output masks (and logits) for all views (if available) and add them to the result if "mask" in self.scene_rep_type: # Get the output masks output_masks = dense_final_outputs.mask # Reshape masks to (B * V, H, W) output_masks = ( output_masks.permute(0, 2, 3, 1).squeeze(-1).contiguous() ) # Threshold the masks at 0.5 to get binary masks (0: ambiguous, 1: non-ambiguous) output_masks = output_masks > 0.5 # Split the predicted masks back to their respective views output_masks_per_view = output_masks.chunk(num_views, dim=0) # Get the output mask logits (for loss) output_mask_logits = dense_final_outputs.logits # Reshape mask logits to (B * V, H, W) output_mask_logits = ( output_mask_logits.permute(0, 2, 3, 1) .squeeze(-1) .contiguous() ) # Split the predicted mask logits back to their respective views output_mask_logits_per_view = output_mask_logits.chunk( num_views, dim=0 ) # Add the masks and logits to the result for i in range(num_views): res[i]["non_ambiguous_mask"] = output_masks_per_view[i] res[i]["non_ambiguous_mask_logits"] = ( output_mask_logits_per_view[i] ) return res def _configure_geometric_input_config( self, use_calibration: bool, use_depth: bool, use_pose: bool, use_depth_scale: bool, use_pose_scale: bool, ): """ Configure the geometric input configuration """ # Store original config for restoration if not hasattr(self, "_original_geometric_config"): self._original_geometric_config = dict(self.geometric_input_config) # Set the geometric input configuration if not (use_calibration or use_depth or use_pose): # No geometric inputs (images-only mode) self.geometric_input_config.update( { "overall_prob": 0.0, "dropout_prob": 1.0, "ray_dirs_prob": 0.0, "depth_prob": 0.0, "cam_prob": 0.0, "sparse_depth_prob": 0.0, "depth_scale_norm_all_prob": 0.0, "pose_scale_norm_all_prob": 0.0, } ) else: # Enable geometric inputs with deterministic behavior self.geometric_input_config.update( { "overall_prob": 1.0, "dropout_prob": 0.0, "ray_dirs_prob": 1.0 if use_calibration else 0.0, "depth_prob": 1.0 if use_depth else 0.0, "cam_prob": 1.0 if use_pose else 0.0, "sparse_depth_prob": 0.0, "depth_scale_norm_all_prob": ( 0.0 if use_depth_scale else 1.0 ), "pose_scale_norm_all_prob": 0.0 if use_pose_scale else 1.0, } ) def _restore_original_geometric_input_config(self): """ Restore original geometric input configuration """ if hasattr(self, "_original_geometric_config"): self.geometric_input_config.update(self._original_geometric_config) @torch.inference_mode() def infer( self, views: List[Dict[str, Any]], memory_efficient_inference: bool = False, use_amp: bool = True, amp_dtype: str = "bf16", apply_mask: bool = True, mask_edges: bool = True, edge_normal_threshold: float = 5.0, edge_depth_threshold: float = 0.03, apply_confidence_mask: bool = False, confidence_percentile: float = 10, ignore_calibration_inputs: bool = False, ignore_depth_inputs: bool = False, ignore_pose_inputs: bool = False, ignore_depth_scale_inputs: bool = False, ignore_pose_scale_inputs: bool = False, ) -> List[Dict[str, Tensor]]: """ User-friendly inference with strict input validation and automatic conversion. Args: views: List of view dictionaries. Each dict can contain: Required: - 'img': Tensor of shape (B, 3, H, W) - normalized RGB images - 'data_norm_type': str - normalization type used to normalize the images (must be equal to self.model.encoder.data_norm_type) Optional Geometric Inputs (only one of intrinsics OR ray_directions): - 'intrinsics': Tensor of shape (B, 3, 3) - will be converted to ray directions - 'ray_directions': Tensor of shape (B, H, W, 3) - ray directions in camera frame - 'depth_z': Tensor of shape (B, H, W, 1) - Z depth in camera frame (intrinsics or ray_directions must be provided) - 'camera_poses': Tensor of shape (B, 4, 4) or tuple of (quats - (B, 4), trans - (B, 3)) - can be any world frame - 'is_metric_scale': bool or Tensor of shape (B,) - if not provided, defaults to True Optional Additional Info: - 'instance': List[str] where length of list is B - instance info for each view - 'idx': List[int] where length of list is B - index info for each view - 'true_shape': List[tuple] where length of list is B - true shape info (H, W) for each view memory_efficient_inference: Whether to use memory-efficient inference for dense prediction heads (trades off speed). Defaults to False. use_amp: Whether to use automatic mixed precision for faster inference. Defaults to True. amp_dtype: The dtype to use for mixed precision. Defaults to "bf16" (bfloat16). Options: "fp16", "bf16", "fp32". apply_mask: Whether to apply the non-ambiguous mask to the output. Defaults to True. mask_edges: Whether to compute an edge mask based on normals and depth and apply it to the output. Defaults to True. edge_normal_threshold: Tolerance threshold for normals-based edge detection. Defaults to 5.0. edge_depth_threshold: Relative tolerance threshold for depth-based edge detection. Defaults to 0.03. apply_confidence_mask: Whether to apply the confidence mask to the output. Defaults to False. confidence_percentile: The percentile to use for the confidence threshold. Defaults to 10. ignore_calibration_inputs: Whether to ignore the calibration inputs (intrinsics and ray_directions). Defaults to False. ignore_depth_inputs: Whether to ignore the depth inputs. Defaults to False. ignore_pose_inputs: Whether to ignore the pose inputs. Defaults to False. ignore_depth_scale_inputs: Whether to ignore the depth scale inputs. Defaults to False. ignore_pose_scale_inputs: Whether to ignore the pose scale inputs. Defaults to False. IMPORTANT CONSTRAINTS: - Cannot provide both 'intrinsics' and 'ray_directions' (they represent the same information) - If 'depth' is provided, then 'intrinsics' or 'ray_directions' must also be provided - If ANY view has 'camera_poses', then view 0 (first view) MUST also have 'camera_poses' Returns: List of prediction dictionaries, one per view. Each dict contains: - 'img_no_norm': Tensor of shape (B, H, W, 3) - denormalized rgb images - 'pts3d': Tensor of shape (B, H, W, 3) - predicted points in world frame - 'pts3d_cam': Tensor of shape (B, H, W, 3) - predicted points in camera frame - 'ray_directions': Tensor of shape (B, H, W, 3) - ray directions in camera frame - 'intrinsics': Tensor of shape (B, 3, 3) - pinhole camera intrinsics recovered from ray directions - 'depth_along_ray': Tensor of shape (B, H, W, 1) - depth along ray in camera frame - 'depth_z': Tensor of shape (B, H, W, 1) - Z depth in camera frame - 'cam_trans': Tensor of shape (B, 3) - camera translation in world frame - 'cam_quats': Tensor of shape (B, 4) - camera quaternion in world frame - 'camera_poses': Tensor of shape (B, 4, 4) - camera pose in world frame - 'metric_scaling_factor': Tensor of shape (B,) - applied metric scaling factor - 'mask': Tensor of shape (B, H, W, 1) - combo of non-ambiguous mask, edge mask and confidence-based mask if used - 'non_ambiguous_mask': Tensor of shape (B, H, W) - non-ambiguous mask - 'non_ambiguous_mask_logits': Tensor of shape (B, H, W) - non-ambiguous mask logits - 'conf': Tensor of shape (B, H, W) - confidence Raises: ValueError: For invalid inputs, missing required keys, conflicting modalities, or constraint violations """ # Determine the mixed precision floating point type if use_amp: if amp_dtype == "fp16": amp_dtype = torch.float16 elif amp_dtype == "bf16": if torch.cuda.is_bf16_supported(): amp_dtype = torch.bfloat16 else: warnings.warn( "bf16 is not supported on this device. Using fp16 instead." ) amp_dtype = torch.float16 elif amp_dtype == "fp32": amp_dtype = torch.float32 else: amp_dtype = torch.float32 # Validate the input views validated_views = validate_input_views_for_inference(views) # Transfer the views to the same device as the model ignore_keys = set( [ "instance", "idx", "true_shape", "data_norm_type", ] ) for view in validated_views: for name in view.keys(): if name in ignore_keys: continue view[name] = view[name].to(self.device, non_blocking=True) # Pre-process the input views processed_views = preprocess_input_views_for_inference(validated_views) # Set the model input probabilities based on input args for ignoring inputs self._configure_geometric_input_config( use_calibration=not ignore_calibration_inputs, use_depth=not ignore_depth_inputs, use_pose=not ignore_pose_inputs, use_depth_scale=not ignore_depth_scale_inputs, use_pose_scale=not ignore_pose_scale_inputs, ) # Run the model with torch.autocast("cuda", enabled=bool(use_amp), dtype=amp_dtype): preds = self.forward( processed_views, memory_efficient_inference=memory_efficient_inference, ) # Post-process the model outputs preds = postprocess_model_outputs_for_inference( raw_outputs=preds, input_views=processed_views, apply_mask=apply_mask, mask_edges=mask_edges, edge_normal_threshold=edge_normal_threshold, edge_depth_threshold=edge_depth_threshold, apply_confidence_mask=apply_confidence_mask, confidence_percentile=confidence_percentile, ) # Restore the original configuration self._restore_original_geometric_input_config() return preds