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

# 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: сегментатор + SAM3 backbone + BPE vocab
• src/nn/depth_anything_3/ -- полная копия пакета из Depth-Anything-3

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

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

pipeline.gin

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_rgb < 0.24) → water
• seg_reclassify_wetland=True (только GTA-UAV) — wetland → vegetation (зелёный) / bare soil (коричневый)

Подробный анализ: docs/segmentation_class_analysis.md

hardware.gin

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

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 (рекомендуемый способ)

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)

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), мигрируйте:

# Сначала проверить (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 суммарно):

# 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 ч.

Тесты

# Все тесты (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	Unified 17 классов: анализ World-UAV (392 локации), UAV_VisLoc, GTA-UAV
docs/segearth_ov3_architecture.md	Архитектура SegEarth-OV3 + SAM 3.1, pipeline инференса, профиль производительности
docs/analysis_optimization.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/ -- отдельная установка не требуется.