commit 2ce4017ebd128d779d95506895487a456192ea40 Author: Pikaliov Date: Fri Apr 17 00:04:46 2026 +0300 Initial commit: caption quality test on UAV-VisLoc 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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ca6a871 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..aa6459a --- /dev/null +++ b/README.md @@ -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) diff --git a/conf/balanced.gin b/conf/balanced.gin new file mode 100644 index 0000000..81a489c --- /dev/null +++ b/conf/balanced.gin @@ -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" diff --git a/conf/baseline_no_text.gin b/conf/baseline_no_text.gin new file mode 100644 index 0000000..a2fddc3 --- /dev/null +++ b/conf/baseline_no_text.gin @@ -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" diff --git a/conf/text_heavy.gin b/conf/text_heavy.gin new file mode 100644 index 0000000..1a69cbd --- /dev/null +++ b/conf/text_heavy.gin @@ -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" diff --git a/scripts/compare_runs.py b/scripts/compare_runs.py new file mode 100644 index 0000000..233098e --- /dev/null +++ b/scripts/compare_runs.py @@ -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() diff --git a/scripts/generate_captions.py b/scripts/generate_captions.py new file mode 100644 index 0000000..ba2e310 --- /dev/null +++ b/scripts/generate_captions.py @@ -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() diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..be1bf21 --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +"""Caption quality test package.""" diff --git a/src/datasets/__init__.py b/src/datasets/__init__.py new file mode 100644 index 0000000..ff27a90 --- /dev/null +++ b/src/datasets/__init__.py @@ -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"] diff --git a/src/datasets/visloc_with_captions.py b/src/datasets/visloc_with_captions.py new file mode 100644 index 0000000..9fff8dd --- /dev/null +++ b/src/datasets/visloc_with_captions.py @@ -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], + } diff --git a/src/eval/__init__.py b/src/eval/__init__.py new file mode 100644 index 0000000..e26b79e --- /dev/null +++ b/src/eval/__init__.py @@ -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"] diff --git a/src/eval/evaluate.py b/src/eval/evaluate.py new file mode 100644 index 0000000..68936c8 --- /dev/null +++ b/src/eval/evaluate.py @@ -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 diff --git a/src/losses/__init__.py b/src/losses/__init__.py new file mode 100644 index 0000000..d74b5ee --- /dev/null +++ b/src/losses/__init__.py @@ -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"] diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py new file mode 100644 index 0000000..33ee38e --- /dev/null +++ b/src/losses/multi_infonce.py @@ -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), + } diff --git a/src/models/__init__.py b/src/models/__init__.py new file mode 100644 index 0000000..c0bc028 --- /dev/null +++ b/src/models/__init__.py @@ -0,0 +1,4 @@ +"""Model components for caption quality test.""" +from src.models.dual_encoder import DualEncoderCaptionTest, ProjectionHead + +__all__ = ["DualEncoderCaptionTest", "ProjectionHead"] diff --git a/src/models/dual_encoder.py b/src/models/dual_encoder.py new file mode 100644 index 0000000..76dad95 --- /dev/null +++ b/src/models/dual_encoder.py @@ -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] diff --git a/src/training/__init__.py b/src/training/__init__.py new file mode 100644 index 0000000..d994270 --- /dev/null +++ b/src/training/__init__.py @@ -0,0 +1 @@ +"""Training loop for caption quality test.""" diff --git a/src/training/train.py b/src/training/train.py new file mode 100644 index 0000000..1b475f3 --- /dev/null +++ b/src/training/train.py @@ -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()