656 lines
25 KiB
Python
656 lines
25 KiB
Python
from __future__ import annotations
|
||
|
||
"""Asymmetric dual encoder for CVGL caption test on GTA-UAV.
|
||
|
||
Architecture:
|
||
Query: DINOv3 ViT-L/16 (LVD, frozen) + LRSCLIP text (L1/L2/L3) -> GatedFusion -> query
|
||
Gallery: DINOv3 ViT-L/16 (SAT, frozen) -> gallery
|
||
Loss: InfoNCE(query, gallery)
|
||
|
||
DINOv3 checkpoints use a custom key layout (not HuggingFace transformers).
|
||
LRSCLIP (DGTRS-CLIP ViT-L-14) uses open_clip layout with KPS positional embeddings.
|
||
"""
|
||
|
||
import logging
|
||
import math
|
||
import warnings
|
||
from pathlib import Path
|
||
|
||
import coloredlogs
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
LOGGER = logging.getLogger("caption_test.model")
|
||
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||
from safetensors.torch import load_file as load_safetensors
|
||
|
||
from src.models.adapters import inject_lora_into_dgtrs, inject_mona_into_dinov3
|
||
from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
|
||
from src.models.dual_encoder import GatedFusion, ProjectionHead
|
||
from src.models.stripnet import inject_conv_mona_into_stripnet
|
||
from src.models.stripnet_encoder import StripNetEncoder
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# DINOv3 ViT-L/16 — minimal implementation matching checkpoint key layout
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class DINOv3Attention(nn.Module):
|
||
"""Multi-head self-attention with separate Q/K/V projections."""
|
||
|
||
def __init__(self, dim: int = 1024, num_heads: int = 16) -> None:
|
||
super().__init__()
|
||
self.num_heads = num_heads
|
||
self.head_dim = dim // num_heads
|
||
self.scale = self.head_dim ** -0.5
|
||
self.q_proj = nn.Linear(dim, dim)
|
||
self.k_proj = nn.Linear(dim, dim, bias=False)
|
||
self.v_proj = nn.Linear(dim, dim)
|
||
self.o_proj = nn.Linear(dim, dim)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
B, N, C = x.shape
|
||
q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
||
k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
||
v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
||
attn = F.scaled_dot_product_attention(q, k, v)
|
||
x = attn.permute(0, 2, 1, 3).reshape(B, N, C)
|
||
return self.o_proj(x)
|
||
|
||
|
||
class DINOv3LayerScale(nn.Module):
|
||
"""Per-channel learnable scale (lambda)."""
|
||
|
||
def __init__(self, dim: int) -> None:
|
||
super().__init__()
|
||
self.lambda1 = nn.Parameter(torch.ones(dim))
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
return x * self.lambda1
|
||
|
||
|
||
class DINOv3MLP(nn.Module):
|
||
"""SwiGLU-like MLP: up_proj + GELU + down_proj."""
|
||
|
||
def __init__(self, dim: int = 1024, mlp_dim: int = 4096) -> None:
|
||
super().__init__()
|
||
self.up_proj = nn.Linear(dim, mlp_dim)
|
||
self.down_proj = nn.Linear(mlp_dim, dim)
|
||
self.act = nn.GELU()
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
return self.down_proj(self.act(self.up_proj(x)))
|
||
|
||
|
||
class DINOv3Block(nn.Module):
|
||
"""Single DINOv3 transformer block."""
|
||
|
||
def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_dim: int = 4096) -> None:
|
||
super().__init__()
|
||
self.norm1 = nn.LayerNorm(dim)
|
||
self.attention = DINOv3Attention(dim, num_heads)
|
||
self.layer_scale1 = DINOv3LayerScale(dim)
|
||
self.norm2 = nn.LayerNorm(dim)
|
||
self.mlp = DINOv3MLP(dim, mlp_dim)
|
||
self.layer_scale2 = DINOv3LayerScale(dim)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = x + self.layer_scale1(self.attention(self.norm1(x)))
|
||
x = x + self.layer_scale2(self.mlp(self.norm2(x)))
|
||
return x
|
||
|
||
|
||
class DINOv3Embeddings(nn.Module):
|
||
"""Patch embedding + CLS token + register tokens."""
|
||
|
||
def __init__(
|
||
self,
|
||
dim: int = 1024,
|
||
patch_size: int = 16,
|
||
num_registers: int = 4,
|
||
) -> None:
|
||
super().__init__()
|
||
self.patch_embeddings = nn.Conv2d(3, dim, patch_size, patch_size)
|
||
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
|
||
self.register_tokens = nn.Parameter(torch.zeros(1, num_registers, dim))
|
||
self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
B = x.shape[0]
|
||
patches = self.patch_embeddings(x).flatten(2).transpose(1, 2) # [B, N, D]
|
||
N = patches.shape[1]
|
||
|
||
cls = self.cls_token.expand(B, -1, -1)
|
||
reg = self.register_tokens.expand(B, -1, -1)
|
||
|
||
# DINOv3: [CLS, registers, patches]
|
||
x = torch.cat([cls, reg, patches], dim=1)
|
||
|
||
# Positional embedding: interpolated sincos (RoPE applied in attention
|
||
# in original, but pretrained checkpoints bake it into weights).
|
||
# We use a simple learned-style pos embed computed on the fly.
|
||
pos = self._get_pos_embed(N, x.device, x.dtype)
|
||
# pos covers patches only, skip CLS + registers
|
||
x[:, 1 + reg.shape[1]:] = x[:, 1 + reg.shape[1]:] + pos
|
||
|
||
return x
|
||
|
||
def _get_pos_embed(self, n_patches: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
||
# DINOv3 uses RoPE internally — no additive pos embed needed.
|
||
# Return zeros as placeholder (weights handle positioning via RoPE).
|
||
return torch.zeros(1, n_patches, self.cls_token.shape[-1], device=device, dtype=dtype)
|
||
|
||
|
||
class DINOv3ViT(nn.Module):
|
||
"""DINOv3 ViT-L/16 matching the checkpoint key layout.
|
||
|
||
Checkpoint keys:
|
||
embeddings.cls_token, embeddings.patch_embeddings.{weight,bias},
|
||
embeddings.register_tokens, embeddings.mask_token,
|
||
layer.{i}.attention.{q,k,v,o}_proj.{weight,bias},
|
||
layer.{i}.layer_scale{1,2}.lambda1,
|
||
layer.{i}.mlp.{up,down}_proj.{weight,bias},
|
||
layer.{i}.norm{1,2}.{weight,bias},
|
||
norm.{weight,bias}
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
dim: int = 1024,
|
||
num_heads: int = 16,
|
||
mlp_dim: int = 4096,
|
||
num_layers: int = 24,
|
||
patch_size: int = 16,
|
||
num_registers: int = 4,
|
||
) -> None:
|
||
super().__init__()
|
||
self.embeddings = DINOv3Embeddings(dim, patch_size, num_registers)
|
||
self.layer = nn.ModuleList([
|
||
DINOv3Block(dim, num_heads, mlp_dim) for _ in range(num_layers)
|
||
])
|
||
self.norm = nn.LayerNorm(dim)
|
||
self.embed_dim = dim
|
||
self.gradient_checkpointing = False
|
||
|
||
def set_gradient_checkpointing(self, enable: bool = True) -> None:
|
||
"""Enable/disable gradient checkpointing to trade compute for VRAM."""
|
||
self.gradient_checkpointing = enable
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""Forward pass. Returns CLS token embedding [B, dim]."""
|
||
x = self.embeddings(x)
|
||
for block in self.layer:
|
||
if self.gradient_checkpointing and self.training:
|
||
x = torch.utils.checkpoint.checkpoint(
|
||
block, x, use_reentrant=False,
|
||
)
|
||
else:
|
||
x = block(x)
|
||
x = self.norm(x)
|
||
return x[:, 0] # CLS token
|
||
|
||
@classmethod
|
||
def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
|
||
"""Load from .pth or .safetensors checkpoint."""
|
||
model = cls()
|
||
path = Path(path)
|
||
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
|
||
if path.suffix == ".safetensors":
|
||
state = load_safetensors(str(path))
|
||
else:
|
||
state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||
if "model" in state:
|
||
state = state["model"]
|
||
elif "state_dict" in state:
|
||
state = state["state_dict"]
|
||
model.load_state_dict(state, strict=False)
|
||
n_params = sum(p.numel() for p in model.parameters())
|
||
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
|
||
return model
|
||
|
||
|
||
# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
|
||
# in src/models/dgtrs/model.py (DGTRSTextEncoder)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Text fusion MLP: concat L1/L2/L3 -> project to D
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TextFusionMLP(nn.Module):
|
||
"""Fuse L1/L2/L3 text embeddings via concat + MLP.
|
||
|
||
[B, 3*text_dim] -> [B, out_dim]
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
text_dim: int = 768,
|
||
out_dim: int = 1024,
|
||
) -> None:
|
||
super().__init__()
|
||
self.mlp = nn.Sequential(
|
||
nn.Linear(3 * text_dim, out_dim),
|
||
nn.GELU(),
|
||
nn.Linear(out_dim, out_dim),
|
||
)
|
||
|
||
def forward(
|
||
self,
|
||
z_l1: torch.Tensor,
|
||
z_l2: torch.Tensor,
|
||
z_l3: torch.Tensor,
|
||
) -> torch.Tensor:
|
||
"""Fuse three text embeddings.
|
||
|
||
Args:
|
||
z_l1: L1 overview [B, text_dim].
|
||
z_l2: L2 full description [B, text_dim].
|
||
z_l3: L3 fingerprint [B, text_dim].
|
||
|
||
Returns:
|
||
Fused text embedding [B, out_dim].
|
||
"""
|
||
cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
|
||
return self.mlp(cat)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main model: AsymmetricEncoder
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# ResidualGateFusin experiment
|
||
from .residual_fusions import ResidualGateType, GatedFusionResidual
|
||
|
||
class AsymmetricEncoder(nn.Module):
|
||
"""Dual encoder for CVGL with text fusion on both branches.
|
||
|
||
Supports two modes:
|
||
- **shared** (default): single DINOv3 WEB encoder for both drone and satellite,
|
||
one set of MONA adapters. Saves ~4-5 GB VRAM and halves adapter params.
|
||
- **asymmetric**: separate DINOv3 encoders (LVD for drone, SAT for satellite),
|
||
each with their own MONA adapters (legacy mode).
|
||
|
||
Query branch: DINOv3 (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
|
||
Gallery branch: DINOv3 (sat) + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]
|
||
|
||
No projection layers — retrieval space is DINOv3 native 1024-dim.
|
||
Text fusion MLP is shared between branches (same caption format).
|
||
Two separate GatedFusion gates (drone/sat may weight text differently).
|
||
|
||
For satellite images without captions, GatedFusion passes image features through
|
||
(text_feat=None → gate acts as identity).
|
||
|
||
Args:
|
||
dino_web_path: Path to DINOv3 LVD checkpoint (used for both branches in shared mode).
|
||
dino_sat_path: Path to DINOv3 SAT checkpoint (only used in asymmetric mode).
|
||
lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
|
||
init_gate: Initial fusion gate (image weight).
|
||
baseline_mode: If True, gate = 1.0 (text ignored), DGTRS not loaded.
|
||
shared_encoder: If True, use single DINOv3 WEB for both branches.
|
||
device: Torch device string.
|
||
"""
|
||
|
||
DINO_DIM = 1024
|
||
TEXT_DIM = 768
|
||
|
||
def __init__(
|
||
self,
|
||
gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate,
|
||
dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
|
||
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
|
||
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
|
||
init_gate: float = 0.7,
|
||
baseline_mode: bool = False,
|
||
shared_encoder: bool = False,
|
||
mona_bottleneck: int = 64,
|
||
mona_last_n_blocks: int = 24,
|
||
lora_rank: int = 4,
|
||
device: str = "cuda",
|
||
backbone: str = "dinov3",
|
||
stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
|
||
stripnet_mona_last_n_stages: int = 2,
|
||
stripnet_freeze: bool = True,
|
||
) -> None:
|
||
super().__init__()
|
||
self.embed_dim = self.DINO_DIM # native 1024 (StripNet projects 512 -> 1024)
|
||
self.baseline_mode = baseline_mode
|
||
self.shared_encoder = shared_encoder
|
||
self.backbone = backbone
|
||
self.device = device
|
||
|
||
# Image encoder(s) (frozen + MONA adapters).
|
||
if backbone == "stripnet":
|
||
# StripNet always operates as shared encoder (one CNN for both branches).
|
||
self.shared_encoder = True
|
||
self.image_encoder = StripNetEncoder(checkpoint_path=stripnet_path, out_dim=self.DINO_DIM)
|
||
if stripnet_freeze:
|
||
self._freeze(self.image_encoder.backbone)
|
||
LOGGER.info("StripNet backbone: frozen (Conv-MONA + projection trainable)")
|
||
else:
|
||
LOGGER.info("StripNet backbone: UNFROZEN — full fine-tune (use lower lr factor)")
|
||
if stripnet_mona_last_n_stages > 0:
|
||
inject_conv_mona_into_stripnet(
|
||
self.image_encoder.backbone,
|
||
bottleneck=mona_bottleneck,
|
||
last_n_stages=stripnet_mona_last_n_stages,
|
||
)
|
||
else:
|
||
LOGGER.info("Conv-MONA disabled (stripnet_mona_last_n_stages=0)")
|
||
LOGGER.info("StripNet backbone: shared encoder, projection 512 -> %d", self.DINO_DIM)
|
||
elif shared_encoder:
|
||
self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
|
||
self._freeze(self.image_encoder)
|
||
inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
|
||
LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
|
||
else:
|
||
self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
|
||
self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
|
||
self._freeze(self.drone_encoder)
|
||
self._freeze(self.sat_encoder)
|
||
inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
|
||
inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
|
||
LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")
|
||
|
||
# Text encoder — official DGTRS architecture (frozen + LoRA).
|
||
if not baseline_mode:
|
||
self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
|
||
self._freeze(self.text_encoder)
|
||
inject_lora_into_dgtrs(self.text_encoder, rank=lora_rank)
|
||
else:
|
||
self.text_encoder = None
|
||
|
||
# Shared text fusion MLP: 3×768 -> 1024 (native DINOv3 dim).
|
||
if not baseline_mode:
|
||
self.text_fusion = TextFusionMLP(
|
||
text_dim=self.TEXT_DIM,
|
||
out_dim=self.DINO_DIM,
|
||
)
|
||
|
||
# Separate gated fusion for query and gallery branches.
|
||
#! Experimental Gated fusion on query branch.
|
||
self.fusion_query = GatedFusionResidual(gate_type=gate_type,
|
||
init_gate=init_gate, baseline_mode=baseline_mode)
|
||
self.fusion_gallery = GatedFusionResidual(gate_type=gate_type,
|
||
init_gate=init_gate, baseline_mode=baseline_mode)
|
||
|
||
|
||
@staticmethod
|
||
def _freeze(module: nn.Module) -> None:
|
||
for p in module.parameters():
|
||
p.requires_grad = False
|
||
module.eval()
|
||
|
||
def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
|
||
"""Encode drone images with MONA adapters. Returns [B, 1024]."""
|
||
if self.shared_encoder:
|
||
return self.image_encoder(images)
|
||
return self.drone_encoder(images)
|
||
|
||
def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
|
||
"""Encode satellite images with MONA adapters. Returns [B, 1024]."""
|
||
if self.shared_encoder:
|
||
return self.image_encoder(images)
|
||
return self.sat_encoder(images)
|
||
|
||
def encode_text_levels(
|
||
self,
|
||
l1_texts: list[str],
|
||
l2_texts: list[str],
|
||
l3_texts: list[str],
|
||
) -> torch.Tensor | None:
|
||
"""Encode L1/L2/L3 captions and fuse. Returns [B, 1024] or None.
|
||
|
||
Returns None if all captions are empty (no text available).
|
||
For mixed batches (some have captions, some don't), encodes all
|
||
texts (empty strings tokenize to pad+EOS — their outputs must be
|
||
masked downstream, see `_fuse_with_mask`).
|
||
"""
|
||
# Check if any caption is non-empty.
|
||
if all(t == "" for t in l1_texts):
|
||
return None
|
||
|
||
z_l1 = self._encode_single_text(l1_texts)
|
||
z_l2 = self._encode_single_text(l2_texts)
|
||
z_l3 = self._encode_single_text(l3_texts)
|
||
fused = self.text_fusion(z_l1, z_l2, z_l3)
|
||
return F.normalize(fused, dim=-1)
|
||
|
||
def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
|
||
"""Tokenize and encode a list of strings using DGTRS tokenizer."""
|
||
tokens = tokenize_dgtrs(list(texts)).to(self.device)
|
||
return self.text_encoder(tokens)
|
||
|
||
def _fuse_with_mask(
|
||
self,
|
||
img_feat: torch.Tensor,
|
||
l1_texts: list[str] | None,
|
||
l2_texts: list[str] | None,
|
||
l3_texts: list[str] | None,
|
||
fusion: GatedFusionResidual,
|
||
) -> torch.Tensor:
|
||
"""Fuse image features with optional text, respecting per-sample presence.
|
||
|
||
For samples where caption is an empty string, output falls back to
|
||
pure image features (avoiding noise contamination from empty-string
|
||
text embeddings). For samples with captions, applies the standard
|
||
gated fusion `σ(α)·img + (1-σ(α))·text`.
|
||
|
||
Returns L2-normalized [B, D] embedding.
|
||
"""
|
||
if (
|
||
self.baseline_mode
|
||
or l1_texts is None
|
||
or l2_texts is None
|
||
or l3_texts is None
|
||
):
|
||
return F.normalize(fusion(img_feat, None), dim=-1)
|
||
|
||
has_text = torch.tensor(
|
||
[t != "" for t in l1_texts], dtype=torch.bool, device=img_feat.device,
|
||
)
|
||
if not has_text.any():
|
||
return F.normalize(fusion(img_feat, None), dim=-1)
|
||
|
||
z_text = self.encode_text_levels(l1_texts, l2_texts, l3_texts)
|
||
if z_text is None:
|
||
return F.normalize(fusion(img_feat, None), dim=-1)
|
||
|
||
# Per-sample fusion: text-present samples use full gated fusion,
|
||
# empty-caption samples pass through pure image features.
|
||
fused_with_text = fusion(img_feat, z_text)
|
||
out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat)
|
||
return F.normalize(out, dim=-1)
|
||
|
||
def encode_query(
|
||
self,
|
||
drone_img: torch.Tensor,
|
||
caption_l1: list[str] | None = None,
|
||
caption_l2: list[str] | None = None,
|
||
caption_l3: list[str] | None = None,
|
||
) -> torch.Tensor:
|
||
"""Encode drone → normalized query embedding with per-sample text mask."""
|
||
drone_feat = self.encode_drone(drone_img)
|
||
return self._fuse_with_mask(
|
||
drone_feat, caption_l1, caption_l2, caption_l3, self.fusion_query,
|
||
)
|
||
|
||
def encode_gallery(
|
||
self,
|
||
sat_img: torch.Tensor,
|
||
sat_caption_l1: list[str] | None = None,
|
||
sat_caption_l2: list[str] | None = None,
|
||
sat_caption_l3: list[str] | None = None,
|
||
) -> torch.Tensor:
|
||
"""Encode satellite → normalized gallery embedding with per-sample text mask."""
|
||
sat_feat = self.encode_satellite(sat_img)
|
||
return self._fuse_with_mask(
|
||
sat_feat, sat_caption_l1, sat_caption_l2, sat_caption_l3, self.fusion_gallery,
|
||
)
|
||
|
||
def forward(
|
||
self,
|
||
drone_img: torch.Tensor,
|
||
sat_img: torch.Tensor,
|
||
caption_l1: list[str] | None = None,
|
||
caption_l2: list[str] | None = None,
|
||
caption_l3: list[str] | None = None,
|
||
sat_caption_l1: list[str] | None = None,
|
||
sat_caption_l2: list[str] | None = None,
|
||
sat_caption_l3: list[str] | None = None,
|
||
) -> dict[str, torch.Tensor]:
|
||
"""Forward pass.
|
||
|
||
Both branches use per-sample caption masking: samples with an empty
|
||
caption string fall back to pure image features instead of being
|
||
fused with noise from empty-string text embeddings.
|
||
|
||
Args:
|
||
drone_img: Drone images [B, 3, 256, 256].
|
||
sat_img: Satellite images [B, 3, 256, 256].
|
||
caption_l1/l2/l3: Drone L1/L2/L3 captions.
|
||
sat_caption_l1/l2/l3: Satellite L1/L2/L3 captions.
|
||
|
||
Returns:
|
||
Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
|
||
'gate_q', 'gate_g'.
|
||
"""
|
||
query = self.encode_query(drone_img, caption_l1, caption_l2, caption_l3)
|
||
gallery = self.encode_gallery(sat_img, sat_caption_l1, sat_caption_l2, sat_caption_l3)
|
||
return {
|
||
"query": query,
|
||
"gallery": gallery,
|
||
"gate_q": self.fusion_query.gate_value,
|
||
"gate_g": self.fusion_gallery.gate_value,
|
||
}
|
||
|
||
def trainable_parameters(self) -> list[nn.Parameter]:
|
||
"""Return list of parameters that require grad."""
|
||
return [p for p in self.parameters() if p.requires_grad]
|
||
|
||
def save_checkpoint(self, path: str | Path, **extra) -> None:
|
||
"""Save model checkpoint with metadata."""
|
||
path = Path(path)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
ckpt = {
|
||
"model_state": self.state_dict(),
|
||
"baseline_mode": self.baseline_mode,
|
||
"shared_encoder": self.shared_encoder,
|
||
**extra,
|
||
}
|
||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||
torch.save(ckpt, tmp)
|
||
tmp.replace(path)
|
||
LOGGER.info("💾 Model saved to %s", path)
|
||
|
||
@classmethod
|
||
def load_checkpoint(
|
||
cls,
|
||
path: str | Path,
|
||
dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
|
||
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
|
||
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
|
||
device: str = "cuda",
|
||
) -> tuple[AsymmetricEncoder, dict]:
|
||
"""Load model from checkpoint.
|
||
|
||
First builds the model (loading frozen backbones), then loads
|
||
the saved trainable weights on top.
|
||
|
||
Returns:
|
||
(model, checkpoint_dict) — model ready for eval/resume,
|
||
checkpoint_dict has optimizer_state, epoch, etc.
|
||
"""
|
||
path = Path(path)
|
||
LOGGER.info("📂 Loading checkpoint from %s", path)
|
||
ckpt = torch.load(str(path), map_location="cpu", weights_only=False)
|
||
|
||
model = cls(
|
||
dino_web_path=dino_web_path,
|
||
dino_sat_path=dino_sat_path,
|
||
lrsclip_path=lrsclip_path,
|
||
baseline_mode=ckpt.get("baseline_mode", False),
|
||
shared_encoder=ckpt.get("shared_encoder", False),
|
||
mona_bottleneck=ckpt.get("mona_bottleneck", 64),
|
||
mona_last_n_blocks=ckpt.get("mona_last_n_blocks", 24),
|
||
device=device,
|
||
)
|
||
model.load_state_dict(ckpt["model_state"], strict=False)
|
||
model = model.to(device)
|
||
LOGGER.info("✅ Checkpoint loaded (epoch=%s)", ckpt.get("epoch", "?"))
|
||
return model, ckpt
|
||
|
||
def train(self, mode: bool = True) -> AsymmetricEncoder:
|
||
"""Override to keep frozen encoders in eval mode."""
|
||
super().train(mode)
|
||
if self.shared_encoder:
|
||
self.image_encoder.eval()
|
||
else:
|
||
self.drone_encoder.eval()
|
||
self.sat_encoder.eval()
|
||
if self.text_encoder is not None:
|
||
self.text_encoder.train(mode)
|
||
return self
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Image preprocessing (DINOv3: 256x256, ImageNet normalization)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
||
_IMAGENET_STD = [0.229, 0.224, 0.225]
|
||
|
||
|
||
def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
|
||
"""Build eval/inference image transform for DINOv3 input."""
|
||
from torchvision import transforms
|
||
return transforms.Compose([
|
||
transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
|
||
transforms.CenterCrop(image_size),
|
||
transforms.ToTensor(),
|
||
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
|
||
])
|
||
|
||
|
||
def get_drone_train_transform(image_size: int = 256) -> torch.nn.Module:
|
||
"""Build training augmentation for drone images.
|
||
|
||
Includes: RandomResizedCrop, HFlip, rotation, color jitter,
|
||
grayscale, Gaussian blur.
|
||
"""
|
||
from torchvision import transforms
|
||
return transforms.Compose([
|
||
transforms.RandomResizedCrop(
|
||
image_size, scale=(0.7, 1.0),
|
||
interpolation=transforms.InterpolationMode.BICUBIC,
|
||
),
|
||
transforms.RandomHorizontalFlip(p=0.5),
|
||
transforms.RandomRotation(degrees=15),
|
||
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
|
||
transforms.RandomGrayscale(p=0.05),
|
||
transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
|
||
transforms.ToTensor(),
|
||
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
|
||
])
|
||
|
||
|
||
def get_satellite_train_transform(image_size: int = 256) -> torch.nn.Module:
|
||
"""Build training augmentation for satellite images.
|
||
|
||
Lighter than drone: no rotation or blur (satellite is nadir/consistent).
|
||
Includes: RandomResizedCrop, HFlip, color jitter, grayscale.
|
||
"""
|
||
from torchvision import transforms
|
||
return transforms.Compose([
|
||
transforms.RandomResizedCrop(
|
||
image_size, scale=(0.7, 1.0),
|
||
interpolation=transforms.InterpolationMode.BICUBIC,
|
||
),
|
||
transforms.RandomHorizontalFlip(p=0.5),
|
||
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
|
||
transforms.RandomGrayscale(p=0.05),
|
||
transforms.ToTensor(),
|
||
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
|
||
])
|