# Multi-Modal Annotation Pipeline

Автоматическая генерация 4 модальностей (depth, edges, segmentation, canopy height) из RGB-изображений аэрофотосъёмки. Поддерживает три датасета:

| Модальность | Модель | Выход | Скорость |
|:---|:---|:---|:---|
| **Depth** | DA3-LARGE-1.1 (411M) | grayscale [256x256] | 18.4 img/s |
| **Edges** | Sobel из depth (CPU) | grayscale [256x256] | 419.6 img/s |
| **Segmentation** | SegEarth-OV3 (SAM 3.1) | class IDs [256x256] | ~3.5 img/s |
| **CHMv2** | DINOv3-ViTL16 (337M, FP32) | grayscale [256x256] | 31.7 img/s |
| **Consolidate** | SafeTensors (CPU) | `.safetensors` per image | ~5000 img/s |

| Датасет | Изображения | Сегм. классы | Скрипт |
|:---|:---|:---|:---|
| **World-UAV** | 973K (486K DB + 486K query) | 17 (unified) | `python -m src.main` |
| **UAV_VisLoc** | 81K (75K DB + 6.7K drone) | 17 (unified) | `python scripts/run_uav_visloc.py` |
| **GTA-UAV-LR** | 48K (15K sat + 34K drone) | 17 (unified) | `python scripts/run_gta_uav.py` |

> Все датасеты используют **единый набор 17 классов** (`scripts/seg_classes.py`) для совместимости при transfer learning.

## Quick Start

```bash
# World-UAV (основной датасет)
python -m src.main

# UAV_VisLoc
python scripts/run_uav_visloc.py

# GTA-UAV-LR
python scripts/run_gta_uav.py

# Тесты (149 шт, без GPU)
python -m pytest src/tests/ -v
```

## Структура проекта

```
.
├── in/
│   ├── config_files/              # Gin-конфигурация
│   │   ├── pipeline.gin           # Пути, стадии, save_npy/save_vis, resume, source
│   │   ├── models.gin             # Model IDs, weights_dir
│   │   ├── hardware.gin           # GPU profile, batch_size (None=auto), FP16
│   │   ├── segmentation.gin       # 11 промптов, threshold=0.15
│   │   └── input.gin              # image_size (256)
│   └── weights/                   # Веса моделей (не в git, >50MB)
│       ├── models--depth-anything--DA3-LARGE-1.1/
│       ├── sam3.1/sam3.1_multiplex.pt
│       └── dinov3-chmv2/
├── src/
│   ├── main.py                    # Entry point + pipeline orchestration
│   ├── nn/                        # Вендорированные нейросетевые пакеты
│   │   ├── __init__.py            # Регистрация sys.path при импорте
│   │   ├── segearth_ov3/          # SegEarth-OV-3 + SAM3 (копия репозитория)
│   │   │   ├── segearthov3_segmentor.py
│   │   │   ├── sam3/              # SAM 3.1 backbone (134 .py файла)
│   │   │   │   └── assets/bpe_simple_vocab_16e6.txt.gz
│   │   │   └── pamr.py
│   │   └── depth_anything_3/      # Depth-Anything-3 (копия пакета)
│   │       ├── api.py             # DepthAnything3 class
│   │       ├── model/             # DA3 архитектура (DinoV2 + DPT)
│   │       ├── configs/           # YAML-конфиги моделей
│   │       └── utils/             # I/O, export, geometry
│   ├── augmentor/
│   │   ├── models.py              # Загрузка/выгрузка моделей
│   │   ├── inference.py           # Inference + post-processing (depth, chmv2, edges, segm)
│   │   ├── io_utils.py            # Сохранение файлов (sync + async) + палитра
│   │   └── dataset.py             # Discovery, filtering, PyTorch Dataset
│   ├── conf/                      # Gin-configurable dataclasses
│   ├── utils/                     # Profiler, benchmark, GPU utils
│   └── tests/                     # 149 тестов (pytest)
├── scripts/
│   ├── seg_classes.py             # UNIFIED_PROMPTS — 17 классов (единый источник)
│   ├── run_uav_visloc.py         # Запуск для UAV_VisLoc
│   ├── run_gta_uav.py            # Запуск для GTA-UAV-LR
│   └── migrate_layout.py         # Миграция со старого prefix-формата
└── docs/
    ├── segmentation_class_analysis.md   # Анализ классов сегментации (11 классов)
    ├── segearth_ov3_architecture.md     # Архитектура SegEarth-OV3 + SAM 3.1
    ├── analysis_optimization.md         # Анализ производительности и оптимизации
    └── skills_optimization_io_dl_ml.md  # Справочник приемов оптимизации
```

### src/nn/ -- вендорированные пакеты

Нейросетевые модели **встроены внутрь проекта** в директории `src/nn/`. Не нужно клонировать внешние репозитории или устанавливать пакеты через pip:

- **`src/nn/segearth_ov3/`** -- полная копия [SegEarth-OV-3](https://github.com/earth-insights/SegEarth-OV-3): сегментатор + SAM3 backbone + BPE vocab
- **`src/nn/depth_anything_3/`** -- полная копия пакета из [Depth-Anything-3](https://github.com/ByteDance-Seed/Depth-Anything-3)

При `import src.nn` автоматически регистрируются пути в `sys.path`, и все внутренние импорты обоих пакетов работают без изменений.

## Конфигурация

### pipeline.gin

```python
PipelineConfig.input_root = '/path/to/UAV-GeoLoc'     # Исходный датасет
PipelineConfig.output_root = '/path/to/World-UAV-aug'  # Куда сохранять
PipelineConfig.stages = ['depth', 'edges', 'segmentation', 'chmv2']
PipelineConfig.save_npy = False            # True = float16/uint8 .npy (промежуточные)
PipelineConfig.save_vis = True             # True = .png визуализации
PipelineConfig.save_safetensors = True     # True = .safetensors (для обучения, zero-copy mmap)
PipelineConfig.cleanup_npy = False         # True = удалить .npy после консолидации
PipelineConfig.resume = True               # Пропускать уже обработанные
PipelineConfig.subset = None               # None=все, 'Rot', 'Country', 'Terrain'
PipelineConfig.source = 'db'               # 'db' = спутник, 'query' = БПЛА, None = оба
```

### segmentation.gin (unified 17 классов)

Все датасеты используют **единый набор 17 классов** из `scripts/seg_classes.py` для совместимости при transfer learning (pretrain GTA-UAV → fine-tune UAV_VisLoc/World-UAV). Не каждый датасет содержит пиксели каждого класса — это нормально (0 пикселей = 0 loss).

| ID | Промпт | World-UAV | UAV_VisLoc | GTA-UAV |
|:--:|:---|:---:|:---:|:---:|
| 0 | background | + | + | + |
| 1 | building | + | + | + |
| 2 | road | + | + | + |
| 3 | vegetation | + | + | + |
| 4 | water | + | + | + |
| 5 | sand and gravel ground | + | + | + |
| 6 | rocky terrain | + | + | + |
| 7 | farmland | + | + | + |
| 8 | railway | + | + | + |
| 9 | parking lot | + | + | + |
| 10 | sidewalk | + | + | + |
| 11 | bare soil and plowed field | + | + | + |
| 12 | roof and rooftop | + | + | + |
| 13 | sports field and playground | + | + | редко |
| 14 | muddy ground and wetland | + | + | reclassify* |
| 15 | embankment and levee | + | + | редко |
| 16 | swimming pool | + | редко | + |

\* GTA-UAV: `seg_reclassify_wetland=True` — wetland переклассифицируется в vegetation/bare soil (ложные срабатывания на холмах GTA V).

**Post-processing** (после SegEarth-OV3):
- `seg_fix_dark_water=True` (все датасеты) — background на тёмных изображениях (mean < 0.24, std < 0.18) → water. Satellite GTA-UAV: bg 57% → 5%.
- `seg_reclassify_wetland=True` (только GTA-UAV) — wetland → vegetation (зелёный) / bare soil (коричневый). Drone: ложный wetland 14% → 0%.

> Подробный анализ: [`docs/segmentation_class_analysis.md`](docs/segmentation_class_analysis.md)

### hardware.gin

```python
HardwareConfig.profile_name = 'rtx4090'
HardwareConfig.total_ram_gb = 24.0
HardwareConfig.use_fp16 = True
HardwareConfig.batch_size = None   # None = auto (из свободного VRAM)
HardwareConfig.num_workers = 4
```

## Как работает пайплайн

Стадии выполняются **последовательно** -- одна модель за раз:

```
DEPTH:        загрузка DA3 -> auto_batch_size из VRAM -> все изображения -> выгрузка
EDGES:        загрузка depth PNG/NPY -> Sobel (CPU, batch=32) -> выгрузка
SEGM:         загрузка SegEarth-OV3 -> batched backbone (<=16 img) + per-image grounding -> выгрузка
CHMv2:        загрузка DINOv3 (FP32) -> auto_batch_size из VRAM -> все изображения -> выгрузка
CONSOLIDATE:  сборка depth+edge+segm+chm -> один .safetensors на изображение (CPU)
```

**SegEarth-OV3:** backbone SAM 3.1 выполняется одним forward pass на батч до 16 изображений через `predict_pil_batch()`. Grounding decoder (11 промптов x per-image) -- основной bottleneck (~84% времени). Text embeddings кэшируются при первом вызове. Подробная архитектура: [`docs/segearth_ov3_architecture.md`](docs/segearth_ov3_architecture.md)

**auto_batch_size** после загрузки модели считывает реальный свободный VRAM:

```
free_vram = total - reserved
batch = round_down_pow2(free_vram / act_per_sample * 0.7)
```

**Resume** проверяет существование файлов в соответствующих директориях модальностей. Пайплайн можно прервать Ctrl+C и перезапустить -- готовые пропускаются.

## Формат выхода

Модальность определяется **папкой**, а не суффиксом файла:

```
World-UAV-aug/
├── depth/Rot/SouthernSuburbs/DB/img/crop_12_4.png       # vis
├── edge/Rot/SouthernSuburbs/DB/img/crop_12_4.png        # vis
├── segm/Rot/SouthernSuburbs/DB/img/crop_12_4.png        # vis (palette mode P)
├── chm/Rot/SouthernSuburbs/DB/img/crop_12_4.png         # vis
├── npy/depth/Rot/SouthernSuburbs/DB/img/crop_12_4.npy   # float16 intermediate
├── npy/edge/...
├── npy/segm/...
├── npy/chm/...
├── safetensors/Rot/SouthernSuburbs/DB/img/crop_12_4.safetensors  # для обучения
└── manifest.json
```

### SafeTensors (рекомендуемый формат для обучения)

Один `.safetensors` файл на изображение, содержит все модальности:

| Ключ | Dtype | Shape | Описание |
|:---|:---|:---|:---|
| `depth` | float16 | [1, H, W] | Карта глубины [0, 1] |
| `edge` | float16 | [1, H, W] | Границы (Sobel) [0, 1] |
| `chm` | float16 | [1, H, W] | Canopy height [0, 1] |
| `segm` | uint8 | [1, H, W] | Class IDs [0, 10] (World-UAV) / [0, 16] (unified) |

Преимущества SafeTensors:
- **Zero-copy mmap** -- тензор читается прямо с диска без копирования в RAM (~0.1ms)
- **1 syscall** вместо 4 (один файл = все модальности)
- **Безопасность** -- нет pickle, нет arbitrary code execution
- **Стандарт HuggingFace** -- нативная поддержка в PyTorch

### PNG визуализации (только для просмотра)

| Стадия | Суффикс | PNG формат |
|:---|:---|:---|
| depth | `_depth` | grayscale (L), uint8, `value / 255.0` -> [0,1] |
| edges | `_edge` | grayscale (L), uint8 |
| segmentation | `_segm` | RGB palette, class ID = argmax по палитре |
| chmv2 | `_chm` | grayscale (L), uint8, `value / 255.0` -> [0,1] |

### Палитра сегментации

| ID | Класс | RGB | Датасеты |
|:--:|:---|:---|:---|
| 0 | background | (0, 0, 0) | Black |
| 1 | building | (220, 40, 40) | Red |
| 2 | road | (160, 160, 160) | Gray |
| 3 | vegetation | (30, 180, 30) | Green |
| 4 | water | (30, 120, 220) | Blue |
| 5 | sand and gravel ground | (180, 140, 80) | Tan |
| 6 | rocky terrain | (120, 100, 80) | Brown |
| 7 | farmland | (200, 200, 50) | Yellow |
| 8 | railway | (100, 60, 120) | Purple |
| 9 | parking lot | (255, 165, 0) | Orange |
| 10 | sidewalk | (200, 200, 200) | Light gray |
| 11 | bare soil | (140, 100, 50) | Dark tan |
| 12 | rooftop | (180, 60, 60) | Dark red |
| 13 | sports field | (50, 200, 150) | Teal |
| 14 | muddy/wetland | (80, 100, 70) | Olive |
| 15 | embankment | (170, 130, 100) | Sandy brown |
| 16 | swimming pool | (0, 200, 255) | Cyan |

## Использование для обучения

### SafeTensors (рекомендуемый способ)

```python
from safetensors.torch import load_file

stem = "crop_12_4"
aug_root = Path("World-UAV-aug")
rel_parent = "Rot/SouthernSuburbs/DB/img"

# Zero-copy чтение всех модальностей за ~0.1ms
data = load_file(aug_root / "safetensors" / rel_parent / f"{stem}.safetensors", device="cpu")

depth = data["depth"]  # [1, 256, 256] float16, [0, 1]
edge  = data["edge"]   # [1, 256, 256] float16, [0, 1]
chm   = data["chm"]    # [1, 256, 256] float16, [0, 1]
segm  = data["segm"]   # [1, 256, 256] uint8, class IDs [0, 10]

# Для Teacher NADEZHDA: segm -> one-hot
import torch.nn.functional as F
segm_onehot = F.one_hot(segm.long().squeeze(0), num_classes=17)  # [H, W, 17]
segm_onehot = segm_onehot.permute(2, 0, 1).float()              # [17, H, W]
```

### PNG fallback (для визуализации или legacy)

```python
from PIL import Image
import numpy as np

# Depth / Edge / CHM -- grayscale float [0, 1]
depth = np.array(Image.open(aug_root / "depth" / rel_parent / f"{stem}.png")) / 255.0
edge  = np.array(Image.open(aug_root / "edge"  / rel_parent / f"{stem}.png")) / 255.0
chm   = np.array(Image.open(aug_root / "chm"   / rel_parent / f"{stem}.png")) / 255.0
```

> PNG визуализации квантуют float16 в uint8 (256 уровней). Для обучения используйте SafeTensors.

### Миграция со старого формата

Если данные сгенерированы в старом prefix-формате (`crop_12_4_depth.png`), мигрируйте:

```bash
# Сначала проверить (dry-run)
python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug --dry-run

# Выполнить миграцию
python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug
```

## Скачивание весов

Веса скачиваются один раз в `in/weights/` (~10 GB суммарно):

```bash
# DA3-LARGE-1.1 (HuggingFace, открытый доступ)
python -c "
from huggingface_hub import snapshot_download
snapshot_download('depth-anything/DA3-LARGE-1.1', cache_dir='in/weights')
"

# SAM 3.1 (для SegEarth-OV3)
mkdir -p in/weights/sam3.1
cp /path/to/sam3.1_multiplex.pt in/weights/sam3.1/

# CHMv2 DINOv3 (gated, нужен доступ к facebook/dinov3-vitl16-chmv2-dpt-head)
python -c "
from transformers import CHMv2ForDepthEstimation, CHMv2ImageProcessor
model = CHMv2ForDepthEstimation.from_pretrained('facebook/dinov3-vitl16-chmv2-dpt-head')
proc = CHMv2ImageProcessor.from_pretrained('facebook/dinov3-vitl16-chmv2-dpt-head')
model.save_pretrained('in/weights/dinov3-chmv2')
proc.save_pretrained('in/weights/dinov3-chmv2')
"
```

> BPE vocab (`bpe_simple_vocab_16e6.txt.gz`) уже встроен в проект: `src/nn/segearth_ov3/sam3/assets/`. Отдельно скачивать не нужно.

## Известные особенности

- **CHMv2 работает только в FP32** -- в FP16 выдает NaN. Модель автоматически загружается в FP32 независимо от `use_fp16`
- **SegEarth-OV3 bottleneck** -- grounding decoder (17 промптов x per-image) = ~84% времени инференса. Text embeddings кэшируются. Batch size backbone = 16
- **Post-processing сегментации** -- dark water fix (background → water для тёмных изображений) + wetland reclassify (GTA-UAV: wetland → vegetation/bare soil)
- **16 сцен Country исключены** -- неполные (нет DB-кропов). Фильтруются автоматически через `INCOMPLETE_SCENES`
- **Ледники/снег** -- SegEarth-OV3 классифицирует как `water` (ограничение модели). Класс `snow and ice` убран как неэффективный
- **Verbose логи подавлены** -- DA3, transformers, SAM 3.1, HF Hub. Управляется через `_silence_model_loggers()`

## Оценка времени (RTX 4090, 24 GB, 973K images)

| Стадия | Время | % |
|:---|:---|:---|
| Depth | ~14.7 ч | 16% |
| Edges | ~0.6 ч | <1% |
| Segmentation (bs=16, 17 prompts) | ~120 ч | **~76%** |
| CHMv2 | ~8.5 ч | ~8% |
| Consolidate (.safetensors) | ~0.1 ч | <1% |
| **Итого** | **~144 ч (~6 дней)** | |

> При обработке только DB (спутник, `source='db'`): ~486K изображений, ~50 ч.
> При обработке только query (БПЛА, `source='query'`): ~486K изображений, ~50 ч.

## Тесты

```bash
# Все тесты (149 штук, ~2.5 сек, без GPU)
python -m pytest src/tests/ -v

# Только pipeline integration
python -m pytest src/tests/test_pipeline_integration.py -v

# Только inference
python -m pytest src/tests/test_inference.py -v
```

Все тесты используют mock-модели -- GPU не требуется.

## Документация

| Документ | Описание |
|---|---|
| [`docs/segmentation_class_analysis.md`](docs/segmentation_class_analysis.md) | Unified 17 классов: анализ World-UAV (392 локации), UAV_VisLoc, GTA-UAV |
| [`docs/segearth_ov3_architecture.md`](docs/segearth_ov3_architecture.md) | Архитектура SegEarth-OV3 + SAM 3.1, pipeline инференса, профиль производительности |
| [`docs/analysis_optimization.md`](docs/analysis_optimization.md) | Общий анализ и оптимизация пайплайна |
| [`docs/skills_optimization_io_dl_ml.md`](docs/skills_optimization_io_dl_ml.md) | Справочник приемов оптимизации I/O, DataLoader, ML |

## Зависимости

- Python 3.10+
- PyTorch 2.x + CUDA
- transformers >= 5.5
- huggingface_hub
- safetensors >= 0.4
- gin-config, tqdm, Pillow, coloredlogs, psutil, matplotlib
- omegaconf, einops (зависимости Depth-Anything-3)
- iopath (зависимость SAM3)

> SegEarth-OV-3 и Depth-Anything-3 **вендорированы** в `src/nn/` -- отдельная установка не требуется.