Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: caption quality test on UAV-VisLoc

Browse Source 
Self-contained experimental track validating generated text captions
via retrieval R@1 lift on UAV-VisLoc.

Architecture: GeoRSCLIP ViT-B/32 dual encoder, 512-dim shared space.
Loss: 4-term InfoNCE (img-img + sat-cap + drone-cap + cap-cap)
      with cosine temperature decay, PALW-like curriculum.
Metric: delta R@1 (with text - without text) >= +3% => PASS.

Gin-configured (balanced / baseline_no_text / text_heavy variants).
Follows NADEZHDA code style.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Pikaliov
2026-04-17 00:04:46 +03:00
commit 2ce4017ebd
18 changed files with 1864 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/models/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
"""Model components for caption quality test."""
from src.models.dual_encoder import DualEncoderCaptionTest, ProjectionHead
__all__ = ["DualEncoderCaptionTest", "ProjectionHead"]

243
src/models/dual_encoder.py Normal file
View File

@@ -0,0 +1,243 @@
from __future__ import annotations
"""Dual encoder for caption quality test on UAV-VisLoc.
GeoRSCLIP ViT-B/32 backbone (image + text towers, shared 512-dim space).
Image encoder is frozen, text encoder has partial unfreeze (last block + projection).
Separate trainable projection heads for drone/sat/text branches.
"""
from typing import Literal
import gin
import open_clip
import torch
import torch.nn as nn
import torch.nn.functional as F
@gin.configurable
class ProjectionHead(nn.Module):
"""Single-layer L2-normalized projection head.
Args:
in_dim: Input embedding dimension.
out_dim: Output embedding dimension (512 for GeoRSCLIP space).
use_mlp: If True, use 2-layer MLP with GELU, else Linear.
hidden_dim: Hidden dim when use_mlp=True (defaults to 2*in_dim).
"""
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:
"""Project features and L2-normalize.
Args:
x: Input features [B, in_dim].
Returns:
Normalized embeddings [B, out_dim].
"""
x = self.proj(x)
return F.normalize(x, dim=-1)
@gin.configurable
class DualEncoderCaptionTest(nn.Module):
"""GeoRSCLIP dual encoder for caption quality validation on UAV-VisLoc.
Shared image encoder for drone and satellite views. Text encoder with
partial unfreeze. Three separate trainable projection heads map raw
GeoRSCLIP embeddings into the shared 512-dim retrieval space.
Args:
variant: open_clip model variant name (e.g., 'ViT-B-32').
pretrained_path: Path to GeoRSCLIP checkpoint (RS5M_ViT-B-32.pt).
unfreeze_mode: Which text encoder layers to unfreeze.
embed_dim: Output retrieval dimension (default 512).
use_mlp_heads: If True, projection heads are 2-layer MLPs.
shared_image_head: If True, drone and sat use single projection head.
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", "full"] = "last_block",
embed_dim: int = 512,
use_mlp_heads: bool = False,
shared_image_head: bool = True,
device: str = "cuda",
) -> None:
super().__init__()
self.variant = variant
self.embed_dim = embed_dim
self.shared_image_head = shared_image_head
self.device = device
# Load open_clip model (GeoRSCLIP compatible with open_clip API).
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 GeoRSCLIP embedding dim (for ViT-B/32 = 512).
self._native_dim = self._infer_native_dim()
# Freeze everything by default.
for p in self.model.parameters():
p.requires_grad = False
# Apply unfreeze strategy.
self._apply_unfreeze(unfreeze_mode)
# Projection heads (trainable).
self.proj_text = ProjectionHead(
in_dim=self._native_dim,
out_dim=embed_dim,
use_mlp=use_mlp_heads,
)
if shared_image_head:
self.proj_image = ProjectionHead(
in_dim=self._native_dim,
out_dim=embed_dim,
use_mlp=use_mlp_heads,
)
self.proj_drone = None # type: ignore[assignment]
self.proj_sat = None # type: ignore[assignment]
else:
self.proj_image = None # type: ignore[assignment]
self.proj_drone = ProjectionHead(
in_dim=self._native_dim,
out_dim=embed_dim,
use_mlp=use_mlp_heads,
)
self.proj_sat = ProjectionHead(
in_dim=self._native_dim,
out_dim=embed_dim,
use_mlp=use_mlp_heads,
)
def _infer_native_dim(self) -> int:
"""Infer native embedding dimension from model (typically 512 for ViT-B/32)."""
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", "full"],
) -> None:
"""Selectively enable gradients for text encoder."""
if unfreeze_mode == "none":
return
if unfreeze_mode == "full":
for p in self.model.parameters():
p.requires_grad = True
return
# Always unfreeze text_projection if available.
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
# Additionally unfreeze last transformer block.
if unfreeze_mode == "last_block" and hasattr(self.model, "transformer"):
last_block = self.model.transformer.resblocks[-1]
for p in last_block.parameters():
p.requires_grad = True
def encode_image(self, images: torch.Tensor) -> torch.Tensor:
"""Encode images through GeoRSCLIP image encoder (no projection head).
Args:
images: Preprocessed image tensor [B, 3, H, W].
Returns:
Raw image embeddings [B, native_dim].
"""
feats = self.model.encode_image(images)
return F.normalize(feats, dim=-1)
def encode_text(self, texts: list[str] | torch.Tensor) -> torch.Tensor:
"""Encode text captions through GeoRSCLIP text encoder.
Args:
texts: List of strings or pre-tokenized LongTensor [B, seq_len].
Returns:
Raw text embeddings [B, native_dim].
"""
if isinstance(texts, (list, tuple)):
tokens = self.tokenizer(list(texts)).to(self.device).long()
else:
tokens = 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,
caption_sat: list[str] | None = None,
) -> dict[str, torch.Tensor]:
"""Forward pass producing projected embeddings for all branches.
Args:
drone_img: Drone RGB tensor [B, 3, H, W].
sat_img: Satellite RGB tensor [B, 3, H, W].
caption_drone: List of drone captions, one per batch item.
caption_sat: List of satellite captions, one per batch item.
Returns:
Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat', each
containing [B, embed_dim] L2-normalized embeddings.
Keys for missing captions are absent.
"""
out: dict[str, torch.Tensor] = {}
drone_feat = self.encode_image(drone_img)
sat_feat = self.encode_image(sat_img)
if self.shared_image_head:
out["drone"] = self.proj_image(drone_feat)
out["sat"] = self.proj_image(sat_feat)
else:
out["drone"] = self.proj_drone(drone_feat)
out["sat"] = self.proj_sat(sat_feat)
if caption_drone is not None:
out["cap_drone"] = self.proj_text(self.encode_text(caption_drone))
if caption_sat is not None:
out["cap_sat"] = self.proj_text(self.encode_text(caption_sat))
return out
def trainable_parameters(self) -> list[nn.Parameter]:
"""Return list of trainable parameters for optimizer construction."""
return [p for p in self.parameters() if p.requires_grad]