World-UAV Annotation Pipeline

Автоматическая генерация 4 модальностей из RGB-изображений датасета World-UAV (973K images):

Модальность	Модель	Выход	Скорость
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

Quick Start

# 1. Запуск (из корня проекта)
python -m src.main

# 2. Тесты
python -m pytest src/tests/ -v

Все параметры настраиваются через in/config_files/*.gin. Аргументов командной строки нет.

Структура проекта

.
├── 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/
│   └── 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

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

Палитра сегментации (11 классов)

ID	Класс	Цвет	RGB
0	background	Black	(0, 0, 0)
1	building	Red	(220, 40, 40)
2	road	Gray	(160, 160, 160)
3	vegetation	Green	(30, 180, 30)
4	water	Blue	(30, 120, 220)
5	sand and gravel ground	Tan	(180, 140, 80)
6	rocky terrain	Brown	(120, 100, 80)
7	farmland	Yellow	(200, 200, 50)
8	railway	Purple	(100, 60, 120)
9	parking lot	Orange	(255, 165, 0)
10	sidewalk	Light gray	(200, 200, 200)

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

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

Тесты

# Все тесты (141 штука, ~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/ -- отдельная установка не требуется.