Spaces:
Running on Zero
Running on Zero
File size: 10,233 Bytes
510e990 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | import torch
import torch.nn as nn
from typing import Optional, Dict
from huggingface_hub import PyTorchModelHubMixin
# Import model components
from argus.models.aggregator import Aggregator
from argus.heads.camera_head import CameraHead
from argus.heads.dpt_head import DPTHead
from argus.heads.utils import reorder_by_reference
class Argus(nn.Module, PyTorchModelHubMixin):
"""
Argus multi-task vision model for camera pose estimation, depth prediction, and 3D points.
Integrates an aggregator backbone with task-specific heads for:
- Camera pose encoding
- Depth map prediction
- 3D camera/rotated/world point prediction
Args:
img_size: Input image size (height/width, assumes square) (default: 518)
patch_size: Patch size for vision transformer backbone (default: 14)
embed_dim: Embedding dimension for transformer features (default: 1024)
enable_camera: Enable camera pose estimation head (default: True)
enable_depth: Enable depth prediction head (default: True)
enable_cam_point: Enable camera coordinate 3D point prediction head (default: False)
enable_rotated_point: Enable rotated 3D point prediction head (default: False)
enable_point: Enable world coordinate 3D point prediction head (default: False, Please do not set it to True during training)
Note:
All heads share the same aggregated transformer features from the Aggregator backbone.
Each DPT-based head outputs both predictions and confidence scores.
"""
def __init__(
self,
img_size: int = 518,
patch_size: int = 14,
embed_dim: int = 1024,
enable_camera: bool = True,
enable_depth: bool = True,
enable_cam_point: bool = False,
enable_rotated_point: bool = False,
enable_point: bool = False,
reorder_by_learning_ref: bool = True,
restore_metric_scale: bool = False
) -> None:
super().__init__()
# For inference
self.restore_metric_scale = restore_metric_scale
self.reorder_by_learning_ref = reorder_by_learning_ref
# Backbone and geometry transformer
self.aggregator = Aggregator(
img_size=img_size,
patch_size=patch_size,
embed_dim=embed_dim,
reorder_by_learning_ref=reorder_by_learning_ref,
)
# Task-specific prediction heads (lazy initialization based on flags)
self.camera_head: Optional[CameraHead] = CameraHead(dim_in=2 * embed_dim) if enable_camera else None
self.depth_head: Optional[DPTHead] = DPTHead(
dim_in=2 * embed_dim,
output_dim=2,
activation="exp",
conf_activation="expp1"
) if enable_depth else None
# 3D point prediction heads (shared architecture, different output semantics)
self.cam_point_head: Optional[DPTHead] = DPTHead(
dim_in=2 * embed_dim,
output_dim=4,
activation="inv_log",
conf_activation="expp1"
) if enable_cam_point else None
self.rotated_point_head: Optional[DPTHead] = DPTHead(
dim_in=2 * embed_dim,
output_dim=4,
activation="inv_log",
conf_activation="expp1"
) if enable_rotated_point else None
self.point_head: Optional[DPTHead] = DPTHead(
dim_in=2 * embed_dim,
output_dim=4,
activation="inv_log",
conf_activation="expp1"
) if enable_point else None
def forward(
self,
images: torch.Tensor,
) -> Dict[str, torch.Tensor]:
"""
Forward pass of the Argus model.
Automatically adds batch dimension if missing and processes multi-task predictions.
Args:
images: Input RGB images with shape:
- [S, 3, H, W] (sequence without batch) or
- [B, S, 3, H, W] (batch of sequences)
Values in range [0, 1], where:
- B: batch size
- S: sequence length (number of frames)
- 3: RGB channels
- H/W: image height/width (matches img_size)
Returns:
Dictionary of model predictions with task-specific outputs:
Common outputs:
- covisibility_scores: Covisibility scores from aggregator (shape varies)
- ref_idx: Reference frame indices (shape varies)
Camera head outputs (if enabled):
- pose_enc: Final camera pose encoding [B, S, 9]
- pose_enc_list: List of pose encodings from all iterations [List[torch.Tensor]]
Depth head outputs (if enabled):
- depth: Predicted depth maps [B, S, H, W, 1]
- depth_conf: Depth prediction confidence [B, S, H, W]
Camera point head outputs (if enabled):
- cam_points: 3D camera coordinates per pixel [B, S, H, W, 3]
- cam_points_conf: Camera point confidence [B, S, H, W]
Rotated point head outputs (if enabled):
- rotated_points: Rotated 3D coordinates per pixel [B, S, H, W, 3]
- rotated_points_conf: Rotated point confidence [B, S, H, W]
World point head outputs (if enabled):
- world_points: 3D world coordinates per pixel [B, S, H, W, 3]
- world_points_conf: World point confidence [B, S, H, W]
Inference-only outputs (not training):
- images: Original input images (for visualization) [B, S, 3, H, W]
"""
# Add batch dimension if missing (handle [S,3,H,W] -> [1,S,3,H,W])
if len(images.shape) == 4:
images = images.unsqueeze(0)
# Extract aggregated features from backbone
(
aggregated_tokens_list, # List of aggregated transformer tokens across iterations
patch_start_idx, # Patch start indices for feature reconstruction
covisibility_scores, # Covisibility scores between frames
ref_idx # Reference frame indices
) = self.aggregator(images)
# Initialize prediction dictionary
predictions: Dict[str, torch.Tensor] = {}
# Disable mixed precision for precise prediction calculations
with torch.amp.autocast("cuda", enabled=False):
# Add aggregator outputs to predictions
if covisibility_scores is not None:
predictions["covisibility_scores"] = covisibility_scores
if ref_idx is not None:
predictions["ref_idx"] = ref_idx
# Camera pose prediction (if enabled)
if self.camera_head is not None:
pose_enc_list = self.camera_head(aggregated_tokens_list)
predictions["pose_enc"] = pose_enc_list[-1] # Use final iteration encoding
predictions["pose_enc_list"] = pose_enc_list # Mutil-layer supervision
# Depth prediction (if enabled)
if self.depth_head is not None:
depth, depth_conf = self.depth_head(
aggregated_tokens_list,
images=images,
patch_start_idx=patch_start_idx
)
predictions["depth"] = depth
predictions["depth_conf"] = depth_conf
# Camera 3D point prediction (if enabled)
if self.cam_point_head is not None:
cam_pts3d, cam_pts3d_conf = self.cam_point_head(
aggregated_tokens_list,
images=images,
patch_start_idx=patch_start_idx
)
predictions["cam_points"] = cam_pts3d
predictions["cam_points_conf"] = cam_pts3d_conf
# Rotated 3D point prediction (if enabled)
if self.rotated_point_head is not None:
rotated_pts3d, rotated_pts3d_conf = self.rotated_point_head(
aggregated_tokens_list,
images=images,
patch_start_idx=patch_start_idx
)
predictions["rotated_points"] = rotated_pts3d
predictions["rotated_points_conf"] = rotated_pts3d_conf
# World 3D point prediction (if enabled)
if self.point_head is not None:
world_pts3d, world_pts3d_conf = self.point_head(
aggregated_tokens_list,
images=images,
patch_start_idx=patch_start_idx
)
predictions["world_points"] = world_pts3d
predictions["world_points_conf"] = world_pts3d_conf
# Store input images for visualization during inference (skip in training)
if not self.training:
predictions["images"] = images
if "ref_idx" in predictions:
ref_idx = predictions["ref_idx"].detach()
# Reorder all spatial/temporal data (exclude adjacency matrix and IDs)
predictions["images"] = reorder_by_reference(predictions["images"], ref_idx)
if self.restore_metric_scale:
# Restore metric scale
abs_scale = 10.0
if self.camera_head is not None:
predictions["pose_enc"][...,:3] *= abs_scale
if self.depth_head is not None:
predictions["depth"] *= abs_scale
if self.cam_point_head is not None:
predictions["cam_points"] *= abs_scale
if self.rotated_point_head is not None:
predictions["rotated_points"] *= abs_scale
if self.point_head is not None:
predictions["world_points"] *= abs_scale
return predictions |