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) <noreply@anthropic.com>
This commit is contained in:
pikaliov
208
CLAUDE.md
Normal file
208
CLAUDE.md
Normal file
@@ -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)
|
||||
136
README.md
136
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,52 +49,24 @@ 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 \
|
||||
@@ -94,45 +75,40 @@ python -m scripts.compare_runs \
|
||||
--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
|
||||
|
||||
378
UAV-VisLoc_Dataset_Analysis.md
Normal file
378
UAV-VisLoc_Dataset_Analysis.md
Normal file
@@ -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
|
||||
```
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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("")
|
||||
header = "| Direction | R@1 base → full | R@5 base → full | R@10 base → full |"
|
||||
|
||||
lines.append("## All directions x 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(_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,
|
||||
|
||||
@@ -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],
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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].
|
||||
@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).
|
||||
"""
|
||||
x = self.proj(x)
|
||||
return F.normalize(x, dim=-1)
|
||||
|
||||
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.
|
||||
# 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,
|
||||
)
|
||||
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,
|
||||
self.proj_gallery = ProjectionHead(
|
||||
in_dim=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)):
|
||||
def encode_text(self, texts: list[str]) -> torch.Tensor:
|
||||
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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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,
|
||||
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,
|
||||
device=cfg.device,
|
||||
use_amp=cfg.use_amp,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user