From abb3337f8d8bfd09e791b844d8f0158d8ce9aa4a Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 17 Apr 2026 17:13:00 +0300 Subject: [PATCH] Rewrite: GatedFusion architecture + UAV-GeoLoc dataset Architecture v2: - Query branch: drone + text -> GatedFusion -> proj -> query_emb - Gallery branch: satellite -> proj -> gallery_emb - Single InfoNCE loss (asymmetric 0.6/0.4) - GatedFusion: learnable gated addition (sigma(alpha)*img + (1-sigma(alpha))*text) - Baseline mode: gate=1.0 (text ignored) Dataset: - UAV-GeoLoc loader with template captions from path metadata - 27 terrain types with predefined features - Random positive crop sampling per epoch Configs: balanced (gate=0.7), baseline (no text), text_heavy (gate=0.3) Co-Authored-By: Claude Opus 4.6 (1M context) --- CLAUDE.md | 208 +++++++++++++++ README.md | 140 ++++------ UAV-VisLoc_Dataset_Analysis.md | 378 +++++++++++++++++++++++++++ conf/balanced.gin | 44 ++-- conf/baseline_no_text.gin | 11 +- conf/text_heavy.gin | 10 +- scripts/compare_runs.py | 59 ++--- src/datasets/visloc_with_captions.py | 245 ++++++++++------- src/eval/evaluate.py | 161 +++--------- src/losses/multi_infonce.py | 208 +++------------ src/models/dual_encoder.py | 222 +++++++--------- src/training/train.py | 172 +++++------- 12 files changed, 1077 insertions(+), 781 deletions(-) create mode 100644 CLAUDE.md create mode 100644 UAV-VisLoc_Dataset_Analysis.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..7d64807 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,208 @@ +# Caption Quality Test for Cross-View Geo-Localization + +## Архитектура системы (v2, 2026-04-17) + +``` +QUERY BRANCH (drone + caption): + drone_img --> GeoRSCLIP ViT-B/32 (frozen) --> drone_feat [B,512] + caption --> GeoRSCLIP Text (partial unfreeze) --> text_feat [B,512] + | + GatedFusion: q = sigma(alpha)*drone + (1-sigma(alpha))*text + | + proj_query (Linear 512->512) --> L2-norm --> query [B,512] + +GALLERY BRANCH (satellite only): + sat_img --> GeoRSCLIP ViT-B/32 (frozen) --> sat_feat [B,512] + | + proj_gallery (Linear 512->512) --> L2-norm --> gallery [B,512] + +LOSS: InfoNCE(query, gallery) — symmetric, asymmetric weights (0.6 q->g, 0.4 g->q) +BASELINE: gate = 1.0 (text ignored) +``` + +### Trainable parameters: ~1.2M из ~151M (proj_query + proj_gallery + fusion alpha + text last_block) + +## Ключевые файлы + +| Файл | Назначение | +|------|-----------| +| `src/models/dual_encoder.py` | GeoRSCLIP + GatedFusion + projection heads | +| `src/losses/multi_infonce.py` | InfoNCE с cosine temperature schedule | +| `src/datasets/visloc_with_captions.py` | UAV-GeoLoc loader + template captions из path metadata | +| `src/training/train.py` | Training loop, логирование loss/gate/tau | +| `src/eval/evaluate.py` | R@K metrics, delta_r_at_1 | +| `scripts/compare_runs.py` | Markdown/JSON сравнение baseline vs caption runs | +| `scripts/generate_captions.py` | Offline caption generation (template/VLM/hybrid) | + +## Backbone: GeoRSCLIP ViT-B/32 + +- **Checkpoint:** `checkpoints/RS5M_ViT-B-32.pt` (скачать с github.com/om-ai-lab/RS5M) +- **Image encoder:** ViT-B/32, 224x224, 512-dim, ~86M params — frozen +- **Text encoder:** CLIP text transformer, 77 tokens, 512-dim — partial unfreeze (last_block + text_projection) +- **Throughput:** ~4000 img/s на RTX 4090 (AMP, batch 128) +- **Выбран вместо SigLIP 2** (ViT-SO400M, 384px, ~400M): в 7-10x быстрее, domain-specific (обучен на 5M RS-изображений), больше batch = больше негативов в InfoNCE + +## GatedFusion + +- `query = sigma(alpha) * drone_feat + (1 - sigma(alpha)) * text_feat` +- `alpha` — один learnable scalar в logit-space +- `init_gate = 0.7` → начальный вес image = 70%, text = 30% +- `baseline_mode = True` → gate = 1.0, text полностью игнорируется +- Gate value логируется каждую эпоху для интерпретации вклада текста + +## Датасет: UAV-GeoLoc + +- **Путь:** `/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/` +- **Train:** 206,108 queries, 94,709 DB crops (140 scenes, Terrain split) +- **Val:** 62,368 queries, 26,597 DB crops (40 scenes) +- **Test:** 33,472 queries, 11,684 DB crops (20 scenes) +- **Структура:** `Terrain/{type}/{scene}/query/height{N}_rot{M}/footage/{file}.jpeg` +- **Index:** `Index/train_query.txt` — `query_path 0 pos_crop1 pos_crop2 ...` + +### Template captions (из path metadata) + +Формат: `"Aerial view at {height}m facing {heading} over {terrain} terrain near {scene}. Plan-view features: {features}."` + +Пример: `"Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano. Plan-view features: lava flows, crater edges, volcanic rock."` + +Metadata извлекается из пути: +- `Terrain/Volcano/KilaueaVolcano/query/height100_rot315/...` → terrain=Volcano, scene=KilaueaVolcano, height=100, heading=northwest +- 27 terrain типов с predefined features (Volcano, Mountain, Hill, Desert, Plain, ...) +- Country subset: features = "buildings, roads, urban blocks, rooftops, intersections" + +## Конфигурации (gin) + +| Конфиг | Gate init | Описание | +|--------|-----------|----------| +| `conf/balanced.gin` | 0.7 (30% text) | **Primary test** | +| `conf/baseline_no_text.gin` | 1.0 (no text) | Reference baseline | +| `conf/text_heavy.gin` | 0.3 (70% text) | Stress test | + +Общие параметры: 10 epochs, batch 128, lr=1e-4, AMP, cosine LR schedule, eval every 2 epochs. + +## Запуск + +```bash +# 1. Baseline (no text) +python -m src.training.train --config conf/baseline_no_text.gin + +# 2. With captions (primary test) +python -m src.training.train --config conf/balanced.gin + +# 3. Text-heavy (stress test) +python -m src.training.train --config conf/text_heavy.gin + +# 4. Compare +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 + +**Primary metric:** Delta R@1 (query -> gallery) + +| Delta R@1 | Verdict | +|-----------|---------| +| >= +3% | PASS — captions informative, proceed to production | +| +1% to +3% | MARGINAL — add VLM refinement, re-run | +| 0 to +1% | WEAK — redesign caption pipeline | +| < 0 | HARMFUL — critical bug | + +**Logged per epoch:** loss, temperature (tau), gate value (sigma(alpha)), lr +**Eval metrics:** R@1, R@5, R@10 для query->gallery и gallery->query + +## Бюджет времени (RTX 4090, 24 GB) + +| Фаза | Время | +|------|-------| +| Один training run (10 epochs, 206K queries, batch 128) | ~15-30 мин | +| Full test (3 варианта) | ~1-1.5 ч | +| Evaluation per run | ~2-5 мин | + +## Связанные проекты + +### Text Annotation Pipeline +- **Путь:** `/home/servml/Документы/pikaliov_obsidian/(Полякова ВЕ_Система для генерации текстовых описаний для БПЛА)/2_work/2_text_annotation/code/` +- **VLM:** Qwen3-VL-8B AWQ (1.68 s/img) +- **Scoring:** SigLIP 2 (drone P3) + CLIP-RSICD (satellite P1+P2) +- **Формат описаний:** 3 абзаца (P1 Inventory + P2 Spatial Map + P3 Fingerprint) +- **Метрики:** FDR, FNR, NumAcc, LLaVA-Critic C1-C6 + +### UAV-VisLoc Prepare +- **Путь:** `/home/servml/Документы/code/Yaroslav/UAV-VisLoc-prepare/scripts/prepare_dataset.py` +- **Статус:** код готов, ещё не запускался +- **Задача:** нарезка satellite кропов 512x512, stride 256 для UAV-VisLoc dataset +- **Подробности:** см. ниже + +--- + +## Датасеты (справочник) + +### UAV-VisLoc +- **Путь:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/` +- **Структура:** 11 маршрутов (папки `01`-`11`), каждая содержит: + - `drone/` — drone-снимки (`XX_NNNN.JPG`) + - `satelliteXX.tif` — спутниковая карта + - `XX.csv` — GPS-метаданные drone (num, filename, date, lat, lon, height, Omega, Kappa, Phi1, Phi2) +- **Исключение:** маршрут `09` — спутник разбит на 4 тайла (`satellite09_01-01.tif` и т.д.) +- **Satellite coordinates:** `satellite_ coordinates_range.csv` — bbox каждой карты (LT_lat_map, LT_lon_map, RB_lat_map, RB_lon_map) +- **Splits:** `visloc_train.csv`, `visloc_test.csv` — списки drone-снимков (TSV, full absolute paths) +- **Размеры:** Drone 3976x2652 / 3000x2000 (6774 снимков), Satellite от 3000x170 до 43421x38408 +- **GSD спутника:** ~0.30 м/px (единый zoom level Google Earth для всех карт). GSD по долготе варьируется 0.23-0.27 м/px из-за косинусного эффекта широты (40°N vs 25°N), но это не разная высота съёмки. Кроп 512x512 покрывает ~154x154 м везде. + +### UAV-GeoLoc +- **Путь:** `/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/` +- **Подмножества:** Country (171 scene), Terrain (200 scenes), Rot (1 scene) +- **Формат пар:** `positive.json`, `semi_positive.json`, `db_postion.txt` +- **Index:** `train_query.txt` — `query_path label pos_crop1 pos_crop2 ...` +- **Drone:** синтетика Google Earth Studio 3D, 512x512, FOV 30deg, heights 100/125/150m +- **Satellite:** кропы из merge.tif, доминирующий размер 200x200 + +## Скрипт подготовки UAV-VisLoc + +- **Путь:** `/home/servml/Документы/code/Yaroslav/UAV-VisLoc-prepare/scripts/prepare_dataset.py` +- **Статус:** код готов, ещё не запускался + +### Запуск +```bash +python scripts/prepare_dataset.py \ + --src /home/servml/Документы/datasets/UAV_VisLoc_dataset \ + --dst /home/servml/Документы/datasets/UAV_VisLoc_processed \ + --crop-size 512 --stride 256 --target-size 256 +``` + +### Pipeline +1. Resize drone -> 256x256 (JPEG, quality=95) +2. Stitch satellite tiles для маршрута 09 (4 тайла -> 44800x33280) +3. Нарезка satellite -> кропы 512x512, stride 256, resize -> 256x256 (PNG) +4. GPS для каждого кропа из bbox карты + позиция в grid +5. Match drone->crop через vectorized haversine (positive = ближайший, semi-positive = +/-1 в grid) +6. Metadata: `positive.json`, `semi_positive.json`, `db_postion.txt` (per route) +7. Index: `train_query.txt`, `test_query.txt`, `train_db.txt`, `test_db.txt`, `all_db.txt` + +### Форматы выходных файлов (совместимость с UAV-GeoLoc) +| Файл | Формат | +|------|--------| +| `positive.json` | `{frame_id: [crop_name]}`, ключ = frame ID без route prefix (`"0001"`) | +| `semi_positive.json` | `{frame_id: [crop1, crop2, ...]}`, соседи +/-1 в grid | +| `db_postion.txt` | tab-separated: `name\tlon\tlat\tscale_lon\tscale_lat` | +| `train_query.txt` | `route/drone/file.JPG 0 route/DB/img/crop1.png ...` | +| `train_db.txt` / `test_db.txt` | все кропы всех маршрутов (gallery одинаковая, split по query) | + +### Ожидаемые объёмы +- Drone: 6744 (без маршрута 07: 30 images excluded) +- Satellite кропов: ~74,807 +- Память: до ~4.5 GB RAM (маршрут 09 stitched 44800x33280) + +### Ревью и исправления (2026-04-17) +1. `train_db.txt`/`test_db.txt` содержали только matched кропы -> теперь все ~74K (полная gallery) +2. `db_position.txt` -> `db_postion.txt` (совместимость с UAV-GeoLoc), добавлены scale_lon/scale_lat, tab-separator +3. `positive.json` ключи были filename (`01_0001.JPG`) -> теперь frame_id (`0001`) +4. Semi-positive поиск O(n) -> O(1) через dict lookup по (x,y) grid +5. Удален мёртвый код (`haversine_m`, `defaultdict` import) + +### Известные ограничения +- Нет val split (только train/test, как в оригинальном UAV-VisLoc) +- Большие спутниковые карты загружаются целиком в RAM (до 4.5 GB для route 09) diff --git a/README.md b/README.md index aa6459a..f201bf0 100644 --- a/README.md +++ b/README.md @@ -1,33 +1,42 @@ -# Caption Quality Test on UAV-VisLoc +# Caption Quality Test for Cross-View Geo-Localization -Validate generated text captions by measuring retrieval R@1 lift on UAV-VisLoc. -Uses GeoRSCLIP ViT-B/32 dual encoder with multi-term InfoNCE. +Validate whether generated text captions improve retrieval R@1 in cross-view +geo-localization (drone-to-satellite). Uses GeoRSCLIP ViT-B/32 dual encoder +with GatedFusion on the query branch. -Full analysis: `2_hypotesis/АНАЛИЗ_caption_quality_test_VisLoc.md` +## Architecture + +``` +Query: drone_img + caption -> GatedFusion -> proj -> query_emb +Gallery: sat_img -> proj -> gallery_emb +Loss: InfoNCE(query, gallery) +``` + +Baseline: fusion gate = 1.0 (text ignored). ## 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) +│ ├── balanced.gin # Primary: gate init 0.7 (30% text) +│ ├── baseline_no_text.gin # Reference: gate = 1.0 (no text) +│ └── text_heavy.gin # Stress: gate init 0.3 (70% text) ├── scripts/ -│ ├── generate_captions.py # Offline caption generation (template/VLM/hybrid) -│ └── compare_runs.py # Δ R@1 comparison report builder +│ ├── generate_captions.py # Offline caption generation +│ └── compare_runs.py # Delta R@1 comparison report ├── src/ │ ├── datasets/ -│ │ └── visloc_with_captions.py +│ │ └── visloc_with_captions.py # UAV-GeoLoc loader + template captions │ ├── models/ -│ │ └── dual_encoder.py # GeoRSCLIP wrapper, projection heads +│ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion + projection heads │ ├── losses/ -│ │ └── multi_infonce.py # 4-term InfoNCE + curriculum + cosine τ +│ │ └── multi_infonce.py # InfoNCE with cosine temperature │ ├── training/ -│ │ └── train.py # Main loop +│ │ └── train.py # Main training loop │ └── eval/ -│ └── evaluate.py # R@K metrics, Δ R@1 helper -└── data/ # (user-provided) VisLoc pairs + captions +│ └── evaluate.py # R@K metrics, Delta R@1 +└── checkpoints/ # RS5M_ViT-B-32.pt (user-provided) ``` ## Prerequisites @@ -40,99 +49,66 @@ Pillow numpy ``` -GeoRSCLIP ViT-B/32 checkpoint: download `RS5M_ViT-B-32.pt` from +GeoRSCLIP checkpoint: download `RS5M_ViT-B-32.pt` from `github.com/om-ai-lab/RS5M` and place under `checkpoints/`. ## Workflow -### 1. Generate captions +### 1. Train baseline (no text) ```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) +### 2. Train with captions + +```bash 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 +### 3. 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 + --full_report out/caption_test/balanced/eval_report.json \ + --output out/caption_test/comparison.md ``` -## Decision rule (from `compare_runs.py`) +## Decision rule -| Δ R@1 (drone→sat) | Verdict | +| Delta R@1 (query->gallery) | 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 | +| >= +3% | PASS -- captions informative, proceed to production | +| +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 | +| Single training run (10 epochs, batch 128, 206K queries) | ~15-30 min | +| Full test (baseline + balanced + text_heavy) | ~1-1.5 h | +| Evaluation | ~2-5 min per run | -## Notes on code style +## Dataset -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. +UAV-GeoLoc Terrain split (from `/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/`): +- Train: 206,108 queries, 94,709 DB crops (140 scenes) +- Val: 62,368 queries, 26,597 DB crops (40 scenes) +- Test: 33,472 queries, 11,684 DB crops (20 scenes) -## Relation to NADEZHDA main pipeline +Template captions generated automatically from path metadata: +``` +"Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano. + Plan-view features: lava flows, crater edges, volcanic rock." +``` -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`). +## Code style -## 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) +- `from __future__ import annotations` everywhere +- Type hints on all signatures +- Google-style docstrings +- `@gin.configurable` on top-level classes +- No emojis in code, English-only comments diff --git a/UAV-VisLoc_Dataset_Analysis.md b/UAV-VisLoc_Dataset_Analysis.md new file mode 100644 index 0000000..e1e4457 --- /dev/null +++ b/UAV-VisLoc_Dataset_Analysis.md @@ -0,0 +1,378 @@ +# АНАЛИЗ ДАТАСЕТА: UAV-VisLoc + +**Дата анализа:** 2026-04-17 +**Метод:** Эмпирический анализ данных на диске + статья arXiv:2405.11936 + GitHub-репозиторий авторов +**Путь к данным:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/` + +--- + +## 1. МЕТАДАННЫЕ + +| Поле | Значение | +|------|----------| +| Полное название | UAV-VisLoc: A Large-scale Dataset for UAV Visual Localization | +| Авторы | Wenjia Xu, Yaxuan Yao, Jiaqi Cao, Zhiwei Wei, Chunbo Liu, Jiuniu Wang, Mugen Peng (BUPT + CAS + CityU HK) | +| Год, Venue | 2024, arXiv:2405.11936 [cs.CV] | +| Код | https://github.com/IntelliSensing/UAV-VisLoc (только README + ссылки) | +| Данные | Google Drive / Baidu Net Disk (16.4 GB) | +| Общий объём на диске | ~16.4 GB | + +--- + +## 2. ОБЩАЯ СТАТИСТИКА + +### 2.1. Сводка + +| Параметр | Значение | +|----------|----------| +| Drone-изображений | **6 774** | +| Спутниковых карт | **11** (+ маршрут 09 разбит на 4 тайла) | +| Маршрутов (flights) | 11 | +| Регионов (Китай) | 7 провинций/районов | +| Типов БПЛА | 2 (multi-rotor + fixed-wing) | +| Сезоны съёмки | 2 (лето, осень) | +| Временной охват | 2016–2023 | + +### 2.2. Разбиение train / test + +| Split | Изображений | Доля | +|-------|-------------|------| +| Train | 5 080 | 75.0% | +| Test | 1 694 | 25.0% | +| **Итого** | **6 774** | 100% | + +Разбиение — **случайное по изображениям** внутри каждого маршрута (не по маршрутам!): + +| Маршрут | Всего | Train | Test | Train% | +|---------|-------|-------|------|--------| +| 01 | 817 | 620 | 197 | 75.9% | +| 02 | 1 071 | 829 | 242 | 77.4% | +| 03 | 768 | 566 | 202 | 73.7% | +| 04 | 738 | 543 | 195 | 73.6% | +| 05 | 473 | 345 | 128 | 72.9% | +| 06 | 344 | 261 | 83 | 75.9% | +| 07 | 30 | 20 | 10 | 66.7% | +| 08 | 1 033 | 796 | 237 | 77.1% | +| 09 | 766 | 551 | 215 | 71.9% | +| 10 | 144 | 99 | 45 | 68.8% | +| 11 | 590 | 450 | 140 | 76.3% | + +--- + +## 3. МАРШРУТЫ (FLIGHTS) + +### 3.1. Детализация по маршрутам + +| Маршрут | Регион | Тип БПЛА | Drone (px) | Высота (м) | Heading (Phi) | Спутник (px) | Дата drone | Дата sat | +|---------|--------|----------|------------|------------|---------------|-------------|------------|----------| +| 01 | Changjiang-20 | multi-rotor | 3976x2652 | ~405 | 165° | 9774x26762 | 2018-09 | 2023-11 | +| 02 | Changjiang-23 | multi-rotor | 3976x2652 | ~405 | 5° | 11482x34291 | 2018-09 | 2022-09 | +| 03 | Taizhou-1 | multi-rotor | 3976x2652 | ~466 | -40° | 35092x24308 | 2018-10 | 2021-04 | +| 04 | Taizhou-6 | multi-rotor | 3976x2652 | ~542 | 170° | 18093x38408 | 2018-10 | 2023-03 | +| 05 | Yunnan | fixed-wing | 3000x2000 | ~2313 | 100° | 9394x6144 | 2016-06 | 2022-03 | +| 06 | Zhuxi | multi-rotor | 3976x2652 | ~840 | — | 8082x9780 | — | — | +| 07 | Donghuayuan | fixed-wing | 3000x2000 | ~688 | -1.5° | 3000x170 | 2018-07 | 2023-06 | +| 08 | Huzhou-3 | multi-rotor | 3976x2652 | ~551 | 100° | 43421x16294 | 2019-06 | 2023-07 | +| 09 | Huzhou-6 | multi-rotor | 3976x2652 | ~546 | -50° | 44800x33280* | 2019-06 | 2024-01 | +| 10 | Huailai | fixed-wing | 3000x2000 | ~772 | 170° | 6593x5077 | 2018-09 | 2023-06 | +| 11 | Shandan | multi-rotor | 3976x2652 | ~2572 | 90° | 29592x16582 | 2023-10 | 2021-03 | + +\* Маршрут 09: спутник разбит на 4 тайла (satellite09_01-01.tif, 01-02, 02-01, 02-02). Суммарный размер: 44800x33280 px. + +### 3.2. Географический охват + +| Регион | Маршруты | Провинция | Ландшафт | +|--------|----------|-----------|----------| +| Changjiang | 01, 02 | Цзянси | Города, деревни, фермы, реки (долина Янцзы) | +| Taizhou | 03, 04 | Цзянсу | Города, фермы, каналы, реки | +| Yunnan | 05 | Юньнань | Горы, леса, холмы (высокогорье) | +| Zhuxi | 06 | Хубэй | Горы, леса, река | +| Donghuayuan | 07 | Хэбэй | Равнина (очень узкий маршрут) | +| Huzhou | 08, 09 | Чжэцзян | Города, озеро Тайху, фермы | +| Huailai | 10 | Хэбэй | Равнина, холмы | +| Shandan | 11 | Ганьсу | Пустыня, степь (коридор Хэси) | + +**Координатный охват:** +- Широта: от 24.65°N (Юньнань) до 40.36°N (Хэбэй) — разброс ~15.7° +- Долгота: от 101.01°E (Ганьсу) до 120.25°E (Чжэцзян) — разброс ~19.2° +- Всё в пределах Китая, но с разнообразным ландшафтом + +### 3.3. Типы БПЛА + +| Тип | Маршруты | Разрешение | Высота полёта | Кол-во изображений | +|-----|----------|-----------|---------------|-------------------| +| Multi-rotor | 01, 02, 03, 04, 06, 08, 09, 11 | 3976x2652 | 405–2572 м | 6 127 (90.4%) | +| Fixed-wing | 05, 07, 10 | 3000x2000 | 688–2313 м | 647 (9.6%) | + +--- + +## 4. ИСТОЧНИКИ ИЗОБРАЖЕНИЙ + +### 4.1. Дроновые виды (query) + +| Параметр | Значение | +|----------|----------| +| Платформа | **Реальные БПЛА** (не синтетика!) | +| Тип съёмки | RGB, ground-down view (камера вертикально вниз) | +| Разрешение кадров | 3976x2652 (multi-rotor) / 3000x2000 (fixed-wing) | +| GSD (drone) | 0.1–0.2 м/пиксель (из README); расчётное: 15–97 см/px в зависимости от высоты | +| Высоты полёта | от 405 м до 2572 м | +| Heading angle | Phi1 (высокая уверенность), Phi2 (низкая уверенность) | +| Pose данные | Omega (pitch), Kappa (roll), Phi1/Phi2 (yaw) | +| Формат | JPEG | + +### 4.2. Спутниковые карты (gallery / DB) + +| Параметр | Значение | +|----------|----------| +| Платформа | Google Earth | +| Формат | GeoTIFF (.tif) | +| GSD (спутник) | **0.3 м/пиксель** (из статьи) | +| Размеры карт | от 3000x170 до 43421x38408 px | +| Кропы/патчи | **Отсутствуют** — авторы предоставляют только целые карты | + +### 4.3. Временной разрыв drone/satellite + +| Маршрут | Drone | Satellite | Разрыв | +|---------|-------|-----------|--------| +| 01 | 2018-09 | 2023-11 | **5 лет** | +| 03 | 2018-10 | 2021-04 | 2.5 года | +| 08 | 2019-06 | 2023-07 | 4 года | +| 11 | 2023-10 | 2021-03 | **-2.5 года** (спутник старше!) | + +Временной разрыв создаёт дополнительную сложность (изменения застройки, сезонные различия). + +--- + +## 5. ПАРАМЕТРЫ СЪЁМКИ ДРОНОВ + +### 5.1. Высоты полёта + +| Диапазон высот | Маршруты | Тип | Кол-во изображений | +|----------------|----------|-----|-------------------| +| 400–410 м | 01, 02 | multi-rotor | 1 888 | +| 460–550 м | 03, 04, 08, 09 | multi-rotor | 3 305 | +| 688–840 м | 06, 07, 10 | mixed | 518 | +| 2300–2575 м | 05, 11 | mixed | 1 063 | + +### 5.2. Наземное покрытие drone-кадра (оценка, FOV ~84°) + +| Высота | Footprint (multi-rotor 3976x2652) | Footprint (fixed-wing 3000x2000) | +|--------|-----------------------------------|----------------------------------| +| ~405 м | ~608 x 405 м | — | +| ~466 м | ~699 x 466 м | — | +| ~551 м | ~826 x 551 м | — | +| ~688 м | — | ~1031 x 688 м | +| ~840 м | ~1260 x 840 м | — | +| ~2313 м | — | ~3469 x 2313 м | +| ~2572 м | ~3858 x 2572 м | — | + +### 5.3. Расстояние между кадрами + +| Маршрут | Avg spacing | Высота | Overlap (оценка) | +|---------|-------------|--------|------------------| +| 01 | 80.7 м | 405 м | ~87% (по ширине footprint) | +| 05 | 62.7 м | 2313 м | ~98% | +| 07 | 23.6 м | 688 м | ~98% | +| 08 | 99.7 м | 551 м | ~88% | +| 11 | 142.9 м | 2572 м | ~96% | + +Высокий overlap (87–98%) типичен для аэрофотосъёмки. + +--- + +## 6. ПРОСТРАНСТВЕННАЯ НАРЕЗКА СПУТНИКОВЫХ ПАТЧЕЙ + +### 6.1. Текущее состояние + +**В оригинальном датасете кропы/патчи НЕ предоставлены.** Авторы дают только целые спутниковые карты. Задача определена как "найти координаты на большой карте", а не как retrieval по патчам. + +### 6.2. Планируемая нарезка (по аналогии с UAV-GeoLoc) + +Для совместимости с pipeline на основе UAV-GeoLoc, планируется нарезка спутниковых карт на патчи. + +**Параметры нарезки:** + +| Параметр | Значение | Обоснование | +|----------|----------|-------------| +| Размер кропа | **512x512 px** | ~154x154 м на земле (при GSD 0.3 м/px) | +| Stride | **256 px** | 50% overlap (как в UAV-GeoLoc) | +| Именование | `crop_X_Y.png` | X по ширине (col), Y по высоте (row) | +| Позиция в карте | `sat[Y*stride : Y*stride+crop, X*stride : X*stride+crop]` | — | +| Финальный размер (модель) | **256x256 px** | Resize для входа в модель | + +### 6.3. Ожидаемое количество кропов + +| Маршрут | Размер карты | Кропов (cols x rows) | Итого | Примечание | +|---------|-------------|---------------------|-------|------------| +| 01 | 9774x26762 | 37 x 103 | 3 811 | OK | +| 02 | 11482x34291 | 43 x 132 | 5 676 | OK | +| 03 | 35092x24308 | 136 x 93 | 12 648 | OK | +| 04 | 18093x38408 | 69 x 149 | 10 281 | OK | +| 05 | 9394x6144 | 35 x 23 | 805 | OK | +| 06 | 8082x9780 | 30 x 37 | 1 110 | OK | +| 07 | 3000x170 | — | — | **Исключён** (высота 170 px < 512) | +| 08 | 43421x16294 | 168 x 62 | 10 416 | OK | +| 09 | 44800x33280 | 174 x 129 | 22 446 | Нужна сшивка 4 тайлов | +| 10 | 6593x5077 | 24 x 18 | 432 | OK | +| 11 | 29592x16582 | 114 x 63 | 7 182 | OK | +| **Итого** | | | **~74 807** | Без маршрута 07 | + +### 6.4. Наземное покрытие одного кропа + +При GSD спутника = 0.3 м/px: +- **Кроп 512x512 px** покрывает **~154 x 154 м** на земле +- Это вписывается в footprint drone-кадра на любой высоте (405+ м) +- Stride 256 px = 76.8 м на земле + +### 6.5. Проблемные маршруты + +| Маршрут | Проблема | Решение | +|---------|----------|---------| +| 07 | Спутник 3000x170 px — слишком узкий | Исключить или кропы 170x170 | +| 09 | Спутник разбит на 4 тайла | Сшить в одну карту перед нарезкой | + +--- + +## 7. АННОТАЦИИ И МЕТАДАННЫЕ + +### 7.1. Файлы аннотаций + +| Файл | Содержание | Формат | +|------|-----------|--------| +| `XX.csv` | GPS + pose каждого drone-кадра | CSV: num, filename, date, lat, lon, height, Omega, Kappa, Phi1, Phi2 | +| `satellite_ coordinates_range.csv` | GPS-bbox каждой спутниковой карты | CSV: mapname, LT_lat, LT_lon, RB_lat, RB_lon, region | +| `visloc_train.csv` | Train split | TSV: filename, height, Omega, Kappa, Phi1, Phi2 | +| `visloc_test.csv` | Test split | TSV: filename, height, Omega, Kappa, Phi1, Phi2 | + +### 7.2. Типы аннотаций + +| Тип аннотации | Наличие | Комментарий | +|---------------|---------|-------------| +| GPS-координаты (drone) | **Да** | Из бортового GNSS, точность ~1-3 м | +| GPS-bbox (satellite) | **Да** | Углы спутниковой карты | +| Высота полёта | **Да** | В метрах | +| Heading angle (yaw) | **Да** | Phi1 (надёжный) и Phi2 (менее надёжный) | +| Pitch / Roll | **Да** | Omega (pitch), Kappa (roll) | +| Дата съёмки | **Да** | Для drone-кадров | +| Positive/Semi-positive pairs | **Нет** | Отсутствуют — нужно генерировать | +| Кропы спутника (DB) | **Нет** | Отсутствуют — нужно нарезать | +| Depth maps | Нет | — | +| Segmentation masks | Нет | — | +| Bounding boxes | Нет | — | +| Семантические метки | Нет | Только implicit через регион | + +--- + +## 8. СРАВНЕНИЕ С UAV-GeoLoc + +| Параметр | UAV-VisLoc | UAV-GeoLoc | +|----------|-----------|-----------| +| Drone-изображения | **Реальные** БПЛА | **Синтетика** (Google Earth Studio 3D) | +| Кол-во drone | 6 774 | 652 744 | +| Кол-во спутниковых кропов | 0 (нужно генерировать) | 274 683 | +| Размер drone | 3976x2652 / 3000x2000 | 512x512 | +| GSD drone | 0.1–0.97 м/px (зависит от высоты) | ~0.5 м/px (синтетика) | +| GSD satellite | 0.3 м/px (Google Earth) | Варьируется | +| Высоты полёта | 405–2572 м (реальные) | 100, 125, 150 м (синтетика) | +| Heading angles | Произвольные (реальный полёт) | Дискретные: 8 x 45° | +| Регионы | 7 провинций Китая | 11 стран, 6 континентов | +| Сцен | 11 маршрутов | 372 сцены | +| Train/test split | По изображениям (~75/25) | По сценам (140/40/20) | +| Positive pairs | Нет (нужно по GPS) | Есть (positive.json) | +| Semi-positive pairs | Нет | Есть (semi_positive.json) | +| Temporal gap | 2–5 лет | Нет (одновременно) | +| Лицензия | Не указана | CC BY-NC 4.0 | + +### Ключевые отличия: +1. **Реальные vs синтетические** — UAV-VisLoc содержит реальные фотографии, что даёт реалистичные артефакты (освещение, шум, blur), но меньше контроля +2. **Масштаб** — UAV-GeoLoc на 2 порядка больше по количеству изображений +3. **Пары не предоставлены** — для UAV-VisLoc нужно самостоятельно сопоставить drone GPS с координатами кропов +4. **Вариативность высот** — гораздо шире (400–2600 м vs 100–150 м) +5. **Temporal gap** — drone и satellite сняты в разные годы, что усложняет matching + +--- + +## 9. ПЛАН ГЕНЕРАЦИИ КРОПОВ И ПАР + +### 9.1. Pipeline + +``` +satellite.tif + satellite_coordinates_range.csv + │ + ▼ + [1] Нарезка кропов (512x512, stride 256) + │ + ▼ + [2] Вычисление GPS центра каждого кропа + │ (из bbox карты + позиция кропа в grid) + ▼ + [3] Для каждого drone-кадра: + │ — Найти кроп с минимальным GPS-расстоянием → positive + │ — Найти кропы в радиусе R → semi-positives + ▼ + [4] Генерация positive.json, semi_positive.json, db_postion.txt + │ + ▼ + [5] Генерация Index файлов (train_query.txt, train_db.txt, ...) + │ + ▼ + [6] Resize drone → 256x256, кропов → 256x256 +``` + +### 9.2. Matching drone → crop + +GPS центр кропа `crop_X_Y.png` вычисляется как: +``` +crop_center_lon = LT_lon + (X * stride + crop_size/2) * (RB_lon - LT_lon) / sat_width +crop_center_lat = LT_lat + (Y * stride + crop_size/2) * (RB_lat - LT_lat) / sat_height +``` + +Positive match: кроп с минимальным евклидовым расстоянием до GPS drone-кадра. +Semi-positive: все кропы в радиусе 1 stride (256 px = ~77 м) от positive. + +--- + +## 10. ПОЛНАЯ СТРУКТУРА ДАННЫХ + +``` +UAV_VisLoc_dataset/ # ~16.4 GB +├── satellite_ coordinates_range.csv # GPS-bbox всех 11 карт +├── visloc_train.csv # 5080 train drone images (TSV) +├── visloc_test.csv # 1694 test drone images (TSV) +├── README_dataset.txt # Описание датасета +│ +├── 01/ # Changjiang-20, multi-rotor, 817 imgs +│ ├── drone/ +│ │ ├── 01_0001.JPG # 3976x2652 +│ │ ├── 01_0002.JPG +│ │ └── ... +│ ├── satellite01.tif # 9774x26762 +│ └── 01.csv # num,filename,date,lat,lon,height,... +│ +├── 02/ ... 06/ # Аналогичная структура +│ +├── 07/ # Donghuayuan, fixed-wing, 30 imgs +│ ├── drone/ +│ │ └── 07_XXXX.JPG # 3000x2000 +│ ├── satellite07.tif # 3000x170 (!) +│ └── 07.csv +│ +├── 08/ # Huzhou-3, multi-rotor, 1033 imgs +│ +├── 09/ # Huzhou-6, multi-rotor, 766 imgs +│ ├── drone/ +│ ├── satellite09_01-01.tif # 25600x25600 ─┐ +│ ├── satellite09_01-02.tif # 19200x25600 │ Суммарно: +│ ├── satellite09_02-01.tif # 25600x7680 │ 44800x33280 +│ ├── satellite09_02-02.tif # 19200x7680 ─┘ +│ └── 09.csv +│ +├── 10/ # Huailai, fixed-wing, 144 imgs +│ +└── 11/ # Shandan, multi-rotor, 590 imgs + ├── drone/ + │ └── 11_XXXX.JPG # 3976x2652 + ├── satellite11.tif # 29592x16582 + └── 11.csv +``` diff --git a/conf/balanced.gin b/conf/balanced.gin index 81a489c..029007b 100644 --- a/conf/balanced.gin +++ b/conf/balanced.gin @@ -1,5 +1,5 @@ -# 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 +# Balanced: GatedFusion with text captions enabled. +# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs satellite import src.datasets.visloc_with_captions import src.losses.multi_infonce @@ -12,43 +12,37 @@ 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.baseline_mode = False +DualEncoderCaptionTest.init_gate = 0.7 DualEncoderCaptionTest.device = "cuda" -ProjectionHead.in_dim = 512 -ProjectionHead.out_dim = 512 -ProjectionHead.use_mlp = False +# ---- Fusion ---- +GatedFusion.init_gate = 0.7 +GatedFusion.baseline_mode = 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 +InfoNCELoss.temperature_init = 0.1 +InfoNCELoss.temperature_final = 0.01 +InfoNCELoss.label_smoothing = 0.1 +InfoNCELoss.weight_q2g = 0.6 +InfoNCELoss.weight_g2q = 0.4 # ---- Dataset ---- -VisLocCaptionDataset.caption_strategy = "hybrid" -VisLocCaptionDataset.drop_caption_prob = 0.0 -VisLocCaptionDataset.seed = 42 +GeoLocCaptionDataset.drop_caption_prob = 0.0 +GeoLocCaptionDataset.seed = 42 # ---- Training ---- -TrainConfig.train_manifest = "data/visloc_train.json" -TrainConfig.val_manifest = "data/visloc_val.json" -TrainConfig.image_root = "data/visloc/images" +TrainConfig.train_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/train_query.txt" +TrainConfig.val_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/val_query.txt" +TrainConfig.data_root = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc" TrainConfig.output_dir = "out/caption_test/balanced" -TrainConfig.epochs = 30 +TrainConfig.epochs = 10 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.eval_every = 2 TrainConfig.seed = 42 TrainConfig.device = "cuda" diff --git a/conf/baseline_no_text.gin b/conf/baseline_no_text.gin index a2fddc3..2dd966a 100644 --- a/conf/baseline_no_text.gin +++ b/conf/baseline_no_text.gin @@ -1,11 +1,10 @@ -# Baseline: image-image only, no captions. Reference R@1 for delta computation. -# L = 1.0 * L_img_img +# Baseline: no text fusion (gate forced to 1.0). +# query = drone_only -> InfoNCE vs satellite +# Reference R@1 for delta computation. 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 +DualEncoderCaptionTest.baseline_mode = True +GatedFusion.baseline_mode = True TrainConfig.output_dir = "out/caption_test/baseline_no_text" diff --git a/conf/text_heavy.gin b/conf/text_heavy.gin index 1a69cbd..c6e5daf 100644 --- a/conf/text_heavy.gin +++ b/conf/text_heavy.gin @@ -1,11 +1,9 @@ -# 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 +# Text-heavy: gate initialized low (more text weight). +# query = sigma(0.3) * drone + 0.7 * text 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 +DualEncoderCaptionTest.init_gate = 0.3 +GatedFusion.init_gate = 0.3 TrainConfig.output_dir = "out/caption_test/text_heavy" diff --git a/scripts/compare_runs.py b/scripts/compare_runs.py index 233098e..74f7875 100644 --- a/scripts/compare_runs.py +++ b/scripts/compare_runs.py @@ -1,9 +1,6 @@ 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. +"""Compare baseline (no text) vs caption-fused runs. Compute Delta R@1 report. Usage: python -m scripts.compare_runs \ @@ -18,10 +15,8 @@ from pathlib import Path _DIRECTIONS = ( - "drone_to_sat", - "sat_to_drone", - "text_to_sat", - "text_to_drone", + "query_to_gallery", + "gallery_to_query", ) _KS = (1, 5, 10) @@ -33,73 +28,70 @@ def _load_metrics(report_path: Path) -> dict[str, float]: 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})") + delta = f_ - b if (b == b and f_ == f_) else float("nan") + cells.append(f"{b:.4f} -> {f_:.4f} (D {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%)" + return "PASS -- captions informative (D R@1 >= +3%)" if delta >= 0.01: - return "⚠️ MARGINAL — consider VLM refinement (+1% ≤ Δ < +3%)" + return "MARGINAL -- consider VLM refinement (+1% <= D < +3%)" if delta >= 0: - return "❌ WEAK — captions add little signal (< +1%)" - return "❌❌ HARMFUL — captions confuse model (Δ < 0)" + return "WEAK -- captions add little signal (< +1%)" + return "HARMFUL -- captions confuse model (D < 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 = "query_to_gallery" 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(f"## Primary metric: D R@1 (query->gallery) = {primary_delta:+.4f}") lines.append("") lines.append(f"**Verdict:** {verdict}") lines.append("") - # Full table. - lines.append("## All directions × K") + # Gate value. + gate = full.get("gate", None) + if gate is not None: + lines.append(f"**Fusion gate:** {gate:.4f} (1.0 = text ignored, 0.0 = image ignored)") + lines.append("") + + lines.append("## All directions x K") lines.append("") - header = "| Direction | R@1 base → full | R@5 base → full | R@10 base → full |" + 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(_format_row(direction, baseline, full)) 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("- D R@1 >= +3% -> captions pass, proceed to production") + lines.append("- +1% <= D R@1 < +3% -> add VLM refinement, re-run") + lines.append("- D R@1 < +1% -> redesign caption pipeline") + lines.append("- D R@1 < 0 -> critical bug, investigate") lines.append("") return "\n".join(lines) def main() -> None: - parser = argparse.ArgumentParser( - description="Compare baseline vs full-caption runs." - ) + parser = argparse.ArgumentParser(description="Compare baseline vs 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) @@ -114,7 +106,6 @@ def main() -> None: 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, diff --git a/src/datasets/visloc_with_captions.py b/src/datasets/visloc_with_captions.py index 9fff8dd..6b4b0a3 100644 --- a/src/datasets/visloc_with_captions.py +++ b/src/datasets/visloc_with_captions.py @@ -1,26 +1,20 @@ from __future__ import annotations -"""UAV-VisLoc dataset loader augmented with generated captions. +"""UAV-GeoLoc dataset loader with template captions for CVGL caption test. -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] - }, - ... - ] +Reads UAV-GeoLoc Index format (train_query.txt / train_db.txt) and generates +template captions from path metadata (terrain type, scene, altitude, heading). -Captions are produced offline by scripts/generate_captions.py using one of -three strategies: template, VLM, or hybrid (see АНАЛИЗ_caption_quality_test). +train_query.txt format: + Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg 0 .../DB/img/crop_X_Y.png ... + +Template caption (P3-style fingerprint for cross-view matching): + "Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano. + Plan-view features: dark lava flows, crater edges, sparse vegetation patches." """ -import json import random +import re from pathlib import Path from typing import Any, Callable @@ -29,130 +23,191 @@ import torch from PIL import Image from torch.utils.data import Dataset +# Compass lookup: heading_deg -> direction name. +_COMPASS = ["north", "northeast", "east", "southeast", + "south", "southwest", "west", "northwest"] + +# Simple terrain-to-features mapping for template captions. +_TERRAIN_FEATURES: dict[str, str] = { + "Volcano": "lava flows, crater edges, volcanic rock", + "Mountain": "ridgelines, steep slopes, rocky terrain", + "Hill": "rolling hills, gentle slopes, scattered trees", + "Desert": "sand dunes, arid ground, sparse scrub", + "Plain": "flat open fields, agricultural plots, dirt roads", + "Plateau": "flat elevated terrain, cliffs, mesa edges", + "Basin": "lowland depression, dry lake bed, sediment patterns", + "Delta": "river channels, sediment fans, wetland patches", + "Gorge": "deep canyon, exposed rock layers, narrow valley", + "Island": "shoreline, coastal vegetation, water boundary", + "Wetland": "marshes, water channels, aquatic vegetation", + "Glacier": "ice formations, crevasses, glacial moraine", + "Forest": "dense canopy, tree shadows, clearings", + "Farm": "crop fields, irrigation lines, farm buildings", + "Prairie": "grassland, open meadow, fence lines", + "Finca": "rural estate, orchards, scattered structures", + "Calcification": "mineral deposits, white crusts, barren soil", + "StoneForest": "karst pillars, eroded limestone, sparse vegetation", + "Oasis": "palm clusters, water pool, surrounding desert", + "Flowers": "colorful ground cover, floral patterns, open field", + "Terrace": "terraced hillside, stepped fields, retaining walls", + "Snow": "snow cover, tracks, exposed rock patches", + "Pasture": "grazing land, fenced paddocks, grass", + "Danxia": "red sandstone, layered cliffs, erosion patterns", + "Hylare": "sparse woodland, rocky outcrops", + "Karst": "sinkholes, limestone towers, caves", + "Fall": "autumn foliage, colored canopy, leaf litter", +} + +# Country/city features. +_COUNTRY_FEATURES: str = "buildings, roads, urban blocks, rooftops, intersections" + + +def _parse_query_path(query_path: str) -> dict[str, str]: + """Extract metadata from UAV-GeoLoc query path. + + Example: Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg + Returns: {category, terrain_type, scene, height_m, heading_deg, heading_dir} + """ + parts = query_path.split("/") + + category = parts[0] if parts else "Unknown" + + if category == "Terrain": + terrain_type = parts[1] if len(parts) > 1 else "Unknown" + scene = parts[2] if len(parts) > 2 else "Unknown" + elif category == "Country": + terrain_type = "Urban" + scene = "/".join(parts[1:3]) if len(parts) > 2 else "Unknown" + else: + terrain_type = "Unknown" + scene = parts[1] if len(parts) > 1 else "Unknown" + + # Parse height and rotation from trajectory folder name. + height_m = "100" + heading_deg = "0" + for part in parts: + m = re.match(r"height(\d+)_rot(\d+)", part) + if m: + height_m = m.group(1) + heading_deg = m.group(2) + break + + heading_idx = round(int(heading_deg) / 45) % 8 + heading_dir = _COMPASS[heading_idx] + + return { + "category": category, + "terrain_type": terrain_type, + "scene": scene, + "height_m": height_m, + "heading_deg": heading_deg, + "heading_dir": heading_dir, + } + + +def _make_template_caption(meta: dict[str, str]) -> str: + """Generate a template caption from parsed metadata.""" + terrain = meta["terrain_type"] + features = _TERRAIN_FEATURES.get(terrain, _COUNTRY_FEATURES) + + return ( + f"Aerial view at {meta['height_m']}m facing {meta['heading_dir']} " + f"over {terrain.lower()} terrain near {meta['scene']}. " + f"Plan-view features: {features}." + ) + @gin.configurable -class VisLocCaptionDataset(Dataset): - """UAV-VisLoc pairs with generated captions. +class GeoLocCaptionDataset(Dataset): + """UAV-GeoLoc pairs with template captions. + + Reads train_query.txt, randomly samples one positive crop per query. 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. + query_file: Path to train_query.txt (or test_query.txt). + data_root: Root directory of UAV-GeoLoc dataset. + image_transform: Callable applied to PIL images. + drop_caption_prob: Probability of dropping caption (for ablation). + seed: Random seed. """ def __init__( self, - manifest_path: str, - image_root: str, + query_file: str, + data_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.data_root = Path(data_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.entries: list[dict[str, Any]] = [] + self._load_query_file(Path(query_file)) - self._validate_entries() + def _load_query_file(self, query_file: Path) -> None: + """Parse train_query.txt into list of entries.""" + with open(query_file) as f: + for line in f: + line = line.strip() + if not line: + continue + parts = line.split() + query_path = parts[0] + # parts[1] is label (always 0), parts[2:] are positive crop paths. + positive_crops = parts[2:] + if not positive_crops: + continue - 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} + meta = _parse_query_path(query_path) + caption = _make_template_caption(meta) - 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}" + self.entries.append({ + "query_path": query_path, + "positive_crops": positive_crops, + "caption": caption, + "meta": meta, + }) def _load_image(self, relative_path: str) -> torch.Tensor: - """Load image and apply preprocessing.""" - path = self.image_root / relative_path + path = self.data_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"]) + drone_img = self._load_image(entry["query_path"]) - caption_drone = self._maybe_drop(entry[self._caption_key("drone")]) - caption_sat = self._maybe_drop(entry[self._caption_key("sat")]) + # Randomly sample one positive crop. + crop_path = self._rng.choice(entry["positive_crops"]) + sat_img = self._load_image(crop_path) + + caption = entry["caption"] + if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob: + caption = "" 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}"), + "caption_drone": caption, + "pair_id": entry["query_path"], } 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. - """ + """Collate into batched dict. Captions stay as string 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/evaluate.py b/src/eval/evaluate.py index 68936c8..0324128 100644 --- a/src/eval/evaluate.py +++ b/src/eval/evaluate.py @@ -1,9 +1,9 @@ from __future__ import annotations -"""Evaluation utilities for caption quality test. +"""Evaluation 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. +Recall@K for query(drone+text) -> gallery(satellite). +delta_r_at_1 compares caption-aware vs baseline runs. """ import json @@ -23,19 +23,10 @@ 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]. - """ + """Recall@K assuming positives on the diagonal.""" 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] @@ -49,51 +40,29 @@ 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]. - """ + """Encode all samples into query and gallery embeddings.""" model.eval() - all_drone: list[torch.Tensor] = [] - all_sat: list[torch.Tensor] = [] - all_cap_drone: list[torch.Tensor] = [] - all_cap_sat: list[torch.Tensor] = [] + all_query: list[torch.Tensor] = [] + all_gallery: 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 + caption_drone = batch["caption_drone"] embeddings = model( drone_img=drone_img, sat_img=sat_img, - caption_drone=captions_drone, - caption_sat=captions_sat, + caption_drone=caption_drone, ) - 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()) + all_query.append(embeddings["query"].cpu()) + all_gallery.append(embeddings["gallery"].cpu()) - out = { - "drone": torch.cat(all_drone, dim=0), - "sat": torch.cat(all_sat, dim=0), + return { + "query": torch.cat(all_query, dim=0), + "gallery": torch.cat(all_gallery, 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( @@ -101,51 +70,25 @@ def evaluate_retrieval( 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. + """Compute R@K for query->gallery and gallery->query. Returns: - Flat dict with keys like 'r@1_drone_to_sat', 'r@5_text_to_sat', etc. + Flat dict: r@1_query_to_gallery, r@5_query_to_gallery, etc. """ - feats = _encode_dataset( - model=model, - loader=loader, - device=device, - include_captions=include_captions, - ) + feats = _encode_dataset(model=model, loader=loader, device=device) metrics: dict[str, float] = {} - sim_d2s = feats["drone"] @ feats["sat"].t() - sim_s2d = sim_d2s.t() + sim_q2g = feats["query"] @ feats["gallery"].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 + for k, val in _recall_at_k(sim_q2g, k_values).items(): + metrics[f"r@{k}_query_to_gallery"] = val + for k, val in _recall_at_k(sim_q2g.t(), k_values).items(): + metrics[f"r@{k}_gallery_to_query"] = 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 + # Gate value for diagnostics. + metrics["gate"] = model.fusion.gate_value return metrics @@ -153,64 +96,36 @@ def evaluate_retrieval( 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)}" - ) + """R@1 gain from adding captions: full - baseline.""" + key = "r@1_query_to_gallery" return full_metrics[key] - baseline_metrics[key] @gin.configurable def run_evaluation_from_checkpoint( checkpoint_path: str, - test_manifest: str, - image_root: str, + test_query_file: str, + data_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. - """ + """Standalone evaluation from checkpoint.""" from src.datasets.visloc_with_captions import ( - VisLocCaptionDataset, + GeoLocCaptionDataset, collate_caption_batch, ) model = DualEncoderCaptionTest().to(device) - ckpt = torch.load(checkpoint_path, map_location=device) + ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) model.load_state_dict(ckpt["model_state"]) model.eval() - test_ds = VisLocCaptionDataset( - manifest_path=test_manifest, - image_root=image_root, + test_ds = GeoLocCaptionDataset( + query_file=test_query_file, + data_root=data_root, image_transform=model.preprocess, ) test_loader = DataLoader( @@ -222,15 +137,11 @@ def run_evaluation_from_checkpoint( pin_memory=True, ) - metrics = evaluate_retrieval( - model=model, - loader=test_loader, - device=device, - ) + metrics = evaluate_retrieval(model=model, loader=test_loader, device=device) report = { "checkpoint": checkpoint_path, - "test_manifest": test_manifest, + "test_query_file": test_query_file, "metrics": metrics, } out = Path(output_path) diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py index 33ee38e..2e2fd81 100644 --- a/src/losses/multi_infonce.py +++ b/src/losses/multi_infonce.py @@ -1,14 +1,10 @@ from __future__ import annotations -"""Multi-term InfoNCE loss for caption quality validation. +"""InfoNCE loss for cross-view geo-localization with optional text fusion. -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). +Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite). +Asymmetric weighting: query->gallery weighted higher (real use-case direction). +Cosine temperature schedule for sharper distribution over training. """ import math @@ -27,26 +23,12 @@ def _symmetric_info_nce( 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. - """ + """Weighted symmetric InfoNCE. Positives on the diagonal.""" 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 @@ -56,85 +38,23 @@ def cosine_temperature( 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. - """ + """Cosine-decay schedule for InfoNCE temperature.""" 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. +class InfoNCELoss(nn.Module): + """Symmetric InfoNCE with cosine temperature schedule. 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. + temperature_init: Temperature at epoch 0. + temperature_final: Temperature after cosine decay. + label_smoothing: Cross-entropy label smoothing. + weight_q2g: Weight for query->gallery direction. + weight_g2q: Weight for gallery->query direction. """ def __init__( @@ -142,27 +62,15 @@ class MultiTermInfoNCE(nn.Module): 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, + weight_q2g: float = 0.6, + weight_g2q: float = 0.4, ) -> 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 + self.weight_q2g = weight_q2g + self.weight_g2q = weight_g2q def forward( self, @@ -170,17 +78,16 @@ class MultiTermInfoNCE(nn.Module): epoch: int, total_epochs: int, ) -> dict[str, torch.Tensor]: - """Compute multi-term loss. + """Compute InfoNCE loss. Args: - embeddings: Dict with keys 'drone', 'sat', and optionally - 'cap_drone', 'cap_sat'. Each [B, D] L2-normalized. + embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized, + plus 'gate' (float) from fusion module. 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'. + Dict with 'total', 'temperature', 'gate'. """ tau = cosine_temperature( epoch=epoch, @@ -188,75 +95,20 @@ class MultiTermInfoNCE(nn.Module): 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, + loss = _symmetric_info_nce( + emb_a=embeddings["query"], + emb_b=embeddings["gallery"], temperature=tau, label_smoothing=self.label_smoothing, - weight_a2b=self.asym_drone_to_sat, - weight_b2a=self.asym_sat_to_drone, + weight_a2b=self.weight_q2g, + weight_b2a=self.weight_g2q, ) - 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 - ) + gate = embeddings.get("gate", 1.0) 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), + "total": loss, + "temperature": torch.tensor(tau, device=loss.device), + "gate": torch.tensor(gate, device=loss.device), } diff --git a/src/models/dual_encoder.py b/src/models/dual_encoder.py index 76dad95..e4e0b6b 100644 --- a/src/models/dual_encoder.py +++ b/src/models/dual_encoder.py @@ -1,10 +1,16 @@ from __future__ import annotations -"""Dual encoder for caption quality test on UAV-VisLoc. +"""Dual encoder for caption quality test on cross-view geo-localization. GeoRSCLIP ViT-B/32 backbone (image + text towers, shared 512-dim space). -Image encoder is frozen, text encoder has partial unfreeze (last block + projection). -Separate trainable projection heads for drone/sat/text branches. +Image encoder frozen. Text encoder with partial unfreeze. + +Architecture: + Query branch: GeoRSCLIP_img(drone) + GeoRSCLIP_text(caption) -> GatedFusion -> proj -> query_emb + Gallery branch: GeoRSCLIP_img(sat) -> proj -> gallery_emb + Loss: InfoNCE(query_emb, gallery_emb) + +Baseline mode: fusion gate forced to 1.0 (text ignored). """ from typing import Literal @@ -16,16 +22,8 @@ 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). - """ + """MLP projection head with L2 normalization.""" def __init__( self, @@ -46,33 +44,56 @@ class ProjectionHead(nn.Module): self.proj = nn.Linear(in_dim, out_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: - """Project features and L2-normalize. + return F.normalize(self.proj(x), dim=-1) - 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 GatedFusion(nn.Module): + """Learnable gated fusion of image and text embeddings. + + q = sigma(alpha) * img + (1 - sigma(alpha)) * text + + alpha is a single learnable scalar, initialized so that gate ~ init_gate. + When baseline_mode=True, gate is clamped to 1.0 (text contribution = 0). + """ + + def __init__(self, init_gate: float = 0.7, baseline_mode: bool = False) -> None: + super().__init__() + # alpha is in logit space: sigmoid(alpha) = init_gate + init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) + self.alpha = nn.Parameter(init_alpha) + self.baseline_mode = baseline_mode + + def forward( + self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + if text_feat is None or self.baseline_mode: + return img_feat + gate = torch.sigmoid(self.alpha) + return gate * img_feat + (1.0 - gate) * text_feat + + @property + def gate_value(self) -> float: + """Current gate value (image weight). 1.0 = text ignored.""" + if self.baseline_mode: + return 1.0 + return torch.sigmoid(self.alpha).item() @gin.configurable class DualEncoderCaptionTest(nn.Module): - """GeoRSCLIP dual encoder 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. + """GeoRSCLIP dual encoder with gated text fusion on query branch. 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. + variant: open_clip model variant name. + pretrained_path: Path to GeoRSCLIP checkpoint. + unfreeze_mode: Text encoder unfreeze strategy. + embed_dim: Output retrieval embedding dimension. + use_mlp_heads: Use 2-layer MLP projection heads. + baseline_mode: If True, fusion gate = 1.0 (no text). + init_gate: Initial gate value (image weight). device: torch device. """ @@ -80,19 +101,19 @@ class DualEncoderCaptionTest(nn.Module): self, variant: str = "ViT-B-32", pretrained_path: str = "RS5M_ViT-B-32.pt", - unfreeze_mode: Literal["none", "projection", "last_block", "full"] = "last_block", + unfreeze_mode: Literal["none", "projection", "last_block"] = "last_block", embed_dim: int = 512, use_mlp_heads: bool = False, - shared_image_head: bool = True, + baseline_mode: bool = False, + init_gate: float = 0.7, device: str = "cuda", ) -> None: super().__init__() - self.variant = variant self.embed_dim = embed_dim - self.shared_image_head = shared_image_head self.device = device + self.baseline_mode = baseline_mode - # Load open_clip model (GeoRSCLIP compatible with open_clip API). + # Load GeoRSCLIP via open_clip. self.model, _, self.preprocess = open_clip.create_model_and_transforms( model_name=variant, pretrained=pretrained_path, @@ -100,45 +121,28 @@ class DualEncoderCaptionTest(nn.Module): ) self.tokenizer = open_clip.get_tokenizer(variant) - # Native GeoRSCLIP embedding dim (for ViT-B/32 = 512). - self._native_dim = self._infer_native_dim() + native_dim = self._infer_native_dim() - # Freeze everything by default. + # Freeze everything. for p in self.model.parameters(): p.requires_grad = False - # Apply unfreeze strategy. - self._apply_unfreeze(unfreeze_mode) + # Selectively unfreeze text encoder (only if not baseline). + if not baseline_mode: + 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, + # Gated fusion on query branch. + self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode) + + # Projection heads. + self.proj_query = ProjectionHead( + in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads, + ) + self.proj_gallery = ProjectionHead( + in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads, ) - 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]) @@ -146,17 +150,11 @@ class DualEncoderCaptionTest(nn.Module): def _apply_unfreeze( self, - unfreeze_mode: Literal["none", "projection", "last_block", "full"], + unfreeze_mode: Literal["none", "projection", "last_block"], ) -> 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. + # Unfreeze text_projection. if hasattr(self.model, "text_projection"): tp = self.model.text_projection if isinstance(tp, nn.Parameter): @@ -164,38 +162,17 @@ class DualEncoderCaptionTest(nn.Module): elif isinstance(tp, nn.Module): for p in tp.parameters(): p.requires_grad = True - - # Additionally unfreeze last transformer block. + # 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(): + for p in self.model.transformer.resblocks[-1].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() + def encode_text(self, texts: list[str]) -> torch.Tensor: + tokens = self.tokenizer(list(texts)).to(self.device).long() feats = self.model.encode_text(tokens) return F.normalize(feats, dim=-1) @@ -204,40 +181,37 @@ class DualEncoderCaptionTest(nn.Module): 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. + """Forward pass. 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. + drone_img: Drone images [B, 3, H, W]. + sat_img: Satellite images [B, 3, H, W]. + caption_drone: Drone captions (P3 fingerprint), one per sample. Returns: - Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat', each - containing [B, embed_dim] L2-normalized embeddings. - Keys for missing captions are absent. + Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim], + and 'gate' (scalar) for logging. """ - out: dict[str, torch.Tensor] = {} - - drone_feat = self.encode_image(drone_img) + # Gallery branch: satellite only. sat_feat = self.encode_image(sat_img) + gallery = self.proj_gallery(sat_feat) - 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) + # Query branch: drone + optional text fusion. + drone_feat = self.encode_image(drone_img) - 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)) + text_feat = None + if caption_drone is not None and not self.baseline_mode: + text_feat = self.encode_text(caption_drone) - return out + fused = self.fusion(drone_feat, text_feat) + query = self.proj_query(fused) + + return { + "query": query, + "gallery": gallery, + "gate": self.fusion.gate_value, + } def trainable_parameters(self) -> list[nn.Parameter]: - """Return list of trainable parameters for optimizer construction.""" return [p for p in self.parameters() if p.requires_grad] diff --git a/src/training/train.py b/src/training/train.py index 1b475f3..723d1c4 100644 --- a/src/training/train.py +++ b/src/training/train.py @@ -1,10 +1,9 @@ from __future__ import annotations -"""Training loop for caption quality validation on UAV-VisLoc. +"""Training loop for caption quality test on cross-view geo-localization. -Uses gin-configurable DualEncoderCaptionTest + MultiTermInfoNCE. -Logs per-component losses, temperature, and lambdas each step. -Saves checkpoint + eval snapshot every epoch. +GeoRSCLIP dual encoder with GatedFusion on query branch. +Single InfoNCE loss: query(drone+text) vs gallery(satellite). """ import argparse @@ -22,11 +21,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from torch.utils.data import DataLoader from src.datasets.visloc_with_captions import ( - VisLocCaptionDataset, + GeoLocCaptionDataset, collate_caption_batch, ) from src.eval.evaluate import evaluate_retrieval -from src.losses.multi_infonce import MultiTermInfoNCE +from src.losses.multi_infonce import InfoNCELoss from src.models.dual_encoder import DualEncoderCaptionTest LOGGER = logging.getLogger("caption_test.train") @@ -34,20 +33,20 @@ LOGGER = logging.getLogger("caption_test.train") @gin.configurable class TrainConfig: - """Top-level training configuration (gin-configurable). + """Top-level training configuration. 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. + train_query_file: Path to train_query.txt. + val_query_file: Path to test_query.txt (used as val). + data_root: Root of UAV-GeoLoc dataset. + output_dir: Checkpoint and log output directory. epochs: Number of training epochs. batch_size: Mini-batch size. - num_workers: DataLoader worker count. + num_workers: DataLoader workers. 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. + grad_clip: Max gradient norm (0 disables). + use_amp: Enable fp16 mixed-precision. eval_every: Run validation every N epochs. seed: Random seed. device: torch device. @@ -55,24 +54,24 @@ class TrainConfig: def __init__( self, - train_manifest: str = "data/visloc_train.json", - val_manifest: str = "data/visloc_val.json", - image_root: str = "data/visloc/images", + train_query_file: str = "Index/train_query.txt", + val_query_file: str = "Index/test_query.txt", + data_root: str = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc", output_dir: str = "out/caption_test", - epochs: int = 30, + epochs: int = 10, 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, + eval_every: int = 2, seed: int = 42, device: str = "cuda", ) -> None: - self.train_manifest = train_manifest - self.val_manifest = val_manifest - self.image_root = image_root + self.train_query_file = train_query_file + self.val_query_file = val_query_file + self.data_root = data_root self.output_dir = Path(output_dir) self.epochs = epochs self.batch_size = batch_size @@ -87,11 +86,8 @@ class TrainConfig: 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) @@ -99,50 +95,14 @@ def _set_seed(seed: int) -> None: 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. - """ + """Run full training loop from gin config.""" gin.parse_config_file(config_path) cfg = TrainConfig() @@ -153,21 +113,20 @@ def train(config_path: str) -> None: _set_seed(cfg.seed) cfg.output_dir.mkdir(parents=True, exist_ok=True) - # Model + loss + # Model + loss. model = DualEncoderCaptionTest().to(cfg.device) - loss_fn = MultiTermInfoNCE().to(cfg.device) + loss_fn = InfoNCELoss().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, + train_ds = GeoLocCaptionDataset( + query_file=cfg.train_query_file, + data_root=cfg.data_root, image_transform=preprocess, ) - val_ds = VisLocCaptionDataset( - manifest_path=cfg.val_manifest, - image_root=cfg.image_root, + val_ds = GeoLocCaptionDataset( + query_file=cfg.val_query_file, + data_root=cfg.data_root, image_transform=preprocess, ) @@ -197,6 +156,14 @@ def train(config_path: str) -> None: scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs) scaler = GradScaler(enabled=cfg.use_amp) + n_trainable = sum(p.numel() for p in model.trainable_parameters()) + n_total = sum(p.numel() for p in model.parameters()) + LOGGER.info( + "trainable=%d (%.2f%%) total=%d train=%d val=%d", + n_trainable, 100.0 * n_trainable / n_total, + n_total, len(train_ds), len(val_ds), + ) + history: list[dict] = [] for epoch in range(cfg.epochs): @@ -208,17 +175,25 @@ def train(config_path: str) -> None: 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, - ) + drone_img = batch["drone_img"].to(cfg.device, non_blocking=True) + sat_img = batch["sat_img"].to(cfg.device, non_blocking=True) + caption_drone = batch["caption_drone"] + with autocast(device_type="cuda", enabled=cfg.use_amp): + embeddings = model( + drone_img=drone_img, + sat_img=sat_img, + caption_drone=caption_drone, + ) + loss_dict = loss_fn( + embeddings=embeddings, + epoch=epoch, + total_epochs=cfg.epochs, + ) + + total_loss = loss_dict["total"] scaler.scale(total_loss).backward() + if cfg.grad_clip > 0: scaler.unscale_(optimizer) nn.utils.clip_grad_norm_( @@ -228,9 +203,8 @@ def train(config_path: str) -> None: 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()) + for key, val in loss_dict.items(): + agg[key] = agg.get(key, 0.0) + float(val.item()) n_batches += 1 scheduler.step() @@ -238,20 +212,15 @@ def train(config_path: str) -> None: 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, + "epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate=%.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), + means.get("gate", 1.0), ) - epoch_record = { + epoch_record: dict = { "epoch": epoch, "elapsed_seconds": elapsed, "train": means, @@ -267,18 +236,15 @@ def train(config_path: str) -> None: ) 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", + "val epoch=%d R@1_q2g=%.4f R@5_q2g=%.4f R@10_q2g=%.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), + val_metrics.get("r@1_query_to_gallery", 0.0), + val_metrics.get("r@5_query_to_gallery", 0.0), + val_metrics.get("r@10_query_to_gallery", 0.0), ) history.append(epoch_record) - # Checkpoint per epoch. _atomic_save( obj={ "epoch": epoch, @@ -289,7 +255,6 @@ def train(config_path: str) -> None: 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) @@ -299,12 +264,7 @@ def train(config_path: str) -> None: 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.", - ) + parser.add_argument("--config", type=str, required=True, help="Gin config file.") args = parser.parse_args() train(config_path=args.config)