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)	11 (базовые)	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

UAV_VisLoc и GTA-UAV используют единый набор 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

# Тесты (143 шт, без 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 функции (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/                     # 143 теста (pytest)
├── scripts/
│   ├── run_uav_visloc.py         # Запуск для UAV_VisLoc (16 классов)
│   ├── run_gta_uav.py            # Запуск для GTA-UAV-LR (14 классов)
│   └── 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 (11 классов open-vocabulary)

SegConfig.prompts = [
    'background',              # 0  -- unclassified
    'building',                # 1  -- buildings, rooftops
    'road',                    # 2  -- roads, asphalt
    'vegetation',              # 3  -- trees, bushes, forest canopy
    'water',                   # 4  -- rivers, canals, sea, lakes
    'sand and gravel ground',  # 5  -- soil, gravel, sand, dust, bare earth
    'rocky terrain',           # 6  -- rock, stone, lava, canyon walls
    'farmland',                # 7  -- agricultural terraces, fields
    'railway',                 # 8  -- railway tracks, rails
    'parking lot',             # 9  -- parking areas
    'sidewalk',                # 10 -- sidewalks, pedestrian zones
]
SegConfig.threshold = 0.15
SegConfig.default_resolution = 1008

Unified 17 классов (UAV_VisLoc + GTA-UAV)

Для transfer learning (pretrain GTA-UAV → fine-tune UAV_VisLoc) оба датасета используют единый набор 17 классов из scripts/seg_classes.py. Не каждый датасет содержит пиксели каждого класса — это нормально (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	—	+	—
15	embankment and levee	—	+	—
16	swimming pool	—	редко	+

UAV_VisLoc threshold=0.1, GTA-UAV threshold=0.15. Подробный анализ: 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)	оба
1	building	(220, 40, 40)	оба
2	road	(160, 160, 160)	оба
3	vegetation	(30, 180, 30)	оба
4	water	(30, 120, 220)	оба
5	sand and gravel ground	(180, 140, 80)	оба
6	rocky terrain	(120, 100, 80)	оба
7	farmland	(200, 200, 50)	оба
8	railway	(100, 60, 120)	оба
9	parking lot	(255, 165, 0)	оба
10	sidewalk	(200, 200, 200)	оба
11	bare soil	(140, 100, 50)	UAV_VisLoc
12	rooftop	(180, 60, 60)	UAV_VisLoc
13	sports field	(50, 200, 150)	UAV_VisLoc
14	muddy/wetland	(80, 100, 70)	UAV_VisLoc
15	embankment	(170, 130, 100)	UAV_VisLoc
16	swimming pool	(0, 200, 255)	GTA-UAV

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

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 (11 промптов x per-image) = ~84% времени инференса. Text embeddings кэшируются. Batch size backbone = 16
• 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, 11 prompts)	~77 ч	~70%
CHMv2	~8.5 ч	~8%
Consolidate (.safetensors)	~0.1 ч	<1%
Итого	~101 ч (~4 дня)

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

Тесты

# Все тесты (143 штуки, ~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	Анализ 392 локаций, выбор 11 классов, результаты тестирования
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/ -- отдельная установка не требуется.