README: architecture diagram with tensor dimensions, L1/L2/L3 text hierarchy description, text fusion formula, InfoNCE loss formula with learnable temperature, metrics table, optimizer/scheduler details with per-group LR, augmentation table, model parameter summary. CLAUDE.md: updated to DGTRS-CLIP (official architecture), loss formula, optimizer/scheduler details, text encoder architecture notes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
287 lines
11 KiB
Markdown
287 lines
11 KiB
Markdown
# 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).
|
||
|
||
## Architecture
|
||
|
||
### Overview
|
||
|
||
Asymmetric dual-encoder with domain-specific frozen backbones and hierarchical
|
||
text fusion for drone-to-satellite image retrieval.
|
||
|
||
```
|
||
┌──────────────────────────── QUERY BRANCH ────────────────────────────┐
|
||
│ │
|
||
│ drone_img ──► DINOv3 ViT-L/16 LVD ──► CLS token │
|
||
│ [B,3,256,256] (frozen, 303M) [B,1024] │
|
||
│ │ │
|
||
│ proj_drone: Linear(1024,512) │
|
||
│ │ │
|
||
│ d_img [B,512] │
|
||
│ │ │
|
||
│ L1 (overview) ──► DGTRS-CLIP ──► z₁ [B,768] ─┐ │
|
||
│ L2 (full desc) ──► DGTRS-CLIP ──► z₂ [B,768] ─┼─ cat ──► [B,2304]│
|
||
│ L3 (fingerprint) ──► DGTRS-CLIP ──► z₃ [B,768] ─┘ │ │
|
||
│ (248 tokens, KPS pos. emb.) MLP(2304→768→512) │
|
||
│ │ │
|
||
│ d_txt [B,512] │
|
||
│ │ │
|
||
│ q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion │
|
||
│ │ │
|
||
│ q̂ = q / ‖q‖₂ ──► query [B,512] │
|
||
└───────────────────────────────────────────────────────────────────────┘
|
||
|
||
┌──────────────────────────── GALLERY BRANCH ──────────────────────────┐
|
||
│ │
|
||
│ sat_img ──► DINOv3 ViT-L/16 SAT ──► CLS token │
|
||
│ [B,3,256,256] (frozen, 303M) [B,1024] │
|
||
│ │ │
|
||
│ proj_sat: Linear(1024,512) │
|
||
│ │ │
|
||
│ ĝ = g / ‖g‖₂ ──► gallery [B,512] │
|
||
└───────────────────────────────────────────────────────────────────────┘
|
||
|
||
BASELINE: σ(α) = 1.0 → q = d_img (text branch disabled, DGTRS not loaded)
|
||
```
|
||
|
||
### Text hierarchy (L1 / L2 / L3)
|
||
|
||
Each drone image has a VLM-generated caption (Qwen3-VL) split into 3 levels:
|
||
|
||
| Level | Name | Content | Typical length |
|
||
|-------|------|---------|----------------|
|
||
| **L1** | Overview | First sentence of P1: land-cover summary with class percentages | 15–30 tokens |
|
||
| **L2** | Full description | Complete P1 (inventory) + P2 (spatial layout with 5+ zones) | 100–200 tokens |
|
||
| **L3** | Fingerprint | P3: unique landmarks and spatial signature for matching | 20–50 tokens |
|
||
|
||
All three levels are encoded by a **single DGTRS-CLIP ViT-L-14** text encoder
|
||
(248-token context via KPS positional embedding, 768-dim output).
|
||
|
||
**Text fusion:**
|
||
|
||
```
|
||
z_text = MLP( [z₁ ; z₂ ; z₃] )
|
||
|
||
where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings
|
||
MLP: Linear(2304, 768) → GELU → Linear(768, 512)
|
||
z_text ∈ ℝ^(B×512)
|
||
```
|
||
|
||
**Gated fusion:**
|
||
|
||
```
|
||
q = σ(α) · d_img + (1 − σ(α)) · d_txt
|
||
|
||
where α — learnable scalar in logit-space (init: σ(α) ≈ 0.7)
|
||
σ — sigmoid function
|
||
d_img — projected drone image embedding [B, 512]
|
||
d_txt — fused text embedding [B, 512]
|
||
```
|
||
|
||
### Loss function
|
||
|
||
Symmetric InfoNCE with learnable temperature (CLIP-style `logit_scale`):
|
||
|
||
```
|
||
L = w_q2g · L_q→g + w_g2q · L_g→q
|
||
|
||
L_q→g = CrossEntropy( q̂ · ĝᵀ / τ, targets )
|
||
L_g→q = CrossEntropy( ĝ · q̂ᵀ / τ, targets )
|
||
|
||
where τ = 1 / exp(logit_scale), logit_scale — learnable scalar
|
||
τ ∈ [0.01, 0.5] (clamped)
|
||
τ_init = 0.07
|
||
w_q2g = 0.6, w_g2q = 0.4
|
||
targets = [0, 1, 2, ..., B−1] (positives on diagonal)
|
||
label_smoothing = 0.1
|
||
```
|
||
|
||
Loss is computed in **fp32** (outside AMP autocast) to prevent gradient overflow
|
||
in the learnable temperature.
|
||
|
||
### Metrics
|
||
|
||
| Metric | Formula | Direction |
|
||
|--------|---------|-----------|
|
||
| **R@K** (Recall at K) | fraction of queries where correct gallery is in top-K | drone → satellite (primary) |
|
||
| **Delta R@1** | R@1(with_text) − R@1(baseline) | higher = text helps |
|
||
|
||
Reported: R@1, R@5, R@10 for both q→g and g→q directions.
|
||
|
||
### Optimizer & scheduler
|
||
|
||
```
|
||
Optimizer: AdamW
|
||
- Projection heads (proj_drone, proj_sat, TextFusionMLP, gate α, logit_scale):
|
||
lr = 1e-4, weight_decay = 1e-4
|
||
- DGTRS text encoder (last resblock + ln_final + text_projection):
|
||
lr = 1e-5 (10× lower, --text-lr-factor 0.1)
|
||
|
||
Scheduler: Linear warmup (2 epochs) + cosine annealing
|
||
- Per-step (not per-epoch)
|
||
- warmup: lr linearly ramps from 0 to lr_max over warmup_steps
|
||
- cosine: lr decays from lr_max to 0 over remaining steps
|
||
|
||
Gradient clipping: max_norm = 1.0
|
||
Mixed precision: AMP fp16 for model forward, fp32 for loss
|
||
```
|
||
|
||
### Augmentations
|
||
|
||
| Transform | Drone (train) | Satellite (train) | Eval |
|
||
|-----------|:---:|:---:|:---:|
|
||
| RandomResizedCrop(256, scale=0.7–1.0) | ✓ | ✓ | — |
|
||
| Resize(256) + CenterCrop(256) | — | — | ✓ |
|
||
| RandomHorizontalFlip(0.5) | ✓ | ✓ | — |
|
||
| RandomRotation(15°) | ✓ | — | — |
|
||
| ColorJitter(0.3, 0.3, 0.2, 0.1) | ✓ | ✓ | — |
|
||
| RandomGrayscale(0.05) | ✓ | ✓ | — |
|
||
| GaussianBlur(k=3, σ=0.1–2.0) | ✓ | — | — |
|
||
| ImageNet Normalize | ✓ | ✓ | ✓ |
|
||
|
||
### Model summary
|
||
|
||
| Component | Params | Trainable | Notes |
|
||
|-----------|--------|-----------|-------|
|
||
| DINOv3 ViT-L/16 LVD (drone) | 303M | 0 | frozen |
|
||
| DINOv3 ViT-L/16 SAT (satellite) | 303M | 0 | frozen |
|
||
| DGTRS-CLIP ViT-L-14 (text) | 124M | ~7.6M | last block + ln_final + text_projection |
|
||
| proj_drone | 524K | 524K | Linear(1024, 512) |
|
||
| proj_sat | 524K | 524K | Linear(1024, 512) |
|
||
| TextFusionMLP | 2.2M | 2.2M | Linear(2304,768) + GELU + Linear(768,512) |
|
||
| GatedFusion α | 1 | 1 | scalar |
|
||
| logit_scale | 1 | 1 | learnable temperature |
|
||
| **Total** | **733M** | **10.9M (1.49%)** | |
|
||
|
||
## Experiments
|
||
|
||
### V3 — GTA-UAV + DINOv3 + DGTRS-CLIP (active)
|
||
|
||
**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)
|
||
- Split: 80/20 random (26,966 train / 6,742 test → 24,891/6,252 after seg filter)
|
||
|
||
### 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/ # Gin configs (v2)
|
||
│ ├── balanced.gin
|
||
│ ├── baseline_no_text.gin
|
||
│ └── text_heavy.gin
|
||
├── nn_models/ # Pre-trained checkpoints (v3, gitignored)
|
||
│ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M (.pth)
|
||
│ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M (.safetensors)
|
||
│ └── LRSCLIP/ # DGTRS-CLIP ViT-L-14 (.pt)
|
||
├── meta/ # Generated metadata
|
||
│ ├── train_80.json # 80% train split (26,966 pairs)
|
||
│ ├── test_20.json # 20% test split (6,742 pairs)
|
||
│ └── seg_filter.json # Segmentation filter results
|
||
├── scripts/
|
||
│ ├── make_split.py # Create 80/20 train/test split
|
||
│ ├── filter_segmentation.py # Exclude 90%+ background/water images
|
||
│ ├── compare_runs.py # Delta R@1 comparison report
|
||
│ └── generate_captions.py # Offline caption generation
|
||
├── src/
|
||
│ ├── datasets/
|
||
│ │ ├── gtauav_dataset.py # GTA-UAV-LR loader + L1/L2/L3 parsing (v3)
|
||
│ │ └── visloc_with_captions.py # UAV-GeoLoc loader (v2)
|
||
│ ├── models/
|
||
│ │ ├── dgtrs/ # Official DGTRS-CLIP text encoder (Apache-2.0)
|
||
│ │ │ ├── model.py # DGTRSTextEncoder, build_model, tokenize
|
||
│ │ │ ├── simple_tokenizer.py # BPE tokenizer (248 tokens)
|
||
│ │ │ └── bpe_simple_vocab_16e6.txt.gz
|
||
│ │ ├── asymmetric_encoder.py # DINOv3 + DGTRS + GatedFusion (v3)
|
||
│ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2)
|
||
│ ├── losses/
|
||
│ │ └── multi_infonce.py # InfoNCE with learnable temperature
|
||
│ ├── training/
|
||
│ │ ├── 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/ # GeoRSCLIP RS5M_ViT-B-32.pt (v2)
|
||
```
|
||
|
||
## Prerequisites
|
||
|
||
```
|
||
torch>=2.0
|
||
safetensors
|
||
coloredlogs
|
||
tqdm
|
||
ftfy
|
||
regex
|
||
gin-config
|
||
Pillow
|
||
numpy
|
||
```
|
||
|
||
## Workflow (V3 — GTA-UAV)
|
||
|
||
### 1. Create 80/20 split and filter segmentation
|
||
|
||
```bash
|
||
python -m scripts.make_split --output-dir meta
|
||
python -m scripts.filter_segmentation --output meta/seg_filter.json
|
||
```
|
||
|
||
### 2. Train baseline (no text)
|
||
|
||
```bash
|
||
python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json
|
||
```
|
||
|
||
### 3. Train with captions (L1/L2/L3)
|
||
|
||
```bash
|
||
python -m src.training.train_gtauav --filter-meta meta/seg_filter.json
|
||
```
|
||
|
||
### 4. Resume from checkpoint
|
||
|
||
```bash
|
||
python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt \
|
||
--filter-meta meta/seg_filter.json
|
||
```
|
||
|
||
### 5. Compare and get verdict
|
||
|
||
```bash
|
||
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
|
||
```
|
||
|
||
## Decision rule
|
||
|
||
| Delta R@1 (drone→satellite) | Verdict |
|
||
|---|---|
|
||
| >= +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 |
|
||
|
||
## Code style
|
||
|
||
- `from __future__ import annotations` everywhere
|
||
- Type hints on all signatures
|
||
- Google-style docstrings
|
||
- English-only comments
|