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

55
.gitignore vendored Normal file
View File

@@ -0,0 +1,55 @@
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
.pytest_cache/
.mypy_cache/
.ruff_cache/
.venv/
venv/
env/
# PyTorch / ML artefacts
*.pt
*.pth
*.ckpt
*.onnx
*.engine
checkpoints/
out/
outputs/
runs/
wandb/
lightning_logs/
# Data (too large / confidential — keep manifests only if < 10 MB)
data/
*.json.gz
*.h5
*.hdf5
*.parquet
*.npy
*.npz
# Gin / config logs
*.gin.log
operative_config-*.gin
# Editors / IDE
.vscode/
.idea/
*.swp
*.swo
*~
.DS_Store
Thumbs.db
# Secrets (never commit)
.env
.env.*
*credentials*
*secret*
*.pem
*.key

138
README.md Normal file
View File

@@ -0,0 +1,138 @@
# Caption Quality Test on UAV-VisLoc
Validate generated text captions by measuring retrieval R@1 lift on UAV-VisLoc.
Uses GeoRSCLIP ViT-B/32 dual encoder with multi-term InfoNCE.
Full analysis: `2_hypotesis/АНАЛИЗ_caption_quality_test_VisLoc.md`
## Structure
```
caption_test/
├── conf/
│ ├── balanced.gin # Primary test: λ=(1.0, 0.3, 0.3, 0.1)
│ ├── baseline_no_text.gin # Reference: λ=(1.0, 0, 0, 0)
│ └── text_heavy.gin # Stress test: λ=(0.5, 0.5, 0.5, 0.2)
├── scripts/
│ ├── generate_captions.py # Offline caption generation (template/VLM/hybrid)
│ └── compare_runs.py # Δ R@1 comparison report builder
├── src/
│ ├── datasets/
│ │ └── visloc_with_captions.py
│ ├── models/
│ │ └── dual_encoder.py # GeoRSCLIP wrapper, projection heads
│ ├── losses/
│ │ └── multi_infonce.py # 4-term InfoNCE + curriculum + cosine τ
│ ├── training/
│ │ └── train.py # Main loop
│ └── eval/
│ └── evaluate.py # R@K metrics, Δ R@1 helper
└── data/ # (user-provided) VisLoc pairs + captions
```
## Prerequisites
```
torch>=2.0
open_clip_torch
gin-config
Pillow
numpy
```
GeoRSCLIP ViT-B/32 checkpoint: download `RS5M_ViT-B-32.pt` from
`github.com/om-ai-lab/RS5M` and place under `checkpoints/`.
## Workflow
### 1. Generate captions
```bash
python -m scripts.generate_captions \
--image_root data/visloc/images \
--pairs_csv data/visloc/pairs.csv \
--output data/visloc_train.json \
--strategy hybrid \
--vlm_refine_ratio 0.1
```
Replace `_placeholder_vlm_caption` in `scripts/generate_captions.py` with real
Qwen2.5-VL or InternVL2 inference before running on production data.
### 2. Train three variants (in parallel or sequentially)
```bash
# Baseline (no captions)
python -m src.training.train --config conf/baseline_no_text.gin
# Balanced (primary, with captions)
python -m src.training.train --config conf/balanced.gin
# Text-heavy (stress test)
python -m src.training.train --config conf/text_heavy.gin
```
### 3. Evaluate each on test split
```python
from src.eval.evaluate import run_evaluation_from_checkpoint
run_evaluation_from_checkpoint(
checkpoint_path="out/caption_test/balanced/ckpt_epoch029.pt",
test_manifest="data/visloc_test.json",
image_root="data/visloc/images",
output_path="out/caption_test/balanced/eval_report.json",
)
```
### 4. Compare and get verdict
```bash
python -m scripts.compare_runs \
--baseline_report out/caption_test/baseline_no_text/eval_report.json \
--full_report out/caption_test/balanced/eval_report.json \
--output out/caption_test/comparison.md
```
## Decision rule (from `compare_runs.py`)
| Δ R@1 (drone→sat) | Verdict |
|---|---|
| ≥ +3% | ✅ PASS — captions informative, proceed to World-UAV |
| +1% to +3% | ⚠️ MARGINAL — add VLM refinement, re-run |
| 0 to +1% | ❌ WEAK — redesign caption pipeline |
| < 0 | ❌❌ HARMFUL — critical bug |
## Expected runtime (RTX 4090, 24 GB)
| Phase | Time |
|---|---|
| Caption generation (6K pairs, hybrid) | ~1 h |
| Single training run (30 epochs, batch 128) | ~2–3 h |
| Full test (3 variants × 3 seeds = 9 runs) | ~30 h |
| Evaluation + comparison | ~30 min |
## Notes on code style
Follows NADEZHDA code style:
- `from __future__ import annotations` everywhere.
- Type hints on all signatures.
- Google-style docstrings.
- `@gin.configurable` on top-level classes.
- Atomic checkpoint saves (`_atomic_save` helper).
- No emojis in code, English-only code comments.
## Relation to NADEZHDA main pipeline
This is an **isolated experimental track**: completely separate from the main
Student/Teacher training under `code_nadezhda/`. Once captions pass the
Δ R@1 ≥ +3% gate here, the hybrid caption generation strategy is applied to
World-UAV 927K, and those captions feed into E1 Teacher training with
Multi-FiLM conditioning (see `АНАЛИЗ_fusion_для_NADEZHDA.md`).
## Files referenced
- `2_hypotesis/АНАЛИЗ_caption_quality_test_VisLoc.md` — full experimental design
- `2_hypotesis/АНАЛИЗ_text_encoder_для_NADEZHDA.md` — why GeoRSCLIP
- `2_hypotesis/АНАЛИЗ_fusion_для_NADEZHDA.md` — where captions land in Teacher
- `2_hypotesis/ROADMAP_E0_E9_unified.md` — phase E_caption (parallel to E0)

54
conf/balanced.gin Normal file
View File

@@ -0,0 +1,54 @@
# Balanced configuration — primary test setup.
# L = 1.0 * L_img_img + 0.3 * L_sat_cap + 0.3 * L_drone_cap + 0.1 * L_cap_cap
import src.datasets.visloc_with_captions
import src.losses.multi_infonce
import src.models.dual_encoder
import src.training.train
# ---- Dual encoder ----
DualEncoderCaptionTest.variant = "ViT-B-32"
DualEncoderCaptionTest.pretrained_path = "checkpoints/RS5M_ViT-B-32.pt"
DualEncoderCaptionTest.unfreeze_mode = "last_block"
DualEncoderCaptionTest.embed_dim = 512
DualEncoderCaptionTest.use_mlp_heads = False
DualEncoderCaptionTest.shared_image_head = True
DualEncoderCaptionTest.device = "cuda"
ProjectionHead.in_dim = 512
ProjectionHead.out_dim = 512
ProjectionHead.use_mlp = False
# ---- Loss ----
MultiTermInfoNCE.temperature_init = 0.1
MultiTermInfoNCE.temperature_final = 0.01
MultiTermInfoNCE.label_smoothing = 0.1
MultiTermInfoNCE.asym_drone_to_sat = 0.6
MultiTermInfoNCE.asym_sat_to_drone = 0.4
MultiTermInfoNCE.warmup_epochs = 3
MultiTermInfoNCE.text_ramp_epochs = 10
MultiTermInfoNCE.lambda_ii = 1.0
MultiTermInfoNCE.lambda_sc_max = 0.3
MultiTermInfoNCE.lambda_dc_max = 0.3
MultiTermInfoNCE.lambda_cc_max = 0.1
# ---- Dataset ----
VisLocCaptionDataset.caption_strategy = "hybrid"
VisLocCaptionDataset.drop_caption_prob = 0.0
VisLocCaptionDataset.seed = 42
# ---- Training ----
TrainConfig.train_manifest = "data/visloc_train.json"
TrainConfig.val_manifest = "data/visloc_val.json"
TrainConfig.image_root = "data/visloc/images"
TrainConfig.output_dir = "out/caption_test/balanced"
TrainConfig.epochs = 30
TrainConfig.batch_size = 128
TrainConfig.num_workers = 4
TrainConfig.learning_rate = 1e-4
TrainConfig.weight_decay = 1e-4
TrainConfig.grad_clip = 1.0
TrainConfig.use_amp = True
TrainConfig.eval_every = 1
TrainConfig.seed = 42
TrainConfig.device = "cuda"

11
conf/baseline_no_text.gin Normal file
View File

@@ -0,0 +1,11 @@
# Baseline: image-image only, no captions. Reference R@1 for delta computation.
# L = 1.0 * L_img_img
include 'balanced.gin'
# Disable all caption loss terms.
MultiTermInfoNCE.lambda_sc_max = 0.0
MultiTermInfoNCE.lambda_dc_max = 0.0
MultiTermInfoNCE.lambda_cc_max = 0.0
TrainConfig.output_dir = "out/caption_test/baseline_no_text"

11
conf/text_heavy.gin Normal file
View File

@@ -0,0 +1,11 @@
# Text-heavy configuration — stress test of caption contribution.
# L = 0.5 * L_img_img + 0.5 * L_sat_cap + 0.5 * L_drone_cap + 0.2 * L_cap_cap
include 'balanced.gin'
MultiTermInfoNCE.lambda_ii = 0.5
MultiTermInfoNCE.lambda_sc_max = 0.5
MultiTermInfoNCE.lambda_dc_max = 0.5
MultiTermInfoNCE.lambda_cc_max = 0.2
TrainConfig.output_dir = "out/caption_test/text_heavy"

138
scripts/compare_runs.py Normal file
View File

@@ -0,0 +1,138 @@
from __future__ import annotations
"""Compare baseline vs full-caption runs and compute Delta R@1 report.
Reads eval reports produced by src.eval.evaluate.run_evaluation_from_checkpoint
and produces a markdown + JSON summary.
Usage:
python -m scripts.compare_runs \
--baseline_report out/caption_test/baseline_no_text/eval_report.json \
--full_report out/caption_test/balanced/eval_report.json \
--output out/caption_test/comparison.md
"""
import argparse
import json
from pathlib import Path
_DIRECTIONS = (
"drone_to_sat",
"sat_to_drone",
"text_to_sat",
"text_to_drone",
)
_KS = (1, 5, 10)
def _load_metrics(report_path: Path) -> dict[str, float]:
with report_path.open("r", encoding="utf-8") as f:
data = json.load(f)
return data.get("metrics", data)
def _format_row(name: str, baseline: dict[str, float], full: dict[str, float]) -> str:
"""Render one markdown row for a direction across R@1, R@5, R@10."""
cells = [name]
for k in _KS:
key = f"r@{k}_{name}"
b = baseline.get(key, float("nan"))
f_ = full.get(key, float("nan"))
delta = f_ - b if (b == b and f_ == f_) else float("nan") # NaN-safe
cells.append(f"{b:.4f} → {f_:.4f} (Δ {delta:+.4f})")
return "| " + " | ".join(cells) + " |"
def _interpret_delta(delta: float) -> str:
"""Human-readable caption-quality verdict."""
if delta >= 0.03:
return "✅ PASS — captions informative (Δ R@1 ≥ +3%)"
if delta >= 0.01:
return "⚠️ MARGINAL — consider VLM refinement (+1% ≤ Δ < +3%)"
if delta >= 0:
return "❌ WEAK — captions add little signal (< +1%)"
return "❌❌ HARMFUL — captions confuse model (Δ < 0)"
def build_comparison_markdown(
baseline: dict[str, float],
full: dict[str, float],
) -> str:
"""Compose markdown comparison report."""
lines: list[str] = ["# Caption Quality Test: Comparison Report", ""]
# Headline Δ R@1 on primary direction.
primary = "drone_to_sat"
primary_key = f"r@1_{primary}"
primary_delta = full.get(primary_key, 0.0) - baseline.get(primary_key, 0.0)
verdict = _interpret_delta(primary_delta)
lines.append(f"## Primary metric: Δ R@1 ({primary}) = {primary_delta:+.4f}")
lines.append("")
lines.append(f"**Verdict:** {verdict}")
lines.append("")
# Full table.
lines.append("## All directions × K")
lines.append("")
header = "| Direction | R@1 base → full | R@5 base → full | R@10 base → full |"
sep = "|---|---|---|---|"
lines.extend([header, sep])
for direction in _DIRECTIONS:
row = _format_row(direction, baseline, full)
lines.append(row)
lines.append("")
# Decision rule recap.
lines.append("## Decision rule")
lines.append("")
lines.append("- Δ R@1 ≥ +3% → captions pass, proceed to World-UAV generation")
lines.append("- +1% ≤ Δ R@1 < +3% → add VLM refinement, re-run")
lines.append("- Δ R@1 < +1% → redesign caption pipeline")
lines.append("- Δ R@1 < 0 → critical bug, investigate caption/image alignment")
lines.append("")
return "\n".join(lines)
def main() -> None:
parser = argparse.ArgumentParser(
description="Compare baseline vs full-caption runs."
)
parser.add_argument("--baseline_report", type=Path, required=True)
parser.add_argument("--full_report", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
baseline = _load_metrics(args.baseline_report)
full = _load_metrics(args.full_report)
md = build_comparison_markdown(baseline=baseline, full=full)
args.output.parent.mkdir(parents=True, exist_ok=True)
with args.output.open("w", encoding="utf-8") as f:
f.write(md)
# Also write machine-readable summary.
summary = {
"baseline_metrics": baseline,
"full_metrics": full,
"deltas": {
f"delta_r@{k}_{d}": (
full.get(f"r@{k}_{d}", 0.0) - baseline.get(f"r@{k}_{d}", 0.0)
)
for d in _DIRECTIONS
for k in _KS
},
}
summary_path = args.output.with_suffix(".json")
with summary_path.open("w", encoding="utf-8") as f:
json.dump(summary, f, indent=2)
print(f"Comparison saved: {args.output}")
print(f"Summary saved: {summary_path}")
if __name__ == "__main__":
main()

210
scripts/generate_captions.py Normal file
View File

@@ -0,0 +1,210 @@
from __future__ import annotations
"""Offline caption generation for UAV-VisLoc images.
Supports three strategies:
- template: rule-based from SegEarth-OV3 masks (fastest, generic).
- vlm: Qwen2.5-VL / InternVL2 VLM (slowest, most diverse).
- hybrid: template first, VLM refinement on 10% sample (balanced).
Writes a manifest JSON that is directly consumable by VisLocCaptionDataset.
Usage:
python -m scripts.generate_captions \
--image_root data/visloc/images \
--pairs_csv data/visloc/pairs.csv \
--output data/visloc_train.json \
--strategy hybrid \
--vlm_refine_ratio 0.1
"""
import argparse
import json
import logging
import random
from pathlib import Path
LOGGER = logging.getLogger("caption_test.generate")
_TEMPLATE_SAT_PATTERNS = [
"aerial satellite view of {area_type} with {feature_1} and {feature_2}",
"orthogonal satellite image showing {area_type}, visible {feature_1}",
"top-down satellite photo of {area_type}",
]
_TEMPLATE_DRONE_PATTERNS = [
"low-altitude drone photo of {area_type} with {feature_1} and {feature_2}",
"oblique UAV view of {area_type}, showing {feature_1}",
"aerial drone image of {area_type}",
]
def _template_caption(
view: str,
area_type: str,
features: list[str],
rng: random.Random,
) -> str:
"""Generate rule-based caption from semantic masks."""
patterns = _TEMPLATE_SAT_PATTERNS if view == "sat" else _TEMPLATE_DRONE_PATTERNS
pattern = rng.choice(patterns)
feat_1 = features[0] if len(features) > 0 else "varied terrain"
feat_2 = features[1] if len(features) > 1 else "natural features"
return pattern.format(area_type=area_type, feature_1=feat_1, feature_2=feat_2)
def _placeholder_vlm_caption(image_path: Path, view: str) -> str:
"""Placeholder for VLM caption. Replace with real Qwen2.5-VL inference.
Returns a short deterministic caption for smoke-testing pipelines.
"""
# TODO: integrate Qwen2.5-VL / InternVL2 inference here.
if view == "sat":
return f"satellite aerial view (placeholder for {image_path.name})"
return f"drone low-altitude view (placeholder for {image_path.name})"
def _parse_pairs_csv(pairs_csv: Path) -> list[dict]:
"""Load pair metadata (drone_path, sat_path, gps, optional masks)."""
import csv
entries: list[dict] = []
with pairs_csv.open("r", encoding="utf-8", newline="") as f:
reader = csv.DictReader(f)
for row in reader:
entries.append(
{
"pair_id": row.get("pair_id", f"pair_{len(entries)}"),
"drone_path": row["drone_path"],
"sat_path": row["sat_path"],
"gps": [float(row.get("lat", 0.0)), float(row.get("lon", 0.0))],
"area_type": row.get("area_type", "mixed terrain"),
"features": [
s.strip()
for s in row.get("features", "").split(";")
if s.strip()
],
}
)
return entries
def build_manifest(
image_root: Path,
pairs_csv: Path,
output_path: Path,
strategy: str,
vlm_refine_ratio: float,
seed: int,
) -> None:
"""Build a manifest JSON with captions for all pairs.
Args:
image_root: Directory prefix for images.
pairs_csv: CSV with pair metadata (drone_path, sat_path, ...).
output_path: Output JSON path.
strategy: 'template' | 'vlm' | 'hybrid'.
vlm_refine_ratio: Fraction to refine with VLM when strategy='hybrid'.
seed: Random seed.
"""
rng = random.Random(seed)
entries = _parse_pairs_csv(pairs_csv)
LOGGER.info("loaded %d pairs from %s", len(entries), pairs_csv)
manifest: list[dict] = []
n_vlm_refined = 0
for i, entry in enumerate(entries):
area_type = entry["area_type"]
features = entry["features"]
# Template captions
cap_sat_tpl = _template_caption("sat", area_type, features, rng)
cap_drone_tpl = _template_caption("drone", area_type, features, rng)
# VLM captions (optional, based on strategy)
use_vlm = False
if strategy == "vlm":
use_vlm = True
elif strategy == "hybrid" and rng.random() < vlm_refine_ratio:
use_vlm = True
if use_vlm:
cap_sat_vlm = _placeholder_vlm_caption(
image_root / entry["sat_path"], "sat"
)
cap_drone_vlm = _placeholder_vlm_caption(
image_root / entry["drone_path"], "drone"
)
n_vlm_refined += 1
else:
cap_sat_vlm = cap_sat_tpl
cap_drone_vlm = cap_drone_tpl
# Final hybrid caption prefers VLM when present.
final_sat = cap_sat_vlm if use_vlm else cap_sat_tpl
final_drone = cap_drone_vlm if use_vlm else cap_drone_tpl
manifest.append(
{
"pair_id": entry["pair_id"],
"drone_path": entry["drone_path"],
"sat_path": entry["sat_path"],
"gps": entry["gps"],
# Strategy-specific captions (for ablations).
"caption_sat_template": cap_sat_tpl,
"caption_drone_template": cap_drone_tpl,
"caption_sat_vlm": cap_sat_vlm,
"caption_drone_vlm": cap_drone_vlm,
# Generic 'hybrid' fields used by default dataset.
"caption_sat": final_sat,
"caption_drone": final_drone,
}
)
if (i + 1) % 1000 == 0:
LOGGER.info("processed %d / %d pairs", i + 1, len(entries))
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8") as f:
json.dump(manifest, f, indent=2, ensure_ascii=False)
LOGGER.info(
"wrote %d entries to %s (%d VLM-refined, strategy=%s)",
len(manifest),
output_path,
n_vlm_refined,
strategy,
)
def main() -> None:
parser = argparse.ArgumentParser(description="Generate captions for UAV-VisLoc.")
parser.add_argument("--image_root", type=Path, required=True)
parser.add_argument("--pairs_csv", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--strategy",
choices=["template", "vlm", "hybrid"],
default="hybrid",
)
parser.add_argument("--vlm_refine_ratio", type=float, default=0.1)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(name)s %(levelname)s %(message)s",
)
build_manifest(
image_root=args.image_root,
pairs_csv=args.pairs_csv,
output_path=args.output,
strategy=args.strategy,
vlm_refine_ratio=args.vlm_refine_ratio,
seed=args.seed,
)
if __name__ == "__main__":
main()

1
src/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Caption quality test package."""

7
src/datasets/__init__.py Normal file
View File

@@ -0,0 +1,7 @@
"""Dataset loaders for caption quality test."""
from src.datasets.visloc_with_captions import (
VisLocCaptionDataset,
collate_caption_batch,
)
__all__ = ["VisLocCaptionDataset", "collate_caption_batch"]

158
src/datasets/visloc_with_captions.py Normal file
View File

@@ -0,0 +1,158 @@
from __future__ import annotations
"""UAV-VisLoc dataset loader augmented with generated captions.
Expects a manifest JSON of the form:
[
{
"pair_id": "v001_0042",
"drone_path": "drone/v001_0042.jpg",
"sat_path": "satellite/v001_0042.png",
"caption_drone": "low-altitude photo of residential ...",
"caption_sat": "aerial view of urban area ...",
"gps": [lat, lon]
},
...
]
Captions are produced offline by scripts/generate_captions.py using one of
three strategies: template, VLM, or hybrid (see АНАЛИЗ_caption_quality_test).
"""
import json
import random
from pathlib import Path
from typing import Any, Callable
import gin
import torch
from PIL import Image
from torch.utils.data import Dataset
@gin.configurable
class VisLocCaptionDataset(Dataset):
"""UAV-VisLoc pairs with generated captions.
Args:
manifest_path: Path to JSON manifest with pair entries.
image_root: Directory prefix joined with manifest relative paths.
image_transform: Callable applied to PIL images (e.g., GeoRSCLIP preprocess).
caption_strategy: Which caption field to use ('template', 'vlm', 'hybrid').
The corresponding field must exist in the manifest
(e.g., 'caption_sat_vlm', or the generic 'caption_sat').
drop_caption_prob: Random probability of replacing a caption with ''.
Useful for dropout ablations during training.
seed: Random seed for reproducibility.
"""
def __init__(
self,
manifest_path: str,
image_root: str,
image_transform: Callable[[Image.Image], torch.Tensor],
caption_strategy: str = "hybrid",
drop_caption_prob: float = 0.0,
seed: int = 0,
) -> None:
self.manifest_path = Path(manifest_path)
self.image_root = Path(image_root)
self.image_transform = image_transform
self.caption_strategy = caption_strategy
self.drop_caption_prob = drop_caption_prob
self._rng = random.Random(seed)
with self.manifest_path.open("r", encoding="utf-8") as f:
self.entries: list[dict[str, Any]] = json.load(f)
self._validate_entries()
def _validate_entries(self) -> None:
"""Ensure all entries have required fields for the chosen strategy."""
required = {"drone_path", "sat_path"}
caption_sat_key = self._caption_key("sat")
caption_drone_key = self._caption_key("drone")
required |= {caption_sat_key, caption_drone_key}
for i, entry in enumerate(self.entries):
missing = required - entry.keys()
if missing:
raise KeyError(
f"Entry {i} (pair_id={entry.get('pair_id', '?')}) missing fields: "
f"{sorted(missing)}"
)
def _caption_key(self, view: str) -> str:
"""Resolve caption field name from strategy + view."""
if self.caption_strategy == "hybrid":
return f"caption_{view}"
return f"caption_{view}_{self.caption_strategy}"
def _load_image(self, relative_path: str) -> torch.Tensor:
"""Load image and apply preprocessing."""
path = self.image_root / relative_path
with Image.open(path) as img:
rgb = img.convert("RGB")
return self.image_transform(rgb)
def _maybe_drop(self, caption: str) -> str:
"""Stochastically drop caption to empty string for robustness training."""
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
return ""
return caption
def __len__(self) -> int:
return len(self.entries)
def __getitem__(self, idx: int) -> dict[str, Any]:
"""Return one pair with images and captions.
Args:
idx: Index into the manifest.
Returns:
Dict with:
- 'drone_img': [3, H, W] tensor
- 'sat_img': [3, H, W] tensor
- 'caption_drone': str (possibly empty)
- 'caption_sat': str (possibly empty)
- 'pair_id': str for logging
"""
entry = self.entries[idx]
drone_img = self._load_image(entry["drone_path"])
sat_img = self._load_image(entry["sat_path"])
caption_drone = self._maybe_drop(entry[self._caption_key("drone")])
caption_sat = self._maybe_drop(entry[self._caption_key("sat")])
return {
"drone_img": drone_img,
"sat_img": sat_img,
"caption_drone": caption_drone,
"caption_sat": caption_sat,
"pair_id": entry.get("pair_id", f"idx_{idx}"),
}
def collate_caption_batch(
batch: list[dict[str, Any]],
) -> dict[str, Any]:
"""Collate VisLocCaptionDataset items into a batched dict.
Images are stacked; captions remain Python lists so the tokenizer can
process them inside the model.forward().
Args:
batch: List of samples from VisLocCaptionDataset.__getitem__.
Returns:
Batched dict with stacked image tensors and caption lists.
"""
return {
"drone_img": torch.stack([b["drone_img"] for b in batch], dim=0),
"sat_img": torch.stack([b["sat_img"] for b in batch], dim=0),
"caption_drone": [b["caption_drone"] for b in batch],
"caption_sat": [b["caption_sat"] for b in batch],
"pair_ids": [b["pair_id"] for b in batch],
}

8
src/eval/__init__.py Normal file
View File

@@ -0,0 +1,8 @@
"""Evaluation utilities for caption quality test."""
from src.eval.evaluate import (
delta_r_at_1,
evaluate_retrieval,
run_evaluation_from_checkpoint,
)
__all__ = ["delta_r_at_1", "evaluate_retrieval", "run_evaluation_from_checkpoint"]

242
src/eval/evaluate.py Normal file
View File

@@ -0,0 +1,242 @@
from __future__ import annotations
"""Evaluation utilities for caption quality test.
Implements retrieval metrics across four directions and a
`delta_r_at_1` helper that compares caption-aware vs. image-only runs.
"""
import json
import logging
from pathlib import Path
import gin
import torch
from torch.utils.data import DataLoader
from src.models.dual_encoder import DualEncoderCaptionTest
LOGGER = logging.getLogger("caption_test.eval")
def _recall_at_k(
similarity: torch.Tensor,
k_values: tuple[int, ...] = (1, 5, 10),
) -> dict[int, float]:
"""Compute Recall@K assuming positives on the diagonal.
Args:
similarity: Pairwise similarity matrix [N_query, N_gallery].
k_values: Tuple of K values to compute.
Returns:
Dict mapping K -> recall in [0, 1].
"""
n_query = similarity.size(0)
targets = torch.arange(n_query, device=similarity.device)
sorted_idx = similarity.argsort(dim=1, descending=True)
result: dict[int, float] = {}
for k in k_values:
top_k = sorted_idx[:, :k]
hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
result[k] = float(hit.mean().item())
return result
@torch.no_grad()
def _encode_dataset(
model: DualEncoderCaptionTest,
loader: DataLoader,
device: str,
include_captions: bool,
) -> dict[str, torch.Tensor]:
"""Encode every sample in the loader into the shared embedding space.
Args:
model: Trained dual encoder.
loader: DataLoader yielding collated batches.
device: Target device string.
include_captions: If False, caption embeddings are skipped.
Returns:
Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat' -> [N, D].
"""
model.eval()
all_drone: list[torch.Tensor] = []
all_sat: list[torch.Tensor] = []
all_cap_drone: list[torch.Tensor] = []
all_cap_sat: list[torch.Tensor] = []
for batch in loader:
drone_img = batch["drone_img"].to(device, non_blocking=True)
sat_img = batch["sat_img"].to(device, non_blocking=True)
captions_drone = batch["caption_drone"] if include_captions else None
captions_sat = batch["caption_sat"] if include_captions else None
embeddings = model(
drone_img=drone_img,
sat_img=sat_img,
caption_drone=captions_drone,
caption_sat=captions_sat,
)
all_drone.append(embeddings["drone"].cpu())
all_sat.append(embeddings["sat"].cpu())
if include_captions:
all_cap_drone.append(embeddings["cap_drone"].cpu())
all_cap_sat.append(embeddings["cap_sat"].cpu())
out = {
"drone": torch.cat(all_drone, dim=0),
"sat": torch.cat(all_sat, dim=0),
}
if include_captions:
out["cap_drone"] = torch.cat(all_cap_drone, dim=0)
out["cap_sat"] = torch.cat(all_cap_sat, dim=0)
return out
def evaluate_retrieval(
model: DualEncoderCaptionTest,
loader: DataLoader,
device: str,
k_values: tuple[int, ...] = (1, 5, 10),
include_captions: bool = True,
) -> dict[str, float]:
"""Compute retrieval metrics across four directions.
Directions reported (when captions included):
drone -> sat, sat -> drone, text -> sat, text -> drone.
Args:
model: Trained DualEncoderCaptionTest.
loader: DataLoader over evaluation split.
device: torch device string.
k_values: Recall@K cutoffs.
include_captions: If False, only image-image directions computed.
Returns:
Flat dict with keys like 'r@1_drone_to_sat', 'r@5_text_to_sat', etc.
"""
feats = _encode_dataset(
model=model,
loader=loader,
device=device,
include_captions=include_captions,
)
metrics: dict[str, float] = {}
sim_d2s = feats["drone"] @ feats["sat"].t()
sim_s2d = sim_d2s.t()
for k, val in _recall_at_k(sim_d2s, k_values).items():
metrics[f"r@{k}_drone_to_sat"] = val
for k, val in _recall_at_k(sim_s2d, k_values).items():
metrics[f"r@{k}_sat_to_drone"] = val
if include_captions and "cap_sat" in feats and "cap_drone" in feats:
sim_t2s = feats["cap_sat"] @ feats["sat"].t()
sim_t2d = feats["cap_drone"] @ feats["drone"].t()
sim_tcd2tcs = feats["cap_drone"] @ feats["cap_sat"].t()
for k, val in _recall_at_k(sim_t2s, k_values).items():
metrics[f"r@{k}_text_to_sat"] = val
for k, val in _recall_at_k(sim_t2d, k_values).items():
metrics[f"r@{k}_text_to_drone"] = val
for k, val in _recall_at_k(sim_tcd2tcs, k_values).items():
metrics[f"r@{k}_capdrone_to_capsat"] = val
return metrics
def delta_r_at_1(
full_metrics: dict[str, float],
baseline_metrics: dict[str, float],
direction: str = "drone_to_sat",
) -> float:
"""Compute caption-quality proxy: R@1 gain from adding captions.
Args:
full_metrics: Metrics from training WITH caption losses.
baseline_metrics: Metrics from training WITHOUT caption losses.
direction: Retrieval direction to compare.
Returns:
Δ R@1 in [−1, +1] range (positive = captions help).
"""
key = f"r@1_{direction}"
if key not in full_metrics or key not in baseline_metrics:
raise KeyError(
f"Missing '{key}' in one of the metric dicts. "
f"Available full={list(full_metrics)}, baseline={list(baseline_metrics)}"
)
return full_metrics[key] - baseline_metrics[key]
@gin.configurable
def run_evaluation_from_checkpoint(
checkpoint_path: str,
test_manifest: str,
image_root: str,
output_path: str = "eval_report.json",
batch_size: int = 128,
num_workers: int = 4,
device: str = "cuda",
) -> dict[str, float]:
"""Standalone evaluation entry point (gin-configurable).
Args:
checkpoint_path: Path to .pt checkpoint from training.
test_manifest: Path to test manifest JSON.
image_root: Directory prefix for images.
output_path: Where to write the JSON report.
batch_size: Batch size for encoding.
num_workers: DataLoader workers.
device: torch device.
Returns:
Dict of retrieval metrics.
"""
from src.datasets.visloc_with_captions import (
VisLocCaptionDataset,
collate_caption_batch,
)
model = DualEncoderCaptionTest().to(device)
ckpt = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(ckpt["model_state"])
model.eval()
test_ds = VisLocCaptionDataset(
manifest_path=test_manifest,
image_root=image_root,
image_transform=model.preprocess,
)
test_loader = DataLoader(
test_ds,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
collate_fn=collate_caption_batch,
pin_memory=True,
)
metrics = evaluate_retrieval(
model=model,
loader=test_loader,
device=device,
)
report = {
"checkpoint": checkpoint_path,
"test_manifest": test_manifest,
"metrics": metrics,
}
out = Path(output_path)
out.parent.mkdir(parents=True, exist_ok=True)
with out.open("w", encoding="utf-8") as f:
json.dump(report, f, indent=2)
LOGGER.info("evaluation report saved to %s", out)
return metrics

8
src/losses/__init__.py Normal file
View File

@@ -0,0 +1,8 @@
"""Loss functions for caption quality test."""
from src.losses.multi_infonce import (
MultiTermInfoNCE,
cosine_temperature,
curriculum_lambdas,
)
__all__ = ["MultiTermInfoNCE", "cosine_temperature", "curriculum_lambdas"]

262
src/losses/multi_infonce.py Normal file
View File

@@ -0,0 +1,262 @@
from __future__ import annotations
"""Multi-term InfoNCE loss for caption quality validation.
Four InfoNCE terms over projected embeddings:
L = lambda_ii * L_img_img
+ lambda_sc * L_sat_cap
+ lambda_dc * L_drone_cap
+ lambda_cc * L_cap_cap
where L_img_img is the classical symmetric CVGL contrastive loss
with asymmetric weights (0.6 drone->sat + 0.4 sat->drone).
"""
import math
import gin
import torch
import torch.nn as nn
import torch.nn.functional as F
def _symmetric_info_nce(
emb_a: torch.Tensor,
emb_b: torch.Tensor,
temperature: float,
label_smoothing: float,
weight_a2b: float = 0.5,
weight_b2a: float = 0.5,
) -> torch.Tensor:
"""Compute weighted symmetric InfoNCE between two L2-normalized embeddings.
Args:
emb_a: First embedding set [B, D].
emb_b: Second embedding set [B, D]. Positive pairs are on the diagonal.
temperature: Softmax temperature (smaller = sharper distribution).
label_smoothing: Cross-entropy label smoothing epsilon.
weight_a2b: Weight for A-query direction.
weight_b2a: Weight for B-query direction.
Returns:
Scalar weighted loss.
"""
batch_size = emb_a.size(0)
logits = emb_a @ emb_b.t() / temperature
targets = torch.arange(batch_size, device=emb_a.device)
loss_a2b = F.cross_entropy(logits, targets, label_smoothing=label_smoothing)
loss_b2a = F.cross_entropy(logits.t(), targets, label_smoothing=label_smoothing)
return weight_a2b * loss_a2b + weight_b2a * loss_b2a
def cosine_temperature(
epoch: int,
total_epochs: int,
tau_init: float = 0.1,
tau_final: float = 0.01,
) -> float:
"""Cosine-decay schedule for InfoNCE temperature.
Args:
epoch: Current training epoch (0-indexed).
total_epochs: Total number of epochs.
tau_init: Initial temperature.
tau_final: Final temperature.
Returns:
Temperature value for this epoch.
"""
total_epochs = max(total_epochs, 1)
progress = min(max(epoch / total_epochs, 0.0), 1.0)
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
return tau_final + (tau_init - tau_final) * cosine
def curriculum_lambdas(
epoch: int,
warmup_epochs: int = 3,
text_ramp_epochs: int = 10,
lambda_ii: float = 1.0,
lambda_sc_max: float = 0.3,
lambda_dc_max: float = 0.3,
lambda_cc_max: float = 0.1,
) -> dict[str, float]:
"""Compute per-epoch loss weights under the curriculum schedule.
- Epochs 0..warmup_epochs: image-image only.
- Epochs warmup..text_ramp_epochs: linearly ramp sat-cap and drone-cap.
- Epochs >= text_ramp_epochs: full loss including caption-caption term.
Args:
epoch: Current epoch (0-indexed).
warmup_epochs: Number of warmup epochs (no text losses).
text_ramp_epochs: Epoch when text losses reach max.
lambda_ii: Constant weight for image-image loss.
lambda_sc_max: Max weight for satellite-caption loss.
lambda_dc_max: Max weight for drone-caption loss.
lambda_cc_max: Max weight for caption-caption loss.
Returns:
Dict with keys 'img_img', 'sat_cap', 'drone_cap', 'cap_cap'.
"""
if epoch < warmup_epochs:
ramp = 0.0
elif epoch >= text_ramp_epochs:
ramp = 1.0
else:
denom = max(text_ramp_epochs - warmup_epochs, 1)
ramp = (epoch - warmup_epochs) / denom
return {
"img_img": lambda_ii,
"sat_cap": lambda_sc_max * ramp,
"drone_cap": lambda_dc_max * ramp,
"cap_cap": lambda_cc_max * ramp,
}
@gin.configurable
class MultiTermInfoNCE(nn.Module):
"""Multi-term InfoNCE loss with curriculum and cosine temperature.
Produces total loss and per-component diagnostics. All inputs must be
L2-normalized embeddings of the same dimension.
Args:
temperature_init: Initial temperature (epoch 0).
temperature_final: Final temperature after cosine decay.
label_smoothing: Cross-entropy label smoothing epsilon.
asym_drone_to_sat: Weight for drone->sat InfoNCE direction.
asym_sat_to_drone: Weight for sat->drone InfoNCE direction.
warmup_epochs: Epochs with image-image loss only.
text_ramp_epochs: Epoch at which text losses reach max.
lambda_ii: Constant weight for image-image loss.
lambda_sc_max: Max weight for sat-caption loss.
lambda_dc_max: Max weight for drone-caption loss.
lambda_cc_max: Max weight for caption-caption loss.
"""
def __init__(
self,
temperature_init: float = 0.1,
temperature_final: float = 0.01,
label_smoothing: float = 0.1,
asym_drone_to_sat: float = 0.6,
asym_sat_to_drone: float = 0.4,
warmup_epochs: int = 3,
text_ramp_epochs: int = 10,
lambda_ii: float = 1.0,
lambda_sc_max: float = 0.3,
lambda_dc_max: float = 0.3,
lambda_cc_max: float = 0.1,
) -> None:
super().__init__()
self.temperature_init = temperature_init
self.temperature_final = temperature_final
self.label_smoothing = label_smoothing
self.asym_drone_to_sat = asym_drone_to_sat
self.asym_sat_to_drone = asym_sat_to_drone
self.warmup_epochs = warmup_epochs
self.text_ramp_epochs = text_ramp_epochs
self.lambda_ii = lambda_ii
self.lambda_sc_max = lambda_sc_max
self.lambda_dc_max = lambda_dc_max
self.lambda_cc_max = lambda_cc_max
def forward(
self,
embeddings: dict[str, torch.Tensor],
epoch: int,
total_epochs: int,
) -> dict[str, torch.Tensor]:
"""Compute multi-term loss.
Args:
embeddings: Dict with keys 'drone', 'sat', and optionally
'cap_drone', 'cap_sat'. Each [B, D] L2-normalized.
epoch: Current epoch (0-indexed).
total_epochs: Total epochs for temperature schedule.
Returns:
Dict with scalar tensors: 'total', 'img_img', 'sat_cap',
'drone_cap', 'cap_cap', plus 'temperature' and 'lambdas'.
"""
tau = cosine_temperature(
epoch=epoch,
total_epochs=total_epochs,
tau_init=self.temperature_init,
tau_final=self.temperature_final,
)
lambdas = curriculum_lambdas(
epoch=epoch,
warmup_epochs=self.warmup_epochs,
text_ramp_epochs=self.text_ramp_epochs,
lambda_ii=self.lambda_ii,
lambda_sc_max=self.lambda_sc_max,
lambda_dc_max=self.lambda_dc_max,
lambda_cc_max=self.lambda_cc_max,
)
drone = embeddings["drone"]
sat = embeddings["sat"]
# Image-image symmetric InfoNCE with asymmetric weights.
loss_ii = _symmetric_info_nce(
emb_a=drone,
emb_b=sat,
temperature=tau,
label_smoothing=self.label_smoothing,
weight_a2b=self.asym_drone_to_sat,
weight_b2a=self.asym_sat_to_drone,
)
loss_sc = torch.zeros_like(loss_ii)
loss_dc = torch.zeros_like(loss_ii)
loss_cc = torch.zeros_like(loss_ii)
if "cap_sat" in embeddings and lambdas["sat_cap"] > 0:
loss_sc = _symmetric_info_nce(
emb_a=sat,
emb_b=embeddings["cap_sat"],
temperature=tau,
label_smoothing=self.label_smoothing,
)
if "cap_drone" in embeddings and lambdas["drone_cap"] > 0:
loss_dc = _symmetric_info_nce(
emb_a=drone,
emb_b=embeddings["cap_drone"],
temperature=tau,
label_smoothing=self.label_smoothing,
)
if (
"cap_drone" in embeddings
and "cap_sat" in embeddings
and lambdas["cap_cap"] > 0
):
loss_cc = _symmetric_info_nce(
emb_a=embeddings["cap_drone"],
emb_b=embeddings["cap_sat"],
temperature=tau,
label_smoothing=self.label_smoothing,
)
total = (
lambdas["img_img"] * loss_ii
+ lambdas["sat_cap"] * loss_sc
+ lambdas["drone_cap"] * loss_dc
+ lambdas["cap_cap"] * loss_cc
)
return {
"total": total,
"img_img": loss_ii.detach(),
"sat_cap": loss_sc.detach(),
"drone_cap": loss_dc.detach(),
"cap_cap": loss_cc.detach(),
"temperature": torch.tensor(tau, device=total.device),
"lambda_ii": torch.tensor(lambdas["img_img"], device=total.device),
"lambda_sc": torch.tensor(lambdas["sat_cap"], device=total.device),
"lambda_dc": torch.tensor(lambdas["drone_cap"], device=total.device),
"lambda_cc": torch.tensor(lambdas["cap_cap"], device=total.device),
}

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]

1
src/training/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Training loop for caption quality test."""

313
src/training/train.py Normal file
View File

@@ -0,0 +1,313 @@
from __future__ import annotations
"""Training loop for caption quality validation on UAV-VisLoc.
Uses gin-configurable DualEncoderCaptionTest + MultiTermInfoNCE.
Logs per-component losses, temperature, and lambdas each step.
Saves checkpoint + eval snapshot every epoch.
"""
import argparse
import json
import logging
import time
from pathlib import Path
import gin
import torch
import torch.nn as nn
from torch.amp import GradScaler, autocast
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader
from src.datasets.visloc_with_captions import (
VisLocCaptionDataset,
collate_caption_batch,
)
from src.eval.evaluate import evaluate_retrieval
from src.losses.multi_infonce import MultiTermInfoNCE
from src.models.dual_encoder import DualEncoderCaptionTest
LOGGER = logging.getLogger("caption_test.train")
@gin.configurable
class TrainConfig:
"""Top-level training configuration (gin-configurable).
Args:
train_manifest: Path to training manifest JSON.
val_manifest: Path to validation manifest JSON.
image_root: Directory prefix for images.
output_dir: Where to save checkpoints and logs.
epochs: Number of training epochs.
batch_size: Mini-batch size.
num_workers: DataLoader worker count.
learning_rate: AdamW initial LR.
weight_decay: AdamW weight decay.
grad_clip: Max gradient norm for clipping (0 disables).
use_amp: Enable fp16 mixed-precision training.
eval_every: Run validation every N epochs.
seed: Random seed.
device: torch device.
"""
def __init__(
self,
train_manifest: str = "data/visloc_train.json",
val_manifest: str = "data/visloc_val.json",
image_root: str = "data/visloc/images",
output_dir: str = "out/caption_test",
epochs: int = 30,
batch_size: int = 128,
num_workers: int = 4,
learning_rate: float = 1e-4,
weight_decay: float = 1e-4,
grad_clip: float = 1.0,
use_amp: bool = True,
eval_every: int = 1,
seed: int = 42,
device: str = "cuda",
) -> None:
self.train_manifest = train_manifest
self.val_manifest = val_manifest
self.image_root = image_root
self.output_dir = Path(output_dir)
self.epochs = epochs
self.batch_size = batch_size
self.num_workers = num_workers
self.learning_rate = learning_rate
self.weight_decay = weight_decay
self.grad_clip = grad_clip
self.use_amp = use_amp
self.eval_every = eval_every
self.seed = seed
self.device = device
def _set_seed(seed: int) -> None:
"""Seed Python, NumPy and PyTorch RNGs."""
import random as _random
import numpy as _np
_random.seed(seed)
_np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def _atomic_save(obj: dict, path: Path) -> None:
"""Write torch checkpoint atomically (temp file + rename)."""
path.parent.mkdir(parents=True, exist_ok=True)
tmp_path = path.with_suffix(path.suffix + ".tmp")
torch.save(obj, tmp_path)
tmp_path.replace(path)
def _step_loss(
model: DualEncoderCaptionTest,
loss_fn: MultiTermInfoNCE,
batch: dict,
epoch: int,
total_epochs: int,
device: str,
use_amp: bool,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
"""Single training forward pass returning (total_loss, diagnostics)."""
drone_img = batch["drone_img"].to(device, non_blocking=True)
sat_img = batch["sat_img"].to(device, non_blocking=True)
caption_drone = batch["caption_drone"]
caption_sat = batch["caption_sat"]
with autocast(device_type="cuda", enabled=use_amp):
embeddings = model(
drone_img=drone_img,
sat_img=sat_img,
caption_drone=caption_drone,
caption_sat=caption_sat,
)
loss_dict = loss_fn(
embeddings=embeddings,
epoch=epoch,
total_epochs=total_epochs,
)
return loss_dict["total"], loss_dict
def train(config_path: str) -> None:
"""Run the full training loop driven by gin configuration.
Args:
config_path: Path to .gin config file.
"""
gin.parse_config_file(config_path)
cfg = TrainConfig()
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(name)s %(levelname)s %(message)s",
)
_set_seed(cfg.seed)
cfg.output_dir.mkdir(parents=True, exist_ok=True)
# Model + loss
model = DualEncoderCaptionTest().to(cfg.device)
loss_fn = MultiTermInfoNCE().to(cfg.device)
# Datasets use the same preprocess function the model already holds.
preprocess = model.preprocess
train_ds = VisLocCaptionDataset(
manifest_path=cfg.train_manifest,
image_root=cfg.image_root,
image_transform=preprocess,
)
val_ds = VisLocCaptionDataset(
manifest_path=cfg.val_manifest,
image_root=cfg.image_root,
image_transform=preprocess,
)
train_loader = DataLoader(
train_ds,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers,
collate_fn=collate_caption_batch,
pin_memory=True,
drop_last=True,
)
val_loader = DataLoader(
val_ds,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers,
collate_fn=collate_caption_batch,
pin_memory=True,
)
optimizer = AdamW(
model.trainable_parameters(),
lr=cfg.learning_rate,
weight_decay=cfg.weight_decay,
)
scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
scaler = GradScaler(enabled=cfg.use_amp)
history: list[dict] = []
for epoch in range(cfg.epochs):
model.train()
epoch_start = time.time()
agg: dict[str, float] = {}
n_batches = 0
for batch in train_loader:
optimizer.zero_grad(set_to_none=True)
total_loss, loss_dict = _step_loss(
model=model,
loss_fn=loss_fn,
batch=batch,
epoch=epoch,
total_epochs=cfg.epochs,
device=cfg.device,
use_amp=cfg.use_amp,
)
scaler.scale(total_loss).backward()
if cfg.grad_clip > 0:
scaler.unscale_(optimizer)
nn.utils.clip_grad_norm_(
model.trainable_parameters(),
max_norm=cfg.grad_clip,
)
scaler.step(optimizer)
scaler.update()
# Accumulate diagnostics.
for key, tensor_val in loss_dict.items():
agg[key] = agg.get(key, 0.0) + float(tensor_val.item())
n_batches += 1
scheduler.step()
elapsed = time.time() - epoch_start
means = {k: v / max(n_batches, 1) for k, v in agg.items()}
LOGGER.info(
"epoch=%d time=%.1fs lr=%.2e total=%.4f img_img=%.4f "
"sat_cap=%.4f drone_cap=%.4f cap_cap=%.4f tau=%.4f",
epoch,
elapsed,
optimizer.param_groups[0]["lr"],
means.get("total", 0.0),
means.get("img_img", 0.0),
means.get("sat_cap", 0.0),
means.get("drone_cap", 0.0),
means.get("cap_cap", 0.0),
means.get("temperature", 0.0),
)
epoch_record = {
"epoch": epoch,
"elapsed_seconds": elapsed,
"train": means,
}
# Validation.
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
model.eval()
val_metrics = evaluate_retrieval(
model=model,
loader=val_loader,
device=cfg.device,
)
epoch_record["val"] = val_metrics
LOGGER.info(
"val epoch=%d R@1_d2s=%.4f R@1_s2d=%.4f "
"R@1_t2s=%.4f R@1_t2d=%.4f",
epoch,
val_metrics.get("r@1_drone_to_sat", 0.0),
val_metrics.get("r@1_sat_to_drone", 0.0),
val_metrics.get("r@1_text_to_sat", 0.0),
val_metrics.get("r@1_text_to_drone", 0.0),
)
history.append(epoch_record)
# Checkpoint per epoch.
_atomic_save(
obj={
"epoch": epoch,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"config_path": config_path,
},
path=cfg.output_dir / f"ckpt_epoch{epoch:03d}.pt",
)
# Save training history.
history_path = cfg.output_dir / "history.json"
with history_path.open("w", encoding="utf-8") as f:
json.dump(history, f, indent=2)
LOGGER.info("training complete, history saved to %s", history_path)
def main() -> None:
parser = argparse.ArgumentParser(description="Caption quality test training.")
parser.add_argument(
"--config",
type=str,
required=True,
help="Path to gin configuration file.",
)
args = parser.parse_args()
train(config_path=args.config)
if __name__ == "__main__":
main()