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)
│   │   └── pipeline_*.gin         # Датасет-специфичные оверрайды (напр. pipeline_uav_visloc.gin);
│   │                              #   НЕ подгружаются `python -m src.main` — только своими скриптами
│   └── 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 = None               # None = оба (по умолчанию), 'db' = спутник, 'query' = БПЛА

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

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 файл на изображение, содержит все модальности как dense spatial maps [1, H, W] — финальные predictions моделей, не промежуточные feature maps. Подробная спецификация: docs/tensor_format_spec.md.

Ключ	Dtype	Shape	Описание	Модель
depth	float16	[1, H, W]	Relative depth, [0, 1], per-frame min-max normalized	DA3-LARGE-1.1 (DinoV2-ViTL)
edge	float16	[1, H, W]	Sobel magnitude из depth, [0, 1], per-frame max normalized	Sobel (CPU)
chm	float16	[1, H, W]	Relative canopy height, [0, 1], per-frame min-max normalized	DINOv3-ViTL16 CHMv2
segm	uint8	[1, H, W]	Dense class ID map (argmax), значения [0, 16] (17 классов)	SegEarth-OV3 (ViTDet-32L)

Что НЕ сохраняется: backbone features (ViT tokens), soft probabilities/logits сегментации, абсолютные значения глубины/высоты в метрах, text embeddings промптов, направление градиента (только magnitude).

Преимущества 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

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

Все модальности хранятся как dense spatial maps [1, H, W] — финальные predictions моделей, не RGB-рендеры и не промежуточные feature maps. Тензоры сохраняют полную информацию без потерь при квантовании/colormapping и не являются OOD-входом для DINOv3. Полная спецификация форматов: docs/tensor_format_spec.md.

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

from safetensors.torch import load_file
import torch.nn.functional as F

# Zero-copy чтение всех модальностей за ~0.1ms
data = load_file("World-UAV-aug/safetensors/Rot/.../crop_12_4.safetensors")

# Все модальности — dense spatial maps, готовые для injection в backbone
depth = data["depth"]  # [1, H, W] float16, relative depth [0, 1], per-frame min-max
edge  = data["edge"]   # [1, H, W] float16, Sobel magnitude [0, 1], per-frame max
chm   = data["chm"]    # [1, H, W] float16, relative canopy height [0, 1], per-frame min-max
segm  = data["segm"]   # [1, H, W] uint8, class IDs (argmax) [0, 16], no normalization

Подача в Teacher NADEZHDA

Каждая модальность подаётся в свой lightweight aux-encoder, затем через FiLM injection в DINOv3 patch tokens:

# Depth / Edge / CHM → [B, 1, H, W] float → Conv aux-encoder → FiLM injection
# Все три уже в [0, 1] — рекомендуемый FiLM init: gamma=1, beta=0 (identity)
aux_depth = depth_encoder(depth.float())   # [1, H, W] → [C, H, W]
aux_edge  = edge_encoder(edge.float())
aux_chm   = chm_encoder(chm.float())

# Segmentation → dense class ID map → embedding → spatial feature map
# Вариант 1: one-hot → Conv (deterministic init, explicit)
segm_onehot = F.one_hot(segm.long().squeeze(0), num_classes=17)  # [H, W, 17]
segm_features = seg_conv(segm_onehot.permute(2, 0, 1).float())  # [17, H, W] → [C, H, W]

# Вариант 2: learned per-class embedding (компактнее)
# seg_emb = nn.Embedding(17, 32)
# segm_features = seg_emb(segm.long().squeeze(0)).permute(2, 0, 1)  # [H, W] → [32, H, W]

Нормирование входных тензоров

Depth, edge, chm уже per-frame normalized в [0, 1]. Для FiLM init gamma=1, beta=0 (identity) — безопасный старт. При необходимости dataset-level статистики — пройти по subset (1K-5K изображений) для running mean/std. Segmentation (uint8 class IDs) требует embedding перед FiLM — нельзя подавать как continuous input.

Почему тензоры, а не RGB-рендеры

Формат	Пример depth	Потеря информации	Для DINOv3
float16 тензор (хранится)	[0.4231, 0.4235, ...]	~0%	Прямой вход в aux-encoder
uint8 grayscale PNG	[108, 108, ...]	~0.4%	Приемлемо
turbo colormap RGB PNG	[R=50, G=180, B=220]	~70%	OOD — DINOv3 обучен на натуральных RGB

Для обучения всегда используйте SafeTensors. PNG визуализации — только для просмотра в Obsidian/файловом менеджере.

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

Если данные сгенерированы в старом 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/ -- отдельная установка не требуется.