From 95f41a44011d42860461b85ed401104c798b306d Mon Sep 17 00:00:00 2001 From: pikaliov Date: Sat, 11 Jul 2026 17:27:51 +0300 Subject: [PATCH] Fix gin config loading + source filter, docs updates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - config_loader: add base_gin_files() to exclude dataset-specific pipeline_*.gin variants from default `python -m src.main` load (pipeline_uav_visloc.gin was overriding pipeline.gin) - main.py: use base_gin_files() in the --gin override branch too - pipeline.gin: source None (both DB + query) — was 'query', which silently dropped all DB satellite crops from discovery - README/CLAUDE/docs: sync source default, tensor format spec Co-Authored-By: Claude Opus 4.8 --- CLAUDE.md | 7 +- README.md | 46 +++++----- docs/segmentation_class_analysis.md | 68 ++++++++++++++- docs/tensor_format_spec.md | 126 ++++++++++++++++++++++++++++ in/config_files/pipeline.gin | 6 +- src/augmentor/inference.py | 53 ++++++++++-- src/conf/config_loader.py | 16 +++- src/conf/pipeline_conf.py | 2 + src/main.py | 10 ++- src/tests/test_inference.py | 78 ++++++++++++++++- 10 files changed, 377 insertions(+), 35 deletions(-) create mode 100644 docs/tensor_format_spec.md diff --git a/CLAUDE.md b/CLAUDE.md index b88f889..ebc18f3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -30,7 +30,9 @@ python -m pytest src/tests/ -v ## Ключевые решения -- **Формат выхода:** SafeTensors с **dense tensor maps** (zero-copy mmap, ~0.1ms). Все модальности — прямые тензоры (float16/uint8), не RGB-рендеры. +- **Формат выхода:** SafeTensors с **dense spatial maps [1, H, W]** — финальные predictions, не промежуточные feature maps (zero-copy mmap, ~0.1ms). Подробная спецификация: `docs/tensor_format_spec.md`. +- **Что сохраняется:** depth (relative, float16), edge (Sobel magnitude, float16), segm (class IDs argmax, uint8), chm (relative height, float16). Всё в [0,1] кроме segm [0,16]. Per-frame normalized. Captions/text embeddings НЕ сохраняются. +- **Что НЕ сохраняется:** backbone features (ViT tokens), soft probabilities/logits, абсолютные значения глубины/высоты, text embeddings промптов. - **Структура директорий:** модальность = папка (`depth/`, `edge/`, `segm/`, `chm/`, `safetensors/`), не суффикс файла. - **Стадии последовательно** — одна модель в GPU за раз (экономия VRAM). - **Сегментация:** SegEarth-OV3 (SAM 3.1 + open-vocabulary prompts). **17 unified классов** для всех датасетов (единые ID для transfer learning). @@ -49,6 +51,7 @@ src/nn/ — вендорированные DA3 + SegEarth-OV3 scripts/seg_classes.py — UNIFIED_PROMPTS (17 классов, единый источник) scripts/run_*.py — скрипты запуска для каждого датасета in/config_files/ — gin-конфиги +docs/tensor_format_spec.md — спецификация тензорных форматов (shapes, dtypes, normalization) docs/ — документация ``` @@ -60,6 +63,7 @@ docs/ — документация Ключевые флаги pipeline: - `seg_fix_dark_water=True` — автоматически исправлять тёмную воду (по умолчанию вкл.) +- `seg_fix_uniform_water=True` — исправлять однородную воду без контекста суши (по умолчанию вкл.). Два условия: (A) rgb_std < 0.10 для мутной речной воды, (B) mean_ch_std < 0.05 AND mean < 0.60 для синего океана со спутника. - `seg_reclassify_wetland=False` — переклассификация wetland в vegetation/bare soil (вкл. для GTA-UAV) ## Что НЕ делать @@ -69,3 +73,4 @@ docs/ — документация - Не рендерить 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. +- Не менять `uniform_water_ch_std_thr` (0.05) — калибровано на GTA-UAV satellite ocean. Выше 0.05 — false positives на terrain (скалы, растительность). diff --git a/README.md b/README.md index 43835fc..aab051f 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,9 @@ python -m pytest src/tests/ -v │ │ ├── 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) +│ │ ├── input.gin # image_size (256) +│ │ └── pipeline_*.gin # Датасет-специфичные оверрайды (напр. pipeline_uav_visloc.gin); +│ │ # НЕ подгружаются `python -m src.main` — только своими скриптами │ └── weights/ # Веса моделей (не в git, >50MB) │ ├── models--depth-anything--DA3-LARGE-1.1/ │ ├── sam3.1/sam3.1_multiplex.pt @@ -106,7 +108,7 @@ PipelineConfig.save_safetensors = True # True = .safetensors (для обу PipelineConfig.cleanup_npy = False # True = удалить .npy после консолидации PipelineConfig.resume = True # Пропускать уже обработанные PipelineConfig.subset = None # None=все, 'Rot', 'Country', 'Terrain' -PipelineConfig.source = 'db' # 'db' = спутник, 'query' = БПЛА, None = оба +PipelineConfig.source = None # None = оба (по умолчанию), 'db' = спутник, 'query' = БПЛА ``` ### segmentation.gin (unified 17 классов) @@ -194,14 +196,16 @@ World-UAV-aug/ ### SafeTensors (рекомендуемый формат для обучения) -Один `.safetensors` файл на изображение, содержит все модальности: +Один `.safetensors` файл на изображение, содержит все модальности как **dense spatial maps [1, H, W]** — финальные predictions моделей, не промежуточные feature maps. Подробная спецификация: [docs/tensor_format_spec.md](docs/tensor_format_spec.md). -| Ключ | 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 классов) | +| Ключ | Dtype | Shape | Описание | Модель | +|:---|:---|:---|:---|:---| +| `depth` | float16 | [1, H, W] | Relative depth, [0, 1], per-frame min-max normalized | DA3-LARGE-1.1 (DinoV2-ViTL) | +| `edge` | float16 | [1, H, W] | Sobel magnitude из depth, [0, 1], per-frame max normalized | Sobel (CPU) | +| `chm` | float16 | [1, H, W] | Relative canopy height, [0, 1], per-frame min-max normalized | DINOv3-ViTL16 CHMv2 | +| `segm` | uint8 | [1, H, W] | Dense class ID map (argmax), значения [0, 16] (17 классов) | SegEarth-OV3 (ViTDet-32L) | + +**Что НЕ сохраняется:** backbone features (ViT tokens), soft probabilities/logits сегментации, абсолютные значения глубины/высоты в метрах, text embeddings промптов, направление градиента (только magnitude). Преимущества SafeTensors: - **Zero-copy mmap** -- тензор читается прямо с диска без копирования в RAM (~0.1ms) @@ -242,7 +246,7 @@ World-UAV-aug/ ## Использование для обучения -Все модальности хранятся как **dense tensor maps** — прямые тензоры, не RGB-рендеры. Это ключевое решение (см. [dialog_fusion_modalities](docs/segmentation_class_analysis.md)): тензоры сохраняют полную информацию без потерь при квантовании/colormapping и не являются OOD-входом для DINOv3. +Все модальности хранятся как **dense spatial maps [1, H, W]** — финальные predictions моделей, не RGB-рендеры и не промежуточные feature maps. Тензоры сохраняют полную информацию без потерь при квантовании/colormapping и не являются OOD-входом для DINOv3. Полная спецификация форматов: [docs/tensor_format_spec.md](docs/tensor_format_spec.md). ### SafeTensors (рекомендуемый способ) @@ -254,33 +258,37 @@ import torch.nn.functional as F 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] +depth = data["depth"] # [1, H, W] float16, relative depth [0, 1], per-frame min-max +edge = data["edge"] # [1, H, W] float16, Sobel magnitude [0, 1], per-frame max +chm = data["chm"] # [1, H, W] float16, relative canopy height [0, 1], per-frame min-max +segm = data["segm"] # [1, H, W] uint8, class IDs (argmax) [0, 16], no normalization ``` ### Подача в Teacher NADEZHDA -Каждая модальность подаётся в свой lightweight aux-encoder, затем через FiLM/Conv1x1 injection в DINOv3 patch tokens: +Каждая модальность подаётся в свой lightweight aux-encoder, затем через FiLM injection в DINOv3 patch tokens: ```python # Depth / Edge / CHM → [B, 1, H, W] float → Conv aux-encoder → FiLM injection -# Прямые тензоры, НЕ RGB-рендеры (turbo colormap = потеря 70% информации + OOD) +# Все три уже в [0, 1] — рекомендуемый FiLM init: gamma=1, beta=0 (identity) 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 +# Segmentation → dense class ID map → embedding → spatial feature map +# Вариант 1: one-hot → Conv (deterministic init, explicit) 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) +# Вариант 2: learned per-class embedding (компактнее) # seg_emb = nn.Embedding(17, 32) # segm_features = seg_emb(segm.long().squeeze(0)).permute(2, 0, 1) # [H, W] → [32, H, W] ``` +### Нормирование входных тензоров + +Depth, edge, chm уже per-frame normalized в [0, 1]. Для FiLM init `gamma=1, beta=0` (identity) — безопасный старт. При необходимости dataset-level статистики — пройти по subset (1K-5K изображений) для running mean/std. Segmentation (uint8 class IDs) требует embedding перед FiLM — нельзя подавать как continuous input. + ### Почему тензоры, а не RGB-рендеры | Формат | Пример depth | Потеря информации | Для DINOv3 | diff --git a/docs/segmentation_class_analysis.md b/docs/segmentation_class_analysis.md index 30df4e3..068e532 100644 --- a/docs/segmentation_class_analysis.md +++ b/docs/segmentation_class_analysis.md @@ -118,7 +118,7 @@ SegConfig.default_resolution = 1008 |---|---| | `snow and ice` | SegEarth-OV3 классифицирует лёд/снег как `water` при виде сверху — промпт не работает. Проверено на Glacier (Athabasca) | | `beach` / `sand` | Покрыт `sand and gravel ground` | -| `bridge` | Слишком редко, перекрывается с `road` | +| `bridge` | Слишком редко, перекрывается с `road`. Мосты уже покрываются комбинацией `road` (2) + `water` (4) — для cross-view geolocation достаточно. Open-vocabulary SAM конфликтует между `bridge` и `road` на одних пикселях (нестабильные предсказания). Доли процента площади в аэрофото — near-zero representation, шум без вклада в loss. Добавление 18-го класса потребует перегенерации SafeTensors для всех датасетов (~1.1M изображений). Комбинация road + water на depth map — достаточно уникальная сигнатура моста | | `grassland` / `lawn` | Перекрывается с `vegetation` | | `river` / `sea` / `lake` | Всё покрыто единым `water` | | `sports field` | Слишком редко на полном датасете | @@ -333,6 +333,72 @@ SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS) --- +## Анализ текущего набора 17 классов (2026-04-18) + +### Покрытие по датасетам + +| ID | Класс | World-UAV | UAV_VisLoc | GTA-UAV | Примечания | +|:--:|:---|:--:|:--:|:--:|:---| +| 0 | background | + | + | + | Catch-all | +| 1 | building | + | + | + | Универсальный | +| 2 | road | + | + | + | Универсальный | +| 3 | vegetation | + | + | + | Универсальный | +| 4 | water | + | + | + | + лёд/снег (ограничение модели) | +| 5 | sand and gravel ground | + | + | + | Desert, Busan, пляжи | +| 6 | rocky terrain | + | — | + | Volcano, Karst, Gorge | +| 7 | farmland | + | + | + | Terrace, Plain, с/х сцены | +| 8 | railway | + | + | — | Munich, Sydney, городские сцены | +| 9 | parking lot | + | + | + | Городские сцены | +| 10 | sidewalk | + | + | + | Городские сцены | +| 11 | bare soil and plowed field | + | + | + | Главная потеря bg в UAV_VisLoc (~15-20%) | +| 12 | roof and rooftop | + | + | + | Дополняет building для вида сверху | +| 13 | sports field and playground | — | + | — | Уникальная геометрия, landmark | +| 14 | muddy ground and wetland | + | + | **reclassify** | В GTA-UAV 14% false positives → 0% | +| 15 | embankment and levee | — | + | — | Прибрежные сцены UAV_VisLoc | +| 16 | swimming pool | — | — | + | Пригороды GTA V (Vinewood/Rockford) | + +### Известные перекрытия между классами + +| Пара классов | Проблема | Критичность | +|:---|:---|:---| +| `farmland` (7) ↔ `bare soil` (11) | Поле без всходов — bare soil, с всходами — farmland. Граница размытая, SAM колеблется. | Низкая — оба покрывают «землю», для CVGL matching семантически близки | +| `sidewalk` (10) ↔ `embankment` (15) | Набережная = тротуар или дамба? Промпты конкурируют. | Низкая — embankment активен только в UAV_VisLoc | +| `building` (1) ↔ `roof` (12) | Фасад vs крыша сверху. При наклонном виде граница нечёткая. | Низкая — при виде строго сверху roof доминирует, при наклоне — building | +| `sand and gravel` (5) ↔ `bare soil` (11) | Песчаная почва — sand или soil? Зависит от контекста. | Низкая — оба покрывают грунт без растительности | + +### Проблемные классы + +1. **`muddy ground and wetland` (14)** — наиболее нестабильный. В GTA-UAV полностью рекласcифицируется (14% ложных → 0%). В World-UAV Wetland (Danube) есть в terrain, но SAM плохо различает от vegetation + bare soil. Класс с наибольшим числом false positives среди всех. + +2. **`rocky terrain` (6)** — нестабилен на красных скалах. Danxia/GrandCanyon плохо определяются как rocky terrain. Известное ограничение SegEarth-OV3, не решаемое сменой промпта. + +3. **`water` (4)** — поглощает лёд/снег. Glacier = water. Промпт `snow and ice` не работает (проверено на Athabasca). Ограничение модели, не набора классов. + +### Узкоспециализированные классы + +- **`swimming pool` (16)** — актуален только для GTA-UAV. В реальных датасетах бассейнов почти нет. Не мешает (0 пикселей = 0 вклад в loss), но занимает слот. +- **`embankment` (15)** — актуален преимущественно для UAV_VisLoc (прибрежные сцены). В World-UAV роль минимальна. +- **`sports field` (13)** — актуален преимущественно для UAV_VisLoc. В GTA-UAV стадионы закрыты сверху. + +### Отклонённые кандидаты на добавление + +| Класс | Причина отказа | Дата | +|:---|:---|:---| +| `bridge` | Покрывается комбинацией road (2) + water (4). SAM конфликтует bridge/road на одних пикселях. Доли % площади — near-zero representation. Потребовал бы перегенерации ~1.1M изображений. Комбинация road + water на depth map — достаточная сигнатура моста для CVGL. | 2026-04-18 | +| `snow and ice` | SegEarth-OV3 классифицирует лёд/снег как water при виде сверху — промпт не работает. Проверено на Glacier (Athabasca). | — | +| `highway` | Покрыт road — SegEarth не различает шоссе и обычную дорогу. | — | +| `beach` / `sand` | Покрыт sand and gravel ground. | — | +| `grassland` / `lawn` | Перекрывается с vegetation. | — | +| `courtyard` | Семантически = bare ground + building. | — | +| `fence` / `road marking` | Слишком мелкие объекты для SegEarth на 1008px. | — | +| `solar panel` | Редко, мелко, часть rooftop (12). | — | + +### Вывод + +Набор из 17 классов сбалансирован для задачи cross-view geolocation. Основные риски — перекрытия farmland/bare_soil и sidewalk/embankment, но они не критичны для matching. **Расширять набор не рекомендуется** — каждый новый класс замедляет инференс (~0.1-0.2x на промпт) и увеличивает конфликты между промптами SAM. + +--- + ## Оценка влияния на производительность - Инференс: ~1.8–2.2x медленнее (11 vs 5 text prompts) diff --git a/docs/tensor_format_spec.md b/docs/tensor_format_spec.md new file mode 100644 index 0000000..6459209 --- /dev/null +++ b/docs/tensor_format_spec.md @@ -0,0 +1,126 @@ +# Спецификация тензорных форматов + +Ответы на ключевые вопросы о формате извлечённых данных для downstream использования в NADEZHDA (cross-view geolocation). + +--- + +## 1. Формат: dense spatial maps [1, H, W] + +Все 4 модальности — **dense spatial maps** с сохранением пространственной размерности. Не pooled descriptors, не feature maps промежуточных слоёв. + +| Модальность | Key | Shape | Dtype | Диапазон | Тип данных | +|:---|:---|:---|:---|:---|:---| +| depth | `"depth"` | [1, H, W] | float16 | [0, 1] | Непрерывная relative depth | +| edge | `"edge"` | [1, H, W] | float16 | [0, 1] | Непрерывная Sobel magnitude | +| chm | `"chm"` | [1, H, W] | float16 | [0, 1] | Непрерывная relative canopy height | +| segm | `"segm"` | [1, H, W] | uint8 | [0, 16] | Дискретные class IDs (argmax) | + +**Что это значит для downstream:** + +- Depth, edge, chm — одноканальные continuous maps. Подаются напрямую в lightweight aux-encoder (Conv1x1 или shallow CNN), затем через FiLM injection в DINOv3 patch tokens. +- Segmentation — дискретная карта из 17 классов (argmax, не soft probabilities, не logits). Требует embedding layer (one-hot → Conv или `nn.Embedding(17, C)`) перед injection. +- Промежуточные feature maps моделей (backbone ViT tokens ~1024 dim) **не сохраняются**. Сохраняются только финальные predictions. Обоснование: backbone features трёх ViT-Large моделей сильно коррелируют с features основного DINOv3 в NADEZHDA — мало complementary информации. Ценность auxiliary модальностей именно в **другом представлении** (геометрия, границы, семантика), а не в дублировании ViT features. + +--- + +## 2. Какие слои cached + +### Depth (DA3-LARGE-1.1, 411M параметров) + +- **Тип:** Relative depth, не метрическая (не Metric3D) +- **Модель:** Depth Anything 3 с DinoV2-ViTL backbone (24 блока, patch 14, embed 1024) +- **Что сохраняется:** Финальный output после DualDPT head — одноканальная карта глубины +- **Нормирование:** Per-frame min-max → [0, 1]. Абсолютные значения глубины (в метрах) **теряются** при нормализации +- **Что НЕ сохраняется:** Промежуточные features из layers [11, 15, 19, 23], DPT pyramid features (256-1024 ch), confidence maps + +### Edges (Sobel из depth, CPU) + +- **Тип:** Sobel magnitude — непрерывный gradient magnitude, не Canny binary и не learned edges +- **Вычисление:** `sqrt(dz_dx² + dz_dy²)` из depth map, ядра Sobel 3×3 нормализованные (/8.0) +- **Что сохраняется:** Финальная magnitude map после per-frame max normalization → [0, 1] +- **Что НЕ сохраняется:** Направление градиента (dx, dy компоненты по отдельности), ненормализованные значения + +### Segmentation (SegEarth-OV3, SAM 3.1 backbone) + +- **Тип:** Dense class ID map (argmax), не soft probabilities и не raw logits +- **Модель:** ViTDet-32L backbone (32 блока, patch 14, embed 1024) + mask decoder + text grounding +- **Что сохраняется:** Финальный argmax из 17-class logits → uint8 class IDs [0, 16] +- **Что НЕ сохраняется:** + - Soft probabilities / logits (потеря информации о confidence и inter-class uncertainty) + - Backbone features [B, 1024, H/14, W/14] из global attention blocks [7, 15, 23, 31] + - Neck pyramid features [B, 256, H/14..H/1.75, W/14..W/1.75] + - Pixel decoder features [B, 256, H, W] + - Text embeddings от 17 промптов (вычисляются в runtime, кешируются только в GPU memory) +- **Post-processing:** Опционально применяются dark water fix и wetland reclassify **до** сохранения — результат уже post-processed + +### CHM (DINOv3-ViTL16 CHMv2, 337M параметров) + +- **Тип:** Relative canopy height, не абсолютная высота в метрах +- **Модель:** DINOv3-ViTL16 backbone (24 блока, patch 16, embed 1024) + DPT head +- **Что сохраняется:** Финальный output DPT head → per-frame min-max normalized [0, 1] +- **Что НЕ сохраняется:** Backbone features, DPT pyramid, абсолютные значения высоты +- **Ограничение:** Инференс **только FP32** (NaN в FP16), конвертация в float16 только при сохранении + +--- + +## 3. Captions и text embeddings + +**Не сохраняются.** Ни raw text, ни text embeddings. + +- **SegEarth-OV3** использует text embeddings от 17 промптов (через VETextEncoder) для open-vocabulary grounding. Эти embeddings вычисляются в runtime из строковых промптов и кешируются **только в GPU memory** во время инференса. В SafeTensors файлы не попадают. +- **Raw text промптов** хранится в `scripts/seg_classes.py:UNIFIED_PROMPTS` — это исходный код, не часть выходных данных. +- Никакие caption/description для изображений не генерируются и не сохраняются. + +**Следствия для downstream:** +- Если нужны text-conditioned fusion варианты (F0/F1/F2), text embeddings нужно вычислять заново из raw промптов в `seg_classes.py`. Это не проблема — промпты фиксированы (17 строк), embeddings вычисляются за доли секунды. +- Привязки к конкретному text encoder нет — можно использовать любой (CLIP, SigLIP, и т.д.) для вычисления embeddings из тех же промптов. + +--- + +## 4. Нормирование + +| Модальность | Нормирование | Per-sample | Диапазон | Исходные значения восстановимы? | +|:---|:---|:---:|:---|:---| +| depth | min-max: `(x - min) / (max - min + 1e-8)` | Да | [0, 1] | Нет — min/max теряются | +| edge | max-normalize: `x / max(x)` | Да | [0, 1] | Нет — max теряется | +| chm | min-max: `(x - min) / (max - min + 1e-8)` | Да | [0, 1] | Нет — min/max теряются | +| segm | Нет (class IDs) | N/A | {0, 1, ..., 16} | N/A | + +### Влияние на FiLM heads initialization + +**Depth, edge, chm (float16, [0, 1]):** +- Диапазон стабилен между изображениями (всегда [0, 1]), но **распределение внутри диапазона сильно варьируется**: пустыня — равномерное, город — бимодальное (здания vs земля), вода — сконцентрировано около одного значения. +- Рекомендуемый init: `gamma=1, beta=0` (identity FiLM) — безопасный старт, нормализация уже в [0, 1]. +- Альтернатива: предварительно вычислить dataset-level mean/std по каждой модальности для BatchNorm перед FiLM. + +**Segmentation (uint8, class IDs):** +- Дискретные значения [0, 16] — нельзя подавать как continuous input в FiLM напрямую. +- Варианты encoding перед FiLM injection: + 1. `nn.Embedding(17, C)` → lookup table → spatial [C, H, W] — компактно, learned + 2. `F.one_hot(segm, 17)` → Conv1x1 → [C, H, W] — явное, deterministic init возможен +- Для варианта 2 init: Xavier uniform для Conv1x1 weights, zero bias. + +### Статистика входных тензоров (для калибровки) + +Все три continuous модальности нормализованы per-frame в [0, 1], поэтому: +- **Global mean** по датасету ≈ 0.3–0.5 (зависит от модальности и датасета) +- **Global std** ≈ 0.2–0.35 +- Для точной калибровки рекомендуется пройти по subset (1000-5000 изображений) и вычислить running mean/std + +--- + +## Сводная таблица + +| Свойство | Depth | Edge | Segm | CHM | +|:---|:---|:---|:---|:---| +| **SafeTensors key** | `"depth"` | `"edge"` | `"segm"` | `"chm"` | +| **Shape** | [1, H, W] | [1, H, W] | [1, H, W] | [1, H, W] | +| **Dtype** | float16 | float16 | uint8 | float16 | +| **Value range** | [0, 1] | [0, 1] | [0, 16] | [0, 1] | +| **Type** | Relative depth | Sobel magnitude | Class IDs (argmax) | Relative height | +| **Model** | DA3-LARGE-1.1 | Sobel from depth | SegEarth-OV3 | DINOv3 CHMv2 | +| **Backbone** | DinoV2-ViTL (24L) | — | ViTDet (32L) | DINOv3-ViTL (24L) | +| **Normalization** | Per-frame min-max | Per-frame max | None | Per-frame min-max | +| **Dense/Pooled** | Dense spatial | Dense spatial | Dense spatial | Dense spatial | +| **Features cached** | Final output only | Final output only | Final output only | Final output only | +| **Captions/embeddings** | — | — | Not saved | — | diff --git a/in/config_files/pipeline.gin b/in/config_files/pipeline.gin index 9775c87..2efe355 100644 --- a/in/config_files/pipeline.gin +++ b/in/config_files/pipeline.gin @@ -1,6 +1,6 @@ # 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.input_root = '/media/servml/SSD_2_2TB1/datasets/cvgl_datasets/UAV-GeoLoc/' +PipelineConfig.output_root = '/media/servml/SSD_2_2TB1/datasets/cvgl_datasets/UAV-GeoLoc-aug/' PipelineConfig.stages = ['depth', 'edges', 'segmentation', 'chmv2'] PipelineConfig.save_npy = False PipelineConfig.save_vis = True @@ -10,5 +10,5 @@ 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.source = None # None = both DB (satellite) + query (drone); 'db' or 'query' to restrict PipelineConfig.log_level = 'INFO' diff --git a/src/augmentor/inference.py b/src/augmentor/inference.py index 6f07d0c..79d5711 100644 --- a/src/augmentor/inference.py +++ b/src/augmentor/inference.py @@ -312,6 +312,11 @@ def postprocess_segmentation( water_class: int = 4, dark_water_mean_thr: float = 0.24, dark_water_std_thr: float = 0.18, + fix_uniform_water: bool = True, + uniform_water_bg_thr: float = 0.80, + uniform_water_std_thr: float = 0.10, + uniform_water_ch_std_thr: float = 0.05, + uniform_water_mean_thr: float = 0.60, reclassify_wetland: bool = False, wetland_class: int = 14, vegetation_class: int = 3, @@ -323,10 +328,18 @@ def postprocess_segmentation( 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. + 1a. Dark water: if background pixels have mean_rgb < dark_water_mean_thr + and std < dark_water_std_thr, reclassify all background pixels as water. + 1b. Uniform water: if background ratio > uniform_water_bg_thr AND + (whole-image std < uniform_water_std_thr OR + mean-per-channel std < uniform_water_ch_std_thr AND + mean brightness < uniform_water_mean_thr), + reclassify background as water. + First condition catches muddy/grey river water (UAV_VisLoc). + Second condition catches blue ocean where cross-channel variance is + high but each channel is spatially uniform (GTA-UAV satellite). + 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. @@ -334,6 +347,13 @@ def postprocess_segmentation( 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. + fix_uniform_water: if True, apply Rule 1b for uniform unclassified water. + uniform_water_bg_thr: background ratio above which Rule 1b triggers. + uniform_water_std_thr: whole-image std below which Rule 1b triggers. + uniform_water_ch_std_thr: mean per-channel std below which Rule 1b triggers + (second condition, for blue ocean with high cross-channel variance). + uniform_water_mean_thr: mean brightness below which the ch_std condition + applies (filters out bright terrain like sand/concrete). reclassify_wetland: if True, split wetland into vegetation/bare_soil. wetland_class: class ID for muddy/wetland. vegetation_class: class ID for vegetation. @@ -350,7 +370,7 @@ def postprocess_segmentation( s = seg[i, 0] # [H, W] uint8 rgb = images_raw[i] # [3, H, W] float32 - # Rule 1: dark uniform images → background becomes water + # Rule 1a: dark uniform images → background becomes water bg_mask = s == 0 if bg_mask.any(): bg_pixels = rgb[:, bg_mask] # [3, N] @@ -358,6 +378,29 @@ def postprocess_segmentation( 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 1b: uniform unclassified water + elif fix_uniform_water: + bg_ratio = bg_mask.sum().item() / bg_mask.numel() + if bg_ratio > uniform_water_bg_thr: + img_std = rgb.std().item() + # Condition A: low overall std (muddy/grey river water) + is_uniform = img_std < uniform_water_std_thr + # Condition B: low per-channel std + not bright + # (blue ocean — high cross-channel variance but + # each channel is spatially uniform) + if not is_uniform: + ch_std = ( + rgb[0].std().item() + + rgb[1].std().item() + + rgb[2].std().item() + ) / 3.0 + mean_val = rgb.mean().item() + is_uniform = ( + ch_std < uniform_water_ch_std_thr + and mean_val < uniform_water_mean_thr + ) + if is_uniform: + s[bg_mask] = water_class # Rule 2: reclassify wetland by color if reclassify_wetland: diff --git a/src/conf/config_loader.py b/src/conf/config_loader.py index 436a43b..797c23f 100644 --- a/src/conf/config_loader.py +++ b/src/conf/config_loader.py @@ -15,6 +15,20 @@ from src.conf.seg_conf import SegConfig logger = logging.getLogger(__name__) +def base_gin_files(cfg_dir: Path) -> list[Path]: + """Return base .gin files, excluding dataset-specific pipeline variants. + + Files named ``pipeline_.gin`` (e.g. ``pipeline_uav_visloc.gin``) + are per-dataset overrides loaded selectively by their own run scripts — + they must NOT be globbed into the default World-UAV config, or their + ``PipelineConfig.*`` bindings would override ``pipeline.gin``. + """ + return sorted( + f for f in cfg_dir.glob("*.gin") + if not f.name.startswith("pipeline_") + ) + + def load_all_configs(path2cfg: str) -> dict[str, Any]: """Parse ALL .gin files at once and return all config objects. @@ -34,7 +48,7 @@ def load_all_configs(path2cfg: str) -> dict[str, Any]: if not cfg_dir.is_dir(): raise FileNotFoundError(f"Config directory not found: {cfg_dir}") - gin_files = sorted(cfg_dir.glob("*.gin")) + gin_files = base_gin_files(cfg_dir) if not gin_files: raise FileNotFoundError(f"No .gin files in {cfg_dir}") diff --git a/src/conf/pipeline_conf.py b/src/conf/pipeline_conf.py index 8815ee5..ebbc3e8 100644 --- a/src/conf/pipeline_conf.py +++ b/src/conf/pipeline_conf.py @@ -18,6 +18,7 @@ class PipelineConfig: save_safetensors: bool = True, cleanup_npy: bool = False, seg_fix_dark_water: bool = True, + seg_fix_uniform_water: bool = True, seg_reclassify_wetland: bool = False, resume: bool = True, subset: str | None = None, @@ -33,6 +34,7 @@ class PipelineConfig: self.save_safetensors = save_safetensors self.cleanup_npy = cleanup_npy self.seg_fix_dark_water = seg_fix_dark_water + self.seg_fix_uniform_water = seg_fix_uniform_water self.seg_reclassify_wetland = seg_reclassify_wetland self.resume = resume self.subset = subset diff --git a/src/main.py b/src/main.py index d739fc2..2cd981d 100644 --- a/src/main.py +++ b/src/main.py @@ -22,7 +22,7 @@ import torch from torch.utils.data import DataLoader from tqdm import tqdm -from src.conf.config_loader import load_all_configs +from src.conf.config_loader import base_gin_files, load_all_configs from src.conf.pipeline_conf import PipelineConfig from src.conf.hardware_conf import HardwareConfig from src.conf.models_conf import ModelsConfig @@ -317,9 +317,12 @@ def run_segmentation_stage( segs = infer_segmentation_batch( model, seg_config, batch["image_raw"], device, ) - if pipeline_conf.seg_fix_dark_water or pipeline_conf.seg_reclassify_wetland: + if (pipeline_conf.seg_fix_dark_water + or pipeline_conf.seg_fix_uniform_water + or pipeline_conf.seg_reclassify_wetland): segs = postprocess_segmentation( segs, batch["image_raw"], + fix_uniform_water=pipeline_conf.seg_fix_uniform_water, reclassify_wetland=pipeline_conf.seg_reclassify_wetland, ) for j in range(segs.shape[0]): @@ -377,6 +380,7 @@ def run_pipeline( # System profiling at startup. log_system_info() + print(Path(pipeline_conf.input_root)) log_disk_info(Path(pipeline_conf.input_root), Path(pipeline_conf.output_root)) # Discover images. @@ -532,7 +536,7 @@ def main() -> None: if args.gin: import gin as _gin cfg_dir = Path(path2cfg) - gin_files = sorted(cfg_dir.glob("*.gin")) + gin_files = base_gin_files(cfg_dir) _gin.clear_config() _gin.parse_config_files_and_bindings( config_files=[str(f) for f in gin_files], diff --git a/src/tests/test_inference.py b/src/tests/test_inference.py index d326416..81c4b98 100644 --- a/src/tests/test_inference.py +++ b/src/tests/test_inference.py @@ -239,8 +239,8 @@ class TestPostprocessSegmentation: assert (result == 4).all() # water def test_bright_background_unchanged(self) -> None: - """Bright background should NOT become water.""" - rgb = torch.full((1, 3, 32, 32), 0.6) + """Bright varied background should NOT become water.""" + rgb = torch.rand(1, 3, 32, 32) * 0.4 + 0.4 # mean ~0.6, std ~0.12 seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) result = postprocess_segmentation(seg, rgb) assert (result == 0).all() # stays background @@ -277,3 +277,77 @@ class TestPostprocessSegmentation: seg = torch.full((1, 1, 32, 32), 14, dtype=torch.uint8) result = postprocess_segmentation(seg, rgb, reclassify_wetland=False) assert (result == 14).all() + + def test_uniform_water_reclassified(self) -> None: + """Uniform muddy water (high mean, low std, high bg) → water.""" + # Muddy river: mean ~0.46, std ~0.05 — SegEarth misses without land context + rgb = torch.full((1, 3, 32, 32), 0.46) + rgb += torch.randn_like(rgb) * 0.02 + rgb.clamp_(0, 1) + seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) # all background + result = postprocess_segmentation(seg, rgb, fix_uniform_water=True) + assert (result == 4).all() # water + + def test_uniform_water_partial_bg(self) -> None: + """Uniform water with some classified pixels — bg part → water.""" + rgb = torch.full((1, 3, 32, 32), 0.46) + rgb += torch.randn_like(rgb) * 0.02 + rgb.clamp_(0, 1) + seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) + # 10% non-background (e.g. boats) + seg[0, 0, :3, :] = 1 # building (boats) + result = postprocess_segmentation(seg, rgb, fix_uniform_water=True) + assert (result[0, 0, :3, :] == 1).all() # boats stay + assert (result[0, 0, 3:, :] == 4).all() # bg → water + + def test_uniform_water_disabled(self) -> None: + """Uniform water fix disabled → bg stays.""" + rgb = torch.full((1, 3, 32, 32), 0.46) + seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) + result = postprocess_segmentation(seg, rgb, fix_uniform_water=False) + assert (result == 0).all() # stays background + + def test_uniform_water_high_std_no_fix(self) -> None: + """High-std image with bg → NOT reclassified (not uniform water).""" + rgb = torch.rand(1, 3, 32, 32) # std ~0.29 + seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) + result = postprocess_segmentation(seg, rgb, fix_uniform_water=True) + assert (result == 0).all() # stays background + + def test_uniform_water_low_bg_ratio_no_fix(self) -> None: + """Low bg ratio + low std → NOT reclassified.""" + rgb = torch.full((1, 3, 32, 32), 0.46) + rgb += torch.randn_like(rgb) * 0.02 + rgb.clamp_(0, 1) + seg = torch.full((1, 1, 32, 32), 3, dtype=torch.uint8) # all vegetation + # Only 20% background + seg[0, 0, :6, :] = 0 + result = postprocess_segmentation(seg, rgb, fix_uniform_water=True) + assert (result[0, 0, :6, :] == 0).all() # bg stays — ratio too low + + def test_blue_ocean_reclassified(self) -> None: + """Blue ocean with high cross-channel variance → water via ch_std.""" + # Blue ocean: R=0.10, G=0.26, B=0.40, per-channel std ~0.015 + rgb = torch.zeros(1, 3, 32, 32) + rgb[0, 0] = 0.10 + torch.randn(32, 32) * 0.01 # R + rgb[0, 1] = 0.26 + torch.randn(32, 32) * 0.01 # G + rgb[0, 2] = 0.40 + torch.randn(32, 32) * 0.01 # B + rgb.clamp_(0, 1) + seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) # all background + # rgb.std() is ~0.13 (cross-channel), but per-channel std ~0.01 + result = postprocess_segmentation(seg, rgb, fix_uniform_water=True) + assert (result == 4).all() # water + + def test_bright_terrain_not_reclassified(self) -> None: + """Bright terrain with low per-channel std → NOT water (mean > 0.60).""" + # Bright terrain with strong cross-channel variance: R=0.90, G=0.65, B=0.50 + # Per-channel std < 0.05 but mean ~0.68 > 0.60, so Condition B rejected. + # rgb_std ~0.16 > 0.10, so Condition A also doesn't trigger. + rgb = torch.zeros(1, 3, 32, 32) + rgb[0, 0] = 0.90 + torch.randn(32, 32) * 0.01 + rgb[0, 1] = 0.65 + torch.randn(32, 32) * 0.01 + rgb[0, 2] = 0.50 + torch.randn(32, 32) * 0.01 + rgb.clamp_(0, 1) + seg = torch.zeros(1, 1, 32, 32, dtype=torch.uint8) + result = postprocess_segmentation(seg, rgb, fix_uniform_water=True) + assert (result == 0).all() # stays background