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>
2026-04-21 17:54:27 +03:00
parent 5da791801c
commit 6ad9c4d149
10 changed files with 50043 additions and 101 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -17,7 +17,9 @@ env/
 *.ckpt
 *.onnx
 *.engine
+*.safetensors
 checkpoints/
+nn_models/
 out/
 outputs/
 runs/
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -1,29 +1,49 @@
 # Caption Quality Test for Cross-View Geo-Localization

-## Архитектура системы (v2, 2026-04-17)
+## Архитектура системы (v3, 2026-04-21) — GTA-UAV эксперимент

 ```
-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]
+QUERY BRANCH (drone + L1/L2/L3 captions):
+  drone_img --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024]
+                        |
+                 proj_drone (1024->512)
+                        |
+  L1 (overview)    --> LRSCLIP (248 tok) --> z_L1 [768] --\
+  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):
-  sat_img --> GeoRSCLIP ViT-B/32 (frozen) --> sat_feat [B,512]
-                     |
-              proj_gallery (Linear 512->512) --> L2-norm --> gallery [B,512]
+  sat_img --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024]
+                      |
+               proj_sat (1024->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)
+### 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 |
@@ -34,13 +54,42 @@ BASELINE: gate = 1.0 (text ignored)
 | `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
+### 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
- **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
+- **Text encoder:** 77 tokens, 512-dim — partial unfreeze

 ## GatedFusion

@@ -50,49 +99,91 @@ BASELINE: gate = 1.0 (text ignored)
 - `baseline_mode = True` → gate = 1.0, text полностью игнорируется
 - 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/`
 - **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)
+### V3 (GTA-UAV)
+Параметры: 10 epochs, batch 64, lr=1e-4, AMP, cosine LR, image 256x256.
+Train: 15,693 pairs (cross-area), Test: 18,015 pairs.
+Eval на cross-area split (primary).

+### V2 (UAV-GeoLoc, 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.
-
 ## Запуск

+### V3 (GTA-UAV)
 ```bash
-# 1. Baseline (no text)
-python -m src.training.train --config conf/baseline_no_text.gin
+# 1. Filter segmentation (exclude 90%+ background/water)
+python -m scripts.filter_segmentation --output meta/seg_filter.json

-# 2. With captions (primary test)
-python -m src.training.train --config conf/balanced.gin
+# 2. Baseline (no text)
+python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json

-# 3. Text-heavy (stress test)
-python -m src.training.train --config conf/text_heavy.gin
+# 3. With captions (L1/L2/L3)
+python -m src.training.train_gtauav --filter-meta meta/seg_filter.json

 # 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 \
    --baseline_report out/caption_test/baseline_no_text/eval_report.json \
    --full_report out/caption_test/balanced/eval_report.json \
@@ -101,7 +192,7 @@ python -m scripts.compare_runs \

 ## Метрики и Decision rule

-**Primary metric:** Delta R@1 (query -> gallery)
+**Primary metric:** Delta R@1 (drone -> satellite)

 | Delta R@1 | Verdict |
 |-----------|---------|
@@ -110,11 +201,21 @@ python -m scripts.compare_runs \
 | 0 to +1% | WEAK — redesign caption pipeline |
 | < 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
-**Eval metrics:** R@1, R@5, R@10 для query->gallery и gallery->query

 ## Бюджет времени (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 мин |
@@ -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
@@ -197,7 +334,7 @@ python scripts/prepare_dataset.py \
 - Размер на диске: 25 GB
 - Median distance drone->crop: 25.9m, P99: 45.7m
 - Память при генерации: до ~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)
 1. `train_db.txt`/`test_db.txt` содержали только matched кропы -> теперь все ~74K (полная gallery)
--- a/README.md
+++ b/README.md
@@ -1,42 +1,77 @@
 # Caption Quality Test for Cross-View Geo-Localization

 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.
+geo-localization (drone-to-satellite).

-## 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
-Gallery: sat_img -> proj -> gallery_emb
+Query:   drone_img (DINOv3 LVD) + L1/L2/L3 captions (LRSCLIP) -> GatedFusion -> query
+Gallery: sat_img (DINOv3 SAT) -> 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

 ```
-caption_test/
-├── conf/
-│   ├── 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)
+caption-test/
+├── conf/                         # Gin configs (v2)
+│   ├── balanced.gin
+│   ├── baseline_no_text.gin
+│   └── 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/
-│   ├── generate_captions.py  # Offline caption generation
-│   └── compare_runs.py       # Delta R@1 comparison report
+│   ├── filter_segmentation.py    # Meta-file: exclude 90%+ background/water
+│   ├── compare_runs.py           # Delta R@1 comparison report
+│   └── generate_captions.py      # Offline caption generation
 ├── src/
 │   ├── 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/
-│   │   └── dual_encoder.py   # GeoRSCLIP + GatedFusion + projection heads
+│   │   ├── asymmetric_encoder.py # DINOv3 + LRSCLIP + GatedFusion (v3)
+│   │   └── dual_encoder.py       # GeoRSCLIP + GatedFusion (v2)
 │   ├── losses/
-│   │   └── multi_infonce.py  # InfoNCE with cosine temperature
+│   │   └── multi_infonce.py      # InfoNCE with cosine temperature
 │   ├── training/
-│   │   └── train.py          # Main training loop
+│   │   ├── train_gtauav.py       # Training loop GTA-UAV (v3)
+│   │   └── train.py              # Training loop UAV-GeoLoc (v2)
 │   └── eval/
-│       └── evaluate.py       # R@K metrics, Delta R@1
-└── checkpoints/              # RS5M_ViT-B-32.pt (user-provided)
+│       └── evaluate.py           # R@K metrics, Delta R@1
+└── checkpoints/                  # GeoRSCLIP RS5M_ViT-B-32.pt (v2)
 ```

 ## Prerequisites
@@ -44,71 +79,54 @@ caption_test/
 ```
 torch>=2.0
 open_clip_torch
+safetensors
+timm
 gin-config
 Pillow
 numpy
 ```

-GeoRSCLIP checkpoint: download `RS5M_ViT-B-32.pt` from
-`github.com/om-ai-lab/RS5M` and place under `checkpoints/`.
+## Workflow (V3 — GTA-UAV)

-## Workflow
-
-### 1. Train baseline (no text)
+### 1. Filter segmentation (exclude uninformative images)

 ```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
-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
 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
+    --baseline_report out/gtauav/baseline/eval_report.json \
+    --full_report out/gtauav/with_text/eval_report.json \
+    --output out/gtauav/comparison.md
 ```

 ## 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 |
 | 0 to +1% | WEAK -- redesign caption pipeline |
 | < 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

 - `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
--- a/meta/seg_filter.json
+++ b/meta/seg_filter.json
--- a/scripts/filter_segmentation.py
+++ b/scripts/filter_segmentation.py
@@ -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()
--- a/src/datasets/init.py
+++ b/src/datasets/init.py
@@ -1,7 +1,7 @@
 """Dataset loaders for caption quality test."""
 from src.datasets.visloc_with_captions import (
-    VisLocCaptionDataset,
+    GeoLocCaptionDataset,
    collate_caption_batch,
 )

-__all__ = ["VisLocCaptionDataset", "collate_caption_batch"]
+__all__ = ["GeoLocCaptionDataset", "collate_caption_batch"]
--- a/src/datasets/gtauav_dataset.py
+++ b/src/datasets/gtauav_dataset.py
@@ -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],
+    }
--- a/src/losses/init.py
+++ b/src/losses/init.py
@@ -1,8 +1,7 @@
 """Loss functions for caption quality test."""
 from src.losses.multi_infonce import (
-    MultiTermInfoNCE,
+    InfoNCELoss,
    cosine_temperature,
-    curriculum_lambdas,
 )

-__all__ = ["MultiTermInfoNCE", "cosine_temperature", "curriculum_lambdas"]
+__all__ = ["InfoNCELoss", "cosine_temperature"]
--- a/src/models/asymmetric_encoder.py
+++ b/src/models/asymmetric_encoder.py
@@ -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],
+        ),
+    ])
--- a/src/training/train_gtauav.py
+++ b/src/training/train_gtauav.py
@@ -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()