Initial commit: World-UAV annotation pipeline
4-modality annotation pipeline (depth, edges, segmentation, chmv2) for 973K drone/satellite images. SegEarth-OV3 open-vocabulary segmentation with 11 classes optimized for cross-view geo-localization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
294
README.md
Normal file
294
README.md
Normal file
@@ -0,0 +1,294 @@
|
||||
# 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) | RGB palette [256x256] | ~3.5 img/s |
|
||||
| **CHMv2** | DINOv3-ViTL16 (337M, FP32) | grayscale [256x256] | 31.7 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/ # 125 тестов (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.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 -> все изображения -> выгрузка
|
||||
```
|
||||
|
||||
**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`) для каждого изображения. Пайплайн можно прервать Ctrl+C и перезапустить -- готовые пропускаются.
|
||||
|
||||
## Формат выхода
|
||||
|
||||
Структура директорий **зеркалит** исходный датасет. Исходные изображения не копируются:
|
||||
|
||||
```
|
||||
World-UAV-aug/
|
||||
├── Rot/SouthernSuburbs/DB/img/
|
||||
│ ├── crop_12_4_depth.png # grayscale, 1 канал
|
||||
│ ├── crop_12_4_edge.png # grayscale, 1 канал
|
||||
│ ├── crop_12_4_segm.png # RGB palette (11 классов)
|
||||
│ └── crop_12_4_chm.png # grayscale, 1 канал
|
||||
├── Country/...
|
||||
└── Terrain/...
|
||||
```
|
||||
|
||||
### Суффиксы
|
||||
|
||||
| Стадия | Суффикс | 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) |
|
||||
|
||||
## Использование для обучения
|
||||
|
||||
Depth, edge, chm -- **grayscale 1-канальные**. Загружать как float [0, 1]:
|
||||
|
||||
```python
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
stem = "crop_12_4"
|
||||
aug_dir = Path("World-UAV-aug/Rot/SouthernSuburbs/DB/img")
|
||||
|
||||
# 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
|
||||
|
||||
# Segmentation -- class index [0, 10]
|
||||
# Если save_npy=True: seg = np.load(aug_dir / f"{stem}_segm.npy") # [1, H, W] uint8
|
||||
# Если только PNG, используй LUT для обратного маппинга из RGB
|
||||
|
||||
# Конкатенация: RGB(3) + depth(1) + edge(1) + chm(1) = 6 каналов
|
||||
aux = np.stack([depth, edge, chm], axis=0) # [3, H, W] float32
|
||||
```
|
||||
|
||||
> Для сегментации рекомендуется включить `save_npy = True` -- обратный маппинг из RGB палитры в class ID ненадежен.
|
||||
|
||||
## Скачивание весов
|
||||
|
||||
Веса скачиваются один раз в `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% |
|
||||
| **Итого** | **~101 ч (~4 дня)** | |
|
||||
|
||||
> При обработке только DB (спутник, `source='db'`): ~486K изображений, ~50 ч.
|
||||
> При обработке только query (БПЛА, `source='query'`): ~486K изображений, ~50 ч.
|
||||
|
||||
## Тесты
|
||||
|
||||
```bash
|
||||
# Все тесты (125 штук, ~0.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
|
||||
- gin-config, tqdm, Pillow, coloredlogs, psutil, matplotlib
|
||||
- omegaconf, einops (зависимости Depth-Anything-3)
|
||||
- iopath (зависимость SAM3)
|
||||
|
||||
> SegEarth-OV-3 и Depth-Anything-3 **вендорированы** в `src/nn/` -- отдельная установка не требуется.
|
||||
Reference in New Issue
Block a user