Unify segmentation classes (17) across GTA-UAV and UAV_VisLoc

Extract shared UNIFIED_PROMPTS (17 classes, ID 0-16) into scripts/seg_classes.py for transfer learning compatibility. Both run_gta_uav.py and run_uav_visloc.py now import from it. Key change: swimming pool moved from ID 13 → ID 16, so sports field (ID 13), muddy ground (14), embankment (15) have stable IDs across both datasets. Missing classes in a dataset = 0 pixels = 0 loss. Updated: README, CLAUDE.md, segmentation_class_analysis.md, palette. Deleted old UAV_VisLoc segmentations (need regeneration with 17 classes). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-17 22:06:10 +03:00
parent 143a837c03
commit 3b5778e303
7 changed files with 88 additions and 107 deletions
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -25,7 +25,7 @@ python -m pytest src/tests/ -v
 - **Формат выхода:** SafeTensors (zero-copy mmap, ~0.1ms чтение). Не NPY, не PNG, не `.pt`.
 - **Структура директорий:** модальность = папка (`depth/`, `edge/`, `segm/`, `chm/`, `safetensors/`), не суффикс файла.
 - **Стадии последовательно** — одна модель в GPU за раз (экономия VRAM).
- **Сегментация:** SegEarth-OV3 (SAM 3.1 + open-vocabulary prompts). 11 классов для World-UAV, 16 для UAV_VisLoc, 14 для GTA-UAV.
+- **Сегментация:** SegEarth-OV3 (SAM 3.1 + open-vocabulary prompts). 11 классов для World-UAV, **17 unified** для UAV_VisLoc + GTA-UAV (единые ID для transfer learning).
 - **CHMv2 только FP32** — в FP16 NaN.

 ## Структура кода
@@ -48,6 +48,6 @@ docs/                — документация

 ## Что НЕ делать

- Не добавлять классы сегментации в `segmentation.gin` для конкретного датасета — классы специфичны для датасета, задавать в скрипте.
+- Не добавлять классы сегментации в `segmentation.gin` для конкретного датасета — использовать `scripts/seg_classes.py` (единый набор 17 классов для transfer learning).
 - Не использовать `.pt` (torch.save) для хранения тензоров — pickle, медленно.
 - Не рендерить depth/seg в RGB colormap для обучения — OOD для DINOv3, потеря информации.
--- a/README.md
+++ b/README.md
@@ -10,11 +10,13 @@
 | **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) | 11 | `python -m src.main` |
-| **UAV_VisLoc** | 81K (75K DB + 6.7K drone) | 16 | `python scripts/run_uav_visloc.py` |
-| **GTA-UAV-LR** | 48K (15K sat + 34K drone) | 14 | `python scripts/run_gta_uav.py` |
+| **World-UAV** | 973K (486K DB + 486K query) | 11 (базовые) | `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` |
+
+> UAV_VisLoc и GTA-UAV используют **единый набор 17 классов** (`scripts/seg_classes.py`) для совместимости при transfer learning.

 ## Quick Start

@@ -126,36 +128,31 @@ SegConfig.threshold = 0.15
 SegConfig.default_resolution = 1008
 ```

-### UAV_VisLoc: 16 классов (расширенные)
+### Unified 17 классов (UAV_VisLoc + GTA-UAV)

-UAV_VisLoc содержит больше сельских/прибрежных сцен. Анализ scene 01 (2496 images) показал 29% background при 11 классах. Добавлены 5 классов:
+Для transfer learning (pretrain GTA-UAV → fine-tune UAV_VisLoc) оба датасета используют **единый набор 17 классов** из `scripts/seg_classes.py`. Не каждый датасет содержит пиксели каждого класса — это нормально (0 пикселей = 0 loss).

-| ID | Промпт | Зачем |
-|:--:|:---|:---|
-| 11 | `bare soil and plowed field` | Распаханные поля без посевов — 15-20% бывшего background |
-| 12 | `roof and rooftop` | Тёмные крыши, солнечные панели — `building` ловит стены, не крыши сверху |
-| 13 | `sports field and playground` | Уникальная геометрия, сильный landmark для CVGL |
-| 14 | `muddy ground and wetland` | Болотистый грунт между водоёмами (дамбы, берега) |
-| 15 | `embankment and levee` | Линейные земляные сооружения вдоль каналов |
+| 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 | — | + | — |
+| 15 | embankment and levee | — | + | — |
+| 16 | swimming pool | — | редко | + |

-> Threshold снижен до 0.1 (vs 0.15 для World-UAV) — UAV_VisLoc содержит больше тёмных/низкоконтрастных сцен.
-
-### GTA-UAV-LR: 14 классов
-
-Синтетический датасет из GTA V. Особенности:
- Satellite в RGBA (alpha = граница карты) → автоматически конвертируется в RGB
- Drone 512x384 (не квадратные) → resize в 512x512
- 6 высот полёта: 100, 200, 300, 400, 500, 600 м
-
-Классы = 11 базовых + 3 специфических для GTA:
-
-| ID | Промпт | Зачем |
-|:--:|:---|:---|
-| 11 | `bare soil and plowed field` | Пустыри, стройплощадки (много в пригородах GTA) |
-| 12 | `roof and rooftop` | Плоские крыши, вертолётные площадки |
-| 13 | `swimming pool` | Бассейны — характерный landmark пригородов GTA |
-
-Подробный анализ выбора классов: [`docs/segmentation_class_analysis.md`](docs/segmentation_class_analysis.md)
+> UAV_VisLoc threshold=0.1, GTA-UAV threshold=0.15. Подробный анализ: [`docs/segmentation_class_analysis.md`](docs/segmentation_class_analysis.md)

 ### hardware.gin

@@ -217,7 +214,7 @@ World-UAV-aug/
 | `depth` | float16 | [1, H, W] | Карта глубины [0, 1] |
 | `edge` | float16 | [1, H, W] | Границы (Sobel) [0, 1] |
 | `chm` | float16 | [1, H, W] | Canopy height [0, 1] |
-| `segm` | uint8 | [1, H, W] | Class IDs [0, 10] |
+| `segm` | uint8 | [1, H, W] | Class IDs [0, 10] (World-UAV) / [0, 16] (unified) |

 Преимущества SafeTensors:
 - **Zero-copy mmap** -- тензор читается прямо с диска без копирования в RAM (~0.1ms)
@@ -254,6 +251,7 @@ World-UAV-aug/
 | 13 | sports field | (50, 200, 150) | UAV_VisLoc |
 | 14 | muddy/wetland | (80, 100, 70) | UAV_VisLoc |
 | 15 | embankment | (170, 130, 100) | UAV_VisLoc |
+| 16 | swimming pool | (0, 200, 255) | GTA-UAV |

 ## Использование для обучения

@@ -276,8 +274,8 @@ segm  = data["segm"]   # [1, 256, 256] uint8, class IDs [0, 10]

 # Для Teacher NADEZHDA: segm -> one-hot
 import torch.nn.functional as F
-segm_onehot = F.one_hot(segm.long().squeeze(0), num_classes=11)  # [H, W, 11]
-segm_onehot = segm_onehot.permute(2, 0, 1).float()              # [11, H, W]
+segm_onehot = F.one_hot(segm.long().squeeze(0), num_classes=17)  # [H, W, 17]
+segm_onehot = segm_onehot.permute(2, 0, 1).float()              # [17, H, W]
 ```

 ### PNG fallback (для визуализации или legacy)
--- a/docs/segmentation_class_analysis.md
+++ b/docs/segmentation_class_analysis.md
@@ -179,7 +179,15 @@ SegConfig.default_resolution = 1008

 ---

-## Конфигурация для UAV_VisLoc: 16 классов
+## Unified 17 классов для transfer learning
+
+UAV_VisLoc и GTA-UAV используют **единый набор 17 классов** (`scripts/seg_classes.py`). Это обеспечивает совместимость class IDs при transfer learning (pretrain GTA-UAV → fine-tune UAV_VisLoc). Не каждый датасет содержит пиксели каждого класса — класс с 0 пикселей просто не даёт вклад в loss.
+
+Источник единой конфигурации: `scripts/seg_classes.py:UNIFIED_PROMPTS`.
+
+---
+
+## Конфигурация для UAV_VisLoc: 17 классов (unified)

 UAV_VisLoc_processed — другой датасет (81K images, 11 сцен, DB + drone, 256x256). Анализ сегментации на scene 01 (2496 images) выявил 29% background — значительно больше, чем у World-UAV после настройки.

@@ -227,27 +235,13 @@ UAV_VisLoc содержит больше экстремальных сцен:
 - Однородные текстуры без семантических признаков
 - Это нормально — background здесь семантически корректен

-### Полная конфигурация UAV_VisLoc (16 классов)
+### Конфигурация (unified 17 классов)

 ```python
-SegConfig(threshold=0.1, prompts=[
-    "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
-])
+from scripts.seg_classes import UNIFIED_PROMPTS
+SegConfig(threshold=0.1, prompts=UNIFIED_PROMPTS)
+# UNIFIED_PROMPTS содержит 17 классов (ID 0-16)
+# UAV_VisLoc: активны все кроме 16 (swimming pool — редко)
 ```

 ### Расширенная палитра (16 классов)
@@ -315,25 +309,13 @@ GTA-UAV-LR — синтетический датасет, сгенерирова
 | `sports field` | Есть в GTA V, но редко виден сверху (стадионы закрыты) |
 | `muddy ground` | Не актуален для GTA V (нет болот в Лос-Сантосе) |

-### Конфигурация
+### Конфигурация (unified 17 классов)

 ```python
-SegConfig(threshold=0.15, prompts=[
-    "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
-    "swimming pool",               # 13
-])
+from scripts.seg_classes import UNIFIED_PROMPTS
+SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS)
+# UNIFIED_PROMPTS содержит 17 классов (ID 0-16)
+# GTA-UAV: активны все кроме 14, 15 (muddy/embankment — нет в GTA V)
 ```

 ---
--- a/scripts/run_gta_uav.py
+++ b/scripts/run_gta_uav.py
@@ -32,6 +32,7 @@ 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"
@@ -82,22 +83,7 @@ def main() -> None:

    # 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=[
-        "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
-        "swimming pool",               # 13
-    ])
+    seg_conf = SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS)

    models_conf = ModelsConfig(weights_dir=str(_PROJECT_ROOT / "in" / "weights"))

--- a/scripts/run_uav_visloc.py
+++ b/scripts/run_uav_visloc.py
@@ -27,6 +27,7 @@ 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/UAV_VisLoc_processed"
@@ -73,24 +74,7 @@ def main() -> None:
    )

    input_conf = InputConfig(image_size=256)
-    seg_conf = SegConfig(threshold=0.1, prompts=[
-        "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
-    ])
+    seg_conf = SegConfig(threshold=0.1, prompts=UNIFIED_PROMPTS)
    models_conf = ModelsConfig(weights_dir=str(_PROJECT_ROOT / "in" / "weights"))

    setup_logging(
--- a/scripts/seg_classes.py
+++ b/scripts/seg_classes.py
@@ -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
--- a/src/augmentor/io_utils.py
+++ b/src/augmentor/io_utils.py
@@ -99,6 +99,7 @@ _FIXED_PALETTE = np.array([
    [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)