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
2026-04-16 11:22:01 +03:00
commit 686db62c25
312 changed files with 83978 additions and 0 deletions

View File

@@ -0,0 +1,163 @@
# Анализ и оптимизация пайплайна depth_edges_annotate_worlduav
## Обзор
Пайплайн для аннотации датасета World-UAV: четыре стадии (depth, edges, segmentation, chmv2), конфигурация через Gin, fallback-модели, атомарное сохранение, resume-логика. Нейросетевые модели (SegEarth-OV-3, Depth-Anything-3) вендорированы в `src/nn/`.
---
## 1. Инференс GPU -- самые крупные выигрыши
### 1.1 `torch.compile()` для моделей
В `models.py` после загрузки модели можно обернуть в `torch.compile(model, mode="reduce-overhead")`. Для DepthAnything и SegFormer это может дать **20-40% ускорения** на repeated batches за счет fusion ядер и устранения Python overhead. Особенно эффективно при большом количестве батчей.
### 1.2 AMP через `torch.autocast` вместо ручного `.half()`
Сейчас в `inference.py:79` и `inference.py:203` делается ручной каст `x = x.half()` при условии `model.parameters().dtype == torch.float16`. Это неоптимально:
- `torch.autocast("cuda", dtype=torch.float16)` автоматически выберет FP16 для matmul/conv, но оставит FP32 для нормализации и softmax -- меньше потерь точности и обычно быстрее
- Убирает необходимость в ручных проверках dtype
### 1.3 DA3: batch-инференс вместо поштучного цикла
`inference.py:59-74` -- DA3 обрабатывается через `model.inference()` API, который внутри делает preprocess + forward + postprocess. Для максимальной производительности можно вызывать низкоуровневый `model.forward()` напрямую (весь батч тензором на GPU), минуя конвертации tensor -> numpy -> PIL -> обратно.
### 1.4 SegEarth-OV3: batched backbone (реализовано в v3.2)
`predict_pil_batch()` батчит backbone SAM 3.1 (~80% времени) на до 8 изображений. Grounding decoder остается per-image. Дальнейшие оптимизации:
- `torch.cuda.Stream` для overlap compute/data-transfer
- Увеличение max_batch если VRAM позволяет (текущий лимит 8 hardcoded)
---
## 2. Data Pipeline -- ускорение загрузки
### 2.1 `persistent_workers=True` (реализовано)
DataLoader создается с `persistent_workers=True` при `num_workers > 0`. Воркеры не пересоздаются между батчами.
### 2.2 `prefetch_factor` (реализовано)
`prefetch_factor=4` для лучшего overlap загрузки с инференсом.
### 2.3 Decode на GPU через `torchvision.io` или `nvidia-dali`
`dataset.py` -- `Image.open().convert("RGB")` + `transforms.Resize()` работает на CPU через PIL. Для больших датасетов это bottleneck. Варианты:
- `torchvision.io.decode_image` + `torchvision.transforms.v2` -- resize на GPU
- NVIDIA DALI pipeline -- полный decode+resize+normalize на GPU, особенно выгоден при images > 10k
### 2.4 Предвычисление списка файлов
`dataset.py` -- `rglob("*")` обходит файловую систему при каждом запуске. Для больших датасетов (тысячи папок) это минуты. Можно кешировать список в `.file_cache.json` и обновлять только при изменении mtime корня.
---
## 3. I/O -- запись результатов
### 3.1 Асинхронная запись (реализовано)
`io_utils.py` -- `ThreadPoolExecutor` с 4 workers для неблокирующей записи файлов. GPU inference продолжается пока предыдущий батч пишется на диск.
### 3.2 Атомарная запись (реализовано)
Temp file + `os.replace()` для crash-safety. Resume-логика корректно обрабатывает прерванные записи.
### 3.3 Визуализации -- отложить или отключить
`save_vis=False` отключает PNG-визуализации. Для ускорения можно генерировать визуализации отдельным скриптом после всего инференса.
### 3.4 Pre-create output dirs (реализовано)
Все output_dir создаются одним проходом до начала обработки (`run_pipeline()` в `main.py:333-339`), а не per-image.
---
## 4. Edges стадия
### 4.1 Batched Sobel (реализовано)
Edges обрабатываются батчами по 32 изображения. `compute_edges_from_depth` поддерживает `[B, 1, H, W]`.
### 4.2 Sobel-ядра как module-level константы (реализовано)
`_SOBEL_X` и `_SOBEL_Y` -- module-level константы в `inference.py`.
### 4.3 Edges как побочный продукт depth
Если не нужен resume между стадиями -- считать edges прямо в `run_depth_stage` из только что вычисленного depth, не сохраняя и не загружая `.npy`. Текущая архитектура предпочитает гранулярный resume.
---
## 5. Resume и discovery
### 5.1 Completion manifest вместо per-file проверок
`filter_completed()` -- для каждого изображения вызывается `Path.exists()` (syscall). При 100k изображений x 4 стадии = 400k stat-вызовов. Альтернатива:
- `completed.json` / SQLite per stage
- Проверка по set-lookup вместо filesystem
### 5.2 `filter_completed` -- per-stage
Каждая стадия фильтрует отдельно. Можно за один проход собрать статусы всех стадий.
---
## 6. Память и типы данных
### 6.1 FP32 для CHMv2 (вынужденная мера)
CHMv2 (DINOv3 DPT head) выдает NaN в FP16. Всегда FP32 -- это увеличивает VRAM на ~0.65 GB vs FP16, но гарантирует корректность.
### 6.2 uint8 для сегментации (реализовано)
`infer_segmentation_batch` возвращает `uint8` для class IDs. При 5 классах это оптимально.
### 6.3 float16 для .npy (реализовано)
Depth, edges, CHM сохраняются в `.npy` как `float16` (если `save_npy=True`).
---
## 7. Вендорированные пакеты (src/nn/)
### 7.1 Текущая архитектура
SegEarth-OV-3 и Depth-Anything-3 **встроены в проект** как вендорированные пакеты в `src/nn/`. При `import src.nn` автоматически регистрируются пути в `sys.path`:
- `src/nn/` -- для `depth_anything_3.*`
- `src/nn/segearth_ov3/` -- для `sam3.*` и `segearthov3_segmentor`
**Преимущества:**
- Нет зависимости от внешних репозиториев
- Воспроизводимость -- фиксированная версия кода
- Нет конфликтов с системными пакетами
**Ограничения:**
- Обновление до новой версии модели требует ручного копирования
- Дублирование кода (~6.5 MB) -- допустимо для проекта с ~200 GB данных
### 7.2 BPE vocab
Файл `bpe_simple_vocab_16e6.txt.gz` встроен в `src/nn/segearth_ov3/sam3/assets/`. Дефолтный путь в `segearthov3_segmentor.py` вычисляется через `os.path.dirname(__file__)`. Копия в `in/weights/` используется приоритетно (если существует) для совместимости с существующими инсталляциями.
---
## Приоритет по соотношению усилие/выигрыш
| # | Оптимизация | Усилие | Выигрыш | Статус |
|---|-------------|--------|---------|--------|
| 1 | Async I/O через ThreadPool | Низкое | **Высокий** (overlap GPU/disk) | **Реализовано** |
| 2 | Batched backbone SegEarth-OV3 | Среднее | **Высокий** (~2.5-3x segm) | **Реализовано** |
| 3 | Pre-create output dirs | Минимальное | Средний | **Реализовано** |
| 4 | Batched Sobel edges | Низкое | Средний | **Реализовано** |
| 5 | dtype: uint8 для seg, float16 для npy | Низкое | Средний | **Реализовано** |
| 6 | `persistent_workers` + `prefetch_factor` | Минимальное | Низкий-Средний | **Реализовано** |
| 7 | Вендоринг моделей (src/nn/) | Среднее | Средний (reliability) | **Реализовано** |
| 8 | `torch.compile()` | Низкое | Средний (20-40%) | Не реализовано |
| 9 | `torch.autocast` вместо `.half()` | Низкое | Низкий-Средний | Не реализовано |
| 10 | DA3 прямой forward вместо `inference()` | Среднее | Средний | Не реализовано |
| 11 | Completion manifest вместо per-file exists | Среднее | Средний (при >100k img) | Не реализовано |
| 12 | GPU decode (DALI / torchvision.io) | Высокое | Средний | Не реализовано |