228 lines
7.6 KiB
Python
228 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
"""Dual encoder for caption quality test on cross-view geo-localization.
|
|
|
|
GeoRSCLIP ViT-B/32 backbone (image + text towers, shared 512-dim space).
|
|
Image encoder frozen. Text encoder with partial unfreeze.
|
|
|
|
Architecture:
|
|
Query branch: GeoRSCLIP_img(drone) + GeoRSCLIP_text(caption) -> GatedFusion -> proj -> query_emb
|
|
Gallery branch: GeoRSCLIP_img(sat) -> proj -> gallery_emb
|
|
Loss: InfoNCE(query_emb, gallery_emb)
|
|
|
|
Baseline mode: fusion gate forced to 1.0 (text ignored).
|
|
"""
|
|
|
|
from typing import Literal
|
|
|
|
import gin
|
|
import open_clip
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# residual fusions exp
|
|
from .residual_fusions import ResidualGateType, GatedFusionResidual
|
|
# from residual_fusions import ResidualGateType, GatedFusionResidual
|
|
|
|
|
|
class ProjectionHead(nn.Module):
|
|
"""MLP projection head with L2 normalization."""
|
|
|
|
def __init__(
|
|
self,
|
|
in_dim: int = 512,
|
|
out_dim: int = 512,
|
|
use_mlp: bool = False,
|
|
hidden_dim: int | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
if use_mlp:
|
|
hidden_dim = hidden_dim or (2 * in_dim)
|
|
self.proj = nn.Sequential(
|
|
nn.Linear(in_dim, hidden_dim),
|
|
nn.GELU(),
|
|
nn.Linear(hidden_dim, out_dim),
|
|
)
|
|
else:
|
|
self.proj = nn.Linear(in_dim, out_dim)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return F.normalize(self.proj(x), dim=-1)
|
|
|
|
#! GATE-FUSION ORIG ---------------------------------
|
|
|
|
@gin.configurable
|
|
class GatedFusion(nn.Module):
|
|
"""Learnable gated fusion of image and text embeddings.
|
|
|
|
q = sigma(alpha) * img + (1 - sigma(alpha)) * text
|
|
|
|
alpha is a single learnable scalar, initialized so that gate ~ init_gate.
|
|
When baseline_mode=True, gate is clamped to 1.0 (text contribution = 0).
|
|
"""
|
|
|
|
def __init__(self, init_gate: float = 0.7, baseline_mode: bool = False) -> None:
|
|
super().__init__()
|
|
# alpha is in logit space: sigmoid(alpha) = init_gate
|
|
init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
|
|
self.alpha = nn.Parameter(init_alpha)
|
|
self.baseline_mode = baseline_mode
|
|
|
|
def forward(
|
|
self,
|
|
img_feat: torch.Tensor,
|
|
text_feat: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
if text_feat is None or self.baseline_mode:
|
|
return img_feat
|
|
gate = torch.sigmoid(self.alpha)
|
|
return gate * img_feat + (1.0 - gate) * text_feat
|
|
|
|
@property
|
|
def gate_value(self) -> float:
|
|
"""Current gate value (image weight). 1.0 = text ignored."""
|
|
if self.baseline_mode:
|
|
return 1.0
|
|
return torch.sigmoid(self.alpha).item()
|
|
|
|
|
|
@gin.configurable
|
|
class DualEncoderCaptionTest(nn.Module):
|
|
"""GeoRSCLIP dual encoder with gated text fusion on query branch.
|
|
|
|
Args:
|
|
variant: open_clip model variant name.
|
|
pretrained_path: Path to GeoRSCLIP checkpoint.
|
|
unfreeze_mode: Text encoder unfreeze strategy.
|
|
embed_dim: Output retrieval embedding dimension.
|
|
use_mlp_heads: Use 2-layer MLP projection heads.
|
|
baseline_mode: If True, fusion gate = 1.0 (no text).
|
|
init_gate: Initial gate value (image weight).
|
|
device: torch device.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
variant: str = "ViT-B-32",
|
|
pretrained_path: str = "RS5M_ViT-B-32.pt",
|
|
unfreeze_mode: Literal["none", "projection", "last_block"] = "last_block",
|
|
embed_dim: int = 512,
|
|
use_mlp_heads: bool = False,
|
|
baseline_mode: bool = False,
|
|
init_gate: float = 0.7,
|
|
device: str = "cuda",
|
|
gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate
|
|
) -> None:
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.device = device
|
|
self.baseline_mode = baseline_mode
|
|
|
|
# Load GeoRSCLIP via open_clip.
|
|
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
|
|
model_name=variant,
|
|
pretrained=pretrained_path,
|
|
device=device,
|
|
)
|
|
self.tokenizer = open_clip.get_tokenizer(variant)
|
|
|
|
native_dim = self._infer_native_dim()
|
|
|
|
# Freeze everything.
|
|
for p in self.model.parameters():
|
|
p.requires_grad = False
|
|
|
|
# Selectively unfreeze text encoder (only if not baseline).
|
|
if not baseline_mode:
|
|
self._apply_unfreeze(unfreeze_mode)
|
|
|
|
# Gated fusion on query branch.
|
|
# self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
|
|
|
|
#! Experimental Gated fusion on query branch.
|
|
self.fusion = GatedFusionResidual(gate_type=gate_type,
|
|
init_gate=init_gate, baseline_mode=baseline_mode)
|
|
|
|
# Projection heads.
|
|
self.proj_query = ProjectionHead(
|
|
in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads,
|
|
)
|
|
self.proj_gallery = ProjectionHead(
|
|
in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads,
|
|
)
|
|
|
|
def _infer_native_dim(self) -> int:
|
|
if hasattr(self.model, "text_projection"):
|
|
shape = self.model.text_projection.shape
|
|
return int(shape[1] if shape.ndim == 2 else shape[0])
|
|
return 512
|
|
|
|
def _apply_unfreeze(
|
|
self,
|
|
unfreeze_mode: Literal["none", "projection", "last_block"],
|
|
) -> None:
|
|
if unfreeze_mode == "none":
|
|
return
|
|
# Unfreeze text_projection.
|
|
if hasattr(self.model, "text_projection"):
|
|
tp = self.model.text_projection
|
|
if isinstance(tp, nn.Parameter):
|
|
tp.requires_grad = True
|
|
elif isinstance(tp, nn.Module):
|
|
for p in tp.parameters():
|
|
p.requires_grad = True
|
|
# Unfreeze last transformer block.
|
|
if unfreeze_mode == "last_block" and hasattr(self.model, "transformer"):
|
|
for p in self.model.transformer.resblocks[-1].parameters():
|
|
p.requires_grad = True
|
|
|
|
def encode_image(self, images: torch.Tensor) -> torch.Tensor:
|
|
feats = self.model.encode_image(images)
|
|
return F.normalize(feats, dim=-1)
|
|
|
|
def encode_text(self, texts: list[str]) -> torch.Tensor:
|
|
tokens = self.tokenizer(list(texts)).to(self.device).long()
|
|
feats = self.model.encode_text(tokens)
|
|
return F.normalize(feats, dim=-1)
|
|
|
|
def forward(
|
|
self,
|
|
drone_img: torch.Tensor,
|
|
sat_img: torch.Tensor,
|
|
caption_drone: list[str] | None = None,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Forward pass.
|
|
|
|
Args:
|
|
drone_img: Drone images [B, 3, H, W].
|
|
sat_img: Satellite images [B, 3, H, W].
|
|
caption_drone: Drone captions (P3 fingerprint), one per sample.
|
|
|
|
Returns:
|
|
Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
|
|
and 'gate' (scalar) for logging.
|
|
"""
|
|
# Gallery branch: satellite only.
|
|
sat_feat = self.encode_image(sat_img)
|
|
gallery = self.proj_gallery(sat_feat)
|
|
|
|
# Query branch: drone + optional text fusion.
|
|
drone_feat = self.encode_image(drone_img)
|
|
|
|
text_feat = None
|
|
if caption_drone is not None and not self.baseline_mode:
|
|
text_feat = self.encode_text(caption_drone)
|
|
|
|
fused = self.fusion(drone_feat, text_feat)
|
|
query = self.proj_query(fused)
|
|
|
|
return {
|
|
"query": query,
|
|
"gallery": gallery,
|
|
"gate": self.fusion.gate_value,
|
|
}
|
|
|
|
def trainable_parameters(self) -> list[nn.Parameter]:
|
|
return [p for p in self.parameters() if p.requires_grad]
|