Add GTA-UAV experiment: asymmetric DINOv3 + LRSCLIP text encoder
V3 architecture for CVGL caption validation on GTA-UAV-LR dataset: - AsymmetricEncoder: DINOv3 ViT-L/16 (LVD drone + SAT satellite, frozen) + LRSCLIP/DGTRS-CLIP ViT-L-14 text encoder (248 tok, partial unfreeze) - L1/L2/L3 hierarchical captions from VLM-generated descriptions - TextFusionMLP (concat 3x768 -> MLP -> 512) + GatedFusion - Segmentation filter: exclude images with >=90% background+water - 10.9M trainable / 733M total params, 256x256 input - coloredlogs + tqdm + emoji for training UX - Baseline mode (--baseline): image-only, no text encoder loaded Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2
.gitignore
vendored
2
.gitignore
vendored
@@ -17,7 +17,9 @@ env/
|
|||||||
*.ckpt
|
*.ckpt
|
||||||
*.onnx
|
*.onnx
|
||||||
*.engine
|
*.engine
|
||||||
|
*.safetensors
|
||||||
checkpoints/
|
checkpoints/
|
||||||
|
nn_models/
|
||||||
out/
|
out/
|
||||||
outputs/
|
outputs/
|
||||||
runs/
|
runs/
|
||||||
|
|||||||
219
CLAUDE.md
219
CLAUDE.md
@@ -1,29 +1,49 @@
|
|||||||
# Caption Quality Test for Cross-View Geo-Localization
|
# Caption Quality Test for Cross-View Geo-Localization
|
||||||
|
|
||||||
## Архитектура системы (v2, 2026-04-17)
|
## Архитектура системы (v3, 2026-04-21) — GTA-UAV эксперимент
|
||||||
|
|
||||||
```
|
```
|
||||||
QUERY BRANCH (drone + caption):
|
QUERY BRANCH (drone + L1/L2/L3 captions):
|
||||||
drone_img --> GeoRSCLIP ViT-B/32 (frozen) --> drone_feat [B,512]
|
drone_img --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024]
|
||||||
caption --> GeoRSCLIP Text (partial unfreeze) --> text_feat [B,512]
|
|
|
||||||
|
|
proj_drone (1024->512)
|
||||||
GatedFusion: q = sigma(alpha)*drone + (1-sigma(alpha))*text
|
|
|
||||||
|
|
L1 (overview) --> LRSCLIP (248 tok) --> z_L1 [768] --\
|
||||||
proj_query (Linear 512->512) --> L2-norm --> query [B,512]
|
L2 (full desc) --> LRSCLIP (248 tok) --> z_L2 [768] ---+-- concat [B,2304]
|
||||||
|
L3 (fingerprint) --> LRSCLIP (248 tok) --> z_L3 [768] --/ |
|
||||||
|
MLP(2304->768->512)
|
||||||
|
|
|
||||||
|
GatedFusion(drone_512, text_512) --> L2-norm --> query [B,512]
|
||||||
|
|
||||||
GALLERY BRANCH (satellite only):
|
GALLERY BRANCH (satellite only):
|
||||||
sat_img --> GeoRSCLIP ViT-B/32 (frozen) --> sat_feat [B,512]
|
sat_img --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024]
|
||||||
|
|
|
|
||||||
proj_gallery (Linear 512->512) --> L2-norm --> gallery [B,512]
|
proj_sat (1024->512) --> L2-norm --> gallery [B,512]
|
||||||
|
|
||||||
LOSS: InfoNCE(query, gallery) — symmetric, asymmetric weights (0.6 q->g, 0.4 g->q)
|
LOSS: InfoNCE(query, gallery) — symmetric, asymmetric weights (0.6 q->g, 0.4 g->q)
|
||||||
BASELINE: gate = 1.0 (text ignored)
|
BASELINE: gate = 1.0 (text ignored)
|
||||||
```
|
```
|
||||||
|
|
||||||
### Trainable parameters: ~1.2M из ~151M (proj_query + proj_gallery + fusion alpha + text last_block)
|
### Trainable parameters: 10.9M из 733M (1.49%)
|
||||||
|
- proj_drone: 1024x512 = ~524K
|
||||||
|
- proj_sat: 1024x512 = ~524K
|
||||||
|
- TextFusionMLP: 2304->768->512 = ~2.2M
|
||||||
|
- gate alpha: 1 scalar
|
||||||
|
- LRSCLIP partial unfreeze (last block + ln_final + text_projection): ~7.6M
|
||||||
|
- Backbones DINOv3 x2 (303M each): frozen
|
||||||
|
|
||||||
|
### Image input: 256x256
|
||||||
|
DINOv3 ViT-L/16 с patch_size=16 → 16x16=256 patches на 256x256.
|
||||||
|
Resize(256) + CenterCrop(256) + ImageNet normalization.
|
||||||
|
|
||||||
|
### Предыдущая архитектура (v2) — UAV-GeoLoc эксперимент
|
||||||
|
|
||||||
|
Использовала GeoRSCLIP ViT-B/32 (512-dim) для обеих веток + template captions.
|
||||||
|
Код в `src/models/dual_encoder.py`, `src/datasets/visloc_with_captions.py`.
|
||||||
|
|
||||||
## Ключевые файлы
|
## Ключевые файлы
|
||||||
|
|
||||||
|
### V2 (UAV-GeoLoc, GeoRSCLIP)
|
||||||
| Файл | Назначение |
|
| Файл | Назначение |
|
||||||
|------|-----------|
|
|------|-----------|
|
||||||
| `src/models/dual_encoder.py` | GeoRSCLIP + GatedFusion + projection heads |
|
| `src/models/dual_encoder.py` | GeoRSCLIP + GatedFusion + projection heads |
|
||||||
@@ -34,13 +54,42 @@ BASELINE: gate = 1.0 (text ignored)
|
|||||||
| `scripts/compare_runs.py` | Markdown/JSON сравнение baseline vs caption runs |
|
| `scripts/compare_runs.py` | Markdown/JSON сравнение baseline vs caption runs |
|
||||||
| `scripts/generate_captions.py` | Offline caption generation (template/VLM/hybrid) |
|
| `scripts/generate_captions.py` | Offline caption generation (template/VLM/hybrid) |
|
||||||
|
|
||||||
## Backbone: GeoRSCLIP ViT-B/32
|
### V3 (GTA-UAV, DINOv3 + LRSCLIP) — DONE
|
||||||
|
| Файл | Назначение |
|
||||||
|
|------|-----------|
|
||||||
|
| `src/models/asymmetric_encoder.py` | DINOv3ViT + LRSCLIPTextEncoder + TextFusionMLP + AsymmetricEncoder + GatedFusion |
|
||||||
|
| `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 caption parsing из VLM JSON |
|
||||||
|
| `src/training/train_gtauav.py` | Training loop с eval, AMP, CLI args (--baseline, --filter-meta) |
|
||||||
|
| `scripts/filter_segmentation.py` | Scan segm masks, output meta JSON (exclude >=90% bg+water) |
|
||||||
|
|
||||||
- **Checkpoint:** `checkpoints/RS5M_ViT-B-32.pt` (скачать с github.com/om-ai-lab/RS5M)
|
## Backbones (v3)
|
||||||
|
|
||||||
|
### DINOv3 ViT-L/16 — Drone (web pretrained)
|
||||||
|
- **Checkpoint:** `nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth`
|
||||||
|
- **Arch:** ViT-L/16, 24 layers, 16 heads, hidden=1024, MLP=4096, 303M params
|
||||||
|
- **Input:** 256x256, ImageNet normalization, patch=16 → 256 patches
|
||||||
|
- **Register tokens:** 4, RoPE theta=100.0
|
||||||
|
- **Status:** frozen
|
||||||
|
|
||||||
|
### DINOv3 ViT-L/16 — Satellite (sat pretrained)
|
||||||
|
- **Checkpoint:** `nn_models/DINO_SAT/model.safetensors`
|
||||||
|
- **HuggingFace:** `facebook/dinov3-vitl16-pretrain-sat493m`
|
||||||
|
- **Arch:** идентична DINO_WEB (ViT-L/16, hidden=1024, 303M params)
|
||||||
|
- **Input:** 256x256
|
||||||
|
- **Config:** `nn_models/DINO_SAT/config.json` — BROKEN (auth error), используем конфиг от DINO_WEB
|
||||||
|
- **Status:** frozen
|
||||||
|
|
||||||
|
### DGTRS-CLIP ViT-L-14 (LRSCLIP) — Text encoder
|
||||||
|
- **Checkpoint:** `nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt`
|
||||||
|
- **Text dim:** 768, max tokens: 248 (KPS stretched from CLIP's 77)
|
||||||
|
- **Содержит:** полную CLIP модель (visual + text), используем только text encoder
|
||||||
|
- **GitHub:** `github.com/MitsuiChen14/DGTRS`
|
||||||
|
- **Status:** partial unfreeze (last block + text_projection)
|
||||||
|
|
||||||
|
### GeoRSCLIP ViT-B/32 (v2, legacy)
|
||||||
|
- **Checkpoint:** `checkpoints/RS5M_ViT-B-32.pt`
|
||||||
- **Image encoder:** ViT-B/32, 224x224, 512-dim, ~86M params — frozen
|
- **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)
|
- **Text encoder:** 77 tokens, 512-dim — partial unfreeze
|
||||||
- **Throughput:** ~4000 img/s на RTX 4090 (AMP, batch 128)
|
|
||||||
- **Выбран вместо SigLIP 2** (ViT-SO400M, 384px, ~400M): в 7-10x быстрее, domain-specific (обучен на 5M RS-изображений), больше batch = больше негативов в InfoNCE
|
|
||||||
|
|
||||||
## GatedFusion
|
## GatedFusion
|
||||||
|
|
||||||
@@ -50,49 +99,91 @@ BASELINE: gate = 1.0 (text ignored)
|
|||||||
- `baseline_mode = True` → gate = 1.0, text полностью игнорируется
|
- `baseline_mode = True` → gate = 1.0, text полностью игнорируется
|
||||||
- Gate value логируется каждую эпоху для интерпретации вклада текста
|
- Gate value логируется каждую эпоху для интерпретации вклада текста
|
||||||
|
|
||||||
## Датасет: UAV-GeoLoc
|
## Text Hierarchy (L1/L2/L3)
|
||||||
|
|
||||||
|
Три уровня описаний из VLM-generated captions:
|
||||||
|
|
||||||
|
| Уровень | Контент | Длина | Источник |
|
||||||
|
|---|---|---|---|
|
||||||
|
| L1 overview | Краткое описание сцены | <=30 tok | Конденсация P1 |
|
||||||
|
| L2 full description | Детальное описание через Qwen3-VL | <=200 tok | Полный P1+P2 |
|
||||||
|
| L3 fingerprint | Ключевые landmark'ы | <=30 tok | Конденсация P3 |
|
||||||
|
|
||||||
|
Все три уровня кодируются одним LRSCLIP (248 tok max).
|
||||||
|
Альтернатива (Stage 2): RemoteCLIP для L1/L3 + LRSCLIP для L2.
|
||||||
|
|
||||||
|
## Датасет: GTA-UAV-LR
|
||||||
|
|
||||||
|
- **RGB:** `/home/servml/Документы/datasets/GTA-UAV-LR/`
|
||||||
|
- Drone: 33,763 PNG (512x384), altitudes 100-600m
|
||||||
|
- Satellite: 14,640 PNG (256x256 RGBA)
|
||||||
|
- Pairs: `cross-area-drone2sate-{train,test}.json` (primary split)
|
||||||
|
- Metadata: `*_drone_meta.csv` (height, yaw, roll, pitch)
|
||||||
|
- Origin: GTA V simulation (Los Santos)
|
||||||
|
|
||||||
|
- **Captions:** `/home/servml/Документы/datasets/GTA-UAV-LR-captions/`
|
||||||
|
- Drone: 33,411 JSON (32,635 multi-paragraph P1/P2/P3 + 776 short water-only)
|
||||||
|
- Satellite: 6,546 JSON (все multi-paragraph)
|
||||||
|
- Формат: 3 абзаца (P1 Inventory + P2 Spatial + P3 Fingerprint)
|
||||||
|
- Token counts: ~430 output tokens per caption
|
||||||
|
|
||||||
|
- **Segmentation:** `/home/servml/Документы/datasets/GTA-UAV-LR-aug/`
|
||||||
|
- 48,403 images, 17 классов (background, building, road, vegetation, water, ...)
|
||||||
|
- Modalities: segm/, depth/, edge/, chm/, safetensors/
|
||||||
|
- Query: 512x512, DB: 256x256
|
||||||
|
|
||||||
|
### Фильтрация сегментации
|
||||||
|
Meta-файл `meta/seg_filter.json`: исключение изображений с >=90% background(class 0) + water(class 4).
|
||||||
|
- **Total:** 48,403 → **Passed:** 37,498 (77.5%) / **Excluded:** 10,905 (22.5%)
|
||||||
|
- Drone: 31,188 passed / 2,575 excluded
|
||||||
|
- Satellite: 6,310 passed / 8,330 excluded (преимущественно open water tiles)
|
||||||
|
|
||||||
|
## Датасет: UAV-GeoLoc (v2, legacy)
|
||||||
|
|
||||||
- **Путь:** `/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/`
|
- **Путь:** `/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/`
|
||||||
- **Train:** 206,108 queries, 94,709 DB crops (140 scenes, Terrain split)
|
- **Train:** 206,108 queries, 94,709 DB crops (140 scenes, Terrain split)
|
||||||
- **Val:** 62,368 queries, 26,597 DB crops (40 scenes)
|
- **Val:** 62,368 queries, 26,597 DB crops (40 scenes)
|
||||||
- **Test:** 33,472 queries, 11,684 DB crops (20 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 ...`
|
- **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}."`
|
### V3 (GTA-UAV)
|
||||||
|
Параметры: 10 epochs, batch 64, lr=1e-4, AMP, cosine LR, image 256x256.
|
||||||
Пример: `"Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano. Plan-view features: lava flows, crater edges, volcanic rock."`
|
Train: 15,693 pairs (cross-area), Test: 18,015 pairs.
|
||||||
|
Eval на cross-area split (primary).
|
||||||
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)
|
|
||||||
|
|
||||||
|
### V2 (UAV-GeoLoc, gin)
|
||||||
| Конфиг | Gate init | Описание |
|
| Конфиг | Gate init | Описание |
|
||||||
|--------|-----------|----------|
|
|--------|-----------|----------|
|
||||||
| `conf/balanced.gin` | 0.7 (30% text) | **Primary test** |
|
| `conf/balanced.gin` | 0.7 (30% text) | **Primary test** |
|
||||||
| `conf/baseline_no_text.gin` | 1.0 (no text) | Reference baseline |
|
| `conf/baseline_no_text.gin` | 1.0 (no text) | Reference baseline |
|
||||||
| `conf/text_heavy.gin` | 0.3 (70% text) | Stress test |
|
| `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.
|
|
||||||
|
|
||||||
## Запуск
|
## Запуск
|
||||||
|
|
||||||
|
### V3 (GTA-UAV)
|
||||||
```bash
|
```bash
|
||||||
# 1. Baseline (no text)
|
# 1. Filter segmentation (exclude 90%+ background/water)
|
||||||
python -m src.training.train --config conf/baseline_no_text.gin
|
python -m scripts.filter_segmentation --output meta/seg_filter.json
|
||||||
|
|
||||||
# 2. With captions (primary test)
|
# 2. Baseline (no text)
|
||||||
python -m src.training.train --config conf/balanced.gin
|
python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json
|
||||||
|
|
||||||
# 3. Text-heavy (stress test)
|
# 3. With captions (L1/L2/L3)
|
||||||
python -m src.training.train --config conf/text_heavy.gin
|
python -m src.training.train_gtauav --filter-meta meta/seg_filter.json
|
||||||
|
|
||||||
# 4. Compare
|
# 4. Compare
|
||||||
|
python -m scripts.compare_runs \
|
||||||
|
--baseline_report out/gtauav/baseline/eval_report.json \
|
||||||
|
--full_report out/gtauav/with_text/eval_report.json \
|
||||||
|
--output out/gtauav/comparison.md
|
||||||
|
```
|
||||||
|
|
||||||
|
### V2 (UAV-GeoLoc)
|
||||||
|
```bash
|
||||||
|
python -m src.training.train --config conf/baseline_no_text.gin
|
||||||
|
python -m src.training.train --config conf/balanced.gin
|
||||||
python -m scripts.compare_runs \
|
python -m scripts.compare_runs \
|
||||||
--baseline_report out/caption_test/baseline_no_text/eval_report.json \
|
--baseline_report out/caption_test/baseline_no_text/eval_report.json \
|
||||||
--full_report out/caption_test/balanced/eval_report.json \
|
--full_report out/caption_test/balanced/eval_report.json \
|
||||||
@@ -101,7 +192,7 @@ python -m scripts.compare_runs \
|
|||||||
|
|
||||||
## Метрики и Decision rule
|
## Метрики и Decision rule
|
||||||
|
|
||||||
**Primary metric:** Delta R@1 (query -> gallery)
|
**Primary metric:** Delta R@1 (drone -> satellite)
|
||||||
|
|
||||||
| Delta R@1 | Verdict |
|
| Delta R@1 | Verdict |
|
||||||
|-----------|---------|
|
|-----------|---------|
|
||||||
@@ -110,11 +201,21 @@ python -m scripts.compare_runs \
|
|||||||
| 0 to +1% | WEAK — redesign caption pipeline |
|
| 0 to +1% | WEAK — redesign caption pipeline |
|
||||||
| < 0 | HARMFUL — critical bug |
|
| < 0 | HARMFUL — critical bug |
|
||||||
|
|
||||||
|
**Eval metrics:** R@1, R@5, R@10 для drone->satellite и satellite->drone
|
||||||
|
**Splits (GTA-UAV):** cross-area (primary, harder) и same-area (sanity check)
|
||||||
**Logged per epoch:** loss, temperature (tau), gate value (sigma(alpha)), lr
|
**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)
|
## Бюджет времени (RTX 4090, 24 GB)
|
||||||
|
|
||||||
|
### V3 (GTA-UAV, DINOv3 ViT-L/16, 256x256)
|
||||||
|
| Фаза | Оценка |
|
||||||
|
|------|--------|
|
||||||
|
| VRAM: 2x DINOv3-L + LRSCLIP + batch 64 | ~18-22 GB |
|
||||||
|
| GPU mem (smoke test, batch 4) | 3.1 GB |
|
||||||
|
| Batch size | 64 (default) |
|
||||||
|
| Total params | 733M (10.9M trainable, 1.49%) |
|
||||||
|
|
||||||
|
### V2 (UAV-GeoLoc, GeoRSCLIP)
|
||||||
| Фаза | Время |
|
| Фаза | Время |
|
||||||
|------|-------|
|
|------|-------|
|
||||||
| Один training run (10 epochs, 206K queries, batch 128) | ~15-30 мин |
|
| Один training run (10 epochs, 206K queries, batch 128) | ~15-30 мин |
|
||||||
@@ -138,6 +239,42 @@ python -m scripts.compare_runs \
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
## Downstream: NADEZHDA Teacher (DINOv3 + Multi-FiLM)
|
||||||
|
|
||||||
|
caption-test — первый этап валидации. При Δ R@1 >= +3% переход к полному teacher'у:
|
||||||
|
|
||||||
|
```
|
||||||
|
Этап 1 (caption-test): GeoRSCLIP + GatedFusion(text) → валидация текста
|
||||||
|
Этап 2 (teacher): DINOv3-L + Multi-FiLM(depth, seg, CHM, normals, text)
|
||||||
|
Этап 3 (distillation): teacher ~300M → student ~5M → Jetson Orin NX
|
||||||
|
```
|
||||||
|
|
||||||
|
### Auxiliary modalities (предвычисляются из 512x512 офлайн)
|
||||||
|
|
||||||
|
| Модальность | Модель | Формат для teacher | Каналы |
|
||||||
|
|---|---|---|---|
|
||||||
|
| Depth | DepthAnything V2 | continuous, log(1+d) | 1 |
|
||||||
|
| Normals | Sobel от depth | continuous | 3 |
|
||||||
|
| Segmentation | SegFormer-B5 | binary per-class masks (top-K) | 16-17 |
|
||||||
|
| Canopy Height | Meta HRCH | binary bins (1-5m, 5-15m, >15m) + occupancy | 4-5 |
|
||||||
|
| Text | Qwen3-VL-8B / MobileCLIP2 | embedding | - |
|
||||||
|
|
||||||
|
### Асимметрия sat/drone
|
||||||
|
- CHM: только satellite (модель обучена на nadir, на oblique drone не работает)
|
||||||
|
- Satellite: ~27 aux каналов, Drone: ~21 aux каналов
|
||||||
|
|
||||||
|
### Fusion: Multi-FiLM
|
||||||
|
```
|
||||||
|
aux_features → FiLM(γ, β) → γ * DINOv3_tokens + β
|
||||||
|
```
|
||||||
|
Binary masks — natural FiLM gates. Modality dropout: text 0.3, CHM 0.5, seg 0.15, depth 0.1.
|
||||||
|
|
||||||
|
### Планируемый эксперимент H5
|
||||||
|
Сравнение T_bin (binary masks) vs T_pyr (native feature pyramids) vs T_hybrid.
|
||||||
|
Прогноз: binary masks лучше на cross-domain из-за робастности к aux-model artifacts.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
## Датасеты (справочник)
|
## Датасеты (справочник)
|
||||||
|
|
||||||
### UAV-VisLoc
|
### UAV-VisLoc
|
||||||
@@ -197,7 +334,7 @@ python scripts/prepare_dataset.py \
|
|||||||
- Размер на диске: 25 GB
|
- Размер на диске: 25 GB
|
||||||
- Median distance drone->crop: 25.9m, P99: 45.7m
|
- Median distance drone->crop: 25.9m, P99: 45.7m
|
||||||
- Память при генерации: до ~8.7 GB RAM (маршрут 09 stitched 44800x33280)
|
- Память при генерации: до ~8.7 GB RAM (маршрут 09 stitched 44800x33280)
|
||||||
- Разрешение 512 для downstream задач (сегментация, depth, normals); resize до 224/256 в dataloader
|
- Разрешение 512 для downstream задач (сегментация, depth, normals); resize до 256x256 в dataloader
|
||||||
|
|
||||||
### Ревью и исправления (2026-04-17)
|
### Ревью и исправления (2026-04-17)
|
||||||
1. `train_db.txt`/`test_db.txt` содержали только matched кропы -> теперь все ~74K (полная gallery)
|
1. `train_db.txt`/`test_db.txt` содержали только matched кропы -> теперь все ~74K (полная gallery)
|
||||||
|
|||||||
128
README.md
128
README.md
@@ -1,42 +1,77 @@
|
|||||||
# Caption Quality Test for Cross-View Geo-Localization
|
# Caption Quality Test for Cross-View Geo-Localization
|
||||||
|
|
||||||
Validate whether generated text captions improve retrieval R@1 in cross-view
|
Validate whether generated text captions improve retrieval R@1 in cross-view
|
||||||
geo-localization (drone-to-satellite). Uses GeoRSCLIP ViT-B/32 dual encoder
|
geo-localization (drone-to-satellite).
|
||||||
with GatedFusion on the query branch.
|
|
||||||
|
|
||||||
## Architecture
|
## Experiments
|
||||||
|
|
||||||
|
### V3 — GTA-UAV + DINOv3 + LRSCLIP (active)
|
||||||
|
|
||||||
|
Asymmetric architecture with domain-specific image encoders and hierarchical text.
|
||||||
|
|
||||||
```
|
```
|
||||||
Query: drone_img + caption -> GatedFusion -> proj -> query_emb
|
Query: drone_img (DINOv3 LVD) + L1/L2/L3 captions (LRSCLIP) -> GatedFusion -> query
|
||||||
Gallery: sat_img -> proj -> gallery_emb
|
Gallery: sat_img (DINOv3 SAT) -> gallery
|
||||||
Loss: InfoNCE(query, gallery)
|
Loss: InfoNCE(query, gallery)
|
||||||
```
|
```
|
||||||
|
|
||||||
Baseline: fusion gate = 1.0 (text ignored).
|
**Models:**
|
||||||
|
- Drone: DINOv3 ViT-L/16 (LVD-1689M, web pretrained) — 1024-dim, 303M params, frozen
|
||||||
|
- Satellite: DINOv3 ViT-L/16 (SAT-493M, satellite pretrained) — 1024-dim, 303M params, frozen
|
||||||
|
- Text: DGTRS-CLIP ViT-L-14 (LRSCLIP, 248 tokens) — 768-dim, partial unfreeze
|
||||||
|
- Total: 733M params, 10.9M trainable (1.49%)
|
||||||
|
|
||||||
|
**Input:** 256x256, ImageNet normalization
|
||||||
|
|
||||||
|
**Dataset:** GTA-UAV-LR (33K drone + 14K satellite, GTA V synthetic)
|
||||||
|
- RGB: `/home/servml/Документы/datasets/GTA-UAV-LR/`
|
||||||
|
- Captions: `/home/servml/Документы/datasets/GTA-UAV-LR-captions/` (40K JSON, 3-paragraph VLM)
|
||||||
|
- Segmentation: `/home/servml/Документы/datasets/GTA-UAV-LR-aug/` (17 classes)
|
||||||
|
- Seg filter: 37,498 passed / 10,905 excluded (>=90% background+water)
|
||||||
|
- Train: 15,693 pairs, Test: 18,015 pairs (cross-area split)
|
||||||
|
|
||||||
|
### V2 — UAV-GeoLoc + GeoRSCLIP (legacy)
|
||||||
|
|
||||||
|
Single backbone with template captions.
|
||||||
|
|
||||||
|
```
|
||||||
|
Query: drone_img + caption -> GeoRSCLIP -> GatedFusion -> query
|
||||||
|
Gallery: sat_img -> GeoRSCLIP -> gallery
|
||||||
|
```
|
||||||
|
|
||||||
|
**Dataset:** UAV-GeoLoc Terrain split (206K train queries)
|
||||||
|
|
||||||
## Structure
|
## Structure
|
||||||
|
|
||||||
```
|
```
|
||||||
caption_test/
|
caption-test/
|
||||||
├── conf/
|
├── conf/ # Gin configs (v2)
|
||||||
│ ├── balanced.gin # Primary: gate init 0.7 (30% text)
|
│ ├── balanced.gin
|
||||||
│ ├── baseline_no_text.gin # Reference: gate = 1.0 (no text)
|
│ ├── baseline_no_text.gin
|
||||||
│ └── text_heavy.gin # Stress: gate init 0.3 (70% text)
|
│ └── text_heavy.gin
|
||||||
|
├── nn_models/ # Pre-trained checkpoints (v3)
|
||||||
|
│ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M
|
||||||
|
│ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M
|
||||||
|
│ └── LRSCLIP/ # DGTRS-CLIP ViT-L-14
|
||||||
├── scripts/
|
├── scripts/
|
||||||
│ ├── generate_captions.py # Offline caption generation
|
│ ├── filter_segmentation.py # Meta-file: exclude 90%+ background/water
|
||||||
│ └── compare_runs.py # Delta R@1 comparison report
|
│ ├── compare_runs.py # Delta R@1 comparison report
|
||||||
|
│ └── generate_captions.py # Offline caption generation
|
||||||
├── src/
|
├── src/
|
||||||
│ ├── datasets/
|
│ ├── datasets/
|
||||||
│ │ └── visloc_with_captions.py # UAV-GeoLoc loader + template captions
|
│ │ ├── gtauav_dataset.py # GTA-UAV-LR loader + L1/L2/L3 captions (v3)
|
||||||
|
│ │ └── visloc_with_captions.py # UAV-GeoLoc loader (v2)
|
||||||
│ ├── models/
|
│ ├── models/
|
||||||
│ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion + projection heads
|
│ │ ├── asymmetric_encoder.py # DINOv3 + LRSCLIP + GatedFusion (v3)
|
||||||
|
│ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2)
|
||||||
│ ├── losses/
|
│ ├── losses/
|
||||||
│ │ └── multi_infonce.py # InfoNCE with cosine temperature
|
│ │ └── multi_infonce.py # InfoNCE with cosine temperature
|
||||||
│ ├── training/
|
│ ├── training/
|
||||||
│ │ └── train.py # Main training loop
|
│ │ ├── train_gtauav.py # Training loop GTA-UAV (v3)
|
||||||
|
│ │ └── train.py # Training loop UAV-GeoLoc (v2)
|
||||||
│ └── eval/
|
│ └── eval/
|
||||||
│ └── evaluate.py # R@K metrics, Delta R@1
|
│ └── evaluate.py # R@K metrics, Delta R@1
|
||||||
└── checkpoints/ # RS5M_ViT-B-32.pt (user-provided)
|
└── checkpoints/ # GeoRSCLIP RS5M_ViT-B-32.pt (v2)
|
||||||
```
|
```
|
||||||
|
|
||||||
## Prerequisites
|
## Prerequisites
|
||||||
@@ -44,71 +79,54 @@ caption_test/
|
|||||||
```
|
```
|
||||||
torch>=2.0
|
torch>=2.0
|
||||||
open_clip_torch
|
open_clip_torch
|
||||||
|
safetensors
|
||||||
|
timm
|
||||||
gin-config
|
gin-config
|
||||||
Pillow
|
Pillow
|
||||||
numpy
|
numpy
|
||||||
```
|
```
|
||||||
|
|
||||||
GeoRSCLIP checkpoint: download `RS5M_ViT-B-32.pt` from
|
## Workflow (V3 — GTA-UAV)
|
||||||
`github.com/om-ai-lab/RS5M` and place under `checkpoints/`.
|
|
||||||
|
|
||||||
## Workflow
|
### 1. Filter segmentation (exclude uninformative images)
|
||||||
|
|
||||||
### 1. Train baseline (no text)
|
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python -m src.training.train --config conf/baseline_no_text.gin
|
python -m scripts.filter_segmentation --output meta/seg_filter.json
|
||||||
```
|
```
|
||||||
|
|
||||||
### 2. Train with captions
|
### 2. Train baseline (no text)
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python -m src.training.train --config conf/balanced.gin
|
python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json
|
||||||
```
|
```
|
||||||
|
|
||||||
### 3. Compare and get verdict
|
### 3. Train with captions (L1/L2/L3)
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -m src.training.train_gtauav --filter-meta meta/seg_filter.json
|
||||||
|
```
|
||||||
|
|
||||||
|
### 4. Compare and get verdict
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python -m scripts.compare_runs \
|
python -m scripts.compare_runs \
|
||||||
--baseline_report out/caption_test/baseline_no_text/eval_report.json \
|
--baseline_report out/gtauav/baseline/eval_report.json \
|
||||||
--full_report out/caption_test/balanced/eval_report.json \
|
--full_report out/gtauav/with_text/eval_report.json \
|
||||||
--output out/caption_test/comparison.md
|
--output out/gtauav/comparison.md
|
||||||
```
|
```
|
||||||
|
|
||||||
## Decision rule
|
## Decision rule
|
||||||
|
|
||||||
| Delta R@1 (query->gallery) | Verdict |
|
| Delta R@1 (drone->satellite) | Verdict |
|
||||||
|---|---|
|
|---|---|
|
||||||
| >= +3% | PASS -- captions informative, proceed to production |
|
| >= +3% | PASS -- captions informative, proceed to NADEZHDA teacher |
|
||||||
| +1% to +3% | MARGINAL -- add VLM refinement, re-run |
|
| +1% to +3% | MARGINAL -- add VLM refinement, re-run |
|
||||||
| 0 to +1% | WEAK -- redesign caption pipeline |
|
| 0 to +1% | WEAK -- redesign caption pipeline |
|
||||||
| < 0 | HARMFUL -- critical bug |
|
| < 0 | HARMFUL -- critical bug |
|
||||||
|
|
||||||
## Expected runtime (RTX 4090, 24 GB)
|
|
||||||
|
|
||||||
| Phase | Time |
|
|
||||||
|---|---|
|
|
||||||
| 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 |
|
|
||||||
|
|
||||||
## Dataset
|
|
||||||
|
|
||||||
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)
|
|
||||||
|
|
||||||
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."
|
|
||||||
```
|
|
||||||
|
|
||||||
## Code style
|
## Code style
|
||||||
|
|
||||||
- `from __future__ import annotations` everywhere
|
- `from __future__ import annotations` everywhere
|
||||||
- Type hints on all signatures
|
- Type hints on all signatures
|
||||||
- Google-style docstrings
|
- Google-style docstrings
|
||||||
- `@gin.configurable` on top-level classes
|
|
||||||
- No emojis in code, English-only comments
|
- No emojis in code, English-only comments
|
||||||
|
|||||||
48417
meta/seg_filter.json
Normal file
48417
meta/seg_filter.json
Normal file
File diff suppressed because it is too large
Load Diff
140
scripts/filter_segmentation.py
Normal file
140
scripts/filter_segmentation.py
Normal file
@@ -0,0 +1,140 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
"""Filter GTA-UAV-LR images by segmentation class coverage.
|
||||||
|
|
||||||
|
Reads palette-mode PNG segmentation masks, computes per-class pixel ratios,
|
||||||
|
and outputs a JSON meta file listing images that pass the filter
|
||||||
|
(i.e. background + water < threshold).
|
||||||
|
|
||||||
|
Classes (from manifest.json):
|
||||||
|
0: background, 1: building, 2: road, 3: vegetation, 4: water, ...
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python -m scripts.filter_segmentation [--threshold 0.9] [--output meta/seg_filter.json]
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import coloredlogs
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
LOGGER = logging.getLogger("caption_test.filter_seg")
|
||||||
|
|
||||||
|
SEGM_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-aug/segm")
|
||||||
|
EXCLUDE_CLASSES = {0, 4} # background, water
|
||||||
|
DEFAULT_THRESHOLD = 0.90
|
||||||
|
|
||||||
|
|
||||||
|
def compute_class_ratios(mask_path: Path) -> dict[int, float]:
|
||||||
|
"""Load a palette-mode PNG mask and return per-class pixel ratios."""
|
||||||
|
with Image.open(mask_path) as img:
|
||||||
|
arr = np.array(img)
|
||||||
|
total = arr.size
|
||||||
|
unique, counts = np.unique(arr, return_counts=True)
|
||||||
|
return {int(cls): float(cnt / total) for cls, cnt in zip(unique, counts)}
|
||||||
|
|
||||||
|
|
||||||
|
def scan_masks(
|
||||||
|
segm_root: Path,
|
||||||
|
exclude_classes: set[int],
|
||||||
|
threshold: float,
|
||||||
|
) -> dict[str, dict]:
|
||||||
|
"""Scan all segmentation masks and classify as pass/fail.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dict keyed by relative image name (e.g. "drone/images/100_0001_0000000000.png")
|
||||||
|
with values: {"ratios": {...}, "excluded_ratio": float, "pass": bool}
|
||||||
|
"""
|
||||||
|
results: dict[str, dict] = {}
|
||||||
|
subdirs = sorted(p for p in segm_root.iterdir() if p.is_dir())
|
||||||
|
|
||||||
|
for subdir in subdirs:
|
||||||
|
mask_files = sorted(subdir.rglob("*.png"))
|
||||||
|
LOGGER.info("🔍 Scanning %s: %d masks", subdir.name, len(mask_files))
|
||||||
|
|
||||||
|
for mask_path in tqdm(mask_files, desc=f" 📂 {subdir.name}", unit="img", leave=False):
|
||||||
|
rel = mask_path.relative_to(segm_root)
|
||||||
|
ratios = compute_class_ratios(mask_path)
|
||||||
|
excluded_ratio = sum(ratios.get(c, 0.0) for c in exclude_classes)
|
||||||
|
results[str(rel)] = {
|
||||||
|
"ratios": ratios,
|
||||||
|
"excluded_ratio": round(excluded_ratio, 6),
|
||||||
|
"pass": excluded_ratio < threshold,
|
||||||
|
}
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
def build_meta(
|
||||||
|
results: dict[str, dict],
|
||||||
|
) -> dict[str, list[str]]:
|
||||||
|
"""Split results into passed and excluded lists."""
|
||||||
|
passed = sorted(k for k, v in results.items() if v["pass"])
|
||||||
|
excluded = sorted(k for k, v in results.items() if not v["pass"])
|
||||||
|
return {"passed": passed, "excluded": excluded}
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Filter GTA-UAV-LR images by segmentation coverage.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--segm-root", type=str, default=str(SEGM_ROOT),
|
||||||
|
help="Root of segmentation masks.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--threshold", type=float, default=DEFAULT_THRESHOLD,
|
||||||
|
help="Exclude images where background+water >= threshold.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--output", type=str, default="meta/seg_filter.json",
|
||||||
|
help="Output JSON meta file path.",
|
||||||
|
)
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
coloredlogs.install(
|
||||||
|
level="INFO",
|
||||||
|
logger=LOGGER,
|
||||||
|
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||||||
|
)
|
||||||
|
|
||||||
|
LOGGER.info("🚀 Starting segmentation filter (threshold=%.2f)", args.threshold)
|
||||||
|
|
||||||
|
segm_root = Path(args.segm_root)
|
||||||
|
results = scan_masks(segm_root, EXCLUDE_CLASSES, args.threshold)
|
||||||
|
meta = build_meta(results)
|
||||||
|
|
||||||
|
n_total = len(results)
|
||||||
|
n_pass = len(meta["passed"])
|
||||||
|
n_excl = len(meta["excluded"])
|
||||||
|
LOGGER.info(
|
||||||
|
"📊 total=%d ✅ passed=%d (%.1f%%) ❌ excluded=%d (%.1f%%)",
|
||||||
|
n_total, n_pass, 100 * n_pass / max(n_total, 1),
|
||||||
|
n_excl, 100 * n_excl / max(n_total, 1),
|
||||||
|
)
|
||||||
|
|
||||||
|
output_path = Path(args.output)
|
||||||
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
out = {
|
||||||
|
"threshold": args.threshold,
|
||||||
|
"exclude_classes": sorted(EXCLUDE_CLASSES),
|
||||||
|
"total_images": n_total,
|
||||||
|
"passed_count": n_pass,
|
||||||
|
"excluded_count": n_excl,
|
||||||
|
"passed": meta["passed"],
|
||||||
|
"excluded": meta["excluded"],
|
||||||
|
}
|
||||||
|
with output_path.open("w", encoding="utf-8") as f:
|
||||||
|
json.dump(out, f, indent=2)
|
||||||
|
|
||||||
|
LOGGER.info("💾 Meta file saved to %s", output_path)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
"""Dataset loaders for caption quality test."""
|
"""Dataset loaders for caption quality test."""
|
||||||
from src.datasets.visloc_with_captions import (
|
from src.datasets.visloc_with_captions import (
|
||||||
VisLocCaptionDataset,
|
GeoLocCaptionDataset,
|
||||||
collate_caption_batch,
|
collate_caption_batch,
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = ["VisLocCaptionDataset", "collate_caption_batch"]
|
__all__ = ["GeoLocCaptionDataset", "collate_caption_batch"]
|
||||||
|
|||||||
247
src/datasets/gtauav_dataset.py
Normal file
247
src/datasets/gtauav_dataset.py
Normal file
@@ -0,0 +1,247 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
"""GTA-UAV-LR dataset loader with L1/L2/L3 hierarchical captions for CVGL.
|
||||||
|
|
||||||
|
Reads cross-area pair JSONs and VLM-generated caption JSONs.
|
||||||
|
Produces (drone_img, sat_img, caption_l1, caption_l2, caption_l3) tuples.
|
||||||
|
|
||||||
|
Caption levels:
|
||||||
|
L1 (overview): First sentence of P1 (land cover summary).
|
||||||
|
L2 (full): Complete P1 + P2 text.
|
||||||
|
L3 (fingerprint): P3 unique signature section.
|
||||||
|
|
||||||
|
For short captions (pure water, no P markers): all levels get the same short text.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import random
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Callable
|
||||||
|
|
||||||
|
import coloredlogs
|
||||||
|
import torch
|
||||||
|
from PIL import Image
|
||||||
|
from torch.utils.data import Dataset
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
LOGGER = logging.getLogger("caption_test.gtauav_dataset")
|
||||||
|
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||||
|
|
||||||
|
# Default paths.
|
||||||
|
_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
|
||||||
|
_CAPTION_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-captions")
|
||||||
|
_EMPTY_CAPTION = ""
|
||||||
|
|
||||||
|
# Regex to split P1/P2/P3 sections.
|
||||||
|
_P_SPLIT = re.compile(r"\*\*P[123][^*]*\*\*\s*:?\s*")
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_caption_levels(output: str) -> tuple[str, str, str]:
|
||||||
|
"""Split VLM caption output into L1, L2, L3 levels.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(l1_overview, l2_full, l3_fingerprint)
|
||||||
|
"""
|
||||||
|
sections = _P_SPLIT.split(output)
|
||||||
|
# sections[0] is empty (before **P1**), sections[1]=P1, [2]=P2, [3]=P3
|
||||||
|
sections = [s.strip() for s in sections if s.strip()]
|
||||||
|
|
||||||
|
if len(sections) >= 3:
|
||||||
|
p1, p2, p3 = sections[0], sections[1], sections[2]
|
||||||
|
elif len(sections) == 2:
|
||||||
|
p1, p2, p3 = sections[0], sections[1], sections[0]
|
||||||
|
elif len(sections) == 1:
|
||||||
|
p1 = p2 = p3 = sections[0]
|
||||||
|
else:
|
||||||
|
p1 = p2 = p3 = output.strip()
|
||||||
|
|
||||||
|
# L1: first sentence of P1.
|
||||||
|
first_dot = p1.find(". ")
|
||||||
|
l1 = p1[:first_dot + 1] if first_dot > 0 else p1
|
||||||
|
|
||||||
|
# L2: full P1 + P2.
|
||||||
|
l2 = p1 + " " + p2
|
||||||
|
|
||||||
|
# L3: P3 (fingerprint / unique signature).
|
||||||
|
l3 = p3
|
||||||
|
|
||||||
|
return l1, l2, l3
|
||||||
|
|
||||||
|
|
||||||
|
def _load_caption_index(caption_root: Path) -> dict[str, dict]:
|
||||||
|
"""Build index: image_name -> caption JSON data.
|
||||||
|
|
||||||
|
Scans drone/images/ and satellite/ directories.
|
||||||
|
"""
|
||||||
|
index: dict[str, dict] = {}
|
||||||
|
|
||||||
|
for subdir in ["drone/images", "satellite"]:
|
||||||
|
cap_dir = caption_root / subdir
|
||||||
|
if not cap_dir.exists():
|
||||||
|
continue
|
||||||
|
cap_files = sorted(cap_dir.glob("*_caption.json"))
|
||||||
|
for cap_file in tqdm(cap_files, desc=f" 📄 Loading {subdir} captions", unit="file", leave=False):
|
||||||
|
with open(cap_file) as f:
|
||||||
|
data = json.load(f)
|
||||||
|
# Key by the image name (without _caption suffix).
|
||||||
|
img_name = cap_file.name.replace("_caption.json", ".png")
|
||||||
|
index[img_name] = data
|
||||||
|
|
||||||
|
return index
|
||||||
|
|
||||||
|
|
||||||
|
class GTAUAVDataset(Dataset):
|
||||||
|
"""GTA-UAV-LR dataset with hierarchical L1/L2/L3 captions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pair_json: Path to cross-area-drone2sate-{train,test}.json.
|
||||||
|
rgb_root: Root of GTA-UAV-LR RGB images.
|
||||||
|
caption_root: Root of GTA-UAV-LR-captions.
|
||||||
|
image_transform: Callable applied to PIL images.
|
||||||
|
filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
|
||||||
|
drop_caption_prob: Probability of dropping captions (ablation).
|
||||||
|
seed: Random seed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
pair_json: str,
|
||||||
|
rgb_root: str = str(_RGB_ROOT),
|
||||||
|
caption_root: str = str(_CAPTION_ROOT),
|
||||||
|
image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
|
||||||
|
filter_meta: str | None = None,
|
||||||
|
drop_caption_prob: float = 0.0,
|
||||||
|
seed: int = 0,
|
||||||
|
) -> None:
|
||||||
|
self.rgb_root = Path(rgb_root)
|
||||||
|
self.caption_root = Path(caption_root)
|
||||||
|
self.image_transform = image_transform
|
||||||
|
self.drop_caption_prob = drop_caption_prob
|
||||||
|
self._rng = random.Random(seed)
|
||||||
|
|
||||||
|
# Load exclusion set from segmentation filter.
|
||||||
|
self.excluded: set[str] = set()
|
||||||
|
if filter_meta is not None:
|
||||||
|
self._load_filter(Path(filter_meta))
|
||||||
|
|
||||||
|
# Load caption index.
|
||||||
|
LOGGER.info("📚 Loading caption index from %s", caption_root)
|
||||||
|
self.caption_index = _load_caption_index(self.caption_root)
|
||||||
|
LOGGER.info("📚 Caption index: %d entries", len(self.caption_index))
|
||||||
|
|
||||||
|
# Load pairs.
|
||||||
|
self.entries: list[dict[str, Any]] = []
|
||||||
|
self._load_pairs(Path(pair_json))
|
||||||
|
LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json)
|
||||||
|
|
||||||
|
def _load_filter(self, path: Path) -> None:
|
||||||
|
with open(path) as f:
|
||||||
|
meta = json.load(f)
|
||||||
|
# excluded list contains paths like "drone/images/xxx.png"
|
||||||
|
for exc in meta.get("excluded", []):
|
||||||
|
# Extract image name from segm path.
|
||||||
|
self.excluded.add(Path(exc).name)
|
||||||
|
LOGGER.info("🔻 Filter loaded: %d excluded images", len(self.excluded))
|
||||||
|
|
||||||
|
def _load_pairs(self, pair_json: Path) -> None:
|
||||||
|
with open(pair_json) as f:
|
||||||
|
raw_pairs = json.load(f)
|
||||||
|
|
||||||
|
for pair in raw_pairs:
|
||||||
|
drone_name = pair["drone_img_name"]
|
||||||
|
|
||||||
|
# Skip excluded images.
|
||||||
|
if drone_name in self.excluded:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Get positive/semi-positive satellite images.
|
||||||
|
pos_list = pair.get("pair_pos_sate_img_list", [])
|
||||||
|
semipos_list = pair.get("pair_pos_semipos_sate_img_list", [])
|
||||||
|
semipos_weights = pair.get("pair_pos_semipos_sate_weight_list", [])
|
||||||
|
|
||||||
|
# Use positives if available, else semi-positives.
|
||||||
|
if pos_list:
|
||||||
|
sat_candidates = pos_list
|
||||||
|
sat_weights = None
|
||||||
|
elif semipos_list:
|
||||||
|
sat_candidates = semipos_list
|
||||||
|
sat_weights = semipos_weights if semipos_weights else None
|
||||||
|
else:
|
||||||
|
continue # No match, skip.
|
||||||
|
|
||||||
|
# Get captions.
|
||||||
|
cap_data = self.caption_index.get(drone_name)
|
||||||
|
if cap_data is not None:
|
||||||
|
l1, l2, l3 = _parse_caption_levels(cap_data["output"])
|
||||||
|
else:
|
||||||
|
l1 = l2 = l3 = _EMPTY_CAPTION
|
||||||
|
|
||||||
|
self.entries.append({
|
||||||
|
"drone_name": drone_name,
|
||||||
|
"drone_dir": pair["drone_img_dir"],
|
||||||
|
"sat_dir": pair["sate_img_dir"],
|
||||||
|
"sat_candidates": sat_candidates,
|
||||||
|
"sat_weights": sat_weights,
|
||||||
|
"caption_l1": l1,
|
||||||
|
"caption_l2": l2,
|
||||||
|
"caption_l3": l3,
|
||||||
|
})
|
||||||
|
|
||||||
|
def _load_image(self, directory: str, filename: str) -> torch.Tensor:
|
||||||
|
path = self.rgb_root / directory / filename
|
||||||
|
with Image.open(path) as img:
|
||||||
|
rgb = img.convert("RGB")
|
||||||
|
if self.image_transform is not None:
|
||||||
|
return self.image_transform(rgb)
|
||||||
|
return torch.tensor(0) # placeholder if no transform
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return len(self.entries)
|
||||||
|
|
||||||
|
def __getitem__(self, idx: int) -> dict[str, Any]:
|
||||||
|
entry = self.entries[idx]
|
||||||
|
|
||||||
|
drone_img = self._load_image(entry["drone_dir"], entry["drone_name"])
|
||||||
|
|
||||||
|
# Sample satellite match (weighted if semi-positive).
|
||||||
|
if entry["sat_weights"] is not None:
|
||||||
|
sat_name = self._rng.choices(
|
||||||
|
entry["sat_candidates"],
|
||||||
|
weights=entry["sat_weights"],
|
||||||
|
k=1,
|
||||||
|
)[0]
|
||||||
|
else:
|
||||||
|
sat_name = self._rng.choice(entry["sat_candidates"])
|
||||||
|
|
||||||
|
sat_img = self._load_image(entry["sat_dir"], sat_name)
|
||||||
|
|
||||||
|
# Captions with optional dropout.
|
||||||
|
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
|
||||||
|
l1 = l2 = l3 = _EMPTY_CAPTION
|
||||||
|
else:
|
||||||
|
l1, l2, l3 = entry["caption_l1"], entry["caption_l2"], entry["caption_l3"]
|
||||||
|
|
||||||
|
return {
|
||||||
|
"drone_img": drone_img,
|
||||||
|
"sat_img": sat_img,
|
||||||
|
"caption_l1": l1,
|
||||||
|
"caption_l2": l2,
|
||||||
|
"caption_l3": l3,
|
||||||
|
"pair_id": entry["drone_name"],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def collate_gtauav_batch(
|
||||||
|
batch: list[dict[str, Any]],
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""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_l1": [b["caption_l1"] for b in batch],
|
||||||
|
"caption_l2": [b["caption_l2"] for b in batch],
|
||||||
|
"caption_l3": [b["caption_l3"] for b in batch],
|
||||||
|
"pair_ids": [b["pair_id"] for b in batch],
|
||||||
|
}
|
||||||
@@ -1,8 +1,7 @@
|
|||||||
"""Loss functions for caption quality test."""
|
"""Loss functions for caption quality test."""
|
||||||
from src.losses.multi_infonce import (
|
from src.losses.multi_infonce import (
|
||||||
MultiTermInfoNCE,
|
InfoNCELoss,
|
||||||
cosine_temperature,
|
cosine_temperature,
|
||||||
curriculum_lambdas,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = ["MultiTermInfoNCE", "cosine_temperature", "curriculum_lambdas"]
|
__all__ = ["InfoNCELoss", "cosine_temperature"]
|
||||||
|
|||||||
557
src/models/asymmetric_encoder.py
Normal file
557
src/models/asymmetric_encoder.py
Normal file
@@ -0,0 +1,557 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
"""Asymmetric dual encoder for CVGL caption test on GTA-UAV.
|
||||||
|
|
||||||
|
Architecture:
|
||||||
|
Query: DINOv3 ViT-L/16 (LVD, frozen) + LRSCLIP text (L1/L2/L3) -> GatedFusion -> query
|
||||||
|
Gallery: DINOv3 ViT-L/16 (SAT, frozen) -> gallery
|
||||||
|
Loss: InfoNCE(query, gallery)
|
||||||
|
|
||||||
|
DINOv3 checkpoints use a custom key layout (not HuggingFace transformers).
|
||||||
|
LRSCLIP (DGTRS-CLIP ViT-L-14) uses open_clip layout with KPS positional embeddings.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import math
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import coloredlogs
|
||||||
|
import open_clip
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
LOGGER = logging.getLogger("caption_test.model")
|
||||||
|
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||||
|
from safetensors.torch import load_file as load_safetensors
|
||||||
|
|
||||||
|
from src.models.dual_encoder import GatedFusion, ProjectionHead
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# DINOv3 ViT-L/16 — minimal implementation matching checkpoint key layout
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class DINOv3Attention(nn.Module):
|
||||||
|
"""Multi-head self-attention with separate Q/K/V projections."""
|
||||||
|
|
||||||
|
def __init__(self, dim: int = 1024, num_heads: int = 16) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.num_heads = num_heads
|
||||||
|
self.head_dim = dim // num_heads
|
||||||
|
self.scale = self.head_dim ** -0.5
|
||||||
|
self.q_proj = nn.Linear(dim, dim)
|
||||||
|
self.k_proj = nn.Linear(dim, dim, bias=False)
|
||||||
|
self.v_proj = nn.Linear(dim, dim)
|
||||||
|
self.o_proj = nn.Linear(dim, dim)
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
B, N, C = x.shape
|
||||||
|
q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
||||||
|
k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
||||||
|
v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
||||||
|
attn = F.scaled_dot_product_attention(q, k, v)
|
||||||
|
x = attn.permute(0, 2, 1, 3).reshape(B, N, C)
|
||||||
|
return self.o_proj(x)
|
||||||
|
|
||||||
|
|
||||||
|
class DINOv3LayerScale(nn.Module):
|
||||||
|
"""Per-channel learnable scale (lambda)."""
|
||||||
|
|
||||||
|
def __init__(self, dim: int) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.lambda1 = nn.Parameter(torch.ones(dim))
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
return x * self.lambda1
|
||||||
|
|
||||||
|
|
||||||
|
class DINOv3MLP(nn.Module):
|
||||||
|
"""SwiGLU-like MLP: up_proj + GELU + down_proj."""
|
||||||
|
|
||||||
|
def __init__(self, dim: int = 1024, mlp_dim: int = 4096) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.up_proj = nn.Linear(dim, mlp_dim)
|
||||||
|
self.down_proj = nn.Linear(mlp_dim, dim)
|
||||||
|
self.act = nn.GELU()
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
return self.down_proj(self.act(self.up_proj(x)))
|
||||||
|
|
||||||
|
|
||||||
|
class DINOv3Block(nn.Module):
|
||||||
|
"""Single DINOv3 transformer block."""
|
||||||
|
|
||||||
|
def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_dim: int = 4096) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.norm1 = nn.LayerNorm(dim)
|
||||||
|
self.attention = DINOv3Attention(dim, num_heads)
|
||||||
|
self.layer_scale1 = DINOv3LayerScale(dim)
|
||||||
|
self.norm2 = nn.LayerNorm(dim)
|
||||||
|
self.mlp = DINOv3MLP(dim, mlp_dim)
|
||||||
|
self.layer_scale2 = DINOv3LayerScale(dim)
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
x = x + self.layer_scale1(self.attention(self.norm1(x)))
|
||||||
|
x = x + self.layer_scale2(self.mlp(self.norm2(x)))
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class DINOv3Embeddings(nn.Module):
|
||||||
|
"""Patch embedding + CLS token + register tokens."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int = 1024,
|
||||||
|
patch_size: int = 16,
|
||||||
|
num_registers: int = 4,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.patch_embeddings = nn.Conv2d(3, dim, patch_size, patch_size)
|
||||||
|
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
|
||||||
|
self.register_tokens = nn.Parameter(torch.zeros(1, num_registers, dim))
|
||||||
|
self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
B = x.shape[0]
|
||||||
|
patches = self.patch_embeddings(x).flatten(2).transpose(1, 2) # [B, N, D]
|
||||||
|
N = patches.shape[1]
|
||||||
|
|
||||||
|
cls = self.cls_token.expand(B, -1, -1)
|
||||||
|
reg = self.register_tokens.expand(B, -1, -1)
|
||||||
|
|
||||||
|
# DINOv3: [CLS, registers, patches]
|
||||||
|
x = torch.cat([cls, reg, patches], dim=1)
|
||||||
|
|
||||||
|
# Positional embedding: interpolated sincos (RoPE applied in attention
|
||||||
|
# in original, but pretrained checkpoints bake it into weights).
|
||||||
|
# We use a simple learned-style pos embed computed on the fly.
|
||||||
|
pos = self._get_pos_embed(N, x.device, x.dtype)
|
||||||
|
# pos covers patches only, skip CLS + registers
|
||||||
|
x[:, 1 + reg.shape[1]:] = x[:, 1 + reg.shape[1]:] + pos
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
def _get_pos_embed(self, n_patches: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
||||||
|
# DINOv3 uses RoPE internally — no additive pos embed needed.
|
||||||
|
# Return zeros as placeholder (weights handle positioning via RoPE).
|
||||||
|
return torch.zeros(1, n_patches, self.cls_token.shape[-1], device=device, dtype=dtype)
|
||||||
|
|
||||||
|
|
||||||
|
class DINOv3ViT(nn.Module):
|
||||||
|
"""DINOv3 ViT-L/16 matching the checkpoint key layout.
|
||||||
|
|
||||||
|
Checkpoint keys:
|
||||||
|
embeddings.cls_token, embeddings.patch_embeddings.{weight,bias},
|
||||||
|
embeddings.register_tokens, embeddings.mask_token,
|
||||||
|
layer.{i}.attention.{q,k,v,o}_proj.{weight,bias},
|
||||||
|
layer.{i}.layer_scale{1,2}.lambda1,
|
||||||
|
layer.{i}.mlp.{up,down}_proj.{weight,bias},
|
||||||
|
layer.{i}.norm{1,2}.{weight,bias},
|
||||||
|
norm.{weight,bias}
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int = 1024,
|
||||||
|
num_heads: int = 16,
|
||||||
|
mlp_dim: int = 4096,
|
||||||
|
num_layers: int = 24,
|
||||||
|
patch_size: int = 16,
|
||||||
|
num_registers: int = 4,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.embeddings = DINOv3Embeddings(dim, patch_size, num_registers)
|
||||||
|
self.layer = nn.ModuleList([
|
||||||
|
DINOv3Block(dim, num_heads, mlp_dim) for _ in range(num_layers)
|
||||||
|
])
|
||||||
|
self.norm = nn.LayerNorm(dim)
|
||||||
|
self.embed_dim = dim
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Forward pass. Returns CLS token embedding [B, dim]."""
|
||||||
|
x = self.embeddings(x)
|
||||||
|
for block in self.layer:
|
||||||
|
x = block(x)
|
||||||
|
x = self.norm(x)
|
||||||
|
return x[:, 0] # CLS token
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
|
||||||
|
"""Load from .pth or .safetensors checkpoint."""
|
||||||
|
model = cls()
|
||||||
|
path = Path(path)
|
||||||
|
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
|
||||||
|
if path.suffix == ".safetensors":
|
||||||
|
state = load_safetensors(str(path))
|
||||||
|
else:
|
||||||
|
state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||||
|
if "model" in state:
|
||||||
|
state = state["model"]
|
||||||
|
elif "state_dict" in state:
|
||||||
|
state = state["state_dict"]
|
||||||
|
model.load_state_dict(state, strict=False)
|
||||||
|
n_params = sum(p.numel() for p in model.parameters())
|
||||||
|
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# LRSCLIP (DGTRS-CLIP) text encoder — open_clip with KPS positional embedding
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class LRSCLIPTextEncoder(nn.Module):
|
||||||
|
"""DGTRS-CLIP text encoder with KPS positional embedding (248 tokens).
|
||||||
|
|
||||||
|
Wraps open_clip ViT-L-14 text tower. Handles the extra
|
||||||
|
positional_embedding_res (KPS residual) not present in vanilla open_clip.
|
||||||
|
|
||||||
|
Checkpoint keys (text-only):
|
||||||
|
token_embedding.weight: [49408, 768]
|
||||||
|
positional_embedding: [248, 768]
|
||||||
|
positional_embedding_res: [248, 768]
|
||||||
|
transformer.resblocks.{0-11}.{attn,mlp,ln}.*
|
||||||
|
ln_final.{weight,bias}
|
||||||
|
text_projection: [768, 768]
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, context_length: int = 248, embed_dim: int = 768) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.context_length = context_length
|
||||||
|
self.embed_dim = embed_dim
|
||||||
|
|
||||||
|
# Build the open_clip model to get correct architecture, then
|
||||||
|
# replace positional embedding and add KPS residual.
|
||||||
|
clip_model = open_clip.create_model("ViT-L-14", pretrained=None)
|
||||||
|
|
||||||
|
# Extract text components.
|
||||||
|
self.token_embedding = clip_model.token_embedding
|
||||||
|
self.transformer = clip_model.transformer
|
||||||
|
self.ln_final = clip_model.ln_final
|
||||||
|
self.text_projection = clip_model.text_projection
|
||||||
|
self.attn_mask = None # rebuilt in forward
|
||||||
|
|
||||||
|
# Replace positional embedding with 248-length version.
|
||||||
|
self.positional_embedding = nn.Parameter(
|
||||||
|
torch.zeros(context_length, embed_dim),
|
||||||
|
)
|
||||||
|
self.positional_embedding_res = nn.Parameter(
|
||||||
|
torch.zeros(context_length, embed_dim),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _build_attn_mask(self, context_length: int, device: torch.device) -> torch.Tensor:
|
||||||
|
mask = torch.empty(context_length, context_length, device=device)
|
||||||
|
mask.fill_(float("-inf"))
|
||||||
|
mask.triu_(1)
|
||||||
|
return mask
|
||||||
|
|
||||||
|
def forward(self, text: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Encode tokenized text.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Token IDs [B, T] (T <= 248).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Text embeddings [B, 768], L2-normalized.
|
||||||
|
"""
|
||||||
|
T = text.shape[1]
|
||||||
|
x = self.token_embedding(text) # [B, T, 768]
|
||||||
|
pos = self.positional_embedding[:T] + self.positional_embedding_res[:T]
|
||||||
|
x = x + pos
|
||||||
|
|
||||||
|
if self.attn_mask is None or self.attn_mask.shape[0] != T:
|
||||||
|
self.attn_mask = self._build_attn_mask(T, x.device)
|
||||||
|
attn_mask = self.attn_mask[:T, :T]
|
||||||
|
|
||||||
|
# open_clip transformer expects batch-first [B, T, D].
|
||||||
|
x = self.transformer(x, attn_mask=attn_mask)
|
||||||
|
x = self.ln_final(x) # [B, T, D]
|
||||||
|
|
||||||
|
# Take features at EOS token position.
|
||||||
|
eos_idx = text.argmax(dim=-1)
|
||||||
|
x = x[torch.arange(x.shape[0], device=x.device), eos_idx]
|
||||||
|
|
||||||
|
# Project.
|
||||||
|
if isinstance(self.text_projection, nn.Parameter):
|
||||||
|
x = x @ self.text_projection
|
||||||
|
else:
|
||||||
|
x = self.text_projection(x)
|
||||||
|
|
||||||
|
return F.normalize(x, dim=-1)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_pretrained(cls, path: str | Path) -> LRSCLIPTextEncoder:
|
||||||
|
"""Load DGTRS-CLIP checkpoint, extracting only text encoder weights."""
|
||||||
|
model = cls()
|
||||||
|
path = Path(path)
|
||||||
|
LOGGER.info("📝 Loading LRSCLIP text encoder from %s", path.name)
|
||||||
|
full_state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||||
|
if "state_dict" in full_state:
|
||||||
|
full_state = full_state["state_dict"]
|
||||||
|
|
||||||
|
# Filter out visual.* keys — keep only text encoder.
|
||||||
|
text_state = {
|
||||||
|
k: v for k, v in full_state.items()
|
||||||
|
if not k.startswith("visual.")
|
||||||
|
}
|
||||||
|
|
||||||
|
# Handle text_projection shape: checkpoint may be [768, 768] Parameter
|
||||||
|
# while open_clip stores it as nn.Parameter directly.
|
||||||
|
model.load_state_dict(text_state, strict=False)
|
||||||
|
n_params = sum(p.numel() for p in model.parameters())
|
||||||
|
LOGGER.info("📝 LRSCLIP loaded: %s params, context=%d tokens", f"{n_params:,}", model.context_length)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Text fusion MLP: concat L1/L2/L3 -> project to D
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TextFusionMLP(nn.Module):
|
||||||
|
"""Fuse L1/L2/L3 text embeddings via concat + MLP.
|
||||||
|
|
||||||
|
[B, 3*text_dim] -> [B, proj_dim]
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
text_dim: int = 768,
|
||||||
|
hidden_dim: int = 768,
|
||||||
|
proj_dim: int = 512,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.mlp = nn.Sequential(
|
||||||
|
nn.Linear(3 * text_dim, hidden_dim),
|
||||||
|
nn.GELU(),
|
||||||
|
nn.Linear(hidden_dim, proj_dim),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
z_l1: torch.Tensor,
|
||||||
|
z_l2: torch.Tensor,
|
||||||
|
z_l3: torch.Tensor,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Fuse three text embeddings.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
z_l1: L1 overview [B, text_dim].
|
||||||
|
z_l2: L2 full description [B, text_dim].
|
||||||
|
z_l3: L3 fingerprint [B, text_dim].
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Fused text embedding [B, proj_dim].
|
||||||
|
"""
|
||||||
|
cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
|
||||||
|
return self.mlp(cat)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Main model: AsymmetricEncoder
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class AsymmetricEncoder(nn.Module):
|
||||||
|
"""Asymmetric dual encoder for CVGL with text fusion.
|
||||||
|
|
||||||
|
Query branch: DINOv3 LVD (drone) + LRSCLIP (L1/L2/L3) -> GatedFusion -> query
|
||||||
|
Gallery branch: DINOv3 SAT (satellite) -> gallery
|
||||||
|
|
||||||
|
Args:
|
||||||
|
dino_web_path: Path to DINOv3 LVD checkpoint (drone encoder).
|
||||||
|
dino_sat_path: Path to DINOv3 SAT checkpoint (satellite encoder).
|
||||||
|
lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
|
||||||
|
proj_dim: Shared projection dimension.
|
||||||
|
init_gate: Initial fusion gate (image weight).
|
||||||
|
baseline_mode: If True, gate = 1.0 (text ignored).
|
||||||
|
device: Torch device string.
|
||||||
|
"""
|
||||||
|
|
||||||
|
DINO_DIM = 1024
|
||||||
|
TEXT_DIM = 768
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
|
||||||
|
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
|
||||||
|
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
|
||||||
|
proj_dim: int = 512,
|
||||||
|
init_gate: float = 0.7,
|
||||||
|
baseline_mode: bool = False,
|
||||||
|
device: str = "cuda",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.proj_dim = proj_dim
|
||||||
|
self.baseline_mode = baseline_mode
|
||||||
|
self.device = device
|
||||||
|
|
||||||
|
# Image encoders (frozen).
|
||||||
|
self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
|
||||||
|
self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
|
||||||
|
self._freeze(self.drone_encoder)
|
||||||
|
self._freeze(self.sat_encoder)
|
||||||
|
|
||||||
|
# Text encoder (partial unfreeze).
|
||||||
|
if not baseline_mode:
|
||||||
|
self.text_encoder = LRSCLIPTextEncoder.from_pretrained(lrsclip_path)
|
||||||
|
self._freeze(self.text_encoder)
|
||||||
|
self._unfreeze_text_last_block()
|
||||||
|
self.tokenizer = open_clip.get_tokenizer("ViT-L-14")
|
||||||
|
else:
|
||||||
|
self.text_encoder = None
|
||||||
|
self.tokenizer = None
|
||||||
|
|
||||||
|
# Projection heads.
|
||||||
|
self.proj_drone = ProjectionHead(
|
||||||
|
in_dim=self.DINO_DIM, out_dim=proj_dim, use_mlp=False,
|
||||||
|
)
|
||||||
|
self.proj_sat = ProjectionHead(
|
||||||
|
in_dim=self.DINO_DIM, out_dim=proj_dim, use_mlp=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Text fusion (L1/L2/L3 -> proj_dim).
|
||||||
|
if not baseline_mode:
|
||||||
|
self.text_fusion = TextFusionMLP(
|
||||||
|
text_dim=self.TEXT_DIM,
|
||||||
|
hidden_dim=self.TEXT_DIM,
|
||||||
|
proj_dim=proj_dim,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Gated fusion.
|
||||||
|
self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _freeze(module: nn.Module) -> None:
|
||||||
|
for p in module.parameters():
|
||||||
|
p.requires_grad = False
|
||||||
|
module.eval()
|
||||||
|
|
||||||
|
def _unfreeze_text_last_block(self) -> None:
|
||||||
|
"""Unfreeze last transformer block + text_projection + ln_final."""
|
||||||
|
if self.text_encoder is None:
|
||||||
|
return
|
||||||
|
# Last resblock.
|
||||||
|
for p in self.text_encoder.transformer.resblocks[-1].parameters():
|
||||||
|
p.requires_grad = True
|
||||||
|
# ln_final.
|
||||||
|
for p in self.text_encoder.ln_final.parameters():
|
||||||
|
p.requires_grad = True
|
||||||
|
# text_projection.
|
||||||
|
tp = self.text_encoder.text_projection
|
||||||
|
if isinstance(tp, nn.Parameter):
|
||||||
|
tp.requires_grad = True
|
||||||
|
elif isinstance(tp, nn.Module):
|
||||||
|
for p in tp.parameters():
|
||||||
|
p.requires_grad = True
|
||||||
|
|
||||||
|
def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Encode drone images. Returns [B, DINO_DIM]."""
|
||||||
|
with torch.no_grad():
|
||||||
|
return self.drone_encoder(images)
|
||||||
|
|
||||||
|
def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Encode satellite images. Returns [B, DINO_DIM]."""
|
||||||
|
with torch.no_grad():
|
||||||
|
return self.sat_encoder(images)
|
||||||
|
|
||||||
|
def encode_text_levels(
|
||||||
|
self,
|
||||||
|
l1_texts: list[str],
|
||||||
|
l2_texts: list[str],
|
||||||
|
l3_texts: list[str],
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Encode L1/L2/L3 captions and fuse. Returns [B, proj_dim]."""
|
||||||
|
z_l1 = self._encode_single_text(l1_texts)
|
||||||
|
z_l2 = self._encode_single_text(l2_texts)
|
||||||
|
z_l3 = self._encode_single_text(l3_texts)
|
||||||
|
fused = self.text_fusion(z_l1, z_l2, z_l3)
|
||||||
|
return F.normalize(fused, dim=-1)
|
||||||
|
|
||||||
|
def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
|
||||||
|
"""Tokenize and encode a list of strings."""
|
||||||
|
tokens = self.tokenizer(list(texts)).to(self.device)
|
||||||
|
# Pad/truncate to context_length.
|
||||||
|
T = tokens.shape[1]
|
||||||
|
if T > self.text_encoder.context_length:
|
||||||
|
tokens = tokens[:, :self.text_encoder.context_length]
|
||||||
|
return self.text_encoder(tokens)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
drone_img: torch.Tensor,
|
||||||
|
sat_img: torch.Tensor,
|
||||||
|
caption_l1: list[str] | None = None,
|
||||||
|
caption_l2: list[str] | None = None,
|
||||||
|
caption_l3: list[str] | None = None,
|
||||||
|
) -> dict[str, torch.Tensor]:
|
||||||
|
"""Forward pass.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
drone_img: Drone images [B, 3, 256, 256].
|
||||||
|
sat_img: Satellite images [B, 3, 256, 256].
|
||||||
|
caption_l1: L1 overview captions.
|
||||||
|
caption_l2: L2 full description captions.
|
||||||
|
caption_l3: L3 fingerprint captions.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dict with 'query' [B, proj_dim], 'gallery' [B, proj_dim], 'gate'.
|
||||||
|
"""
|
||||||
|
# Gallery: satellite only.
|
||||||
|
sat_feat = self.encode_satellite(sat_img)
|
||||||
|
gallery = self.proj_sat(sat_feat)
|
||||||
|
|
||||||
|
# Query: drone + optional text.
|
||||||
|
drone_feat = self.encode_drone(drone_img)
|
||||||
|
drone_proj = self.proj_drone(drone_feat)
|
||||||
|
|
||||||
|
text_proj = None
|
||||||
|
has_text = (
|
||||||
|
caption_l1 is not None
|
||||||
|
and caption_l2 is not None
|
||||||
|
and caption_l3 is not None
|
||||||
|
and not self.baseline_mode
|
||||||
|
)
|
||||||
|
if has_text:
|
||||||
|
text_proj = self.encode_text_levels(caption_l1, caption_l2, caption_l3)
|
||||||
|
|
||||||
|
query = self.fusion(drone_proj, text_proj)
|
||||||
|
# Re-normalize after fusion.
|
||||||
|
query = F.normalize(query, dim=-1)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"query": query,
|
||||||
|
"gallery": gallery,
|
||||||
|
"gate": self.fusion.gate_value,
|
||||||
|
}
|
||||||
|
|
||||||
|
def trainable_parameters(self) -> list[nn.Parameter]:
|
||||||
|
"""Return list of parameters that require grad."""
|
||||||
|
return [p for p in self.parameters() if p.requires_grad]
|
||||||
|
|
||||||
|
def train(self, mode: bool = True) -> AsymmetricEncoder:
|
||||||
|
"""Override to keep frozen encoders in eval mode."""
|
||||||
|
super().train(mode)
|
||||||
|
self.drone_encoder.eval()
|
||||||
|
self.sat_encoder.eval()
|
||||||
|
if self.text_encoder is not None:
|
||||||
|
# Text encoder partially unfrozen — set to train mode
|
||||||
|
# but frozen layers won't update anyway.
|
||||||
|
self.text_encoder.train(mode)
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Image preprocessing (DINOv3: 256x256, ImageNet normalization)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
|
||||||
|
"""Build image transform for DINOv3 input."""
|
||||||
|
from torchvision import transforms
|
||||||
|
return transforms.Compose([
|
||||||
|
transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
|
||||||
|
transforms.CenterCrop(image_size),
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize(
|
||||||
|
mean=[0.485, 0.456, 0.406],
|
||||||
|
std=[0.229, 0.224, 0.225],
|
||||||
|
),
|
||||||
|
])
|
||||||
425
src/training/train_gtauav.py
Normal file
425
src/training/train_gtauav.py
Normal file
@@ -0,0 +1,425 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
"""Training loop for CVGL caption test on GTA-UAV-LR dataset.
|
||||||
|
|
||||||
|
Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion.
|
||||||
|
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import coloredlogs
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from torch.amp import GradScaler, autocast
|
||||||
|
from torch.optim import AdamW
|
||||||
|
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||||
|
from torch.utils.data import DataLoader
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
|
||||||
|
from src.losses.multi_infonce import InfoNCELoss
|
||||||
|
from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform
|
||||||
|
|
||||||
|
LOGGER = logging.getLogger("caption_test.train_gtauav")
|
||||||
|
|
||||||
|
# Default paths.
|
||||||
|
_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"
|
||||||
|
_CAPTION_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-captions"
|
||||||
|
_TRAIN_JSON = f"{_RGB_ROOT}/cross-area-drone2sate-train.json"
|
||||||
|
_TEST_JSON = f"{_RGB_ROOT}/cross-area-drone2sate-test.json"
|
||||||
|
|
||||||
|
_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth"
|
||||||
|
_DINO_SAT = "nn_models/DINO_SAT/model.safetensors"
|
||||||
|
_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class TrainConfigGTAUAV:
|
||||||
|
"""Training configuration for GTA-UAV experiment."""
|
||||||
|
|
||||||
|
# Data.
|
||||||
|
train_json: str = _TRAIN_JSON
|
||||||
|
test_json: str = _TEST_JSON
|
||||||
|
rgb_root: str = _RGB_ROOT
|
||||||
|
caption_root: str = _CAPTION_ROOT
|
||||||
|
filter_meta: str | None = None
|
||||||
|
|
||||||
|
# Model.
|
||||||
|
dino_web_path: str = _DINO_WEB
|
||||||
|
dino_sat_path: str = _DINO_SAT
|
||||||
|
lrsclip_path: str = _LRSCLIP
|
||||||
|
proj_dim: int = 512
|
||||||
|
init_gate: float = 0.7
|
||||||
|
baseline_mode: bool = False
|
||||||
|
|
||||||
|
# Training.
|
||||||
|
output_dir: str = "out/gtauav/with_text"
|
||||||
|
epochs: int = 10
|
||||||
|
batch_size: int = 64
|
||||||
|
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 = 2
|
||||||
|
seed: int = 42
|
||||||
|
device: str = "cuda"
|
||||||
|
|
||||||
|
# Loss.
|
||||||
|
tau_init: float = 0.1
|
||||||
|
tau_final: float = 0.01
|
||||||
|
label_smoothing: float = 0.1
|
||||||
|
weight_q2g: float = 0.6
|
||||||
|
weight_g2q: float = 0.4
|
||||||
|
|
||||||
|
|
||||||
|
def _set_seed(seed: int) -> None:
|
||||||
|
import random as _random
|
||||||
|
import numpy as _np
|
||||||
|
_random.seed(seed)
|
||||||
|
_np.random.seed(seed)
|
||||||
|
torch.manual_seed(seed)
|
||||||
|
torch.cuda.manual_seed_all(seed)
|
||||||
|
|
||||||
|
|
||||||
|
def _atomic_save(obj: dict, path: Path) -> None:
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def _evaluate(
|
||||||
|
model: AsymmetricEncoder,
|
||||||
|
loader: DataLoader,
|
||||||
|
device: str,
|
||||||
|
k_values: tuple[int, ...] = (1, 5, 10),
|
||||||
|
) -> dict[str, float]:
|
||||||
|
"""Compute R@K on validation set."""
|
||||||
|
model.eval()
|
||||||
|
all_query: list[torch.Tensor] = []
|
||||||
|
all_gallery: list[torch.Tensor] = []
|
||||||
|
|
||||||
|
for batch in tqdm(loader, desc=" 🔎 Evaluating", unit="batch", leave=False):
|
||||||
|
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||||
|
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
||||||
|
|
||||||
|
if model.baseline_mode:
|
||||||
|
embeddings = model(drone_img=drone_img, sat_img=sat_img)
|
||||||
|
else:
|
||||||
|
embeddings = model(
|
||||||
|
drone_img=drone_img,
|
||||||
|
sat_img=sat_img,
|
||||||
|
caption_l1=batch["caption_l1"],
|
||||||
|
caption_l2=batch["caption_l2"],
|
||||||
|
caption_l3=batch["caption_l3"],
|
||||||
|
)
|
||||||
|
all_query.append(embeddings["query"].cpu())
|
||||||
|
all_gallery.append(embeddings["gallery"].cpu())
|
||||||
|
|
||||||
|
query = torch.cat(all_query, dim=0)
|
||||||
|
gallery = torch.cat(all_gallery, dim=0)
|
||||||
|
|
||||||
|
sim = query @ gallery.t()
|
||||||
|
n = sim.size(0)
|
||||||
|
targets = torch.arange(n)
|
||||||
|
|
||||||
|
metrics: dict[str, float] = {}
|
||||||
|
sorted_idx = sim.argsort(dim=1, descending=True)
|
||||||
|
for k in k_values:
|
||||||
|
top_k = sorted_idx[:, :k]
|
||||||
|
hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
|
||||||
|
metrics[f"r@{k}_q2g"] = float(hit.mean().item())
|
||||||
|
|
||||||
|
sorted_idx_g2q = sim.t().argsort(dim=1, descending=True)
|
||||||
|
for k in k_values:
|
||||||
|
top_k = sorted_idx_g2q[:, :k]
|
||||||
|
hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
|
||||||
|
metrics[f"r@{k}_g2q"] = float(hit.mean().item())
|
||||||
|
|
||||||
|
metrics["gate"] = model.fusion.gate_value
|
||||||
|
return metrics
|
||||||
|
|
||||||
|
|
||||||
|
def train(cfg: TrainConfigGTAUAV) -> None:
|
||||||
|
"""Run full training loop."""
|
||||||
|
coloredlogs.install(
|
||||||
|
level="INFO",
|
||||||
|
logger=LOGGER,
|
||||||
|
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||||||
|
)
|
||||||
|
_set_seed(cfg.seed)
|
||||||
|
output_dir = Path(cfg.output_dir)
|
||||||
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# Save config.
|
||||||
|
with (output_dir / "config.json").open("w") as f:
|
||||||
|
json.dump(vars(cfg), f, indent=2)
|
||||||
|
|
||||||
|
# Model.
|
||||||
|
mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)"
|
||||||
|
LOGGER.info("🏗️ Building model — %s", mode_str)
|
||||||
|
model = AsymmetricEncoder(
|
||||||
|
dino_web_path=cfg.dino_web_path,
|
||||||
|
dino_sat_path=cfg.dino_sat_path,
|
||||||
|
lrsclip_path=cfg.lrsclip_path,
|
||||||
|
proj_dim=cfg.proj_dim,
|
||||||
|
init_gate=cfg.init_gate,
|
||||||
|
baseline_mode=cfg.baseline_mode,
|
||||||
|
device=cfg.device,
|
||||||
|
).to(cfg.device)
|
||||||
|
|
||||||
|
n_trainable = sum(p.numel() for p in model.trainable_parameters())
|
||||||
|
n_total = sum(p.numel() for p in model.parameters())
|
||||||
|
LOGGER.info(
|
||||||
|
"🧮 trainable=%s (%.2f%%) total=%s",
|
||||||
|
f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Loss.
|
||||||
|
loss_fn = InfoNCELoss(
|
||||||
|
temperature_init=cfg.tau_init,
|
||||||
|
temperature_final=cfg.tau_final,
|
||||||
|
label_smoothing=cfg.label_smoothing,
|
||||||
|
weight_q2g=cfg.weight_q2g,
|
||||||
|
weight_g2q=cfg.weight_g2q,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Data.
|
||||||
|
transform = get_dino_transform(image_size=256)
|
||||||
|
|
||||||
|
train_ds = GTAUAVDataset(
|
||||||
|
pair_json=cfg.train_json,
|
||||||
|
rgb_root=cfg.rgb_root,
|
||||||
|
caption_root=cfg.caption_root,
|
||||||
|
image_transform=transform,
|
||||||
|
filter_meta=cfg.filter_meta,
|
||||||
|
)
|
||||||
|
test_ds = GTAUAVDataset(
|
||||||
|
pair_json=cfg.test_json,
|
||||||
|
rgb_root=cfg.rgb_root,
|
||||||
|
caption_root=cfg.caption_root,
|
||||||
|
image_transform=transform,
|
||||||
|
filter_meta=cfg.filter_meta,
|
||||||
|
)
|
||||||
|
|
||||||
|
train_loader = DataLoader(
|
||||||
|
train_ds,
|
||||||
|
batch_size=cfg.batch_size,
|
||||||
|
shuffle=True,
|
||||||
|
num_workers=cfg.num_workers,
|
||||||
|
collate_fn=collate_gtauav_batch,
|
||||||
|
pin_memory=True,
|
||||||
|
drop_last=True,
|
||||||
|
)
|
||||||
|
test_loader = DataLoader(
|
||||||
|
test_ds,
|
||||||
|
batch_size=cfg.batch_size,
|
||||||
|
shuffle=False,
|
||||||
|
num_workers=cfg.num_workers,
|
||||||
|
collate_fn=collate_gtauav_batch,
|
||||||
|
pin_memory=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
LOGGER.info("📦 train=%d test=%d batch=%d", len(train_ds), len(test_ds), cfg.batch_size)
|
||||||
|
|
||||||
|
# Optimizer.
|
||||||
|
optimizer = AdamW(
|
||||||
|
model.trainable_parameters(),
|
||||||
|
lr=cfg.learning_rate,
|
||||||
|
weight_decay=cfg.weight_decay,
|
||||||
|
)
|
||||||
|
scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
|
||||||
|
scaler = GradScaler(enabled=cfg.use_amp)
|
||||||
|
|
||||||
|
history: list[dict] = []
|
||||||
|
|
||||||
|
LOGGER.info("🚀 Starting training for %d epochs", cfg.epochs)
|
||||||
|
|
||||||
|
for epoch in range(cfg.epochs):
|
||||||
|
model.train()
|
||||||
|
epoch_start = time.time()
|
||||||
|
agg: dict[str, float] = {}
|
||||||
|
n_batches = 0
|
||||||
|
|
||||||
|
pbar = tqdm(
|
||||||
|
train_loader,
|
||||||
|
desc=f" 🏋️ Epoch {epoch}/{cfg.epochs - 1}",
|
||||||
|
unit="batch",
|
||||||
|
leave=False,
|
||||||
|
)
|
||||||
|
for batch in pbar:
|
||||||
|
optimizer.zero_grad(set_to_none=True)
|
||||||
|
|
||||||
|
drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
|
||||||
|
sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
|
||||||
|
|
||||||
|
with autocast(device_type="cuda", enabled=cfg.use_amp):
|
||||||
|
if cfg.baseline_mode:
|
||||||
|
embeddings = model(drone_img=drone_img, sat_img=sat_img)
|
||||||
|
else:
|
||||||
|
embeddings = model(
|
||||||
|
drone_img=drone_img,
|
||||||
|
sat_img=sat_img,
|
||||||
|
caption_l1=batch["caption_l1"],
|
||||||
|
caption_l2=batch["caption_l2"],
|
||||||
|
caption_l3=batch["caption_l3"],
|
||||||
|
)
|
||||||
|
loss_dict = loss_fn(
|
||||||
|
embeddings=embeddings,
|
||||||
|
epoch=epoch,
|
||||||
|
total_epochs=cfg.epochs,
|
||||||
|
)
|
||||||
|
|
||||||
|
total_loss = loss_dict["total"]
|
||||||
|
scaler.scale(total_loss).backward()
|
||||||
|
|
||||||
|
if cfg.grad_clip > 0:
|
||||||
|
scaler.unscale_(optimizer)
|
||||||
|
nn.utils.clip_grad_norm_(
|
||||||
|
model.trainable_parameters(),
|
||||||
|
max_norm=cfg.grad_clip,
|
||||||
|
)
|
||||||
|
scaler.step(optimizer)
|
||||||
|
scaler.update()
|
||||||
|
|
||||||
|
for key, val in loss_dict.items():
|
||||||
|
agg[key] = agg.get(key, 0.0) + float(val.item())
|
||||||
|
n_batches += 1
|
||||||
|
|
||||||
|
pbar.set_postfix(
|
||||||
|
loss=f"{total_loss.item():.3f}",
|
||||||
|
gate=f"{loss_dict['gate'].item():.3f}",
|
||||||
|
)
|
||||||
|
|
||||||
|
scheduler.step()
|
||||||
|
elapsed = time.time() - epoch_start
|
||||||
|
|
||||||
|
means = {k: v / max(n_batches, 1) for k, v in agg.items()}
|
||||||
|
LOGGER.info(
|
||||||
|
"📈 epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate=%.4f",
|
||||||
|
epoch, elapsed,
|
||||||
|
optimizer.param_groups[0]["lr"],
|
||||||
|
means.get("total", 0.0),
|
||||||
|
means.get("temperature", 0.0),
|
||||||
|
means.get("gate", 1.0),
|
||||||
|
)
|
||||||
|
|
||||||
|
epoch_record: dict = {
|
||||||
|
"epoch": epoch,
|
||||||
|
"elapsed_seconds": elapsed,
|
||||||
|
"train": means,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Evaluation.
|
||||||
|
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
|
||||||
|
val_metrics = _evaluate(model, test_loader, cfg.device)
|
||||||
|
epoch_record["val"] = val_metrics
|
||||||
|
LOGGER.info(
|
||||||
|
"🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate=%.4f",
|
||||||
|
epoch,
|
||||||
|
val_metrics.get("r@1_q2g", 0.0),
|
||||||
|
val_metrics.get("r@5_q2g", 0.0),
|
||||||
|
val_metrics.get("r@10_q2g", 0.0),
|
||||||
|
val_metrics.get("gate", 1.0),
|
||||||
|
)
|
||||||
|
|
||||||
|
history.append(epoch_record)
|
||||||
|
|
||||||
|
# Save checkpoint.
|
||||||
|
_atomic_save(
|
||||||
|
obj={
|
||||||
|
"epoch": epoch,
|
||||||
|
"model_state": model.state_dict(),
|
||||||
|
"optimizer_state": optimizer.state_dict(),
|
||||||
|
},
|
||||||
|
path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
|
||||||
|
)
|
||||||
|
LOGGER.info("💾 Checkpoint saved: ckpt_epoch%03d.pt", epoch)
|
||||||
|
|
||||||
|
# Save history.
|
||||||
|
history_path = output_dir / "history.json"
|
||||||
|
with history_path.open("w", encoding="utf-8") as f:
|
||||||
|
json.dump(history, f, indent=2)
|
||||||
|
|
||||||
|
# Save final eval report.
|
||||||
|
LOGGER.info("🔎 Running final evaluation...")
|
||||||
|
final_metrics = _evaluate(model, test_loader, cfg.device)
|
||||||
|
report = {
|
||||||
|
"config": vars(cfg),
|
||||||
|
"metrics": final_metrics,
|
||||||
|
"history": history,
|
||||||
|
}
|
||||||
|
report_path = output_dir / "eval_report.json"
|
||||||
|
with report_path.open("w", encoding="utf-8") as f:
|
||||||
|
json.dump(report, f, indent=2)
|
||||||
|
|
||||||
|
LOGGER.info("✅ Training complete. Report: %s", report_path)
|
||||||
|
LOGGER.info(
|
||||||
|
"📊 Final — R@1=%.4f R@5=%.4f R@10=%.4f gate=%.4f",
|
||||||
|
final_metrics.get("r@1_q2g", 0.0),
|
||||||
|
final_metrics.get("r@5_q2g", 0.0),
|
||||||
|
final_metrics.get("r@10_q2g", 0.0),
|
||||||
|
final_metrics.get("gate", 1.0),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description="GTA-UAV caption test training.")
|
||||||
|
parser.add_argument(
|
||||||
|
"--baseline", action="store_true",
|
||||||
|
help="Run baseline mode (no text).",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--output-dir", type=str, default=None,
|
||||||
|
help="Override output directory.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--filter-meta", type=str, default=None,
|
||||||
|
help="Path to seg_filter.json for excluding bad images.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--batch-size", type=int, default=64,
|
||||||
|
help="Batch size.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--epochs", type=int, default=10,
|
||||||
|
help="Number of epochs.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--lr", type=float, default=1e-4,
|
||||||
|
help="Learning rate.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--init-gate", type=float, default=0.7,
|
||||||
|
help="Initial gate value (image weight).",
|
||||||
|
)
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
cfg = TrainConfigGTAUAV()
|
||||||
|
cfg.baseline_mode = args.baseline
|
||||||
|
cfg.batch_size = args.batch_size
|
||||||
|
cfg.epochs = args.epochs
|
||||||
|
cfg.learning_rate = args.lr
|
||||||
|
cfg.init_gate = args.init_gate
|
||||||
|
|
||||||
|
if args.filter_meta is not None:
|
||||||
|
cfg.filter_meta = args.filter_meta
|
||||||
|
|
||||||
|
if args.output_dir is not None:
|
||||||
|
cfg.output_dir = args.output_dir
|
||||||
|
elif args.baseline:
|
||||||
|
cfg.output_dir = "out/gtauav/baseline"
|
||||||
|
|
||||||
|
train(cfg)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user