f3cb18ac4d
Bundle all per-image modalities (depth, edge, chm, segm) into a single .safetensors file for fast training DataLoader reads (~0.1ms zero-copy mmap vs ~5ms for 4x PNG). Adds consolidate stage after main pipeline stages, save_safetensors/cleanup_npy config flags, resume support, and 10 new tests. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
331 lines
17 KiB
Markdown
331 lines
17 KiB
Markdown
# 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
|
||
|
||
```bash
|
||
# 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/ # 141 тест (pytest)
|
||
└── 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](https://github.com/earth-insights/SegEarth-OV-3): сегментатор + SAM3 backbone + BPE vocab
|
||
- **`src/nn/depth_anything_3/`** -- полная копия пакета из [Depth-Anything-3](https://github.com/ByteDance-Seed/Depth-Anything-3)
|
||
|
||
При `import src.nn` автоматически регистрируются пути в `sys.path`, и все внутренние импорты обоих пакетов работают без изменений.
|
||
|
||
## Конфигурация
|
||
|
||
### pipeline.gin
|
||
|
||
```python
|
||
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)
|
||
|
||
```python
|
||
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`](docs/segmentation_class_analysis.md)
|
||
|
||
### hardware.gin
|
||
|
||
```python
|
||
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`](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** проверяет существование `{stem}_{suffix}.png` (или `.npy`) для каждого изображения и `{stem}.safetensors` для этапа консолидации. Пайплайн можно прервать Ctrl+C и перезапустить -- готовые пропускаются.
|
||
|
||
## Формат выхода
|
||
|
||
Структура директорий **зеркалит** исходный датасет. Исходные изображения не копируются:
|
||
|
||
```
|
||
World-UAV-aug/
|
||
├── Rot/SouthernSuburbs/DB/img/
|
||
│ ├── crop_12_4.safetensors # ВСЕ модальности (для обучения, zero-copy mmap)
|
||
│ ├── crop_12_4_depth.png # grayscale визуализация
|
||
│ ├── crop_12_4_edge.png # grayscale визуализация
|
||
│ ├── crop_12_4_segm.png # RGB palette визуализация (11 классов)
|
||
│ └── crop_12_4_chm.png # grayscale визуализация
|
||
├── Country/...
|
||
└── Terrain/...
|
||
```
|
||
|
||
### 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 (рекомендуемый способ)
|
||
|
||
```python
|
||
from safetensors.torch import load_file
|
||
|
||
stem = "crop_12_4"
|
||
aug_dir = Path("World-UAV-aug/Rot/SouthernSuburbs/DB/img")
|
||
|
||
# Zero-copy чтение всех модальностей за ~0.1ms
|
||
data = load_file(aug_dir / 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)
|
||
|
||
```python
|
||
from PIL import Image
|
||
import numpy as np
|
||
|
||
# Depth / Edge / CHM -- grayscale float [0, 1]
|
||
depth = np.array(Image.open(aug_dir / f"{stem}_depth.png")) / 255.0 # [H, W]
|
||
edge = np.array(Image.open(aug_dir / f"{stem}_edge.png")) / 255.0
|
||
chm = np.array(Image.open(aug_dir / f"{stem}_chm.png")) / 255.0
|
||
```
|
||
|
||
> PNG визуализации квантуют float16 в uint8 (256 уровней). Для обучения используйте SafeTensors.
|
||
|
||
## Скачивание весов
|
||
|
||
Веса скачиваются один раз в `in/weights/` (~10 GB суммарно):
|
||
|
||
```bash
|
||
# 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 ч.
|
||
|
||
## Тесты
|
||
|
||
```bash
|
||
# Все тесты (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`](docs/segmentation_class_analysis.md) | Анализ 392 локаций, выбор 11 классов, результаты тестирования |
|
||
| [`docs/segearth_ov3_architecture.md`](docs/segearth_ov3_architecture.md) | Архитектура SegEarth-OV3 + SAM 3.1, pipeline инференса, профиль производительности |
|
||
| [`docs/analysis_optimization.md`](docs/analysis_optimization.md) | Общий анализ и оптимизация пайплайна |
|
||
| [`docs/skills_optimization_io_dl_ml.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/` -- отдельная установка не требуется.
|