forked from Pikaliov/fuze_task
fuse_proj: Initial operational package for 3 researchers (Pavlenko/Blizno/Moroz)
Multimodal fusion research on StripNet+GTA-UAV proxy: - 3 independent fusion tracks: condition-aware (A), token/bottleneck (B), role-aware (C) - Shared interfaces, protocol, dataset audit, baseline benchmarks - Canonical version-chain references to vault (SPEC, ANALYSIS, TRIAGE) - Personalized task plans and decision tables for each researcher - 3 generated DOCX task assignment files with milestones and DoD checklist - Full modality dropout diagnostics and missing-modality robustness requirements - Data contract, benchmark registry, experiment tracking infrastructure Operational documents: - docs/00_project/: MERIDIAN context, protocol, repository reuse guide, experiment specification - docs/01_tasks/: Master assignment + 3 individual researcher tracks + joint integration - docs/02_references/: Core literature, version-chain bases, code maps - docs/03_codebase_guides/: Existing code snapshots from vault - scripts/: gen_task_plans.js (DOCX generation), placeholder infrastructure - vendor_reference/: Snapshots of caption_test, depth_edges_annotate, existing SOFIA/SegModel code - reports/, results/, experiments/: Shared output structure for all 3 researchers 3 DOCX files generated from gen_task_plans.js (Times New Roman 14pt, GOST format): - План_заданий_Павленко_БВ.docx (Condition-Aware track, fusion API owner) - План_заданий_Близно_МВ.docx (Token/Bottleneck track, benchmark owner) - План_заданий_Мороз_ЕС.docx (Role-Aware track, data contract owner) Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
Pikaliov
71
vendor_reference/depth_edges_annotate_worlduav/CLAUDE.md
Normal file
71
vendor_reference/depth_edges_annotate_worlduav/CLAUDE.md
Normal file
@@ -0,0 +1,71 @@
|
||||
# CLAUDE.md
|
||||
|
||||
## Что это за проект
|
||||
|
||||
Пайплайн автоматической генерации 4 вспомогательных модальностей (depth, edges, segmentation, canopy height) из RGB-изображений аэрофотоснимков. Используется для подготовки обучающих данных для NADEZHDA — системы cross-view geolocalization (БПЛА ↔ спутник).
|
||||
|
||||
## Быстрый старт
|
||||
|
||||
```bash
|
||||
# World-UAV (973K images, основной датасет)
|
||||
python -m src.main
|
||||
|
||||
# UAV_VisLoc (81K images)
|
||||
python scripts/run_uav_visloc.py
|
||||
|
||||
# GTA-UAV-LR (48K images, synthetic GTA V)
|
||||
python scripts/run_gta_uav.py
|
||||
|
||||
# Тесты (149 шт, без GPU)
|
||||
python -m pytest src/tests/ -v
|
||||
```
|
||||
|
||||
## Поддерживаемые датасеты
|
||||
|
||||
| Датасет | Изображения | Тип | Скрипт |
|
||||
|---|---|---|---|
|
||||
| World-UAV | 973K | Реальные аэрофото, 27 terrain, 11 стран | `python -m src.main` |
|
||||
| UAV_VisLoc | 81K | Реальные, 11 сцен, DB + drone | `python scripts/run_uav_visloc.py` |
|
||||
| GTA-UAV-LR | 48K | Синтетика GTA V, 6 высот полёта | `python scripts/run_gta_uav.py` |
|
||||
|
||||
## Ключевые решения
|
||||
|
||||
- **Формат выхода:** SafeTensors с **dense tensor maps** (zero-copy mmap, ~0.1ms). Все модальности — прямые тензоры (float16/uint8), не RGB-рендеры.
|
||||
- **Структура директорий:** модальность = папка (`depth/`, `edge/`, `segm/`, `chm/`, `safetensors/`), не суффикс файла.
|
||||
- **Стадии последовательно** — одна модель в GPU за раз (экономия VRAM).
|
||||
- **Сегментация:** SegEarth-OV3 (SAM 3.1 + open-vocabulary prompts). **17 unified классов** для всех датасетов (единые ID для transfer learning).
|
||||
- **Post-processing:** два правила после SegEarth — dark water fix (mean<0.24, std<0.18 → water; satellite bg 57%→5%) и wetland reclassify (GTA-UAV: ложный wetland 14%→0%).
|
||||
- **CHMv2 только FP32** — в FP16 NaN.
|
||||
|
||||
## Структура кода
|
||||
|
||||
```
|
||||
src/main.py — точка входа, оркестрация стадий
|
||||
src/augmentor/inference.py — инференс + postprocess_segmentation()
|
||||
src/augmentor/io_utils.py — I/O, SafeTensors, палитра 17 классов
|
||||
src/augmentor/dataset.py — discovery, filtering (DB/query/drone/satellite)
|
||||
src/conf/ — gin-configurable dataclasses
|
||||
src/nn/ — вендорированные DA3 + SegEarth-OV3
|
||||
scripts/seg_classes.py — UNIFIED_PROMPTS (17 классов, единый источник)
|
||||
scripts/run_*.py — скрипты запуска для каждого датасета
|
||||
in/config_files/ — gin-конфиги
|
||||
docs/ — документация
|
||||
```
|
||||
|
||||
## Конфигурация
|
||||
|
||||
Все параметры через gin. CLI override: `--gin "PipelineConfig.source = 'db'"`.
|
||||
|
||||
Для нового датасета — создать скрипт в `scripts/` (пример: `run_gta_uav.py`).
|
||||
|
||||
Ключевые флаги pipeline:
|
||||
- `seg_fix_dark_water=True` — автоматически исправлять тёмную воду (по умолчанию вкл.)
|
||||
- `seg_reclassify_wetland=False` — переклассификация wetland в vegetation/bare soil (вкл. для GTA-UAV)
|
||||
|
||||
## Что НЕ делать
|
||||
|
||||
- Не менять порядок/ID классов в `scripts/seg_classes.py` — все датасеты зависят от фиксированных ID.
|
||||
- Не использовать `.pt` (torch.save) для хранения тензоров — pickle, медленно.
|
||||
- Не рендерить depth/seg в RGB colormap для обучения — OOD для DINOv3, потеря ~70% информации. Использовать dense tensor maps из SafeTensors.
|
||||
- Не снижать threshold ниже 0.1 — увеличивает false positives без значимого улучшения recall.
|
||||
- Не менять `dark_water_std_thr` (0.18) — калибровано на GTA-UAV ocean (std 0.10-0.15). Ниже — не ловит, выше — false positives на normal images.
|
||||
391
vendor_reference/depth_edges_annotate_worlduav/README.md
Normal file
391
vendor_reference/depth_edges_annotate_worlduav/README.md
Normal file
@@ -0,0 +1,391 @@
|
||||
# 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) | 17 (unified) | `python -m src.main` |
|
||||
| **UAV_VisLoc** | 81K (75K DB + 6.7K drone) | 17 (unified) | `python scripts/run_uav_visloc.py` |
|
||||
| **GTA-UAV-LR** | 48K (15K sat + 34K drone) | 17 (unified) | `python scripts/run_gta_uav.py` |
|
||||
|
||||
> Все датасеты используют **единый набор 17 классов** (`scripts/seg_classes.py`) для совместимости при transfer learning.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# World-UAV (основной датасет)
|
||||
python -m src.main
|
||||
|
||||
# UAV_VisLoc
|
||||
python scripts/run_uav_visloc.py
|
||||
|
||||
# GTA-UAV-LR
|
||||
python scripts/run_gta_uav.py
|
||||
|
||||
# Тесты (149 шт, без 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 + post-processing (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/ # 149 тестов (pytest)
|
||||
├── scripts/
|
||||
│ ├── seg_classes.py # UNIFIED_PROMPTS — 17 классов (единый источник)
|
||||
│ ├── run_uav_visloc.py # Запуск для UAV_VisLoc
|
||||
│ ├── run_gta_uav.py # Запуск для GTA-UAV-LR
|
||||
│ └── 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](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 (unified 17 классов)
|
||||
|
||||
Все датасеты используют **единый набор 17 классов** из `scripts/seg_classes.py` для совместимости при transfer learning (pretrain GTA-UAV → fine-tune UAV_VisLoc/World-UAV). Не каждый датасет содержит пиксели каждого класса — это нормально (0 пикселей = 0 loss).
|
||||
|
||||
| ID | Промпт | World-UAV | UAV_VisLoc | GTA-UAV |
|
||||
|:--:|:---|:---:|:---:|:---:|
|
||||
| 0 | background | + | + | + |
|
||||
| 1 | building | + | + | + |
|
||||
| 2 | road | + | + | + |
|
||||
| 3 | vegetation | + | + | + |
|
||||
| 4 | water | + | + | + |
|
||||
| 5 | sand and gravel ground | + | + | + |
|
||||
| 6 | rocky terrain | + | + | + |
|
||||
| 7 | farmland | + | + | + |
|
||||
| 8 | railway | + | + | + |
|
||||
| 9 | parking lot | + | + | + |
|
||||
| 10 | sidewalk | + | + | + |
|
||||
| 11 | bare soil and plowed field | + | + | + |
|
||||
| 12 | roof and rooftop | + | + | + |
|
||||
| 13 | sports field and playground | + | + | редко |
|
||||
| 14 | muddy ground and wetland | + | + | reclassify* |
|
||||
| 15 | embankment and levee | + | + | редко |
|
||||
| 16 | swimming pool | + | редко | + |
|
||||
|
||||
\* GTA-UAV: `seg_reclassify_wetland=True` — wetland переклассифицируется в vegetation/bare soil (ложные срабатывания на холмах GTA V).
|
||||
|
||||
**Post-processing** (после SegEarth-OV3):
|
||||
- `seg_fix_dark_water=True` (все датасеты) — background на тёмных изображениях (mean < 0.24, std < 0.18) → water. Satellite GTA-UAV: bg 57% → 5%.
|
||||
- `seg_reclassify_wetland=True` (только GTA-UAV) — wetland → vegetation (зелёный) / bare soil (коричневый). Drone: ложный wetland 14% → 0%.
|
||||
|
||||
> Подробный анализ: [`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** проверяет существование файлов в соответствующих директориях модальностей. Пайплайн можно прервать 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] | Dense depth map, непрерывная [0, 1], per-frame normalized |
|
||||
| `edge` | float16 | [1, H, W] | Dense edge map (Sobel magnitude), [0, 1] |
|
||||
| `chm` | float16 | [1, H, W] | Dense canopy height map, [0, 1], per-frame normalized |
|
||||
| `segm` | uint8 | [1, H, W] | Dense class ID map, значения [0, 16] (17 unified классов) |
|
||||
|
||||
Преимущества 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) | Black |
|
||||
| 1 | building | (220, 40, 40) | Red |
|
||||
| 2 | road | (160, 160, 160) | Gray |
|
||||
| 3 | vegetation | (30, 180, 30) | Green |
|
||||
| 4 | water | (30, 120, 220) | Blue |
|
||||
| 5 | sand and gravel ground | (180, 140, 80) | Tan |
|
||||
| 6 | rocky terrain | (120, 100, 80) | Brown |
|
||||
| 7 | farmland | (200, 200, 50) | Yellow |
|
||||
| 8 | railway | (100, 60, 120) | Purple |
|
||||
| 9 | parking lot | (255, 165, 0) | Orange |
|
||||
| 10 | sidewalk | (200, 200, 200) | Light gray |
|
||||
| 11 | bare soil | (140, 100, 50) | Dark tan |
|
||||
| 12 | rooftop | (180, 60, 60) | Dark red |
|
||||
| 13 | sports field | (50, 200, 150) | Teal |
|
||||
| 14 | muddy/wetland | (80, 100, 70) | Olive |
|
||||
| 15 | embankment | (170, 130, 100) | Sandy brown |
|
||||
| 16 | swimming pool | (0, 200, 255) | Cyan |
|
||||
|
||||
## Использование для обучения
|
||||
|
||||
Все модальности хранятся как **dense tensor maps** — прямые тензоры, не RGB-рендеры. Это ключевое решение (см. [dialog_fusion_modalities](docs/segmentation_class_analysis.md)): тензоры сохраняют полную информацию без потерь при квантовании/colormapping и не являются OOD-входом для DINOv3.
|
||||
|
||||
### SafeTensors (рекомендуемый способ)
|
||||
|
||||
```python
|
||||
from safetensors.torch import load_file
|
||||
import torch.nn.functional as F
|
||||
|
||||
# Zero-copy чтение всех модальностей за ~0.1ms
|
||||
data = load_file("World-UAV-aug/safetensors/Rot/.../crop_12_4.safetensors")
|
||||
|
||||
# Все модальности — dense spatial maps, готовые для injection в backbone
|
||||
depth = data["depth"] # [1, H, W] float16, непрерывная глубина [0, 1]
|
||||
edge = data["edge"] # [1, H, W] float16, Sobel magnitude [0, 1]
|
||||
chm = data["chm"] # [1, H, W] float16, canopy height [0, 1]
|
||||
segm = data["segm"] # [1, H, W] uint8, dense class ID map [0, 16]
|
||||
```
|
||||
|
||||
### Подача в Teacher NADEZHDA
|
||||
|
||||
Каждая модальность подаётся в свой lightweight aux-encoder, затем через FiLM/Conv1x1 injection в DINOv3 patch tokens:
|
||||
|
||||
```python
|
||||
# Depth / Edge / CHM → [B, 1, H, W] float → Conv aux-encoder → FiLM injection
|
||||
# Прямые тензоры, НЕ RGB-рендеры (turbo colormap = потеря 70% информации + OOD)
|
||||
aux_depth = depth_encoder(depth.float()) # [1, H, W] → [C, H, W]
|
||||
aux_edge = edge_encoder(edge.float())
|
||||
aux_chm = chm_encoder(chm.float())
|
||||
|
||||
# Segmentation → dense class ID map → per-class embedding → spatial feature map
|
||||
# Вариант 1: one-hot → Conv
|
||||
segm_onehot = F.one_hot(segm.long().squeeze(0), num_classes=17) # [H, W, 17]
|
||||
segm_features = seg_conv(segm_onehot.permute(2, 0, 1).float()) # [17, H, W] → [C, H, W]
|
||||
|
||||
# Вариант 2: learned per-class embedding (SegAuxEncoder)
|
||||
# seg_emb = nn.Embedding(17, 32)
|
||||
# segm_features = seg_emb(segm.long().squeeze(0)).permute(2, 0, 1) # [H, W] → [32, H, W]
|
||||
```
|
||||
|
||||
### Почему тензоры, а не RGB-рендеры
|
||||
|
||||
| Формат | Пример depth | Потеря информации | Для DINOv3 |
|
||||
|---|---|---|---|
|
||||
| `float16` тензор (хранится) | `[0.4231, 0.4235, ...]` | ~0% | Прямой вход в aux-encoder |
|
||||
| `uint8` grayscale PNG | `[108, 108, ...]` | ~0.4% | Приемлемо |
|
||||
| `turbo colormap` RGB PNG | `[R=50, G=180, B=220]` | **~70%** | **OOD** — DINOv3 обучен на натуральных RGB |
|
||||
|
||||
> Для обучения **всегда** используйте SafeTensors. PNG визуализации — только для просмотра в Obsidian/файловом менеджере.
|
||||
|
||||
### Миграция со старого формата
|
||||
|
||||
Если данные сгенерированы в старом prefix-формате (`crop_12_4_depth.png`), мигрируйте:
|
||||
|
||||
```bash
|
||||
# Сначала проверить (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 суммарно):
|
||||
|
||||
```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 (17 промптов x per-image) = ~84% времени инференса. Text embeddings кэшируются. Batch size backbone = 16
|
||||
- **Post-processing сегментации** -- dark water fix (background → water для тёмных изображений) + wetland reclassify (GTA-UAV: wetland → vegetation/bare soil)
|
||||
- **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, 17 prompts) | ~120 ч | **~76%** |
|
||||
| CHMv2 | ~8.5 ч | ~8% |
|
||||
| Consolidate (.safetensors) | ~0.1 ч | <1% |
|
||||
| **Итого** | **~144 ч (~6 дней)** | |
|
||||
|
||||
> При обработке только DB (спутник, `source='db'`): ~486K изображений, ~50 ч.
|
||||
> При обработке только query (БПЛА, `source='query'`): ~486K изображений, ~50 ч.
|
||||
|
||||
## Тесты
|
||||
|
||||
```bash
|
||||
# Все тесты (149 штук, ~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) | Unified 17 классов: анализ World-UAV (392 локации), UAV_VisLoc, GTA-UAV |
|
||||
| [`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/` -- отдельная установка не требуется.
|
||||
@@ -0,0 +1,7 @@
|
||||
# Hardware profile: GPU, precision, batch size
|
||||
HardwareConfig.profile_name = 'rtx4090'
|
||||
HardwareConfig.total_ram_gb = 24.0
|
||||
HardwareConfig.reserve_gb = 2.0
|
||||
HardwareConfig.use_fp16 = True
|
||||
HardwareConfig.batch_size = None
|
||||
HardwareConfig.num_workers = 4
|
||||
@@ -0,0 +1,7 @@
|
||||
# Image preprocessing parameters
|
||||
InputConfig.image_size = 256 # DB (satellite) resolution
|
||||
InputConfig.query_image_size = 512 # Query (drone) resolution
|
||||
InputConfig.sobel_kernel_size = 3
|
||||
InputConfig.edge_normalize = True
|
||||
InputConfig.imagenet_mean = [0.485, 0.456, 0.406]
|
||||
InputConfig.imagenet_std = [0.229, 0.224, 0.225]
|
||||
@@ -0,0 +1,9 @@
|
||||
# Model identifiers and fallback strategies
|
||||
ModelsConfig.depth_model_id = 'DA3-LARGE-1.1'
|
||||
ModelsConfig.depth_fallback_id = 'depth-anything/Depth-Anything-V2-Large-hf'
|
||||
ModelsConfig.chmv2_model_id = 'facebook/dinov3-vitl16-chmv2-dpt-head'
|
||||
ModelsConfig.seg_model_type = 'segearth-ov3'
|
||||
ModelsConfig.seg_fallback_type = 'segformer-b5'
|
||||
ModelsConfig.seg_fallback_id = 'nvidia/segformer-b5-finetuned-ade-640-640'
|
||||
# Local directory for downloading and caching model weights (leave empty for HF default cache)
|
||||
ModelsConfig.weights_dir = 'in/weights'
|
||||
@@ -0,0 +1,14 @@
|
||||
# Pipeline configuration: what to process and where to save
|
||||
PipelineConfig.input_root = '/mnt/data1tb/cvgl_datasets/UAV-GeoLoc'
|
||||
PipelineConfig.output_root = '/mnt/data1tb/cvgl_datasets/World-UAV-aug'
|
||||
PipelineConfig.stages = ['depth', 'edges', 'segmentation', 'chmv2']
|
||||
PipelineConfig.save_npy = False
|
||||
PipelineConfig.save_vis = True
|
||||
PipelineConfig.save_concat = False
|
||||
PipelineConfig.save_safetensors = True
|
||||
PipelineConfig.cleanup_npy = False
|
||||
PipelineConfig.resume = True
|
||||
PipelineConfig.subset = None
|
||||
# Source filter: 'db' = satellite only, 'query' = drone/UAV only, None = both
|
||||
PipelineConfig.source = 'query' #'db'
|
||||
PipelineConfig.log_level = 'INFO'
|
||||
@@ -0,0 +1,23 @@
|
||||
# Unified 17-class open-vocabulary segmentation (shared across all datasets)
|
||||
# See scripts/seg_classes.py for canonical source, docs/segmentation_class_analysis.md for rationale
|
||||
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
|
||||
'bare soil and plowed field', # 11 — plowed fields, construction sites
|
||||
'roof and rooftop', # 12 — rooftops, solar panels
|
||||
'sports field and playground', # 13 — courts, pitches
|
||||
'muddy ground and wetland', # 14 — wet soil, marshes, levee banks
|
||||
'embankment and levee', # 15 — earthen dams, canal walls
|
||||
'swimming pool', # 16 — pools (GTA-UAV suburbs)
|
||||
]
|
||||
SegConfig.threshold = 0.15
|
||||
SegConfig.default_resolution = 1008
|
||||
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run annotation pipeline for GTA-UAV-LR dataset.
|
||||
|
||||
GTA-UAV-LR: synthetic dataset from GTA V engine.
|
||||
- drone/images/: 33763 images, 512x384, RGB PNG
|
||||
- satellite/: 14640 images, 256x256, RGBA PNG (alpha = map boundary)
|
||||
- Total: 48403 images
|
||||
- 6 flight heights: 100, 200, 300, 400, 500, 600 meters
|
||||
|
||||
Usage:
|
||||
python scripts/run_gta_uav.py
|
||||
python scripts/run_gta_uav.py --source db # only satellite (14.6K)
|
||||
python scripts/run_gta_uav.py --source drone # only drone (33.8K)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from src.conf.hardware_conf import HardwareConfig
|
||||
from src.conf.input_conf import InputConfig
|
||||
from src.conf.models_conf import ModelsConfig
|
||||
from src.conf.pipeline_conf import PipelineConfig
|
||||
from src.conf.seg_conf import SegConfig
|
||||
from src.augmentor.io_utils import setup_logging
|
||||
from src.main import run_pipeline
|
||||
from scripts.seg_classes import UNIFIED_PROMPTS
|
||||
|
||||
|
||||
INPUT_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"
|
||||
OUTPUT_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-aug"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Annotate GTA-UAV-LR")
|
||||
parser.add_argument("--source", choices=["db", "drone", "all"], default="all",
|
||||
help="Process only db (satellite), drone, or all (default)")
|
||||
parser.add_argument("--stages", nargs="+",
|
||||
default=["depth", "edges", "segmentation", "chmv2"],
|
||||
help="Stages to run")
|
||||
args = parser.parse_args()
|
||||
|
||||
import gin
|
||||
gin.clear_config()
|
||||
|
||||
source = None if args.source == "all" else args.source
|
||||
if source == "drone":
|
||||
source = "query"
|
||||
|
||||
pipeline_conf = PipelineConfig(
|
||||
input_root=INPUT_ROOT,
|
||||
output_root=OUTPUT_ROOT,
|
||||
stages=args.stages,
|
||||
save_npy=False,
|
||||
save_vis=True,
|
||||
save_safetensors=True,
|
||||
cleanup_npy=True,
|
||||
seg_fix_dark_water=True,
|
||||
seg_reclassify_wetland=True,
|
||||
resume=True,
|
||||
source=source,
|
||||
log_level="INFO",
|
||||
)
|
||||
|
||||
hw_conf = HardwareConfig(
|
||||
profile_name="rtx4090",
|
||||
total_ram_gb=24.0,
|
||||
reserve_gb=2.0,
|
||||
use_fp16=True,
|
||||
batch_size=None,
|
||||
num_workers=4,
|
||||
)
|
||||
|
||||
# GTA-UAV: satellite 256x256, drone 512x384
|
||||
# Use 256 for satellite, 512 for drone (non-square → resize to square)
|
||||
input_conf = InputConfig(image_size=256, query_image_size=512)
|
||||
|
||||
# GTA V synthetic scenes: urban, suburban, rural, coastal, mountainous
|
||||
# 11 base classes + pool (swimming pools common in GTA suburbs)
|
||||
seg_conf = SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS)
|
||||
|
||||
models_conf = ModelsConfig(weights_dir=str(_PROJECT_ROOT / "in" / "weights"))
|
||||
|
||||
setup_logging(
|
||||
pipeline_conf.log_level,
|
||||
log_file=Path(OUTPUT_ROOT) / "pipeline.log",
|
||||
)
|
||||
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Unified segmentation classes shared across all datasets.
|
||||
|
||||
All datasets MUST use the same prompt list and class IDs to enable
|
||||
transfer learning (e.g., pretrain on GTA-UAV → fine-tune on UAV_VisLoc).
|
||||
|
||||
Not every dataset will have pixels for every class — that's fine.
|
||||
A class with 0 pixels simply won't contribute to training loss.
|
||||
"""
|
||||
|
||||
UNIFIED_PROMPTS: list[str] = [
|
||||
"background", # 0
|
||||
"building", # 1
|
||||
"road", # 2
|
||||
"vegetation", # 3
|
||||
"water", # 4
|
||||
"sand and gravel ground", # 5
|
||||
"rocky terrain", # 6
|
||||
"farmland", # 7
|
||||
"railway", # 8
|
||||
"parking lot", # 9
|
||||
"sidewalk", # 10
|
||||
"bare soil and plowed field", # 11
|
||||
"roof and rooftop", # 12
|
||||
"sports field and playground", # 13
|
||||
"muddy ground and wetland", # 14
|
||||
"embankment and levee", # 15
|
||||
"swimming pool", # 16
|
||||
]
|
||||
|
||||
NUM_CLASSES = len(UNIFIED_PROMPTS) # 17
|
||||
@@ -0,0 +1,375 @@
|
||||
"""Batched inference functions for depth (DA3), edges (Sobel), and segmentation.
|
||||
|
||||
All functions accept explicit parameters — no global config imports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_IMGNET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
|
||||
_IMGNET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
|
||||
_cached_device: torch.device | None = None
|
||||
_cached_mean: torch.Tensor | None = None
|
||||
_cached_std: torch.Tensor | None = None
|
||||
|
||||
|
||||
def _get_imgnet_stats(
|
||||
device: torch.device, dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Return cached ImageNet mean/std on device."""
|
||||
global _cached_device, _cached_mean, _cached_std
|
||||
if _cached_device != device:
|
||||
_cached_mean = _IMGNET_MEAN.to(device, dtype=dtype)
|
||||
_cached_std = _IMGNET_STD.to(device, dtype=dtype)
|
||||
_cached_device = device
|
||||
return _cached_mean, _cached_std # type: ignore[return-value]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Depth
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@torch.inference_mode()
|
||||
def infer_depth_batch(
|
||||
model: nn.Module,
|
||||
images_raw: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Run depth estimation on a batch.
|
||||
|
||||
Args:
|
||||
images_raw: [B, 3, H, W] float32 [0, 1].
|
||||
|
||||
Returns:
|
||||
depth: [B, 1, H, W] float32 [0, 1] (per-frame normalized).
|
||||
"""
|
||||
B, _, H, W = images_raw.shape
|
||||
|
||||
# DA3 API: model.inference([PIL/np images]) → Prediction with .depth [N, H, W].
|
||||
if hasattr(model, "inference"):
|
||||
img_list = [
|
||||
Image.fromarray(
|
||||
(images_raw[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
|
||||
)
|
||||
for i in range(B)
|
||||
]
|
||||
prediction = model.inference(img_list, process_res=H)
|
||||
depth_np = np.asarray(prediction.depth) # [N, H', W']
|
||||
depths_list = []
|
||||
for i in range(B):
|
||||
d = torch.from_numpy(depth_np[i]).float()
|
||||
if d.ndim == 2:
|
||||
d = d.unsqueeze(0)
|
||||
if d.shape[-2:] != (H, W):
|
||||
d = F.interpolate(
|
||||
d.unsqueeze(0), size=(H, W), mode="bilinear", align_corners=False,
|
||||
).squeeze(0)
|
||||
d_min, d_max = d.min(), d.max()
|
||||
d = (d - d_min) / (d_max - d_min + 1e-8)
|
||||
depths_list.append(d)
|
||||
return torch.stack(depths_list)
|
||||
|
||||
# DA V2 fallback (transformers API).
|
||||
mean, std = _get_imgnet_stats(device, images_raw.dtype)
|
||||
x = (images_raw.to(device) - mean) / std
|
||||
if next(model.parameters()).dtype == torch.float16:
|
||||
x = x.half()
|
||||
|
||||
pred = model(pixel_values=x).predicted_depth
|
||||
depth = F.interpolate(
|
||||
pred.unsqueeze(1).float(), size=(H, W), mode="bilinear", align_corners=False,
|
||||
)
|
||||
for i in range(B):
|
||||
d = depth[i]
|
||||
d_min, d_max = d.min(), d.max()
|
||||
depth[i] = (d - d_min) / (d_max - d_min + 1e-8)
|
||||
return depth.cpu()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CHMv2 (DINOv3 cross-view height map)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@torch.inference_mode()
|
||||
def infer_chmv2_batch(
|
||||
model: nn.Module,
|
||||
processor: Any,
|
||||
images_raw: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Run CHMv2 depth estimation on a batch.
|
||||
|
||||
Args:
|
||||
model: CHMv2ForDepthEstimation.
|
||||
processor: CHMv2ImageProcessor (used for post-processing).
|
||||
images_raw: [B, 3, H, W] float32 [0, 1].
|
||||
|
||||
Returns:
|
||||
depth: [B, 1, H, W] float32 [0, 1] (per-frame normalized).
|
||||
"""
|
||||
B, _, H, W = images_raw.shape
|
||||
|
||||
# Convert to PIL for processor
|
||||
pil_images = [
|
||||
Image.fromarray(
|
||||
(images_raw[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
|
||||
)
|
||||
for i in range(B)
|
||||
]
|
||||
inputs = processor(images=pil_images, return_tensors="pt")
|
||||
# CHMv2 must run in FP32 (NaN in FP16).
|
||||
pixel_values = inputs["pixel_values"].to(device, dtype=torch.float32)
|
||||
|
||||
outputs = model(pixel_values=pixel_values)
|
||||
|
||||
# Post-process to get depth maps at original resolution
|
||||
target_sizes = [(H, W)] * B
|
||||
results = processor.post_process_depth_estimation(
|
||||
outputs, target_sizes=target_sizes,
|
||||
)
|
||||
|
||||
depths_list = []
|
||||
for r in results:
|
||||
d = r["predicted_depth"].float()
|
||||
if d.ndim == 2:
|
||||
d = d.unsqueeze(0)
|
||||
d_min, d_max = d.min(), d.max()
|
||||
d = (d - d_min) / (d_max - d_min + 1e-8)
|
||||
depths_list.append(d)
|
||||
return torch.stack(depths_list).cpu()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Edges
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_SOBEL_X = torch.tensor(
|
||||
[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32,
|
||||
).view(1, 1, 3, 3) / 8.0
|
||||
_SOBEL_Y = torch.tensor(
|
||||
[[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32,
|
||||
).view(1, 1, 3, 3) / 8.0
|
||||
|
||||
|
||||
def compute_edges_from_depth(depth: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute structural edges from depth via Sobel filters.
|
||||
|
||||
Args:
|
||||
depth: [B, 1, H, W] float32 [0, 1].
|
||||
|
||||
Returns:
|
||||
edges: [B, 1, H, W] float32 [0, 1] (edge magnitude).
|
||||
"""
|
||||
depth_padded = F.pad(depth, (1, 1, 1, 1), mode="replicate")
|
||||
dz_dx = F.conv2d(depth_padded, _SOBEL_X)
|
||||
dz_dy = F.conv2d(depth_padded, _SOBEL_Y)
|
||||
edges = torch.sqrt(dz_dx ** 2 + dz_dy ** 2)
|
||||
|
||||
for i in range(edges.shape[0]):
|
||||
e = edges[i]
|
||||
e_max = e.max()
|
||||
if e_max > 0:
|
||||
edges[i] = e / e_max
|
||||
return edges
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Segmentation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@torch.inference_mode()
|
||||
def infer_segmentation_batch(
|
||||
model: Any,
|
||||
seg_config: dict[str, Any],
|
||||
images_raw: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Run semantic segmentation on a batch.
|
||||
|
||||
Args:
|
||||
model: SegEarth-OV3 pipeline or SegFormer nn.Module.
|
||||
seg_config: Must contain 'type' and optionally 'prompts', 'processor'.
|
||||
images_raw: [B, 3, H, W] float32 [0, 1].
|
||||
|
||||
Returns:
|
||||
seg_ids: [B, 1, H, W] uint8.
|
||||
"""
|
||||
if seg_config.get("type") == "segearth-ov3":
|
||||
prompts = seg_config.get("prompts", [])
|
||||
return _infer_segearth_ov3(model, images_raw, prompts)
|
||||
else:
|
||||
processor = seg_config["processor"]
|
||||
proc_mean = torch.tensor(processor.image_mean).view(1, 3, 1, 1)
|
||||
proc_std = torch.tensor(processor.image_std).view(1, 3, 1, 1)
|
||||
return _infer_segformer(model, proc_mean, proc_std, images_raw, device)
|
||||
|
||||
|
||||
def _infer_segearth_ov3(
|
||||
model: Any,
|
||||
images_raw: torch.Tensor,
|
||||
prompts: list[str],
|
||||
) -> torch.Tensor:
|
||||
"""Run SegEarth-OV3 on a batch of images.
|
||||
|
||||
Uses model.predict_pil_batch() for batched backbone inference when available,
|
||||
falls back to per-image predict_pil().
|
||||
"""
|
||||
B, _, H, W = images_raw.shape
|
||||
|
||||
# Convert all tensors to PIL up front
|
||||
pil_images = []
|
||||
for i in range(B):
|
||||
img_np = (images_raw[i].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
|
||||
pil_images.append(Image.fromarray(img_np))
|
||||
|
||||
ori_shape = (H, W)
|
||||
|
||||
# Batched backbone path — chunk into sub-batches of up to 16
|
||||
_MAX_SEG_BATCH = 16
|
||||
if hasattr(model, "predict_pil_batch"):
|
||||
try:
|
||||
seg_list = []
|
||||
for start in range(0, B, _MAX_SEG_BATCH):
|
||||
chunk = pil_images[start:start + _MAX_SEG_BATCH]
|
||||
chunk_results = model.predict_pil_batch(
|
||||
chunk, ori_shapes=[ori_shape] * len(chunk),
|
||||
)
|
||||
for seg_pred, _ in chunk_results:
|
||||
t = seg_pred.cpu().to(torch.uint8)
|
||||
if t.ndim == 2:
|
||||
t = t.unsqueeze(0)
|
||||
seg_list.append(t)
|
||||
return torch.stack(seg_list)
|
||||
except Exception as exc:
|
||||
logger.warning("⚠️ predict_pil_batch failed, falling back to per-image: %s", exc)
|
||||
|
||||
# Fallback: per-image inference
|
||||
seg_list = []
|
||||
for i, pil_img in enumerate(pil_images):
|
||||
try:
|
||||
if hasattr(model, "predict_pil"):
|
||||
seg_pred, _ = model.predict_pil(pil_img, ori_shape=ori_shape)
|
||||
t = seg_pred.cpu().to(torch.uint8)
|
||||
elif hasattr(model, "predict"):
|
||||
seg_map = model.predict(pil_img, text_prompts=prompts)
|
||||
seg_np = np.asarray(seg_map).squeeze()
|
||||
t = torch.from_numpy(seg_np.astype(np.uint8))
|
||||
else:
|
||||
raise AttributeError(f"Unknown SegEarth API: {type(model)}")
|
||||
|
||||
if t.ndim == 2 and t.shape != (H, W):
|
||||
t = F.interpolate(
|
||||
t.float().unsqueeze(0).unsqueeze(0), size=(H, W), mode="nearest",
|
||||
).squeeze(0).squeeze(0).to(torch.uint8)
|
||||
if t.ndim == 2:
|
||||
t = t.unsqueeze(0)
|
||||
seg_list.append(t)
|
||||
except Exception as exc:
|
||||
logger.warning("⚠️ SegEarth-OV3 failed on image %d: %s", i, exc)
|
||||
seg_list.append(torch.zeros(1, H, W, dtype=torch.uint8))
|
||||
return torch.stack(seg_list)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _infer_segformer(
|
||||
model: nn.Module,
|
||||
proc_mean: torch.Tensor,
|
||||
proc_std: torch.Tensor,
|
||||
images_raw: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Run SegFormer-B5 on a batch (fallback)."""
|
||||
_, _, H, W = images_raw.shape
|
||||
mean = proc_mean.to(device)
|
||||
std = proc_std.to(device)
|
||||
x = (images_raw.to(device) - mean) / std
|
||||
if next(model.parameters()).dtype == torch.float16:
|
||||
x = x.half()
|
||||
logits = model(pixel_values=x).logits
|
||||
upsampled = F.interpolate(
|
||||
logits.float(), size=(H, W), mode="bilinear", align_corners=False,
|
||||
)
|
||||
return upsampled.argmax(dim=1, keepdim=True).cpu().to(torch.uint8)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Post-processing heuristics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def postprocess_segmentation(
|
||||
seg_ids: torch.Tensor,
|
||||
images_raw: torch.Tensor,
|
||||
water_class: int = 4,
|
||||
dark_water_mean_thr: float = 0.24,
|
||||
dark_water_std_thr: float = 0.18,
|
||||
reclassify_wetland: bool = False,
|
||||
wetland_class: int = 14,
|
||||
vegetation_class: int = 3,
|
||||
bare_soil_class: int = 11,
|
||||
green_thr: float = 0.35,
|
||||
) -> torch.Tensor:
|
||||
"""Fix known segmentation failure modes with simple heuristics.
|
||||
|
||||
Applied per-image after model inference.
|
||||
|
||||
Rules:
|
||||
1. Dark water: if a background (0) image has mean_rgb < dark_water_mean_thr
|
||||
and std < dark_water_std_thr, reclassify all background pixels as water.
|
||||
2. Wetland reclassification (optional, for GTA-UAV): reclassify wetland pixels
|
||||
by local color — green-dominant → vegetation, else → bare soil.
|
||||
|
||||
Args:
|
||||
seg_ids: [B, 1, H, W] uint8 class IDs.
|
||||
images_raw: [B, 3, H, W] float32 [0, 1] original RGB.
|
||||
water_class: class ID for water.
|
||||
dark_water_mean_thr: mean RGB threshold (0-1) below which bg → water.
|
||||
dark_water_std_thr: std threshold below which bg → water.
|
||||
reclassify_wetland: if True, split wetland into vegetation/bare_soil.
|
||||
wetland_class: class ID for muddy/wetland.
|
||||
vegetation_class: class ID for vegetation.
|
||||
bare_soil_class: class ID for bare soil.
|
||||
green_thr: green-channel ratio threshold for vegetation vs bare soil.
|
||||
|
||||
Returns:
|
||||
seg_ids: [B, 1, H, W] uint8, corrected in-place.
|
||||
"""
|
||||
B = seg_ids.shape[0]
|
||||
seg = seg_ids.clone()
|
||||
|
||||
for i in range(B):
|
||||
s = seg[i, 0] # [H, W] uint8
|
||||
rgb = images_raw[i] # [3, H, W] float32
|
||||
|
||||
# Rule 1: dark uniform images → background becomes water
|
||||
bg_mask = s == 0
|
||||
if bg_mask.any():
|
||||
bg_pixels = rgb[:, bg_mask] # [3, N]
|
||||
mean_val = bg_pixels.mean().item()
|
||||
std_val = bg_pixels.std().item()
|
||||
if mean_val < dark_water_mean_thr and std_val < dark_water_std_thr:
|
||||
s[bg_mask] = water_class
|
||||
|
||||
# Rule 2: reclassify wetland by color
|
||||
if reclassify_wetland:
|
||||
wet_mask = s == wetland_class
|
||||
if wet_mask.any():
|
||||
g = rgb[1, wet_mask] # green channel
|
||||
r = rgb[0, wet_mask]
|
||||
# Green-dominant → vegetation, else → bare soil
|
||||
is_green = g > (r + green_thr * 0.5)
|
||||
reclassed = torch.where(is_green, vegetation_class, bare_soil_class)
|
||||
s[wet_mask] = reclassed.to(torch.uint8)
|
||||
|
||||
seg[i, 0] = s
|
||||
|
||||
return seg
|
||||
@@ -0,0 +1,400 @@
|
||||
"""I/O utilities: saving depth / edges / segmentation / safetensors.
|
||||
|
||||
Directory-based output layout — modality determines the folder, not file suffix:
|
||||
|
||||
output_root/
|
||||
├── depth/{rel_parent}/{stem}.png # vis
|
||||
├── edge/{rel_parent}/{stem}.png
|
||||
├── segm/{rel_parent}/{stem}.png
|
||||
├── chm/{rel_parent}/{stem}.png
|
||||
├── npy/depth/{rel_parent}/{stem}.npy # intermediate float16/uint8
|
||||
├── npy/edge/{rel_parent}/{stem}.npy
|
||||
├── npy/segm/{rel_parent}/{stem}.npy
|
||||
├── npy/chm/{rel_parent}/{stem}.npy
|
||||
└── safetensors/{rel_parent}/{stem}.safetensors
|
||||
|
||||
No global config imports — all parameters passed explicitly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from safetensors.torch import save_file as _st_save_file, load_file as st_load_file
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_palette_cache: dict[int, np.ndarray] = {}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async I/O pool
|
||||
# ---------------------------------------------------------------------------
|
||||
_io_pool: ThreadPoolExecutor | None = None
|
||||
_IO_WORKERS = 4
|
||||
|
||||
|
||||
def get_io_pool() -> ThreadPoolExecutor:
|
||||
"""Return (lazily created) shared thread pool for async saves."""
|
||||
global _io_pool
|
||||
if _io_pool is None:
|
||||
_io_pool = ThreadPoolExecutor(max_workers=_IO_WORKERS)
|
||||
atexit.register(shutdown_io_pool)
|
||||
return _io_pool
|
||||
|
||||
|
||||
def shutdown_io_pool() -> None:
|
||||
"""Wait for all pending writes and shut down the pool."""
|
||||
global _io_pool
|
||||
if _io_pool is not None:
|
||||
_io_pool.shutdown(wait=True)
|
||||
_io_pool = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def vis_path(output_root: Path, modality: str, rel_parent: str, stem: str) -> Path:
|
||||
"""Build: output_root / modality / rel_parent / stem.png"""
|
||||
return output_root / modality / rel_parent / f"{stem}.png"
|
||||
|
||||
|
||||
def npy_path(output_root: Path, modality: str, rel_parent: str, stem: str) -> Path:
|
||||
"""Build: output_root / npy / modality / rel_parent / stem.npy"""
|
||||
return output_root / "npy" / modality / rel_parent / f"{stem}.npy"
|
||||
|
||||
|
||||
def safetensors_path(output_root: Path, rel_parent: str, stem: str) -> Path:
|
||||
"""Build: output_root / safetensors / rel_parent / stem.safetensors"""
|
||||
return output_root / "safetensors" / rel_parent / f"{stem}.safetensors"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Palette
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Intuitive RS segmentation palette: index → RGB.
|
||||
_FIXED_PALETTE = np.array([
|
||||
[0, 0, 0], # 0: background — black
|
||||
[220, 40, 40], # 1: building — red
|
||||
[160, 160, 160], # 2: road — gray
|
||||
[30, 180, 30], # 3: vegetation — green
|
||||
[30, 120, 220], # 4: water — blue
|
||||
[180, 140, 80], # 5: bare ground — tan
|
||||
[120, 100, 80], # 6: rock — brown
|
||||
[200, 200, 50], # 7: farmland — yellow
|
||||
[100, 60, 120], # 8: railway — purple
|
||||
[255, 165, 0], # 9: parking lot — orange
|
||||
[200, 200, 200], # 10: sidewalk — light gray
|
||||
[140, 100, 50], # 11: bare soil — dark tan
|
||||
[180, 60, 60], # 12: rooftop — dark red
|
||||
[50, 200, 150], # 13: sports field — teal
|
||||
[80, 100, 70], # 14: muddy/wetland — olive
|
||||
[170, 130, 100], # 15: embankment — sandy brown
|
||||
[0, 200, 255], # 16: pool — cyan
|
||||
], dtype=np.uint8)
|
||||
|
||||
|
||||
def make_palette(num_classes: int, seed: int = 42) -> np.ndarray:
|
||||
"""Return color palette for segmentation visualization."""
|
||||
if num_classes in _palette_cache:
|
||||
return _palette_cache[num_classes]
|
||||
if num_classes <= len(_FIXED_PALETTE):
|
||||
palette = _FIXED_PALETTE[:num_classes].copy()
|
||||
else:
|
||||
rng = np.random.RandomState(seed)
|
||||
palette = rng.randint(0, 255, (num_classes, 3), dtype=np.uint8)
|
||||
palette[:len(_FIXED_PALETTE)] = _FIXED_PALETTE
|
||||
_palette_cache[num_classes] = palette
|
||||
return palette
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Low-level atomic save
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _atomic_save_npy(arr: np.ndarray, path: Path) -> None:
|
||||
"""Write .npy atomically via temp file + rename."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp = tempfile.mkstemp(suffix=".npy", dir=path.parent)
|
||||
os.close(fd)
|
||||
try:
|
||||
np.save(tmp, arr)
|
||||
os.replace(tmp, path)
|
||||
except BaseException:
|
||||
if os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
raise
|
||||
|
||||
|
||||
_COLORMAP_CACHE: dict[str, np.ndarray] = {}
|
||||
|
||||
|
||||
def _apply_colormap(gray: np.ndarray, cmap_name: str = "turbo") -> np.ndarray:
|
||||
"""Apply matplotlib colormap to [H, W] float32 [0, 1] → [H, W, 3] uint8."""
|
||||
if cmap_name not in _COLORMAP_CACHE:
|
||||
import matplotlib.cm as cm
|
||||
cmap = cm.get_cmap(cmap_name)
|
||||
lut = (cmap(np.linspace(0, 1, 256))[:, :3] * 255).astype(np.uint8)
|
||||
_COLORMAP_CACHE[cmap_name] = lut
|
||||
lut = _COLORMAP_CACHE[cmap_name]
|
||||
idx = (gray.clip(0, 1) * 255).astype(np.uint8)
|
||||
return lut[idx]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Save float16 maps (depth, edge, chm)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _save_float16_map(
|
||||
data: torch.Tensor,
|
||||
output_root: Path,
|
||||
rel_parent: str,
|
||||
stem: str,
|
||||
modality: str,
|
||||
save_npy: bool = True,
|
||||
save_vis: bool = True,
|
||||
colormap: str | None = None,
|
||||
) -> None:
|
||||
"""Save a [1, H, W] float tensor as .npy (float16) + optional vis .png."""
|
||||
arr = data.half().numpy()
|
||||
if save_npy:
|
||||
p = npy_path(output_root, modality, rel_parent, stem)
|
||||
_atomic_save_npy(arr, p)
|
||||
if save_vis:
|
||||
gray = arr.squeeze(0).astype(np.float32)
|
||||
if colormap:
|
||||
vis = _apply_colormap(gray, colormap)
|
||||
else:
|
||||
vis = (gray * 255).clip(0, 255).astype(np.uint8)
|
||||
p = vis_path(output_root, modality, rel_parent, stem)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
Image.fromarray(vis).save(p)
|
||||
|
||||
|
||||
def save_depth(depth: torch.Tensor, output_root: Path, rel_parent: str,
|
||||
stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
|
||||
_save_float16_map(depth, output_root, rel_parent, stem, "depth", save_npy, save_vis)
|
||||
|
||||
|
||||
def save_depth_async(depth: torch.Tensor, output_root: Path, rel_parent: str,
|
||||
stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
|
||||
get_io_pool().submit(save_depth, depth.clone().cpu(), output_root, rel_parent,
|
||||
stem, save_npy, save_vis)
|
||||
|
||||
|
||||
def save_chmv2(depth: torch.Tensor, output_root: Path, rel_parent: str,
|
||||
stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
|
||||
_save_float16_map(depth, output_root, rel_parent, stem, "chm", save_npy, save_vis)
|
||||
|
||||
|
||||
def save_chmv2_async(depth: torch.Tensor, output_root: Path, rel_parent: str,
|
||||
stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
|
||||
get_io_pool().submit(save_chmv2, depth.clone().cpu(), output_root, rel_parent,
|
||||
stem, save_npy, save_vis)
|
||||
|
||||
|
||||
def save_edges(edges: torch.Tensor, output_root: Path, rel_parent: str,
|
||||
stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
|
||||
_save_float16_map(edges, output_root, rel_parent, stem, "edge", save_npy, save_vis)
|
||||
|
||||
|
||||
def save_edges_async(edges: torch.Tensor, output_root: Path, rel_parent: str,
|
||||
stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
|
||||
get_io_pool().submit(save_edges, edges.clone().cpu(), output_root, rel_parent,
|
||||
stem, save_npy, save_vis)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Save segmentation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def save_segmentation(
|
||||
seg_ids: torch.Tensor,
|
||||
output_root: Path,
|
||||
rel_parent: str,
|
||||
stem: str,
|
||||
save_npy: bool = True,
|
||||
save_vis: bool = True,
|
||||
num_classes: int = 150,
|
||||
) -> None:
|
||||
"""Save segmentation map [1, H, W] uint8."""
|
||||
arr = seg_ids.byte().numpy()
|
||||
if save_npy:
|
||||
_atomic_save_npy(arr, npy_path(output_root, "segm", rel_parent, stem))
|
||||
if save_vis:
|
||||
palette = make_palette(num_classes)
|
||||
seg_np = arr.squeeze(0).astype(np.uint8)
|
||||
seg_clamped = np.clip(seg_np, 0, num_classes - 1).astype(np.uint8)
|
||||
img = Image.fromarray(seg_clamped).convert("P")
|
||||
flat_pal = np.zeros(768, dtype=np.uint8)
|
||||
flat_pal[: num_classes * 3] = palette.flatten()
|
||||
img.putpalette(flat_pal.tolist())
|
||||
p = vis_path(output_root, "segm", rel_parent, stem)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
img.save(p)
|
||||
|
||||
|
||||
def save_segmentation_async(
|
||||
seg_ids: torch.Tensor,
|
||||
output_root: Path,
|
||||
rel_parent: str,
|
||||
stem: str,
|
||||
save_npy: bool = True,
|
||||
save_vis: bool = True,
|
||||
num_classes: int = 150,
|
||||
) -> None:
|
||||
get_io_pool().submit(
|
||||
save_segmentation, seg_ids.clone().cpu(), output_root, rel_parent,
|
||||
stem, save_npy, save_vis, num_classes,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SafeTensors: consolidate all modalities into one file per image
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_MODALITY_SPEC: dict[str, tuple[torch.dtype, str]] = {
|
||||
"depth": (torch.float16, "depth"),
|
||||
"edge": (torch.float16, "edge"),
|
||||
"chm": (torch.float16, "chm"),
|
||||
"segm": (torch.uint8, "segm"),
|
||||
}
|
||||
|
||||
|
||||
def _load_modality_tensor(
|
||||
output_root: Path, rel_parent: str, stem: str,
|
||||
modality: str, dtype: torch.dtype,
|
||||
) -> torch.Tensor | None:
|
||||
"""Load a single modality from .npy or .png, return [1, H, W] tensor or None."""
|
||||
np_p = npy_path(output_root, modality, rel_parent, stem)
|
||||
vis_p = vis_path(output_root, modality, rel_parent, stem)
|
||||
|
||||
if np_p.exists():
|
||||
arr = np.load(np_p)
|
||||
t = torch.from_numpy(arr.astype(np.float32 if dtype != torch.uint8 else np.uint8))
|
||||
if t.ndim == 2:
|
||||
t = t.unsqueeze(0)
|
||||
return t.to(dtype)
|
||||
|
||||
if vis_p.exists():
|
||||
if modality == "segm":
|
||||
pil = Image.open(vis_p)
|
||||
if pil.mode == "P":
|
||||
img = np.array(pil)
|
||||
else:
|
||||
logger.debug("Skipping %s segm.png (RGB, no class IDs).", stem)
|
||||
return None
|
||||
t = torch.from_numpy(img.astype(np.uint8))
|
||||
if t.ndim == 2:
|
||||
t = t.unsqueeze(0)
|
||||
return t
|
||||
else:
|
||||
img = np.array(Image.open(vis_p))
|
||||
arr = img.astype(np.float32) / 255.0
|
||||
if arr.ndim == 2:
|
||||
arr = arr[np.newaxis]
|
||||
elif arr.ndim == 3:
|
||||
arr = arr[:, :, 0:1].transpose(2, 0, 1)
|
||||
return torch.from_numpy(arr).to(dtype)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def consolidate_safetensors(
|
||||
output_root: Path,
|
||||
rel_parent: str,
|
||||
stem: str,
|
||||
cleanup_npy: bool = False,
|
||||
) -> bool:
|
||||
"""Bundle available modalities into one .safetensors file.
|
||||
|
||||
Returns True if the file was written, False if no modalities found.
|
||||
"""
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
npy_paths_to_clean: list[Path] = []
|
||||
|
||||
for modality, (dtype, _) in _MODALITY_SPEC.items():
|
||||
t = _load_modality_tensor(output_root, rel_parent, stem, modality, dtype)
|
||||
if t is not None:
|
||||
tensors[modality] = t
|
||||
np_p = npy_path(output_root, modality, rel_parent, stem)
|
||||
if np_p.exists():
|
||||
npy_paths_to_clean.append(np_p)
|
||||
|
||||
if not tensors:
|
||||
return False
|
||||
|
||||
st_p = safetensors_path(output_root, rel_parent, stem)
|
||||
st_p.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
fd, tmp = tempfile.mkstemp(suffix=".safetensors", dir=st_p.parent)
|
||||
os.close(fd)
|
||||
try:
|
||||
_st_save_file(tensors, tmp)
|
||||
os.replace(tmp, st_p)
|
||||
except BaseException:
|
||||
if os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
raise
|
||||
|
||||
if cleanup_npy:
|
||||
for p in npy_paths_to_clean:
|
||||
p.unlink(missing_ok=True)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def consolidate_safetensors_async(
|
||||
output_root: Path,
|
||||
rel_parent: str,
|
||||
stem: str,
|
||||
cleanup_npy: bool = False,
|
||||
) -> None:
|
||||
get_io_pool().submit(consolidate_safetensors, output_root, rel_parent,
|
||||
stem, cleanup_npy)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Logging
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def setup_logging(log_level: str = "INFO", log_file: Path | None = None) -> None:
|
||||
"""Configure root logger with coloredlogs for console + optional file handler."""
|
||||
import coloredlogs
|
||||
|
||||
fmt = "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s"
|
||||
datefmt = "%H:%M:%S"
|
||||
level = getattr(logging, log_level)
|
||||
|
||||
coloredlogs.install(
|
||||
level=level,
|
||||
fmt=fmt,
|
||||
datefmt=datefmt,
|
||||
level_styles={
|
||||
"debug": {"color": "cyan"},
|
||||
"info": {"color": "green"},
|
||||
"warning": {"color": "yellow", "bold": True},
|
||||
"error": {"color": "red", "bold": True},
|
||||
"critical": {"color": "red", "bold": True, "background": "white"},
|
||||
},
|
||||
field_styles={
|
||||
"asctime": {"color": "white", "faint": True},
|
||||
"levelname": {"color": "magenta", "bold": True},
|
||||
"name": {"color": "blue"},
|
||||
},
|
||||
)
|
||||
|
||||
if log_file is not None:
|
||||
log_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_handler = logging.FileHandler(log_file, encoding="utf-8")
|
||||
file_handler.setFormatter(logging.Formatter(fmt, datefmt=datefmt))
|
||||
logging.root.addHandler(file_handler)
|
||||
@@ -0,0 +1,241 @@
|
||||
"""Model loading / unloading for depth (DA3) and segmentation (SegEarth-OV3).
|
||||
|
||||
Model IDs and prompts come from config objects — nothing hardcoded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import logging
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from src.conf.hardware_conf import HardwareConfig
|
||||
from src.conf.models_conf import ModelsConfig
|
||||
from src.conf.seg_conf import SegConfig
|
||||
from src.utils.profiler import profile_model, log_vram_snapshot
|
||||
|
||||
import src.nn # noqa: F401 — registers vendored packages on sys.path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Depth
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _resolve_cache_dir(models_conf: ModelsConfig) -> str | None:
|
||||
"""Return absolute cache_dir for from_pretrained, or None for HF default."""
|
||||
if not models_conf.weights_dir:
|
||||
return None
|
||||
cache = Path(models_conf.weights_dir)
|
||||
if not cache.is_absolute():
|
||||
cache = Path(__file__).resolve().parents[2] / cache
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
return str(cache)
|
||||
|
||||
|
||||
def load_depth_model(
|
||||
models_conf: ModelsConfig,
|
||||
hw_conf: HardwareConfig,
|
||||
device: torch.device,
|
||||
) -> nn.Module:
|
||||
"""Load depth estimation model from config.
|
||||
|
||||
Args:
|
||||
models_conf: Model IDs from gin config.
|
||||
hw_conf: FP16 setting.
|
||||
device: Target CUDA device.
|
||||
|
||||
Returns:
|
||||
Loaded depth model on device.
|
||||
"""
|
||||
model_id = models_conf.depth_model_id
|
||||
cache_dir = _resolve_cache_dir(models_conf)
|
||||
logger.info("Loading depth model: %s (cache_dir=%s)", model_id, cache_dir)
|
||||
|
||||
try:
|
||||
from depth_anything_3.api import DepthAnything3
|
||||
|
||||
kwargs = {"cache_dir": cache_dir} if cache_dir else {}
|
||||
model = DepthAnything3.from_pretrained(f"depth-anything/{model_id}", **kwargs)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
# DA3 forward is too complex for fvcore/thop (optional kwargs, nested models).
|
||||
# Profile params + VRAM only.
|
||||
profile_model(model, None, device, model_name=f"Depth ({model_id})")
|
||||
log_vram_snapshot("after depth load")
|
||||
return model
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"⚠️ DA3 not found, falling back to %s", models_conf.depth_fallback_id,
|
||||
)
|
||||
from transformers import AutoModelForDepthEstimation
|
||||
|
||||
dtype = torch.float16 if hw_conf.use_fp16 else torch.float32
|
||||
model = AutoModelForDepthEstimation.from_pretrained(
|
||||
models_conf.depth_fallback_id, torch_dtype=dtype,
|
||||
cache_dir=cache_dir,
|
||||
).to(device).eval()
|
||||
profile_model(model, (1, 3, 256, 256), device,
|
||||
model_name=f"Depth ({models_conf.depth_fallback_id})")
|
||||
log_vram_snapshot("after depth load")
|
||||
return model
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CHMv2 (DINOv3 cross-view height map)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_chmv2_model(
|
||||
models_conf: ModelsConfig,
|
||||
hw_conf: HardwareConfig,
|
||||
device: torch.device,
|
||||
) -> tuple[nn.Module, Any]:
|
||||
"""Load CHMv2 depth model and processor.
|
||||
|
||||
Returns:
|
||||
(model, processor) tuple.
|
||||
"""
|
||||
cache_dir = _resolve_cache_dir(models_conf)
|
||||
# Prefer local weights if available.
|
||||
local_dir = Path(cache_dir) / "dinov3-chmv2" if cache_dir else None
|
||||
model_path = str(local_dir) if local_dir and local_dir.exists() else models_conf.chmv2_model_id
|
||||
|
||||
logger.info("Loading CHMv2 model: %s", model_path)
|
||||
from transformers import CHMv2ForDepthEstimation, CHMv2ImageProcessor
|
||||
|
||||
# CHMv2 (DINOv3 DPT) produces NaN in FP16 — always use FP32.
|
||||
processor = CHMv2ImageProcessor.from_pretrained(model_path)
|
||||
model = CHMv2ForDepthEstimation.from_pretrained(
|
||||
model_path, torch_dtype=torch.float32,
|
||||
).to(device).eval()
|
||||
profile_model(model, (1, 3, 518, 518), device, model_name="CHMv2 (DINOv3)")
|
||||
log_vram_snapshot("after chmv2 load")
|
||||
return model, processor
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Segmentation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_segmentation_model(
|
||||
models_conf: ModelsConfig,
|
||||
hw_conf: HardwareConfig,
|
||||
seg_conf: SegConfig,
|
||||
device: torch.device,
|
||||
) -> tuple[Any, dict[str, Any]]:
|
||||
"""Load segmentation model from config.
|
||||
|
||||
Args:
|
||||
models_conf: Model type and fallback IDs.
|
||||
hw_conf: FP16 setting.
|
||||
seg_conf: Text prompts for open-vocabulary segmentation.
|
||||
device: Target CUDA device.
|
||||
|
||||
Returns:
|
||||
(model_or_pipeline, config_dict).
|
||||
"""
|
||||
prompts = seg_conf.prompts
|
||||
logger.info("Loading segmentation (%s, %d classes) ...",
|
||||
models_conf.seg_model_type, len(prompts))
|
||||
|
||||
if models_conf.seg_model_type == "segearth-ov3":
|
||||
try:
|
||||
from segearthov3_segmentor import SegEarthOV3Segmentation
|
||||
|
||||
# Generate classname file from prompts.
|
||||
classname_file = tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".txt", delete=False,
|
||||
)
|
||||
for prompt in prompts:
|
||||
classname_file.write(f"{prompt}\n")
|
||||
classname_file.close()
|
||||
|
||||
seg_kwargs = dict(
|
||||
classname_path=classname_file.name,
|
||||
device=device,
|
||||
prob_thd=seg_conf.threshold,
|
||||
confidence_threshold=0.5,
|
||||
use_sem_seg=True,
|
||||
use_presence_score=True,
|
||||
use_transformer_decoder=True,
|
||||
)
|
||||
# Resolve bpe_path from weights_dir.
|
||||
cache_dir = _resolve_cache_dir(models_conf)
|
||||
if cache_dir:
|
||||
bpe = Path(cache_dir) / "bpe_simple_vocab_16e6.txt.gz"
|
||||
if bpe.exists():
|
||||
seg_kwargs["bpe_path"] = str(bpe)
|
||||
if cache_dir:
|
||||
# Prefer SAM 3.1 weights, fall back to SAM 3.
|
||||
sam31_ckpt = Path(cache_dir) / "sam3.1" / "sam3.1_multiplex.pt"
|
||||
sam3_ckpt = Path(cache_dir) / "sam3" / "sam3.pt"
|
||||
if sam31_ckpt.exists():
|
||||
seg_kwargs["checkpoint_path"] = str(sam31_ckpt)
|
||||
logger.info(" Using SAM3.1 checkpoint: %s", sam31_ckpt)
|
||||
elif sam3_ckpt.exists():
|
||||
seg_kwargs["checkpoint_path"] = str(sam3_ckpt)
|
||||
logger.info(" Using SAM3 checkpoint: %s", sam3_ckpt)
|
||||
model = SegEarthOV3Segmentation(**seg_kwargs)
|
||||
logger.info(" 🗺️ SegEarth-OV3 loaded. Prompts: %s", prompts)
|
||||
log_vram_snapshot("after segearth load")
|
||||
|
||||
# Clean up temp file.
|
||||
Path(classname_file.name).unlink(missing_ok=True)
|
||||
|
||||
return model, {
|
||||
"type": "segearth-ov3",
|
||||
"prompts": prompts,
|
||||
"num_classes": len(prompts),
|
||||
}
|
||||
except ImportError:
|
||||
logger.warning("⚠️ SegEarth-OV3 not found, falling back to SegFormer.")
|
||||
except Exception as exc:
|
||||
logger.warning("⚠️ SegEarth-OV3 load failed: %s, falling back.", exc)
|
||||
|
||||
# SegFormer fallback.
|
||||
from transformers import (
|
||||
SegformerForSemanticSegmentation,
|
||||
SegformerImageProcessor,
|
||||
)
|
||||
|
||||
cache_dir = _resolve_cache_dir(models_conf)
|
||||
model_id = models_conf.seg_fallback_id
|
||||
logger.info("Loading fallback: %s (cache_dir=%s)", model_id, cache_dir)
|
||||
dtype = torch.float16 if hw_conf.use_fp16 else torch.float32
|
||||
processor = SegformerImageProcessor.from_pretrained(model_id, cache_dir=cache_dir)
|
||||
model = SegformerForSemanticSegmentation.from_pretrained(
|
||||
model_id, torch_dtype=dtype, cache_dir=cache_dir,
|
||||
).to(device).eval()
|
||||
profile_model(model, (1, 3, 640, 640), device,
|
||||
model_name=f"Segmentation ({model_id})")
|
||||
log_vram_snapshot("after segformer load")
|
||||
return model, {
|
||||
"type": "segformer",
|
||||
"processor": processor,
|
||||
"num_classes": 150,
|
||||
"prompts": [],
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unload
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def unload_model(model: Any | None) -> None:
|
||||
"""Delete model and free GPU memory."""
|
||||
if model is None:
|
||||
return
|
||||
if hasattr(model, "model"):
|
||||
del model.model
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
log_vram_snapshot("after unload")
|
||||
logger.debug("🗑️ Model unloaded, CUDA cache cleared.")
|
||||
Reference in New Issue
Block a user