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)	16	python scripts/run_uav_visloc.py
GTA-UAV-LR	48K (15K sat + 34K drone)	14	python scripts/run_gta_uav.py

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

UAV_VisLoc: 16 классов (расширенные)

UAV_VisLoc содержит больше сельских/прибрежных сцен. Анализ scene 01 (2496 images) показал 29% background при 11 классах. Добавлены 5 классов:

ID	Промпт	Зачем
11	bare soil and plowed field	Распаханные поля без посевов — 15-20% бывшего background
12	roof and rooftop	Тёмные крыши, солнечные панели — building ловит стены, не крыши сверху
13	sports field and playground	Уникальная геометрия, сильный landmark для CVGL
14	muddy ground and wetland	Болотистый грунт между водоёмами (дамбы, берега)
15	embankment and levee	Линейные земляные сооружения вдоль каналов

Threshold снижен до 0.1 (vs 0.15 для World-UAV) — UAV_VisLoc содержит больше тёмных/низкоконтрастных сцен.

GTA-UAV-LR: 14 классов

Синтетический датасет из GTA V. Особенности:

• Satellite в RGBA (alpha = граница карты) → автоматически конвертируется в RGB
• Drone 512x384 (не квадратные) → resize в 512x512
• 6 высот полёта: 100, 200, 300, 400, 500, 600 м

Классы = 11 базовых + 3 специфических для GTA:

ID	Промпт	Зачем
11	bare soil and plowed field	Пустыри, стройплощадки (много в пригородах GTA)
12	roof and rooftop	Плоские крыши, вертолётные площадки
13	swimming pool	Бассейны — характерный landmark пригородов GTA

Подробный анализ выбора классов: 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]

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

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

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=11)  # [H, W, 11]
segm_onehot = segm_onehot.permute(2, 0, 1).float()              # [11, 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/ -- отдельная установка не требуется.