# SegEarth-OV3: архитектура, pipeline и оптимизация

## Обзор

SegEarth-OV3 — модель **open-vocabulary** семантической сегментации для дистанционного зондирования, построенная на базе **SAM 3.1** (Segment Anything Model 3.1, Meta). Позволяет сегментировать изображения по произвольным текстовым описаниям классов без дообучения.

В нашем пайплайне используется для генерации семантических карт из аэрофотоснимков (дрон) и спутниковых изображений датасета World-UAV (973K изображений).

---

## Архитектура модели

### Общая схема

```
                    ┌──────────────────────────────────────────┐
                    │           SAM3VLBackbone                 │
                    │  ┌─────────────────┐  ┌───────────────┐  │
  Image ──────────►│  │  Vision Backbone │  │ Text Encoder  │◄─── Text Prompt
  [B,3,1008,1008]  │  │  (ViT + Neck)   │  │ (VE, 24-layer)│  │
                    │  └────────┬────────┘  └───────┬───────┘  │
                    │           │                    │          │
                    │     vision_features      language_features│
                    │     [B,256,72,72]        [32,1,256]       │
                    └──────────┬────────────────────┬───────────┘
                               │                    │
                               ▼                    ▼
                    ┌──────────────────────────────────────────┐
                    │      Transformer Encoder (6 layers)      │
                    │   Cross-attention: vision × language      │
                    │   Self-attention: vision features          │
                    └──────────────────┬───────────────────────┘
                                       │
                                       ▼
                    ┌──────────────────────────────────────────┐
                    │      Transformer Decoder (6 layers)      │
                    │   200 object queries                      │
                    │   Cross-attention: queries × encoder out  │
                    │   Text cross-attention                    │
                    │   Box refinement (DAC)                    │
                    │   Presence token                          │
                    └──────────┬───────────────────────────────┘
                               │
                    ┌──────────┴──────────┐
                    │                     │
                    ▼                     ▼
           ┌──────────────┐    ┌───────────────────┐
           │ Instance Head │    │ Segmentation Head  │
           │ pred_boxes    │    │ (PixelDecoder +    │
           │ pred_masks    │    │  semantic_seg)     │
           │ pred_logits   │    │                    │
           │ presence_score│    │ semantic_mask_logits│
           └──────────────┘    └───────────────────┘
                    │                     │
                    └──────────┬──────────┘
                               ▼
                    ┌──────────────────────────────────────────┐
                    │           Post-processing                │
                    │  argmax → class index map                 │
                    │  threshold filtering → background         │
                    └──────────────────────────────────────────┘
```

### Компоненты

#### 1. Vision Backbone (ViT + FPN Neck)

- **ViT** (Vision Transformer): 32 слоя, embed_dim=1024, 16 heads
  - Входное разрешение: **1008×1008** (patch_size=14 → 72×72 patches)
  - RoPE (Rotary Position Embeddings) с интерполяцией
  - Window attention (window=24) + 4 global attention blocks (слои 7, 15, 23, 31)
  - Tiled absolute position embeddings (pretrain=336 → tile до 1008)
- **Sam3DualViTDetNeck**: FPN-neck с 4 масштабами (×4, ×2, ×1, ×0.5)
  - Выход: `[B, 256, 72, 72]` (после scalp=1, отбрасывается низшее разрешение)
  - Также генерирует SAM2-совместимые features для instance interactivity

#### 2. Text Encoder (VETextEncoder)

- 24-слойный Transformer, width=1024, 16 heads
- BPE токенизатор (vocab: 16M tokens из `bpe_simple_vocab_16e6.txt.gz`)
- Вход: текстовый промпт (например, `"railway"`)
- Выход:
  - `language_features`: `[32, 1, 256]` — проекция в 256-dim space
  - `language_mask`: `[1, 32]` — маска внимания
  - `language_embeds`: `[32, 1, 1024]` — полные эмбеддинги

#### 3. Transformer Encoder (6 layers)

- Fusion: cross-attention между vision features и language features
- Self-attention на vision features
- d_model=256, dim_feedforward=2048, 8 heads
- Activation checkpointing включён

#### 4. Transformer Decoder (6 layers)

- **200 object queries** — learnable queries для обнаружения объектов
- Cross-attention: queries × encoder output
- Text cross-attention: queries × language features
- **DAC** (Decoupled Attention for Classification) — разделение box regression и classification
- **Box refinement** на каждом слое (iterative)
- **Presence token** — для оценки наличия объекта в сцене
- Resolution=1008, stride=14

#### 5. Scoring & Heads

- **DotProductScoring**: MLP (256→2048→256) + dot product для class scores
- **Instance Head**: pred_boxes, pred_masks, pred_logits, presence_score
- **Segmentation Head** (UniversalSegmentationHead):
  - PixelDecoder: 3-stage upsampling (nearest interpolation), hidden_dim=256
  - Выход: `semantic_seg` — пиксельная семантическая маска

---

## Pipeline инференса

### Для одного изображения (`predict_pil`)

```
1. set_image(pil_image)
   └── transform → [1, 3, 1008, 1008]
   └── backbone.forward_image() → backbone_out (vision features)
   
2. Для каждого из 11 промптов:
   a. reset_all_prompts(state)
   b. set_text_prompt(prompt, state)
      └── backbone.forward_text([prompt]) → text_outputs
      └── state["backbone_out"].update(text_outputs)
      └── _forward_grounding(state)
          ├── model.forward_grounding(backbone_out, geometric_prompt)
          ├── pred_boxes, pred_masks, pred_logits → filter by confidence
          └── semantic_seg → interpolate to (H, W)
   c. Агрегация:
      - instance: seg_logits[i] = max(seg_logits[i], mask * score)
      - semantic: seg_logits[i] = max(seg_logits[i], semantic_logits)
      - presence: seg_logits[i] *= presence_score
   
3. Post-processing:
   └── argmax(seg_logits, dim=0) → class map
   └── max_vals < threshold → set to background (class 0)
```

**Итого:** 1 backbone pass + **11 × (text_encoder + grounding_decoder)** forward passes.

### Для батча (`predict_pil_batch`)

```
1. set_image_batch(images)
   └── transform all → [B, 3, 1008, 1008]
   └── backbone.forward_image(batch) → batch_state  ← ОДНА операция на весь батч
   
2. Для каждого изображения i в батче:
   a. _slice_backbone_out(batch_state, i) → per-image state
   b. Для каждого из 11 промптов:
      └── (используются кэшированные text embeddings)
      └── _forward_grounding(state) → instance + semantic masks
   c. argmax + threshold → class map
```

**Батчинг backbone** экономит ~15% времени vs per-image.
**Кэширование text embeddings** экономит ~2% (text encoder очень быстрый).

---

## Профиль производительности (RTX 4090, 24 GB VRAM)

### Разбивка по времени (на 1 изображение, 256×256, 11 промптов)

| Этап | Время | Доля |
|---|---|---|
| Vision backbone (ViT + Neck) | ~25 ms | ~9% |
| Text encoder (× 11 промптов) | ~5 ms | ~2% |
| Grounding decoder (× 11 промптов) | ~240 ms | ~84% |
| Post-processing (argmax, threshold) | ~2 ms | ~1% |
| Overhead (PIL convert, transfer) | ~14 ms | ~5% |
| **Итого** | **~286 ms** | **100%** |

### Throughput при разных batch size

| Batch size | Throughput | VRAM | ms/img |
|---|---|---|---|
| 8 | 3.4 img/s | 6.2 GB | 292 |
| 10 | 3.5 img/s | 6.2 GB | 286 |
| 16 | 3.3 img/s | 8.1 GB | 301 |
| 24 | 3.5 img/s | 8.8 GB | 286 |
| 32 | 3.5 img/s | 8.8 GB | 284 |

**Вывод:** throughput **не масштабируется** с batch size, т.к. bottleneck — **per-image grounding decoder** (11 последовательных forward passes на каждое изображение). Backbone батчится, но это лишь 9% времени.

### Оценка на полный датасет

| Подмножество | Кол-во изображений | Время при 3.5 img/s |
|---|---|---|
| DB (спутник) | ~486K | ~38.5 часов |
| Query (дрон) | ~486K | ~38.5 часов |
| Всё | ~973K | ~77 часов |

---

## Применённые оптимизации

### 1. Кэширование text embeddings

**Файл:** `src/nn/segearth_ov3/segearthov3_segmentor.py`

Text encoder вызывается с одними и теми же 11 промптами для каждого изображения. Кэширование результатов `forward_text()` при первом вызове и повторное использование для всех последующих изображений.

```python
def _cache_text_embeddings(self):
    """Pre-compute and cache text embeddings for all prompts (run once)."""
    if hasattr(self, '_text_cache'):
        return
    self._text_cache = []
    for query_word in self.query_words:
        text_out = self.processor.model.backbone.forward_text(
            [query_word], device=self.device,
        )
        cached = {k: v.clone() if isinstance(v, torch.Tensor) else v
                  for k, v in text_out.items()}
        self._text_cache.append(cached)
```

**Эффект:** ~2% ускорение (54 ms на батч из 8). Text encoder итак быстрый (~0.5 ms/промпт).

### 2. Увеличение batch size (8 → 16)

**Файлы:** `src/main.py`, `src/augmentor/inference.py`

- `_MAX_SEG_BATCH`: 8 → 16
- Segmentation stage `bs`: 8 → 16

**Эффект:** VRAM вырос 6.2 → 8.8 GB (из 24 доступных). Throughput стабилен, но меньше overhead на создание батчей и DataLoader.

### 3. Autocast bfloat16

Уже включён в оригинальном коде:
```python
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
```

### 4. TF32 для Ampere+ GPU

Включён автоматически в `model_builder.py`:
```python
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
```

---

## Фундаментальные ограничения (нельзя ускорить без переписывания модели)

### 1. Последовательный grounding per prompt

Grounding decoder принимает **один текстовый промпт за раз** и прогоняет весь decoder (6 layers, 200 queries). При 11 промптах это 11 × decoder forward = **84% общего времени**.

**Почему нельзя батчить:** архитектура decoder требует `language_features` в cross-attention — разные промпты дают разные `language_features`, что меняет attention pattern. Для батчинга нужно переписать decoder для поддержки multi-prompt inference (нетривиально).

### 2. Последовательный grounding per image

Даже при батчинге backbone, grounding всё равно выполняется **per-image** (slice backbone features → per-image state → 11 decoder passes). Причина: decoder выдаёт per-image instance predictions (boxes, masks, scores), которые нельзя батчить из-за переменного числа detected instances.

### 3. Высокое разрешение ViT

ViT работает на 1008×1008 (72×72 patches), что требует значительного compute. Снижение `default_resolution` ускорит backbone, но ухудшит качество сегментации мелких объектов.

---

## Альтернативные пути ускорения (не реализованы)

| Подход | Ожидаемый эффект | Сложность | Риск |
|---|---|---|---|
| `torch.compile(model)` | 10-30% на decoder | Средняя | Dynamic shapes могут сломать |
| Снижение `num_queries` (200 → 100) | ~20% на decoder | Нужно переучить | Потеря мелких объектов |
| Снижение `default_resolution` (1008 → 504) | ~4x backbone | Тривиально (config) | Ухудшение качества |
| Multi-GPU inference | ~2x при 2 GPU | Средняя | Нужен второй GPU |
| ONNX/TensorRT export | 2-5x overall | Высокая | SAM3 dynamic shapes |
| Замена на SegFormer | ~3x быстрее | Тривиально (fallback есть) | Нет open-vocab, фиксированные 150 классов |

---

## Конфигурация в проекте

### Файлы

| Файл | Назначение |
|---|---|
| `in/config_files/segmentation.gin` | Промпты, threshold, разрешение |
| `src/nn/segearth_ov3/segearthov3_segmentor.py` | Обёртка SegEarth-OV3 (predict_pil, predict_pil_batch) |
| `src/nn/segearth_ov3/sam3/model_builder.py` | Сборка модели (build_sam3_image_model) |
| `src/nn/segearth_ov3/sam3/model/sam3_image_processor.py` | Inference processor (set_image, set_text_prompt) |
| `src/nn/segearth_ov3/sam3/model/sam3_image.py` | Основная модель Sam3Image |
| `src/nn/segearth_ov3/sam3/model/vl_combiner.py` | Vision-Language backbone |
| `src/augmentor/models.py` | Загрузка модели (load_segmentation_model) |
| `src/augmentor/inference.py` | Батчевый инференс (infer_segmentation_batch) |
| `in/weights/sam3.1/sam3.1_multiplex.pt` | Чекпоинт модели (~2.5 GB) |

### Текущая конфигурация

```gin
SegConfig.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
]
SegConfig.threshold = 0.15
SegConfig.default_resolution = 1008
```

### Параметры модели

| Параметр | Значение | Описание |
|---|---|---|
| `prob_thd` | 0.15 | Минимальная уверенность для не-background класса |
| `confidence_threshold` | 0.5 | Порог для instance detection (boxes/masks) |
| `use_sem_seg` | True | Использовать semantic segmentation head |
| `use_presence_score` | True | Масштабировать logits на presence score |
| `use_transformer_decoder` | True | Использовать instance-level predictions |
| `bg_idx` | 0 | Индекс класса background |