commit 3e95f4f618145d8d42cc757f2f9a10b7333bc693 Author: pikaliov Date: Tue Jul 7 16:22:49 2026 +0300 initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3952fb7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +outputs/ +results/ +cache/ +backtranslate.py +debug_augmentation.py +text_augmentation.py + +__pycache__/ +*.py[cod] +*.so +.ipynb_checkpoints/ + +.venv/ venv/ env/ + +cache/ .cache/ *.log + +results/ runs/ wandb/ checkpoints/ +*.ckpt *.pth *.pt + +datasets/ *.csv + +.env *.key + +.DS_Store .vscode/ .idea/ \ No newline at end of file diff --git a/collect_results.py b/collect_results.py new file mode 100644 index 0000000..cd5289e --- /dev/null +++ b/collect_results.py @@ -0,0 +1,286 @@ +"""Сборка результатов экспериментов: таблица + CSV + графики. + +Использование: + python collect_results.py # v1 + v2 + python collect_results.py v1 # только v1 + python collect_results.py v2 # только v2 +""" +from __future__ import annotations + +import csv +import json +import sys +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + + +SCRIPT_DIR = Path(__file__).resolve().parent +OUTPUTS_DIR = SCRIPT_DIR / "outputs" +RESULTS_DIR = SCRIPT_DIR / "results" + + +# --------------------------------------------------------------------------- +# Загрузка +# --------------------------------------------------------------------------- + +def load_experiment(exp_dir: Path) -> dict | None: + history_path = exp_dir / "history.json" + config_path = exp_dir / "config.json" + if not history_path.exists(): + return None + + with open(history_path) as f: + history = json.load(f) + config = {} + if config_path.exists(): + with open(config_path) as f: + config = json.load(f) + if not history: + return None + + best = max(history, key=lambda r: r.get("eval_recall@1", 0)) + latest = history[-1] + + return { + "dir": exp_dir.name, + "text_levels": " + ".join(config.get("text_levels", ["?"])), + "epochs_done": latest["epoch"], + "epochs_total": config.get("epochs", "?"), + "best_epoch": best["epoch"], + "best_R@1": best.get("eval_recall@1", 0), + "best_R@5": best.get("eval_recall@5", 0), + "best_R@10": best.get("eval_recall@10", 0), + "best_AP": best.get("eval_AP", 0), + "latest_loss": latest.get("train_loss", 0), + "latest_R@1": latest.get("eval_recall@1", 0), + "avg_epoch_time": sum(r.get("elapsed_s", 0) for r in history) / len(history), + "_history": history, + } + + +def collect_version(version: str) -> list[dict]: + """Собрать результаты одной версии (v1 или v2).""" + version_dir = OUTPUTS_DIR / version + if not version_dir.exists(): + print(f"⚠️ Папка не найдена: {version_dir}") + return [] + + results = [] + for exp_dir in sorted(version_dir.iterdir()): + if exp_dir.is_dir() and (exp_dir / "history.json").exists(): + data = load_experiment(exp_dir) + if data: + data["version"] = version + results.append(data) + return results + + +# --------------------------------------------------------------------------- +# Таблица в консоль +# --------------------------------------------------------------------------- + +def print_table(results: list[dict], version: str) -> None: + if not results: + print(f" {version}: нет данных") + return + + results.sort(key=lambda r: -r["best_R@1"]) + + header = ( + f"{'Levels':<24} {'Prog':<8} " + f"{'BestEp':>6} {'R@1':>7} {'R@5':>7} {'R@10':>7} " + f"{'AP':>7} {'Loss':>8} {'Time':>6}" + ) + sep = "─" * len(header) + + print(sep) + print(f" {version.upper()} Results") + print(sep) + print(header) + print(sep) + + for r in results: + prog = f"{r['epochs_done']}/{r['epochs_total']}" + print( + f"{r['text_levels']:<24} {prog:<8} " + f"{r['best_epoch']:>6} {r['best_R@1']:>7.4f} {r['best_R@5']:>7.4f} " + f"{r['best_R@10']:>7.4f} {r['best_AP']:>7.4f} " + f"{r['latest_loss']:>8.4f} {r['avg_epoch_time']:>5.0f}s" + ) + + print(sep) + print(f" {len(results)} experiments") + print() + + +# --------------------------------------------------------------------------- +# CSV +# --------------------------------------------------------------------------- + +def save_csv(results: list[dict], version: str) -> None: + fields = [ + "version", "text_levels", "epochs_done", "epochs_total", + "best_epoch", "best_R@1", "best_R@5", "best_R@10", "best_AP", + "latest_loss", "latest_R@1", "avg_epoch_time", "dir", + ] + path = RESULTS_DIR / f"results_{version}.csv" + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore") + writer.writeheader() + for r in results: + writer.writerow(r) + print(f" CSV: {path}") + + +# --------------------------------------------------------------------------- +# Графики +# --------------------------------------------------------------------------- + +def save_plots(results: list[dict], version: str) -> None: + if not results: + return + + colors = plt.cm.tab10(np.linspace(0, 1, max(len(results), 1))) + + # --- 1. Loss --- + fig, ax = plt.subplots(figsize=(8, 5)) + for r, c in zip(results, colors): + h = r["_history"] + ax.plot([e["epoch"] for e in h], + [e.get("train_loss", 0) for e in h], + label=r["text_levels"], color=c, linewidth=1.5) + ax.set_xlabel("Epoch") + ax.set_ylabel("Train Loss") + ax.set_title(f"{version.upper()} — Training Loss") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + plt.tight_layout() + path = RESULTS_DIR / f"loss_{version}.png" + plt.savefig(path, dpi=150) + plt.close() + print(f" PNG: {path}") + + # --- 2. Recall@1 --- + fig, ax = plt.subplots(figsize=(8, 5)) + for r, c in zip(results, colors): + h = r["_history"] + eps = [e["epoch"] for e in h if e.get("eval_recall@1") is not None] + r1s = [e["eval_recall@1"] for e in h if e.get("eval_recall@1") is not None] + if eps: + ax.plot(eps, r1s, label=r["text_levels"], color=c, + linewidth=1.5, marker=".", markersize=3) + ax.set_xlabel("Epoch") + ax.set_ylabel("Recall@1") + ax.set_title(f"{version.upper()} — Recall@1") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + plt.tight_layout() + path = RESULTS_DIR / f"recall1_{version}.png" + plt.savefig(path, dpi=150) + plt.close() + print(f" PNG: {path}") + + # --- 3. Recall@5 --- + fig, ax = plt.subplots(figsize=(8, 5)) + for r, c in zip(results, colors): + h = r["_history"] + eps = [e["epoch"] for e in h if e.get("eval_recall@5") is not None] + r5s = [e["eval_recall@5"] for e in h if e.get("eval_recall@5") is not None] + if eps: + ax.plot(eps, r5s, label=r["text_levels"], color=c, + linewidth=1.5, marker=".", markersize=3) + ax.set_xlabel("Epoch") + ax.set_ylabel("Recall@5") + ax.set_title(f"{version.upper()} — Recall@5") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + plt.tight_layout() + path = RESULTS_DIR / f"recall5_{version}.png" + plt.savefig(path, dpi=150) + plt.close() + print(f" PNG: {path}") + + # --- 4. Recall@10 --- + fig, ax = plt.subplots(figsize=(8, 5)) + for r, c in zip(results, colors): + h = r["_history"] + eps = [e["epoch"] for e in h if e.get("eval_recall@10") is not None] + r10s = [e["eval_recall@10"] for e in h if e.get("eval_recall@10") is not None] + if eps: + ax.plot(eps, r10s, label=r["text_levels"], color=c, + linewidth=1.5, marker=".", markersize=3) + ax.set_xlabel("Epoch") + ax.set_ylabel("Recall@10") + ax.set_title(f"{version.upper()} — Recall@10") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + plt.tight_layout() + path = RESULTS_DIR / f"recall10_{version}.png" + plt.savefig(path, dpi=150) + plt.close() + print(f" PNG: {path}") + + # --- 3. Bar chart лучших --- + fig, ax = plt.subplots(figsize=(8, 5)) + labels = [r["text_levels"] for r in results] + r1 = [r["best_R@1"] for r in results] + r5 = [r["best_R@5"] for r in results] + r10 = [r["best_R@10"] for r in results] + + x = np.arange(len(labels)) + w = 0.25 + ax.bar(x - w, r1, w, label="R@1", color="#4C78A8") + ax.bar(x, r5, w, label="R@5", color="#54A24B") + ax.bar(x + w, r10, w, label="R@10", color="#E45756") + ax.set_xticks(x) + ax.set_xticklabels(labels, fontsize=8, rotation=20, ha="right") + ax.set_ylabel("Score") + ax.set_title(f"{version.upper()} — Best Recall") + ax.legend() + ax.grid(True, alpha=0.3, axis="y") + for i, v in enumerate(r1): + ax.text(i - w, v + 0.005, f"{v:.3f}", ha="center", fontsize=7) + plt.tight_layout() + path = RESULTS_DIR / f"best_recall_{version}.png" + plt.savefig(path, dpi=150) + plt.close() + print(f" PNG: {path}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + # Парсинг: python collect_results.py [v1|v2] + if len(sys.argv) > 1 and sys.argv[1] in ("v1", "v2"): + versions = [sys.argv[1]] + else: + versions = ["v1", "v2"] + + if not OUTPUTS_DIR.exists(): + print(f"❌ Папка outputs не найдена: {OUTPUTS_DIR}") + return + + RESULTS_DIR.mkdir(exist_ok=True) + + for version in versions: + results = collect_version(version) + + # Таблица в консоль + print_table(results, version) + + if results: + # CSV и графики в results/ + save_csv(results, version) + save_plots(results, version) + + print(f"\n📁 Все результаты: {RESULTS_DIR}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/measure_truncation.py b/measure_truncation.py new file mode 100644 index 0000000..9882d26 --- /dev/null +++ b/measure_truncation.py @@ -0,0 +1,119 @@ +"""Измерение процента обрезаемых описаний по комбинациям уровней. + +Обрезка (truncation до 248 токенов) зависит ТОЛЬКО от текста и токенизатора, +не от обучения. Поэтому процент обрезанных сэмплов можно посчитать по готовым +данным — результат идентичен тому, что было во время обучения. + +Для каждой из 6 комбинаций уровней и каждого набора (v1/v2) считает: + - сколько сэмплов превышает 248 токенов (обрезается) + - средняя/максимальная длина в токенах + - средний «перебор» у обрезанных (на сколько токенов текст длиннее 248) + +Запуск: + python measure_truncation.py \ + --descriptions_path "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1" \ + --version v1 +""" +from __future__ import annotations + +import argparse +import sys + +sys.path.insert(0, ".") + +from src.data.gta_uav import load_text_descriptions, combine_text_levels +from src.models.dgtrs.model import tokenize_dgtrs + +CONTEXT_LENGTH = 248 + +# те же 6 комбинаций, что в экспериментах +COMBINATIONS = [ + ["level1"], + ["level2"], + ["level3"], + ["level1", "level2"], + ["level1", "level3"], + ["level1", "level2", "level3"], +] + + +def count_tokens(text: str) -> int: + """Реальное число ненулевых токенов ДО обрезки. + + tokenize с truncate=False даёт полную длину; если функция не поддерживает + truncate=False, считаем через увеличенный context_length. + """ + if not text.strip(): + return 0 + # токенизируем с большим запасом, чтобы увидеть полную длину без обрезки + toks = tokenize_dgtrs(text, context_length=1024, truncate=True) + return int((toks != 0).sum()) + + +def measure(descriptions: dict, combo: list[str]) -> dict: + """Посчитать статистику обрезки для одной комбинации уровней.""" + n_total = 0 + n_truncated = 0 + lengths = [] + overflows = [] + + for img_name, desc in descriptions.items(): + text = combine_text_levels(desc, combo) + if not text.strip(): + continue + n = count_tokens(text) + n_total += 1 + lengths.append(n) + if n > CONTEXT_LENGTH: + n_truncated += 1 + overflows.append(n - CONTEXT_LENGTH) + + pct = 100.0 * n_truncated / n_total if n_total else 0.0 + avg_len = sum(lengths) / len(lengths) if lengths else 0 + max_len = max(lengths) if lengths else 0 + avg_overflow = sum(overflows) / len(overflows) if overflows else 0 + + return { + "combo": " + ".join(combo), + "n_total": n_total, + "n_truncated": n_truncated, + "pct_truncated": pct, + "avg_len": avg_len, + "max_len": max_len, + "avg_overflow": avg_overflow, + } + + +def main(): + args = parse_args() + + descriptions = load_text_descriptions(args.descriptions_path, view_type="drone") + print(f"Загружено {len(descriptions)} описаний ({args.version})\n") + + print(f"{'Комбинация':<26} {'Всего':>7} {'Обрезано':>9} {'%':>7} " + f"{'СрДлина':>8} {'МаксДлина':>10} {'СрПеребор':>10}") + print("-" * 82) + + rows = [] + for combo in COMBINATIONS: + r = measure(descriptions, combo) + rows.append(r) + print(f"{r['combo']:<26} {r['n_total']:>7} {r['n_truncated']:>9} " + f"{r['pct_truncated']:>6.1f}% {r['avg_len']:>8.0f} " + f"{r['max_len']:>10} {r['avg_overflow']:>10.0f}") + + print("-" * 82) + print(f"\nЛимит контекста: {CONTEXT_LENGTH} токенов") + print("Обрезано = число сэмплов, где склеенный текст длиннее лимита " + "(хвост, включая level3-якорь, отсекается).") + + +def parse_args(): + p = argparse.ArgumentParser(description="Измерение обрезки по комбинациям уровней") + p.add_argument("--descriptions_path", type=str, required=True) + p.add_argument("--version", type=str, default="v1") + return p.parse_args() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/run_experiments.sh b/scripts/run_experiments.sh new file mode 100644 index 0000000..0e0b112 --- /dev/null +++ b/scripts/run_experiments.sh @@ -0,0 +1,143 @@ +#!/bin/bash +# ============================================================ +# Запуск 12 экспериментов на GTA-UAV (2 варианта описаний × 6 комбинаций) +# ============================================================ +# +# 6 комбинаций text_levels для каждого варианта: +# 1) level1 — только краткое описание +# 2) level2 — только пространственное +# 3) level3 — только ключевые паттерны +# 4) level1 level2 — краткое + пространственное +# 5) level1 level2 level3 — все три уровня +# 6) level1 level3 — краткое + паттерны (ablation) +# +# Использование: +# bash scripts/run_experiments.sh # все 12 экспериментов +# bash scripts/run_experiments.sh v1 # только v1 (6 экспериментов) +# bash scripts/run_experiments.sh v2 # только v2 +# bash scripts/run_experiments.sh v1 3 # v1, эксперимент 3 +# bash scripts/run_experiments.sh v2 1 4 # v2, эксперименты 1 и 4 +# ============================================================ + +set -e # Остановиться при ошибке + +# ====================== НАСТРОЙКИ =========================== +DATA_ROOT="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR" +DESCRIPTIONS_v1="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1" +DESCRIPTIONS_v2="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_v2" +DGTRS_CKPT="/media/servml/SSD_2_2TB/weights/LRSCLIP/DGTRS-CLIP-ViT-B-16.pt" +STRIPNET_CKPT="/media/servml/SSD_2_2TB/weights/stripnet_small.pth" + +# Общие гиперпараметры +EPOCHS=50 +BATCH_SIZE=64 +MICRO_BATCH=64 +LR=0.0001 +OUTPUT="outputs" +# ============================================================ + +# Массив экспериментов: номер → уровни текста +declare -A EXPERIMENTS +EXPERIMENTS[1]="level1" # Только краткое описание +EXPERIMENTS[2]="level2" # Только пространственное +EXPERIMENTS[3]="level3" # Только ключевые паттерны +EXPERIMENTS[4]="level1 level2" # Краткое + пространственное +EXPERIMENTS[5]="level1 level2 level3" # Все три уровня +EXPERIMENTS[6]="level1 level3" # Краткое + паттерны (ablation) + +# ---------- Парсинг аргументов ---------- +# Определить какие версии и эксперименты запускать +VERSIONS=() +EXP_NUMS=() + +if [ $# -eq 0 ]; then + # Без аргументов → все версии, все эксперименты + VERSIONS=(v1 v2) + EXP_NUMS=(1 2 3 4 5 6) +elif [ "$1" = "v1" ] || [ "$1" = "v2" ]; then + # Первый аргумент — версия + VERSIONS=("$1") + shift + if [ $# -eq 0 ]; then + EXP_NUMS=(1 2 3 4 5 6) + else + EXP_NUMS=("$@") + fi +else + # Числа без версии → обе версии, указанные эксперименты + VERSIONS=(v1 v2) + EXP_NUMS=("$@") +fi + +TOTAL=$(( ${#VERSIONS[@]} * ${#EXP_NUMS[@]} )) + +echo "========================================" +echo " CVGL Experiments on GTA-UAV" +echo " Epochs: $EPOCHS | Batch: $BATCH_SIZE" +echo " Versions: ${VERSIONS[*]}" +echo " Experiments: ${EXP_NUMS[*]}" +echo " Total runs: $TOTAL" +echo "========================================" + +# ---------- Проверки перед запуском ---------- +if [[ "$DGTRS_CKPT" == "/path/to/"* ]]; then + echo "❌ ОШИБКА: Заполни DGTRS_CKPT в скрипте!" + exit 1 +fi +if [[ "$STRIPNET_CKPT" == "/path/to/"* ]]; then + echo "❌ ОШИБКА: Заполни STRIPNET_CKPT в скрипте!" + exit 1 +fi + +RUN_IDX=0 + +for VERSION in "${VERSIONS[@]}"; do + # Выбрать путь к описаниям по версии + if [ "$VERSION" = "v1" ]; then + DESCRIPTIONS="$DESCRIPTIONS_v1" + else + DESCRIPTIONS="$DESCRIPTIONS_v2" + fi + + # Проверить что директория существует + if [ ! -d "$DESCRIPTIONS" ]; then + echo "⚠️ Папка описаний не найдена: $DESCRIPTIONS" + echo " Пропускаю версию $VERSION" + continue + fi + + for EXP_NUM in "${EXP_NUMS[@]}"; do + LEVELS="${EXPERIMENTS[$EXP_NUM]}" + RUN_IDX=$((RUN_IDX + 1)) + + echo "" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo " [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM: levels=[$LEVELS]" + echo " Descriptions: $DESCRIPTIONS" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + + python train.py \ + --data_root "$DATA_ROOT" \ + --descriptions_path "$DESCRIPTIONS" \ + --text_levels $LEVELS \ + --dgtrs_checkpoint "$DGTRS_CKPT" \ + --stripnet_checkpoint "$STRIPNET_CKPT" \ + --epochs $EPOCHS \ + --batch_size $BATCH_SIZE \ + --micro_batch_size $MICRO_BATCH \ + --lr $LR \ + --output_dir "${OUTPUT}/${VERSION}" \ + --inject_mona + + echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM complete" + done +done + +echo "" +echo "========================================" +echo " All $TOTAL experiments complete!" +echo " Results:" +for VERSION in "${VERSIONS[@]}"; do + echo " ${VERSION}: ${OUTPUT}/${VERSION}/" +done +echo "========================================" \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/data/__init__.py b/src/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/data/gta_uav.py b/src/data/gta_uav.py new file mode 100644 index 0000000..ae7d930 --- /dev/null +++ b/src/data/gta_uav.py @@ -0,0 +1,435 @@ +"""DataLoader для GTA-UAV с текстовыми описаниями для ОБЕИХ сторон. + +ВАЖНОЕ ИЗМЕНЕНИЕ (симметричная фьюзия): теперь каждый сэмпл содержит +ЧЕТЫРЕ компонента вместо трёх: + drone_image — картинка с дрона + drone_tokens — токенизированный текст описания дрона + satellite_image — спутниковый снимок (галерея, цель поиска) + satellite_tokens — токенизированный текст описания СПУТНИКА (НОВОЕ!) + +И дрон, и спутник имеют собственное текстовое описание (сгенерированное +одной и той же VLM по своей картинке). Модель сливает (картинка+текст) +СИММЕТРИЧНО на обеих сторонах через два отдельных TextFusionMLP, и только +затем сравнивает два слитых вектора. + +Задача: symmetric fusion drone↔satellite retrieval (NGCG), где обе стороны +представлены парой (картинка, текст), слитой в единый вектор. +""" +from __future__ import annotations + +import json +import logging +import random +import re +from pathlib import Path + +import torch +from torch.utils.data import Dataset, DataLoader +from PIL import Image +from torchvision import transforms + +from src.models.dgtrs.model import tokenize_dgtrs + +LOGGER = logging.getLogger("cvgl.data.gta_uav") + + +# --------------------------------------------------------------------------- +# Загрузка и очистка текстовых описаний +# --------------------------------------------------------------------------- + +def _clean_caption_text(raw_output: str) -> list[str]: + """Очистить output от markdown-разметки (v1) и разбить на абзацы. + + V1 формат содержит markdown-заголовки (**P1 — ...**) и зонные метки + (upper-left: ...). V2 формат — чистая проза. Функция обрабатывает оба. + Применяется одинаково и к drone-, и к satellite-описаниям — формат + разметки общий для обеих сторон. + """ + text = re.sub(r'\*\*P\d[^*]*\*\*\s*', '', raw_output) + text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text) + text = re.sub( + r'^(upper|middle|lower|center)[-‑]?(left|center|right)?:\s*', + '', text, flags=re.MULTILINE, + ) + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + cleaned = [] + for p in paragraphs: + merged = " ".join(line.strip() for line in p.split("\n") if line.strip()) + if merged: + cleaned.append(merged) + return cleaned + + +def load_text_descriptions( + descriptions_path: str | Path, + view_type: str = "drone", +) -> dict[str, dict[str, str]]: + """Загрузить текстовые описания сцен заданного типа (drone или satellite). + + Формат: *_caption.json файлы рекурсивно в descriptions_path, каждый + содержит поле "view_type" ("drone" или "satellite") и поле "output" + с 1-3 абзацами, разделёнными \\n\\n, которые маппятся на level1/2/3. + + Вызывается ДВАЖДЫ в Dataset.__init__: один раз с view_type="drone", + один раз с view_type="satellite" — обе стороны фьюзятся одинаковым + образом, поэтому обеим нужны свои описания. + + Args: + descriptions_path: Путь к директории с *_caption.json. + view_type: "drone" или "satellite" — какие описания оставить. + + Returns: + dict: {"image_name.png": {"level1": "...", "level2": "...", "level3": "..."}} + """ + path = Path(descriptions_path) + caption_files = sorted(path.glob("**/*_caption.json")) + + if not caption_files: + raise FileNotFoundError( + f"No *_caption.json files found in {path}. " + f"Expected structure: {path}/drone/images/*_caption.json " + f"and {path}/satellite/*_caption.json" + ) + + data = {} + n_single = 0 + n_skipped = 0 + + for caption_file in caption_files: + with open(caption_file, "r", encoding="utf-8") as f: + entry = json.load(f) + + if entry.get("view_type") != view_type: + n_skipped += 1 + continue + + if "rel_path" in entry: + img_name = Path(entry["rel_path"]).name + else: + img_name = caption_file.name.replace("_caption.json", ".png") + + output_text = entry.get("output", "") + paragraphs = _clean_caption_text(output_text) + + levels = { + "level1": paragraphs[0] if len(paragraphs) >= 1 else "", + "level2": paragraphs[1] if len(paragraphs) >= 2 else "", + "level3": paragraphs[2] if len(paragraphs) >= 3 else "", + } + + if len(paragraphs) < 2: + n_single += 1 + + data[img_name] = levels + + LOGGER.info( + "📝 Loaded %d %s descriptions from %s (%d single-paragraph, %d other-view skipped)", + len(data), view_type, path.name, n_single, n_skipped, + ) + + return data + + +def combine_text_levels( + descriptions: dict[str, str], + levels: list[str], + separator: str = " ", +) -> str: + """Объединить указанные уровни описаний в один текст. + + Используется одинаково для drone- и satellite-описаний — выбранные + text_levels (level1/level2/level3) применяются СИММЕТРИЧНО к обеим + сторонам в рамках одного эксперимента. + """ + parts = [] + for level in levels: + text = descriptions.get(level, "") + if text: + parts.append(text.strip()) + + combined = separator.join(parts) + + if not combined: + LOGGER.debug("Empty text after combining levels %s", levels) + combined = "No description available." + + return combined + + +# --------------------------------------------------------------------------- +# Image transforms (общие для drone и satellite изображений) +# --------------------------------------------------------------------------- + +def get_image_transforms( + image_size: int = 384, + is_train: bool = True, +) -> transforms.Compose: + """Трансформации для изображений (используются и для drone, и для satellite). + + Train: augmentations (flip, color jitter, resize+crop). + Test: только resize + normalize. + """ + normalize = transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ) + + if is_train: + return transforms.Compose([ + transforms.Resize((image_size, image_size)), + transforms.RandomHorizontalFlip(p=0.5), + transforms.RandomVerticalFlip(p=0.3), + transforms.ColorJitter( + brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05, + ), + transforms.ToTensor(), + normalize, + ]) + else: + return transforms.Compose([ + transforms.Resize((image_size, image_size)), + transforms.ToTensor(), + normalize, + ]) + + +# --------------------------------------------------------------------------- +# Dataset +# --------------------------------------------------------------------------- + +class GTAUAVTextDataset(Dataset): + """GTA-UAV dataset для symmetric drone(image+text) ↔ satellite(image+text). + + Каждый сэмпл: (drone_image, drone_tokens, satellite_image, + satellite_tokens, weight, metadata). + + Обе стороны несут пару (картинка, текст), которая сливается моделью + (TextFusionMLP) в единый вектор перед сравнением. + + Args: + data_root: Корневая папка GTA-UAV. + pairs_meta_file: JSON-файл с парами. + descriptions_path: Путь к текстовым описаниям (drone И satellite). + text_levels: Какие уровни текста использовать — одинаково + применяются к обеим сторонам. + image_size: Размер изображения (drone и satellite). + is_train: Train/test режим. + use_semipos: Использовать semi-positive пары (с весами). + context_length: Максимальная длина токенов. + """ + + def __init__( + self, + data_root: str | Path, + pairs_meta_file: str, + descriptions_path: str | Path, + text_levels: list[str] = ("level1",), + image_size: int = 384, + is_train: bool = True, + use_semipos: bool = True, + context_length: int = 248, + ) -> None: + super().__init__() + self.data_root = Path(data_root) + self.text_levels = list(text_levels) + self.is_train = is_train + self.use_semipos = use_semipos + self.context_length = context_length + self.image_size = image_size + + meta_path = self.data_root / pairs_meta_file + with open(meta_path, "r") as f: + self.pairs_meta = json.load(f) + LOGGER.info( + "📦 GTA-UAV %s: %d entries from %s", + "train" if is_train else "test", + len(self.pairs_meta), + pairs_meta_file, + ) + + # Описания загружаются ДВАЖДЫ — для дрона и для спутника отдельно, + # т.к. это разные картинки с разными описаниями. + self.drone_descriptions = load_text_descriptions(descriptions_path, view_type="drone") + self.satellite_descriptions = load_text_descriptions(descriptions_path, view_type="satellite") + + covered_drone = sum( + 1 for entry in self.pairs_meta + if entry["drone_img_name"] in self.drone_descriptions + ) + LOGGER.info( + "📝 Drone text coverage: %d/%d entries (%.1f%%)", + covered_drone, len(self.pairs_meta), + 100.0 * covered_drone / max(len(self.pairs_meta), 1), + ) + LOGGER.info( + "📝 Satellite descriptions available: %d images", + len(self.satellite_descriptions), + ) + + self.drone_transform = get_image_transforms(image_size, is_train) + self.satellite_transform = get_image_transforms(image_size, is_train) + + def __len__(self) -> int: + return len(self.pairs_meta) + + def _load_drone_image(self, entry: dict) -> torch.Tensor: + """Загрузить картинку дрона.""" + drone_img_dir = entry.get("drone_img_dir", "drone/images") + drone_img_name = entry["drone_img_name"] + drone_path = self.data_root / drone_img_dir / drone_img_name + + try: + image = Image.open(drone_path).convert("RGB") + return self.drone_transform(image) + except (FileNotFoundError, OSError) as e: + LOGGER.debug("Failed to load drone image %s: %s", drone_path, e) + return torch.zeros(3, self.image_size, self.image_size) + + def _load_satellite_image(self, entry: dict) -> tuple[torch.Tensor, float, str]: + """Загрузить спутниковый снимок (цель поиска, галерея). + + Returns: + (image_tensor, weight, sate_img_name) — имя нужно для того, + чтобы найти соответствующее текстовое описание спутника. + """ + key = "pair_pos_semipos_sate_img_list" if self.use_semipos else "pair_pos_sate_img_list" + weight_key = "pair_pos_semipos_sate_weight_list" if self.use_semipos else "pair_pos_sate_weight_list" + + sate_list = entry.get(key, entry.get("pair_pos_sate_img_list", [])) + weight_list = entry.get(weight_key, entry.get("pair_pos_sate_weight_list", [])) + + if not sate_list: + LOGGER.debug("No satellite matches for %s, using placeholder", entry.get("drone_img_name")) + return torch.zeros(3, self.image_size, self.image_size), 0.0, "" + + if self.is_train: + if weight_list and sum(weight_list) > 0: + probs = torch.tensor(weight_list, dtype=torch.float) + probs = probs / probs.sum() + chosen_idx = torch.multinomial(probs, 1).item() + else: + chosen_idx = random.randint(0, len(sate_list) - 1) + else: + chosen_idx = 0 + + sate_img_name = sate_list[chosen_idx] + weight = weight_list[chosen_idx] if weight_list else 1.0 + sate_img_dir = entry.get("sate_img_dir", "satellite") + sate_path = self.data_root / sate_img_dir / sate_img_name + + try: + image = Image.open(sate_path).convert("RGB") + image = self.satellite_transform(image) + except (FileNotFoundError, OSError) as e: + LOGGER.debug("Failed to load %s: %s", sate_path, e) + image = torch.zeros(3, self.image_size, self.image_size) + weight = 0.0 + + return image, weight, sate_img_name + + def __getitem__(self, idx: int) -> dict: + entry = self.pairs_meta[idx] + drone_img_name = entry["drone_img_name"] + + # --- Дрон: картинка + текст --- + drone_image = self._load_drone_image(entry) + drone_desc = self.drone_descriptions.get(drone_img_name, {}) + drone_text = combine_text_levels(drone_desc, self.text_levels) + drone_tokens = tokenize_dgtrs(drone_text, context_length=self.context_length, truncate=True) + drone_tokens = drone_tokens.squeeze(0) # [context_length] + + # --- Спутник: картинка + текст (НОВОЕ: текст для спутника тоже) --- + satellite_image, weight, sate_img_name = self._load_satellite_image(entry) + sat_desc = self.satellite_descriptions.get(sate_img_name, {}) + sat_text = combine_text_levels(sat_desc, self.text_levels) + satellite_tokens = tokenize_dgtrs(sat_text, context_length=self.context_length, truncate=True) + satellite_tokens = satellite_tokens.squeeze(0) # [context_length] + + drone_loc = entry.get("drone_loc_lat_lon", [0.0, 0.0]) + + return { + "drone_image": drone_image, # [3, H, W] + "drone_tokens": drone_tokens, # [248] + "satellite_image": satellite_image, # [3, H, W] + "satellite_tokens": satellite_tokens, # [248] — НОВОЕ ПОЛЕ + "weight": torch.tensor(weight, dtype=torch.float32), + "drone_img_name": drone_img_name, + "drone_lat": drone_loc[0], + "drone_lon": drone_loc[1], + } + + +# --------------------------------------------------------------------------- +# Collate и DataLoader factory +# --------------------------------------------------------------------------- + +def collate_fn(batch: list[dict]) -> dict: + """Custom collate — теперь включает satellite_tokens.""" + return { + "drone_image": torch.stack([b["drone_image"] for b in batch]), + "drone_tokens": torch.stack([b["drone_tokens"] for b in batch]), + "satellite_image": torch.stack([b["satellite_image"] for b in batch]), + "satellite_tokens": torch.stack([b["satellite_tokens"] for b in batch]), + "weight": torch.stack([b["weight"] for b in batch]), + "drone_img_name": [b["drone_img_name"] for b in batch], + "drone_lat": torch.tensor([b["drone_lat"] for b in batch]), + "drone_lon": torch.tensor([b["drone_lon"] for b in batch]), + } + + +def build_dataloaders( + data_root: str | Path, + descriptions_path: str | Path, + text_levels: list[str], + train_meta: str = "cross-area-drone2sate-train.json", + test_meta: str = "cross-area-drone2sate-test.json", + batch_size: int = 64, + num_workers: int = 4, + image_size: int = 384, +) -> tuple[DataLoader, DataLoader]: + """Создать train и test DataLoader.""" + train_dataset = GTAUAVTextDataset( + data_root=data_root, + pairs_meta_file=train_meta, + descriptions_path=descriptions_path, + text_levels=text_levels, + image_size=image_size, + is_train=True, + use_semipos=True, + ) + + test_dataset = GTAUAVTextDataset( + data_root=data_root, + pairs_meta_file=test_meta, + descriptions_path=descriptions_path, + text_levels=text_levels, + image_size=image_size, + is_train=False, + use_semipos=False, + ) + + train_loader = DataLoader( + train_dataset, + batch_size=batch_size, + shuffle=True, + num_workers=num_workers, + collate_fn=collate_fn, + pin_memory=True, + drop_last=True, + ) + + test_loader = DataLoader( + test_dataset, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + collate_fn=collate_fn, + pin_memory=True, + drop_last=False, + ) + + LOGGER.info( + "📦 DataLoaders ready: train=%d batches, test=%d batches (bs=%d, levels=%s)", + len(train_loader), len(test_loader), batch_size, text_levels, + ) + \ No newline at end of file diff --git a/src/data/gta_uav_eval.py b/src/data/gta_uav_eval.py new file mode 100644 index 0000000..095a87e --- /dev/null +++ b/src/data/gta_uav_eval.py @@ -0,0 +1,257 @@ +"""Multi-positive eval для GTA-UAV cross-view retrieval (протокол §6.2). + +В отличие от обучающего загрузчика (одна спутниковая пара на дрон, +диагональная оценка), здесь строится КОРРЕКТНАЯ retrieval-постановка: + + * галерея = УНИКАЛЬНЫЕ спутниковые тайлы теста (без дублей); + * для каждого дрона-запроса — список ВСЕХ валидных positive-тайлов + (positive + semi-positive, IoU ≥ 0.14) как индексы в галерее; + * метрика hit-if-any: запрос засчитан, если ХОТЯ БЫ ОДИН из его + positive-тайлов попал в top-K. + +Галерея собирается как объединение всех pos/semi-pos тайлов по всем +дронам теста — это все достижимые reference-тайлы данной области. + +ВАЖНО (задокументировано): отдельного официального списка галереи +для этого набора НЕТ (в папке набора — только split описаний +drone/satellite и pipeline.log). Поэтому объединение pos/semi-pos — +канонический и воспроизводимый источник галереи здесь. Единственное +следствие: тайлы тест-области, не являющиеся positive НИ для одного +дрона (чистые дистракторы), в галерею не попадают → абсолютный recall +слегка оптимистичен. Смещение ОДИНАКОВО для всех вариантов, поэтому +ΔR@1 между вариантами (основа decision rule) остаётся честным. +""" +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import torch +from PIL import Image +from torch.utils.data import DataLoader, Dataset + +from src.data.gta_uav import ( + combine_text_levels, + get_image_transforms, + load_text_descriptions, +) +from src.models.dgtrs.model import tokenize_dgtrs + +LOGGER = logging.getLogger("cvgl.eval") + +# Поля meta с positive/semi-positive списками (pos + semi-pos, IoU ≥ 0.14). +_SEMIPOS_KEY = "pair_pos_semipos_sate_img_list" +_POS_KEY = "pair_pos_sate_img_list" + + +def _entry_positive_names(entry: dict) -> list[str]: + """Список positive+semi-positive спутниковых тайлов для одного дрона.""" + names = entry.get(_SEMIPOS_KEY) or entry.get(_POS_KEY) or [] + # Уникализируем, сохраняя порядок (у одного дрона тайл не должен дублиться). + seen: dict[str, None] = {} + for name in names: + seen.setdefault(name, None) + return list(seen.keys()) + + +def collect_gallery_names(pairs_meta: list[dict]) -> tuple[list[str], dict[str, str]]: + """Собрать уникальную галерею тайлов из meta теста. + + Returns: + gallery_names: отсортированный список уникальных имён тайлов. + name_to_dir: имя тайла → каталог (sate_img_dir) для загрузки. + """ + name_to_dir: dict[str, str] = {} + for entry in pairs_meta: + sate_dir = entry.get("sate_img_dir", "satellite") + for name in _entry_positive_names(entry): + name_to_dir.setdefault(name, sate_dir) + gallery_names = sorted(name_to_dir.keys()) + return gallery_names, name_to_dir + + +class _SatelliteGalleryDataset(Dataset): + """Уникальные спутниковые тайлы (картинка + текст) — галерея.""" + + def __init__( + self, + data_root: Path, + gallery_names: list[str], + name_to_dir: dict[str, str], + satellite_descriptions: dict[str, dict], + text_levels: list[str], + image_size: int, + context_length: int, + ) -> None: + self.data_root = data_root + self.gallery_names = gallery_names + self.name_to_dir = name_to_dir + self.satellite_descriptions = satellite_descriptions + self.text_levels = text_levels + self.image_size = image_size + self.context_length = context_length + self.transform = get_image_transforms(image_size, is_train=False) + + def __len__(self) -> int: + return len(self.gallery_names) + + def __getitem__(self, idx: int) -> dict: + name = self.gallery_names[idx] + path = self.data_root / self.name_to_dir[name] / name + try: + image = self.transform(Image.open(path).convert("RGB")) + except (FileNotFoundError, OSError) as e: + LOGGER.debug("Failed to load gallery tile %s: %s", path, e) + image = torch.zeros(3, self.image_size, self.image_size) + + desc = self.satellite_descriptions.get(name, {}) + text = combine_text_levels(desc, self.text_levels) + tokens = tokenize_dgtrs(text, context_length=self.context_length, truncate=True).squeeze(0) + return {"image": image, "tokens": tokens, "name": name} + + +class _DroneQueryDataset(Dataset): + """Дрон-запросы (картинка + текст).""" + + def __init__( + self, + data_root: Path, + entries: list[dict], + drone_descriptions: dict[str, dict], + text_levels: list[str], + image_size: int, + context_length: int, + ) -> None: + self.data_root = data_root + self.entries = entries + self.drone_descriptions = drone_descriptions + self.text_levels = text_levels + self.image_size = image_size + self.context_length = context_length + self.transform = get_image_transforms(image_size, is_train=False) + + def __len__(self) -> int: + return len(self.entries) + + def __getitem__(self, idx: int) -> dict: + entry = self.entries[idx] + name = entry["drone_img_name"] + img_dir = entry.get("drone_img_dir", "drone/images") + path = self.data_root / img_dir / name + try: + image = self.transform(Image.open(path).convert("RGB")) + except (FileNotFoundError, OSError) as e: + LOGGER.debug("Failed to load drone query %s: %s", path, e) + image = torch.zeros(3, self.image_size, self.image_size) + + desc = self.drone_descriptions.get(name, {}) + text = combine_text_levels(desc, self.text_levels) + tokens = tokenize_dgtrs(text, context_length=self.context_length, truncate=True).squeeze(0) + return {"image": image, "tokens": tokens, "name": name} + + +def _collate(batch: list[dict]) -> dict: + return { + "image": torch.stack([b["image"] for b in batch]), + "tokens": torch.stack([b["tokens"] for b in batch]), + "name": [b["name"] for b in batch], + } + + +def invert_positives( + positives_q2g: list[list[int]], + n_gallery: int, +) -> list[list[int]]: + """Инвертировать positive-карту: для каждого тайла — список дрон-запросов.""" + positives_g2q: list[list[int]] = [[] for _ in range(n_gallery)] + for query_idx, gallery_indices in enumerate(positives_q2g): + for g in gallery_indices: + positives_g2q[g].append(query_idx) + return positives_g2q + + +def build_multipos_eval( + data_root: str | Path, + test_meta: str, + descriptions_path: str | Path, + text_levels: list[str], + image_size: int = 384, + batch_size: int = 64, + num_workers: int = 4, + context_length: int = 248, +) -> dict: + """Собрать multi-positive eval: query-loader, gallery-loader, positive-карты. + + Дроны без единого валидного positive-тайла в галерее исключаются из + запросов (recall для них не определён), с логированием количества. + + Returns: + dict с ключами: + drone_loader: DataLoader дрон-запросов. + gallery_loader: DataLoader уникальной спутниковой галереи. + positives_q2g: list[list[int]] — дрон → индексы тайлов галереи. + positives_g2q: list[list[int]] — тайл → индексы дрон-запросов. + gallery_names: list[str] — имена тайлов галереи (порядок = индекс). + n_queries / n_gallery / n_skipped: счётчики. + """ + data_root = Path(data_root) + meta_path = data_root / test_meta + with open(meta_path, "r") as f: + pairs_meta = json.load(f) + + drone_descriptions = load_text_descriptions(descriptions_path, view_type="drone") + satellite_descriptions = load_text_descriptions(descriptions_path, view_type="satellite") + + # 1. Уникальная галерея тайлов + индекс. + gallery_names, name_to_dir = collect_gallery_names(pairs_meta) + name_to_idx = {name: i for i, name in enumerate(gallery_names)} + + # 2. Запросы + positive-карта (только дроны с ≥1 positive в галерее). + kept_entries: list[dict] = [] + positives_q2g: list[list[int]] = [] + n_skipped = 0 + for entry in pairs_meta: + pos_idx = [name_to_idx[n] for n in _entry_positive_names(entry) if n in name_to_idx] + if not pos_idx: + n_skipped += 1 + continue + kept_entries.append(entry) + positives_q2g.append(pos_idx) + + positives_g2q = invert_positives(positives_q2g, len(gallery_names)) + + LOGGER.info( + "🔎 Multi-positive eval: %d queries (%d skipped, no positive), " + "%d unique gallery tiles", + len(kept_entries), n_skipped, len(gallery_names), + ) + + gallery_ds = _SatelliteGalleryDataset( + data_root, gallery_names, name_to_dir, satellite_descriptions, + text_levels, image_size, context_length, + ) + drone_ds = _DroneQueryDataset( + data_root, kept_entries, drone_descriptions, + text_levels, image_size, context_length, + ) + + drone_loader = DataLoader( + drone_ds, batch_size=batch_size, shuffle=False, + num_workers=num_workers, collate_fn=_collate, + ) + gallery_loader = DataLoader( + gallery_ds, batch_size=batch_size, shuffle=False, + num_workers=num_workers, collate_fn=_collate, + ) + + return { + "drone_loader": drone_loader, + "gallery_loader": gallery_loader, + "positives_q2g": positives_q2g, + "positives_g2q": positives_g2q, + "gallery_names": gallery_names, + "n_queries": len(kept_entries), + "n_gallery": len(gallery_names), + "n_skipped": n_skipped, + } \ No newline at end of file diff --git a/src/losses.py b/src/losses.py new file mode 100644 index 0000000..9b73fe1 --- /dev/null +++ b/src/losses.py @@ -0,0 +1,143 @@ +"""Contrastive loss functions для dual-encoder retrieval. + +InfoNCE (он же NT-Xent, CLIP loss) — стандарт для contrastive learning. +Также включён вариант weighted-InfoNCE из Game4Loc для GTA-UAV. + +Как работает InfoNCE: + У тебя батч из B пар (text_i, image_i). Строим матрицу сходства B×B. + Правильная пара — на диагонали (text_0↔image_0, text_1↔image_1, ...). + Все остальные B-1 пар в строке — «негативные примеры». + + Loss = CrossEntropy по строкам (text→image) + CrossEntropy по столбцам (image→text). + + Это заставляет модель: + - Приближать правильные пары (высокое сходство) + - Отдалять неправильные (низкое сходство) + + Чем больше B, тем больше негативов, тем сильнее сигнал. + Поэтому batch_size=64 важен. +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class InfoNCELoss(nn.Module): + """Symmetric InfoNCE loss (как в CLIP). + + Считает loss в обе стороны: + - text→image: для каждого текста, какое изображение правильное? + - image→text: для каждого изображения, какой текст правильный? + Усредняет оба направления. + + Args: + label_smoothing: Сглаживание меток. 0.0 = жёсткие метки (стандарт CLIP). + Небольшое значение (0.1) может помочь при шумных описаниях. + """ + + def __init__(self, label_smoothing: float = 0.0) -> None: + super().__init__() + self.label_smoothing = label_smoothing + + def forward(self, logits: torch.Tensor) -> dict[str, torch.Tensor]: + """ + Args: + logits: Матрица сходства [B, B] = text_emb @ image_emb.T / τ. + Правильные пары на диагонали. + + Returns: + dict: + loss: Средний loss (скаляр). + loss_t2i: Text-to-image loss. + loss_i2t: Image-to-text loss. + acc_t2i: Accuracy text→image (для мониторинга). + acc_i2t: Accuracy image→text. + """ + B = logits.shape[0] + # Метки: правильная пара для i-го текста — i-е изображение + labels = torch.arange(B, device=logits.device) + + # Text → Image: строки матрицы + loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing) + # Image → Text: столбцы матрицы (транспонируем) + loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing) + + loss = (loss_t2i + loss_i2t) / 2.0 + + # Accuracy для мониторинга (не для backprop) + with torch.no_grad(): + acc_t2i = (logits.argmax(dim=1) == labels).float().mean() + acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean() + + return { + "loss": loss, + "loss_t2i": loss_t2i, + "loss_i2t": loss_i2t, + "acc_t2i": acc_t2i, + "acc_i2t": acc_i2t, + } + + +class WeightedInfoNCELoss(nn.Module): + """Weighted InfoNCE из Game4Loc для GTA-UAV. + + В GTA-UAV не все пары одинаково «правильные» — есть positive (полное + совпадение) и semi-positive (частичное перекрытие по IoU). Веса + отражают степень совпадения. + + Вместо жёстких меток [0,0,1,0,...] используются мягкие [0.1, 0.0, 0.7, 0.2,...], + где вес пропорционален IoU перекрытия drone-satellite пар. + + Args: + label_smoothing: Дополнительное сглаживание. + """ + + def __init__(self, label_smoothing: float = 0.0) -> None: + super().__init__() + self.label_smoothing = label_smoothing + + def forward( + self, + logits: torch.Tensor, + weights: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """ + Args: + logits: [B, B] матрица сходства. + weights: [B, B] матрица весов (soft labels). Если None → обычный InfoNCE. + """ + B = logits.shape[0] + + if weights is None: + # Fallback на обычный InfoNCE + labels = torch.arange(B, device=logits.device) + loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing) + loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing) + else: + # Нормализуем веса в распределение вероятностей по каждой строке + targets_t2i = F.softmax(weights, dim=1) # для text→image + targets_i2t = F.softmax(weights.T, dim=1) # для image→text + + # KL-divergence вместо cross-entropy (для soft labels) + log_probs_t2i = F.log_softmax(logits, dim=1) + log_probs_i2t = F.log_softmax(logits.T, dim=1) + + loss_t2i = F.kl_div(log_probs_t2i, targets_t2i, reduction="batchmean") + loss_i2t = F.kl_div(log_probs_i2t, targets_i2t, reduction="batchmean") + + loss = (loss_t2i + loss_i2t) / 2.0 + + with torch.no_grad(): + labels = torch.arange(B, device=logits.device) + acc_t2i = (logits.argmax(dim=1) == labels).float().mean() + acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean() + + return { + "loss": loss, + "loss_t2i": loss_t2i, + "loss_i2t": loss_i2t, + "acc_t2i": acc_t2i, + "acc_i2t": acc_i2t, + } \ No newline at end of file diff --git a/src/metrics.py b/src/metrics.py new file mode 100644 index 0000000..f3854a9 --- /dev/null +++ b/src/metrics.py @@ -0,0 +1,163 @@ +"""Метрики для оценки retrieval качества. + +После обучения мы прогоняем все тексты и изображения через энкодеры, +строим матрицу сходства и для каждого запроса ранжируем кандидатов. + +Recall@K — основная метрика: + «Какая доля запросов нашла правильный ответ в топ-K?» + Recall@1 = 0.65 значит: в 65% случаев правильное изображение — первое в списке. + +AP (Average Precision): + Усреднённая точность по всем позициям ранжирования. + Учитывает не только «попал ли в топ-K», но и на какой именно позиции. + +Meter-level distance (специфика GTA-UAV): + Каждое изображение имеет GPS-координаты. Мы можем посчитать + расстояние в метрах между предсказанной и реальной позицией. +""" +from __future__ import annotations + +import numpy as np +import torch + + +@torch.no_grad() +def compute_retrieval_metrics( + text_embeddings: torch.Tensor, + image_embeddings: torch.Tensor, + ks: list[int] = (1, 5, 10), + compute_ap: bool = True, +) -> dict[str, float]: + """Вычислить retrieval метрики. + + Предполагается что text_embeddings[i] соответствует image_embeddings[i] + (i-й текст описывает i-е изображение). + + Args: + text_embeddings: [N, D] нормализованные текстовые эмбеддинги. + image_embeddings: [N, D] нормализованные визуальные эмбеддинги. + ks: Значения K для Recall@K. + compute_ap: Считать ли Average Precision (медленнее). + + Returns: + dict с метриками: recall@1, recall@5, recall@10, AP, ... + """ + N = text_embeddings.shape[0] + + # Матрица сходства: [N, N]. sim[i][j] = косинусное сходство текста i и изображения j. + sim = text_embeddings @ image_embeddings.T # уже L2-нормализованы → cosine sim + + # Для каждого текста ранжируем все изображения по убыванию сходства. + # ranks[i] = на какой позиции (0-indexed) стоит правильное изображение i. + # Правильное изображение для текста i — это image i (диагональ). + sorted_indices = sim.argsort(dim=1, descending=True) # [N, N] + + # Найти позицию правильного ответа (i-е изображение для i-го текста) + ranks = torch.zeros(N, dtype=torch.long, device=sim.device) + for i in range(N): + rank = (sorted_indices[i] == i).nonzero(as_tuple=True)[0].item() + ranks[i] = rank + + results = {} + + # Recall@K: доля запросов, где правильный ответ в топ-K + for k in ks: + recall = (ranks < k).float().mean().item() + results[f"recall@{k}"] = recall + + # Recall@1%: правильный ответ в топ-1% базы + k_1pct = max(1, N // 100) + results["recall@1%"] = (ranks < k_1pct).float().mean().item() + + # Mean Rank и Median Rank (для анализа — чем меньше, тем лучше) + results["mean_rank"] = ranks.float().mean().item() + results["median_rank"] = ranks.float().median().item() + + # Average Precision + if compute_ap: + # AP = среднее по всем запросам от 1/(rank+1) + # Это упрощённый AP для случая одного правильного ответа на запрос. + ap = (1.0 / (ranks.float() + 1)).mean().item() + results["AP"] = ap + + return results + + +@torch.no_grad() +def compute_bidirectional_metrics( + text_embeddings: torch.Tensor, + image_embeddings: torch.Tensor, + ks: list[int] = (1, 5, 10), +) -> dict[str, float]: + """Метрики в обе стороны: text→image и image→text. + + Returns: + dict с префиксами t2i_ и i2t_ для каждого направления. + """ + t2i = compute_retrieval_metrics(text_embeddings, image_embeddings, ks=ks) + # Для image→text: переворачиваем — ищем текст по изображению + i2t = compute_retrieval_metrics(image_embeddings, text_embeddings, ks=ks) + + results = {} + for key, val in t2i.items(): + results[f"t2i_{key}"] = val + for key, val in i2t.items(): + results[f"i2t_{key}"] = val + + # Средний Recall@K по обоим направлениям + for k in ks: + results[f"avg_recall@{k}"] = (t2i[f"recall@{k}"] + i2t[f"recall@{k}"]) / 2 + + return results + + +def compute_meter_distance( + predicted_indices: np.ndarray, + query_coords: np.ndarray, + gallery_coords: np.ndarray, +) -> dict[str, float]: + """Расстояние в метрах между предсказанной и реальной позицией. + + Специфика GTA-UAV: каждое изображение имеет GPS-координаты. + Мы берём координаты top-1 предсказания и считаем расстояние до GT. + + Args: + predicted_indices: [N] индексы top-1 предсказанных изображений. + query_coords: [N, 2] GPS-координаты запросов (lat, lon). + gallery_coords: [M, 2] GPS-координаты базы изображений. + + Returns: + dict: mean_distance_m, median_distance_m, recall_at_Xm. + """ + predicted_coords = gallery_coords[predicted_indices] # [N, 2] + + # Haversine distance (приблизительно, для малых расстояний ≈ Euclidean × scale) + # Для GTA-UAV координаты в игровых юнитах, не GPS → используем Euclidean + diffs = predicted_coords - query_coords # [N, 2] + distances = np.linalg.norm(diffs, axis=1) # [N] в метрах (или юнитах) + + results = { + "mean_distance": float(distances.mean()), + "median_distance": float(np.median(distances)), + } + + # Recall at distance thresholds + for threshold in (5, 10, 25, 50, 100): + results[f"recall@{threshold}m"] = float((distances < threshold).mean()) + + return results + + +def format_metrics(metrics: dict[str, float], prefix: str = "") -> str: + """Форматировать метрики для логирования.""" + lines = [] + for key, val in metrics.items(): + if "recall" in key or "AP" in key or "acc" in key: + lines.append(f"{prefix}{key}: {val:.4f} ({val*100:.1f}%)") + elif "rank" in key: + lines.append(f"{prefix}{key}: {val:.1f}") + elif "distance" in key: + lines.append(f"{prefix}{key}: {val:.1f}m") + else: + lines.append(f"{prefix}{key}: {val:.4f}") + return "\n".join(lines) diff --git a/src/models/__init__.py b/src/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/dgtrs/__init__.py b/src/models/dgtrs/__init__.py new file mode 100644 index 0000000..5461f09 --- /dev/null +++ b/src/models/dgtrs/__init__.py @@ -0,0 +1,6 @@ +from src.models.dgtrs.model import ( + DGTRSTextEncoder, + load_dgtrs_text_encoder, + tokenize_dgtrs, + build_model, +) diff --git a/src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz b/src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000..7b5088a Binary files /dev/null and b/src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz differ diff --git a/src/models/dgtrs/model.py b/src/models/dgtrs/model.py new file mode 100644 index 0000000..2f238fb --- /dev/null +++ b/src/models/dgtrs/model.py @@ -0,0 +1,284 @@ +"""DGTRS-CLIP model — text encoder components. + +Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0). +Only text-encoder classes are kept; vision encoder removed. +""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from pathlib import Path +from typing import Union + +import coloredlogs +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from src.models.dgtrs.simple_tokenizer import SimpleTokenizer as _Tokenizer + +LOGGER = logging.getLogger("caption_test.dgtrs") +coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s") + +_tokenizer = _Tokenizer() + + +# --------------------------------------------------------------------------- +# Transformer blocks (DGTRS original — sequence-first LND format) +# --------------------------------------------------------------------------- + +class LayerNorm(nn.LayerNorm): + """LayerNorm that handles fp16 by casting to fp32 internally.""" + + def forward(self, x: torch.Tensor): + orig_type = x.dtype + ret = super().forward(x.type(torch.float32)) + return ret.type(orig_type) + + +class QuickGELU(nn.Module): + def forward(self, x: torch.Tensor): + return x * torch.sigmoid(1.702 * x) + + +class ResidualAttentionBlock(nn.Module): + def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): + super().__init__() + self.attn = nn.MultiheadAttention(d_model, n_head) + self.ln_1 = LayerNorm(d_model) + self.mlp = nn.Sequential(OrderedDict([ + ("c_fc", nn.Linear(d_model, d_model * 4)), + ("gelu", QuickGELU()), + ("c_proj", nn.Linear(d_model * 4, d_model)), + ])) + self.ln_2 = LayerNorm(d_model) + self.attn_mask = attn_mask + + def attention(self, x: torch.Tensor): + self.attn_mask = ( + self.attn_mask.to(dtype=x.dtype, device=x.device) + if self.attn_mask is not None else None + ) + return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] + + def forward(self, x: torch.Tensor): + x = x + self.attention(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +class Transformer(nn.Module): + def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None): + super().__init__() + self.width = width + self.layers = layers + self.resblocks = nn.Sequential( + *[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)] + ) + self.gradient_checkpointing = False + + def forward(self, x: torch.Tensor): + if self.gradient_checkpointing and self.training: + for block in self.resblocks: + x = torch.utils.checkpoint.checkpoint( + block, x, use_reentrant=False, + ) + return x + return self.resblocks(x) + + +# --------------------------------------------------------------------------- +# DGTRS-CLIP text encoder +# --------------------------------------------------------------------------- + +class DGTRSTextEncoder(nn.Module): + """DGTRS-CLIP text encoder with KPS dual positional embeddings. + + Context length: 248 tokens. + Uses mask1 (positions 0-19) for positional_embedding and + mask2 (positions 20-247) for positional_embedding_res. + + This is the official architecture from + https://github.com/MitsuiChen14/DGTRS/blob/main/model/model_longclip.py + """ + + CONTEXT_LENGTH = 248 + + def __init__( + self, + vocab_size: int = 49408, + transformer_width: int = 768, + transformer_heads: int = 12, + transformer_layers: int = 12, + embed_dim: int = 768, + ) -> None: + super().__init__() + self.context_length = self.CONTEXT_LENGTH + self.embed_dim = embed_dim + + self.token_embedding = nn.Embedding(vocab_size, transformer_width) + self.positional_embedding = nn.Parameter( + torch.empty(self.CONTEXT_LENGTH, transformer_width), + ) + self.positional_embedding_res = nn.Parameter( + torch.empty(self.CONTEXT_LENGTH, transformer_width), + ) + + self.transformer = Transformer( + width=transformer_width, + layers=transformer_layers, + heads=transformer_heads, + attn_mask=self._build_attention_mask(), + ) + + self.ln_final = LayerNorm(transformer_width) + self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim)) + + # KPS masks: positions 0-19 use positional_embedding, + # positions 20-247 use positional_embedding_res. + mask1 = torch.zeros(self.CONTEXT_LENGTH, 1) + mask1[:20, :] = 1 + mask2 = torch.zeros(self.CONTEXT_LENGTH, 1) + mask2[20:, :] = 1 + self.register_buffer("mask1", mask1) + self.register_buffer("mask2", mask2) + + def _build_attention_mask(self) -> torch.Tensor: + mask = torch.empty(self.CONTEXT_LENGTH, self.CONTEXT_LENGTH) + mask.fill_(float("-inf")) + mask.triu_(1) + return mask + + @property + def dtype(self) -> torch.dtype: + return self.token_embedding.weight.dtype + + def forward(self, text: torch.Tensor) -> torch.Tensor: + """Encode tokenized text. + + Args: + text: Token IDs [B, 248] (int/long). + + Returns: + Text embeddings [B, embed_dim]. + """ + x = self.token_embedding(text).type(self.dtype) # [B, 248, D] + + # Dual masked positional embeddings (KPS). + x = x + ( + self.positional_embedding * self.mask1 + + self.positional_embedding_res * self.mask2 + ).type(self.dtype) + + x = x.permute(1, 0, 2) # NLD -> LND (sequence-first for nn.MultiheadAttention) + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + x = self.ln_final(x).type(self.dtype) + + # Take features from EOT token (highest token ID in each sequence). + x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection + return x + + +# --------------------------------------------------------------------------- +# Loading and tokenization +# --------------------------------------------------------------------------- + +def _extract_text_state(full_state: dict) -> dict: + """Extract text encoder keys from a full DGTRS-CLIP state dict.""" + text_keys = { + "token_embedding", "positional_embedding", "positional_embedding_res", + "transformer", "ln_final", "text_projection", + } + return { + k: v for k, v in full_state.items() + if any(k == prefix or k.startswith(prefix + ".") for prefix in text_keys) + } + + +def build_model(state_dict: dict) -> DGTRSTextEncoder: + """Build DGTRSTextEncoder from a DGTRS-CLIP state dict. + + Auto-detects architecture dimensions from the state dict. + """ + embed_dim = state_dict["text_projection"].shape[1] + vocab_size = state_dict["token_embedding.weight"].shape[0] + transformer_width = state_dict["ln_final.weight"].shape[0] + transformer_heads = transformer_width // 64 + transformer_layers = len(set( + k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks") + )) + + model = DGTRSTextEncoder( + vocab_size=vocab_size, + transformer_width=transformer_width, + transformer_heads=transformer_heads, + transformer_layers=transformer_layers, + embed_dim=embed_dim, + ) + # mask1/mask2 are buffers created in __init__, not in checkpoint. + model.load_state_dict(state_dict, strict=False) + return model.eval() + + +def load_dgtrs_text_encoder( + checkpoint_path: str | Path, + device: str = "cpu", +) -> DGTRSTextEncoder: + """Load DGTRS-CLIP text encoder from checkpoint. + + Extracts only text encoder weights from the full CLIP checkpoint. + """ + path = Path(checkpoint_path) + LOGGER.info("📝 Loading DGTRS-CLIP 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"] + + text_state = _extract_text_state(full_state) + model = build_model(text_state) + model = model.to(device) + + n_params = sum(p.numel() for p in model.parameters()) + LOGGER.info( + "📝 DGTRS text encoder loaded: %s params, context=%d tokens", + f"{n_params:,}", model.context_length, + ) + return model + + +def tokenize_dgtrs( + texts: str | list[str], + context_length: int = 248, + truncate: bool = True, +) -> torch.Tensor: + """Tokenize text for DGTRS-CLIP (248 token context). + + Args: + texts: Input string or list of strings. + context_length: Output sequence length (default 248). + truncate: Whether to truncate long sequences. + + Returns: + Token IDs [B, context_length] (int tensor). + """ + if isinstance(texts, str): + texts = [texts] + + sot_token = _tokenizer.encoder["<|startoftext|>"] + eot_token = _tokenizer.encoder["<|endoftext|>"] + all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] + + result = torch.zeros(len(all_tokens), context_length, dtype=torch.int) + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + if truncate: + tokens = tokens[:context_length] + tokens[-1] = eot_token + else: + raise RuntimeError(f"Input too long ({len(tokens)} > {context_length})") + result[i, :len(tokens)] = torch.tensor(tokens) + + return result diff --git a/src/models/dgtrs/simple_tokenizer.py b/src/models/dgtrs/simple_tokenizer.py new file mode 100644 index 0000000..9bae2b9 --- /dev/null +++ b/src/models/dgtrs/simple_tokenizer.py @@ -0,0 +1,125 @@ +"""BPE tokenizer for DGTRS-CLIP / LongCLIP. + +Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0). +""" + +import gzip +import html +import os +from functools import lru_cache + +import ftfy +import regex as re + + +@lru_cache() +def default_bpe(): + return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz") + + +@lru_cache() +def bytes_to_unicode(): + bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +class SimpleTokenizer: + def __init__(self, bpe_path: str = default_bpe()): + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + merges = gzip.open(bpe_path).read().decode("utf-8").split("\n") + merges = merges[1:49152 - 256 - 2 + 1] + merges = [tuple(merge.split()) for merge in merges] + vocab = list(bytes_to_unicode().values()) + vocab = vocab + [v + "" for v in vocab] + for merge in merges: + vocab.append("".join(merge)) + vocab.extend(["<|startoftext|>", "<|endoftext|>"]) + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.decoder = {v: k for k, v in self.encoder.items()} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {"<|startoftext|>": "<|startoftext|>", "<|endoftext|>": "<|endoftext|>"} + self.pat = re.compile( + r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", + re.IGNORECASE, + ) + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + (token[-1] + "",) + pairs = get_pairs(word) + if not pairs: + return token + "" + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + new_word.extend(word[i:j]) + i = j + except ValueError: + new_word.extend(word[i:]) + break + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = " ".join(word) + self.cache[token] = word + return word + + def encode(self, text): + bpe_tokens = [] + text = whitespace_clean(basic_clean(text)).lower() + for token in re.findall(self.pat, text): + token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) + bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" ")) + return bpe_tokens + + def decode(self, tokens): + text = "".join([self.decoder[token] for token in tokens]) + text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors="replace").replace("", " ") + return text diff --git a/src/models/dual_encoder.py b/src/models/dual_encoder.py new file mode 100644 index 0000000..c79be39 --- /dev/null +++ b/src/models/dual_encoder.py @@ -0,0 +1,431 @@ +"""Dual-encoder модель для symmetric text-guided cross-view geo-localization. + +ИСТОРИЯ ИЗМЕНЕНИЙ АРХИТЕКТУРЫ (по правкам от Ярослава): + v1: текст и картинка дрона кодировались раздельно и сравнивались + между собой напрямую (CLIP-style text↔image retrieval). + v2: картинка дрона и текст-описание СЛИВАЛИСЬ в один вектор + (TextFusionMLP), спутник оставался чисто визуальным. + v3 (текущая): И дрон, И спутник имеют собственное текстовое описание. + Слияние происходит СИММЕТРИЧНО на обеих сторонах: + + Drone (картинка + текст → слияние) ──┐ + ├─→ cosine similarity → InfoNCE + Satellite (картинка + текст → слияние) ┘ + +Зачем симметричное слияние: + У спутникового снимка тоже есть текстовое описание (та же VLM, тот же + формат), сгенерированное по картинке спутника. Раз у обеих сторон есть + пара (картинка, текст), логично использовать оба сигнала одинаковым + образом — иначе одна сторона получает преимущество (доступ к двум + модальностям), а другая остаётся обеднённой (только картинка), что не + оправдано, если данные для текста есть и там, и там. + +Зачем ДВЕ отдельные TextFusionMLP (а не одна общая для обоих видов): + Визуальные признаки дрона и спутника лежат в разных доменах (наклонный + вид с дрона vs нормированный надирный спутниковый снимок), поэтому то, + как нужно сочетать "картинка + текст" в один вектор, тоже отличается. + Текстовый энкодер (DGTRS-CLIP) используется один на двоих — это одна и + та же языковая модель независимо от того, какой вид она описывает. + +Зачем residual-gate в TextFusionMLP (для обеих сторон): + Картинка — "якорь" (всегда присутствует и валидна). + Текст — дополняющий сигнал переменной полезности (level1 несёт больше + сигнала чем level3, см. результаты экспериментов). Residual-gate + позволяет модели самой регулировать вклад текста, а не сливать 50/50. +""" +from __future__ import annotations + +import logging +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from src.models.dgtrs.model import ( + DGTRSTextEncoder, + load_dgtrs_text_encoder, + tokenize_dgtrs, +) +from src.models.stripnet_encoder import StripNetEncoder + +LOGGER = logging.getLogger("cvgl.dual_encoder") + + +# --------------------------------------------------------------------------- +# Слияние картинки и текста (TextFusionMLP) — используется для ОБЕИХ сторон +# --------------------------------------------------------------------------- + +class TextFusionMLP(nn.Module): + """Слияние визуальных признаков (дрона ИЛИ спутника) с текстовым описанием. + + Архитектура: + image_feat ──→ Linear ──┐ + ├─→ concat ──→ MLP ──┐ + text_feat ──→ Linear ──┘ │ + ├─→ gate (sigmoid) + image_proj ─────────────────────────(residual)─┴─→ fused = image_proj + gate * candidate + + image_proj выступает "якорем" (residual connection): даже если текст + неинформативен (gate → 0), слитый эмбеддинг не теряет визуальный сигнал. + + Используется как для дрона, так и для спутника — но как ДВА РАЗНЫХ + экземпляра с независимыми весами (см. DualEncoder), поскольку визуальные + домены различаются. + + Args: + image_dim: Размерность визуальных признаков (вход). + text_dim: Размерность текстовых признаков (вход). + fused_dim: Размерность слитого вектора (выход). + hidden_dim: Размерность скрытого слоя MLP (по умолчанию = fused_dim). + """ + + def __init__( + self, + image_dim: int, + text_dim: int, + fused_dim: int, + hidden_dim: int | None = None, + ) -> None: + super().__init__() + hidden_dim = hidden_dim or fused_dim + + self.image_proj = nn.Linear(image_dim, fused_dim) + self.text_proj = nn.Linear(text_dim, fused_dim) + + self.fusion_mlp = nn.Sequential( + nn.Linear(fused_dim * 2, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, fused_dim), + ) + self.gate = nn.Sequential( + nn.Linear(fused_dim * 2, fused_dim), + nn.Sigmoid(), + ) + + for m in (self.image_proj, self.text_proj): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.zeros_(m.bias) + for m in self.fusion_mlp.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.zeros_(m.bias) + # Gate инициализируем так, чтобы изначально пропускать мало текста + # (bias смещён в отрицательную сторону → sigmoid ≈ 0 на старте). + last_gate_linear = self.gate[0] + nn.init.trunc_normal_(last_gate_linear.weight, std=0.02) + nn.init.constant_(last_gate_linear.bias, -2.0) + + def forward(self, image_feat: torch.Tensor, text_feat: torch.Tensor) -> torch.Tensor: + """Слить визуальные и текстовые признаки в один вектор. + + Args: + image_feat: [B, image_dim] — визуальные признаки (до проекции). + text_feat: [B, text_dim] — текстовые признаки (до проекции). + + Returns: + [B, fused_dim] — слитый вектор (картинка как якорь + gated текст). + """ + img = self.image_proj(image_feat) + txt = self.text_proj(text_feat) + + concat = torch.cat([img, txt], dim=-1) + candidate = self.fusion_mlp(concat) + gate = self.gate(concat) + + fused = img + gate * candidate + return fused + + +# --------------------------------------------------------------------------- +# DualEncoder: drone (fused) vs satellite (fused) — СИММЕТРИЧНО +# --------------------------------------------------------------------------- + +class DualEncoder(nn.Module): + """Symmetric dual-encoder: слитый drone-эмбеддинг vs слитый satellite-эмбеддинг. + + И дрон, и спутник проходят одинаковую по структуре, но раздельную по + весам процедуру: картинка + текст → TextFusionMLP → проекция → + L2-норма. Полученные эмбеддинги сравниваются через cosine similarity. + + Args: + text_encoder: DGTRS-CLIP текстовый энкодер — ОБЩИЙ для + текста дрона и текста спутника (это одна + и та же языковая модель). + drone_image_encoder: StripNet энкодер для картинок дрона. + satellite_image_encoder: StripNet энкодер для спутниковых снимков + (отдельный экземпляр — домены различны). + fused_dim: Размерность вектора после слияния. + shared_dim: Итоговая размерность общего пространства. + temperature_init: Начальное значение temperature для InfoNCE. + """ + + def __init__( + self, + text_encoder: DGTRSTextEncoder, + drone_image_encoder: StripNetEncoder, + satellite_image_encoder: StripNetEncoder, + fused_dim: int = 512, + shared_dim: int = 512, + temperature_init: float = 0.07, + ) -> None: + super().__init__() + + self.text_encoder = text_encoder + self.drone_image_encoder = drone_image_encoder + self.satellite_image_encoder = satellite_image_encoder + self.fused_dim = fused_dim + self.shared_dim = shared_dim + + # --- Слияние картинки + текста — ДВА отдельных модуля --- + # Раздельные веса: визуальные домены дрона и спутника различны, + # поэтому оптимальный способ "смешивания" картинки с текстом тоже + # должен быть индивидуальным для каждой стороны. + self.drone_fusion = TextFusionMLP( + image_dim=drone_image_encoder.out_dim, + text_dim=text_encoder.embed_dim, + fused_dim=fused_dim, + ) + self.satellite_fusion = TextFusionMLP( + image_dim=satellite_image_encoder.out_dim, + text_dim=text_encoder.embed_dim, + fused_dim=fused_dim, + ) + + # --- Проекционные головы в общее пространство сравнения --- + self.drone_proj = nn.Sequential( + nn.Linear(fused_dim, shared_dim), + nn.GELU(), + nn.Linear(shared_dim, shared_dim), + ) + self.satellite_proj = nn.Sequential( + nn.Linear(fused_dim, shared_dim), + nn.GELU(), + nn.Linear(shared_dim, shared_dim), + ) + for proj in (self.drone_proj, self.satellite_proj): + for m in proj.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # --- Learnable temperature --- + self.log_temperature = nn.Parameter( + torch.tensor(temperature_init).log() + ) + + @property + def temperature(self) -> torch.Tensor: + return self.log_temperature.exp().clamp(min=0.01, max=100.0) + + def encode_drone( + self, + drone_images: torch.Tensor, + drone_tokens: torch.Tensor, + ) -> torch.Tensor: + """Дрон (картинка + текст) → слитый, нормализованный эмбеддинг. + + Args: + drone_images: [B, 3, H, W] — картинки дрона. + drone_tokens: [B, 248] — токенизированный текст описания дрона. + """ + image_feat = self.drone_image_encoder(drone_images) + text_feat = self.text_encoder(drone_tokens) + fused = self.drone_fusion(image_feat, text_feat) + projected = self.drone_proj(fused) + return F.normalize(projected, dim=-1) + + def encode_satellite( + self, + satellite_images: torch.Tensor, + satellite_tokens: torch.Tensor, + ) -> torch.Tensor: + """Спутник (картинка + текст) → слитый, нормализованный эмбеддинг. + + Args: + satellite_images: [B, 3, H, W] — спутниковые снимки галереи. + satellite_tokens: [B, 248] — токенизированный текст описания + спутникового снимка. + """ + image_feat = self.satellite_image_encoder(satellite_images) + text_feat = self.text_encoder(satellite_tokens) + fused = self.satellite_fusion(image_feat, text_feat) + projected = self.satellite_proj(fused) + return F.normalize(projected, dim=-1) + + def forward( + self, + drone_images: torch.Tensor, + drone_tokens: torch.Tensor, + satellite_images: torch.Tensor, + satellite_tokens: torch.Tensor, + ) -> dict[str, torch.Tensor]: + """Прямой проход: дрон (слитый) vs спутник (слитый). + + Returns: + dict: + drone_emb: [B, shared_dim] + satellite_emb: [B, shared_dim] + logits: [B, B] — drone_emb @ satellite_emb.T / temperature + temperature: скаляр + """ + drone_emb = self.encode_drone(drone_images, drone_tokens) + satellite_emb = self.encode_satellite(satellite_images, satellite_tokens) + + # logits[i][j] = cosine_sim(drone_i, satellite_j) / temperature + logits = drone_emb @ satellite_emb.T / self.temperature + + return { + "drone_emb": drone_emb, + "satellite_emb": satellite_emb, + "logits": logits, + "temperature": self.temperature, + } + + +# --------------------------------------------------------------------------- +# Управление заморозкой параметров +# --------------------------------------------------------------------------- + +def freeze_encoder(encoder: nn.Module) -> int: + """Заморозить все параметры энкодера. Возвращает кол-во замороженных.""" + count = 0 + for p in encoder.parameters(): + p.requires_grad = False + count += 1 + return count + + +def get_trainable_params(model: DualEncoder) -> dict[str, dict]: + """Подсчёт обучаемых vs замороженных параметров по компонентам.""" + components = { + "text_encoder": model.text_encoder, + "drone_image_encoder": model.drone_image_encoder, + "satellite_image_encoder": model.satellite_image_encoder, + "drone_fusion": model.drone_fusion, + "satellite_fusion": model.satellite_fusion, + "drone_proj": model.drone_proj, + "satellite_proj": model.satellite_proj, + } + result = {} + for name, module in components.items(): + trainable = sum(p.numel() for p in module.parameters() if p.requires_grad) + total = sum(p.numel() for p in module.parameters()) + result[name] = {"trainable": trainable, "total": total} + result["temperature"] = {"trainable": 1, "total": 1} + return result + + +# --------------------------------------------------------------------------- +# Фабричная функция: собрать всё из чекпоинтов +# --------------------------------------------------------------------------- + +def build_dual_encoder( + dgtrs_checkpoint: str | Path, + stripnet_checkpoint: str | Path, + fused_dim: int = 512, + shared_dim: int = 512, + freeze_text: bool = True, + freeze_image_backbone: bool = True, + inject_mona: bool = True, + mona_bottleneck: int = 64, + device: str = "cpu", +) -> DualEncoder: + """Собрать симметричный DualEncoder из чекпоинтов. + + Текстовый энкодер — ОДИН общий экземпляр, используется для текста + и дрона, и спутника (это одна и та же языковая модель; разный текст + подаётся на вход одного и того же набора весов). + + Визуальные энкодеры — ДВА отдельных экземпляра StripNetEncoder, по + одному на домен (дрон / спутник), с отдельными Conv-MONA адаптерами. + + Слияние — ДВА отдельных TextFusionMLP (drone_fusion, satellite_fusion), + с отдельными весами, так как оптимальное смешивание "картинка+текст" + индивидуально для каждого визуального домена. + + Обучаются: оба TextFusionMLP, обе проекционные головы, оба набора + Conv-MONA адаптеров, temperature. + Замораживаются: текстовый энкодер целиком, backbone StripNet в обоих + визуальных энкодерах. + + Args: + dgtrs_checkpoint: Путь к чекпоинту DGTRS-CLIP. + stripnet_checkpoint: Путь к pretrained StripNet (общий чекпоинт + для инициализации обоих визуальных энкодеров). + fused_dim: Размерность вектора после слияния. + shared_dim: Размерность общего пространства сравнения. + freeze_text: Заморозить текстовый энкодер. + freeze_image_backbone: Заморозить backbone StripNet (в обоих). + inject_mona: Инжектировать Conv-MONA в оба энкодера. + mona_bottleneck: Bottleneck dim для Conv-MONA. + device: Устройство. + """ + # 1. Текстовый энкодер (общий, заморожен, используется для обеих сторон) + text_encoder = load_dgtrs_text_encoder(dgtrs_checkpoint, device="cpu") + + # 2. Два отдельных визуальных энкодера + drone_image_encoder = StripNetEncoder( + checkpoint_path=stripnet_checkpoint, + out_dim=1024, + load_pretrained=True, + ) + satellite_image_encoder = StripNetEncoder( + checkpoint_path=stripnet_checkpoint, + out_dim=1024, + load_pretrained=True, + ) + + # 3. Заморозка + if freeze_text: + n = freeze_encoder(text_encoder) + LOGGER.info("❄️ Text encoder frozen: %d parameters (shared for drone+satellite)", n) + + if freeze_image_backbone: + n1 = freeze_encoder(drone_image_encoder.backbone) + n2 = freeze_encoder(satellite_image_encoder.backbone) + LOGGER.info( + "❄️ StripNet backbones frozen: drone=%d, satellite=%d parameters", + n1, n2, + ) + + # 4. Conv-MONA адаптеры — отдельно для каждого визуального энкодера + if inject_mona: + from src.models.stripnet.conv_mona import inject_conv_mona_into_stripnet + inject_conv_mona_into_stripnet( + drone_image_encoder.backbone, + bottleneck=mona_bottleneck, + last_n_stages=2, + ) + inject_conv_mona_into_stripnet( + satellite_image_encoder.backbone, + bottleneck=mona_bottleneck, + last_n_stages=2, + ) + + # 5. Собрать DualEncoder (TextFusionMLP создаются внутри __init__) + model = DualEncoder( + text_encoder=text_encoder, + drone_image_encoder=drone_image_encoder, + satellite_image_encoder=satellite_image_encoder, + fused_dim=fused_dim, + shared_dim=shared_dim, + ) + model = model.to(device) + + # 6. Логирование + params = get_trainable_params(model) + total_trainable = sum(v["trainable"] for v in params.values()) + total_all = sum(v["total"] for v in params.values()) + LOGGER.info( + "🔗 DualEncoder (symmetric fusion) built: %s trainable / %s total params", + f"{total_trainable:,}", f"{total_all:,}", + ) + for name, info in params.items(): + LOGGER.info( + " %-26s %10s trainable / %10s total", + name, f"{info['trainable']:,}", f"{info['total']:,}", + ) + + return model \ No newline at end of file diff --git a/src/models/stripnet/__init__.py b/src/models/stripnet/__init__.py new file mode 100644 index 0000000..b486594 --- /dev/null +++ b/src/models/stripnet/__init__.py @@ -0,0 +1 @@ +from src.models.stripnet.model import StripNet, get_stripnet_small, load_stripnet_small_pretrained diff --git a/src/models/stripnet/conv_mona.py b/src/models/stripnet/conv_mona.py new file mode 100644 index 0000000..c13c0eb --- /dev/null +++ b/src/models/stripnet/conv_mona.py @@ -0,0 +1,106 @@ +"""Conv-MONA: 2D adaptation of MONA (CVPR 2025) for hierarchical CNN backbones. + +MONA paper applies sequence-form adapters after MSA / MLP in ViT blocks. Here we +mirror that idea in [B, C, H, W] form: BN → 1×1 Down(C→bn) → multi-scale DWConv +{3,5,7} mean → +residual → GELU → 1×1 Up(bn→C). Layer scale (γ) channel-wise, +init 1e-6 for near-identity start. Two adapters per StripNet Block: post-attn +and post-mlp. +""" +from __future__ import annotations + +import logging +import torch +import torch.nn as nn + +from src.models.stripnet.model import StripNet, Block + +LOGGER = logging.getLogger("caption_test.stripnet.adapters") + + +class ConvMona(nn.Module): + """Single Conv-MONA adapter. + + Args: + dim: input channel dim. + bottleneck: bottleneck channel dim (e.g. 64). + gamma_init: layer-scale init value (1e-6 for near-identity at start). + """ + + def __init__(self, dim: int, bottleneck: int = 64, gamma_init: float = 1e-6) -> None: + super().__init__() + self.norm = nn.BatchNorm2d(dim) + self.down = nn.Conv2d(dim, bottleneck, kernel_size=1, bias=True) + self.dw3 = nn.Conv2d(bottleneck, bottleneck, kernel_size=3, padding=1, groups=bottleneck, bias=True) + self.dw5 = nn.Conv2d(bottleneck, bottleneck, kernel_size=5, padding=2, groups=bottleneck, bias=True) + self.dw7 = nn.Conv2d(bottleneck, bottleneck, kernel_size=7, padding=3, groups=bottleneck, bias=True) + self.act = nn.GELU() + self.up = nn.Conv2d(bottleneck, dim, kernel_size=1, bias=True) + # Channel-wise layer scale (γ), broadcast across H, W. + self.gamma = nn.Parameter(gamma_init * torch.ones(dim), requires_grad=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.norm(x) + h = self.down(h) + h = (self.dw3(h) + self.dw5(h) + self.dw7(h)) / 3.0 + h + h = self.act(h) + h = self.up(h) + return self.gamma.view(1, -1, 1, 1) * h + + +def _patched_block_forward(block: Block, mona_attn: ConvMona, mona_mlp: ConvMona): + """Closure that wraps a Block.forward with two Conv-MONA residuals.""" + orig_attn = block.attn + orig_mlp = block.mlp + orig_norm1 = block.norm1 + orig_norm2 = block.norm2 + orig_drop = block.drop_path + ls1 = block.layer_scale_1 + ls2 = block.layer_scale_2 + + def forward(x: torch.Tensor) -> torch.Tensor: + x = x + orig_drop(ls1.unsqueeze(-1).unsqueeze(-1) * orig_attn(orig_norm1(x))) + mona_attn(x) + x = x + orig_drop(ls2.unsqueeze(-1).unsqueeze(-1) * orig_mlp(orig_norm2(x))) + mona_mlp(x) + return x + + return forward + + +def inject_conv_mona_into_stripnet( + model: StripNet, + bottleneck: int = 64, + last_n_stages: int = 2, + use_bf16: bool = False, +) -> int: + """Inject Conv-MONA adapters into the deepest `last_n_stages` of StripNet. + + Each Block in the targeted stages gets two adapters (post-attn, post-mlp). + Returns the number of adapters injected. + + Stages are 1-indexed in StripNet (block1..block4). With `last_n_stages=2` + we adapt block3 and block4 — the deepest, semantically richest features. + """ + n_stages = model.num_stages + target_stages = list(range(max(1, n_stages - last_n_stages + 1), n_stages + 1)) + n_added = 0 + + for stage_idx in target_stages: + blocks: nn.ModuleList = getattr(model, f"block{stage_idx}") + dim = model.embed_dims[stage_idx - 1] + for blk_idx, block in enumerate(blocks): + mona_a = ConvMona(dim=dim, bottleneck=bottleneck) + mona_m = ConvMona(dim=dim, bottleneck=bottleneck) + if use_bf16: + mona_a.to(dtype=torch.bfloat16) + mona_m.to(dtype=torch.bfloat16) + # Register as submodules so they get moved by .to(device) / .train() etc. + block.add_module(f"mona_attn", mona_a) + block.add_module(f"mona_mlp", mona_m) + block.forward = _patched_block_forward(block, mona_a, mona_m) + n_added += 2 + + n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + LOGGER.info( + "🔧 Conv-MONA injected: %d adapters in stages %s, %d trainable params (bottleneck=%d)", + n_added, target_stages, n_trainable, bottleneck, + ) + return n_added diff --git a/src/models/stripnet/model.py b/src/models/stripnet/model.py new file mode 100644 index 0000000..edf4fc9 --- /dev/null +++ b/src/models/stripnet/model.py @@ -0,0 +1,263 @@ +"""StripNet (small) backbone — adapted from Strip-R-CNN (HVision-NKU). + +Self-contained: no external utils. State-dict naming follows the upstream +ImageNet-pretrained checkpoint (`conv_spatial1/2` for the strip kernels). +""" +from __future__ import annotations + +import logging +import math +from functools import partial +from pathlib import Path +from typing import List + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +LOGGER = logging.getLogger("caption_test.stripnet") + + +def _to_2tuple(x): + if isinstance(x, (tuple, list)): + return tuple(x) + return (x, x) + + +def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: + if drop_prob == 0.0 or not training: + return x + keep = 1.0 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + rand = x.new_empty(shape).bernoulli_(keep) + if keep > 0: + rand.div_(keep) + return x * rand + + +class DropPath(nn.Module): + def __init__(self, p: float = 0.0) -> None: + super().__init__() + self.p = p + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return drop_path(x, self.p, self.training) + + +class DWConv(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.dwconv(x) + + +class Mlp(nn.Module): + def __init__(self, in_features: int, hidden_features: int, drop: float = 0.0) -> None: + super().__init__() + self.fc1 = nn.Conv2d(in_features, hidden_features, 1) + self.dwconv = DWConv(hidden_features) + self.act = nn.GELU() + self.fc2 = nn.Conv2d(hidden_features, in_features, 1) + self.drop = nn.Dropout(drop) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc1(x) + x = self.dwconv(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class StripGatingUnit(nn.Module): + """Strip spatial gating: 5x5 DWConv -> (1, k2) -> (k2, 1) -> 1x1 -> gate.""" + + def __init__(self, dim: int, k1: int, k2: int) -> None: + super().__init__() + self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim) + # Names match upstream pretrained checkpoint: conv_spatial1 / conv_spatial2. + self.conv_spatial1 = nn.Conv2d(dim, dim, kernel_size=(k1, k2), stride=1, + padding=(k1 // 2, k2 // 2), groups=dim) + self.conv_spatial2 = nn.Conv2d(dim, dim, kernel_size=(k2, k1), stride=1, + padding=(k2 // 2, k1 // 2), groups=dim) + self.conv1 = nn.Conv2d(dim, dim, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + attn = self.conv0(x) + attn = self.conv_spatial1(attn) + attn = self.conv_spatial2(attn) + attn = self.conv1(attn) + return x * attn + + +class StripAttention(nn.Module): + def __init__(self, dim: int, k1: int, k2: int) -> None: + super().__init__() + self.proj_1 = nn.Conv2d(dim, dim, 1) + self.activation = nn.GELU() + self.spatial_gating_unit = StripGatingUnit(dim, k1, k2) + self.proj_2 = nn.Conv2d(dim, dim, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.proj_1(x) + x = self.activation(x) + x = self.spatial_gating_unit(x) + x = self.proj_2(x) + return x + residual + + +class Block(nn.Module): + def __init__(self, dim: int, mlp_ratio: float, k1: int, k2: int, drop: float, drop_path: float) -> None: + super().__init__() + self.norm1 = nn.BatchNorm2d(dim) + self.norm2 = nn.BatchNorm2d(dim) + self.attn = StripAttention(dim, k1, k2) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.mlp = Mlp(dim, int(dim * mlp_ratio), drop=drop) + ls_init = 1e-2 + self.layer_scale_1 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True) + self.layer_scale_2 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.drop_path( + self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x)) + ) + x = x + self.drop_path( + self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x)) + ) + return x + + +class OverlapPatchEmbed(nn.Module): + def __init__(self, patch_size: int, stride: int, in_chans: int, embed_dim: int) -> None: + super().__init__() + ph, pw = _to_2tuple(patch_size) + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(ph, pw), stride=stride, + padding=(ph // 2, pw // 2)) + self.norm = nn.BatchNorm2d(embed_dim) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]: + x = self.proj(x) + _, _, H, W = x.shape + x = self.norm(x) + return x, H, W + + +class StripNet(nn.Module): + """Strip-R-CNN backbone: 4-stage hierarchical CNN with strip-shaped DWConv attention. + + Output: list of [B, C_i, H/s_i, W/s_i] per stage. Use `forward_last_features` for + the deepest stage only. + """ + + def __init__( + self, + embed_dims: List[int] = [64, 128, 320, 512], + mlp_ratios: List[int] = [8, 8, 4, 4], + k1s: List[int] = [1, 1, 1, 1], + k2s: List[int] = [19, 19, 19, 19], + depths: List[int] = [2, 2, 4, 2], + drop_rate: float = 0.1, + drop_path_rate: float = 0.15, + in_chans: int = 3, + ) -> None: + super().__init__() + self.depths = depths + self.num_stages = len(depths) + self.embed_dims = embed_dims + self.gradient_checkpointing = False + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(self.num_stages): + patch_embed = OverlapPatchEmbed( + patch_size=7 if i == 0 else 3, + stride=4 if i == 0 else 2, + in_chans=in_chans if i == 0 else embed_dims[i - 1], + embed_dim=embed_dims[i], + ) + block = nn.ModuleList([ + Block(dim=embed_dims[i], mlp_ratio=mlp_ratios[i], k1=k1s[i], k2=k2s[i], + drop=drop_rate, drop_path=dpr[cur + j]) + for j in range(depths[i]) + ]) + norm = nn.LayerNorm(embed_dims[i], eps=1e-6) + cur += depths[i] + setattr(self, f"patch_embed{i + 1}", patch_embed) + setattr(self, f"block{i + 1}", block) + setattr(self, f"norm{i + 1}", norm) + + def set_gradient_checkpointing(self, enable: bool = True) -> None: + """Recompute Strip blocks during backward to lower VRAM (full fine-tune).""" + self.gradient_checkpointing = bool(enable) + + def forward_features(self, x: torch.Tensor) -> List[torch.Tensor]: + B = x.shape[0] + outs: List[torch.Tensor] = [] + for i in range(self.num_stages): + patch_embed = getattr(self, f"patch_embed{i + 1}") + block = getattr(self, f"block{i + 1}") + norm = getattr(self, f"norm{i + 1}") + x, H, W = patch_embed(x) + for blk in block: + if self.gradient_checkpointing and self.training: + x = torch.utils.checkpoint.checkpoint( + blk, x, use_reentrant=False, + ) + else: + x = blk(x) + x = x.flatten(2).transpose(1, 2) + x = norm(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + return outs + + def forward_last_features(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_features(x)[-1] + + +def get_stripnet_small() -> StripNet: + return StripNet( + embed_dims=[64, 128, 320, 512], + mlp_ratios=[8, 8, 4, 4], + k1s=[1, 1, 1, 1], + k2s=[19, 19, 19, 19], + depths=[2, 2, 4, 2], + drop_rate=0.1, + drop_path_rate=0.15, + ) + + +def load_stripnet_small_pretrained(checkpoint_path: str | Path) -> StripNet: + """Build StripNet-small and load ImageNet-pretrained weights. + + Strips the classification `head.*` keys. Tolerates missing/extra keys + (norm{N}.* are LayerNorm here vs BatchNorm in some forks — we keep LN). + """ + LOGGER.info("📐 Loading StripNet-small from %s", checkpoint_path) + model = get_stripnet_small() + raw = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + state = raw.get("state_dict", raw) if isinstance(raw, dict) else raw + + # Drop classification head + the BN-form norm{N} keys if present (we use LN here). + drop_prefixes = ("head.",) + cleaned = {k: v for k, v in state.items() if not any(k.startswith(p) for p in drop_prefixes)} + + # The pretrained checkpoint stores norm{N} as BatchNorm2d (running_mean/var/num_batches_tracked). + # Our code uses LayerNorm at this position. Strip BN running stats if found; copy weight/bias. + for n in (1, 2, 3, 4): + for suffix in ("running_mean", "running_var", "num_batches_tracked"): + cleaned.pop(f"norm{n}.{suffix}", None) + + missing, unexpected = model.load_state_dict(cleaned, strict=False) + if missing: + LOGGER.info("StripNet missing keys (expected for newly-init layers): %d", len(missing)) + if unexpected: + LOGGER.info("StripNet unexpected keys (ignored): %d", len(unexpected)) + LOGGER.info("📐 StripNet-small loaded: %d params", sum(p.numel() for p in model.parameters())) + return model diff --git a/src/models/stripnet_encoder.py b/src/models/stripnet_encoder.py new file mode 100644 index 0000000..5f22331 --- /dev/null +++ b/src/models/stripnet_encoder.py @@ -0,0 +1,61 @@ +"""StripNet image encoder wrapper for the caption-test pipeline. + +Exposes the same interface as DINOv3ViT: `forward(images) -> [B, embed_dim]`. +StripNet's deepest stage produces [B, 512, H/32, W/32]; we apply global average +pooling (GAP) and project to the target retrieval dimension via Linear(512→1024) +to match DINOv3 native dim and keep GatedFusion unchanged. +""" +from __future__ import annotations + +import logging +import torch +import torch.nn as nn + +from src.models.stripnet import StripNet, get_stripnet_small, load_stripnet_small_pretrained + +LOGGER = logging.getLogger("caption_test.stripnet_encoder") + + +class StripNetEncoder(nn.Module): + """StripNet-small + GAP + projection to `out_dim`. + + Frozen backbone (BatchNorm in eval mode); only the projection head and + any injected Conv-MONA adapters are trainable. + """ + + LAST_STAGE_DIM = 512 # StripNet-small last stage embed dim + + def __init__( + self, + checkpoint_path: str, + out_dim: int = 1024, + *, + load_pretrained: bool = True, + ) -> None: + super().__init__() + self.out_dim = out_dim + if load_pretrained: + self.backbone: StripNet = load_stripnet_small_pretrained(checkpoint_path) + else: + LOGGER.info("StripNet-small: random init (stripnet_load_pretrained=False)") + self.backbone = get_stripnet_small() + self.pool = nn.AdaptiveAvgPool2d(1) + self.projection = nn.Linear(self.LAST_STAGE_DIM, out_dim) + nn.init.trunc_normal_(self.projection.weight, std=0.02) + nn.init.zeros_(self.projection.bias) + + def set_gradient_checkpointing(self, enable: bool = True) -> None: + self.backbone.set_gradient_checkpointing(enable) + + def train(self, mode: bool = True): + """Override: keep frozen backbone in eval mode (BN running stats stable).""" + super().train(mode) + # Frozen backbone always in eval; trainable adapters/projection follow `mode`. + if not any(p.requires_grad for p in self.backbone.parameters()): + self.backbone.eval() + return self + + def forward(self, images: torch.Tensor) -> torch.Tensor: + feat = self.backbone.forward_last_features(images) # [B, 512, H/32, W/32] + pooled = self.pool(feat).flatten(1) # [B, 512] + return self.projection(pooled) # [B, out_dim] \ No newline at end of file diff --git a/train.py b/train.py new file mode 100644 index 0000000..d83acd0 --- /dev/null +++ b/train.py @@ -0,0 +1,455 @@ +"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура). + +ВАЖНОЕ ИЗМЕНЕНИЕ: модель теперь принимает ЧЕТЫРЕ входа вместо трёх: + drone_images, drone_tokens → сливаются в drone_emb (TextFusionMLP) + satellite_images, satellite_tokens → сливаются в satellite_emb (TextFusionMLP) + drone_emb и satellite_emb сравниваются между собой → cosine similarity → InfoNCE + +И дрон, и спутник имеют собственное текстовое описание; слияние +(картинка+текст) происходит СИММЕТРИЧНО на обеих сторонах, каждая — +со своим экземпляром TextFusionMLP (веса не общие, т.к. визуальные +домены различаются). + +Оптимизировано под RTX 4090 (24 GB VRAM, Ada Lovelace): BF16 AMP, +micro_batch=64 по умолчанию (effective batch = micro_batch при отсутствии +gradient accumulation). + +Использование: + python train.py \\ + --data_root /path/to/GTA-UAV \\ + --descriptions_path /path/to/descriptions.json \\ + --text_levels level1 \\ + --dgtrs_checkpoint /path/to/DGTRS-CLIP-ViT-B-16 \\ + --stripnet_checkpoint /path/to/stripnet_small.pth \\ + --epochs 50 --batch_size 64 --micro_batch_size 64 --bf16 +""" +from __future__ import annotations + +import argparse +import json +import logging +import sys +import time +from contextlib import nullcontext +from pathlib import Path + +import torch +import torch.nn as nn +from torch.optim import AdamW +from torch.optim.lr_scheduler import CosineAnnealingLR + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from src.models.dual_encoder import build_dual_encoder, get_trainable_params +from src.losses import InfoNCELoss +from src.metrics import compute_retrieval_metrics, format_metrics +from src.data.gta_uav import build_dataloaders + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(name)s %(levelname)s %(message)s", +) +LOGGER = logging.getLogger("cvgl.train") + + +# --------------------------------------------------------------------------- +# GPU info +# --------------------------------------------------------------------------- + +def log_gpu_info(device: torch.device) -> None: + if device.type != "cuda": + return + name = torch.cuda.get_device_name(device) + total = torch.cuda.get_device_properties(device).total_memory / 1024**3 + LOGGER.info("🖥️ GPU: %s (%.1f GB VRAM)", name, total) + cap = torch.cuda.get_device_capability(device) + bf16_ok = cap[0] >= 8 + LOGGER.info( + " Compute capability: %d.%d | BF16: %s", + cap[0], cap[1], "✅ supported" if bf16_ok else "❌ use FP16 instead", + ) + + +def log_vram_usage(prefix: str = "") -> None: + if not torch.cuda.is_available(): + return + allocated = torch.cuda.memory_allocated() / 1024**3 + reserved = torch.cuda.memory_reserved() / 1024**3 + LOGGER.info(" %sVRAM: %.2f GB allocated / %.2f GB reserved", prefix, allocated, reserved) + + +def get_amp_context(use_bf16: bool, use_fp16: bool, device: torch.device): + if use_bf16: + return torch.autocast(device_type=device.type, dtype=torch.bfloat16) + elif use_fp16: + return torch.autocast(device_type=device.type, dtype=torch.float16) + else: + return nullcontext() + + +# --------------------------------------------------------------------------- +# Evaluation +# --------------------------------------------------------------------------- + +@torch.no_grad() +def evaluate( + model: nn.Module, + test_loader, + device: torch.device, + amp_ctx, +) -> dict[str, float]: + """Прогнать test set: drone (fused) vs satellite (fused) — симметрично.""" + model.eval() + + all_drone_emb = [] + all_satellite_emb = [] + + for batch in test_loader: + drone_images = batch["drone_image"].to(device) + drone_tokens = batch["drone_tokens"].to(device) + satellite_images = batch["satellite_image"].to(device) + satellite_tokens = batch["satellite_tokens"].to(device) + + with amp_ctx: + drone_emb = model.encode_drone(drone_images, drone_tokens) + satellite_emb = model.encode_satellite(satellite_images, satellite_tokens) + + all_drone_emb.append(drone_emb.float().cpu()) + all_satellite_emb.append(satellite_emb.float().cpu()) + + all_drone_emb = torch.cat(all_drone_emb, dim=0) + all_satellite_emb = torch.cat(all_satellite_emb, dim=0) + + metrics = compute_retrieval_metrics( + all_drone_emb, all_satellite_emb, ks=[1, 5, 10], + ) + + model.train() + return metrics + + +# --------------------------------------------------------------------------- +# Training loop +# --------------------------------------------------------------------------- + +def train_one_epoch( + model: nn.Module, + train_loader, + criterion: nn.Module, + optimizer: torch.optim.Optimizer, + device: torch.device, + epoch: int, + grad_accumulate_steps: int = 1, + max_grad_norm: float = 1.0, + amp_ctx=None, + scaler: torch.amp.GradScaler | None = None, +) -> dict[str, float]: + model.train() + if amp_ctx is None: + amp_ctx = nullcontext() + + total_loss = 0.0 + total_acc_t2i = 0.0 + total_acc_i2t = 0.0 + n_steps = 0 + + optimizer.zero_grad() + + for batch_idx, batch in enumerate(train_loader): + drone_images = batch["drone_image"].to(device) + drone_tokens = batch["drone_tokens"].to(device) + satellite_images = batch["satellite_image"].to(device) + satellite_tokens = batch["satellite_tokens"].to(device) + + with amp_ctx: + outputs = model(drone_images, drone_tokens, satellite_images, satellite_tokens) + logits = outputs["logits"] + loss_dict = criterion(logits) + loss = loss_dict["loss"] + scaled_loss = loss / grad_accumulate_steps + + if scaler is not None: + scaler.scale(scaled_loss).backward() + else: + scaled_loss.backward() + + total_loss += loss.item() + total_acc_t2i += loss_dict["acc_t2i"].item() + total_acc_i2t += loss_dict["acc_i2t"].item() + n_steps += 1 + + if (batch_idx + 1) % grad_accumulate_steps == 0: + if scaler is not None: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_( + [p for p in model.parameters() if p.requires_grad], + max_grad_norm, + ) + scaler.step(optimizer) + scaler.update() + else: + torch.nn.utils.clip_grad_norm_( + [p for p in model.parameters() if p.requires_grad], + max_grad_norm, + ) + optimizer.step() + optimizer.zero_grad() + + if (batch_idx + 1) % 50 == 0: + LOGGER.info( + " [Epoch %d] Step %d/%d | loss=%.4f | acc_d2s=%.3f | " + "acc_s2d=%.3f | τ=%.4f", + epoch, batch_idx + 1, len(train_loader), + loss.item(), + loss_dict["acc_t2i"].item(), + loss_dict["acc_i2t"].item(), + outputs["temperature"].item(), + ) + + if len(train_loader) % grad_accumulate_steps != 0: + if scaler is not None: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_( + [p for p in model.parameters() if p.requires_grad], + max_grad_norm, + ) + scaler.step(optimizer) + scaler.update() + else: + torch.nn.utils.clip_grad_norm_( + [p for p in model.parameters() if p.requires_grad], + max_grad_norm, + ) + optimizer.step() + optimizer.zero_grad() + + return { + "loss": total_loss / max(n_steps, 1), + "acc_t2i": total_acc_t2i / max(n_steps, 1), + "acc_i2t": total_acc_i2t / max(n_steps, 1), + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(args): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + LOGGER.info("🚀 Device: %s", device) + log_gpu_info(device) + + exp_name = f"exp_{'-'.join(args.text_levels)}_ep{args.epochs}_bs{args.batch_size}" + output_dir = Path(args.output_dir) / exp_name + output_dir.mkdir(parents=True, exist_ok=True) + LOGGER.info("📁 Output: %s", output_dir) + + config = vars(args) + with open(output_dir / "config.json", "w") as f: + json.dump(config, f, indent=2, default=str) + + train_loader, test_loader = build_dataloaders( + data_root=args.data_root, + descriptions_path=args.descriptions_path, + text_levels=args.text_levels, + train_meta=args.train_meta, + test_meta=args.test_meta, + batch_size=args.micro_batch_size, + num_workers=args.num_workers, + image_size=args.image_size, + ) + + model = build_dual_encoder( + dgtrs_checkpoint=args.dgtrs_checkpoint, + stripnet_checkpoint=args.stripnet_checkpoint, + fused_dim=args.fused_dim, + shared_dim=args.shared_dim, + freeze_text=True, + freeze_image_backbone=True, + inject_mona=args.inject_mona, + mona_bottleneck=args.mona_bottleneck, + device=str(device), + ) + log_vram_usage("After model load: ") + + if args.compile and hasattr(torch, "compile"): + LOGGER.info("⚡ Compiling model with torch.compile (mode=%s)", args.compile_mode) + model = torch.compile(model, mode=args.compile_mode) + + use_bf16 = args.bf16 and device.type == "cuda" + use_fp16 = args.fp16 and device.type == "cuda" and not use_bf16 + amp_ctx = get_amp_context(use_bf16, use_fp16, device) + scaler = torch.amp.GradScaler("cuda") if use_fp16 else None + + precision_str = "BF16" if use_bf16 else ("FP16" if use_fp16 else "FP32") + LOGGER.info("🔢 Precision: %s", precision_str) + + trainable_params = [p for p in model.parameters() if p.requires_grad] + optimizer = AdamW( + trainable_params, + lr=args.lr, + weight_decay=args.weight_decay, + betas=(0.9, 0.98), + ) + + scheduler = CosineAnnealingLR( + optimizer, + T_max=args.epochs, + eta_min=args.lr * 0.01, + ) + + criterion = InfoNCELoss(label_smoothing=args.label_smoothing) + + grad_accumulate_steps = max(1, args.batch_size // args.micro_batch_size) + LOGGER.info( + "⚙️ Effective batch=%d (micro=%d × accumulate=%d)", + args.batch_size, args.micro_batch_size, grad_accumulate_steps, + ) + + start_epoch = 1 + if args.resume and (output_dir / "latest_model.pth").exists(): + ckpt = torch.load(output_dir / "latest_model.pth", map_location=device) + model.load_state_dict(ckpt["model_state_dict"]) + optimizer.load_state_dict(ckpt["optimizer_state_dict"]) + if "scheduler_state_dict" in ckpt: + scheduler.load_state_dict(ckpt["scheduler_state_dict"]) + start_epoch = ckpt["epoch"] + 1 + LOGGER.info("🔄 Resumed from epoch %d", start_epoch - 1) + + best_recall1 = 0.0 + history = [] + + history_path = output_dir / "history.json" + if args.resume and history_path.exists(): + with open(history_path) as f: + history = json.load(f) + + log_vram_usage("Before training: ") + + for epoch in range(start_epoch, args.epochs + 1): + t0 = time.time() + + train_metrics = train_one_epoch( + model, train_loader, criterion, optimizer, + device, epoch, + grad_accumulate_steps=grad_accumulate_steps, + max_grad_norm=args.max_grad_norm, + amp_ctx=amp_ctx, + scaler=scaler, + ) + + scheduler.step() + + if epoch % args.eval_every == 0 or epoch == args.epochs: + eval_metrics = evaluate(model, test_loader, device, amp_ctx) + else: + eval_metrics = {} + + elapsed = time.time() - t0 + + LOGGER.info( + "📈 Epoch %d/%d (%.0fs) | loss=%.4f | " + "R@1=%.3f R@5=%.3f R@10=%.3f | AP=%.3f", + epoch, args.epochs, elapsed, + train_metrics["loss"], + eval_metrics.get("recall@1", 0), + eval_metrics.get("recall@5", 0), + eval_metrics.get("recall@10", 0), + eval_metrics.get("AP", 0), + ) + + if epoch == 1: + log_vram_usage("After first epoch: ") + + record = { + "epoch": epoch, + "lr": scheduler.get_last_lr()[0], + **{f"train_{k}": v for k, v in train_metrics.items()}, + **{f"eval_{k}": v for k, v in eval_metrics.items()}, + "elapsed_s": elapsed, + } + history.append(record) + + recall1 = eval_metrics.get("recall@1", 0) + if recall1 > best_recall1: + best_recall1 = recall1 + torch.save( + { + "epoch": epoch, + "model_state_dict": model.state_dict(), + "eval_metrics": eval_metrics, + "config": config, + }, + output_dir / "best_model.pth", + ) + LOGGER.info("💾 New best model (R@1=%.4f)", recall1) + + torch.save( + { + "epoch": epoch, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict(), + }, + output_dir / "latest_model.pth", + ) + + with open(history_path, "w") as f: + json.dump(history, f, indent=2) + + LOGGER.info("=" * 60) + LOGGER.info("🏁 Training complete. Best R@1: %.4f", best_recall1) + LOGGER.info("📁 Results: %s", output_dir) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def parse_args(): + p = argparse.ArgumentParser(description="Train CVGL fusion dual-encoder on GTA-UAV") + + # Data + p.add_argument("--data_root", type=str, required=True) + p.add_argument("--descriptions_path", type=str, required=True) + p.add_argument("--text_levels", nargs="+", default=["level1"]) + p.add_argument("--train_meta", default="cross-area-drone2sate-train.json") + p.add_argument("--test_meta", default="cross-area-drone2sate-test.json") + p.add_argument("--image_size", type=int, default=384) + p.add_argument("--num_workers", type=int, default=8) + + # Model + p.add_argument("--dgtrs_checkpoint", type=str, required=True) + p.add_argument("--stripnet_checkpoint", type=str, required=True) + p.add_argument("--fused_dim", type=int, default=512, + help="Размерность вектора после слияния картинки и текста") + p.add_argument("--shared_dim", type=int, default=512) + p.add_argument("--inject_mona", action="store_true", default=True) + p.add_argument("--mona_bottleneck", type=int, default=64) + + # Training + p.add_argument("--epochs", type=int, default=50) + p.add_argument("--batch_size", type=int, default=64) + p.add_argument("--micro_batch_size", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-4) + p.add_argument("--weight_decay", type=float, default=0.01) + p.add_argument("--max_grad_norm", type=float, default=1.0) + p.add_argument("--label_smoothing", type=float, default=0.0) + p.add_argument("--eval_every", type=int, default=1) + + # Performance + p.add_argument("--bf16", action="store_true", default=True) + p.add_argument("--fp16", action="store_true", default=False) + p.add_argument("--compile", action="store_true", default=False) + p.add_argument("--compile_mode", default="reduce-overhead", + choices=["default", "reduce-overhead", "max-autotune"]) + + # Resume / output + p.add_argument("--resume", action="store_true", default=False) + p.add_argument("--output_dir", type=str, default="outputs") + + return p.parse_args() + + +if __name__ == "__main__": + main(parse_args()) \ No newline at end of file