# import sys # sys.path.append("..") import torch from torch import nn import torch.nn.init as init import torch.nn.functional as F from paths import * from typing import Dict, List, Optional, Set, Tuple, Union import os from contextlib import nullcontext from vggt.models.vggt import VGGT from vggt.utils.pose_enc import pose_encoding_to_extri_intri from vggt.layers import Mlp from vggt.layers.block import Block from vggt.heads.head_act import activate_pose class OriAny_CameraHead(nn.Module): """ CameraHead predicts camera parameters from token representations using iterative refinement. It applies a series of transformer blocks (the "trunk") to dedicated camera tokens. """ def __init__( self, dim_in: int = 2048, trunk_depth: int = 4, pose_encoding_type: str = "OriAny", num_heads: int = 16, mlp_ratio: int = 4, init_values: float = 0.01, ): super().__init__() if pose_encoding_type == "OriAny": self.target_dim = 360+180+360+2 else: raise ValueError(f"Unsupported camera encoding type: {pose_encoding_type}") self.trunk_depth = trunk_depth # Build the trunk using a sequence of transformer blocks. self.trunk = nn.Sequential( *[ Block( dim=dim_in, num_heads=num_heads, mlp_ratio=mlp_ratio, init_values=init_values, ) for _ in range(trunk_depth) ] ) # Normalizations for camera token and trunk output. self.token_norm = nn.LayerNorm(dim_in) self.trunk_norm = nn.LayerNorm(dim_in) # Learnable empty camera pose token. self.empty_pose_tokens = nn.Parameter(torch.zeros(1, 1, self.target_dim)) self.embed_pose = nn.Linear(self.target_dim, dim_in) # Module for producing modulation parameters: shift, scale, and a gate. self.poseLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim_in, 3 * dim_in, bias=True)) # Adaptive layer normalization without affine parameters. self.adaln_norm = nn.LayerNorm(dim_in, elementwise_affine=False, eps=1e-6) self.pose_branch = Mlp( in_features=dim_in, hidden_features=dim_in // 2, out_features=self.target_dim, drop=0, ) def forward(self, aggregated_tokens_list: list, num_iterations: int = 4) -> list: """ Forward pass to predict camera parameters. Args: aggregated_tokens_list (list): List of token tensors from the network; the last tensor is used for prediction. num_iterations (int, optional): Number of iterative refinement steps. Defaults to 4. Returns: list: A list of predicted camera encodings (post-activation) from each iteration. """ # Use tokens from the last block for camera prediction. tokens = aggregated_tokens_list[-1] # Extract the camera tokens pose_tokens = tokens[:, :, 0] pose_tokens = self.token_norm(pose_tokens) pred_pose_enc_list = self.trunk_fn(pose_tokens, num_iterations) return pred_pose_enc_list def trunk_fn(self, pose_tokens: torch.Tensor, num_iterations: int) -> list: """ Iteratively refine camera pose predictions. Args: pose_tokens (torch.Tensor): Normalized camera tokens with shape [B, 1, C]. num_iterations (int): Number of refinement iterations. Returns: list: List of activated camera encodings from each iteration. """ B, S, C = pose_tokens.shape # S is expected to be 1. pred_pose_enc = None pred_pose_enc_list = [] for _ in range(num_iterations): # Use a learned empty pose for the first iteration. if pred_pose_enc is None: module_input = self.embed_pose(self.empty_pose_tokens.expand(B, S, -1)) else: # Detach the previous prediction to avoid backprop through time. pred_pose_enc = pred_pose_enc.detach() module_input = self.embed_pose(pred_pose_enc) # Generate modulation parameters and split them into shift, scale, and gate components. shift_msa, scale_msa, gate_msa = self.poseLN_modulation(module_input).chunk(3, dim=-1) # Adaptive layer normalization and modulation. pose_tokens_modulated = gate_msa * modulate(self.adaln_norm(pose_tokens), shift_msa, scale_msa) pose_tokens_modulated = pose_tokens_modulated + pose_tokens pose_tokens_modulated = self.trunk(pose_tokens_modulated) # Compute the delta update for the pose encoding. pred_pose_enc_delta = self.pose_branch(self.trunk_norm(pose_tokens_modulated)) if pred_pose_enc is None: pred_pose_enc = pred_pose_enc_delta else: pred_pose_enc = pred_pose_enc + pred_pose_enc_delta # Apply final activation functions for translation, quaternion, and field-of-view. # activated_pose = activate_pose( # pred_pose_enc, # trans_act=self.trans_act, # quat_act=self.quat_act, # fl_act=self.fl_act, # ) # pred_pose_enc_list.append(activated_pose) pred_pose_enc_list.append(pred_pose_enc) return pred_pose_enc_list def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: """ Modulate the input tensor using scaling and shifting parameters. """ # modified from https://github.com/facebookresearch/DiT/blob/796c29e532f47bba17c5b9c5eb39b9354b8b7c64/models.py#L19 return x * (1 + scale) + shift def load_patch_embed_weights(model, checkpoint_path): # 1. 加载 checkpoint checkpoint = torch.load(checkpoint_path, map_location="cpu") # 2. 获取 state_dict state_dict = checkpoint.get("state_dict", checkpoint) # 3. 过滤只包含 aggregator.patch_embed 的参数 patch_embed_state = { k.replace("aggregator.patch_embed.", ""): v for k, v in state_dict.items() if k.startswith("aggregator.patch_embed.") } # 4. 加载到目标模块 missing_keys, unexpected_keys = model.aggregator.patch_embed.load_state_dict( patch_embed_state, strict=False ) print("Loaded patch_embed weights.") print("Missing keys:", missing_keys) print("Unexpected keys:", unexpected_keys) class VGGT_OriAny_Ref(nn.Module): def __init__(self, dtype, out_dim, nopretrain ) -> None: super().__init__() self.vggt = VGGT() self.dtype = dtype self.ref_sampler = MLP_dim(in_dim=2048, out_dim=out_dim) self.ref_sampler.apply(init_weights) self.tgt_sampler = MLP_dim(in_dim=2048, out_dim=out_dim) self.tgt_sampler.apply(init_weights) def forward(self, img_inputs): device = self.get_device() with torch.amp.autocast(device_type='cuda', dtype=self.dtype): if img_inputs.shape == 4: img_inputs = img_inputs[None] aggregated_tokens_list, ps_idx = self.vggt.aggregator(img_inputs) # Predict Cameras # pose_enc = self.oriany_camera_head(aggregated_tokens_list)[-1] # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world) # extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, images.shape[-2:]) # Use tokens from the last block for camera prediction. tokens = aggregated_tokens_list[-1] # Extract the camera tokens pose_tokens = tokens[:, :, 0] # tokens = aggregated_tokens_list[-1] B, S, C = pose_tokens.shape if S>1: # 分离每个 batch 的第一个 token 和其余 token ref_tokens = pose_tokens[:, 0, :] # shape: (B, C) tgt_tokens = pose_tokens[:, 1:, :] # shape: (B, S-1, C) # 下采样 ref_feat = self.ref_sampler(ref_tokens) # shape: (B, C')，假设输出 channel 为 C' tgt_feat = self.tgt_sampler(tgt_tokens.reshape(B * (S - 1), C)) # shape: (B*(S-1), C') # 合并结果 pose_enc = torch.cat([ ref_feat.unsqueeze(1), # (B, 1, C') tgt_feat.view(B, S - 1, -1) # (B, S-1, C') ], dim=1) # 最终 shape: (B*S, C') else: pose_enc = self.ref_sampler(pose_tokens.view(B*S,C)) return pose_enc def get_device(self): return next(self.parameters()).device def init_weights(m): if isinstance(m, nn.Linear): init.xavier_uniform_(m.weight) if m.bias is not None: init.constant_(m.bias, 0) def get_activation(activation): if activation.lower() == 'gelu': return nn.GELU() elif activation.lower() == 'rrelu': return nn.RReLU(inplace=True) elif activation.lower() == 'selu': return nn.SELU(inplace=True) elif activation.lower() == 'silu': return nn.SiLU(inplace=True) elif activation.lower() == 'hardswish': return nn.Hardswish(inplace=True) elif activation.lower() == 'leakyrelu': return nn.LeakyReLU(inplace=True) elif activation.lower() == 'sigmoid': return nn.Sigmoid() elif activation.lower() == 'tanh': return nn.Tanh() else: return nn.ReLU(inplace=True) class MLP_dim(nn.Module): def __init__( self, in_dim=512, out_dim=1024, bias=True, activation='relu'): super().__init__() self.act = get_activation(activation) self.net1 = nn.Sequential( nn.Linear(in_dim, int(out_dim), bias=bias), nn.BatchNorm1d(int(out_dim)), self.act ) self.net2 = nn.Sequential( nn.Linear(int(out_dim), out_dim, bias=bias), nn.BatchNorm1d(out_dim) ) def forward(self, x): return self.net2(self.net1(x))