Verified on 128 satellite images after dark water fix (std threshold 0.08 → 0.18). Document calibrated thresholds and measured improvements. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
380 lines
19 KiB
Markdown
380 lines
19 KiB
Markdown
# 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/` -- отдельная установка не требуется.
|