commit 686db62c2506c12d1617f82eced1ab9c193e60e5 Author: pikaliov Date: Thu Apr 16 11:22:01 2026 +0300 Initial commit: World-UAV annotation pipeline 4-modality annotation pipeline (depth, edges, segmentation, chmv2) for 973K drone/satellite images. SegEarth-OV3 open-vocabulary segmentation with 11 classes optimized for cross-view geo-localization. Co-Authored-By: Claude Opus 4.6 (1M context) diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0efd9f4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,48 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +*.egg +dist/ +build/ +.eggs/ + +# Virtual environments +.venv/ +venv/ +env/ + +# IDE +.idea/ +.vscode/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Pytest / coverage +.pytest_cache/ +.coverage +htmlcov/ + +# Jupyter +.ipynb_checkpoints/ + +# Claude Code +.claude/ + +# Model weights (>50MB each, download separately) +in/weights/ + +# Test outputs +test_seg_output/ + +# Archives +*.zip + +# Byte-compiled / cached +*.pyc diff --git a/README.md b/README.md new file mode 100644 index 0000000..6debead --- /dev/null +++ b/README.md @@ -0,0 +1,294 @@ +# World-UAV Annotation Pipeline + +Автоматическая генерация 4 модальностей из RGB-изображений датасета World-UAV (973K images): + +| Модальность | Модель | Выход | Скорость | +|:---|:---|:---|:---| +| **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) | RGB palette [256x256] | ~3.5 img/s | +| **CHMv2** | DINOv3-ViTL16 (337M, FP32) | grayscale [256x256] | 31.7 img/s | + +## Quick Start + +```bash +# 1. Запуск (из корня проекта) +python -m src.main + +# 2. Тесты +python -m pytest src/tests/ -v +``` + +Все параметры настраиваются через `in/config_files/*.gin`. Аргументов командной строки нет. + +## Структура проекта + +``` +. +├── 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 функции (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/ # 125 тестов (pytest) +└── 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.resume = True # Пропускать уже обработанные +PipelineConfig.subset = None # None=все, 'Rot', 'Country', 'Terrain' +PipelineConfig.source = 'db' # 'db' = спутник, 'query' = БПЛА, None = оба +``` + +### segmentation.gin (11 классов open-vocabulary) + +```python +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 +] +SegConfig.threshold = 0.15 +SegConfig.default_resolution = 1008 +``` + +Подробный анализ выбора классов: [`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 -> все изображения -> выгрузка +``` + +**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** проверяет существование `{stem}_{suffix}.png` (или `.npy`) для каждого изображения. Пайплайн можно прервать Ctrl+C и перезапустить -- готовые пропускаются. + +## Формат выхода + +Структура директорий **зеркалит** исходный датасет. Исходные изображения не копируются: + +``` +World-UAV-aug/ +├── Rot/SouthernSuburbs/DB/img/ +│ ├── crop_12_4_depth.png # grayscale, 1 канал +│ ├── crop_12_4_edge.png # grayscale, 1 канал +│ ├── crop_12_4_segm.png # RGB palette (11 классов) +│ └── crop_12_4_chm.png # grayscale, 1 канал +├── Country/... +└── Terrain/... +``` + +### Суффиксы + +| Стадия | Суффикс | 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] | + +### Палитра сегментации (11 классов) + +| ID | Класс | Цвет | RGB | +|:--:|:---|:---|:---| +| 0 | background | Black | (0, 0, 0) | +| 1 | building | Red | (220, 40, 40) | +| 2 | road | Gray | (160, 160, 160) | +| 3 | vegetation | Green | (30, 180, 30) | +| 4 | water | Blue | (30, 120, 220) | +| 5 | sand and gravel ground | Tan | (180, 140, 80) | +| 6 | rocky terrain | Brown | (120, 100, 80) | +| 7 | farmland | Yellow | (200, 200, 50) | +| 8 | railway | Purple | (100, 60, 120) | +| 9 | parking lot | Orange | (255, 165, 0) | +| 10 | sidewalk | Light gray | (200, 200, 200) | + +## Использование для обучения + +Depth, edge, chm -- **grayscale 1-канальные**. Загружать как float [0, 1]: + +```python +from PIL import Image +import numpy as np + +stem = "crop_12_4" +aug_dir = Path("World-UAV-aug/Rot/SouthernSuburbs/DB/img") + +# Depth / Edge / CHM -- grayscale float [0, 1] +depth = np.array(Image.open(aug_dir / f"{stem}_depth.png")) / 255.0 # [H, W] +edge = np.array(Image.open(aug_dir / f"{stem}_edge.png")) / 255.0 +chm = np.array(Image.open(aug_dir / f"{stem}_chm.png")) / 255.0 + +# Segmentation -- class index [0, 10] +# Если save_npy=True: seg = np.load(aug_dir / f"{stem}_segm.npy") # [1, H, W] uint8 +# Если только PNG, используй LUT для обратного маппинга из RGB + +# Конкатенация: RGB(3) + depth(1) + edge(1) + chm(1) = 6 каналов +aux = np.stack([depth, edge, chm], axis=0) # [3, H, W] float32 +``` + +> Для сегментации рекомендуется включить `save_npy = True` -- обратный маппинг из RGB палитры в class ID ненадежен. + +## Скачивание весов + +Веса скачиваются один раз в `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 (11 промптов x per-image) = ~84% времени инференса. Text embeddings кэшируются. Batch size backbone = 16 +- **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, 11 prompts) | ~77 ч | **~70%** | +| CHMv2 | ~8.5 ч | ~8% | +| **Итого** | **~101 ч (~4 дня)** | | + +> При обработке только DB (спутник, `source='db'`): ~486K изображений, ~50 ч. +> При обработке только query (БПЛА, `source='query'`): ~486K изображений, ~50 ч. + +## Тесты + +```bash +# Все тесты (125 штук, ~0.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) | Анализ 392 локаций, выбор 11 классов, результаты тестирования | +| [`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 +- gin-config, tqdm, Pillow, coloredlogs, psutil, matplotlib +- omegaconf, einops (зависимости Depth-Anything-3) +- iopath (зависимость SAM3) + +> SegEarth-OV-3 и Depth-Anything-3 **вендорированы** в `src/nn/` -- отдельная установка не требуется. diff --git a/docs/analysis_optimization.md b/docs/analysis_optimization.md new file mode 100644 index 0000000..af08066 --- /dev/null +++ b/docs/analysis_optimization.md @@ -0,0 +1,163 @@ +# Анализ и оптимизация пайплайна depth_edges_annotate_worlduav + +## Обзор + +Пайплайн для аннотации датасета World-UAV: четыре стадии (depth, edges, segmentation, chmv2), конфигурация через Gin, fallback-модели, атомарное сохранение, resume-логика. Нейросетевые модели (SegEarth-OV-3, Depth-Anything-3) вендорированы в `src/nn/`. + +--- + +## 1. Инференс GPU -- самые крупные выигрыши + +### 1.1 `torch.compile()` для моделей + +В `models.py` после загрузки модели можно обернуть в `torch.compile(model, mode="reduce-overhead")`. Для DepthAnything и SegFormer это может дать **20-40% ускорения** на repeated batches за счет fusion ядер и устранения Python overhead. Особенно эффективно при большом количестве батчей. + +### 1.2 AMP через `torch.autocast` вместо ручного `.half()` + +Сейчас в `inference.py:79` и `inference.py:203` делается ручной каст `x = x.half()` при условии `model.parameters().dtype == torch.float16`. Это неоптимально: + +- `torch.autocast("cuda", dtype=torch.float16)` автоматически выберет FP16 для matmul/conv, но оставит FP32 для нормализации и softmax -- меньше потерь точности и обычно быстрее +- Убирает необходимость в ручных проверках dtype + +### 1.3 DA3: batch-инференс вместо поштучного цикла + +`inference.py:59-74` -- DA3 обрабатывается через `model.inference()` API, который внутри делает preprocess + forward + postprocess. Для максимальной производительности можно вызывать низкоуровневый `model.forward()` напрямую (весь батч тензором на GPU), минуя конвертации tensor -> numpy -> PIL -> обратно. + +### 1.4 SegEarth-OV3: batched backbone (реализовано в v3.2) + +`predict_pil_batch()` батчит backbone SAM 3.1 (~80% времени) на до 8 изображений. Grounding decoder остается per-image. Дальнейшие оптимизации: + +- `torch.cuda.Stream` для overlap compute/data-transfer +- Увеличение max_batch если VRAM позволяет (текущий лимит 8 hardcoded) + +--- + +## 2. Data Pipeline -- ускорение загрузки + +### 2.1 `persistent_workers=True` (реализовано) + +DataLoader создается с `persistent_workers=True` при `num_workers > 0`. Воркеры не пересоздаются между батчами. + +### 2.2 `prefetch_factor` (реализовано) + +`prefetch_factor=4` для лучшего overlap загрузки с инференсом. + +### 2.3 Decode на GPU через `torchvision.io` или `nvidia-dali` + +`dataset.py` -- `Image.open().convert("RGB")` + `transforms.Resize()` работает на CPU через PIL. Для больших датасетов это bottleneck. Варианты: + +- `torchvision.io.decode_image` + `torchvision.transforms.v2` -- resize на GPU +- NVIDIA DALI pipeline -- полный decode+resize+normalize на GPU, особенно выгоден при images > 10k + +### 2.4 Предвычисление списка файлов + +`dataset.py` -- `rglob("*")` обходит файловую систему при каждом запуске. Для больших датасетов (тысячи папок) это минуты. Можно кешировать список в `.file_cache.json` и обновлять только при изменении mtime корня. + +--- + +## 3. I/O -- запись результатов + +### 3.1 Асинхронная запись (реализовано) + +`io_utils.py` -- `ThreadPoolExecutor` с 4 workers для неблокирующей записи файлов. GPU inference продолжается пока предыдущий батч пишется на диск. + +### 3.2 Атомарная запись (реализовано) + +Temp file + `os.replace()` для crash-safety. Resume-логика корректно обрабатывает прерванные записи. + +### 3.3 Визуализации -- отложить или отключить + +`save_vis=False` отключает PNG-визуализации. Для ускорения можно генерировать визуализации отдельным скриптом после всего инференса. + +### 3.4 Pre-create output dirs (реализовано) + +Все output_dir создаются одним проходом до начала обработки (`run_pipeline()` в `main.py:333-339`), а не per-image. + +--- + +## 4. Edges стадия + +### 4.1 Batched Sobel (реализовано) + +Edges обрабатываются батчами по 32 изображения. `compute_edges_from_depth` поддерживает `[B, 1, H, W]`. + +### 4.2 Sobel-ядра как module-level константы (реализовано) + +`_SOBEL_X` и `_SOBEL_Y` -- module-level константы в `inference.py`. + +### 4.3 Edges как побочный продукт depth + +Если не нужен resume между стадиями -- считать edges прямо в `run_depth_stage` из только что вычисленного depth, не сохраняя и не загружая `.npy`. Текущая архитектура предпочитает гранулярный resume. + +--- + +## 5. Resume и discovery + +### 5.1 Completion manifest вместо per-file проверок + +`filter_completed()` -- для каждого изображения вызывается `Path.exists()` (syscall). При 100k изображений x 4 стадии = 400k stat-вызовов. Альтернатива: + +- `completed.json` / SQLite per stage +- Проверка по set-lookup вместо filesystem + +### 5.2 `filter_completed` -- per-stage + +Каждая стадия фильтрует отдельно. Можно за один проход собрать статусы всех стадий. + +--- + +## 6. Память и типы данных + +### 6.1 FP32 для CHMv2 (вынужденная мера) + +CHMv2 (DINOv3 DPT head) выдает NaN в FP16. Всегда FP32 -- это увеличивает VRAM на ~0.65 GB vs FP16, но гарантирует корректность. + +### 6.2 uint8 для сегментации (реализовано) + +`infer_segmentation_batch` возвращает `uint8` для class IDs. При 5 классах это оптимально. + +### 6.3 float16 для .npy (реализовано) + +Depth, edges, CHM сохраняются в `.npy` как `float16` (если `save_npy=True`). + +--- + +## 7. Вендорированные пакеты (src/nn/) + +### 7.1 Текущая архитектура + +SegEarth-OV-3 и Depth-Anything-3 **встроены в проект** как вендорированные пакеты в `src/nn/`. При `import src.nn` автоматически регистрируются пути в `sys.path`: +- `src/nn/` -- для `depth_anything_3.*` +- `src/nn/segearth_ov3/` -- для `sam3.*` и `segearthov3_segmentor` + +**Преимущества:** +- Нет зависимости от внешних репозиториев +- Воспроизводимость -- фиксированная версия кода +- Нет конфликтов с системными пакетами + +**Ограничения:** +- Обновление до новой версии модели требует ручного копирования +- Дублирование кода (~6.5 MB) -- допустимо для проекта с ~200 GB данных + +### 7.2 BPE vocab + +Файл `bpe_simple_vocab_16e6.txt.gz` встроен в `src/nn/segearth_ov3/sam3/assets/`. Дефолтный путь в `segearthov3_segmentor.py` вычисляется через `os.path.dirname(__file__)`. Копия в `in/weights/` используется приоритетно (если существует) для совместимости с существующими инсталляциями. + +--- + +## Приоритет по соотношению усилие/выигрыш + +| # | Оптимизация | Усилие | Выигрыш | Статус | +|---|-------------|--------|---------|--------| +| 1 | Async I/O через ThreadPool | Низкое | **Высокий** (overlap GPU/disk) | **Реализовано** | +| 2 | Batched backbone SegEarth-OV3 | Среднее | **Высокий** (~2.5-3x segm) | **Реализовано** | +| 3 | Pre-create output dirs | Минимальное | Средний | **Реализовано** | +| 4 | Batched Sobel edges | Низкое | Средний | **Реализовано** | +| 5 | dtype: uint8 для seg, float16 для npy | Низкое | Средний | **Реализовано** | +| 6 | `persistent_workers` + `prefetch_factor` | Минимальное | Низкий-Средний | **Реализовано** | +| 7 | Вендоринг моделей (src/nn/) | Среднее | Средний (reliability) | **Реализовано** | +| 8 | `torch.compile()` | Низкое | Средний (20-40%) | Не реализовано | +| 9 | `torch.autocast` вместо `.half()` | Низкое | Низкий-Средний | Не реализовано | +| 10 | DA3 прямой forward вместо `inference()` | Среднее | Средний | Не реализовано | +| 11 | Completion manifest вместо per-file exists | Среднее | Средний (при >100k img) | Не реализовано | +| 12 | GPU decode (DALI / torchvision.io) | Высокое | Средний | Не реализовано | diff --git a/docs/segearth_ov3_architecture.md b/docs/segearth_ov3_architecture.md new file mode 100644 index 0000000..9647f39 --- /dev/null +++ b/docs/segearth_ov3_architecture.md @@ -0,0 +1,326 @@ +# 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 | diff --git a/docs/segmentation_class_analysis.md b/docs/segmentation_class_analysis.md new file mode 100644 index 0000000..5032d16 --- /dev/null +++ b/docs/segmentation_class_analysis.md @@ -0,0 +1,194 @@ +# Анализ классов сегментации для UAV-GeoLoc + +## Обзор датасета + +- **392 локации**, 76 кадров каждая (height100_rot90) = ~29 800 изображений +- **11 стран** (Country/): Australia, Brazil, English, French, German, Italy, Japan, Korea, Poland, Spain, USA +- **27 типов terrain** (Terrain/): Basin, Calcification, Danxia, Delta, Desert, Fall, Farm, Finca, Flowers, Glacier, Gorge, Hill, Hylare, Island, Karst, Lakes-ignore, Mountain, Oasis, Pasture, Plain, Plateau, Prairie, Snow, StoneForest, Terrace, Volcano, Wetland + +## Проблема с исходной конфигурацией (5 классов) + +Исходные промпты: +``` +SegConfig.prompts = ['background', 'building', 'road', 'vegetation', 'water'] +SegConfig.threshold = 0.3 +``` + +При 5 классах **60-90% площади terrain-кадров и 30-50% городских кадров уходит в background**, теряя информативные объекты (ж/д пути, грунт, скалы и т.д.). + +--- + +## Инвентаризация объектов + +### URBAN (Country/) — 11 стран, ~300 локаций + +| Объект | Где встречается | Частота | +|---|---|---| +| Здания (крыши: черепица, плоские, металл, стекло) | Все города | Доминирует | +| Дорога / асфальт | Все города | Очень высокая | +| Растительность (деревья, кустарники) | Все города | Высокая | +| Тротуар / пешеходная зона / набережная | Все города, особенно Европа | Высокая | +| Парковка | Мюнхен, SF, Варшава, Аделаида | Средняя | +| Железная дорога (рельсы, балласт, стрелки) | Мюнхен, Сидней, Варшава | Средняя | +| Вода (река, канал, море) | Париж/Сена, Венеция, Киото, Бусан, Копакабана | Средняя | +| Мост / эстакада | Сидней, Париж | Низкая-средняя | +| Внутренний двор | Все европейские города | Средняя | +| Открытый грунт / гравий | Мюнхен, Варшава | Средняя | +| Газон / парк | Лондон/Westminster, Rio | Средняя | +| Спортивная площадка | Мюнхен | Низкая | +| Пляж / песок | Копакабана, Бусан/Gwangalli | Низкая | +| Дорожная разметка (зебры, полосы) | Все города | Высокая, но мелкая | +| Солнечные панели | SF, Мюнхен | Низкая | + +### TERRAIN — 27 типов + +| Terrain | Что на кадрах | Ключевые поверхности | +|---|---|---| +| Desert (Gobi) | Песок, эрозионные борозды, скалы | bare ground, sand, rock | +| Farm (Bohemian) | Густой лес, кроны деревьев | vegetation | +| Glacier (Athabasca) | Лёд, снег, трещины | snow, ice, water | +| Island (Aldabra) | Бирюзовая вода, мелководье | water | +| Wetland (Danube) | Болотистая почва, кустарник | vegetation, bare ground, water | +| Mountain (Andes) | Песок/пыль, следы троп | bare ground, rock | +| Delta (Congo) | Лавовые/скальные поверхности | rock, bare ground | +| Volcano (Kilauea) | Тёмная лава, скальные породы | rock, bare ground | +| Snow (Aconcagua) | Белый снег, точки камней | snow | +| Danxia (GrandCanyon) | Красная порода + растительность | rock, vegetation | +| Oasis | Саванна, кустарники на песке | bare ground, vegetation | +| Finca (Althorp) | Хвойный лес, плантации | vegetation | +| Terrace (Banaway) | Сельхоз террасы, тропы | farmland, bare ground | +| Plain (Alberta) | Луг, речная пойма, грунтовая дорога | grassland, bare ground, road | +| Plateau | Кустарник на каменистой почве | vegetation, rock | +| Prairie (Etosha) | Саванна, точки кустов на песке | bare ground, vegetation | +| Gorge (Antelope) | Каньон, красные скалы | rock, bare ground | +| Karst (Mammoth) | Скальные плиты, трещины | rock | +| Calcification (ElTatio) | Гейзерные поля, песчаная порода | bare ground, rock | +| Flowers (BlueHotSpring) | Горячие источники, минеральные отложения | rock, water | +| Hylare (Amazon) | Густой тропический лес | vegetation | +| StoneForest (GrandCanyon) | Каньонные скалы + кустарник | rock, vegetation | +| Hill (Sedona) | Красные скалы + кустарник | rock, vegetation | +| Basin | — | — | +| Fall | — | — | +| Pasture | — | — | +| Lakes-ignore | — | — | + +--- + +## Финальная конфигурация: 11 классов + +```python +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, embankments +] +SegConfig.threshold = 0.15 +SegConfig.default_resolution = 1008 +``` + +### Эволюция конфигурации + +1. **v1** (исходная): 5 классов, threshold=0.3 — слишком много background +2. **v2** (первая итерация): 12 классов (`bare ground`, `rock`, `snow`, `farmland`, `railway`, `parking lot`, `sidewalk`), threshold=0.3 — terrain по-прежнему плохо +3. **v3** (вторая итерация): 12 классов с переименованными промптами (`sand and gravel ground`, `rocky terrain`, `snow and ice`), threshold=0.15 — значительное улучшение terrain +4. **v4** (финальная): 11 классов — убран `snow and ice` (SegEarth-OV3 не различает снег/лёд от воды при виде сверху), threshold=0.15 + +### Обоснование каждого нового класса + +| Класс | Промпт | Обоснование | Покрытие | +|---|---|---|---| +| bare ground | `sand and gravel ground` | Без него пустыня, гейзеры, плато, прерия, пляж уходят в background | Критичен для terrain | +| rock | `rocky terrain` | Вулканы, каньоны, карст — уникальная текстура | Критичен для terrain | +| farmland | `farmland` | Террасы, поля — геометрически уникальны для matching | Важен для Terrace/Plain/Farm | +| railway | `railway` | Ж/д пути — линейный ориентир с уникальной топологией | Важен для городов | +| parking lot | `parking lot` | Чёткая прямоугольная геометрия, видна и с дрона и со спутника | Средний для городов | +| sidewalk | `sidewalk` | Структурные границы в городах, набережные | Средний для городов | + +### Отброшенные классы и причины + +| Класс | Причина отказа | +|---|---| +| `snow and ice` | SegEarth-OV3 классифицирует лёд/снег как `water` при виде сверху — промпт не работает. Проверено на Glacier (Athabasca) | +| `beach` / `sand` | Покрыт `sand and gravel ground` | +| `bridge` | Слишком редко, перекрывается с `road` | +| `grassland` / `lawn` | Перекрывается с `vegetation` | +| `river` / `sea` / `lake` | Всё покрыто единым `water` | +| `sports field` | Слишком редко на полном датасете | +| `courtyard` | Семантически = `bare ground` + `building` | +| `fence` | Слишком тонкий объект для SegEarth на 1008px | +| `road marking` | Слишком мелкий, лучше оставить как часть `road` | +| `solar panel` | Редко, мелко, часть крыши | + +--- + +## Результаты тестирования на 10 изображениях + +Тестовые изображения: 6 городских (Munich, Paris, Venice, New York, Sydney, Busan) + 4 terrain (Desert, Glacier, Volcano, Danxia). + +### Городские сцены + +| Локация | Обнаруженные классы | Качество | +|---|---|---| +| Munich (ж/д пути) | railway, sand/gravel, vegetation, road, water | Отлично — ж/д = фиолетовый, балласт = бежевый | +| Paris (набережная Сены) | road, vegetation, water, sidewalk, parking lot | Отлично — река, тротуар, дорога разделены | +| Venice (старый город) | building, road, water, parking lot | Хорошо — каналы появились как water | +| New York (Manhattan) | building, road, vegetation, parking lot, sidewalk | Отлично — все городские классы | +| Sydney (эстакада) | road, building, vegetation, parking lot, sidewalk | Хорошо | +| Busan (пляж) | water, sand and gravel ground | Отлично — пляж = бежевый, море = синий | + +### Terrain + +| Локация | Обнаруженные классы | Качество | +|---|---|---| +| Desert (Gobi) | road, rocky terrain | Приемлемо — пустыня = road (серый близок к песку) | +| Glacier (Athabasca) | water | Ограничение модели — лёд = water | +| Volcano (Kilauea) | rocky terrain | Отлично — было 100% background, стало 100% rocky terrain | +| Danxia (GrandCanyon) | vegetation | Частично — кусты определены, скалы в background | + +### Известные ограничения SegEarth-OV3 + +- Ледники/снег классифицируются как `water` — визуально неразличимы сверху для модели +- Красные скалы (Danxia) плохо определяются как `rocky terrain` +- Пустынный грунт может путаться с `road` + +--- + +## Палитра цветов + +| 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 | + +## Оценка влияния на производительность + +- Инференс: ~1.8–2.2x медленнее (11 vs 5 text prompts) +- На 973K изображений: дополнительные ~8–12 часов на RTX 4090 +- Background на terrain: сокращение с 60–90% до ~10–20% +- Background на urban: сокращение с 30–50% до ~5–15% + +## Методология анализа + +Просмотрено ~50 кадров с равномерной выборкой: +- По 1–3 кадра из каждой страны (разные города/районы) +- По 1 кадру из каждого типа terrain (первая доступная локация) +- Кадры выбирались из середины траектории (frame 30) для репрезентативности +- Дополнительно просмотрены кадры из начала/конца траектории для Мюнхена (полный обзор 76 кадров) +- Тестирование проведено в 3 итерации с подбором формулировок промптов и порога confidence diff --git a/docs/skills_optimization_io_dl_ml.md b/docs/skills_optimization_io_dl_ml.md new file mode 100644 index 0000000..9196609 --- /dev/null +++ b/docs/skills_optimization_io_dl_ml.md @@ -0,0 +1,477 @@ +# Skills: Оптимизация I/O и DL/ML инференса + +Справочник приёмов оптимизации для пайплайнов обработки данных и инференса нейросетей. +Используется как чеклист при code review и проектировании новых пайплайнов. + +--- + +## Часть 1. Оптимизация I/O + +### 1.1 Асинхронная запись на диск + +**Проблема:** синхронный `np.save()` / `Image.save()` блокирует основной поток — GPU простаивает. + +**Решение:** + +```python +from concurrent.futures import ThreadPoolExecutor + +_io_pool = ThreadPoolExecutor(max_workers=4) + +def save_async(fn, *args, **kwargs): + """Отправить операцию сохранения в фоновый поток.""" + _io_pool.submit(fn, *args, **kwargs) + +# В inference-цикле: +for batch in loader: + result = model(batch) + save_async(save_result, result, path) # не блокирует GPU +``` + +**Когда применять:** всегда, когда запись идёт между батчами инференса. + +**Подводные камни:** +- Контролировать размер очереди — если запись медленнее инференса, очередь растёт и съедает RAM +- При аварийном завершении незаписанные данные теряются — использовать atomic save (temp + rename) +- `_io_pool.shutdown(wait=True)` в конце пайплайна + +--- + +### 1.2 Атомарная запись файлов + +**Проблема:** прерванная запись оставляет битый файл, resume-логика считает его валидным. + +**Решение:** + +```python +import tempfile, os + +def atomic_save_npy(arr, path): + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(suffix=".tmp", dir=path.parent) + os.close(fd) + try: + np.save(tmp, arr) + os.replace(tmp, path) # атомарная операция на одной ФС + except BaseException: + os.remove(tmp) if os.path.exists(tmp) else None + raise +``` + +**Когда применять:** при любом сохранении промежуточных результатов с resume-логикой. + +--- + +### 1.3 Минимизация syscalls + +**Проблема:** `Path.mkdir()`, `Path.exists()` — это syscalls. При 100k файлов × 3 стадии = 300k+ вызовов. + +**Решения:** + +| Приём | Описание | +|-------|----------| +| **Pre-create dirs** | Один проход `mkdir` для всех output_dir до начала обработки | +| **Completion manifest** | `set()` в памяти вместо `Path.exists()` на каждый файл | +| **Batch stat** | `os.scandir()` вместо поштучного `exists()` | + +```python +# Вместо per-file exists(): +completed = set() +for entry in os.scandir(output_root): + if entry.is_dir(): + for f in os.scandir(entry.path): + if f.name == "depth.npy": + completed.add(entry.name) + +# Фильтрация: +pending = [r for r in records if r.stem not in completed] +``` + +--- + +### 1.4 Выбор формата сохранения + +| Формат | Скорость записи | Размер | Загрузка | Когда использовать | +|--------|----------------|--------|----------|--------------------| +| `.npy` (float32) | Быстро | Большой | Быстро | Промежуточные результаты | +| `.npy` (float16) | Быстро | 2× меньше | Быстро | Depth, edges [0,1] | +| `.npy` (uint8) | Быстро | 8× меньше int64 | Быстро | Class IDs (< 256 классов) | +| `.npz` compressed | Медленно | Маленький | Медленно | Архивация, не для пайплайнов | +| `.png` | Медленно (компрессия) | Маленький | Средне | Только визуализация | +| `.safetensors` | Быстро | Маленький | Быстро, mmap | Тензоры для обучения | + +**Правило:** для промежуточных данных — минимальный достаточный dtype без компрессии. Визуализации — отдельный шаг по требованию. + +--- + +### 1.5 Оптимизация PNG-записи + +```python +# PIL — медленный (zlib compression level 6 по умолчанию) +Image.fromarray(vis).save("out.png") + +# OpenCV — быстрее, контроль компрессии +cv2.imwrite("out.png", vis, [cv2.IMWRITE_PNG_COMPRESSION, 1]) # 1 = минимум + +# Самый быстрый: отложить визуализацию +# Генерировать PNG отдельным скриптом после всего инференса +``` + +--- + +### 1.6 Кеширование файловых списков + +**Проблема:** `rglob("*")` на больших датасетах — минуты. + +**Решение:** + +```python +import json, os + +CACHE = root / ".file_cache.json" + +def discover_cached(root): + if CACHE.exists(): + mtime_cache = CACHE.stat().st_mtime + mtime_root = root.stat().st_mtime + if mtime_cache > mtime_root: + return json.loads(CACHE.read_text()) + + files = [str(p.relative_to(root)) for p in root.rglob("*") if p.is_file()] + CACHE.write_text(json.dumps(files)) + return files +``` + +**Когда применять:** датасет > 10k файлов, особенно на HDD или сетевых ФС. + +--- + +## Часть 2. Оптимизация DL/ML инференса + +### 2.1 `torch.inference_mode()` вместо `torch.no_grad()` + +```python +# Хорошо: +@torch.inference_mode() +def predict(model, x): + return model(x) + +# Менее эффективно: +with torch.no_grad(): + return model(x) +``` + +`inference_mode` отключает version counting и autograd tracking полностью — на 5–10% быстрее `no_grad` и меньше потребление памяти. + +--- + +### 2.2 `torch.compile()` + +```python +model = load_model() +model = torch.compile(model, mode="reduce-overhead") +# первый вызов — долгий (компиляция), далее — 20–40% быстрее +``` + +| Режим | Скорость компиляции | Ускорение рантайма | Когда | +|-------|--------------------|--------------------|-------| +| `"default"` | Средняя | Среднее | Общий случай | +| `"reduce-overhead"` | Медленная | Максимальное | Много батчей, стационарные формы | +| `"max-autotune"` | Очень медленная | Максимальное + подбор ядер | Продакшн, фиксированные shapes | + +**Ограничения:** +- Не работает с dynamic shapes без `dynamic=True` +- Некоторые кастомные операции не компилируются — используйте `torch._dynamo.config.suppress_errors = True` при отладке +- Первый вызов медленный — прогрев обязателен + +--- + +### 2.3 Automatic Mixed Precision (AMP) + +```python +# Правильно: autocast выбирает dtype per-operation +with torch.autocast("cuda", dtype=torch.float16): + output = model(input_tensor) + +# Неправильно: ручной .half() на всё +input_tensor = input_tensor.half() +output = model(input_tensor) # нормализация и softmax тоже в fp16 — потеря точности +``` + +**Что autocast делает автоматически:** +- matmul, conv → FP16 (выигрыш скорости) +- layernorm, softmax, loss → FP32 (сохранение точности) +- Accumulations → FP32 + +**Для обучения** дополнительно нужен `GradScaler`: + +```python +scaler = torch.amp.GradScaler() +with torch.autocast("cuda", dtype=torch.float16): + loss = model(x) +scaler.scale(loss).backward() +scaler.step(optimizer) +scaler.update() +``` + +--- + +### 2.4 Batch-инференс vs. per-image цикл + +**Проблема:** + +```python +# Медленно: per-image на CPU с конвертациями +for i in range(B): + img_np = tensor_to_numpy(batch[i]) # GPU→CPU + result = model.infer_image(img_np) # одно изображение + output.append(numpy_to_tensor(result)) # CPU→GPU +``` + +**Решение:** + +```python +# Быстро: батч на GPU без конвертаций +x = preprocess(batch).to(device) # один transfer +output = model(x) # весь батч за один forward +``` + +**Ключевые правила:** +- Избегать циклов `for i in range(B)` внутри inference-функций +- Если API модели принимает только одно изображение — обращаться к внутреннему `model.forward()` напрямую +- Минимизировать CPU↔GPU трансферы (каждый `tensor.cpu()` / `.to(device)` — задержка) +- Numpy конвертации (`tensor → numpy → PIL → numpy → tensor`) — красный флаг + +--- + +### 2.5 CUDA Streams для overlap + +```python +stream_compute = torch.cuda.Stream() +stream_transfer = torch.cuda.Stream() + +for batch in loader: + with torch.cuda.stream(stream_transfer): + next_batch = next_batch.to(device, non_blocking=True) + + with torch.cuda.stream(stream_compute): + output = model(current_batch) + + torch.cuda.synchronize() +``` + +**Когда применять:** per-image модели (SegEarth/SAM), где нельзя батчить — overlap загрузки следующего изображения с инференсом текущего. + +--- + +### 2.6 DataLoader — скрытые параметры производительности + +```python +DataLoader( + dataset, + batch_size=bs, + num_workers=4, # параллельная загрузка + pin_memory=True, # ускоряет CPU→GPU transfer + persistent_workers=True, # не пересоздавать процессы между эпохами + prefetch_factor=4, # буфер загрузки (дефолт 2) + drop_last=False, # True для обучения, False для инференса +) +``` + +**Диагностика bottleneck загрузки:** + +```python +import time +for batch in loader: + t0 = time.perf_counter() + output = model(batch) + gpu_time = time.perf_counter() - t0 + + t0 = time.perf_counter() + # просто итерация — если это медленно, bottleneck в загрузке +load_time = time.perf_counter() - t0 +``` + +Если `load_time > gpu_time` — увеличить `num_workers`, добавить `prefetch_factor`, перейти на DALI. + +--- + +### 2.7 Decode и preprocessing на GPU + +```python +# CPU (медленно для больших датасетов): +img = Image.open(path).convert("RGB") +tensor = transforms.ToTensor()(transforms.Resize((H, W))(img)) + +# GPU через torchvision v2: +from torchvision.io import decode_image +from torchvision.transforms import v2 + +tensor = decode_image(path) # decode на CPU +tensor = tensor.to(device) +tensor = v2.Resize((H, W))(tensor) # resize на GPU + +# NVIDIA DALI (полностью на GPU): +# decode JPEG → resize → normalize → выход как CUDA тензор +# Требует отдельного pipeline definition, но 3–5× быстрее для >10k images +``` + +--- + +### 2.8 Оптимизация размера батча + +```python +def auto_batch_size(total_vram_mb, weights_mb, per_sample_mb, overhead_mb=200): + """Вычислить безопасный batch size по VRAM бюджету.""" + free = total_vram_mb - weights_mb - overhead_mb + if free <= 0: + return 1 + max_b = int(free / per_sample_mb) + # Округление вниз до степени 2 (лучше для GPU) + safe = max(1, 2 ** int(math.log2(max(1, int(max_b * 0.7))))) + return safe +``` + +**Правила:** +- Степени двойки — лучшая утилизация GPU tensor cores +- 70% от теоретического максимума — запас на пиковое потребление +- `torch.cuda.mem_get_info()` — актуальная свободная память в рантайме (лучше статических оценок) +- При OOM — fallback на `batch_size // 2` с retry + +--- + +### 2.9 Оптимизация конвертаций данных + +**Красные флаги (каждый — потеря производительности):** + +| Паттерн | Проблема | Решение | +|---------|----------|---------| +| `tensor.numpy()` в цикле | GPU→CPU sync | Батчить, делать `.cpu()` один раз | +| `np.uint8 → float32 → device` | Двойная конвертация | Конвертировать сразу на device | +| `Image.fromarray()` в цикле | PIL overhead | cv2 или отложить | +| `tensor.item()` в цикле | Sync point | Собирать в тензор, `.item()` один раз | +| Пересоздание тензоров-констант | Аллокация памяти | Module-level или `register_buffer` | + +--- + +### 2.10 Профилирование инференса + +```python +# Быстрая диагностика: где время? +import torch.utils.benchmark as benchmark + +timer = benchmark.Timer( + stmt="model(x)", + globals={"model": model, "x": sample_batch}, + num_threads=1, +) +print(timer.blocked_autorange()) + +# Детальный профиль: +with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + with_stack=True, +) as prof: + model(sample_batch) + +print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) +# Экспорт в Chrome trace: +prof.export_chrome_trace("trace.json") +``` + +**Что искать:** +- Операции с большим `cuda_time` — кандидаты для `torch.compile` / fusion +- Большое `cpu_time` при маленьком `cuda_time` — Python overhead, data transfer +- Много мелких CUDA kernels — нужен fusion (compile / scripting) + +--- + +### 2.11 Кеширование повторяющихся тензоров + +```python +# Плохо: пересоздание при каждом вызове +def process(x): + kernel = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32) + return F.conv2d(x, kernel.view(1, 1, 3, 3)) + +# Хорошо: module-level константа +_SOBEL_X = torch.tensor( + [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32 +).view(1, 1, 3, 3) / 8.0 + +def process(x): + kernel = _SOBEL_X.to(x.device) # перенос на device только при смене + return F.conv2d(x, kernel) +``` + +--- + +### 2.12 Экспорт и рантаймы для продакшн-инференса + +| Рантайм | Ускорение | Сложность | Когда | +|---------|-----------|-----------|-------| +| `torch.compile` | 1.2–1.5× | Низкая | Первый шаг оптимизации | +| TorchScript (`torch.jit.trace`) | 1.1–1.3× | Средняя | Деплой без Python | +| ONNX Runtime | 1.3–2× | Средняя | Кросс-платформенный инференс | +| TensorRT | 2–5× | Высокая | NVIDIA GPU, фиксированные shapes | +| OpenVINO | 1.5–3× | Средняя | Intel CPU/GPU | + +```python +# ONNX экспорт: +torch.onnx.export(model, sample_input, "model.onnx", + input_names=["image"], output_names=["depth"], + dynamic_axes={"image": {0: "batch"}, "depth": {0: "batch"}}) + +# TensorRT через torch_tensorrt: +import torch_tensorrt +trt_model = torch_tensorrt.compile(model, + inputs=[torch_tensorrt.Input(shape=[bs, 3, H, W], dtype=torch.float16)], + enabled_precisions={torch.float16}) +``` + +--- + +## Часть 3. Чеклист для нового пайплайна + +### Перед написанием кода + +- [ ] Определить bottleneck: GPU compute, CPU preprocessing, или disk I/O? +- [ ] Выбрать минимально достаточные dtype для хранения +- [ ] Спланировать overlap: загрузка ↔ инференс ↔ сохранение + +### Инференс + +- [ ] `torch.inference_mode()` на всех predict-функциях +- [ ] `torch.compile()` если > 100 батчей +- [ ] `torch.autocast` вместо ручного `.half()` +- [ ] Batch forward вместо per-image циклов +- [ ] Нет лишних CPU↔GPU трансферов +- [ ] Нет numpy конвертаций внутри inference loop +- [ ] Тензоры-константы вынесены на module level + +### Data Pipeline + +- [ ] `pin_memory=True` +- [ ] `persistent_workers=True` +- [ ] `prefetch_factor` ≥ 2 +- [ ] `num_workers` подобран (обычно 4–8) +- [ ] Decode/resize на GPU при > 10k images + +### I/O + +- [ ] Асинхронная запись через `ThreadPoolExecutor` +- [ ] Atomic save для crash safety +- [ ] Минимум syscalls (pre-create dirs, batch stat) +- [ ] Визуализации отделены от основного пайплайна +- [ ] Completion tracking через manifest, не per-file `exists()` + +### Память + +- [ ] `gc.collect()` + `torch.cuda.empty_cache()` между стадиями +- [ ] Модели выгружаются когда не нужны +- [ ] Промежуточные тензоры не удерживаются в памяти +- [ ] dtype соответствует данным (uint8 для class IDs, float16 для [0,1]) diff --git a/in/config_files/hardware.gin b/in/config_files/hardware.gin new file mode 100644 index 0000000..458bc7a --- /dev/null +++ b/in/config_files/hardware.gin @@ -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 diff --git a/in/config_files/input.gin b/in/config_files/input.gin new file mode 100644 index 0000000..4a03a11 --- /dev/null +++ b/in/config_files/input.gin @@ -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] diff --git a/in/config_files/models.gin b/in/config_files/models.gin new file mode 100644 index 0000000..0dbac2d --- /dev/null +++ b/in/config_files/models.gin @@ -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' diff --git a/in/config_files/pipeline.gin b/in/config_files/pipeline.gin new file mode 100644 index 0000000..06cc2b5 --- /dev/null +++ b/in/config_files/pipeline.gin @@ -0,0 +1,12 @@ +# 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.resume = True +PipelineConfig.subset = None +# Source filter: 'db' = satellite only, 'query' = drone/UAV only, None = both +PipelineConfig.source = 'query' #'db' +PipelineConfig.log_level = 'INFO' diff --git a/in/config_files/segmentation.gin b/in/config_files/segmentation.gin new file mode 100644 index 0000000..ab71063 --- /dev/null +++ b/in/config_files/segmentation.gin @@ -0,0 +1,18 @@ +# Open-vocabulary segmentation parameters +# 12 cross-view invariant classes for geo-localization +# See docs/segmentation_class_analysis.md for full 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, embankments +] +SegConfig.threshold = 0.15 +SegConfig.default_resolution = 1008 diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..8ab660b --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +"""Source package for depth/edges/segmentation augmentation pipeline.""" diff --git a/src/augmentor/__init__.py b/src/augmentor/__init__.py new file mode 100644 index 0000000..9caa785 --- /dev/null +++ b/src/augmentor/__init__.py @@ -0,0 +1 @@ +"""Augmentor: depth/edges/segmentation augmentation pipeline for CVGL datasets.""" diff --git a/src/augmentor/dataset.py b/src/augmentor/dataset.py new file mode 100644 index 0000000..8afb5dd --- /dev/null +++ b/src/augmentor/dataset.py @@ -0,0 +1,199 @@ +"""Dataset discovery, completion filtering, and PyTorch Dataset for augmentation.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, NamedTuple + +import torch +from PIL import Image +from torch.utils.data import Dataset +from torchvision import transforms + +logger = logging.getLogger(__name__) + +EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp"} + +EXCLUDE_NAMES = {"merge.tif"} +EXCLUDE_DIRS = {"Index", "__MACOSX", "charts", "__pycache__"} + +# Incomplete World-UAV Country scenes (16 of 171): no DB crops, no positive.json. +INCOMPLETE_SCENES: set[str] = { + "CastleHill", "Dalry", "Haymarket", "NewTown", "Stockbridge", + "CamdenTown", "CoventGarden", "Fitzrovia", "Mayfair", "SoHo", + "Ancoats", "Castlefield", "Deansgate", "NorthernQuarter", "Piccadilly", + "JewelleryQuarter", +} + + +class ImageRecord(NamedTuple): + """Lightweight descriptor for a single dataset image.""" + + abs_path: Path + rel_path: str + stem: str + output_dir: Path + + +def is_query_record(record: ImageRecord) -> bool: + """Return True if the record belongs to a query (drone) image.""" + return "query" in Path(record.rel_path).parts + + +def split_by_view( + records: list[ImageRecord], +) -> tuple[list[ImageRecord], list[ImageRecord]]: + """Split records into (db_records, query_records).""" + db: list[ImageRecord] = [] + query: list[ImageRecord] = [] + for r in records: + if is_query_record(r): + query.append(r) + else: + db.append(r) + return db, query + + +# --------------------------------------------------------------------------- +# Discovery +# --------------------------------------------------------------------------- + +def discover_images( + root: Path, + subset: str | None = None, + source: str | None = None, +) -> list[ImageRecord]: + """Recursively find images under *root*, preserving relative paths. + + Args: + root: Dataset root directory. + subset: Limit to a World-UAV subset (Country, Terrain, Rot). + source: Filter by source — 'query' (drone) or 'db' (satellite). + + Returns: + Sorted list of ImageRecord. + """ + search_root = root / subset if subset else root + if not search_root.exists(): + logger.warning("Search root does not exist: %s", search_root) + return [] + + records: list[ImageRecord] = [] + n_skipped_incomplete = 0 + + for p in sorted(search_root.rglob("*")): + if not p.is_file(): + continue + if p.name in EXCLUDE_NAMES: + continue + if p.suffix.lower() not in EXTENSIONS: + continue + if any(d in p.parts for d in EXCLUDE_DIRS): + continue + if any(scene in p.parts for scene in INCOMPLETE_SCENES): + n_skipped_incomplete += 1 + continue + + if source is not None: + rel_parts = p.relative_to(root).parts + if source == "query" and "DB" in rel_parts: + continue + if source == "db" and "query" in rel_parts: + continue + + rel = p.relative_to(root) + records.append(ImageRecord( + abs_path=p, rel_path=str(rel), stem=p.stem, output_dir=Path(), + )) + + if n_skipped_incomplete > 0: + logger.info( + "Skipped %d images from %d incomplete scenes.", + n_skipped_incomplete, len(INCOMPLETE_SCENES), + ) + return records + + +def attach_output_dirs( + records: list[ImageRecord], + output_root: Path, +) -> list[ImageRecord]: + """Set output_dir for each record: output_root / /.""" + out: list[ImageRecord] = [] + for r in records: + rel = Path(r.rel_path) + odir = output_root / rel.parent + out.append(r._replace(output_dir=odir)) + return out + + +# Suffix appended to stem for each stage: {stem}_{suffix}.npy +STAGE_SUFFIX: dict[str, str] = { + "depth": "depth", + "edges": "edge", + "segmentation": "segm", + "chmv2": "chm", +} + + +def stage_filename(stem: str, stage: str, ext: str = ".npy") -> str: + """Build output filename: e.g. crop_12_4_depth.npy""" + suffix = STAGE_SUFFIX.get(stage, stage) + return f"{stem}_{suffix}{ext}" + + +def filter_completed( + records: list[ImageRecord], + stage: str, +) -> tuple[list[ImageRecord], int]: + """Return (pending_records, n_skipped) for a given stage.""" + suffix = STAGE_SUFFIX.get(stage) + if suffix is None: + return records, 0 + pending: list[ImageRecord] = [] + skipped = 0 + for r in records: + # Check both .npy and .png — either means the stage is done. + npy = r.output_dir / f"{r.stem}_{suffix}.npy" + png = r.output_dir / f"{r.stem}_{suffix}.png" + if npy.exists() or png.exists(): + skipped += 1 + else: + pending.append(r) + return pending, skipped + + +# --------------------------------------------------------------------------- +# PyTorch Dataset +# --------------------------------------------------------------------------- + +class AugmentDataset(Dataset): + """Loads RGB images at image_size x image_size for the augmentation pipeline. + + Args: + records: List of ImageRecord to load. + image_size: Target spatial resolution (default 256). + """ + + def __init__(self, records: list[ImageRecord], image_size: int = 256) -> None: + self.records = records + self.resize = transforms.Resize( + (image_size, image_size), + interpolation=transforms.InterpolationMode.BILINEAR, + ) + self.to_tensor = transforms.ToTensor() + + def __len__(self) -> int: + return len(self.records) + + def __getitem__(self, idx: int) -> dict[str, Any]: + r = self.records[idx] + img = Image.open(r.abs_path).convert("RGB") + tensor = self.to_tensor(self.resize(img)) + return { + "image_raw": tensor, + "rel_path": r.rel_path, + "stem": r.stem, + "output_dir": str(r.output_dir), + } diff --git a/src/augmentor/inference.py b/src/augmentor/inference.py new file mode 100644 index 0000000..cde5d4b --- /dev/null +++ b/src/augmentor/inference.py @@ -0,0 +1,302 @@ +"""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) diff --git a/src/augmentor/io_utils.py b/src/augmentor/io_utils.py new file mode 100644 index 0000000..ac6695f --- /dev/null +++ b/src/augmentor/io_utils.py @@ -0,0 +1,252 @@ +"""I/O utilities: saving depth / edges / segmentation / 6-ch concat. + +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 + +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 + + +# 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 +], 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 + + +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] + + +def _save_float16_map( + data: torch.Tensor, + output_dir: Path, + stem: str, + suffix: str, + save_npy: bool = True, + save_vis: bool = True, + colormap: str | None = None, +) -> None: + """Save a [1, H, W] float tensor as {stem}_{suffix}.npy (float16) + optional vis. + + Args: + colormap: If set (e.g. "turbo"), apply colormap for RGB visualization. + If None, save grayscale. + """ + arr = data.half().numpy() + if save_npy: + _atomic_save_npy(arr, output_dir / f"{stem}_{suffix}.npy") + 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) + Image.fromarray(vis).save(output_dir / f"{stem}_{suffix}.png") + + +def save_depth(depth: torch.Tensor, output_dir: Path, stem: str, + save_npy: bool = True, save_vis: bool = True) -> None: + _save_float16_map(depth, output_dir, stem, "depth", save_npy, save_vis) + + +def save_depth_async(depth: torch.Tensor, output_dir: Path, stem: str, + save_npy: bool = True, save_vis: bool = True) -> None: + get_io_pool().submit(save_depth, depth.clone().cpu(), output_dir, stem, save_npy, save_vis) + + +def save_chmv2(depth: torch.Tensor, output_dir: Path, stem: str, + save_npy: bool = True, save_vis: bool = True) -> None: + _save_float16_map(depth, output_dir, stem, "chm", save_npy, save_vis) + + +def save_chmv2_async(depth: torch.Tensor, output_dir: Path, stem: str, + save_npy: bool = True, save_vis: bool = True) -> None: + get_io_pool().submit(save_chmv2, depth.clone().cpu(), output_dir, stem, save_npy, save_vis) + + +def save_edges(edges: torch.Tensor, output_dir: Path, stem: str, + save_npy: bool = True, save_vis: bool = True) -> None: + _save_float16_map(edges, output_dir, stem, "edge", save_npy, save_vis) + + +def save_edges_async(edges: torch.Tensor, output_dir: Path, stem: str, + save_npy: bool = True, save_vis: bool = True) -> None: + get_io_pool().submit(save_edges, edges.clone().cpu(), output_dir, stem, save_npy, save_vis) + + +def save_segmentation( + seg_ids: torch.Tensor, + output_dir: Path, + stem: str, + save_npy: bool = True, + save_vis: bool = True, + num_classes: int = 150, +) -> None: + """Save segmentation map [1, H, W] uint8 as {stem}_segm.npy.""" + arr = seg_ids.byte().numpy() + if save_npy: + _atomic_save_npy(arr, output_dir / f"{stem}_segm.npy") + if save_vis: + palette = make_palette(num_classes) + seg_np = arr.squeeze(0).astype(np.int32) + h, w = seg_np.shape + seg_clamped = np.clip(seg_np.flatten(), 0, num_classes - 1) + vis = palette[seg_clamped].reshape(h, w, 3) + Image.fromarray(vis).save(output_dir / f"{stem}_segm.png") + + +def save_segmentation_async( + seg_ids: torch.Tensor, + output_dir: Path, + 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_dir, stem, + save_npy, save_vis, num_classes, + ) + + +def save_concat_6ch( + rgb: torch.Tensor, + output_dir: Path, + stem: str, + num_classes: int = 150, +) -> None: + """Assemble 6-ch tensor from saved .npy files.""" + depth_path = output_dir / f"{stem}_depth.npy" + edges_path = output_dir / f"{stem}_edge.npy" + seg_path = output_dir / f"{stem}_segm.npy" + + if not (depth_path.exists() and edges_path.exists() and seg_path.exists()): + logger.warning("⚠️ Missing modality .npy for %s, skipping concat.", stem) + return + + depth = torch.from_numpy(np.load(depth_path).astype(np.float32)) + edges = torch.from_numpy(np.load(edges_path).astype(np.float32)) + seg_ids = torch.from_numpy(np.load(seg_path).astype(np.float32)) + seg_float = seg_ids / max(float(num_classes), 1.0) + + concat = torch.cat([rgb, depth, edges, seg_float], dim=0) + _atomic_save_npy(concat.numpy(), output_dir / f"{stem}_concat.npy") + + +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) diff --git a/src/augmentor/models.py b/src/augmentor/models.py new file mode 100644 index 0000000..70fd635 --- /dev/null +++ b/src/augmentor/models.py @@ -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.") diff --git a/src/conf/__init__.py b/src/conf/__init__.py new file mode 100644 index 0000000..86f999b --- /dev/null +++ b/src/conf/__init__.py @@ -0,0 +1 @@ +"""Configuration package for the depth/edges/segmentation augmentation pipeline.""" diff --git a/src/conf/config_loader.py b/src/conf/config_loader.py new file mode 100644 index 0000000..436a43b --- /dev/null +++ b/src/conf/config_loader.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +import gin + +from src.conf.pipeline_conf import PipelineConfig +from src.conf.hardware_conf import HardwareConfig +from src.conf.models_conf import ModelsConfig +from src.conf.input_conf import InputConfig +from src.conf.seg_conf import SegConfig + +logger = logging.getLogger(__name__) + + +def load_all_configs(path2cfg: str) -> dict[str, Any]: + """Parse ALL .gin files at once and return all config objects. + + Clears gin global state first, then loads all .gin files in sorted + order, then creates config instances from the fully-populated state. + + Args: + path2cfg: Path to config directory (WITH trailing slash). + + Returns: + Dictionary with config objects keyed by short name. + + Raises: + FileNotFoundError: If path2cfg does not exist or has no .gin files. + """ + cfg_dir = Path(path2cfg) + if not cfg_dir.is_dir(): + raise FileNotFoundError(f"Config directory not found: {cfg_dir}") + + gin_files = sorted(cfg_dir.glob("*.gin")) + if not gin_files: + raise FileNotFoundError(f"No .gin files in {cfg_dir}") + + gin.clear_config() + gin.parse_config_files_and_bindings( + config_files=[str(f) for f in gin_files], + bindings=[], + ) + logger.info("Loaded %d gin files from %s", len(gin_files), cfg_dir) + + configs = { + "pipeline": PipelineConfig(), + "hardware": HardwareConfig(), + "models": ModelsConfig(), + "input": InputConfig(), + "seg": SegConfig(), + } + logger.info("Created config objects: %s", list(configs.keys())) + return configs diff --git a/src/conf/hardware_conf.py b/src/conf/hardware_conf.py new file mode 100644 index 0000000..614cdf6 --- /dev/null +++ b/src/conf/hardware_conf.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import gc +import logging +import math +from typing import Any, Callable + +import gin +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +@gin.configurable +class HardwareConfig: + """GPU hardware profile for the augmentation pipeline.""" + + def __init__( + self, + profile_name: str = "rtx4090", + total_ram_gb: float = 24.0, + reserve_gb: float = 2.0, + use_fp16: bool = True, + batch_size: int | None = None, + num_workers: int = 4, + ) -> None: + self.profile_name = profile_name + self.total_ram_gb = total_ram_gb + self.reserve_gb = reserve_gb + self.use_fp16 = use_fp16 + self.batch_size = batch_size + self.num_workers = num_workers + # Derived. + self.available_gb = self.total_ram_gb - self.reserve_gb + + def auto_batch_size(self, act_per_sample_mb: float, + overhead_mb: float = 200.0) -> int: + """Compute safe batch size from actual free VRAM after model loading. + + If CUDA is available, reads real free VRAM. Otherwise falls back + to config-based estimate. + + Returns manual batch_size if set in config. + """ + if self.batch_size: + return self.batch_size + + if torch.cuda.is_available(): + gpu_id = torch.cuda.current_device() + free_driver, _ = torch.cuda.mem_get_info(gpu_id) + cache_free = (torch.cuda.memory_reserved(gpu_id) + - torch.cuda.memory_allocated(gpu_id)) + free_mb = (free_driver + cache_free) / (1024 * 1024) - overhead_mb + else: + free_mb = self.available_gb * 1024 - overhead_mb + + if free_mb <= 0: + return 1 + max_b = int(free_mb / act_per_sample_mb) + safe = max(1, 2 ** int(math.log2(max(1, int(max_b * 0.7))))) + logger.info("🎮 auto_batch_size: %.0f MB free, %.0f MB/sample → batch=%d", + free_mb + overhead_mb, act_per_sample_mb, safe) + return safe + + + def find_batch_size( + self, + inference_fn: Callable[[torch.Tensor], Any], + input_shape: tuple[int, ...], + device: torch.device, + dtype: torch.dtype = torch.float32, + max_batch: int = 64, + safety_factor: float = 0.70, + ) -> int: + """Find optimal batch size by measuring incremental VRAM per sample. + + Runs forward passes with batch=1 and batch=2 to isolate the true + per-sample cost from one-time overhead (CUDA warmup, kernel JIT, etc.). + Then calculates max batch from free VRAM with a safety margin. + + Args: + inference_fn: Callable that accepts [B, C, H, W] tensor (on CPU) + and runs one forward pass. + input_shape: (C, H, W) — shape of a single sample. + device: CUDA device. + dtype: Tensor dtype for the dummy input. + max_batch: Upper bound to try. + safety_factor: Fraction of free VRAM to actually use (0.0–1.0). + + Returns: + Optimal batch size (power of 2, >= 1). + """ + if self.batch_size: + logger.info(" batch_size forced by config: %d", self.batch_size) + return self.batch_size + + if device.type != "cuda" or not torch.cuda.is_available(): + logger.info(" No CUDA device — batch_size=1") + return 1 + + gpu_id = torch.cuda.current_device() + + # --- Step 1: warmup with batch=1 (absorbs one-time overhead) --- + torch.cuda.empty_cache() + gc.collect() + dummy = torch.rand(1, *input_shape, dtype=dtype) + inference_fn(dummy) + torch.cuda.synchronize() + del dummy + torch.cuda.empty_cache() + gc.collect() + + # --- Step 2: measure peak for batch=1 --- + torch.cuda.reset_peak_memory_stats(gpu_id) + dummy1 = torch.rand(1, *input_shape, dtype=dtype) + inference_fn(dummy1) + torch.cuda.synchronize() + peak_bs1 = torch.cuda.max_memory_allocated(gpu_id) + del dummy1 + torch.cuda.empty_cache() + gc.collect() + + # --- Step 3: measure peak for batch=2 --- + torch.cuda.reset_peak_memory_stats(gpu_id) + dummy2 = torch.rand(2, *input_shape, dtype=dtype) + inference_fn(dummy2) + torch.cuda.synchronize() + peak_bs2 = torch.cuda.max_memory_allocated(gpu_id) + del dummy2 + torch.cuda.empty_cache() + gc.collect() + + # --- Step 4: incremental per-sample cost --- + per_sample_bytes = peak_bs2 - peak_bs1 + + # If peak doesn't grow with batch (model processes images + # sequentially and manages its own memory), VRAM is constant + # regardless of batch size → use max_batch for prefetch benefit. + tolerance = 10 * 1024 * 1024 # 10 MB noise margin + if per_sample_bytes <= tolerance: + logger.info( + "🎮 find_batch_size: peak(bs=1)=%.0f MB, peak(bs=2)=%.0f MB — " + "VRAM constant, using max_batch=%d", + peak_bs1 / (1024 * 1024), peak_bs2 / (1024 * 1024), max_batch, + ) + return max_batch + + # --- Step 5: compute batch size from free VRAM --- + free_driver, _ = torch.cuda.mem_get_info(gpu_id) + cache_free = (torch.cuda.memory_reserved(gpu_id) + - torch.cuda.memory_allocated(gpu_id)) + usable_bytes = (free_driver + cache_free) * safety_factor + + per_sample_mb = per_sample_bytes / (1024 * 1024) + usable_mb = usable_bytes / (1024 * 1024) + + logger.info( + "🎮 find_batch_size: peak(bs=1)=%.0f MB, peak(bs=2)=%.0f MB, " + "incremental=%.1f MB/sample", + peak_bs1 / (1024 * 1024), peak_bs2 / (1024 * 1024), + per_sample_mb, + ) + + if per_sample_bytes <= 0: + logger.warning(" per-sample VRAM ≤ 0 (%.1f MB), defaulting to batch=1", + per_sample_mb) + return 1 + + max_b = int(usable_bytes / per_sample_bytes) + # Round down to power of 2 for stable GPU occupancy. + best = max(1, min(max_batch, 2 ** int(math.log2(max(1, max_b))))) + + logger.info( + "🎮 find_batch_size: %.0f MB free × %.0f%% = %.0f MB usable, " + "%.1f MB/sample → batch=%d", + (free_driver + cache_free) / (1024 * 1024), + safety_factor * 100, + usable_mb, + per_sample_mb, + best, + ) + return best + + +def get_hardware_cfg(path2cfg: str) -> HardwareConfig: + """Load ONLY hardware config (for isolated testing).""" + gin.clear_config() + gin.parse_config_file(f"{path2cfg}hardware.gin") + return HardwareConfig() diff --git a/src/conf/input_conf.py b/src/conf/input_conf.py new file mode 100644 index 0000000..9c95f03 --- /dev/null +++ b/src/conf/input_conf.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import gin + + +@gin.configurable +class InputConfig: + """Image preprocessing parameters. + + Attributes: + image_size: Default output resolution (used for DB/satellite images). + query_image_size: Output resolution for query/drone images. + If None, falls back to ``image_size``. + """ + + def __init__( + self, + image_size: int = 256, + query_image_size: int | None = None, + sobel_kernel_size: int = 3, + edge_normalize: bool = True, + imagenet_mean: list[float] | None = None, + imagenet_std: list[float] | None = None, + ) -> None: + self.image_size = image_size + self.query_image_size = query_image_size or image_size + self.sobel_kernel_size = sobel_kernel_size + self.edge_normalize = edge_normalize + self.imagenet_mean = imagenet_mean or [0.485, 0.456, 0.406] + self.imagenet_std = imagenet_std or [0.229, 0.224, 0.225] + + +def get_input_cfg(path2cfg: str) -> InputConfig: + """Load ONLY input config (for isolated testing).""" + gin.clear_config() + gin.parse_config_file(f"{path2cfg}input.gin") + return InputConfig() diff --git a/src/conf/models_conf.py b/src/conf/models_conf.py new file mode 100644 index 0000000..104a357 --- /dev/null +++ b/src/conf/models_conf.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import gin + + +@gin.configurable +class ModelsConfig: + """Model identifiers and fallback strategies.""" + + def __init__( + self, + depth_model_id: str = "DA3-LARGE-1.1", + depth_fallback_id: str = "depth-anything/Depth-Anything-V2-Large-hf", + chmv2_model_id: str = "facebook/dinov3-vitl16-chmv2-dpt-head", + seg_model_type: str = "segearth-ov3", + seg_fallback_type: str = "segformer-b5", + seg_fallback_id: str = "nvidia/segformer-b5-finetuned-ade-640-640", + weights_dir: str = "", + ) -> None: + self.depth_model_id = depth_model_id + self.depth_fallback_id = depth_fallback_id + self.chmv2_model_id = chmv2_model_id + self.seg_model_type = seg_model_type + self.seg_fallback_type = seg_fallback_type + self.seg_fallback_id = seg_fallback_id + self.weights_dir = weights_dir + + +def get_models_cfg(path2cfg: str) -> ModelsConfig: + """Load ONLY models config (for isolated testing).""" + gin.clear_config() + gin.parse_config_file(f"{path2cfg}models.gin") + return ModelsConfig() diff --git a/src/conf/pipeline_conf.py b/src/conf/pipeline_conf.py new file mode 100644 index 0000000..c52f6a5 --- /dev/null +++ b/src/conf/pipeline_conf.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import gin + + +@gin.configurable +class PipelineConfig: + """Pipeline stage configuration: what to process and where to save.""" + + def __init__( + self, + input_root: str = "/data/UAV-GeoLoc", + output_root: str = "/data/UAV-GeoLoc-aug", + stages: list[str] | None = None, + save_npy: bool = True, + save_vis: bool = True, + save_concat: bool = False, + resume: bool = True, + subset: str | None = None, + source: str | None = None, + log_level: str = "INFO", + ) -> None: + self.input_root = input_root + self.output_root = output_root + self.stages = stages or ["depth", "edges", "segmentation"] + self.save_npy = save_npy + self.save_vis = save_vis + self.save_concat = save_concat + self.resume = resume + self.subset = subset + self.source = source + self.log_level = log_level + + +def get_pipeline_cfg(path2cfg: str) -> PipelineConfig: + """Load ONLY pipeline config (for isolated testing).""" + gin.clear_config() + gin.parse_config_file(f"{path2cfg}pipeline.gin") + return PipelineConfig() diff --git a/src/conf/seg_conf.py b/src/conf/seg_conf.py new file mode 100644 index 0000000..86be82d --- /dev/null +++ b/src/conf/seg_conf.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import gin + + +@gin.configurable +class SegConfig: + """Open-vocabulary segmentation parameters.""" + + def __init__( + self, + prompts: list[str] | None = None, + threshold: float = 0.3, + default_resolution: int = 1008, + ) -> None: + self.prompts = prompts # or ["background", "building", "road", "vegetation", "water",] + self.threshold = threshold + self.default_resolution = default_resolution + # Derived. + self.num_classes = len(self.prompts) + + +def get_seg_cfg(path2cfg: str) -> SegConfig: + """Load ONLY segmentation config (for isolated testing).""" + gin.clear_config() + gin.parse_config_file(f"{path2cfg}segmentation.gin") + return SegConfig() diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..6fd79fb --- /dev/null +++ b/src/main.py @@ -0,0 +1,515 @@ +"""Entry point for the depth/edges/segmentation/chmv2 augmentation pipeline. + +All parameters loaded from gin config files — no argparse. +Sequential stage processing: one model at a time, load → process all → unload. + +Usage: + python -m src.main +""" + +from __future__ import annotations + +import gc +import json +import logging +import time +from datetime import datetime +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader +from tqdm import tqdm + +from src.conf.config_loader import load_all_configs +from src.conf.pipeline_conf import PipelineConfig +from src.conf.hardware_conf import HardwareConfig +from src.conf.models_conf import ModelsConfig +from src.conf.input_conf import InputConfig +from src.conf.seg_conf import SegConfig +from src.utils.utils_file_dir import get_proj_dir + +from src.augmentor.dataset import ( + AugmentDataset, ImageRecord, attach_output_dirs, + discover_images, filter_completed, split_by_view, +) +from src.augmentor.inference import ( + compute_edges_from_depth, infer_depth_batch, infer_chmv2_batch, + infer_segmentation_batch, +) +from src.augmentor.io_utils import ( + save_depth_async, save_chmv2_async, save_edges_async, + save_segmentation_async, setup_logging, shutdown_io_pool, +) +from src.augmentor.models import ( + load_depth_model, load_chmv2_model, load_segmentation_model, unload_model, +) +from src.utils.profiler import ( + log_system_info, log_disk_info, log_vram_snapshot, log_ram_snapshot, +) + +logger = logging.getLogger(__name__) + +_STAGE_EMOJI = { + "depth": "🌊", + "edges": "🔪", + "segmentation": "🗺️", + "chmv2": "🦕", + "concat": "🧩", +} + + +def _silence_model_loggers() -> None: + """Suppress verbose inference logs from all models.""" + import os + import warnings + + os.environ["DA3_LOG_LEVEL"] = "ERROR" + os.environ["TRANSFORMERS_VERBOSITY"] = "error" + os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + + for name in ( + "depth_anything_3", "depth_anything_3.api", + "depth_anything_3.utils.logger", + "transformers", "transformers.modeling_utils", + "transformers.configuration_utils", "transformers.image_processing_utils", + "transformers.image_processing_base", + "sam3", "segearthov3_segmentor", + "huggingface_hub", "httpx", "filelock", + "numexpr", "numexpr.utils", + "torch", "py.warnings", + ): + logging.getLogger(name).setLevel(logging.ERROR) + + warnings.filterwarnings("ignore", category=DeprecationWarning) + warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", message=".*not sharded.*") + + +# --------------------------------------------------------------------------- +# Stage runners +# --------------------------------------------------------------------------- + +def _resolve_image_sizes( + records: list[ImageRecord], + input_conf: InputConfig, +) -> list[tuple[list[ImageRecord], int, str]]: + """Split records into groups by target resolution. + + Returns list of (records, image_size, label) tuples. When + ``query_image_size == image_size`` a single group is returned (no split). + """ + if input_conf.query_image_size == input_conf.image_size: + return [(records, input_conf.image_size, "all")] + db_recs, query_recs = split_by_view(records) + groups: list[tuple[list[ImageRecord], int, str]] = [] + if db_recs: + groups.append((db_recs, input_conf.image_size, f"db {input_conf.image_size}")) + if query_recs: + groups.append((query_recs, input_conf.query_image_size, f"query {input_conf.query_image_size}")) + return groups + + +def run_depth_stage( + records: list[ImageRecord], + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + models_conf: ModelsConfig, + input_conf: InputConfig, + device: torch.device, +) -> None: + """🌊 Load DA3, process all images, unload.""" + model = load_depth_model(models_conf, hw_conf, device) + + for group_records, sz, label in _resolve_image_sizes(records, input_conf): + bs = hw_conf.find_batch_size( + inference_fn=lambda x: infer_depth_batch(model, x, device), + input_shape=(3, sz, sz), + device=device, + ) + + ds = AugmentDataset(group_records, image_size=sz) + logger.info("🌊 [depth/%s] DataLoader: batch_size=%d, %d images, %d batches", + label, bs, len(ds), (len(ds) + bs - 1) // bs) + loader = DataLoader( + ds, batch_size=bs, shuffle=False, + num_workers=hw_conf.num_workers, pin_memory=True, + persistent_workers=hw_conf.num_workers > 0, + prefetch_factor=4 if hw_conf.num_workers > 0 else None, + ) + + total_images = len(ds) + pbar = tqdm(loader, desc=f"🌊 depth/{label} (bs={bs})", unit="batch", + total=len(loader), colour="green", + bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]") + processed = 0 + for batch in pbar: + depths = infer_depth_batch(model, batch["image_raw"], device) + for i in range(depths.shape[0]): + save_depth_async(depths[i], Path(batch["output_dir"][i]), + stem=batch["stem"][i], + save_npy=pipeline_conf.save_npy, + save_vis=pipeline_conf.save_vis) + processed += depths.shape[0] + pbar.set_postfix(images=f"{processed}/{total_images}") + + shutdown_io_pool() + + unload_model(model) + + +def run_edges_stage( + records: list[ImageRecord], + pipeline_conf: PipelineConfig, + batch_size: int = 32, +) -> None: + """🔪 Compute Sobel edges from saved depth (CPU, batched).""" + valid: list[ImageRecord] = [] + for r in records: + depth_png = r.output_dir / f"{r.stem}_depth.png" + depth_npy = r.output_dir / f"{r.stem}_depth.npy" + if depth_png.exists() or depth_npy.exists(): + valid.append(r) + else: + logger.warning("⚠️ No depth for %s, skipping edges.", r.rel_path) + + total_images = len(valid) + pbar = tqdm(range(0, len(valid), batch_size), desc="🔪 edges (sobel)", unit="batch", + colour="cyan", + bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]") + processed = 0 + for start in pbar: + chunk = valid[start : start + batch_size] + depth_tensors = [] + for r in chunk: + npy_path = r.output_dir / f"{r.stem}_depth.npy" + png_path = r.output_dir / f"{r.stem}_depth.png" + if npy_path.exists(): + d = np.load(npy_path).astype(np.float32) + else: + from PIL import Image + d = np.array(Image.open(png_path)).astype(np.float32) / 255.0 + if d.ndim == 2: + d = d[np.newaxis] + depth_tensors.append(torch.from_numpy(d)) + depths = torch.stack(depth_tensors) + if depths.ndim == 3: + depths = depths.unsqueeze(1) + edges_batch = compute_edges_from_depth(depths) + for j, r in enumerate(chunk): + save_edges_async(edges_batch[j], r.output_dir, stem=r.stem, + save_npy=pipeline_conf.save_npy, + save_vis=pipeline_conf.save_vis) + processed += len(chunk) + pbar.set_postfix(images=f"{processed}/{total_images}") + + shutdown_io_pool() + + +def run_chmv2_stage( + records: list[ImageRecord], + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + models_conf: ModelsConfig, + input_conf: InputConfig, + device: torch.device, +) -> None: + """🦕 Load CHMv2 (DINOv3), process all images, unload.""" + model, processor = load_chmv2_model(models_conf, hw_conf, device) + + for group_records, sz, label in _resolve_image_sizes(records, input_conf): + bs = hw_conf.find_batch_size( + inference_fn=lambda x: infer_chmv2_batch(model, processor, x, device), + input_shape=(3, sz, sz), + device=device, + ) + + ds = AugmentDataset(group_records, image_size=sz) + logger.info("🦕 [chmv2/%s] DataLoader: batch_size=%d, %d images, %d batches", + label, bs, len(ds), (len(ds) + bs - 1) // bs) + loader = DataLoader( + ds, batch_size=bs, shuffle=False, + num_workers=hw_conf.num_workers, pin_memory=True, + persistent_workers=hw_conf.num_workers > 0, + prefetch_factor=4 if hw_conf.num_workers > 0 else None, + ) + + total_images = len(ds) + pbar = tqdm(loader, desc=f"🦕 chmv2/{label} (bs={bs})", unit="batch", + total=len(loader), colour="blue", + bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]") + processed = 0 + for batch in pbar: + depths = infer_chmv2_batch(model, processor, batch["image_raw"], device) + for i in range(depths.shape[0]): + save_chmv2_async(depths[i], Path(batch["output_dir"][i]), + stem=batch["stem"][i], + save_npy=pipeline_conf.save_npy, + save_vis=pipeline_conf.save_vis) + processed += depths.shape[0] + pbar.set_postfix(images=f"{processed}/{total_images}") + + shutdown_io_pool() + + unload_model(model) + + +def run_segmentation_stage( + records: list[ImageRecord], + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + models_conf: ModelsConfig, + input_conf: InputConfig, + seg_conf: SegConfig, + device: torch.device, +) -> None: + """🗺️ Load segmentation model, process all images, unload.""" + model, seg_config = load_segmentation_model(models_conf, hw_conf, seg_conf, device) + num_classes = seg_config.get("num_classes", 150) + is_segearth = seg_config.get("type") == "segearth-ov3" + + for group_records, sz, label in _resolve_image_sizes(records, input_conf): + if is_segearth: + bs = 16 + else: + bs = hw_conf.find_batch_size( + inference_fn=lambda x: infer_segmentation_batch(model, seg_config, x, device), + input_shape=(3, sz, sz), + device=device, + ) + + ds = AugmentDataset(group_records, image_size=sz) + total_images = len(ds) + logger.info("🗺️ [segmentation/%s] DataLoader: batch_size=%d, %d images, %d batches", + label, bs, total_images, (total_images + bs - 1) // bs) + loader = DataLoader( + ds, batch_size=bs, shuffle=False, + num_workers=hw_conf.num_workers, pin_memory=True, + persistent_workers=hw_conf.num_workers > 0, + prefetch_factor=4 if hw_conf.num_workers > 0 else None, + ) + + seg_type = "SegEarth-OV3" if is_segearth else "SegFormer" + pbar = tqdm(loader, desc=f"🗺️ seg/{label} {seg_type} (bs={bs})", unit="batch", + colour="yellow", + bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]") + processed = 0 + for batch in pbar: + segs = infer_segmentation_batch( + model, seg_config, batch["image_raw"], device, + ) + for j in range(segs.shape[0]): + save_segmentation_async( + segs[j], Path(batch["output_dir"][j]), stem=batch["stem"][j], + save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis, + num_classes=num_classes, + ) + processed += segs.shape[0] + pbar.set_postfix(images=f"{processed}/{total_images}") + + shutdown_io_pool() + + unload_model(model) + + +# --------------------------------------------------------------------------- +# Pipeline orchestration +# --------------------------------------------------------------------------- + +def run_pipeline( + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + models_conf: ModelsConfig, + input_conf: InputConfig, + seg_conf: SegConfig, +) -> None: + """Execute the full augmentation pipeline: one stage at a time.""" + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if device.type != "cuda": + logger.warning("⚠️ CUDA not available, running on CPU (very slow).") + + _silence_model_loggers() + + # System profiling at startup. + log_system_info() + log_disk_info(Path(pipeline_conf.input_root), Path(pipeline_conf.output_root)) + + # Discover images. + input_root = Path(pipeline_conf.input_root) + output_root = Path(pipeline_conf.output_root) + logger.info( + "🔍 Discovering images in %s (subset=%s, source=%s) ...", + input_root, pipeline_conf.subset or "all", pipeline_conf.source or "all", + ) + all_records = discover_images(input_root, subset=pipeline_conf.subset, + source=pipeline_conf.source) + all_records = attach_output_dirs(all_records, output_root) + logger.info("📸 Found %d images.", len(all_records)) + + if not all_records: + logger.error("❌ No images found. Check input_root in pipeline.gin.") + return + + # Pre-create all output directories in one pass. + logger.info("📁 Pre-creating output directories...") + seen_dirs: set[str] = set() + for r in all_records: + d = str(r.output_dir) + if d not in seen_dirs: + r.output_dir.mkdir(parents=True, exist_ok=True) + seen_dirs.add(d) + + # Process each stage sequentially. + stage_times: dict[str, float] = {} + stage_counts: dict[str, int] = {} + failed_stages: set[str] = set() + + for stage in pipeline_conf.stages: + emoji = _STAGE_EMOJI.get(stage, "⚙️") + + if stage == "edges" and "depth" in failed_stages: + logger.error("❌ [edges] Skipped — depth stage failed.") + failed_stages.add(stage) + continue + + pending, skipped = filter_completed(all_records, stage) + + logger.info("%s [%s] %d pending, %d skipped.", emoji, stage, len(pending), skipped) + stage_counts[stage] = len(pending) + + if not pending: + stage_times[stage] = 0.0 + continue + + logger.info("%s [%s] Starting stage...", emoji, stage) + log_vram_snapshot(f"before {stage}") + log_ram_snapshot(f"before {stage}") + t0 = time.perf_counter() + try: + if stage == "depth": + run_depth_stage(pending, pipeline_conf, hw_conf, models_conf, + input_conf, device) + elif stage == "edges": + run_edges_stage(pending, pipeline_conf) + elif stage == "segmentation": + run_segmentation_stage(pending, pipeline_conf, hw_conf, models_conf, + input_conf, seg_conf, device) + elif stage == "chmv2": + run_chmv2_stage(pending, pipeline_conf, hw_conf, models_conf, + input_conf, device) + except Exception: + logger.exception("💥 Stage '%s' failed.", stage) + failed_stages.add(stage) + elapsed = time.perf_counter() - t0 + stage_times[stage] = elapsed + if stage not in failed_stages: + logger.info("✅ [%s] Completed in %.1f s (%d images).", stage, elapsed, len(pending)) + log_vram_snapshot(f"after {stage}") + log_ram_snapshot(f"after {stage}") + + # Manifest. + manifest = { + "pipeline_version": "3.2.0-dual-resolution", + "image_size_db": input_conf.image_size, + "image_size_query": input_conf.query_image_size, + "profile": hw_conf.profile_name, + "models": { + "depth": models_conf.depth_model_id, + "edges": "Sobel from depth (CPU)", + "segmentation": models_conf.seg_model_type, + "chmv2": models_conf.chmv2_model_id, + }, + "seg_prompts": seg_conf.prompts, + "total_images": len(all_records), + "stages": { + s: {"processed": stage_counts.get(s, 0), + "time_sec": round(stage_times.get(s, 0), 1)} + for s in pipeline_conf.stages + }, + "timestamp": datetime.now().isoformat(), + } + manifest_path = output_root / "manifest.json" + manifest_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.write_text( + json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8", + ) + + # Summary. + total = sum(stage_times.values()) + logger.info("=" * 60) + logger.info("🏁 DONE: %d images, %.1f s total", len(all_records), total) + for s, t in stage_times.items(): + cnt = stage_counts.get(s, len(all_records)) + fps = cnt / t if t > 0 else 0 + emoji = _STAGE_EMOJI.get(s, "⚙️") + logger.info(" %s %-13s %6.1f s (%d images, %.0f FPS)", emoji, s, t, cnt, fps) + logger.info("📂 Output: %s", output_root) + logger.info("=" * 60) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +def main() -> None: + """Load all gin configs and run the augmentation pipeline. + + Supports CLI gin overrides for quick mode switches:: + + # Process only query (drone) images: + python -m src.main --gin "PipelineConfig.source = 'query'" + + # Process only db (satellite) images: + python -m src.main --gin "PipelineConfig.source = 'db'" + """ + import argparse + + parser = argparse.ArgumentParser(description="Augmentation pipeline") + parser.add_argument("--gin", action="append", default=[], + help="Gin parameter overrides (repeatable)") + args = parser.parse_args() + + _silence_model_loggers() + + proj_dir = get_proj_dir() + path2cfg = f"{proj_dir}in/config_files/" + + # Load configs with optional CLI overrides. + if args.gin: + import gin as _gin + cfg_dir = Path(path2cfg) + gin_files = sorted(cfg_dir.glob("*.gin")) + _gin.clear_config() + _gin.parse_config_files_and_bindings( + config_files=[str(f) for f in gin_files], + bindings=args.gin, + ) + configs = { + "pipeline": PipelineConfig(), + "hardware": HardwareConfig(), + "models": ModelsConfig(), + "input": InputConfig(), + "seg": SegConfig(), + } + else: + configs = load_all_configs(path2cfg) + + pipeline_conf: PipelineConfig = configs["pipeline"] + setup_logging(pipeline_conf.log_level, + log_file=Path(pipeline_conf.output_root) / "pipeline.log") + + torch.manual_seed(42) + np.random.seed(42) + + run_pipeline( + configs["pipeline"], + configs["hardware"], + configs["models"], + configs["input"], + configs["seg"], + ) + + +if __name__ == "__main__": + main() diff --git a/src/nn/__init__.py b/src/nn/__init__.py new file mode 100644 index 0000000..3935643 --- /dev/null +++ b/src/nn/__init__.py @@ -0,0 +1,23 @@ +"""Vendored neural network packages: SegEarth-OV-3 and Depth-Anything-3. + +On import, this module adds the necessary directories to sys.path so that +the internal imports inside each vendored package work unchanged: + - src/nn/ -> makes `depth_anything_3.*` importable + - src/nn/segearth_ov3/ -> makes `sam3.*` and `segearthov3_segmentor` importable +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent + +_VENDOR_PATHS = [ + str(_THIS_DIR), # depth_anything_3.* + str(_THIS_DIR / "segearth_ov3"), # sam3.*, segearthov3_segmentor +] + +for _p in _VENDOR_PATHS: + if _p not in sys.path: + sys.path.insert(0, _p) diff --git a/src/nn/depth_anything_3/api.py b/src/nn/depth_anything_3/api.py new file mode 100644 index 0000000..d0f18aa --- /dev/null +++ b/src/nn/depth_anything_3/api.py @@ -0,0 +1,446 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Depth Anything 3 API module. + +This module provides the main API for Depth Anything 3, including model loading, +inference, and export capabilities. It supports both single and nested model architectures. +""" + +from __future__ import annotations + +import time +from typing import Optional, Sequence +import numpy as np +import torch +import torch.nn as nn +from huggingface_hub import PyTorchModelHubMixin +from PIL import Image + +from depth_anything_3.cfg import create_object, load_config +from depth_anything_3.registry import MODEL_REGISTRY +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.export import export +from depth_anything_3.utils.geometry import affine_inverse +from depth_anything_3.utils.io.input_processor import InputProcessor +from depth_anything_3.utils.io.output_processor import OutputProcessor +from depth_anything_3.utils.logger import logger +from depth_anything_3.utils.pose_align import align_poses_umeyama + +torch.backends.cudnn.benchmark = False +# logger.info("CUDNN Benchmark Disabled") + +SAFETENSORS_NAME = "model.safetensors" +CONFIG_NAME = "config.json" + + +class DepthAnything3(nn.Module, PyTorchModelHubMixin): + """ + Depth Anything 3 main API class. + + This class provides a high-level interface for depth estimation using Depth Anything 3. + It supports both single and nested model architectures with metric scaling capabilities. + + Features: + - Hugging Face Hub integration via PyTorchModelHubMixin + - Support for multiple model presets (vitb, vitg, nested variants) + - Automatic mixed precision inference + - Export capabilities for various formats (GLB, PLY, NPZ, etc.) + - Camera pose estimation and metric depth scaling + + Usage: + # Load from Hugging Face Hub + model = DepthAnything3.from_pretrained("huggingface/model-name") + + # Or create with specific preset + model = DepthAnything3(preset="vitg") + + # Run inference + prediction = model.inference(images, export_dir="output", export_format="glb") + """ + + _commit_hash: str | None = None # Set by mixin when loading from Hub + + def __init__(self, model_name: str = "da3-large", **kwargs): + """ + Initialize DepthAnything3 with specified preset. + + Args: + model_name: The name of the model preset to use. + Examples: 'da3-giant', 'da3-large', 'da3metric-large', 'da3nested-giant-large'. + **kwargs: Additional keyword arguments (currently unused). + """ + super().__init__() + self.model_name = model_name + + # Build the underlying network + self.config = load_config(MODEL_REGISTRY[self.model_name]) + self.model = create_object(self.config) + self.model.eval() + + # Initialize processors + self.input_processor = InputProcessor() + self.output_processor = OutputProcessor() + + # Device management (set by user) + self.device = None + + @torch.inference_mode() + def forward( + self, + image: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + export_feat_layers: list[int] | None = None, + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> dict[str, torch.Tensor]: + """ + Forward pass through the model. + + Args: + image: Input batch with shape ``(B, N, 3, H, W)`` on the model device. + extrinsics: Optional camera extrinsics with shape ``(B, N, 4, 4)``. + intrinsics: Optional camera intrinsics with shape ``(B, N, 3, 3)``. + export_feat_layers: Layer indices to return intermediate features for. + infer_gs: Enable Gaussian Splatting branch. + use_ray_pose: Use ray-based pose estimation instead of camera decoder. + ref_view_strategy: Strategy for selecting reference view from multiple views. + + Returns: + Dictionary containing model predictions + """ + # Determine optimal autocast dtype + autocast_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + with torch.no_grad(): + with torch.autocast(device_type=image.device.type, dtype=autocast_dtype): + return self.model( + image, extrinsics, intrinsics, export_feat_layers, infer_gs, use_ray_pose, ref_view_strategy + ) + + def inference( + self, + image: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + align_to_input_ext_scale: bool = True, + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + render_exts: np.ndarray | None = None, + render_ixts: np.ndarray | None = None, + render_hw: tuple[int, int] | None = None, + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_dir: str | None = None, + export_format: str = "mini_npz", + export_feat_layers: Sequence[int] | None = None, + # GLB export parameters + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + # Feat_vis export parameters + feat_vis_fps: int = 15, + # Other export parameters, e.g., gs_ply, gs_video + export_kwargs: Optional[dict] = {}, + ) -> Prediction: + """ + Run inference on input images. + + Args: + image: List of input images (numpy arrays, PIL Images, or file paths) + extrinsics: Camera extrinsics (N, 4, 4) + intrinsics: Camera intrinsics (N, 3, 3) + align_to_input_ext_scale: whether to align the input pose scale to the prediction + infer_gs: Enable the 3D Gaussian branch (needed for `gs_ply`/`gs_video` exports) + use_ray_pose: Use ray-based pose estimation instead of camera decoder (default: False) + ref_view_strategy: Strategy for selecting reference view from multiple views. + Options: "first", "middle", "saddle_balanced", "saddle_sim_range". + Default: "saddle_balanced". For single view input (S ≤ 2), no reordering is performed. + render_exts: Optional render extrinsics for Gaussian video export + render_ixts: Optional render intrinsics for Gaussian video export + render_hw: Optional render resolution for Gaussian video export + process_res: Processing resolution + process_res_method: Resize method for processing + export_dir: Directory to export results + export_format: Export format (mini_npz, npz, glb, ply, gs, gs_video) + export_feat_layers: Layer indices to export intermediate features from + conf_thresh_percentile: [GLB] Lower percentile for adaptive confidence threshold (default: 40.0) # noqa: E501 + num_max_points: [GLB] Maximum number of points in the point cloud (default: 1,000,000) + show_cameras: [GLB] Show camera wireframes in the exported scene (default: True) + feat_vis_fps: [FEAT_VIS] Frame rate for output video (default: 15) + export_kwargs: additional arguments to export functions. + + Returns: + Prediction object containing depth maps and camera parameters + """ + if "gs" in export_format: + assert infer_gs, "must set `infer_gs=True` to perform gs-related export." + + if "colmap" in export_format: + assert isinstance(image[0], str), "`image` must be image paths for COLMAP export." + + # Preprocess images + imgs_cpu, extrinsics, intrinsics = self._preprocess_inputs( + image, extrinsics, intrinsics, process_res, process_res_method + ) + + # Prepare tensors for model + imgs, ex_t, in_t = self._prepare_model_inputs(imgs_cpu, extrinsics, intrinsics) + + # Normalize extrinsics + ex_t_norm = self._normalize_extrinsics(ex_t.clone() if ex_t is not None else None) + + # Run model forward pass + export_feat_layers = list(export_feat_layers) if export_feat_layers is not None else [] + + raw_output = self._run_model_forward( + imgs, ex_t_norm, in_t, export_feat_layers, infer_gs, use_ray_pose, ref_view_strategy + ) + + # Convert raw output to prediction + prediction = self._convert_to_prediction(raw_output) + + # Align prediction to extrinsincs + prediction = self._align_to_input_extrinsics_intrinsics( + extrinsics, intrinsics, prediction, align_to_input_ext_scale + ) + + # Add processed images for visualization + prediction = self._add_processed_images(prediction, imgs_cpu) + + # Export if requested + if export_dir is not None: + + if "gs" in export_format: + if infer_gs and "gs_video" not in export_format: + export_format = f"{export_format}-gs_video" + if "gs_video" in export_format: + if "gs_video" not in export_kwargs: + export_kwargs["gs_video"] = {} + export_kwargs["gs_video"].update( + { + "extrinsics": render_exts, + "intrinsics": render_ixts, + "out_image_hw": render_hw, + } + ) + # Add GLB export parameters + if "glb" in export_format: + if "glb" not in export_kwargs: + export_kwargs["glb"] = {} + export_kwargs["glb"].update( + { + "conf_thresh_percentile": conf_thresh_percentile, + "num_max_points": num_max_points, + "show_cameras": show_cameras, + } + ) + # Add Feat_vis export parameters + if "feat_vis" in export_format: + if "feat_vis" not in export_kwargs: + export_kwargs["feat_vis"] = {} + export_kwargs["feat_vis"].update( + { + "fps": feat_vis_fps, + } + ) + # Add COLMAP export parameters + if "colmap" in export_format: + if "colmap" not in export_kwargs: + export_kwargs["colmap"] = {} + export_kwargs["colmap"].update( + { + "image_paths": image, + "conf_thresh_percentile": conf_thresh_percentile, + "process_res_method": process_res_method, + } + ) + self._export_results(prediction, export_format, export_dir, **export_kwargs) + + return prediction + + def _preprocess_inputs( + self, + image: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Preprocess input images using input processor.""" + start_time = time.time() + imgs_cpu, extrinsics, intrinsics = self.input_processor( + image, + extrinsics.copy() if extrinsics is not None else None, + intrinsics.copy() if intrinsics is not None else None, + process_res, + process_res_method, + ) + end_time = time.time() + logger.info( + "Processed Images Done taking", + end_time - start_time, + "seconds. Shape: ", + imgs_cpu.shape, + ) + return imgs_cpu, extrinsics, intrinsics + + def _prepare_model_inputs( + self, + imgs_cpu: torch.Tensor, + extrinsics: torch.Tensor | None, + intrinsics: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Prepare tensors for model input.""" + device = self._get_model_device() + + # Move images to model device + imgs = imgs_cpu.to(device, non_blocking=True)[None].float() + + # Convert camera parameters to tensors + ex_t = ( + extrinsics.to(device, non_blocking=True)[None].float() + if extrinsics is not None + else None + ) + in_t = ( + intrinsics.to(device, non_blocking=True)[None].float() + if intrinsics is not None + else None + ) + + return imgs, ex_t, in_t + + def _normalize_extrinsics(self, ex_t: torch.Tensor | None) -> torch.Tensor | None: + """Normalize extrinsics""" + if ex_t is None: + return None + transform = affine_inverse(ex_t[:, :1]) + ex_t_norm = ex_t @ transform + c2ws = affine_inverse(ex_t_norm) + translations = c2ws[..., :3, 3] + dists = translations.norm(dim=-1) + median_dist = torch.median(dists) + median_dist = torch.clamp(median_dist, min=1e-1) + ex_t_norm[..., :3, 3] = ex_t_norm[..., :3, 3] / median_dist + return ex_t_norm + + def _align_to_input_extrinsics_intrinsics( + self, + extrinsics: torch.Tensor | None, + intrinsics: torch.Tensor | None, + prediction: Prediction, + align_to_input_ext_scale: bool = True, + ransac_view_thresh: int = 10, + ) -> Prediction: + """Align depth map to input extrinsics""" + if extrinsics is None: + return prediction + prediction.intrinsics = intrinsics.numpy() + _, _, scale, aligned_extrinsics = align_poses_umeyama( + prediction.extrinsics, + extrinsics.numpy(), + ransac=len(extrinsics) >= ransac_view_thresh, + return_aligned=True, + random_state=42, + ) + if align_to_input_ext_scale: + prediction.extrinsics = extrinsics[..., :3, :].numpy() + prediction.depth /= scale + else: + prediction.extrinsics = aligned_extrinsics + return prediction + + def _run_model_forward( + self, + imgs: torch.Tensor, + ex_t: torch.Tensor | None, + in_t: torch.Tensor | None, + export_feat_layers: Sequence[int] | None = None, + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> dict[str, torch.Tensor]: + """Run model forward pass.""" + device = imgs.device + need_sync = device.type == "cuda" + if need_sync: + torch.cuda.synchronize(device) + start_time = time.time() + feat_layers = list(export_feat_layers) if export_feat_layers is not None else None + output = self.forward(imgs, ex_t, in_t, feat_layers, infer_gs, use_ray_pose, ref_view_strategy) + if need_sync: + torch.cuda.synchronize(device) + end_time = time.time() + logger.info(f"Model Forward Pass Done. Time: {end_time - start_time} seconds") + return output + + def _convert_to_prediction(self, raw_output: dict[str, torch.Tensor]) -> Prediction: + """Convert raw model output to Prediction object.""" + start_time = time.time() + output = self.output_processor(raw_output) + end_time = time.time() + logger.info(f"Conversion to Prediction Done. Time: {end_time - start_time} seconds") + return output + + def _add_processed_images(self, prediction: Prediction, imgs_cpu: torch.Tensor) -> Prediction: + """Add processed images to prediction for visualization.""" + # Convert from (N, 3, H, W) to (N, H, W, 3) and denormalize + processed_imgs = imgs_cpu.permute(0, 2, 3, 1).cpu().numpy() # (N, H, W, 3) + + # Denormalize from ImageNet normalization + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + processed_imgs = processed_imgs * std + mean + processed_imgs = np.clip(processed_imgs, 0, 1) + processed_imgs = (processed_imgs * 255).astype(np.uint8) + + prediction.processed_images = processed_imgs + return prediction + + def _export_results( + self, prediction: Prediction, export_format: str, export_dir: str, **kwargs + ) -> None: + """Export results to specified format and directory.""" + start_time = time.time() + export(prediction, export_format, export_dir, **kwargs) + end_time = time.time() + logger.info(f"Export Results Done. Time: {end_time - start_time} seconds") + + def _get_model_device(self) -> torch.device: + """ + Get the device where the model is located. + + Returns: + Device where the model parameters are located + + Raises: + ValueError: If no tensors are found in the model + """ + if self.device is not None: + return self.device + + # Find device from parameters + for param in self.parameters(): + self.device = param.device + return param.device + + # Find device from buffers + for buffer in self.buffers(): + self.device = buffer.device + return buffer.device + + raise ValueError("No tensor found in model") diff --git a/src/nn/depth_anything_3/app/css_and_html.py b/src/nn/depth_anything_3/app/css_and_html.py new file mode 100644 index 0000000..d414df9 --- /dev/null +++ b/src/nn/depth_anything_3/app/css_and_html.py @@ -0,0 +1,594 @@ +# flake8: noqa: E501 + +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +CSS and HTML content for the Depth Anything 3 Gradio application. +This module contains all the CSS styles and HTML content blocks +used in the Gradio interface. +""" + +# CSS Styles for the Gradio interface +GRADIO_CSS = """ +/* Add Font Awesome CDN with all styles including brands and colors */ +@import url('https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css'); + +/* Add custom styles for colored icons */ +.fa-color-blue { + color: #3b82f6; +} + +.fa-color-purple { + color: #8b5cf6; +} + +.fa-color-cyan { + color: #06b6d4; +} + +.fa-color-green { + color: #10b981; +} + +.fa-color-yellow { + color: #f59e0b; +} + +.fa-color-red { + color: #ef4444; +} + +.link-btn { + display: inline-flex; + align-items: center; + gap: 8px; + text-decoration: none; + padding: 12px 24px; + border-radius: 50px; + font-weight: 500; + transition: all 0.3s ease; +} + +/* Dark mode tech theme */ +@media (prefers-color-scheme: dark) { + html, body { + background: #1e293b; + color: #ffffff; + } + + .gradio-container { + background: #1e293b; + color: #ffffff; + } + + .link-btn { + background: rgba(255, 255, 255, 0.2); + color: white; + backdrop-filter: blur(10px); + border: 1px solid rgba(255, 255, 255, 0.3); + } + + .link-btn:hover { + background: rgba(255, 255, 255, 0.3); + transform: translateY(-2px); + box-shadow: 0 8px 25px rgba(0, 0, 0, 0.2); + } + + .tech-bg { + background: linear-gradient(135deg, #0f172a, #1e293b); /* Darker colors */ + position: relative; + overflow: hidden; + } + + .tech-bg::before { + content: ''; + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: + radial-gradient(circle at 20% 80%, rgba(59, 130, 246, 0.15) 0%, transparent 50%), /* Reduced opacity */ + radial-gradient(circle at 80% 20%, rgba(139, 92, 246, 0.15) 0%, transparent 50%), /* Reduced opacity */ + radial-gradient(circle at 40% 40%, rgba(18, 194, 233, 0.1) 0%, transparent 50%); /* Reduced opacity */ + animation: techPulse 8s ease-in-out infinite; + } + + .gradio-container .panel, + .gradio-container .block, + .gradio-container .form { + background: rgba(0, 0, 0, 0.3); + border: 1px solid rgba(59, 130, 246, 0.2); + border-radius: 10px; + } + + .gradio-container * { + color: #ffffff; + } + + .gradio-container label { + color: #e0e0e0; + } + + .gradio-container .markdown { + color: #e0e0e0; + } +} + +/* Light mode tech theme */ +@media (prefers-color-scheme: light) { + html, body { + background: #ffffff; + color: #1e293b; + } + + .gradio-container { + background: #ffffff; + color: #1e293b; + } + + .tech-bg { + background: linear-gradient(135deg, #ffffff, #f1f5f9); + position: relative; + overflow: hidden; + } + + .link-btn { + background: rgba(59, 130, 246, 0.15); + color: var(--body-text-color); + border: 1px solid rgba(59, 130, 246, 0.3); + } + + .link-btn:hover { + background: rgba(59, 130, 246, 0.25); + transform: translateY(-2px); + box-shadow: 0 8px 25px rgba(59, 130, 246, 0.2); + } + + .tech-bg::before { + content: ''; + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: + radial-gradient(circle at 20% 80%, rgba(59, 130, 246, 0.1) 0%, transparent 50%), + radial-gradient(circle at 80% 20%, rgba(139, 92, 246, 0.1) 0%, transparent 50%), + radial-gradient(circle at 40% 40%, rgba(18, 194, 233, 0.08) 0%, transparent 50%); + animation: techPulse 8s ease-in-out infinite; + } + + .gradio-container .panel, + .gradio-container .block, + .gradio-container .form { + background: rgba(255, 255, 255, 0.8); + border: 1px solid rgba(59, 130, 246, 0.3); + border-radius: 10px; + box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); + } + + .gradio-container * { + color: #1e293b; + } + + .gradio-container label { + color: #334155; + } + + .gradio-container .markdown { + color: #334155; + } +} + + + + +@keyframes techPulse { + 0%, 100% { opacity: 0.5; } + 50% { opacity: 0.8; } +} + +/* Custom log with tech gradient */ +.custom-log * { + font-style: italic; + font-size: 22px !important; + background: linear-gradient(135deg, #3b82f6, #8b5cf6); + background-size: 400% 400%; + -webkit-background-clip: text; + background-clip: text; + font-weight: bold !important; + color: transparent !important; + text-align: center !important; + animation: techGradient 3s ease infinite; +} + +@keyframes techGradient { + 0% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } + 100% { background-position: 0% 50%; } +} + +@keyframes metricPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +@keyframes pointcloudPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +@keyframes camerasPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +@keyframes gaussiansPulse { + 0%, 100% { background-position: 0% 50%; } + 50% { background-position: 100% 50%; } +} + +/* Special colors for key terms - Global styles */ +.metric-text { + background: linear-gradient(45deg, #ff6b6b, #ff8e53, #ff6b6b); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: metricPulse 2s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(255, 107, 107, 0.5); +} + +.pointcloud-text { + background: linear-gradient(45deg, #4ecdc4, #44a08d, #4ecdc4); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: pointcloudPulse 2.5s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(78, 205, 196, 0.5); +} + +.cameras-text { + background: linear-gradient(45deg, #667eea, #764ba2, #667eea); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: camerasPulse 3s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(102, 126, 234, 0.5); +} + +.gaussians-text { + background: linear-gradient(45deg, #f093fb, #f5576c, #f093fb); + background-size: 200% 200%; + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; + animation: gaussiansPulse 2.2s ease-in-out infinite; + font-weight: 700; + text-shadow: 0 0 10px rgba(240, 147, 251, 0.5); +} + +.example-log * { + font-style: italic; + font-size: 16px !important; + background: linear-gradient(135deg, #3b82f6, #8b5cf6); + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; +} + +#my_radio .wrap { + display: flex; + flex-wrap: nowrap; + justify-content: center; + align-items: center; +} + +#my_radio .wrap label { + display: flex; + width: 50%; + justify-content: center; + align-items: center; + margin: 0; + padding: 10px 0; + box-sizing: border-box; +} + +/* Align navigation buttons with dropdown bottom */ +.navigation-row { + display: flex !important; + align-items: flex-end !important; + gap: 8px !important; +} + +.navigation-row > div:nth-child(1), +.navigation-row > div:nth-child(3) { + align-self: flex-end !important; +} + +.navigation-row > div:nth-child(2) { + flex: 1 !important; +} + +/* Make thumbnails clickable with pointer cursor */ +.clickable-thumbnail img { + cursor: pointer !important; +} + +.clickable-thumbnail:hover img { + cursor: pointer !important; + opacity: 0.8; + transition: opacity 0.3s ease; +} + +/* Make thumbnail containers narrower horizontally */ +.clickable-thumbnail { + padding: 5px 2px !important; + margin: 0 2px !important; +} + +.clickable-thumbnail .image-container { + margin: 0 !important; + padding: 0 !important; +} + +.scene-info { + text-align: center !important; + padding: 5px 2px !important; + margin: 0 !important; +} +""" + + +def get_header_html(logo_base64=None): + """ + Generate the main header HTML with logo and title. + + Args: + logo_base64 (str, optional): Base64 encoded logo image + + Returns: + str: HTML string for the header + """ + return """ +
+
+

+ Depth Anything 3 +

+

+ Recovering the Visual Space from Any Views +

+ +
+
+ + + """ + + +def get_description_html(): + """ + Generate the main description and getting started HTML. + + Returns: + str: HTML string for the description + """ + return """ +
+

+ What This Demo Does +

+
+

+ Upload images or videos → Get Metric Point Clouds, Cameras and Novel Views → Explore in 3D +

+
+ +
+

+ Tip: Landscape-oriented images or videos are preferred for best 3D recovering. +

+
+
+ + + """ + + +def get_acknowledgements_html(): + """ + Generate the acknowledgements section HTML. + + Returns: + str: HTML string for the acknowledgements + """ + return """ +
+

+ Research Credits & Acknowledgments +

+ +
+ +
+

Original Research

+

+ + Depth Anything 3 + +

+
+ + +
+

Previous Versions

+ +
+
+ + +
+

+ HF demo adapted from Map Anything +

+
+
+ """ + + +def get_gradio_theme(): + """ + Get the configured Gradio theme with adaptive tech colors. + + Returns: + gr.themes.Base: Configured Gradio theme + """ + import gradio as gr + + return gr.themes.Base( + primary_hue=gr.themes.Color( + c50="#eff6ff", + c100="#dbeafe", + c200="#bfdbfe", + c300="#93c5fd", + c400="#60a5fa", + c500="#3b82f6", + c600="#2563eb", + c700="#1d4ed8", + c800="#1e40af", + c900="#1e3a8a", + c950="#172554", + ), + secondary_hue=gr.themes.Color( + c50="#f5f3ff", + c100="#ede9fe", + c200="#ddd6fe", + c300="#c4b5fd", + c400="#a78bfa", + c500="#8b5cf6", + c600="#7c3aed", + c700="#6d28d9", + c800="#5b21b6", + c900="#4c1d95", + c950="#2e1065", + ), + neutral_hue=gr.themes.Color( + c50="#f8fafc", + c100="#f1f5f9", + c200="#e2e8f0", + c300="#cbd5e1", + c400="#94a3b8", + c500="#64748b", + c600="#475569", + c700="#334155", + c800="#1e293b", + c900="#0f172a", + c950="#020617", + ), + ) + + +# Measure tab instructions HTML +MEASURE_INSTRUCTIONS_HTML = """ +### Click points on the image to compute distance. +> Metric scale estimation is difficult on aerial/drone images. +""" diff --git a/src/nn/depth_anything_3/app/gradio_app.py b/src/nn/depth_anything_3/app/gradio_app.py new file mode 100644 index 0000000..ef188c0 --- /dev/null +++ b/src/nn/depth_anything_3/app/gradio_app.py @@ -0,0 +1,724 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Refactored Gradio App for Depth Anything 3. + +This is the main application file that orchestrates all components. +The original functionality has been split into modular components for better maintainability. +""" + +import argparse +import os +from typing import Any, Dict, List +import gradio as gr + +from depth_anything_3.app.css_and_html import GRADIO_CSS, get_gradio_theme +from depth_anything_3.app.modules.event_handlers import EventHandlers +from depth_anything_3.app.modules.ui_components import UIComponents + +# Set environment variables +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + + +class DepthAnything3App: + """ + Main application class for Depth Anything 3 Gradio app. + """ + + def __init__(self, model_dir: str = None, workspace_dir: str = None, gallery_dir: str = None): + """ + Initialize the application. + + Args: + model_dir: Path to the model directory + workspace_dir: Path to the workspace directory + gallery_dir: Path to the gallery directory + """ + self.model_dir = model_dir + self.workspace_dir = workspace_dir + self.gallery_dir = gallery_dir + + # Set environment variables for directories + if self.model_dir: + os.environ["DA3_MODEL_DIR"] = self.model_dir + if self.workspace_dir: + os.environ["DA3_WORKSPACE_DIR"] = self.workspace_dir + if self.gallery_dir: + os.environ["DA3_GALLERY_DIR"] = self.gallery_dir + + self.event_handlers = EventHandlers() + self.ui_components = UIComponents() + + def cache_examples( + self, + show_cam: bool = True, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + save_percentage: float = 20.0, + num_max_points: int = 1000, + cache_gs_tag: str = "", + gs_trj_mode: str = "smooth", + gs_video_quality: str = "low", + ) -> None: + """ + Pre-cache all example scenes at startup. + + Args: + show_cam: Whether to show camera in visualization + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + save_percentage: Filter percentage for point cloud + num_max_points: Maximum number of points + cache_gs_tag: Tag to match scene names for high-res+3DGS caching (e.g., "dl3dv") + gs_trj_mode: Trajectory mode for 3DGS + gs_video_quality: Video quality for 3DGS + """ + from depth_anything_3.app.modules.utils import get_scene_info + + examples_dir = os.path.join(self.workspace_dir, "examples") + if not os.path.exists(examples_dir): + print(f"Examples directory not found: {examples_dir}") + return + + scenes = get_scene_info(examples_dir) + if not scenes: + print("No example scenes found to cache.") + return + + print(f"\n{'='*60}") + print(f"Caching {len(scenes)} example scenes...") + print(f"{'='*60}\n") + + for i, scene in enumerate(scenes, 1): + scene_name = scene["name"] + + # Check if scene name matches the gs tag for high-res+3DGS caching + use_high_res_gs = cache_gs_tag and cache_gs_tag.lower() in scene_name.lower() + + if use_high_res_gs: + print(f"[{i}/{len(scenes)}] Caching scene: {scene_name} (HIGH-RES + 3DGS)") + print(f" - Number of images: {scene['num_images']}") + print(f" - Matched tag: '{cache_gs_tag}' - using high_res + 3DGS") + else: + print(f"[{i}/{len(scenes)}] Caching scene: {scene_name} (LOW-RES)") + print(f" - Number of images: {scene['num_images']}") + + try: + # Load example scene + _, target_dir, _, _, _, _, _, _, _ = self.event_handlers.load_example_scene( + scene_name + ) + + if target_dir and target_dir != "None": + # Run reconstruction with appropriate settings + print(" - Running reconstruction...") + result = self.event_handlers.gradio_demo( + target_dir=target_dir, + show_cam=show_cam, + filter_black_bg=filter_black_bg, + filter_white_bg=filter_white_bg, + process_res_method="high_res" if use_high_res_gs else "low_res", + save_percentage=save_percentage, + num_max_points=num_max_points, + infer_gs=use_high_res_gs, + ref_view_strategy="saddle_balanced", + gs_trj_mode=gs_trj_mode, + gs_video_quality=gs_video_quality, + ) + + # Check if successful + if result[0] is not None: # reconstruction_output + print(f" ✓ Scene '{scene_name}' cached successfully") + else: + print(f" ✗ Scene '{scene_name}' caching failed: {result[1]}") + else: + print(f" ✗ Scene '{scene_name}' loading failed") + + except Exception as e: + print(f" ✗ Error caching scene '{scene_name}': {str(e)}") + + print() + + print("=" * 60) + print("Example scene caching completed!") + print("=" * 60 + "\n") + + def create_app(self) -> gr.Blocks: + """ + Create and configure the Gradio application. + + Returns: + Configured Gradio Blocks interface + """ + + # Initialize theme + def get_theme(): + return get_gradio_theme() + + with gr.Blocks(theme=get_theme(), css=GRADIO_CSS) as demo: + # State variables for the tabbed interface + is_example = gr.Textbox(label="is_example", visible=False, value="None") + processed_data_state = gr.State(value=None) + measure_points_state = gr.State(value=[]) + selected_image_index_state = gr.State(value=0) # Track selected image index + # current_view_index = gr.State(value=0) # noqa: F841 Track current view index + + # Header and description + self.ui_components.create_header_section() + self.ui_components.create_description_section() + + target_dir_output = gr.Textbox(label="Target Dir", visible=False, value="None") + + # Main content area + with gr.Row(): + with gr.Column(scale=2): + # Upload section + ( + input_video, + s_time_interval, + input_images, + image_gallery, + ) = self.ui_components.create_upload_section() + + with gr.Column(scale=4): + with gr.Column(): + # gr.Markdown("**Metric 3D Reconstruction (Point Cloud and Camera Poses)**") + # Reconstruction control section (buttons) - moved below tabs + + log_output = gr.Markdown( + "Please upload a video or images, then click Reconstruct.", + elem_classes=["custom-log"], + ) + + # Tabbed interface + with gr.Tabs(): + with gr.Tab("Point Cloud & Cameras"): + reconstruction_output = ( + self.ui_components.create_3d_viewer_section() + ) + + with gr.Tab("Metric Depth"): + ( + prev_measure_btn, + measure_view_selector, + next_measure_btn, + measure_image, + measure_depth_image, + measure_text, + ) = self.ui_components.create_measure_section() + + with gr.Tab("3DGS Rendered Novel Views"): + gs_video, gs_info = self.ui_components.create_nvs_video() + + # Inference control section (before inference) + (process_res_method_dropdown, infer_gs, ref_view_strategy_dropdown) = ( + self.ui_components.create_inference_control_section() + ) + + # Display control section - includes 3DGS options, buttons, and Visualization Options # noqa: E501 + ( + show_cam, + filter_black_bg, + filter_white_bg, + save_percentage, + num_max_points, + gs_trj_mode, + gs_video_quality, + submit_btn, + clear_btn, + ) = self.ui_components.create_display_control_section() + + # bind visibility of gs_trj_mode to infer_gs + infer_gs.change( + fn=lambda checked: ( + gr.update(visible=checked), + gr.update(visible=checked), + gr.update(visible=checked), + gr.update(visible=(not checked)), + ), + inputs=infer_gs, + outputs=[gs_trj_mode, gs_video_quality, gs_video, gs_info], + ) + + # Example scenes section + gr.Markdown("## Example Scenes") + + scenes = self.ui_components.create_example_scenes_section() + scene_components = self.ui_components.create_example_scene_grid(scenes) + + # Set up event handlers + self._setup_event_handlers( + demo, + is_example, + processed_data_state, + measure_points_state, + target_dir_output, + input_video, + input_images, + s_time_interval, + image_gallery, + reconstruction_output, + log_output, + show_cam, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + save_percentage, + submit_btn, + clear_btn, + num_max_points, + infer_gs, + ref_view_strategy_dropdown, + selected_image_index_state, + measure_view_selector, + measure_image, + measure_depth_image, + measure_text, + prev_measure_btn, + next_measure_btn, + scenes, + scene_components, + gs_video, + gs_info, + gs_trj_mode, + gs_video_quality, + ) + + # Acknowledgements + self.ui_components.create_acknowledgements_section() + + return demo + + def _setup_event_handlers( + self, + demo: gr.Blocks, + is_example: gr.Textbox, + processed_data_state: gr.State, + measure_points_state: gr.State, + target_dir_output: gr.Textbox, + input_video: gr.Video, + input_images: gr.File, + s_time_interval: gr.Slider, + image_gallery: gr.Gallery, + reconstruction_output: gr.Model3D, + log_output: gr.Markdown, + show_cam: gr.Checkbox, + filter_black_bg: gr.Checkbox, + filter_white_bg: gr.Checkbox, + process_res_method_dropdown: gr.Dropdown, + save_percentage: gr.Slider, + submit_btn: gr.Button, + clear_btn: gr.ClearButton, + num_max_points: gr.Slider, + infer_gs: gr.Checkbox, + ref_view_strategy_dropdown: gr.Dropdown, + selected_image_index_state: gr.State, + measure_view_selector: gr.Dropdown, + measure_image: gr.Image, + measure_depth_image: gr.Image, + measure_text: gr.Markdown, + prev_measure_btn: gr.Button, + next_measure_btn: gr.Button, + scenes: List[Dict[str, Any]], + scene_components: List[gr.Image], + gs_video: gr.Video, + gs_info: gr.Markdown, + gs_trj_mode: gr.Dropdown, + gs_video_quality: gr.Dropdown, + ) -> None: + """ + Set up all event handlers for the application. + + Args: + demo: Gradio Blocks interface + All other arguments: Gradio components to connect + """ + # Configure clear button + clear_btn.add( + [ + input_video, + input_images, + reconstruction_output, + log_output, + target_dir_output, + image_gallery, + gs_video, + ] + ) + + # Main reconstruction button + submit_btn.click( + fn=self.event_handlers.clear_fields, inputs=[], outputs=[reconstruction_output] + ).then(fn=self.event_handlers.update_log, inputs=[], outputs=[log_output]).then( + fn=self.event_handlers.gradio_demo, + inputs=[ + target_dir_output, + show_cam, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + save_percentage, + # pass num_max_points + num_max_points, + infer_gs, + ref_view_strategy_dropdown, + gs_trj_mode, + gs_video_quality, + ], + outputs=[ + reconstruction_output, + log_output, + processed_data_state, + measure_image, + measure_depth_image, + measure_text, + measure_view_selector, + gs_video, + gs_video, # gs_video visibility + gs_info, # gs_info visibility + ], + ).then( + fn=lambda: "False", + inputs=[], + outputs=[is_example], # set is_example to "False" + ) + + # Real-time visualization updates + self._setup_visualization_handlers( + show_cam, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + target_dir_output, + is_example, + reconstruction_output, + log_output, + ) + + # File upload handlers + input_video.change( + fn=self.event_handlers.handle_uploads, + inputs=[input_video, input_images, s_time_interval], + outputs=[reconstruction_output, target_dir_output, image_gallery, log_output], + ) + input_images.change( + fn=self.event_handlers.handle_uploads, + inputs=[input_video, input_images, s_time_interval], + outputs=[reconstruction_output, target_dir_output, image_gallery, log_output], + ) + + # Navigation handlers + self._setup_navigation_handlers( + prev_measure_btn, + next_measure_btn, + measure_view_selector, + measure_image, + measure_depth_image, + measure_points_state, + processed_data_state, + ) + + # Measurement handler + measure_image.select( + fn=self.event_handlers.measure, + inputs=[processed_data_state, measure_points_state, measure_view_selector], + outputs=[measure_image, measure_depth_image, measure_points_state, measure_text], + ) + + # Example scene handlers + self._setup_example_scene_handlers( + scenes, + scene_components, + reconstruction_output, + target_dir_output, + image_gallery, + log_output, + is_example, + processed_data_state, + measure_view_selector, + measure_image, + measure_depth_image, + gs_video, + gs_info, + ) + + def _setup_visualization_handlers( + self, + show_cam: gr.Checkbox, + filter_black_bg: gr.Checkbox, + filter_white_bg: gr.Checkbox, + process_res_method_dropdown: gr.Dropdown, + target_dir_output: gr.Textbox, + is_example: gr.Textbox, + reconstruction_output: gr.Model3D, + log_output: gr.Markdown, + ) -> None: + """Set up visualization update handlers.""" + # Common inputs for visualization updates + viz_inputs = [ + target_dir_output, + show_cam, + is_example, + filter_black_bg, + filter_white_bg, + process_res_method_dropdown, + ] + + # Set up change handlers for all visualization controls + for component in [show_cam, filter_black_bg, filter_white_bg]: + component.change( + fn=self.event_handlers.update_visualization, + inputs=viz_inputs, + outputs=[reconstruction_output, log_output], + ) + + def _setup_navigation_handlers( + self, + prev_measure_btn: gr.Button, + next_measure_btn: gr.Button, + measure_view_selector: gr.Dropdown, + measure_image: gr.Image, + measure_depth_image: gr.Image, + measure_points_state: gr.State, + processed_data_state: gr.State, + ) -> None: + """Set up navigation handlers for measure tab.""" + # Measure tab navigation + prev_measure_btn.click( + fn=lambda processed_data, current_selector: self.event_handlers.navigate_measure_view( + processed_data, current_selector, -1 + ), + inputs=[processed_data_state, measure_view_selector], + outputs=[ + measure_view_selector, + measure_image, + measure_depth_image, + measure_points_state, + ], + ) + + next_measure_btn.click( + fn=lambda processed_data, current_selector: self.event_handlers.navigate_measure_view( + processed_data, current_selector, 1 + ), + inputs=[processed_data_state, measure_view_selector], + outputs=[ + measure_view_selector, + measure_image, + measure_depth_image, + measure_points_state, + ], + ) + + measure_view_selector.change( + fn=lambda processed_data, selector_value: ( + self.event_handlers.update_measure_view( + processed_data, int(selector_value.split()[1]) - 1 + ) + if selector_value + else (None, None, []) + ), + inputs=[processed_data_state, measure_view_selector], + outputs=[measure_image, measure_depth_image, measure_points_state], + ) + + def _setup_example_scene_handlers( + self, + scenes: List[Dict[str, Any]], + scene_components: List[gr.Image], + reconstruction_output: gr.Model3D, + target_dir_output: gr.Textbox, + image_gallery: gr.Gallery, + log_output: gr.Markdown, + is_example: gr.Textbox, + processed_data_state: gr.State, + measure_view_selector: gr.Dropdown, + measure_image: gr.Image, + measure_depth_image: gr.Image, + gs_video: gr.Video, + gs_info: gr.Markdown, + ) -> None: + """Set up example scene handlers.""" + + def load_and_update_measure(name): + result = self.event_handlers.load_example_scene(name) + # result = (reconstruction_output, target_dir, image_paths, log_message, processed_data, measure_view_selector, gs_video, gs_video_vis, gs_info_vis) # noqa: E501 + + # Update measure view if processed_data is available + measure_img = None + measure_depth = None + if result[4] is not None: # processed_data exists + measure_img, measure_depth, _ = ( + self.event_handlers.visualization_handler.update_measure_view(result[4], 0) + ) + + return result + ("True", measure_img, measure_depth) + + for i, scene in enumerate(scenes): + if i < len(scene_components): + scene_components[i].select( + fn=lambda name=scene["name"]: load_and_update_measure(name), + outputs=[ + reconstruction_output, + target_dir_output, + image_gallery, + log_output, + processed_data_state, + measure_view_selector, + gs_video, + gs_video, # gs_video_visibility + gs_info, # gs_info_visibility + is_example, + measure_image, + measure_depth_image, + ], + ) + + def launch(self, host: str = "127.0.0.1", port: int = 7860, **kwargs) -> None: + """ + Launch the application. + + Args: + host: Host address to bind to + port: Port number to bind to + **kwargs: Additional arguments for demo.launch() + """ + demo = self.create_app() + demo.queue(max_size=20).launch( + show_error=True, ssr_mode=False, server_name=host, server_port=port, **kwargs + ) + + +def main(): + """Main function to run the application.""" + parser = argparse.ArgumentParser( + description="Depth Anything 3 Gradio Application", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Basic usage + python gradio_app.py --help + python gradio_app.py --host 0.0.0.0 --port 8080 + python gradio_app.py --model-dir /path/to/model --workspace-dir /path/to/workspace + + # Cache examples at startup (all low-res) + python gradio_app.py --cache-examples + + # Cache with selective high-res+3DGS for scenes matching tag + python gradio_app.py --cache-examples --cache-gs-tag dl3dv + # This will use high-res + 3DGS for scenes containing "dl3dv" in their name, + # and low-res only for other scenes + """, + ) + + # Server configuration + parser.add_argument( + "--host", default="127.0.0.1", help="Host address to bind to (default: 127.0.0.1)" + ) + parser.add_argument( + "--port", type=int, default=7860, help="Port number to bind to (default: 7860)" + ) + + # Directory configuration + parser.add_argument( + "--model-dir", + default="depth-anything/DA3NESTED-GIANT-LARGE", + help="Path to the model directory (default: depth-anything/DA3NESTED-GIANT-LARGE)", + ) + parser.add_argument( + "--workspace-dir", + default="workspace/gradio", # noqa: E501 + help="Path to the workspace directory (default: workspace/gradio)", # noqa: E501 + ) + parser.add_argument( + "--gallery-dir", + default="workspace/gallery", + help="Path to the gallery directory (default: workspace/gallery)", # noqa: E501 + ) + + # Additional Gradio options + parser.add_argument("--share", action="store_true", help="Create a public link for the app") + parser.add_argument("--debug", action="store_true", help="Enable debug mode") + + # Example caching options + parser.add_argument( + "--cache-examples", + action="store_true", + help="Pre-cache all example scenes at startup for faster loading", + ) + parser.add_argument( + "--cache-gs-tag", + type=str, + default="", + help="Tag to match scene names for high-res+3DGS caching (e.g., 'dl3dv'). Scenes containing this tag will use high_res and infer_gs=True; others will use low_res only.", # noqa: E501 + ) + + args = parser.parse_args() + + # Create directories if they don't exist + os.makedirs(args.workspace_dir, exist_ok=True) + os.makedirs(args.gallery_dir, exist_ok=True) + + # Initialize and launch the application + app = DepthAnything3App( + model_dir=args.model_dir, workspace_dir=args.workspace_dir, gallery_dir=args.gallery_dir + ) + + # Prepare launch arguments + launch_kwargs = {"share": args.share, "debug": args.debug} + + print("Starting Depth Anything 3 Gradio App...") + print(f"Host: {args.host}") + print(f"Port: {args.port}") + print(f"Model Directory: {args.model_dir}") + print(f"Workspace Directory: {args.workspace_dir}") + print(f"Gallery Directory: {args.gallery_dir}") + print(f"Share: {args.share}") + print(f"Debug: {args.debug}") + print(f"Cache Examples: {args.cache_examples}") + if args.cache_examples: + if args.cache_gs_tag: + print( + f"Cache GS Tag: '{args.cache_gs_tag}' (scenes matching this tag will use high-res + 3DGS)" # noqa: E501 + ) # noqa: E501 + else: + print("Cache GS Tag: None (all scenes will use low-res only)") + + # Pre-cache examples if requested + if args.cache_examples: + print("\n" + "=" * 60) + print("Pre-caching mode enabled") + if args.cache_gs_tag: + print(f"Scenes containing '{args.cache_gs_tag}' will use HIGH-RES + 3DGS") + print("Other scenes will use LOW-RES only") + else: + print("All scenes will use LOW-RES only") + print("=" * 60) + app.cache_examples( + show_cam=True, + filter_black_bg=False, + filter_white_bg=False, + save_percentage=5.0, + num_max_points=1000, + cache_gs_tag=args.cache_gs_tag, + gs_trj_mode="smooth", + gs_video_quality="low", + ) + + app.launch(host=args.host, port=args.port, **launch_kwargs) + + +if __name__ == "__main__": + main() diff --git a/src/nn/depth_anything_3/app/modules/__init__.py b/src/nn/depth_anything_3/app/modules/__init__.py new file mode 100644 index 0000000..dd0b717 --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Modules package for Depth Anything 3 Gradio app. + +This package contains all the modular components for the Gradio application. +""" + +from depth_anything_3.app.modules.event_handlers import EventHandlers +from depth_anything_3.app.modules.file_handlers import FileHandler +from depth_anything_3.app.modules.model_inference import ModelInference +from depth_anything_3.app.modules.ui_components import UIComponents +from depth_anything_3.app.modules.utils import ( + create_depth_visualization, + get_logo_base64, + get_scene_info, + save_to_gallery_func, +) +from depth_anything_3.app.modules.visualization import VisualizationHandler + +__all__ = [ + "ModelInference", + "FileHandler", + "VisualizationHandler", + "EventHandlers", + "UIComponents", + "create_depth_visualization", + "save_to_gallery_func", + "get_scene_info", + "get_logo_base64", +] diff --git a/src/nn/depth_anything_3/app/modules/event_handlers.py b/src/nn/depth_anything_3/app/modules/event_handlers.py new file mode 100644 index 0000000..c494edb --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/event_handlers.py @@ -0,0 +1,619 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Event handling module for Depth Anything 3 Gradio app. + +This module handles all event callbacks and user interactions. +""" + +import os +import time +from glob import glob +from typing import Any, Dict, List, Optional, Tuple +import gradio as gr +import numpy as np +import torch + +from depth_anything_3.app.modules.file_handlers import FileHandler +from depth_anything_3.app.modules.model_inference import ModelInference +from depth_anything_3.utils.memory import cleanup_cuda_memory +from depth_anything_3.app.modules.visualization import VisualizationHandler + + +class EventHandlers: + """ + Handles all event callbacks and user interactions for the Gradio app. + """ + + def __init__(self): + """Initialize the event handlers.""" + self.model_inference = ModelInference() + self.file_handler = FileHandler() + self.visualization_handler = VisualizationHandler() + + def clear_fields(self) -> None: + """ + Clears the 3D viewer, the stored target_dir, and empties the gallery. + """ + return None + + def update_log(self) -> str: + """ + Display a quick log message while waiting. + """ + return "Loading and Reconstructing..." + + def save_current_visualization( + self, + target_dir: str, + save_percentage: float, + show_cam: bool, + filter_black_bg: bool, + filter_white_bg: bool, + processed_data: Optional[Dict], + scene_name: str = "", + ) -> str: + """ + Save current visualization results to gallery with specified save percentage. + + Args: + target_dir: Directory containing results + save_percentage: Percentage of points to save (0-100) + show_cam: Whether to show cameras + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + processed_data: Processed data from reconstruction + + Returns: + Status message + """ + if not target_dir or target_dir == "None" or not os.path.isdir(target_dir): + return "No reconstruction available. Please run 'Reconstruct' first." + + if processed_data is None: + return "No processed data available. Please run 'Reconstruct' first." + + try: + # Add debug information + print("[DEBUG] save_current_visualization called with:") + print(f" target_dir: {target_dir}") + print(f" save_percentage: {save_percentage}") + print(f" show_cam: {show_cam}") + print(f" filter_black_bg: {filter_black_bg}") + print(f" filter_white_bg: {filter_white_bg}") + print(f" processed_data: {processed_data is not None}") + + # Import the gallery save function + # Create gallery name with user input or auto-generated + import datetime + + from .utils import save_to_gallery_func + + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + if scene_name and scene_name.strip(): + gallery_name = f"{scene_name.strip()}_{timestamp}_pct{save_percentage:.0f}" + else: + gallery_name = f"save_{timestamp}_pct{save_percentage:.0f}" + + print(f"[DEBUG] Saving to gallery with name: {gallery_name}") + + # Save entire process folder to gallery + success, message = save_to_gallery_func( + target_dir=target_dir, processed_data=processed_data, gallery_name=gallery_name + ) + + if success: + print(f"[DEBUG] Gallery save completed successfully: {message}") + return ( + "Successfully saved to gallery!\n" + f"Gallery name: {gallery_name}\n" + f"Save percentage: {save_percentage}%\n" + f"Show cameras: {show_cam}\n" + f"Filter black bg: {filter_black_bg}\n" + f"Filter white bg: {filter_white_bg}\n\n" + f"{message}" + ) + else: + print(f"[DEBUG] Gallery save failed: {message}") + return f"Failed to save to gallery: {message}" + + except Exception as e: + return f"Error saving visualization: {str(e)}" + + def gradio_demo( + self, + target_dir: str, + show_cam: bool = True, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + process_res_method: str = "upper_bound_resize", + save_percentage: float = 30.0, + num_max_points: int = 1_000_000, + infer_gs: bool = False, + ref_view_strategy: str = "saddle_balanced", + gs_trj_mode: str = "extend", + gs_video_quality: str = "high", + ) -> Tuple[ + Optional[str], + str, + Optional[Dict], + Optional[np.ndarray], + Optional[np.ndarray], + str, + gr.Dropdown, + Optional[str], # gs video path + gr.update, # gs video visibility update + gr.update, # gs info visibility update + ]: + """ + Perform reconstruction using the already-created target_dir/images. + + Args: + target_dir: Directory containing images + show_cam: Whether to show camera + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + process_res_method: Method for resizing input images + save_percentage: Filter percentage for point cloud + num_max_points: Maximum number of points + infer_gs: Whether to infer 3D Gaussian Splatting + ref_view_strategy: Reference view selection strategy + + Returns: + Tuple of reconstruction results + """ + if not os.path.isdir(target_dir) or target_dir == "None": + return ( + None, + "No valid target directory found. Please upload first.", + None, + None, + None, + "", + None, + None, + gr.update(visible=False), # gs_video + gr.update(visible=True), # gs_info + ) + + start_time = time.time() + cleanup_cuda_memory() + + # Get image files for logging + target_dir_images = os.path.join(target_dir, "images") + all_files = ( + sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else [] + ) + + print("Running DepthAnything3 model...") + print(f"Reference view strategy: {ref_view_strategy}") + + with torch.no_grad(): + prediction, processed_data = self.model_inference.run_inference( + target_dir, + process_res_method=process_res_method, + show_camera=show_cam, + save_percentage=save_percentage, + num_max_points=int(num_max_points * 1000), # Convert K to actual count + infer_gs=infer_gs, + ref_view_strategy=ref_view_strategy, + gs_trj_mode=gs_trj_mode, + gs_video_quality=gs_video_quality, + ) + + # The GLB file is already generated by the API + glbfile = os.path.join(target_dir, "scene.glb") + + # Handle 3DGS video based on infer_gs flag + gsvideo_path = None + gs_video_visible = False + gs_info_visible = True + + if infer_gs: + try: + gsvideo_path = sorted(glob(os.path.join(target_dir, "gs_video", "*.mp4")))[-1] + gs_video_visible = True + gs_info_visible = False + except IndexError: + gsvideo_path = None + print("3DGS video not found, but infer_gs was enabled") + + # Cleanup + cleanup_cuda_memory() + + end_time = time.time() + print(f"Total time: {end_time - start_time:.2f} seconds") + log_msg = f"Reconstruction Success ({len(all_files)} frames). Waiting for visualization." + + # Populate visualization tabs with processed data + depth_vis, measure_img, measure_depth_vis, measure_pts = ( + self.visualization_handler.populate_visualization_tabs(processed_data) + ) + + # Update view selectors based on available views + depth_selector, measure_selector = self.visualization_handler.update_view_selectors( + processed_data + ) + + return ( + glbfile, + log_msg, + processed_data, + measure_img, # measure_image + measure_depth_vis, # measure_depth_image + "", # measure_text (empty initially) + measure_selector, # measure_view_selector + gsvideo_path, + gr.update(visible=gs_video_visible), # gs_video visibility + gr.update(visible=gs_info_visible), # gs_info visibility + ) + + def update_visualization( + self, + target_dir: str, + show_cam: bool, + is_example: str, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + process_res_method: str = "upper_bound_resize", + ) -> Tuple[gr.update, str]: + """ + Reload saved predictions from npz, create (or reuse) the GLB for new parameters, + and return it for the 3D viewer. + + Args: + target_dir: Directory containing results + show_cam: Whether to show camera + is_example: Whether this is an example scene + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + process_res_method: Method for resizing input images + + Returns: + Tuple of (glb_file, log_message) + """ + if not target_dir or target_dir == "None" or not os.path.isdir(target_dir): + return ( + gr.update(), + "No reconstruction available. Please click the Reconstruct button first.", + ) + + # Check if GLB exists (could be cached example or reconstructed scene) + glbfile = os.path.join(target_dir, "scene.glb") + if os.path.exists(glbfile): + return ( + glbfile, + ( + "Visualization loaded from cache." + if is_example == "True" + else "Visualization updated." + ), + ) + + # If no GLB but it's an example that hasn't been reconstructed yet + if is_example == "True": + return ( + gr.update(), + "No reconstruction available. Please click the Reconstruct button first.", + ) + + # For non-examples, check predictions.npz + predictions_path = os.path.join(target_dir, "predictions.npz") + if not os.path.exists(predictions_path): + error_message = ( + f"No reconstruction available at {predictions_path}. " + "Please run 'Reconstruct' first." + ) + return gr.update(), error_message + + loaded = np.load(predictions_path, allow_pickle=True) + predictions = {key: loaded[key] for key in loaded.keys()} # noqa: F841 + + return ( + glbfile, + "Visualization updated.", + ) + + def handle_uploads( + self, + input_video: Optional[str], + input_images: Optional[List], + s_time_interval: float = 10.0, + ) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[str]]: + """ + Handle file uploads and update gallery. + + Args: + input_video: Path to input video file + input_images: List of input image files + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message) + """ + return self.file_handler.update_gallery_on_upload( + input_video, input_images, s_time_interval + ) + + def load_example_scene(self, scene_name: str, examples_dir: str = None) -> Tuple[ + Optional[str], + Optional[str], + Optional[List], + str, + Optional[Dict], + gr.Dropdown, + Optional[str], + gr.update, + gr.update, + ]: + """ + Load a scene from examples directory. + + Args: + scene_name: Name of the scene to load + examples_dir: Path to examples directory (if None, uses workspace_dir/examples) + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message, processed_data, measure_view_selector, gs_video, gs_video_vis, gs_info_vis) # noqa: E501 + """ + if examples_dir is None: + # Get workspace directory from environment variable + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + examples_dir = os.path.join(workspace_dir, "examples") + + reconstruction_output, target_dir, image_paths, log_message = ( + self.file_handler.load_example_scene(scene_name, examples_dir) + ) + + # Try to load cached processed data if available + processed_data = None + measure_view_selector = gr.Dropdown(choices=["View 1"], value="View 1") + gs_video_path = None + gs_video_visible = False + gs_info_visible = True + + if target_dir and target_dir != "None": + predictions_path = os.path.join(target_dir, "predictions.npz") + if os.path.exists(predictions_path): + try: + # Load predictions from cache + loaded = np.load(predictions_path, allow_pickle=True) + predictions = {key: loaded[key] for key in loaded.keys()} + + # Reconstruct processed_data structure + num_images = len(predictions.get("images", [])) + processed_data = {} + + for i in range(num_images): + processed_data[i] = { + "image": predictions["images"][i] if "images" in predictions else None, + "depth": predictions["depths"][i] if "depths" in predictions else None, + "depth_image": os.path.join( + target_dir, "depth_vis", f"{i:04d}.jpg" # Fixed: use .jpg not .png + ), + "intrinsics": ( + predictions["intrinsics"][i] + if "intrinsics" in predictions + and i < len(predictions["intrinsics"]) + else None + ), + "mask": None, + } + + # Update measure view selector + choices = [f"View {i + 1}" for i in range(num_images)] + measure_view_selector = gr.Dropdown(choices=choices, value=choices[0]) + + except Exception as e: + print(f"Error loading cached data: {e}") + + # Check for cached 3DGS video + gs_video_dir = os.path.join(target_dir, "gs_video") + if os.path.exists(gs_video_dir): + try: + from glob import glob + + gs_videos = sorted(glob(os.path.join(gs_video_dir, "*.mp4"))) + if gs_videos: + gs_video_path = gs_videos[-1] + gs_video_visible = True + gs_info_visible = False + print(f"Loaded cached 3DGS video: {gs_video_path}") + except Exception as e: + print(f"Error loading cached 3DGS video: {e}") + + return ( + reconstruction_output, + target_dir, + image_paths, + log_message, + processed_data, + measure_view_selector, + gs_video_path, + gr.update(visible=gs_video_visible), + gr.update(visible=gs_info_visible), + ) + + def navigate_depth_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector: str, + direction: int, + ) -> Tuple[str, Optional[str]]: + """ + Navigate depth view. + + Args: + processed_data: Processed data dictionary + current_selector: Current selector value + direction: Direction to navigate + + Returns: + Tuple of (new_selector_value, depth_vis) + """ + return self.visualization_handler.navigate_depth_view( + processed_data, current_selector, direction + ) + + def update_depth_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Optional[str]: + """ + Update depth view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Path to depth visualization image or None + """ + return self.visualization_handler.update_depth_view(processed_data, view_index) + + def navigate_measure_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector: str, + direction: int, + ) -> Tuple[str, Optional[np.ndarray], Optional[np.ndarray], List]: + """ + Navigate measure view. + + Args: + processed_data: Processed data dictionary + current_selector: Current selector value + direction: Direction to navigate + + Returns: + Tuple of (new_selector_value, measure_image, depth_right_half, measure_points) + """ + return self.visualization_handler.navigate_measure_view( + processed_data, current_selector, direction + ) + + def update_measure_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List]: + """ + Update measure view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Tuple of (measure_image, depth_right_half, measure_points) + """ + return self.visualization_handler.update_measure_view(processed_data, view_index) + + def measure( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + measure_points: List, + current_view_selector: str, + event: gr.SelectData, + ) -> List: + """ + Handle measurement on images. + + Args: + processed_data: Processed data dictionary + measure_points: List of current measure points + current_view_selector: Current view selector value + event: Gradio select event + + Returns: + List of [image, depth_right_half, measure_points, text] + """ + return self.visualization_handler.measure( + processed_data, measure_points, current_view_selector, event + ) + + def select_first_frame( + self, image_gallery: List, selected_index: int = 0 + ) -> Tuple[List, str, str]: + """ + Select the first frame from the image gallery. + + Args: + image_gallery: List of images in the gallery + selected_index: Index of the selected image (default: 0) + + Returns: + Tuple of (updated_image_gallery, log_message, selected_frame_path) + """ + try: + if not image_gallery or len(image_gallery) == 0: + return image_gallery, "No images available to select as first frame.", "" + + # Handle None or invalid selected_index + if ( + selected_index is None + or selected_index < 0 + or selected_index >= len(image_gallery) + ): + selected_index = 0 + print(f"Invalid selected_index: {selected_index}, using default: 0") + + # Get the selected image based on index + selected_image = image_gallery[selected_index] + print(f"Selected image index: {selected_index}") + print(f"Total images: {len(image_gallery)}") + + # Extract the file path from the selected image + selected_frame_path = "" + print(f"Selected image type: {type(selected_image)}") + print(f"Selected image: {selected_image}") + + if isinstance(selected_image, tuple): + # Gradio Gallery returns tuple (path, None) + selected_frame_path = selected_image[0] + elif isinstance(selected_image, str): + selected_frame_path = selected_image + elif hasattr(selected_image, "name"): + selected_frame_path = selected_image.name + elif isinstance(selected_image, dict): + if "name" in selected_image: + selected_frame_path = selected_image["name"] + elif "path" in selected_image: + selected_frame_path = selected_image["path"] + elif "src" in selected_image: + selected_frame_path = selected_image["src"] + else: + # Try to convert to string + selected_frame_path = str(selected_image) + + print(f"Extracted path: {selected_frame_path}") + + # Extract filename from the path for matching + import os + + selected_filename = os.path.basename(selected_frame_path) + print(f"Selected filename: {selected_filename}") + + # Move the selected image to the front + updated_gallery = [selected_image] + [ + img for img in image_gallery if img != selected_image + ] + + log_message = ( + f"Selected frame: {selected_filename}. " + f"Moved to first position. Total frames: {len(updated_gallery)}" + ) + return updated_gallery, log_message, selected_filename + + except Exception as e: + print(f"Error selecting first frame: {e}") + return image_gallery, f"Error selecting first frame: {e}", "" diff --git a/src/nn/depth_anything_3/app/modules/file_handlers.py b/src/nn/depth_anything_3/app/modules/file_handlers.py new file mode 100644 index 0000000..d738bfe --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/file_handlers.py @@ -0,0 +1,304 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +File handling module for Depth Anything 3 Gradio app. + +This module handles file uploads, video processing, and file operations. +""" + +import os +import shutil +import time +from datetime import datetime +from typing import List, Optional, Tuple +import cv2 +from PIL import Image +from pillow_heif import register_heif_opener + +register_heif_opener() + + +class FileHandler: + """ + Handles file uploads and processing for the Gradio app. + """ + + def __init__(self): + """Initialize the file handler.""" + + def handle_uploads( + self, + input_video: Optional[str], + input_images: Optional[List], + s_time_interval: float = 10.0, + ) -> Tuple[str, List[str]]: + """ + Create a new 'target_dir' + 'images' subfolder, and place user-uploaded + images or extracted frames from video into it. + + Args: + input_video: Path to input video file + input_images: List of input image files + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + Tuple of (target_dir, image_paths) + """ + start_time = time.time() + + # Get workspace directory from environment variable or use default + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + if not os.path.exists(workspace_dir): + os.makedirs(workspace_dir) + + # Create input_images subdirectory + input_images_dir = os.path.join(workspace_dir, "input_images") + if not os.path.exists(input_images_dir): + os.makedirs(input_images_dir) + + # Create a unique folder name within input_images + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + target_dir = os.path.join(input_images_dir, f"session_{timestamp}") + target_dir_images = os.path.join(target_dir, "images") + + # Clean up if somehow that folder already exists + if os.path.exists(target_dir): + shutil.rmtree(target_dir) + os.makedirs(target_dir) + os.makedirs(target_dir_images) + + image_paths = [] + + # Handle images + if input_images is not None: + image_paths.extend(self._process_images(input_images, target_dir_images)) + + # Handle video + if input_video is not None: + image_paths.extend( + self._process_video(input_video, target_dir_images, s_time_interval) + ) + + # Sort final images for gallery + image_paths = sorted(image_paths) + + end_time = time.time() + print(f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds") + return target_dir, image_paths + + def _process_images(self, input_images: List, target_dir_images: str) -> List[str]: + """ + Process uploaded images. + + Args: + input_images: List of input image files + target_dir_images: Target directory for images + + Returns: + List of processed image paths + """ + image_paths = [] + + for file_data in input_images: + if isinstance(file_data, dict) and "name" in file_data: + file_path = file_data["name"] + else: + file_path = file_data + + # Check if the file is a HEIC image + file_ext = os.path.splitext(file_path)[1].lower() + if file_ext in [".heic", ".heif"]: + # Convert HEIC to JPEG for better gallery compatibility + try: + with Image.open(file_path) as img: + # Convert to RGB if necessary (HEIC can have different color modes) + if img.mode not in ("RGB", "L"): + img = img.convert("RGB") + + # Create JPEG filename + base_name = os.path.splitext(os.path.basename(file_path))[0] + dst_path = os.path.join(target_dir_images, f"{base_name}.jpg") + + # Save as JPEG with high quality + img.save(dst_path, "JPEG", quality=95) + image_paths.append(dst_path) + print( + f"Converted HEIC to JPEG: {os.path.basename(file_path)} -> " + f"{os.path.basename(dst_path)}" + ) + except Exception as e: + print(f"Error converting HEIC file {file_path}: {e}") + # Fall back to copying as is + dst_path = os.path.join(target_dir_images, os.path.basename(file_path)) + shutil.copy(file_path, dst_path) + image_paths.append(dst_path) + else: + # Regular image files - copy as is + dst_path = os.path.join(target_dir_images, os.path.basename(file_path)) + shutil.copy(file_path, dst_path) + image_paths.append(dst_path) + + return image_paths + + def _process_video( + self, input_video: str, target_dir_images: str, s_time_interval: float + ) -> List[str]: + """ + Process video file and extract frames. + + Args: + input_video: Path to input video file + target_dir_images: Target directory for extracted frames + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + List of extracted frame paths + """ + image_paths = [] + + if isinstance(input_video, dict) and "name" in input_video: + video_path = input_video["name"] + else: + video_path = input_video + + vs = cv2.VideoCapture(video_path) + fps = vs.get(cv2.CAP_PROP_FPS) + frame_interval = max(1, int(fps / s_time_interval)) # Convert FPS to frame interval + + count = 0 + video_frame_num = 0 + while True: + gotit, frame = vs.read() + if not gotit: + break + count += 1 + if count % frame_interval == 0: + image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png") + cv2.imwrite(image_path, frame) + image_paths.append(image_path) + video_frame_num += 1 + + return image_paths + + def update_gallery_on_upload( + self, + input_video: Optional[str], + input_images: Optional[List], + s_time_interval: float = 10.0, + ) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[str]]: + """ + Handle file uploads and update gallery. + + Args: + input_video: Path to input video file + input_images: List of input image files + s_time_interval: Sampling FPS (frames per second) for frame extraction + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message) + """ + if not input_video and not input_images: + return None, None, None, None + + target_dir, image_paths = self.handle_uploads(input_video, input_images, s_time_interval) + return ( + None, + target_dir, + image_paths, + "Upload complete. Click 'Reconstruct' to begin 3D processing.", + ) + + def load_example_scene( + self, scene_name: str, examples_dir: str = "examples" + ) -> Tuple[Optional[str], Optional[str], Optional[List], str]: + """ + Load a scene from examples directory. + + Args: + scene_name: Name of the scene to load + examples_dir: Path to examples directory + + Returns: + Tuple of (reconstruction_output, target_dir, image_paths, log_message) + """ + from depth_anything_3.app.modules.utils import get_scene_info + + scenes = get_scene_info(examples_dir) + + # Find the selected scene + selected_scene = None + for scene in scenes: + if scene["name"] == scene_name: + selected_scene = scene + break + + if selected_scene is None: + return None, None, None, "Scene not found" + + # Use fixed directory name for examples (not timestamp-based) + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + input_images_dir = os.path.join(workspace_dir, "input_images") + if not os.path.exists(input_images_dir): + os.makedirs(input_images_dir) + + # Create a fixed folder name based on scene name + target_dir = os.path.join(input_images_dir, f"example_{scene_name}") + target_dir_images = os.path.join(target_dir, "images") + + # Check if already cached (GLB file exists) + glb_path = os.path.join(target_dir, "scene.glb") + is_cached = os.path.exists(glb_path) + + # Create directory if it doesn't exist + if not os.path.exists(target_dir): + os.makedirs(target_dir) + os.makedirs(target_dir_images) + + # Copy images if directory is new or empty + if not os.path.exists(target_dir_images) or len(os.listdir(target_dir_images)) == 0: + os.makedirs(target_dir_images, exist_ok=True) + image_paths = [] + for file_path in selected_scene["image_files"]: + dst_path = os.path.join(target_dir_images, os.path.basename(file_path)) + shutil.copy(file_path, dst_path) + image_paths.append(dst_path) + else: + # Use existing images + image_paths = sorted( + [ + os.path.join(target_dir_images, f) + for f in os.listdir(target_dir_images) + if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif")) + ] + ) + + # Return cached GLB if available + if is_cached: + return ( + glb_path, # Return cached reconstruction + target_dir, # Set target directory + image_paths, # Set gallery + f"Loaded cached scene '{scene_name}' with {selected_scene['num_images']} images.", + ) + else: + return ( + None, # No cached reconstruction + target_dir, # Set target directory + image_paths, # Set gallery + ( + f"Loaded scene '{scene_name}' with {selected_scene['num_images']} images. " + "Click 'Reconstruct' to begin 3D processing." + ), + ) diff --git a/src/nn/depth_anything_3/app/modules/model_inference.py b/src/nn/depth_anything_3/app/modules/model_inference.py new file mode 100644 index 0000000..6ce77b4 --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/model_inference.py @@ -0,0 +1,260 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model inference module for Depth Anything 3 Gradio app. + +This module handles all model-related operations including inference, +data processing, and result preparation. +""" + +import glob +import os +from typing import Any, Dict, Optional, Tuple +import numpy as np +import torch + +from depth_anything_3.api import DepthAnything3 +from depth_anything_3.utils.memory import cleanup_cuda_memory +from depth_anything_3.utils.export.glb import export_to_glb +from depth_anything_3.utils.export.gs import export_to_gs_video + + +class ModelInference: + """ + Handles model inference and data processing for Depth Anything 3. + """ + + def __init__(self): + """Initialize the model inference handler.""" + self.model = None + + def initialize_model(self, device: str = "cuda") -> None: + """ + Initialize the DepthAnything3 model. + + Args: + device: Device to load the model on + """ + if self.model is None: + # Get model directory from environment variable or use default + model_dir = os.environ.get( + "DA3_MODEL_DIR", "/dev/shm/da3_models/DA3HF-VITG-METRIC_VITL" + ) + self.model = DepthAnything3.from_pretrained(model_dir) + self.model = self.model.to(device) + else: + self.model = self.model.to(device) + + self.model.eval() + + def run_inference( + self, + target_dir: str, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + process_res_method: str = "upper_bound_resize", + show_camera: bool = True, + save_percentage: float = 30.0, + num_max_points: int = 1_000_000, + infer_gs: bool = False, + ref_view_strategy: str = "saddle_balanced", + gs_trj_mode: str = "extend", + gs_video_quality: str = "high", + ) -> Tuple[Any, Dict[int, Dict[str, Any]]]: + """ + Run DepthAnything3 model inference on images. + + Args: + target_dir: Directory containing images + filter_black_bg: Whether to filter black background + filter_white_bg: Whether to filter white background + process_res_method: Method for resizing input images + show_camera: Whether to show camera in 3D view + save_percentage: Percentage of points to save (0-100) + num_max_points: Maximum number of points in point cloud + infer_gs: Whether to infer 3D Gaussian Splatting + ref_view_strategy: Reference view selection strategy + gs_trj_mode: Trajectory mode for 3DGS + gs_video_quality: Video quality for 3DGS + + Returns: + Tuple of (prediction, processed_data) + """ + print(f"Processing images from {target_dir}") + + # Device check + device = "cuda" if torch.cuda.is_available() else "cpu" + device = torch.device(device) + + # Initialize model if needed + self.initialize_model(device) + + # Get image paths + print("Loading images...") + image_folder_path = os.path.join(target_dir, "images") + all_image_paths = sorted(glob.glob(os.path.join(image_folder_path, "*"))) + + # Filter for image files + image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"] + all_image_paths = [ + path + for path in all_image_paths + if any(path.lower().endswith(ext) for ext in image_extensions) + ] + + print(f"Found {len(all_image_paths)} images") + print(f"All image paths: {all_image_paths}") + + # Use sorted image order (reference view will be selected automatically) + image_paths = all_image_paths + print(f"Reference view selection strategy: {ref_view_strategy}") + + if len(image_paths) == 0: + raise ValueError("No images found. Check your upload.") + + # Map UI options to actual method names + method_mapping = {"high_res": "lower_bound_resize", "low_res": "upper_bound_resize"} + actual_method = method_mapping.get(process_res_method, "upper_bound_crop") + + # Run model inference + print(f"Running inference with method: {actual_method}") + with torch.no_grad(): + prediction = self.model.inference( + image_paths, + export_dir=None, + process_res_method=actual_method, + infer_gs=infer_gs, + ref_view_strategy=ref_view_strategy, + ) + # num_max_points: int = 1_000_000, + export_to_glb( + prediction, + filter_black_bg=filter_black_bg, + filter_white_bg=filter_white_bg, + export_dir=target_dir, + show_cameras=show_camera, + conf_thresh_percentile=save_percentage, + num_max_points=int(num_max_points), + ) + + # export to gs video if needed + if infer_gs: + mode_mapping = {"extend": "extend", "smooth": "interpolate_smooth"} + print(f"GS mode: {gs_trj_mode}; Backend mode: {mode_mapping[gs_trj_mode]}") + export_to_gs_video( + prediction, + export_dir=target_dir, + chunk_size=4, + trj_mode=mode_mapping.get(gs_trj_mode, "extend"), + enable_tqdm=True, + vis_depth="hcat", + video_quality=gs_video_quality, + ) + + # Save predictions.npz for caching metric depth data + self._save_predictions_cache(target_dir, prediction) + + # Process results + processed_data = self._process_results(target_dir, prediction, image_paths) + + # Clean up using centralized memory utilities for consistency with backend + cleanup_cuda_memory() + + return prediction, processed_data + + def _save_predictions_cache(self, target_dir: str, prediction: Any) -> None: + """ + Save predictions data to predictions.npz for caching. + + Args: + target_dir: Directory to save the cache + prediction: Model prediction object + """ + try: + output_file = os.path.join(target_dir, "predictions.npz") + + # Build save dict with prediction data + save_dict = {} + + # Save processed images if available + if prediction.processed_images is not None: + save_dict["images"] = prediction.processed_images + + # Save depth data + if prediction.depth is not None: + save_dict["depths"] = np.round(prediction.depth, 6) + + # Save confidence if available + if prediction.conf is not None: + save_dict["conf"] = np.round(prediction.conf, 2) + + # Save camera parameters + if prediction.extrinsics is not None: + save_dict["extrinsics"] = prediction.extrinsics + if prediction.intrinsics is not None: + save_dict["intrinsics"] = prediction.intrinsics + + # Save to file + np.savez_compressed(output_file, **save_dict) + print(f"Saved predictions cache to: {output_file}") + + except Exception as e: + print(f"Warning: Failed to save predictions cache: {e}") + + def _process_results( + self, target_dir: str, prediction: Any, image_paths: list + ) -> Dict[int, Dict[str, Any]]: + """ + Process model results into structured data. + + Args: + target_dir: Directory containing results + prediction: Model prediction object + image_paths: List of input image paths + + Returns: + Dictionary containing processed data for each view + """ + processed_data = {} + + # Read generated depth visualization files + depth_vis_dir = os.path.join(target_dir, "depth_vis") + + if os.path.exists(depth_vis_dir): + depth_files = sorted(glob.glob(os.path.join(depth_vis_dir, "*.jpg"))) + for i, depth_file in enumerate(depth_files): + # Use processed images directly from API + processed_image = None + if prediction.processed_images is not None and i < len( + prediction.processed_images + ): + processed_image = prediction.processed_images[i] + + processed_data[i] = { + "depth_image": depth_file, + "image": processed_image, + "original_image_path": image_paths[i] if i < len(image_paths) else None, + "depth": prediction.depth[i] if i < len(prediction.depth) else None, + "intrinsics": ( + prediction.intrinsics[i] + if prediction.intrinsics is not None and i < len(prediction.intrinsics) + else None + ), + "mask": None, # No mask information available + } + + return processed_data + + # cleanup() removed: call cleanup_cuda_memory() directly where needed. diff --git a/src/nn/depth_anything_3/app/modules/ui_components.py b/src/nn/depth_anything_3/app/modules/ui_components.py new file mode 100644 index 0000000..3a3762f --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/ui_components.py @@ -0,0 +1,477 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +UI components module for Depth Anything 3 Gradio app. + +This module contains UI component definitions and layout functions. +""" + +import os +from typing import Any, Dict, List, Tuple +import gradio as gr + +from depth_anything_3.app.modules.utils import get_logo_base64, get_scene_info + + +class UIComponents: + """ + Handles UI component creation and layout for the Gradio app. + """ + + def __init__(self): + """Initialize the UI components handler.""" + + def create_upload_section(self) -> Tuple[gr.Video, gr.Slider, gr.File, gr.Gallery]: + """ + Create the upload section with video, images, and gallery components. + + Returns: + A tuple of Gradio components: (input_video, s_time_interval, input_images, image_gallery). + """ + input_video = gr.Video(label="Upload Video", interactive=True) + s_time_interval = gr.Slider( + minimum=0.1, + maximum=60, + value=10, + step=0.1, + label="Sampling FPS (Frames Per Second)", + interactive=True, + visible=True, + ) + input_images = gr.File(file_count="multiple", label="Upload Images", interactive=True) + image_gallery = gr.Gallery( + label="Preview", + columns=4, + height="300px", + show_download_button=True, + object_fit="contain", + preview=True, + interactive=False, + ) + + return input_video, s_time_interval, input_images, image_gallery + + def create_3d_viewer_section(self) -> gr.Model3D: + """ + Create the 3D viewer component. + + Returns: + 3D model viewer component + """ + return gr.Model3D( + height=520, + zoom_speed=0.5, + pan_speed=0.5, + clear_color=[0.0, 0.0, 0.0, 0.0], + key="persistent_3d_viewer", + elem_id="reconstruction_3d_viewer", + ) + + def create_nvs_video(self) -> Tuple[gr.Video, gr.Markdown]: + """ + Create the 3DGS rendered video display component and info message. + + Returns: + Tuple of (video component, info message component) + """ + with gr.Column(): + gs_info = gr.Markdown( + ( + "‼️ **3D Gaussian Splatting rendering is currently DISABLED.**


" + "To render novel views from 3DGS, " + "enable **Infer 3D Gaussian Splatting** below.
" + "Next, in **Visualization Options**, " + "*optionally* configure the **rendering trajectory** (default: smooth) " + "and **video quality** (default: low), " + "then click **Reconstruct**." + ), + visible=True, + height=520, + ) + gs_video = gr.Video( + height=520, + label="3DGS Rendered NVS Video (depth shown for reference only)", + interactive=False, + visible=False, + ) + return gs_video, gs_info + + def create_depth_section(self) -> Tuple[gr.Button, gr.Dropdown, gr.Button, gr.Image]: + """ + Create the depth visualization section. + + Returns: + A tuple of (prev_depth_btn, depth_view_selector, next_depth_btn, depth_map) + """ + with gr.Row(elem_classes=["navigation-row"]): + prev_depth_btn = gr.Button("◀ Previous", size="sm", scale=1) + depth_view_selector = gr.Dropdown( + choices=["View 1"], + value="View 1", + label="Select View", + scale=2, + interactive=True, + allow_custom_value=True, + ) + next_depth_btn = gr.Button("Next ▶", size="sm", scale=1) + depth_map = gr.Image( + type="numpy", + label="Colorized Depth Map", + format="png", + interactive=False, + ) + + return prev_depth_btn, depth_view_selector, next_depth_btn, depth_map + + def create_measure_section( + self, + ) -> Tuple[gr.Button, gr.Dropdown, gr.Button, gr.Image, gr.Image, gr.Markdown]: + """ + Create the measurement section. + + Returns: + A tuple of (prev_measure_btn, measure_view_selector, next_measure_btn, measure_image, + measure_depth_image, measure_text) + """ + from depth_anything_3.app.css_and_html import MEASURE_INSTRUCTIONS_HTML + + gr.Markdown(MEASURE_INSTRUCTIONS_HTML) + with gr.Row(elem_classes=["navigation-row"]): + prev_measure_btn = gr.Button("◀ Previous", size="sm", scale=1) + measure_view_selector = gr.Dropdown( + choices=["View 1"], + value="View 1", + label="Select View", + scale=2, + interactive=True, + allow_custom_value=True, + ) + next_measure_btn = gr.Button("Next ▶", size="sm", scale=1) + with gr.Row(): + measure_image = gr.Image( + type="numpy", + show_label=False, + format="webp", + interactive=False, + sources=[], + label="RGB Image", + scale=1, + height=275, + ) + measure_depth_image = gr.Image( + type="numpy", + show_label=False, + format="webp", + interactive=False, + sources=[], + label="Depth Visualization (Right Half)", + scale=1, + height=275, + ) + gr.Markdown( + "**Note:** Images have been adjusted to model processing size. " + "Click two points on the RGB image to measure distance." + ) + measure_text = gr.Markdown("") + + return ( + prev_measure_btn, + measure_view_selector, + next_measure_btn, + measure_image, + measure_depth_image, + measure_text, + ) + + def create_inference_control_section(self) -> Tuple[gr.Dropdown, gr.Checkbox, gr.Dropdown]: + """ + Create the inference control section (before inference). + + Returns: + Tuple of (process_res_method_dropdown, infer_gs, ref_view_strategy) + """ + with gr.Row(): + process_res_method_dropdown = gr.Dropdown( + choices=["high_res", "low_res"], + value="low_res", + label="Image Processing Method", + info="low_res for much more images", + scale=1, + ) + # Modify line 220, add color class + infer_gs = gr.Checkbox( + label="Infer 3D Gaussian Splatting", + value=False, + info=( + 'Enable novel view rendering from 3DGS ( requires extra processing time)' + ), + scale=1, + ) + ref_view_strategy = gr.Dropdown( + choices=["saddle_balanced", "saddle_sim_range", "first", "middle"], + value="saddle_balanced", + label="Reference View Strategy", + info="Strategy for selecting reference view from multiple inputs", + scale=1, + ) + + return (process_res_method_dropdown, infer_gs, ref_view_strategy) + + def create_display_control_section( + self, + ) -> Tuple[ + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Slider, + gr.Slider, + gr.Dropdown, + gr.Dropdown, + gr.Button, + gr.ClearButton, + ]: + """ + Create the display control section (options for visualization). + + Returns: + Tuple of display control components including buttons + """ + with gr.Column(): + # 3DGS options at the top + with gr.Row(): + gs_trj_mode = gr.Dropdown( + choices=["smooth", "extend"], + value="smooth", + label=("Rendering trajectory for 3DGS viewpoints (requires n_views ≥ 2)"), + info=("'smooth' for view interpolation; 'extend' for longer trajectory"), + visible=False, # initially hidden + ) + gs_video_quality = gr.Dropdown( + choices=["low", "medium", "high"], + value="low", + label=("Video quality for 3DGS rendered outputs"), + info=("'low' for faster loading speed; 'high' for better visual quality"), + visible=False, # initially hidden + ) + + # Reconstruct and Clear buttons (before Visualization Options) + with gr.Row(): + submit_btn = gr.Button("Reconstruct", scale=1, variant="primary") + clear_btn = gr.ClearButton(scale=1) + + gr.Markdown("### Visualization Options: (Click Reconstruct to update)") + show_cam = gr.Checkbox(label="Show Camera", value=True) + filter_black_bg = gr.Checkbox(label="Filter Black Background", value=False) + filter_white_bg = gr.Checkbox(label="Filter White Background", value=False) + save_percentage = gr.Slider( + minimum=0, + maximum=100, + value=10, + step=1, + label="Filter Percentage", + info="Confidence Threshold (%): Higher values filter more points.", + ) + num_max_points = gr.Slider( + minimum=1000, + maximum=100000, + value=1000, + step=1000, + label="Max Points (K points)", + info="Maximum number of points to export to GLB (in thousands)", + ) + + return ( + show_cam, + filter_black_bg, + filter_white_bg, + save_percentage, + num_max_points, + gs_trj_mode, + gs_video_quality, + submit_btn, + clear_btn, + ) + + def create_control_section( + self, + ) -> Tuple[ + gr.Button, + gr.ClearButton, + gr.Dropdown, + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Checkbox, + gr.Dropdown, + gr.Checkbox, + gr.Textbox, + ]: + """ + Create the control section with buttons and options. + + Returns: + Tuple of control components + """ + with gr.Row(): + submit_btn = gr.Button("Reconstruct", scale=1, variant="primary") + clear_btn = gr.ClearButton( + scale=1, + ) + + with gr.Row(): + frame_filter = gr.Dropdown( + choices=["All"], value="All", label="Show Points from Frame" + ) + with gr.Column(): + gr.Markdown("### Visualization Option: (Click Reconstruct to update)") + show_cam = gr.Checkbox(label="Show Camera", value=True) + show_mesh = gr.Checkbox(label="Show Mesh", value=True) + filter_black_bg = gr.Checkbox(label="Filter Black Background", value=False) + filter_white_bg = gr.Checkbox(label="Filter White Background", value=False) + gr.Markdown("### Reconstruction Options: (updated on next run)") + apply_mask_checkbox = gr.Checkbox( + label="Apply mask for predicted ambiguous depth classes & edges", + value=True, + ) + process_res_method_dropdown = gr.Dropdown( + choices=[ + "upper_bound_resize", + "upper_bound_crop", + "lower_bound_resize", + "lower_bound_crop", + ], + value="upper_bound_resize", + label="Image Processing Method", + info="Method for resizing input images", + ) + save_to_gallery_checkbox = gr.Checkbox( + label="Save to Gallery", + value=False, + info="Save current reconstruction results to gallery directory", + ) + gallery_name_input = gr.Textbox( + label="Gallery Name", + placeholder="Enter a name for the gallery folder", + value="", + info="Leave empty for auto-generated name with timestamp", + ) + + return ( + submit_btn, + clear_btn, + frame_filter, + show_cam, + show_mesh, + filter_black_bg, + filter_white_bg, + apply_mask_checkbox, + process_res_method_dropdown, + save_to_gallery_checkbox, + gallery_name_input, + ) + + def create_example_scenes_section(self) -> List[Dict[str, Any]]: + """ + Create the example scenes section. + + Returns: + List of scene information dictionaries + """ + # Get workspace directory from environment variable + workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace") + examples_dir = os.path.join(workspace_dir, "examples") + + # Get scene information + scenes = get_scene_info(examples_dir) + + return scenes + + def create_example_scene_grid(self, scenes: List[Dict[str, Any]]) -> List[gr.Image]: + """ + Create the example scene grid. + + Args: + scenes: List of scene information dictionaries + + Returns: + List of scene image components + """ + scene_components = [] + + if scenes: + for i in range(0, len(scenes), 4): # Process 4 scenes per row + with gr.Row(): + for j in range(4): + scene_idx = i + j + if scene_idx < len(scenes): + scene = scenes[scene_idx] + with gr.Column(scale=1, elem_classes=["clickable-thumbnail"]): + # Clickable thumbnail + scene_img = gr.Image( + value=scene["thumbnail"], + height=150, + interactive=False, + show_label=False, + elem_id=f"scene_thumb_{scene['name']}", + sources=[], + ) + scene_components.append(scene_img) + + # Scene name and image count as text below thumbnail + gr.Markdown( + f"**{scene['name']}** \n {scene['num_images']} images", + elem_classes=["scene-info"], + ) + else: + # Empty column to maintain grid structure + with gr.Column(scale=1): + pass + + return scene_components + + def create_header_section(self) -> gr.HTML: + """ + Create the header section with logo and title. + + Returns: + Header HTML component + """ + from depth_anything_3.app.css_and_html import get_header_html + + return gr.HTML(get_header_html(get_logo_base64())) + + def create_description_section(self) -> gr.HTML: + """ + Create the description section. + + Returns: + Description HTML component + """ + from depth_anything_3.app.css_and_html import get_description_html + + return gr.HTML(get_description_html()) + + def create_acknowledgements_section(self) -> gr.HTML: + """ + Create the acknowledgements section. + + Returns: + Acknowledgements HTML component + """ + from depth_anything_3.app.css_and_html import get_acknowledgements_html + + return gr.HTML(get_acknowledgements_html()) diff --git a/src/nn/depth_anything_3/app/modules/utils.py b/src/nn/depth_anything_3/app/modules/utils.py new file mode 100644 index 0000000..985abc1 --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/utils.py @@ -0,0 +1,207 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utility functions for Depth Anything 3 Gradio app. + +This module contains helper functions for data processing, visualization, +and file operations. +""" + + +import json +import os +import shutil +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple +import numpy as np + +def create_depth_visualization(depth: np.ndarray) -> Optional[np.ndarray]: + """ + Create a colored depth visualization. + + Args: + depth: Depth array + + Returns: + Colored depth visualization or None + """ + if depth is None: + return None + + # Normalize depth to 0-1 range + depth_min = depth[depth > 0].min() if (depth > 0).any() else 0 + depth_max = depth.max() + + if depth_max <= depth_min: + return None + + # Normalize depth + depth_norm = (depth - depth_min) / (depth_max - depth_min) + depth_norm = np.clip(depth_norm, 0, 1) + + # Apply colormap (using matplotlib's viridis colormap) + import matplotlib.cm as cm + + # Convert to colored image + depth_colored = cm.viridis(depth_norm)[:, :, :3] # Remove alpha channel + depth_colored = (depth_colored * 255).astype(np.uint8) + + return depth_colored + + +def save_to_gallery_func( + target_dir: str, processed_data: Dict[int, Dict[str, Any]], gallery_name: Optional[str] = None +) -> Tuple[bool, str]: + """ + Save the current reconstruction results to the gallery directory. + + Args: + target_dir: Source directory containing reconstruction results + processed_data: Processed data dictionary + gallery_name: Name for the gallery folder + + Returns: + Tuple of (success, message) + """ + try: + # Get gallery directory from environment variable or use default + gallery_dir = os.environ.get( + "DA3_GALLERY_DIR", + "workspace/gallery", + ) + if not os.path.exists(gallery_dir): + os.makedirs(gallery_dir) + + # Use provided name or create a unique name + if gallery_name is None or gallery_name.strip() == "": + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + gallery_name = f"reconstruction_{timestamp}" + + gallery_path = os.path.join(gallery_dir, gallery_name) + + # Check if directory already exists + if os.path.exists(gallery_path): + return False, f"Save failed: folder '{gallery_name}' already exists" + + # Create the gallery directory + os.makedirs(gallery_path, exist_ok=True) + + # Copy GLB file + glb_source = os.path.join(target_dir, "scene.glb") + glb_dest = os.path.join(gallery_path, "scene.glb") + if os.path.exists(glb_source): + shutil.copy2(glb_source, glb_dest) + + # Copy depth visualization images + depth_vis_dir = os.path.join(target_dir, "depth_vis") + if os.path.exists(depth_vis_dir): + gallery_depth_vis = os.path.join(gallery_path, "depth_vis") + shutil.copytree(depth_vis_dir, gallery_depth_vis) + + # Copy original images + images_source = os.path.join(target_dir, "images") + if os.path.exists(images_source): + gallery_images = os.path.join(gallery_path, "images") + shutil.copytree(images_source, gallery_images) + + scene_preview_source = os.path.join(target_dir, "scene.jpg") + scene_preview_dest = os.path.join(gallery_path, "scene.jpg") + shutil.copy2(scene_preview_source, scene_preview_dest) + + # Save metadata + metadata = { + "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"), + "num_images": len(processed_data) if processed_data else 0, + "gallery_name": gallery_name, + } + + with open(os.path.join(gallery_path, "metadata.json"), "w") as f: + json.dump(metadata, f, indent=2) + + print(f"Saved reconstruction to gallery: {gallery_path}") + return True, f"Save successful: saved to {gallery_path}" + + except Exception as e: + print(f"Error saving to gallery: {e}") + return False, f"Save failed: {str(e)}" + + +def get_scene_info(examples_dir: str) -> List[Dict[str, Any]]: + """ + Get information about scenes in the examples directory. + + Args: + examples_dir: Path to examples directory + + Returns: + List of scene information dictionaries + """ + import glob + + scenes = [] + if not os.path.exists(examples_dir): + return scenes + + for scene_folder in sorted(os.listdir(examples_dir)): + scene_path = os.path.join(examples_dir, scene_folder) + if os.path.isdir(scene_path): + # Find all image files in the scene folder + image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.tiff", "*.tif"] + image_files = [] + for ext in image_extensions: + image_files.extend(glob.glob(os.path.join(scene_path, ext))) + image_files.extend(glob.glob(os.path.join(scene_path, ext.upper()))) + + if image_files: + # Sort images and get the first one for thumbnail + image_files = sorted(image_files) + first_image = image_files[0] + num_images = len(image_files) + + scenes.append( + { + "name": scene_folder, + "path": scene_path, + "thumbnail": first_image, + "num_images": num_images, + "image_files": image_files, + } + ) + + return scenes + + +# NOTE: cleanup was moved to a single canonical helper in +# `depth_anything_3.utils.memory.cleanup_cuda_memory`. +# Callers should import and call that directly instead of using this module. + + +def get_logo_base64() -> Optional[str]: + """ + Convert WAI logo to base64 for embedding in HTML. + + Returns: + Base64 encoded logo string or None + """ + import base64 + + logo_path = "examples/WAI-Logo/wai_logo.png" + try: + with open(logo_path, "rb") as img_file: + img_data = img_file.read() + base64_str = base64.b64encode(img_data).decode() + return f"data:image/png;base64,{base64_str}" + except FileNotFoundError: + return None diff --git a/src/nn/depth_anything_3/app/modules/visualization.py b/src/nn/depth_anything_3/app/modules/visualization.py new file mode 100644 index 0000000..ada49b2 --- /dev/null +++ b/src/nn/depth_anything_3/app/modules/visualization.py @@ -0,0 +1,434 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Visualization module for Depth Anything 3 Gradio app. + +This module handles visualization updates, navigation, and measurement functionality. +""" + +import os +from typing import Any, Dict, List, Optional, Tuple +import cv2 +import gradio as gr +import numpy as np + + +class VisualizationHandler: + """ + Handles visualization updates and navigation for the Gradio app. + """ + + def __init__(self): + """Initialize the visualization handler.""" + + def update_view_selectors( + self, processed_data: Optional[Dict[int, Dict[str, Any]]] + ) -> Tuple[gr.Dropdown, gr.Dropdown]: + """ + Update view selector dropdowns based on available views. + + Args: + processed_data: Processed data dictionary + + Returns: + Tuple of (depth_view_selector, measure_view_selector) + """ + if processed_data is None or len(processed_data) == 0: + choices = ["View 1"] + else: + num_views = len(processed_data) + choices = [f"View {i + 1}" for i in range(num_views)] + + return ( + gr.Dropdown(choices=choices, value=choices[0]), # depth_view_selector + gr.Dropdown(choices=choices, value=choices[0]), # measure_view_selector + ) + + def get_view_data_by_index( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Optional[Dict[str, Any]]: + """ + Get view data by index, handling bounds. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to get + + Returns: + View data dictionary or None + """ + if processed_data is None or len(processed_data) == 0: + return None + + view_keys = list(processed_data.keys()) + if view_index < 0 or view_index >= len(view_keys): + view_index = 0 + + return processed_data[view_keys[view_index]] + + def update_depth_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Optional[str]: + """ + Update depth view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Path to depth visualization image or None + """ + view_data = self.get_view_data_by_index(processed_data, view_index) + if view_data is None or view_data.get("depth_image") is None: + return None + + # Return the depth visualization image directly + return view_data["depth_image"] + + def navigate_depth_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector_value: str, + direction: int, + ) -> Tuple[str, Optional[str]]: + """ + Navigate depth view (direction: -1 for previous, +1 for next). + + Args: + processed_data: Processed data dictionary + current_selector_value: Current selector value + direction: Direction to navigate (-1 for previous, +1 for next) + + Returns: + Tuple of (new_selector_value, depth_vis) + """ + if processed_data is None or len(processed_data) == 0: + return "View 1", None + + # Parse current view number + try: + current_view = int(current_selector_value.split()[1]) - 1 + except: # noqa + current_view = 0 + + num_views = len(processed_data) + new_view = (current_view + direction) % num_views + + new_selector_value = f"View {new_view + 1}" + depth_vis = self.update_depth_view(processed_data, new_view) + + return new_selector_value, depth_vis + + def update_measure_view( + self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int + ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List]: + """ + Update measure view for a specific view index. + + Args: + processed_data: Processed data dictionary + view_index: Index of the view to update + + Returns: + Tuple of (measure_image, depth_right_half, measure_points) + """ + view_data = self.get_view_data_by_index(processed_data, view_index) + if view_data is None: + return None, None, [] # image, depth_right_half, measure_points + + # Get the processed (resized) image + if "image" in view_data and view_data["image"] is not None: + image = view_data["image"].copy() + else: + return None, None, [] + + # Ensure image is in uint8 format + if image.dtype != np.uint8: + if image.max() <= 1.0: + image = (image * 255).astype(np.uint8) + else: + image = image.astype(np.uint8) + + # Extract right half of the depth visualization (pure depth part) + depth_image_path = view_data.get("depth_image", None) + depth_right_half = None + + if depth_image_path and os.path.exists(depth_image_path): + try: + # Load the combined depth visualization image + depth_combined = cv2.imread(depth_image_path) + depth_combined = cv2.cvtColor(depth_combined, cv2.COLOR_BGR2RGB) + if depth_combined is not None: + height, width = depth_combined.shape[:2] + # Extract right half (depth visualization part) + depth_right_half = depth_combined[:, width // 2 :] + except Exception as e: + print(f"Error extracting depth right half: {e}") + + return image, depth_right_half, [] + + def navigate_measure_view( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + current_selector_value: str, + direction: int, + ) -> Tuple[str, Optional[np.ndarray], Optional[str], List]: + """ + Navigate measure view (direction: -1 for previous, +1 for next). + + Args: + processed_data: Processed data dictionary + current_selector_value: Current selector value + direction: Direction to navigate (-1 for previous, +1 for next) + + Returns: + Tuple of (new_selector_value, measure_image, depth_image_path, measure_points) + """ + if processed_data is None or len(processed_data) == 0: + return "View 1", None, None, [] + + # Parse current view number + try: + current_view = int(current_selector_value.split()[1]) - 1 + except: # noqa + current_view = 0 + + num_views = len(processed_data) + new_view = (current_view + direction) % num_views + + new_selector_value = f"View {new_view + 1}" + measure_image, depth_right_half, measure_points = self.update_measure_view( + processed_data, new_view + ) + + return new_selector_value, measure_image, depth_right_half, measure_points + + def populate_visualization_tabs( + self, processed_data: Optional[Dict[int, Dict[str, Any]]] + ) -> Tuple[Optional[str], Optional[np.ndarray], Optional[str], List]: + """ + Populate the depth and measure tabs with processed data. + + Args: + processed_data: Processed data dictionary + + Returns: + Tuple of (depth_vis, measure_img, depth_image_path, measure_points) + """ + if processed_data is None or len(processed_data) == 0: + return None, None, None, [] + + # Use update function to get depth visualization + depth_vis = self.update_depth_view(processed_data, 0) + measure_img, depth_right_half, _ = self.update_measure_view(processed_data, 0) + + return depth_vis, measure_img, depth_right_half, [] + + def reset_measure( + self, processed_data: Optional[Dict[int, Dict[str, Any]]] + ) -> Tuple[Optional[np.ndarray], List, str]: + """ + Reset measure points. + + Args: + processed_data: Processed data dictionary + + Returns: + Tuple of (image, measure_points, text) + """ + if processed_data is None or len(processed_data) == 0: + return None, [], "" + + # Return the first view image + first_view = list(processed_data.values())[0] + return first_view["image"], [], "" + + def measure( + self, + processed_data: Optional[Dict[int, Dict[str, Any]]], + measure_points: List, + current_view_selector: str, + event: gr.SelectData, + ) -> List: + """ + Handle measurement on images. + + Args: + processed_data: Processed data dictionary + measure_points: List of current measure points + current_view_selector: Current view selector value + event: Gradio select event + + Returns: + List of [image, depth_right_half, measure_points, text] + """ + try: + print(f"Measure function called with selector: {current_view_selector}") + + if processed_data is None or len(processed_data) == 0: + return [None, [], "No data available"] + + # Use the currently selected view instead of always using the first view + try: + current_view_index = int(current_view_selector.split()[1]) - 1 + except: # noqa + current_view_index = 0 + + print(f"Using view index: {current_view_index}") + + # Get view data safely + if current_view_index < 0 or current_view_index >= len(processed_data): + current_view_index = 0 + + view_keys = list(processed_data.keys()) + current_view = processed_data[view_keys[current_view_index]] + + if current_view is None: + return [None, [], "No view data available"] + + point2d = event.index[0], event.index[1] + print(f"Clicked point: {point2d}") + + measure_points.append(point2d) + + # Get image and depth visualization + image, depth_right_half, _ = self.update_measure_view( + processed_data, current_view_index + ) + if image is None: + return [None, [], "No image available"] + + image = image.copy() + + # Ensure image is in uint8 format for proper cv2 operations + try: + if image.dtype != np.uint8: + if image.max() <= 1.0: + # Image is in [0, 1] range, convert to [0, 255] + image = (image * 255).astype(np.uint8) + else: + # Image is already in [0, 255] range + image = image.astype(np.uint8) + except Exception as e: + print(f"Image conversion error: {e}") + return [None, [], f"Image conversion error: {e}"] + + # Draw circles for points + try: + for p in measure_points: + if 0 <= p[0] < image.shape[1] and 0 <= p[1] < image.shape[0]: + image = cv2.circle(image, p, radius=5, color=(255, 0, 0), thickness=2) + except Exception as e: + print(f"Drawing error: {e}") + return [None, [], f"Drawing error: {e}"] + + # Get depth information from processed_data + depth_text = "" + try: + for i, p in enumerate(measure_points): + if ( + current_view["depth"] is not None + and 0 <= p[1] < current_view["depth"].shape[0] + and 0 <= p[0] < current_view["depth"].shape[1] + ): + d = current_view["depth"][p[1], p[0]] + depth_text += f"- **P{i + 1} depth: {d:.2f}m**\n" + else: + depth_text += f"- **P{i + 1}: Click position ({p[0]}, {p[1]}) - No depth information**\n" # noqa: E501 + except Exception as e: + print(f"Depth text error: {e}") + depth_text = f"Error computing depth: {e}\n" + + if len(measure_points) == 2: + try: + point1, point2 = measure_points + # Draw line + if ( + 0 <= point1[0] < image.shape[1] + and 0 <= point1[1] < image.shape[0] + and 0 <= point2[0] < image.shape[1] + and 0 <= point2[1] < image.shape[0] + ): + image = cv2.line(image, point1, point2, color=(255, 0, 0), thickness=2) + + # Compute 3D distance using depth information and camera intrinsics + distance_text = "- **Distance: Unable to calculate 3D distance**" + if ( + current_view["depth"] is not None + and 0 <= point1[1] < current_view["depth"].shape[0] + and 0 <= point1[0] < current_view["depth"].shape[1] + and 0 <= point2[1] < current_view["depth"].shape[0] + and 0 <= point2[0] < current_view["depth"].shape[1] + ): + try: + # Get depth values at the two points + d1 = current_view["depth"][point1[1], point1[0]] + d2 = current_view["depth"][point2[1], point2[0]] + + # Convert 2D pixel coordinates to 3D world coordinates + if current_view["intrinsics"] is not None: + # Get camera intrinsics + K = current_view["intrinsics"] # 3x3 intrinsic matrix + fx, fy = K[0, 0], K[1, 1] # focal lengths + cx, cy = K[0, 2], K[1, 2] # principal point + + # Convert pixel coordinates to normalized camera coordinates + # Point 1: (u1, v1) -> (x1, y1, z1) + u1, v1 = point1[0], point1[1] + x1 = (u1 - cx) * d1 / fx + y1 = (v1 - cy) * d1 / fy + z1 = d1 + + # Point 2: (u2, v2) -> (x2, y2, z2) + u2, v2 = point2[0], point2[1] + x2 = (u2 - cx) * d2 / fx + y2 = (v2 - cy) * d2 / fy + z2 = d2 + + # Calculate 3D Euclidean distance + p1_3d = np.array([x1, y1, z1]) + p2_3d = np.array([x2, y2, z2]) + distance_3d = np.linalg.norm(p1_3d - p2_3d) + + distance_text = f"- **Distance: {distance_3d:.2f}m**" + else: + # Fallback to simplified calculation if no intrinsics + pixel_distance = np.sqrt( + (point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2 + ) + avg_depth = (d1 + d2) / 2 + scale_factor = avg_depth / 1000 # Rough scaling factor + estimated_3d_distance = pixel_distance * scale_factor + distance_text = f"- **Distance: {estimated_3d_distance:.2f}m (estimated, no intrinsics)**" # noqa: E501 + + except Exception as e: + print(f"Distance computation error: {e}") + distance_text = f"- **Distance computation error: {e}**" + + measure_points = [] + text = depth_text + distance_text + print(f"Measurement complete: {text}") + return [image, depth_right_half, measure_points, text] + except Exception as e: + print(f"Final measurement error: {e}") + return [None, [], f"Measurement error: {e}"] + else: + print(f"Single point measurement: {depth_text}") + return [image, depth_right_half, measure_points, depth_text] + + except Exception as e: + print(f"Overall measure function error: {e}") + return [None, [], f"Measure function error: {e}"] diff --git a/src/nn/depth_anything_3/bench/__init__.py b/src/nn/depth_anything_3/bench/__init__.py new file mode 100644 index 0000000..664ee41 --- /dev/null +++ b/src/nn/depth_anything_3/bench/__init__.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Depth Anything 3 Benchmark Evaluation Module. + +This module provides tools for evaluating DepthAnything3 model on various benchmark datasets. +Currently supported datasets: +- DTU (3D Reconstruction) +- DTU-64 (Pose Evaluation Only) +- ETH3D (3D Reconstruction) +- 7Scenes (3D Reconstruction) +- ScanNet++ (3D Reconstruction) +- HiRoom (3D Reconstruction) + +Supported evaluation modes: +- pose: Camera pose estimation evaluation +- recon_unposed: 3D reconstruction with predicted poses +- recon_posed: 3D reconstruction with ground truth poses +""" + +from depth_anything_3.bench.registries import MV_REGISTRY, MONO_REGISTRY + + +def __getattr__(name): + """Lazy import to avoid circular import when running as __main__.""" + if name == "Evaluator": + from depth_anything_3.bench.evaluator import Evaluator + return Evaluator + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = ["Evaluator", "MV_REGISTRY", "MONO_REGISTRY"] + diff --git a/src/nn/depth_anything_3/bench/configs/eval_bench.yaml b/src/nn/depth_anything_3/bench/configs/eval_bench.yaml new file mode 100644 index 0000000..1d9924d --- /dev/null +++ b/src/nn/depth_anything_3/bench/configs/eval_bench.yaml @@ -0,0 +1,98 @@ +# DepthAnything3 Benchmark Evaluation Configuration +# +# This config can be loaded and overridden via command line. +# Example: python -m depth_anything_3.bench.evaluator --model /path/to/model --work_dir /path/to/workspace +# +# See depth_anything_3.cfg for config utility functions. + +# ============================================================================== +# Model Configuration +# ============================================================================== +model: + # Path to model checkpoint or HuggingFace model ID + path: depth-anything/DA3-GIANT + +# ============================================================================== +# Workspace Configuration +# ============================================================================== +workspace: + # Working directory for outputs (model results, metrics, etc.) + work_dir: ./workspace/evaluation + +# ============================================================================== +# Evaluation Configuration +# ============================================================================== +eval: + # Datasets to evaluate + # Options: dtu, dtu64, eth3d, 7scenes (sevenscenes), scannetpp, hiroom + datasets: + - eth3d + - 7scenes + - scannetpp + - hiroom + - dtu + - dtu64 + + # Evaluation modes + # Options: pose, recon_unposed, recon_posed, view_syn + modes: + - pose + - recon_unposed + - recon_posed + + # Reference view selection strategy for inference + # Options: first, saddle_balanced, auto, mid + ref_view_strategy: "first" + + # Specific scenes to evaluate (null = all scenes) + # Example: [courtyard, relief] for eth3d + scenes: null + + # Maximum number of frames per scene (for sampling) + # If a scene has more frames, randomly sample to this limit. + # Set to -1 to disable sampling. + max_frames: 100 + + # Only run evaluation (skip inference) + eval_only: false + + # Only print saved metrics (skip inference and evaluation) + print_only: false + +# ============================================================================== +# Inference Configuration +# ============================================================================== +inference: + # Number of parallel workers for TSDF fusion + num_fusion_workers: 4 + + # Enable debug mode with verbose output + debug: false + +# ============================================================================== +# Preset Configurations +# ============================================================================== +# These can be activated via command line: --preset full_eval + +presets: + # Full evaluation on all 6 datasets + full_eval: + datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu, dtu64] + modes: [pose, recon_unposed, recon_posed] + + # Pose-only evaluation + pose_only: + datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu64] + modes: [pose] + + # Reconstruction-only evaluation (5 datasets, excluding dtu64) + recon_only: + datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu] + modes: [recon_unposed, recon_posed] + + # Quick test (single scene per dataset) + quick_test: + datasets: [eth3d] + modes: [pose, recon_unposed] + scenes: [courtyard] + diff --git a/src/nn/depth_anything_3/bench/dataset.py b/src/nn/depth_anything_3/bench/dataset.py new file mode 100644 index 0000000..2f30a0f --- /dev/null +++ b/src/nn/depth_anything_3/bench/dataset.py @@ -0,0 +1,136 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Base dataset class for benchmark evaluation. + +All dataset implementations should inherit from this class and implement +the required abstract methods. +""" + +import os +import time +from abc import abstractmethod +from typing import Dict as TDict + +import numpy as np +import torch +from addict import Dict + +from depth_anything_3.bench.utils import compute_pose +from depth_anything_3.utils.geometry import as_homogeneous + + +def _wait_for_file_ready(path: str, timeout: float = 3.0, interval: float = 0.2) -> None: + """Wait until file size stabilizes for 2 consecutive checks.""" + last_size = -1 + stable_count = 0 + start = time.time() + while time.time() - start < timeout: + time.sleep(interval) + size = os.path.getsize(path) + if size == last_size and size > 0: + stable_count += 1 + if stable_count >= 2: # Need 2 consecutive stable checks + return + else: + stable_count = 0 + last_size = size + + +class Dataset: + """ + Base class for all benchmark datasets. + + Subclasses must implement: + - SCENES: List of scene identifiers + - data_root: Path to dataset root + - get_data(scene): Return scene data (images, intrinsics, extrinsics, etc.) + - eval3d(scene, fuse_path): Evaluate 3D reconstruction + - fuse3d(scene, result_path, fuse_path, mode): Fuse depth maps into point cloud + + Optional overrides: + - eval_pose(scene, result_path): Evaluate pose estimation (default provided) + """ + + # Subclasses should define these + SCENES: list = [] + data_root: str = "" + + def __init__(self): + pass + + def eval_pose(self, scene: str, result_path: str) -> TDict[str, float]: + """ + Evaluate camera pose estimation accuracy. + + Args: + scene: Scene identifier + result_path: Path to .npz file containing predicted extrinsics + + Returns: + Dict with pose metrics (auc30, auc15, auc05, auc03) + """ + _wait_for_file_ready(result_path) + pred = np.load(result_path) + gt = self.get_data(scene) + return compute_pose( + torch.from_numpy(as_homogeneous(pred["extrinsics"])), + torch.from_numpy(as_homogeneous(gt["extrinsics"])), + ) + + @abstractmethod + def get_data(self, scene: str) -> Dict: + """ + Get scene data including images, camera parameters, and auxiliary info. + + Args: + scene: Scene identifier + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - camera extrinsics (world-to-camera) + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict - auxiliary data (masks, GT paths, etc.) + """ + raise NotImplementedError + + @abstractmethod + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate 3D reconstruction quality against ground truth. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with reconstruction metrics (e.g., acc, comp, overall) + """ + raise NotImplementedError + + @abstractmethod + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depth maps into a single point cloud. + + Args: + scene: Scene identifier + result_path: Path to .npz file with predicted depths and poses + fuse_path: Output path for fused point cloud (.ply) + mode: Fusion mode ("recon_unposed" or "recon_posed") + """ + raise NotImplementedError + diff --git a/src/nn/depth_anything_3/bench/datasets/__init__.py b/src/nn/depth_anything_3/bench/datasets/__init__.py new file mode 100644 index 0000000..2e9f464 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Benchmark dataset implementations. + +Datasets are auto-registered via decorators when imported. +Add new dataset files here and they will be automatically discovered. +""" + diff --git a/src/nn/depth_anything_3/bench/datasets/dtu.py b/src/nn/depth_anything_3/bench/datasets/dtu.py new file mode 100644 index 0000000..4d07c91 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/dtu.py @@ -0,0 +1,681 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +DTU Benchmark dataset implementation. + +DTU is a multi-view stereo benchmark for 3D reconstruction evaluation. +Reference: https://roboimagedata.compute.dtu.dk/ + +Note: DepthAnything3 was never trained on any images from DTU. +""" + +import glob +import os +from typing import Dict as TDict, List + +import numpy as np +import open3d as o3d +import torch +import torch.nn.functional as F +from addict import Dict +from PIL import Image +from plyfile import PlyData +from scipy.io import loadmat +from sklearn import neighbors as skln +from tqdm import tqdm + +from depth_anything_3.bench.dataset import Dataset +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.utils.constants import ( + DTU_DIST_THRESH, + DTU_EVAL_DATA_ROOT, + DTU_MAX_POINTS, + DTU_NUM_CONSIST, + DTU_SCENES, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +@MV_REGISTRY.register(name="dtu") +@MONO_REGISTRY.register(name="dtu") +class DTU(Dataset): + """ + DTU Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (accuracy, completeness, overall) + - Point cloud fusion from depth maps + + The dataset uses MVSNet evaluation protocol: + https://drive.google.com/file/d/1rX0EXlUL4prRxrRu2DgLJv2j7-tpUD4D/view + """ + + data_root = DTU_EVAL_DATA_ROOT + SCENES = DTU_SCENES + + # Evaluation/triangulation hyperparameters from constants + dist_thresh = DTU_DIST_THRESH + num_consist = DTU_NUM_CONSIST + + # ------------------------------ + # Public API + # ------------------------------ + + def read_cam_file(self, filename: str) -> tuple: + """ + Read DTU camera file containing extrinsics and intrinsics. + + Args: + filename: Path to camera text file + + Returns: + Tuple of (intrinsics [3,3], extrinsics [4,4]) + """ + with open(filename) as f: + lines = [line.rstrip() for line in f.readlines()] + extrinsics = np.fromstring(" ".join(lines[1:5]), dtype=np.float32, sep=" ").reshape((4, 4)) + intrinsics = np.fromstring(" ".join(lines[7:10]), dtype=np.float32, sep=" ").reshape((3, 3)) + return intrinsics, extrinsics + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics, and GT masks. + + Args: + scene: Scene identifier (e.g., "scan1") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] + - intrinsics: np.ndarray [N, 3, 3] + - aux.mask_files: List[str] - paths to depth masks + """ + rgb_folder = os.path.join(self.data_root, "Rectified", scene) + camera_folder = os.path.join(self.data_root, "Cameras") + + files = sorted(glob.glob(os.path.join(rgb_folder, "*.png"))) + # Reorder: place index 33 first (reference view convention) + files = [files[33]] + files[:33] + files[34:] + + out = Dict( + { + "image_files": files, + "extrinsics": [], + "intrinsics": [], + "aux": Dict({"mask_files": []}), + } + ) + + for rgb_file in files: + basename = os.path.basename(rgb_file) + file_idx = basename.split("_")[1] + cam_idx = depth_idx = int(file_idx) - 1 + + mask_file = self._depth_mask_path(scene, depth_idx) + proj_mat_filename = os.path.join(camera_folder, f"{cam_idx:0>8}_cam.txt") + + ixt, ext = self.read_cam_file(proj_mat_filename) + out.extrinsics.append(ext) + out.intrinsics.append(ixt) + out.aux.mask_files.append(mask_file) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + return out + + def get_3dgtpath(self, scene: str) -> str: + """Get path to ground truth point cloud for a scene.""" + scene_id = int(scene[4:]) + return os.path.join(self.data_root, f"Points/stl/stl{scene_id:03}_total.ply") + + def eval3d(self, scene: str, fuse_path: str, use_gpu: bool = False) -> TDict[str, float]: + """ + Evaluate fused point cloud against DTU GT with ObsMask/Plane. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud + use_gpu: If True, use GPU-accelerated distance computation (faster but may have minor numerical differences) + + Returns: + Dict with metrics: {"comp": float, "acc": float, "overall": float} + """ + scene_id = int(scene[4:]) + gt_ply = os.path.join(self.data_root, f"Points/stl/stl{scene_id:03}_total.ply") + mask_file = os.path.join( + self.data_root, f"SampleSet/mvs_data/ObsMask/ObsMask{scene_id}_10.mat" + ) + plane_file = os.path.join( + self.data_root, f"SampleSet/mvs_data/ObsMask/Plane{scene_id}.mat" + ) + result = self._evaluate_reconstruction( + scene, fuse_path, gt_ply, mask_file, plane_file, use_gpu=use_gpu + ) + return {"comp": result[0], "acc": result[1], "overall": result[2]} + + def load_masks(self, mask_files: List[str]) -> np.ndarray: + """ + Load DTU depth validity masks. + + Args: + mask_files: List of paths to mask images + + Returns: + Boolean array [N, H, W] indicating valid depth regions + """ + masks = [] + for mask_file in mask_files: + mask = Image.open(mask_file) + mask = np.array(mask, dtype=np.float32) + masks.append(mask > 10) + return np.asarray(masks) + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud and save to PLY. + + Args: + scene: Scene identifier (e.g., "scan114") + result_path: Path to npz file containing predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + gt_data = self.get_data(scene) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + masks = self.load_masks(gt_data.aux.mask_files) + + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed(pred_data, gt_data, masks) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed(pred_data, gt_data, masks) + else: + raise ValueError(f"Invalid mode: {mode}") + + proj_mat = self._build_proj_mats(intrinsics, extrinsics) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float32 + depths_t = torch.from_numpy(depths).to(device=device, dtype=dtype).unsqueeze(1) + proj_t = torch.from_numpy(proj_mat).to(device=device, dtype=dtype) + height, width = depths_t.shape[-2:] + + points: List[np.ndarray] = [] + for idx in range(len(gt_data.image_files)): + if mode == "recon_unposed": + # Simple unfiltered back-projection per frame + cur_p_pcd = self._generate_points_from_depth( + depths_t[idx : idx + 1], proj_t[idx : idx + 1] + ) + mask = (depths_t[idx : idx + 1] > 0.001).squeeze() + cur_p_pcd = cur_p_pcd[:, :, mask] + no_filter_pc = cur_p_pcd.squeeze(0).permute(1, 0).cpu().numpy() + points.append(no_filter_pc) + else: # recon_posed + final_pc = self._fuse_consistent_points(depths_t, proj_t, idx, height, width) + points.append(final_pc) + + # Concatenate and optionally downsample to hard cap + points_np = np.concatenate(points, axis=0) + points_np = self._cap_points(points_np, max_points=DTU_MAX_POINTS) + + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points_np) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Geometry helpers + # ------------------------------ + + def _generate_points_from_depth( + self, depth: torch.Tensor, proj: torch.Tensor + ) -> torch.Tensor: + """ + Back-project depth map into 3D world coordinates. + + Args: + depth: Depth tensor [B, 1, H, W] + proj: Projection matrix [B, 4, 4] = [[K@R, K@t], [0,0,0,1]] + + Returns: + Point cloud tensor [B, 3, H, W] + """ + batch, height, width = depth.shape[0], depth.shape[2], depth.shape[3] + inv_proj = torch.inverse(proj) + rot = inv_proj[:, :3, :3] + trans = inv_proj[:, :3, 3:4] + + y, x = torch.meshgrid( + [ + torch.arange(0, height, dtype=torch.float32, device=depth.device), + torch.arange(0, width, dtype=torch.float32, device=depth.device), + ], + indexing="ij", + ) + y, x = y.contiguous(), x.contiguous() + y, x = y.view(height * width), x.view(height * width) + xyz = torch.stack((x, y, torch.ones_like(x))) + xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1) + rot_xyz = torch.matmul(rot, xyz) + rot_depth_xyz = rot_xyz * depth.view(batch, 1, -1) + proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1) + return proj_xyz.view(batch, 3, height, width) + + def _homo_warping( + self, + src_fea: torch.Tensor, + src_proj: torch.Tensor, + ref_proj: torch.Tensor, + depth_values: torch.Tensor, + ) -> torch.Tensor: + """ + Homography warping for multi-view consistency checking. + + Args: + src_fea: Source features [B, C, H, W] + src_proj: Source projection [B, 4, 4] + ref_proj: Reference projection [B, 4, 4] + depth_values: Depth values [B, Ndepth] or [B, Ndepth, H, W] + + Returns: + Warped features [B, C, H, W] + """ + batch, channels = src_fea.shape[0], src_fea.shape[1] + height, width = src_fea.shape[2], src_fea.shape[3] + + with torch.no_grad(): + proj = torch.matmul(src_proj, torch.inverse(ref_proj)) + rot = proj[:, :3, :3] + trans = proj[:, :3, 3:4] + + y, x = torch.meshgrid( + [ + torch.arange(0, height, dtype=torch.float32, device=src_fea.device), + torch.arange(0, width, dtype=torch.float32, device=src_fea.device), + ], + indexing="ij", + ) + y, x = y.contiguous(), x.contiguous() + y, x = y.view(height * width), x.view(height * width) + xyz = torch.stack((x, y, torch.ones_like(x))) + xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1) + rot_xyz = torch.matmul(rot, xyz) + + rot_depth_xyz = rot_xyz.unsqueeze(2) * depth_values.view(-1, 1, 1, height * width) + proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1) + proj_xy = proj_xyz[:, :2, :, :] / proj_xyz[:, 2:3, :, :] + proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1 + proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1 + grid = torch.stack((proj_x_normalized, proj_y_normalized), dim=3) + + warped_src_fea = F.grid_sample( + src_fea, + grid.view(batch, height, width, 2), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + return warped_src_fea.view(batch, channels, height, width) + + def _filter_depth( + self, + ref_depth: torch.Tensor, + src_depths: torch.Tensor, + ref_proj: torch.Tensor, + src_projs: torch.Tensor, + ) -> tuple: + """ + Compute geometric consistency between reference and source depths. + + Args: + ref_depth: Reference depth [1, 1, H, W] + src_depths: Source depths [B, 1, H, W] + ref_proj: Reference projection [1, 4, 4] + src_projs: Source projections [B, 4, 4] + + Returns: + Tuple of (ref_pc, aligned_pcs, dist) + """ + ref_pc = self._generate_points_from_depth(ref_depth, ref_proj) + src_pcs = self._generate_points_from_depth(src_depths, src_projs) + aligned_pcs = self._homo_warping(src_pcs, src_projs, ref_proj, ref_depth) + x_2 = (ref_pc[:, 0] - aligned_pcs[:, 0]) ** 2 + y_2 = (ref_pc[:, 1] - aligned_pcs[:, 1]) ** 2 + z_2 = (ref_pc[:, 2] - aligned_pcs[:, 2]) ** 2 + dist = torch.sqrt(x_2 + y_2 + z_2).unsqueeze(1) + return ref_pc, aligned_pcs, dist + + def _extract_points( + self, pc: torch.Tensor, mask: torch.Tensor, rgb: np.ndarray = None + ) -> np.ndarray: + """Extract masked points from a dense grid.""" + pc = pc.cpu().numpy() + mask = mask.cpu().numpy().reshape(-1) + pc = pc.reshape(-1, 3) + points = pc[np.where(mask)] + if rgb is not None: + rgb = rgb.reshape(-1, 3) + colors = rgb[np.where(mask)] + return np.concatenate([points, colors], axis=1) + return points + + # ------------------------------ + # 3D Reconstruction Evaluation + # ------------------------------ + + def _evaluate_reconstruction( + self, + scanid: str, + pred_ply: str, + gt_ply: str, + mask_file: str, + plane_file: str, + down_dense: float = 0.2, + patch: int = 60, + max_dist: int = 20, + use_gpu: bool = False, + ) -> tuple: + """ + Compute accuracy, completeness, and overall metrics for one scan. + + Args: + scanid: Scan identifier + pred_ply: Predicted point cloud path or array + gt_ply: Ground truth point cloud path or array + mask_file: ObsMask file path + plane_file: Plane file path + down_dense: Downsample density (min distance between points) + patch: Patch size for boundary + max_dist: Outlier threshold in mm + use_gpu: If True, use GPU-accelerated distance computation + + Returns: + Tuple of (mean_d2s, mean_s2d, overall) + """ + thresh = down_dense + + # Load and downsample predicted point cloud + data_pcd = self._read_ply(pred_ply) if isinstance(pred_ply, str) else pred_ply + # Use fixed seed for reproducibility + shuffle_rng = np.random.default_rng(seed=42) + shuffle_rng.shuffle(data_pcd, axis=0) + + # Downsample point cloud + nn_engine = skln.NearestNeighbors( + n_neighbors=1, radius=thresh, algorithm="kd_tree", n_jobs=-1 + ) + nn_engine.fit(data_pcd) + rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False) + mask = np.ones(data_pcd.shape[0], dtype=np.bool_) + for curr, idxs in enumerate(rnn_idxs): + if mask[curr]: + mask[idxs] = 0 + mask[curr] = 1 + data_down = data_pcd[mask] + + # Restrict to observed volume (ObsMask) + obs_mask_file = loadmat(mask_file) + ObsMask, BB, Res = (obs_mask_file[attr] for attr in ["ObsMask", "BB", "Res"]) + BB = BB.astype(np.float32) + + inbound = ((data_down >= BB[:1] - patch) & (data_down < BB[1:] + patch * 2)).sum( + axis=-1 + ) == 3 + data_in = data_down[inbound] + + data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32) + grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum( + axis=-1 + ) == 3 + data_grid_in = data_grid[grid_inbound] + in_obs = ObsMask[data_grid_in[:, 0], data_grid_in[:, 1], data_grid_in[:, 2]].astype( + np.bool_ + ) + data_in_obs = data_in[grid_inbound][in_obs] + + # Compute accuracy (pred -> GT) and completeness (GT -> pred) + stl = self._read_ply(gt_ply) if isinstance(gt_ply, str) else gt_ply + + if use_gpu and torch.cuda.is_available(): + # GPU-accelerated distance computation + mean_d2s = self._knn_dist_gpu(data_in_obs, stl, max_dist) + else: + # CPU version (original, for exact reproduction) + nn_engine.fit(stl) + dist_d2s, _ = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True) + mean_d2s = dist_d2s[dist_d2s < max_dist].mean() + + ground_plane = loadmat(plane_file)["P"] + stl_hom = np.concatenate([stl, np.ones_like(stl[:, :1])], -1) + above = (ground_plane.reshape((1, 4)) * stl_hom).sum(-1) > 0 + stl_above = stl[above] + + if use_gpu and torch.cuda.is_available(): + # GPU-accelerated distance computation + mean_s2d = self._knn_dist_gpu(stl_above, data_in, max_dist) + else: + # CPU version (original, for exact reproduction) + nn_engine.fit(data_in) + dist_s2d, _ = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True) + mean_s2d = dist_s2d[dist_s2d < max_dist].mean() + + overall = (mean_d2s + mean_s2d) / 2 + return mean_d2s, mean_s2d, overall + + def _knn_dist_gpu( + self, + query: np.ndarray, + target: np.ndarray, + max_dist: float, + batch_size: int = 8192, + target_batch_size: int = 50000, + ) -> float: + """ + GPU-accelerated nearest neighbor distance computation. + + Args: + query: Query points [N, 3] + target: Target points [M, 3] + max_dist: Outlier threshold + batch_size: Batch size for query to avoid OOM (tuned for 16GB GPU) + target_batch_size: Batch size for target to avoid OOM + + Returns: + Mean distance (excluding outliers) + """ + device = torch.device("cuda") + + all_min_dists = [] + n_query_batches = (len(query) + batch_size - 1) // batch_size + n_target_batches = (len(target) + target_batch_size - 1) // target_batch_size + + # Pre-load target batches to GPU to avoid repeated transfers + # Memory: ~50000 pts * 3 coords * 4 bytes * n_batches + target_batches = [] + for j in range(0, len(target), target_batch_size): + target_batch = target[j : j + target_batch_size] + target_t = torch.from_numpy(target_batch).float().to(device) + target_batches.append(target_t) + + with tqdm(total=n_query_batches, desc=" GPU KNN", leave=False, ncols=100) as pbar: + for i in range(0, len(query), batch_size): + batch = query[i : i + batch_size] + query_t = torch.from_numpy(batch).float().to(device) + + # Compute distances to all target batches + # Memory peak: query_batch × target_batch_size × 4 bytes + # = 8192 × 50000 × 4 = ~1.6 GB per cdist call + batch_min_dists = [] + for target_t in target_batches: + dists = torch.cdist(query_t, target_t) + batch_min_dists.append(dists.min(dim=1).values) + del dists # Free immediately + + # Get minimum distance across all target batches + min_dists = torch.stack(batch_min_dists, dim=1).min(dim=1).values + all_min_dists.append(min_dists.cpu().numpy()) + + del query_t, min_dists, batch_min_dists + pbar.update(1) + + # Clean up target batches + for target_t in target_batches: + del target_t + torch.cuda.empty_cache() + + all_min_dists = np.concatenate(all_min_dists) + return all_min_dists[all_min_dists < max_dist].mean() + + def _read_ply(self, file: str) -> np.ndarray: + """Read point cloud from PLY file.""" + data = PlyData.read(file) + vertex = data["vertex"] + return np.stack([vertex["x"], vertex["y"], vertex["z"]], axis=-1) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _depth_mask_path(self, scene: str, depth_idx: int) -> str: + """Get path to depth mask for a scene and frame.""" + return os.path.join( + self.data_root, "depth_raw", "Depths", scene, f"depth_visual_{depth_idx:04d}.png" + ) + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, masks: np.ndarray + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_unposed mode. + + Applies Umeyama scale, rescales intrinsics if depth resolution differs, + and zeroes invalid-mask depths (nearest interpolation as in paper). + """ + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + ransac=True, + return_aligned=True, + random_state=42, + ) + depths = pred_data.depth * scale + intrinsics = pred_data.intrinsics.copy() + + if depths.shape[-2:] != masks.shape[-2:]: + # When resizing depths to mask size, adjust intrinsics accordingly + sx = masks.shape[-1] / depths.shape[-1] + sy = masks.shape[-2] / depths.shape[-2] + intrinsics[:, 0:1] *= sx + intrinsics[:, 1:2] *= sy + depths = F.interpolate( + torch.from_numpy(depths)[None].float(), + size=(masks.shape[-2], masks.shape[-1]), + mode="nearest", + )[0].numpy() + depths[masks == False] = 0.0 # noqa: E712 + + return depths, intrinsics, extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, masks: np.ndarray + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_posed mode. + + Uses GT intrinsics/extrinsics but aligns scale via Umeyama. + Same mask order as other datasets: mask BEFORE scale. + """ + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + ransac=True, + return_aligned=True, + random_state=42, + ) + depths = pred_data.depth.copy() + intrinsics = gt_data.intrinsics.copy() + extrinsics = gt_data.extrinsics.copy() + + if depths.shape[-2:] != masks.shape[-2:]: + depths = F.interpolate( + torch.from_numpy(depths)[None].float(), + size=(masks.shape[-2], masks.shape[-1]), + mode="nearest", + )[0].numpy() + + # Mask BEFORE scale (same as other datasets) + depths[masks == False] = 0.0 # noqa: E712 + depths = depths * scale + + return depths, intrinsics, extrinsics + + def _build_proj_mats( + self, intrinsics: np.ndarray, extrinsics: np.ndarray + ) -> np.ndarray: + """Compute per-view 4x4 projection matrices from K and [R|t].""" + proj_mat_list = [] + for i in range(len(intrinsics)): + proj_mat = np.eye(4, dtype=np.float32) + proj_mat[:3, :4] = np.dot(intrinsics[i], extrinsics[i][:3]) + proj_mat_list.append(proj_mat) + return np.stack(proj_mat_list, axis=0) + + def _fuse_consistent_points( + self, + depths_t: torch.Tensor, + proj_t: torch.Tensor, + idx: int, + H: int, + W: int, + ) -> np.ndarray: + """Fuse points consistent across multiple source views for a reference index.""" + device, dtype = depths_t.device, depths_t.dtype + pc_buff = torch.zeros((3, H, W), device=device, dtype=dtype) + val_cnt = torch.zeros((1, H, W), device=device, dtype=dtype) + + j = 0 + batch_size = 20 + tot_frame = depths_t.shape[0] + while True: + ref_pc, pcs, dist = self._filter_depth( + ref_depth=depths_t[idx : idx + 1], + src_depths=depths_t[j : min(j + batch_size, tot_frame)], + ref_proj=proj_t[idx : idx + 1], + src_projs=proj_t[j : min(j + batch_size, tot_frame)], + ) + masks = (dist < self.dist_thresh).float() + masked_pc = pcs * masks + pc_buff += masked_pc.sum(dim=0, keepdim=False) + val_cnt += masks.sum(dim=0, keepdim=False) + j += batch_size + if j >= tot_frame: + break + + final_mask = (val_cnt >= self.num_consist).squeeze(0) + avg_points = torch.div(pc_buff, val_cnt).permute(1, 2, 0) + final_pc = self._extract_points(avg_points, final_mask) + return final_pc + + def _cap_points(self, points: np.ndarray, max_points: int) -> np.ndarray: + """Downsample points if exceeding max count.""" + if len(points) <= max_points: + return points + # Use fixed seed for reproducibility + rng = np.random.default_rng(seed=42) + random_idx = rng.choice(len(points), max_points, replace=False) + return points[random_idx] + diff --git a/src/nn/depth_anything_3/bench/datasets/dtu64.py b/src/nn/depth_anything_3/bench/datasets/dtu64.py new file mode 100644 index 0000000..36ebb63 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/dtu64.py @@ -0,0 +1,182 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +DTU-64 Dataset implementation for POSE EVALUATION ONLY. + +This is a subset of DTU with 64 images per scene, specifically designed for +camera pose estimation evaluation. It does NOT support 3D reconstruction. + +Note: GT depth loading is not implemented as it's not needed for pose evaluation. +""" + +import glob +import os +from typing import Dict as TDict + +import numpy as np +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.utils.constants import ( + DTU64_CAMERA_ROOT, + DTU64_EVAL_DATA_ROOT, + DTU64_SCENES, +) + + +@MV_REGISTRY.register(name="dtu64") +@MONO_REGISTRY.register(name="dtu64") +class DTU64(Dataset): + """ + DTU-64 Dataset wrapper for DepthAnything3 POSE EVALUATION ONLY. + + This dataset is a subset of DTU with 64 images per scene. + It is specifically designed for camera pose estimation evaluation + and does NOT support 3D reconstruction evaluation. + + Dataset structure: + DTU/scans/ + ├── {scene}/ + │ └── image/ # RGB images (64 per scene) + └── Cameras/ + └── {idx}_cam.txt # Camera parameters + + Supported modes: + - pose: Camera pose estimation evaluation + + NOT supported: + - recon_unposed: 3D reconstruction (no GT depth available) + - recon_posed: 3D reconstruction (no GT depth available) + """ + + data_root = DTU64_EVAL_DATA_ROOT + camera_root = DTU64_CAMERA_ROOT + SCENES = DTU64_SCENES + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Camera file parsing + # ------------------------------ + + def read_cam_file(self, filename: str) -> tuple: + """ + Read DTU camera file containing extrinsics and intrinsics. + + Args: + filename: Path to camera text file + + Returns: + Tuple of (intrinsics [3,3], extrinsics [4,4]) + """ + with open(filename) as f: + lines = [line.rstrip() for line in f.readlines()] + # extrinsics: line [1,5), 4x4 matrix + extrinsics = np.fromstring(" ".join(lines[1:5]), dtype=np.float32, sep=" ").reshape((4, 4)) + # intrinsics: line [7-10), 3x3 matrix + intrinsics = np.fromstring(" ".join(lines[7:10]), dtype=np.float32, sep=" ").reshape((3, 3)) + return intrinsics, extrinsics + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene identifier (e.g., "scan105") + + Returns: + Dict with: + - image_files: List[str] - paths to images (64 per scene) + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict (empty for this dataset) + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + rgb_folder = os.path.join(self.data_root, scene, "image") + + # Get all PNG files sorted + files = sorted(glob.glob(os.path.join(rgb_folder, "*.png"))) + + # Reorder: place index 33 first (reference view convention) + if len(files) > 33: + files = [files[33]] + files[:33] + files[34:] + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({}), + }) + + for rgb_file in files: + basename = os.path.basename(rgb_file) + # File naming: "00000033.png" -> cam_idx = 33 + file_idx = basename.split(".")[0] + cam_idx = int(file_idx) + + # Camera file path + cam_file = os.path.join(self.camera_root, f"{cam_idx:0>8}_cam.txt") + + if not os.path.exists(cam_file): + print(f"[DTU-64] Warning: Camera file not found: {cam_file}") + continue + + intrinsics, extrinsics = self.read_cam_file(cam_file) + + out.image_files.append(rgb_file) + out.extrinsics.append(extrinsics) + out.intrinsics.append(intrinsics) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[DTU-64] {scene}: {len(out.image_files)} images (pose evaluation only)") + + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + NOT SUPPORTED for DTU-64. + + DTU-64 is only for pose evaluation, not 3D reconstruction. + """ + raise NotImplementedError( + "DTU-64 dataset is for POSE EVALUATION ONLY. " + "3D reconstruction evaluation is not supported. " + "Use the standard 'dtu' dataset for 3D reconstruction evaluation." + ) + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + NOT SUPPORTED for DTU-64. + + DTU-64 is only for pose evaluation, not 3D reconstruction. + """ + raise NotImplementedError( + "DTU-64 dataset is for POSE EVALUATION ONLY. " + "3D reconstruction (fuse3d) is not supported. " + "Use the standard 'dtu' dataset for 3D reconstruction." + ) + diff --git a/src/nn/depth_anything_3/bench/datasets/eth3d.py b/src/nn/depth_anything_3/bench/datasets/eth3d.py new file mode 100644 index 0000000..b390695 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/eth3d.py @@ -0,0 +1,594 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +ETH3D Benchmark dataset implementation. + +ETH3D is a multi-view stereo benchmark with high-resolution images and +accurate ground truth geometry from laser scanning. +Reference: https://www.eth3d.net/ + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import glob +import os +from typing import Dict as TDict, List, Optional + +import cv2 +import numpy as np +import open3d as o3d +import torch +import torch.nn.functional as F +from addict import Dict +from PIL import Image + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + evaluate_3d_reconstruction, + fuse_depth_to_tsdf, + quat2rotmat, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + ETH3D_DOWN_SAMPLE, + ETH3D_EVAL_DATA_ROOT, + ETH3D_EVAL_THRESHOLD, + ETH3D_FILTER_KEYS, + ETH3D_MAX_DEPTH, + ETH3D_SAMPLING_NUMBER, + ETH3D_SCENES, + ETH3D_SDF_TRUNC, + ETH3D_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +@MV_REGISTRY.register(name="eth3d") +@MONO_REGISTRY.register(name="eth3d") +class ETH3D(Dataset): + """ + ETH3D Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + eth3d/multiview/ + ├── scene_name/ + │ ├── images/ # RGB images + │ ├── dslr_calibration_jpg/ + │ │ ├── cameras.txt # Camera intrinsics + │ │ └── images.txt # Camera poses + │ ├── combined_mesh.ply # Ground truth mesh + │ └── ground_truth_depth/ # GT depth maps (optional) + """ + + data_root = ETH3D_EVAL_DATA_ROOT + SCENES = ETH3D_SCENES + + # Evaluation hyperparameters from constants + max_depth = ETH3D_MAX_DEPTH + sampling_number = ETH3D_SAMPLING_NUMBER + voxel_length = ETH3D_VOXEL_LENGTH + sdf_trunc = ETH3D_SDF_TRUNC + eval_threshold = ETH3D_EVAL_THRESHOLD + down_sample = ETH3D_DOWN_SAMPLE + + def __init__(self): + super().__init__() + # Pre-load scene data for efficiency + self._scene_cache = {} + + # ------------------------------ + # Camera file parsing + # ------------------------------ + + def _parse_cameras_txt(self, filepath: str) -> dict: + """ + Parse COLMAP-style cameras.txt file. + + Returns: + Dict mapping camera_id to intrinsic parameters + """ + camera_dict = {} + with open(filepath) as f: + lines = f.readlines() + for line in lines[3:]: # Skip header + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split() + if len(parts) < 8: + continue + cam_id = parts[0] + # Format: ID, MODEL, WIDTH, HEIGHT, fx, fy, cx, cy, [distortion params...] + camera_dict[cam_id] = { + "width": float(parts[2]), + "height": float(parts[3]), + "fx": float(parts[4]), + "fy": float(parts[5]), + "cx": float(parts[6]), + "cy": float(parts[7]), + } + return camera_dict + + def _parse_images_txt(self, filepath: str) -> dict: + """ + Parse COLMAP-style images.txt file. + + Returns: + Dict mapping image path to pose parameters + """ + pose_dict = {} + with open(filepath) as f: + lines = f.readlines() + for idx, line in enumerate(lines[4:]): # Skip header + line = line.strip() + if not line or line.startswith("#"): + continue + # Every other line contains pose info + if idx % 2 == 0: + parts = line.split() + if len(parts) < 10: + continue + # Format: IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME + image_id = parts[0] + qw, qx, qy, qz = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4]) + tx, ty, tz = float(parts[5]), float(parts[6]), float(parts[7]) + camera_id = parts[8] + name = parts[9] + pose_dict[name] = { + "image_id": image_id, + "quat": [qw, qx, qy, qz], + "trans": [tx, ty, tz], + "camera_id": camera_id, + } + return pose_dict + + def _should_filter_image(self, scene: str, image_name: str) -> bool: + """Check if image should be filtered out based on known problematic views.""" + filter_keys = ETH3D_FILTER_KEYS.get(scene, []) + for key in filter_keys: + if image_name.endswith(key): + return True + return False + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene identifier (e.g., "courtyard") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_mesh_path + """ + # Check cache + if scene in self._scene_cache: + return self._scene_cache[scene] + + scene_dir = os.path.join(self.data_root, scene) + + # Parse camera files + cameras_file = os.path.join(scene_dir, "dslr_calibration_jpg", "cameras.txt") + images_file = os.path.join(scene_dir, "dslr_calibration_jpg", "images.txt") + camera_dict = self._parse_cameras_txt(cameras_file) + pose_dict = self._parse_images_txt(images_file) + + # Ground truth mesh path + gt_mesh_path = os.path.join(scene_dir, "combined_mesh.ply") + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_mesh_path": gt_mesh_path, + "heights": [], + "widths": [], + }), + }) + + # Process each image (preserve original order from images.txt) + filtered_count = 0 + for image_name, pose_info in pose_dict.items(): + # Filter problematic views + if self._should_filter_image(scene, image_name): + filtered_count += 1 + continue + + image_path = os.path.join(scene_dir, "images", image_name) + if not os.path.exists(image_path): + continue + + cam_info = camera_dict.get(pose_info["camera_id"]) + if cam_info is None: + continue + + # Build intrinsics matrix + ixt = np.array([ + [cam_info["fx"], 0, cam_info["cx"]], + [0, cam_info["fy"], cam_info["cy"]], + [0, 0, 1], + ], dtype=np.float32) + + # Build extrinsics matrix (world-to-camera) + # COLMAP format: world point -> camera point + rot = quat2rotmat(pose_info["quat"]) + ext = np.eye(4, dtype=np.float32) + ext[:3, :3] = rot + ext[:3, 3] = pose_info["trans"] + + out.image_files.append(image_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt) + out.aux.heights.append(cam_info["height"]) + out.aux.widths.append(cam_info["width"]) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + # Print scene info + total_images = len(pose_dict) + used_images = len(out.image_files) + print(f"[ETH3D] {scene}: {used_images}/{total_images} images " + f"(filtered {filtered_count}, missing {total_images - used_images - filtered_count})") + + if used_images < 3: + print(f"[ETH3D] ⚠️ WARNING: {scene} has only {used_images} images - evaluation may fail!") + + # Cache result + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against ETH3D ground truth mesh. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_mesh_path = gt_data.aux.gt_mesh_path + + # Load and sample ground truth mesh + gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path) + gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Evaluate using shared utility function + metrics = evaluate_3d_reconstruction( + pred_pcd, + gt_pcd, + threshold=self.eval_threshold, + down_sample=self.down_sample, + ) + + return metrics + + def _load_gt_meta(self, result_path: str) -> Dict: + """ + Load saved GT meta (extrinsics, intrinsics, image_files) for fusion. + + This is needed when frames are sampled, so fuse3d uses the correct + (sampled) GT instead of full dataset GT. + + Args: + result_path: Path to npz file (used to derive gt_meta.npz path) + + Returns: + Dict with GT data, or None if gt_meta.npz doesn't exist + """ + # gt_meta.npz is in the same exports/ directory as results.npz + export_dir = os.path.dirname(result_path) # exports/mini_npz/ + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": data["image_files"] if "image_files" in data else None, + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Pipeline: + 1. Load original images (keep original size) + 2. Resize depth to original image size (nearest interpolation) + 3. Adjust intrinsics to original image size + 4. Apply scale alignment and mask invalid depths + 5. TSDF fusion + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + else: + gt_data = self.get_data(scene) + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load original images (keep original size) + images = [] + orig_sizes = [] # (H, W) for each image + for img_path in gt_data.image_files: + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + images.append(img) + orig_sizes.append((img.shape[0], img.shape[1])) + + # Prepare depths, intrinsics, extrinsics with resize to original size + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, orig_sizes, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, orig_sizes, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + images = np.stack(images, axis=0) + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str = None + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_unposed mode. + + Pipeline: + 1. Umeyama scale alignment + 2. Load GT mask for each frame + 3. Resize depth to original image size (nearest) + 4. Apply GT mask BEFORE scale + 5. Apply scale + 6. Adjust intrinsics to original image size + """ + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + # Get model output size + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + # Process each frame + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + image_name = os.path.basename(gt_data.image_files[i]) + + # Resize depth to original image size (nearest interpolation) + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask (apply BEFORE scale) + gt_zero_mask = None + if scene is not None: + gt_zero_mask = self._load_gt_mask(scene, image_name, (orig_h, orig_w)) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to original image size + h_ratio = orig_h / model_h + w_ratio = orig_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio # fx, 0, cx + ixt[1, :] *= h_ratio # 0, fy, cy + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str = None + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_posed mode. + + Uses GT intrinsics/extrinsics but aligns depth scale via Umeyama. + Depth is resized to original image size. + """ + # Scale alignment with fixed random_state for reproducibility + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + # Process each frame + depths_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + image_name = os.path.basename(gt_data.image_files[i]) + + # Resize depth to original image size (nearest interpolation) + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask (apply BEFORE scale) + gt_zero_mask = None + if scene is not None: + gt_zero_mask = self._load_gt_mask(scene, image_name, (orig_h, orig_w)) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Use GT intrinsics and extrinsics (already at original image size) + return np.stack(depths_out), gt_data.intrinsics.copy(), gt_data.extrinsics.copy() + + def _load_gt_mask(self, scene: str, image_name: str, shape: tuple) -> np.ndarray: + """ + Load GT mask for masking invalid regions. + + GT mask marks occluded or invalid regions that should be excluded + from depth fusion and evaluation. + + Args: + scene: Scene identifier + image_name: Image filename (e.g., "DSC_0307.JPG") + shape: (height, width) of the image + + Returns: + Boolean mask where True = valid region to keep + """ + h, w = shape + + # GT mask file path + gt_mask_path = os.path.join( + self.data_root, scene, "masks_for_images", "dslr_images", + image_name.replace(".JPG", ".png") + ) + + # GT depth file path (used to determine valid depth regions) + gt_depth_path = os.path.join( + self.data_root, scene, "ground_truth_depth", "dslr_images", image_name + ) + + # Load GT depth + if os.path.exists(gt_depth_path): + gt_depth = np.fromfile(gt_depth_path, dtype=np.float32).reshape(h, w) + else: + gt_depth = np.ones((h, w), dtype=np.float32) + + # Load GT mask + if os.path.exists(gt_mask_path): + gt_mask = cv2.imread(gt_mask_path, cv2.IMREAD_GRAYSCALE) + gt_mask = np.asarray(gt_mask) + else: + gt_mask = np.zeros((h, w), dtype=np.uint8) + + # Compute zero_mask + # gt_mask == 1 means occluded/invalid region + invalid_mask_from_gt = gt_mask == 1 + gt_depth_copy = gt_depth.copy() + gt_depth_copy[gt_mask == 1] = 0 + + invalid_mask_from_gt_depth = np.logical_or(gt_depth_copy == 0, gt_depth_copy == np.inf) + + # zero_mask: valid region that should be kept + zero_mask = np.logical_and( + np.logical_not(invalid_mask_from_gt), + np.logical_not(invalid_mask_from_gt_depth) + ) + + return zero_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """ + Mask invalid depth values by setting them to 0. + + Logic: + 1. Apply GT mask (if provided) - marks occluded/invalid regions + 2. Mask pred invalid values (nan, inf) + + Args: + depth: Depth map to mask + gt_zero_mask: Optional GT mask (True = valid region) + + Returns: + Masked depth map with invalid regions set to 0 + """ + depth = depth.copy() + + # Apply GT mask first (before scale) + if gt_zero_mask is not None: + # Also mask out invalid pred depth + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + # Fallback: only mask pred invalid values + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + diff --git a/src/nn/depth_anything_3/bench/datasets/hiroom.py b/src/nn/depth_anything_3/bench/datasets/hiroom.py new file mode 100644 index 0000000..45619f7 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/hiroom.py @@ -0,0 +1,440 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +HiRoom Dataset implementation. + +HiRoom is an indoor RGB-D dataset containing ground truth camera poses, +depth maps, and fused point clouds. + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import os +from typing import Dict as TDict, List + +import cv2 +import numpy as np +import open3d as o3d +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + evaluate_3d_reconstruction, + fuse_depth_to_tsdf, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + HIROOM_DOWN_SAMPLE, + HIROOM_EVAL_DATA_ROOT, + HIROOM_EVAL_THRESHOLD, + HIROOM_GT_ROOT_PATH, + HIROOM_MAX_DEPTH, + HIROOM_SAMPLING_NUMBER, + HIROOM_SCENE_LIST_PATH, + HIROOM_SDF_TRUNC, + HIROOM_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +def _load_scene_list() -> List[str]: + """Load scene list from file.""" + if os.path.exists(HIROOM_SCENE_LIST_PATH): + with open(HIROOM_SCENE_LIST_PATH, "r") as f: + return f.read().splitlines() + return [] + + +@MV_REGISTRY.register(name="hiroom") +@MONO_REGISTRY.register(name="hiroom") +class HiRoomDataset(Dataset): + """ + HiRoom Dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + HiRoom/ + ├── {scene_path}/ + │ ├── image/ # RGB images + │ ├── depth/ # GT depth maps + │ ├── pose/ # Camera poses (.npy) + │ ├── cam_K.npy # Camera intrinsics + │ └── aliasing_mask/ # Aliasing masks + + fused_pcd/ + └── {scene_name}.ply # Ground truth fused point cloud + """ + + data_root = HIROOM_EVAL_DATA_ROOT + gt_root_path = HIROOM_GT_ROOT_PATH + SCENES = _load_scene_list() + + # Evaluation hyperparameters from constants + max_depth = HIROOM_MAX_DEPTH + sampling_number = HIROOM_SAMPLING_NUMBER + voxel_length = HIROOM_VOXEL_LENGTH + sdf_trunc = HIROOM_SDF_TRUNC + eval_threshold = HIROOM_EVAL_THRESHOLD + down_sample = HIROOM_DOWN_SAMPLE + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene path (e.g., "xxx/yyy/zzz") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_pcd_path, gt_depth_files, aliasing_mask_files + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + scene_dir = os.path.join(self.data_root, scene) + image_dir = os.path.join(scene_dir, "image") + + # Get scene name for GT point cloud + scene_name = "-".join(scene.split("/")[-3:]) + gt_pcd_path = os.path.join(self.gt_root_path, f"{scene_name}.ply") + + # Load shared camera intrinsics + intrin_path = os.path.join(scene_dir, "cam_K.npy") + ixt_shared = np.load(intrin_path).astype(np.float32) + + # Get all image names sorted + image_names = sorted(os.listdir(image_dir)) + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_pcd_path": gt_pcd_path, + "gt_depth_files": [], + "aliasing_mask_files": [], + }), + }) + + for img_name in image_names: + img_path = os.path.join(image_dir, img_name) + frame_name = img_name.split(".")[0] + + # Depth and pose paths + depth_path = os.path.join(scene_dir, "depth", f"{frame_name}.png") + pose_path = os.path.join(scene_dir, "pose", f"{frame_name}.npy") + aliasing_mask_path = os.path.join(scene_dir, "aliasing_mask", f"{frame_name}.png") + + if not os.path.exists(pose_path): + continue + + # Load extrinsics (world-to-camera) + ext = np.load(pose_path).astype(np.float32) + + out.image_files.append(img_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt_shared.copy()) + out.aux.gt_depth_files.append(depth_path) + out.aux.aliasing_mask_files.append(aliasing_mask_path) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[HiRoom] {scene}: {len(out.image_files)} images") + + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against HiRoom ground truth point cloud. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_pcd_path = gt_data.aux.gt_pcd_path + + # Load ground truth point cloud + gt_pcd = o3d.io.read_point_cloud(gt_pcd_path) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Evaluate using shared utility function + metrics = evaluate_3d_reconstruction( + pred_pcd, + gt_pcd, + threshold=self.eval_threshold, + down_sample=self.down_sample, + ) + + return metrics + + def _load_gt_meta(self, result_path: str) -> Dict: + """Load saved GT meta for fusion.""" + export_dir = os.path.dirname(result_path) + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + image_files = list(data["image_files"]) + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": image_files, + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Get full GT data + full_gt_data = self.get_data(scene) + + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + image_indices = [ + full_gt_data.image_files.index(f) + for f in gt_data.image_files + if f in full_gt_data.image_files + ] + else: + gt_data = full_gt_data + image_indices = list(range(len(full_gt_data.image_files))) + + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load images + images = [] + orig_sizes = [] + for img_idx in image_indices: + img_path = full_gt_data.image_files[img_idx] + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + images.append(img) + orig_sizes.append((img.shape[0], img.shape[1])) + + images = np.stack(images, axis=0) + + # Prepare depths, intrinsics, extrinsics + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, full_gt_data, image_indices, orig_sizes, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, full_gt_data, image_indices, orig_sizes, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, orig_sizes: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_unposed mode.""" + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + img_idx = image_indices[i] + + # Resize depth to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask + gt_zero_mask = self._load_gt_mask( + full_gt_data.aux.gt_depth_files[img_idx], + full_gt_data.aux.aliasing_mask_files[img_idx], + ) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to original image size + h_ratio = orig_h / model_h + w_ratio = orig_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio + ixt[1, :] *= h_ratio + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, orig_sizes: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_posed mode.""" + # Scale alignment + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + depths_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + img_idx = image_indices[i] + + # Resize depth to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT mask + gt_zero_mask = self._load_gt_mask( + full_gt_data.aux.gt_depth_files[img_idx], + full_gt_data.aux.aliasing_mask_files[img_idx], + ) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Use GT intrinsics and extrinsics + gt_intrinsics = np.stack([full_gt_data.intrinsics[idx] for idx in image_indices]) + gt_extrinsics = np.stack([full_gt_data.extrinsics[idx] for idx in image_indices]) + + return np.stack(depths_out), gt_intrinsics, gt_extrinsics + + def _load_gt_mask(self, gt_depth_path: str, aliasing_mask_path: str) -> np.ndarray: + """ + Load GT depth and aliasing mask to create valid mask. + + For HiRoom: + - GT depth is stored as 16-bit PNG, scaled to 100m range + - Aliasing mask marks regions to exclude + + Returns: + Boolean mask where True = valid region to keep + """ + # Load GT depth + if os.path.exists(gt_depth_path): + gt_depth = cv2.imread(gt_depth_path, -1) / 65535.0 * 100.0 + else: + return None + + # Load aliasing mask + aliasing_mask = None + if os.path.exists(aliasing_mask_path): + aliasing_mask = cv2.imread(aliasing_mask_path, -1) > 0 + + # Valid mask: depth > 0 and not in aliasing region + valid_mask = gt_depth > 0 + if aliasing_mask is not None: + valid_mask = np.logical_and(valid_mask, np.logical_not(aliasing_mask)) + + return valid_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """Mask invalid depth values by setting them to 0.""" + depth = depth.copy() + + if gt_zero_mask is not None: + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + diff --git a/src/nn/depth_anything_3/bench/datasets/scannetpp.py b/src/nn/depth_anything_3/bench/datasets/scannetpp.py new file mode 100644 index 0000000..f841ad1 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/scannetpp.py @@ -0,0 +1,591 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +ScanNet++ Benchmark dataset implementation. + +ScanNet++ is a high-quality indoor RGB-D dataset with iPhone and DSLR images, +ground truth camera poses from COLMAP, and high-resolution 3D meshes. +Reference: https://kaldir.vc.in.tum.de/scannetpp/ + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import os +from typing import Dict as TDict + +import cv2 +import imageio +import numpy as np +import open3d as o3d +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + fuse_depth_to_tsdf, + nn_correspondance, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + SCANNETPP_DOWN_SAMPLE, + SCANNETPP_EVAL_DATA_ROOT, + SCANNETPP_EVAL_THRESHOLD, + SCANNETPP_INPUT_H, + SCANNETPP_INPUT_W, + SCANNETPP_MAX_DEPTH, + SCANNETPP_SAMPLING_NUMBER, + SCANNETPP_SCENES, + SCANNETPP_SDF_TRUNC, + SCANNETPP_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama +from depth_anything_3.utils.read_write_model import read_model + + +@MV_REGISTRY.register(name="scannetpp") +@MONO_REGISTRY.register(name="scannetpp") +class ScanNetPP(Dataset): + """ + ScanNet++ Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + scannetpp/data/ + ├── {scene_id}/ + │ ├── merge_dslr_iphone/ + │ │ ├── colmap/sparse_render_rgb/ # COLMAP reconstruction + │ │ ├── images/ # RGB images + │ │ └── render_depth/ # GT depth maps + │ └── scans/ + │ └── mesh_aligned_0.05.ply # Ground truth mesh + """ + + data_root = SCANNETPP_EVAL_DATA_ROOT + SCENES = SCANNETPP_SCENES + + # Input resolution after undistortion and resize + input_h = SCANNETPP_INPUT_H + input_w = SCANNETPP_INPUT_W + + # Evaluation hyperparameters from constants + max_depth = SCANNETPP_MAX_DEPTH + sampling_number = SCANNETPP_SAMPLING_NUMBER + voxel_length = SCANNETPP_VOXEL_LENGTH + sdf_trunc = SCANNETPP_SDF_TRUNC + eval_threshold = SCANNETPP_EVAL_THRESHOLD + down_sample = SCANNETPP_DOWN_SAMPLE + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Only uses iPhone images (not DSLR). + + Args: + scene: Scene identifier (e.g., "09c1414f1b") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_mesh_path, dist, roi, cam_hw, etc. + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + input_path = os.path.join(self.data_root, scene, "merge_dslr_iphone") + colmap_path = os.path.join(input_path, "colmap/sparse_render_rgb") + image_path = os.path.join(input_path, "images") + depth_path_dir = os.path.join(input_path, "render_depth") + + # Read COLMAP model + cams, images, points3d = read_model(colmap_path) + + # Map image names to IDs + name2id = {image.name: k for k, image in images.items()} + names = sorted([image.name for k, image in images.items()]) + # Only use iPhone images + names = [name for name in names if "iphone" in name] + + gt_mesh_path = os.path.join( + input_path.replace("merge_dslr_iphone", "scans"), "mesh_aligned_0.05.ply" + ) + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_mesh_path": gt_mesh_path, + "dist_list": [], + "roi_list": [], + "cam_hw_list": [], + "ixt_raw_list": [], + "gt_depth_files": [], + }), + }) + + for name in names: + image = images[name2id[name]] + img_path = os.path.join(image_path, name) + + if not os.path.exists(img_path): + continue + + # Build extrinsics (world-to-camera) + ext = np.eye(4, dtype=np.float32) + ext[:3, :3] = image.qvec2rotmat() + ext[:3, 3] = image.tvec + + # Get camera parameters + cam_id = image.camera_id + camera = cams[cam_id] + cam_height, cam_width = camera.height, camera.width + + # Build intrinsics + ixt = np.eye(3, dtype=np.float32) + ixt[0, 0], ixt[1, 1], ixt[0, 2], ixt[1, 2] = camera.params[:4] + ixt[:2, 2] -= 0.5 # COLMAP convention adjustment + ixt_raw = ixt.copy() + + # Handle distortion (OPENCV model) + dist = np.zeros(5, dtype=np.float32) + roi = (0, 0, cam_width, cam_height) + if camera.model == "OPENCV": + dist[:4] = camera.params[4:] + ixt, roi = cv2.getOptimalNewCameraMatrix( + ixt, dist, (cam_width, cam_height), 1, (cam_width, cam_height) + ) + + # Depth file path + frame_name = os.path.basename(name)[:-4] # Remove .jpg + depth_file = os.path.join(depth_path_dir, f"{frame_name}.png") + + out.image_files.append(img_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt) + out.aux.dist_list.append(dist) + out.aux.roi_list.append(roi) + out.aux.cam_hw_list.append((cam_height, cam_width)) + out.aux.ixt_raw_list.append(ixt_raw) + out.aux.gt_depth_files.append(depth_file) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[ScanNet++] {scene}: {len(out.image_files)} images") + + self._scene_cache[scene] = out + return out + + def load_image(self, img_path: str, idx: int, aux: Dict) -> np.ndarray: + """ + Load and preprocess image with undistortion and cropping. + + Args: + img_path: Path to image file + idx: Index of the image in the dataset + aux: Auxiliary data from get_data + + Returns: + Preprocessed RGB image + """ + image = imageio.imread(img_path).astype(np.uint8) + ixt_raw = aux.ixt_raw_list[idx] + ixt = aux.intrinsics[idx] if hasattr(aux, 'intrinsics') else None + dist = aux.dist_list[idx] + roi = aux.roi_list[idx] + + # Undistort using raw intrinsics + # Use the stored intrinsics from get_data for newCameraMatrix + stored_ixt = self._scene_cache.get(aux.scene, {}).get('intrinsics', [None])[idx] if hasattr(aux, 'scene') else None + if stored_ixt is None: + # Recompute optimal camera matrix for undistortion + cam_h, cam_w = aux.cam_hw_list[idx] + ixt_for_undistort = ixt_raw.copy() + ixt_for_undistort, _ = cv2.getOptimalNewCameraMatrix( + ixt_raw, dist, (cam_w, cam_h), 1, (cam_w, cam_h) + ) + else: + ixt_for_undistort = stored_ixt + + image = cv2.undistort(image, ixt_raw, dist, newCameraMatrix=ixt_for_undistort) + + # Crop to ROI + x, y, w, h = roi + image = image[y:y+h, x:x+w] + + # Resize to target resolution + image = cv2.resize(image, (self.input_w, self.input_h), interpolation=cv2.INTER_AREA) + + return image + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against ScanNet++ ground truth mesh. + + Uses AABB cropping to only evaluate points within GT bounding box. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_mesh_path = gt_data.aux.gt_mesh_path + + # Load ground truth mesh and sample points + gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path) + gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Crop prediction to GT bounding box (with 0.1m margin) + aabb = gt_pcd.get_axis_aligned_bounding_box() + points = np.asarray(pred_pcd.points) + inside_mask = ( + (points[:, 0] >= aabb.min_bound[0] - 0.1) & + (points[:, 0] <= aabb.max_bound[0] + 0.1) & + (points[:, 1] >= aabb.min_bound[1] - 0.1) & + (points[:, 1] <= aabb.max_bound[1] + 0.1) & + (points[:, 2] >= aabb.min_bound[2] - 0.1) & + (points[:, 2] <= aabb.max_bound[2] + 0.1) + ) + pred_pcd = pred_pcd.select_by_index(inside_mask.nonzero()[0]) + + # Downsample + if self.down_sample > 0: + pred_pcd = pred_pcd.voxel_down_sample(self.down_sample) + gt_pcd = gt_pcd.voxel_down_sample(self.down_sample) + + verts_pred = np.asarray(pred_pcd.points) + verts_gt = np.asarray(gt_pcd.points) + + if len(verts_pred) == 0 or len(verts_gt) == 0: + return { + "acc": float("inf"), + "comp": float("inf"), + "overall": float("inf"), + "precision": 0.0, + "recall": 0.0, + "fscore": 0.0, + } + + # Compute distances + dist_pred_to_gt = nn_correspondance(verts_gt, verts_pred) + dist_gt_to_pred = nn_correspondance(verts_pred, verts_gt) + + # Compute metrics + accuracy = float(np.mean(dist_pred_to_gt)) + completeness = float(np.mean(dist_gt_to_pred)) + overall = (accuracy + completeness) / 2 + + precision = float(np.mean((dist_pred_to_gt < self.eval_threshold).astype(float))) + recall = float(np.mean((dist_gt_to_pred < self.eval_threshold).astype(float))) + + if precision + recall > 0: + fscore = 2 * precision * recall / (precision + recall) + else: + fscore = 0.0 + + return { + "acc": accuracy, + "comp": completeness, + "overall": overall, + "precision": precision, + "recall": recall, + "fscore": fscore, + } + + def _load_gt_meta(self, result_path: str) -> Dict: + """Load saved GT meta for fusion.""" + export_dir = os.path.dirname(result_path) + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + image_files = list(data["image_files"]) + + # Reconstruct aux data from image files + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": image_files, + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Get GT data + full_gt_data = self.get_data(scene) + + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + # Need to rebuild aux from full GT data based on image indices + image_indices = [ + full_gt_data.image_files.index(f) + for f in gt_data.image_files + if f in full_gt_data.image_files + ] + else: + gt_data = full_gt_data + image_indices = list(range(len(full_gt_data.image_files))) + + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load and preprocess images + images = [] + for idx, img_idx in enumerate(image_indices): + img_path = full_gt_data.image_files[img_idx] + image = imageio.imread(img_path).astype(np.uint8) + + # Undistort and crop + ixt_raw = full_gt_data.aux.ixt_raw_list[img_idx] + ixt = full_gt_data.intrinsics[img_idx] + dist = full_gt_data.aux.dist_list[img_idx] + roi = full_gt_data.aux.roi_list[img_idx] + + image = cv2.undistort(image, ixt_raw, dist, newCameraMatrix=ixt) + x, y, w, h = roi + image = image[y:y+h, x:x+w] + image = cv2.resize(image, (self.input_w, self.input_h), interpolation=cv2.INTER_AREA) + + images.append(image) + + images = np.stack(images, axis=0) + + # Prepare depths, intrinsics, extrinsics + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, full_gt_data, image_indices, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, full_gt_data, image_indices, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_unposed mode.""" + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + img_idx = image_indices[i] + + # Get original image size (after undistort+crop, before resize to input_h/w) + orig_h, orig_w = full_gt_data.aux.cam_hw_list[img_idx] + + # Step 1: nearest resize to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Step 2: linear resize to target resolution + depth = cv2.resize( + depth, + (self.input_w, self.input_h), + interpolation=cv2.INTER_LINEAR, + ).astype(np.float32) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(full_gt_data.aux.gt_depth_files[img_idx]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to target resolution + h_ratio = self.input_h / model_h + w_ratio = self.input_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio + ixt[1, :] *= h_ratio + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict, + image_indices: list, scene: str = None + ) -> tuple: + """Prepare depths/intrinsics/extrinsics for recon_posed mode.""" + # Scale alignment + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + depths_out = [] + intrinsics_out = [] + extrinsics_out = [] + + for i in range(len(pred_data.depth)): + img_idx = image_indices[i] + + # Get original image size (after undistort+crop, before resize to input_h/w) + orig_h, orig_w = full_gt_data.aux.cam_hw_list[img_idx] + + # Step 1: nearest resize to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Step 2: linear resize to target resolution + depth = cv2.resize( + depth, + (self.input_w, self.input_h), + interpolation=cv2.INTER_LINEAR, + ).astype(np.float32) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(full_gt_data.aux.gt_depth_files[img_idx]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Get GT intrinsics and scale to target resolution + ixt = full_gt_data.intrinsics[img_idx].copy() + cam_h, cam_w = full_gt_data.aux.cam_hw_list[img_idx] + ixt[:2, 2] += 0.5 # Undo COLMAP convention + ixt[0, :] *= self.input_w / cam_w + ixt[1, :] *= self.input_h / cam_h + intrinsics_out.append(ixt) + + extrinsics_out.append(full_gt_data.extrinsics[img_idx]) + + return np.stack(depths_out), np.stack(intrinsics_out), np.stack(extrinsics_out) + + def _load_gt_mask(self, gt_depth_path: str) -> np.ndarray: + """ + Load GT depth and create valid mask. + + For ScanNet++, GT depth is stored as 16-bit PNG in millimeters. + + Returns: + Boolean mask where True = valid region to keep + """ + if not os.path.exists(gt_depth_path): + return None + + gt_depth = imageio.imread(gt_depth_path) / 1000.0 # mm to meters + + # Resize to target resolution + gt_depth = cv2.resize( + gt_depth, + (self.input_w, self.input_h), + interpolation=cv2.INTER_LINEAR, + ).astype(np.float32) + + # Valid mask: depth > 0 and not inf + valid_mask = np.logical_and(gt_depth > 0, gt_depth != np.inf) + return valid_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """Mask invalid depth values by setting them to 0.""" + depth = depth.copy() + + if gt_zero_mask is not None: + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + diff --git a/src/nn/depth_anything_3/bench/datasets/sevenscenes.py b/src/nn/depth_anything_3/bench/datasets/sevenscenes.py new file mode 100644 index 0000000..cce82b3 --- /dev/null +++ b/src/nn/depth_anything_3/bench/datasets/sevenscenes.py @@ -0,0 +1,449 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +7Scenes Benchmark dataset implementation. + +7Scenes is an indoor RGB-D dataset with ground truth camera poses and 3D meshes. +Reference: https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/ + +Evaluation metrics: +- 3D reconstruction: Accuracy, Completeness, F-score +- Camera pose estimation: AUC metrics +""" + +import os +from typing import Dict as TDict + +import cv2 +import numpy as np +import open3d as o3d +from addict import Dict + +from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready +from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY +from depth_anything_3.bench.utils import ( + create_tsdf_volume, + evaluate_3d_reconstruction, + fuse_depth_to_tsdf, + sample_points_from_mesh, +) +from depth_anything_3.utils.constants import ( + SEVENSCENES_CX, + SEVENSCENES_CY, + SEVENSCENES_DOWN_SAMPLE, + SEVENSCENES_EVAL_DATA_ROOT, + SEVENSCENES_EVAL_THRESHOLD, + SEVENSCENES_FX, + SEVENSCENES_FY, + SEVENSCENES_MAX_DEPTH, + SEVENSCENES_SAMPLING_NUMBER, + SEVENSCENES_SCENES, + SEVENSCENES_SDF_TRUNC, + SEVENSCENES_VOXEL_LENGTH, +) +from depth_anything_3.utils.pose_align import align_poses_umeyama + + +@MV_REGISTRY.register(name="7scenes") +@MONO_REGISTRY.register(name="7scenes") +class SevenScenes(Dataset): + """ + 7Scenes Benchmark dataset wrapper for DepthAnything3 evaluation. + + Supports: + - Camera pose estimation evaluation (AUC metrics) + - 3D reconstruction evaluation (Accuracy, Completeness, F-score) + - TSDF-based point cloud fusion + + Dataset structure: + 7scenes/ + ├── 7Scenes/ + │ ├── {scene}/ + │ │ └── seq-01/ (or seq-02 for stairs) + │ │ ├── frame-XXXXXX.color.png + │ │ ├── frame-XXXXXX.depth.png + │ │ └── frame-XXXXXX.pose.txt + │ └── meshes/ + │ └── {scene}.ply # Ground truth mesh + """ + + data_root = SEVENSCENES_EVAL_DATA_ROOT + SCENES = SEVENSCENES_SCENES + + # Evaluation hyperparameters from constants + max_depth = SEVENSCENES_MAX_DEPTH + sampling_number = SEVENSCENES_SAMPLING_NUMBER + voxel_length = SEVENSCENES_VOXEL_LENGTH + sdf_trunc = SEVENSCENES_SDF_TRUNC + eval_threshold = SEVENSCENES_EVAL_THRESHOLD + down_sample = SEVENSCENES_DOWN_SAMPLE + + # Fixed camera intrinsics for all 7Scenes images + fx = SEVENSCENES_FX + fy = SEVENSCENES_FY + cx = SEVENSCENES_CX + cy = SEVENSCENES_CY + + def __init__(self): + super().__init__() + self._scene_cache = {} + + # ------------------------------ + # Public API + # ------------------------------ + + def get_data(self, scene: str) -> Dict: + """ + Collect per-view image paths, intrinsics/extrinsics for a scene. + + Args: + scene: Scene identifier (e.g., "chess") + + Returns: + Dict with: + - image_files: List[str] - paths to images + - extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms + - intrinsics: np.ndarray [N, 3, 3] - camera intrinsics + - aux: Dict with gt_mesh_path, gt_depth_files + """ + if scene in self._scene_cache: + return self._scene_cache[scene] + + # Different sequence for stairs scene + if scene == "stairs": + data_folder = os.path.join(self.data_root, "7Scenes", scene, "seq-02") + n_imgs = 500 + else: + data_folder = os.path.join(self.data_root, "7Scenes", scene, "seq-01") + n_imgs = 1000 + + gt_mesh_path = os.path.join(self.data_root, "7Scenes", "meshes", f"{scene}.ply") + + # Fixed intrinsics for all images + ixt = np.array([ + [self.fx, 0, self.cx], + [0, self.fy, self.cy], + [0, 0, 1], + ], dtype=np.float32) + + out = Dict({ + "image_files": [], + "extrinsics": [], + "intrinsics": [], + "aux": Dict({ + "gt_mesh_path": gt_mesh_path, + "gt_depth_files": [], + }), + }) + + for i in range(0, n_imgs, 1): + img_path = os.path.join(data_folder, f"frame-{i:06d}.color.png") + pose_path = os.path.join(data_folder, f"frame-{i:06d}.pose.txt") + depth_path = os.path.join(data_folder, f"frame-{i:06d}.depth.png") + + if not os.path.exists(img_path) or not os.path.exists(pose_path): + continue + + # Load camera-to-world pose and convert to world-to-camera (extrinsic) + c2w = np.loadtxt(pose_path) + ext = np.linalg.inv(c2w).astype(np.float32) + + out.image_files.append(img_path) + out.extrinsics.append(ext) + out.intrinsics.append(ixt.copy()) + out.aux.gt_depth_files.append(depth_path) + + out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32) + out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32) + + print(f"[7Scenes] {scene}: {len(out.image_files)} images") + + self._scene_cache[scene] = out + return out + + def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]: + """ + Evaluate fused point cloud against 7Scenes ground truth mesh. + + Args: + scene: Scene identifier + fuse_path: Path to fused point cloud (.ply) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + gt_data = self.get_data(scene) + gt_mesh_path = gt_data.aux.gt_mesh_path + + # Load and sample ground truth mesh + gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path) + gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number) + + # Load predicted point cloud + pred_pcd = o3d.io.read_point_cloud(fuse_path) + + # Evaluate using shared utility function + metrics = evaluate_3d_reconstruction( + pred_pcd, + gt_pcd, + threshold=self.eval_threshold, + down_sample=self.down_sample, + ) + + return metrics + + def _load_gt_meta(self, result_path: str) -> Dict: + """ + Load saved GT meta (extrinsics, intrinsics, image_files) for fusion. + + This is needed when frames are sampled, so fuse3d uses the correct + (sampled) GT instead of full dataset GT. + + Args: + result_path: Path to npz file (used to derive gt_meta.npz path) + + Returns: + Dict with GT data, or None if gt_meta.npz doesn't exist + """ + export_dir = os.path.dirname(result_path) # exports/mini_npz/ + gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz") + + if os.path.exists(gt_meta_path): + data = np.load(gt_meta_path, allow_pickle=True) + # Build aux with gt_depth_files derived from image_files + image_files = list(data["image_files"]) + gt_depth_files = [ + img_path.replace("color", "depth").replace(".color.", ".depth.") + for img_path in image_files + ] + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + "image_files": image_files, + "aux": Dict({"gt_depth_files": gt_depth_files}), + }) + return None + + def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None: + """ + Fuse per-view depths into a point cloud using TSDF fusion. + + Args: + scene: Scene identifier + result_path: Path to npz file with predicted depths/poses + fuse_path: Output path for fused point cloud (.ply) + mode: "recon_unposed" or "recon_posed" + """ + # Try to load saved GT meta (handles frame sampling) + gt_meta = self._load_gt_meta(result_path) + if gt_meta is not None: + gt_data = gt_meta + else: + gt_data = self.get_data(scene) + _wait_for_file_ready(result_path) + pred_data = Dict({k: v for k, v in np.load(result_path).items()}) + + # Load original images (keep original size) + images = [] + orig_sizes = [] + for img_path in gt_data.image_files: + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + images.append(img) + orig_sizes.append((img.shape[0], img.shape[1])) + + # Prepare depths, intrinsics, extrinsics + if mode == "recon_unposed": + depths, intrinsics, extrinsics = self._prep_unposed( + pred_data, gt_data, orig_sizes, scene=scene + ) + elif mode == "recon_posed": + depths, intrinsics, extrinsics = self._prep_posed( + pred_data, gt_data, orig_sizes, scene=scene + ) + else: + raise ValueError(f"Invalid mode: {mode}") + + images = np.stack(images, axis=0) + + # Create TSDF volume and fuse + volume = create_tsdf_volume( + voxel_length=self.voxel_length, + sdf_trunc=self.sdf_trunc, + ) + mesh = fuse_depth_to_tsdf( + volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth + ) + + # Sample points from mesh + pcd = sample_points_from_mesh(mesh, self.sampling_number) + + # Save point cloud + os.makedirs(os.path.dirname(fuse_path), exist_ok=True) + o3d.io.write_point_cloud(fuse_path, pcd) + + # ------------------------------ + # Private helpers + # ------------------------------ + + def _prep_unposed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_unposed mode. + + Similar to ETH3D but uses GT depth for masking instead of separate mask files. + """ + # Scale alignment with fixed random_state for reproducibility + _, _, scale, extrinsics = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + intrinsics_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + + # Resize depth to original image size (nearest interpolation) + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(gt_data.aux.gt_depth_files[i]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + # Adjust intrinsics to original image size + h_ratio = orig_h / model_h + w_ratio = orig_w / model_w + ixt = pred_data.intrinsics[i].copy() + ixt[0, :] *= w_ratio + ixt[1, :] *= h_ratio + + depths_out.append(depth) + intrinsics_out.append(ixt) + + return np.stack(depths_out), np.stack(intrinsics_out), extrinsics + + def _prep_posed( + self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str + ) -> tuple: + """ + Prepare depths/intrinsics/extrinsics for recon_posed mode. + Uses GT intrinsics/extrinsics but aligns depth scale via Umeyama. + """ + # Scale alignment with fixed random_state + _, _, scale, _ = align_poses_umeyama( + gt_data.extrinsics.copy(), + pred_data.extrinsics.copy(), + return_aligned=True, + ransac=True, + random_state=42, + ) + + model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2] + + depths_out = [] + for i in range(len(pred_data.depth)): + orig_h, orig_w = orig_sizes[i] + + # Resize depth to original image size + depth = cv2.resize( + pred_data.depth[i], + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ) + + # Load GT depth for masking + gt_zero_mask = self._load_gt_mask(gt_data.aux.gt_depth_files[i]) + + # Mask invalid depths BEFORE scale + depth = self._mask_invalid_depth(depth, gt_zero_mask) + + # Apply scale AFTER mask + depth = depth * scale + + depths_out.append(depth) + + # Use GT intrinsics and extrinsics + return np.stack(depths_out), gt_data.intrinsics.copy(), gt_data.extrinsics.copy() + + def _load_gt_mask(self, gt_depth_path: str) -> np.ndarray: + """ + Load GT depth and create valid mask. + + For 7Scenes, GT depth is stored as 16-bit PNG in millimeters. + Value 65535 indicates invalid depth. + + Returns: + Boolean mask where True = valid region to keep + """ + if not os.path.exists(gt_depth_path): + return None + + gt_depth = cv2.imread(gt_depth_path, -1) + if gt_depth is None: + return None + + # 65535 is invalid depth marker in 7Scenes + gt_depth[gt_depth == 65535] = 0 + # Convert to meters + gt_depth = gt_depth / 1000.0 + + # Valid mask: depth > 0 + valid_mask = gt_depth > 0 + return valid_mask + + def _mask_invalid_depth( + self, depth: np.ndarray, gt_zero_mask: np.ndarray = None + ) -> np.ndarray: + """ + Mask invalid depth values by setting them to 0. + + Args: + depth: Depth map to mask + gt_zero_mask: Optional GT mask (True = valid region) + + Returns: + Masked depth map with invalid regions set to 0 + """ + depth = depth.copy() + + if gt_zero_mask is not None: + # Also mask out invalid pred depth + pred_invalid = np.isnan(depth) | np.isinf(depth) + combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid)) + depth = depth * combined_mask.astype(np.float32) + else: + # Fallback: only mask pred invalid values + invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0) + depth[invalid_mask] = 0.0 + + return depth + + diff --git a/src/nn/depth_anything_3/bench/evaluator.py b/src/nn/depth_anything_3/bench/evaluator.py new file mode 100644 index 0000000..5b93b98 --- /dev/null +++ b/src/nn/depth_anything_3/bench/evaluator.py @@ -0,0 +1,752 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Main Evaluator class for DepthAnything3 benchmark evaluation. + +Supports multiple datasets and evaluation modes: +- pose: Camera pose estimation (AUC metrics) +- recon_unposed: 3D reconstruction with predicted poses +- recon_posed: 3D reconstruction with GT poses +- view_syn: Novel view synthesis (TODO) +""" + +import json +import os +import random +from typing import Dict as TDict, Iterable, List + +import numpy as np +import torch +from addict import Dict +from tqdm import tqdm + +from depth_anything_3.bench.print_metrics import MetricsPrinter +from depth_anything_3.utils.parallel_utils import parallel_execution +from depth_anything_3.bench.registries import MV_REGISTRY +from depth_anything_3.utils.constants import EVAL_REF_VIEW_STRATEGY + + +class Evaluator: + """ + Main evaluation orchestrator for DepthAnything3 benchmarks. + + Usage: + evaluator = Evaluator( + work_dir="./eval_workspace", + datas=["dtu"], + modes=["pose", "recon_unposed", "recon_posed"], + ) + api = DepthAnything3.from_pretrained("...") + evaluator.infer(api) + metrics = evaluator.eval() + evaluator.print_metrics() + """ + + VALID_MODES = {"pose", "recon_unposed", "recon_posed", "view_syn"} + + def __init__( + self, + work_dir: str = "./eval_workspace", + datas: List[str] = ("dtu",), + modes: List[str] = ("recon_unposed",), + ref_view_strategy: str = EVAL_REF_VIEW_STRATEGY, + scenes: List[str] = None, + debug: bool = False, + num_fusion_workers: int = 4, + max_frames: int = 100, + gpu_id: int = 0, + total_gpus: int = 1, + ): + """ + Initialize the evaluator. + + Args: + work_dir: Base directory for model outputs and metric files + datas: List of dataset names (must be registered in MV_REGISTRY) + modes: List of evaluation modes to run + ref_view_strategy: Reference view selection strategy for inference + ("first", "saddle_balanced", etc.) + scenes: Specific scenes to evaluate (None = all scenes) + debug: Enable verbose debug output + num_fusion_workers: Number of parallel workers for TSDF fusion (default: 4) + max_frames: Maximum number of frames per scene (default: 100). + If a scene has more frames, randomly sample to this limit. + Set to -1 to disable sampling. + gpu_id: GPU index for multi-GPU (0-indexed) + total_gpus: Total number of GPUs for task distribution + """ + self.work_dir = work_dir + self.datas = list(datas) + self.modes = set(modes) + self.ref_view_strategy = ref_view_strategy + self.scenes_filter = scenes + self.debug = debug + self.num_fusion_workers = num_fusion_workers + self.max_frames = max_frames + self.gpu_id = gpu_id + self.total_gpus = total_gpus + + # Validate modes + unknown = self.modes - self.VALID_MODES + if unknown: + raise ValueError(f"Unknown modes: {unknown}. Valid: {sorted(self.VALID_MODES)}") + + os.makedirs(self.work_dir, exist_ok=True) + + # Initialize datasets + self.datasets = Dict() + for data in self.datas: + if not MV_REGISTRY.has(data): + available = list(MV_REGISTRY.all().keys()) + raise ValueError(f"Dataset '{data}' not found. Available: {available}") + self.datasets[data] = MV_REGISTRY.get(data)() + + # Initialize metrics printer + self._printer = MetricsPrinter() + + # -------------------- Public APIs -------------------- # + + def all(self, api) -> TDict[str, dict]: + """ + Run complete evaluation pipeline: inference + evaluation. + + Args: + api: DepthAnything3 API instance + + Returns: + Combined metrics dictionary + """ + self.infer(api) + return self.eval() + + def _get_scenes(self, dataset) -> List[str]: + """Get list of scenes to evaluate, optionally filtered.""" + all_scenes = dataset.SCENES + if self.scenes_filter: + scenes = [s for s in all_scenes if s in self.scenes_filter] + if self.debug: + print(f"[DEBUG] Filtered scenes: {scenes} (from {len(all_scenes)} total)") + return scenes + return all_scenes + + def infer(self, api, model_path: str = None) -> None: + """ + Run inference according to requested modes. + + - Unposed export if 'pose' or 'recon_unposed' is in modes + - Posed export if 'recon_posed' or 'view_syn' is in modes + + Multi-GPU: Use --gpu_id and --total_gpus to distribute tasks. + Example: Launch 4 processes with gpu_id=0,1,2,3 and total_gpus=4 + + Args: + api: DepthAnything3 API instance + model_path: Model path (unused, kept for API compatibility) + """ + need_unposed = {"pose", "recon_unposed"} & self.modes + need_posed = {"recon_posed", "view_syn"} & self.modes + export_format = "mini_npz-glb" if self.debug else "mini_npz" + + # Collect all tasks + all_tasks = [] + for data in self.datas: + dataset = self.datasets[data] + for scene in self._get_scenes(dataset): + all_tasks.append((data, scene)) + + # Distribute tasks across GPUs + if self.total_gpus > 1: + tasks = [t for i, t in enumerate(all_tasks) if i % self.total_gpus == self.gpu_id] + print(f"[INFO] GPU {self.gpu_id}/{self.total_gpus}: {len(tasks)}/{len(all_tasks)} tasks") + else: + tasks = all_tasks + print(f"[INFO] Total inference tasks: {len(tasks)}") + + for data, scene in tqdm(tasks, desc=f"Inference (GPU {self.gpu_id})"): + dataset = self.datasets[data] + scene_data = dataset.get_data(scene) + scene_data = self._sample_frames(scene_data, scene) + + if need_unposed: + export_dir = self._export_dir(data, scene, posed=False) + api.inference( + scene_data.image_files, + export_dir=export_dir, + export_format=export_format, + ref_view_strategy=self.ref_view_strategy, + ) + self._save_gt_meta(export_dir, scene_data) + + if need_posed: + export_dir = self._export_dir(data, scene, posed=True) + api.inference( + scene_data.image_files, + scene_data.extrinsics, + scene_data.intrinsics, + export_dir=export_dir, + export_format=export_format, + ref_view_strategy=self.ref_view_strategy, + ) + self._save_gt_meta(export_dir, scene_data) + + def eval(self) -> TDict[str, dict]: + """ + Evaluate for all configured modes and write JSON files. + + Evaluation order by mode (all datasets per mode): + 1. pose - all datasets + 2. recon_unposed - all datasets + 3. recon_posed - all datasets + + Returns: + Summary mapping: {"_": metrics_dict} + """ + summary: TDict[str, dict] = {} + + # Evaluate by mode (all datasets per mode) + if "pose" in self.modes: + print(f"\n{'='*60}") + print(f"📊 Evaluating POSE for all datasets...") + print(f"{'='*60}") + for data, result in self._eval_pose(): + summary[f"{data}_pose"] = result + + if "recon_unposed" in self.modes: + print(f"\n{'='*60}") + print(f"📊 Evaluating RECON_UNPOSED for all datasets...") + print(f"{'='*60}") + for data, result in self._eval_reconstruction("recon_unposed"): + summary[f"{data}_recon_unposed"] = result + + if "recon_posed" in self.modes: + print(f"\n{'='*60}") + print(f"📊 Evaluating RECON_POSED for all datasets...") + print(f"{'='*60}") + for data, result in self._eval_reconstruction("recon_posed"): + summary[f"{data}_recon_posed"] = result + + if "view_syn" in self.modes: + # TODO: Add view synthesis metrics here when available + pass + + return summary + + def print_metrics(self, metrics: TDict[str, dict] = None) -> None: + """ + Print evaluation metrics in a beautiful tabular format. + + Args: + metrics: Metrics dictionary. If None, loads from saved JSON files. + """ + if metrics is None: + metrics = self._load_metrics() + + self._printer.print_results(metrics) + + # -------------------- Evaluation Methods -------------------- # + + def _eval_pose(self) -> Iterable[tuple]: + """Compute pose-estimation metrics for each dataset and scene.""" + os.makedirs(self._metric_dir, exist_ok=True) + + for data in tqdm(self.datas, desc="Datasets (pose eval)"): + dataset = self.datasets[data] + dataset_results = Dict() + scenes = self._get_scenes(dataset) + + for scene in tqdm(scenes, desc=f"{data} scenes", leave=False): + export_dir = self._export_dir(data, scene, posed=False) + result_path = os.path.join(export_dir, "exports", "mini_npz", "results.npz") + + # Check if result file exists and is valid + if not os.path.exists(result_path): + print(f"\n[ERROR] Result file not found: {result_path}") + print(f"[ERROR] CWD: {os.getcwd()}") + print(f"[ERROR] Please run inference first (remove --eval_only)") + continue + + try: + # Use saved GT meta (handles frame sampling correctly) + gt_meta = self._load_gt_meta(export_dir) + if gt_meta is not None: + result = self._compute_pose_with_gt(result_path, gt_meta) + else: + # Fallback to dataset GT (no sampling was done) + result = dataset.eval_pose(scene, result_path) + dataset_results[scene] = self._to_float_dict(result) + except Exception as e: + print(f"\n[ERROR] Failed to evaluate pose for {data}/{scene}: {e}") + print(f"[ERROR] File path: {os.path.abspath(result_path)}") + if self.debug: + import traceback + traceback.print_exc() + continue + + if not dataset_results: + print(f"[WARNING] No valid results for {data}") + continue + + dataset_results["mean"] = self._mean_of_dicts(dataset_results.values()) + out_path = os.path.join(self._metric_dir, f"{data}_pose.json") + self._dump_json(out_path, dataset_results) + yield data, dataset_results + + def _eval_reconstruction(self, mode: str) -> Iterable[tuple]: + """ + Compute reconstruction metrics for each dataset and scene. + + Args: + mode: "recon_unposed" or "recon_posed" + """ + assert mode in {"recon_unposed", "recon_posed"} + os.makedirs(self._metric_dir, exist_ok=True) + + posed_flag = mode == "recon_posed" + + # Filter out datasets that don't support reconstruction (e.g., dtu64) + recon_datas = [d for d in self.datas if d != "dtu64"] + + for data in tqdm(recon_datas, desc=f"Datasets ({mode} eval)"): + dataset = self.datasets[data] + dataset_results = Dict() + scenes = self._get_scenes(dataset) + + # Prepare paths for all scenes + scene_list = [] + result_paths = [] + fuse_paths = [] + for scene in scenes: + export_dir = self._export_dir(data, scene, posed=posed_flag) + result_path = os.path.join(export_dir, "exports", "mini_npz", "results.npz") + fuse_path = os.path.join(export_dir, "exports", "fuse", "pcd.ply") + scene_list.append(scene) + result_paths.append(result_path) + fuse_paths.append(fuse_path) + + # Parallel fusion (default 4 workers) + # DTU uses CUDA operations in fusion, which doesn't work well with ThreadPool + use_sequential = (data == "dtu") + parallel_execution( + scene_list, + result_paths, + fuse_paths, + action=lambda s, rp, fp: dataset.fuse3d(s, rp, fp, mode), + num_processes=self.num_fusion_workers, + print_progress=True, + desc=f"{data} fusion", + sequential=use_sequential, + ) + + # Sequential evaluation (fast, no need to parallelize) + for scene, fuse_path in zip(scene_list, fuse_paths): + # DTU supports CPU-based evaluation + if data == "dtu" and hasattr(dataset, "eval3d"): + result = dataset.eval3d(scene, fuse_path) + else: + result = dataset.eval3d(scene, fuse_path) + dataset_results[scene] = self._to_float_dict(result) + print(f" {mode} | {data} | {scene}: {result}") + + dataset_results["mean"] = self._mean_of_dicts(dataset_results.values()) + out_path = os.path.join(self._metric_dir, f"{data}_{mode}.json") + self._dump_json(out_path, dataset_results) + yield data, dataset_results + + # -------------------- Helpers -------------------- # + + def _save_gt_meta(self, export_dir: str, scene_data: Dict) -> None: + """ + Save GT extrinsics/intrinsics/image_files for evaluation. + + This is needed when frames are sampled, so eval_pose and fuse3d can use + the correct (sampled) GT instead of full dataset GT. + + Args: + export_dir: Export directory for the scene + scene_data: Sampled scene data + """ + meta_path = os.path.join(export_dir, "exports", "gt_meta.npz") + os.makedirs(os.path.dirname(meta_path), exist_ok=True) + np.savez_compressed( + meta_path, + extrinsics=scene_data.extrinsics, + intrinsics=scene_data.intrinsics, + image_files=np.array(scene_data.image_files, dtype=object), + ) + + def _load_gt_meta(self, export_dir: str) -> Dict: + """ + Load saved GT extrinsics/intrinsics for evaluation. + + Returns: + Dict with extrinsics and intrinsics, or None if not found + """ + meta_path = os.path.join(export_dir, "exports", "gt_meta.npz") + if os.path.exists(meta_path): + data = np.load(meta_path) + return Dict({ + "extrinsics": data["extrinsics"], + "intrinsics": data["intrinsics"], + }) + return None + + def _compute_pose_with_gt(self, result_path: str, gt_meta: Dict) -> TDict[str, float]: + """ + Compute pose metrics using saved GT meta (handles frame sampling). + + Args: + result_path: Path to npz with predicted extrinsics + gt_meta: Dict with GT extrinsics from saved meta + + Returns: + Dict with pose metrics + """ + from depth_anything_3.bench.dataset import _wait_for_file_ready + from depth_anything_3.bench.utils import compute_pose + from depth_anything_3.utils.geometry import as_homogeneous + + _wait_for_file_ready(result_path) + pred = np.load(result_path) + return compute_pose( + torch.from_numpy(as_homogeneous(pred["extrinsics"])), + torch.from_numpy(as_homogeneous(gt_meta["extrinsics"])), + ) + + def _sample_frames(self, scene_data: Dict, scene: str) -> Dict: + """ + Sample frames if scene has more than max_frames. + + Uses fixed random seed (42) for reproducibility. + + Args: + scene_data: Scene data dict with image_files, extrinsics, intrinsics, aux + scene: Scene name (for logging) + + Returns: + Sampled scene_data if num_frames > max_frames, otherwise original + """ + if self.max_frames <= 0: + return scene_data + + num_frames = len(scene_data.image_files) + if num_frames <= self.max_frames: + return scene_data + + # Sample with fixed seed for reproducibility + random.seed(42) + indices = list(range(num_frames)) + random.shuffle(indices) + sampled_indices = sorted(indices[:self.max_frames]) + + print(f" [Sampling] {scene}: {num_frames} -> {self.max_frames} frames") + + # Create new scene_data with sampled frames + sampled = Dict() + sampled.image_files = [scene_data.image_files[i] for i in sampled_indices] + sampled.extrinsics = scene_data.extrinsics[sampled_indices] + sampled.intrinsics = scene_data.intrinsics[sampled_indices] + + # Copy aux data, sampling lists if needed + sampled.aux = Dict() + for key, val in scene_data.aux.items(): + if isinstance(val, list) and len(val) == num_frames: + sampled.aux[key] = [val[i] for i in sampled_indices] + elif isinstance(val, np.ndarray) and len(val) == num_frames: + sampled.aux[key] = val[sampled_indices] + else: + sampled.aux[key] = val + + return sampled + + @property + def _metric_dir(self) -> str: + """Directory for storing metric JSON files.""" + return os.path.join(self.work_dir, "metric_results") + + def _export_dir(self, data: str, scene: str, posed: bool) -> str: + """ + Get export directory path. + + Structure: .../model_results/{data}/{scene}/{posed|unposed} + """ + suffix = "posed" if posed else "unposed" + export_dir = os.path.join(self.work_dir, "model_results", data, scene, suffix) + os.makedirs(export_dir, exist_ok=True) + return export_dir + + @staticmethod + def _to_float_dict(d: TDict[str, float]) -> dict: + """Convert numpy scalars to plain Python floats for JSON safety.""" + return {k: float(v) for k, v in d.items()} + + @staticmethod + def _mean_of_dicts(dicts: Iterable[dict]) -> dict: + """Compute elementwise mean across a list of homogeneous metric dicts.""" + dicts = list(dicts) + if not dicts: + return {} + keys = dicts[0].keys() + return {k: float(np.mean([d[k] for d in dicts]).item()) for k in keys} + + @staticmethod + def _dump_json(path: str, obj: dict, indent: int = 4) -> None: + """Write JSON with UTF-8 and pretty indentation.""" + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(obj, f, indent=indent, ensure_ascii=False) + + def _load_metrics(self) -> TDict[str, dict]: + """Load evaluation metrics from JSON files.""" + metrics = {} + metric_dir = self._metric_dir + + if not os.path.exists(metric_dir): + return metrics + + for filename in os.listdir(metric_dir): + if filename.endswith(".json"): + filepath = os.path.join(metric_dir, filename) + try: + with open(filepath, encoding="utf-8") as f: + data = json.load(f) + key = filename[:-5] # Remove .json extension + metrics[key] = data + except Exception as e: + print(f"Warning: Failed to read metrics file: {filename} - {e}") + + return metrics + + +# -------------------- CLI Entry Point -------------------- # + + +if __name__ == "__main__": + import sys + from omegaconf import OmegaConf + from depth_anything_3.cfg import load_config + + # Get default config path (relative to this file) + _default_config = os.path.join( + os.path.dirname(__file__), "configs", "eval_bench.yaml" + ) + + # Check for help flag first (we need to handle this before OmegaConf) + if "--help" in sys.argv or "-h" in sys.argv: + pass # Will handle after config loading + + # Set up argv for OmegaConf processing + argv = sys.argv[1:] + + # Check if user provides custom config + config_path = _default_config + if "--config" in argv: + config_idx = argv.index("--config") + if config_idx + 1 < len(argv): + config_path = argv[config_idx + 1] + # Remove --config and its value + argv = argv[:config_idx] + argv[config_idx + 2:] + + # Print help if requested + if "--help" in sys.argv or "-h" in sys.argv: + print(""" +DepthAnything3 Benchmark Evaluation + +Usage: + python -m depth_anything_3.bench.evaluator [OPTIONS] [KEY=VALUE ...] + +Configuration: + --config PATH Config YAML file (default: bench/configs/eval_bench.yaml) + +Config Overrides (using dotlist notation): + model.path=VALUE Model path or HuggingFace ID + workspace.work_dir=VALUE Working directory for outputs + eval.datasets=[dataset1,dataset2] Datasets to evaluate (eth3d,7scenes,scannetpp,hiroom,dtu,dtu64) + eval.modes=[mode1,mode2] Evaluation modes (pose,recon_unposed,recon_posed) + eval.scenes=[scene1,scene2] Specific scenes to evaluate (null=all) + eval.max_frames=VALUE Max frames per scene (-1=no limit, default: 100) + eval.ref_view_strategy=VALUE Reference view strategy (default: first) + eval.eval_only=VALUE Only run evaluation (skip inference) (true/false) + eval.print_only=VALUE Only print saved metrics (true/false) + inference.num_fusion_workers=VALUE Number of parallel workers (default: 4) + inference.debug=VALUE Enable debug mode (true/false) + +Special Flags: + --help, -h Show this help message + +Multi-GPU: + Use CUDA_VISIBLE_DEVICES to specify GPUs (auto-detected and distributed) + +Examples: + # Use default config + python -m depth_anything_3.bench.evaluator + + # Override model path + python -m depth_anything_3.bench.evaluator model.path=depth-anything/DA3-LARGE + + # Evaluate specific datasets and modes + python -m depth_anything_3.bench.evaluator \\ + eval.datasets=[eth3d,hiroom] \\ + eval.modes=[pose] + + # Use custom config with overrides + python -m depth_anything_3.bench.evaluator \\ + --config my_config.yaml \\ + model.path=/path/to/model \\ + eval.max_frames=50 + + # Multi-GPU inference (auto-distributed) + CUDA_VISIBLE_DEVICES=0,1,2,3 python -m depth_anything_3.bench.evaluator + + # Debug specific scenes + python -m depth_anything_3.bench.evaluator \\ + eval.datasets=[eth3d] \\ + eval.scenes=[courtyard] \\ + inference.debug=true + + # Only evaluate (skip inference) + python -m depth_anything_3.bench.evaluator eval.eval_only=true + + # Only print saved metrics + python -m depth_anything_3.bench.evaluator eval.print_only=true + + """) + sys.exit(0) + + # Load config with CLI overrides using OmegaConf dotlist + # Example: python evaluator.py model.path=/path/to/model eval.datasets=[eth3d,dtu] + config = load_config(config_path, argv=argv) + + # Extract config values + work_dir = config.workspace.work_dir + model_path = config.model.path + datasets = config.eval.datasets + modes = config.eval.modes + ref_view_strategy = config.eval.ref_view_strategy + scenes = config.eval.scenes + max_frames = config.eval.max_frames + eval_only = config.eval.eval_only + print_only = config.eval.print_only + debug = config.inference.debug + num_fusion_workers = config.inference.num_fusion_workers + + # GPU settings: parse from CLI dotlist args (gpu_id=X total_gpus=Y) + # These are passed by the main process when spawning workers + gpu_id = 0 + total_gpus = 1 + for arg in argv: + if arg.startswith("gpu_id="): + gpu_id = int(arg.split("=")[1]) + elif arg.startswith("total_gpus="): + total_gpus = int(arg.split("=")[1]) + + # Override dataset scenes if specified + if scenes: + print(f"[INFO] Running on specific scenes: {scenes}") + + evaluator = Evaluator( + work_dir=work_dir, + datas=datasets, + modes=modes, + ref_view_strategy=ref_view_strategy, + scenes=scenes, + debug=debug, + num_fusion_workers=num_fusion_workers, + max_frames=max_frames, + gpu_id=gpu_id, + total_gpus=total_gpus, + ) + + if print_only: + evaluator.print_metrics() + elif eval_only: + metrics = evaluator.eval() + evaluator.print_metrics(metrics) + else: + # Parse CUDA_VISIBLE_DEVICES to get GPU list + # If not set, use all available GPUs + cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if cuda_devices is not None and cuda_devices.strip(): + gpu_list = [g.strip() for g in cuda_devices.split(",") if g.strip()] + else: + # CUDA_VISIBLE_DEVICES not set, use all available GPUs + num_available = torch.cuda.device_count() + gpu_list = [str(i) for i in range(num_available)] if num_available > 0 else ["0"] + + # Auto multi-GPU: if multiple GPUs and not a worker process + is_worker = os.environ.get("_DA3_WORKER") == "1" + + if len(gpu_list) > 1 and not is_worker: + # Launch worker processes + import subprocess + + num_gpus = len(gpu_list) + print(f"[INFO] Detected {num_gpus} GPUs: {gpu_list}") + print(f"[INFO] Launching {num_gpus} workers...") + + # Build base command + base_cmd = [sys.executable, "-m", "depth_anything_3.bench.evaluator"] + # Pass config via dotlist instead of CLI args + if config_path != _default_config: + base_cmd += ["--config", config_path] + base_cmd += [f"model.path={model_path}"] + base_cmd += [f"workspace.work_dir={work_dir}"] + base_cmd += [f"eval.datasets=[{','.join(datasets)}]"] + base_cmd += [f"eval.modes=[{','.join(modes)}]"] + if scenes: + base_cmd += [f"eval.scenes=[{','.join(scenes)}]"] + base_cmd += [f"eval.max_frames={max_frames}"] + base_cmd += [f"eval.ref_view_strategy={ref_view_strategy}"] + base_cmd += [f"inference.debug={str(debug).lower()}"] + base_cmd += [f"inference.num_fusion_workers={num_fusion_workers}"] + + # Launch workers + processes = [] + for idx, gpu_id in enumerate(gpu_list): + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = gpu_id + env["_DA3_WORKER"] = "1" # Mark as worker process + + cmd = base_cmd.copy() + # GPU-specific worker config + cmd += [f"gpu_id={idx}", f"total_gpus={num_gpus}"] + + print(f"[INFO] Starting worker {idx} on GPU {gpu_id}") + p = subprocess.Popen(cmd, env=env) + processes.append(p) + + # Wait for all workers + for p in processes: + p.wait() + + print(f"[INFO] All {num_gpus} workers completed") + + # Run evaluation after all inference is done + metrics = evaluator.eval() + evaluator.print_metrics(metrics) + else: + # Single GPU or worker process + from depth_anything_3.api import DepthAnything3 + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + api = DepthAnything3.from_pretrained(model_path) + api = api.to(device) + + evaluator.infer(api, model_path=model_path) + + # Only run eval if single GPU mode (workers don't eval) + if not is_worker: + metrics = evaluator.eval() + evaluator.print_metrics(metrics) + diff --git a/src/nn/depth_anything_3/bench/print_metrics.py b/src/nn/depth_anything_3/bench/print_metrics.py new file mode 100644 index 0000000..9964f3b --- /dev/null +++ b/src/nn/depth_anything_3/bench/print_metrics.py @@ -0,0 +1,618 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Beautiful metrics printing utilities for benchmark evaluation. + +Provides colorized, well-formatted tabular output for evaluation results. +Supports highlighting best/worst values and grouping by dataset/mode. +""" + +import argparse +import json +import os +import re +from typing import Dict as TDict, List, Optional + + +# ANSI color codes for terminal output +class Colors: + """ANSI escape codes for terminal colors.""" + + RESET = "\033[0m" + BOLD = "\033[1m" + RED = "\033[31m" + GREEN = "\033[32m" + YELLOW = "\033[33m" + BLUE = "\033[34m" + MAGENTA = "\033[35m" + CYAN = "\033[36m" + WHITE = "\033[37m" + + # Bold variants + BOLD_RED = "\033[1;31m" + BOLD_GREEN = "\033[1;32m" + BOLD_YELLOW = "\033[1;33m" + BOLD_BLUE = "\033[1;34m" + BOLD_MAGENTA = "\033[1;35m" + BOLD_CYAN = "\033[1;36m" + + # Background + BG_DARK = "\033[48;5;236m" + + +def strip_ansi(text: str) -> str: + """Remove ANSI escape sequences from string for length calculation.""" + ansi_escape = re.compile(r"\x1b\[[0-9;]*m") + return ansi_escape.sub("", text) + + +def colorize_value( + value: str, + is_best: bool = False, + is_worst: bool = False, + lower_is_better: bool = False, +) -> str: + """ + Apply color to a metric value based on whether it's best/worst. + + Args: + value: String representation of the value + is_best: Whether this is the best value in its column + is_worst: Whether this is the worst value in its column + lower_is_better: If True, lower values are better (e.g., error metrics) + + Returns: + Colorized string + """ + if lower_is_better: + # For metrics like error/distance, lower is better + if is_best: + return f"{Colors.BOLD_GREEN}{value}{Colors.RESET}" + elif is_worst: + return f"{Colors.BOLD_RED}{value}{Colors.RESET}" + else: + # For metrics like accuracy/AUC, higher is better + if is_best: + return f"{Colors.BOLD_GREEN}{value}{Colors.RESET}" + elif is_worst: + return f"{Colors.BOLD_RED}{value}{Colors.RESET}" + return value + + +class MetricsPrinter: + """ + Beautiful tabular metrics printer with color support. + + Features: + - Colorized best/worst values + - Grouped by dataset and evaluation mode + - Automatic column width calculation + - Support for multiple input directories comparison + """ + + # Metrics where lower values are better + LOWER_IS_BETTER = {"comp", "acc", "overall", "error", "loss", "rmse", "mae"} + + def __init__(self, use_color: bool = True): + """ + Initialize the printer. + + Args: + use_color: Whether to use ANSI colors in output + """ + self.use_color = use_color + + def print_results(self, metrics: TDict[str, dict], summary_only: bool = True) -> None: + """ + Print evaluation metrics in a beautiful tabular format. + + Args: + metrics: Dictionary mapping "dataset_mode" to metric results + summary_only: If True, only print summary table. If False, print per-dataset details too. + """ + if not metrics: + print(f"\n{Colors.BOLD_RED}❌ No evaluation metrics found{Colors.RESET}") + return + + if not summary_only: + self._print_header() + grouped = self._group_by_dataset(metrics) + + for dataset, modes_data in grouped.items(): + self._print_dataset_section(dataset, modes_data) + + # Print summary table with average metrics across datasets + self._print_summary(metrics) + + self._print_footer() + + def print_comparison( + self, + metrics_list: List[TDict[str, dict]], + labels: List[str], + ) -> None: + """ + Print comparison table for multiple evaluation runs. + + Args: + metrics_list: List of metrics dictionaries + labels: Labels for each metrics dictionary + """ + if not metrics_list or not all(metrics_list): + print(f"\n{Colors.BOLD_RED}❌ No metrics to compare{Colors.RESET}") + return + + # Collect all datasets and modes + all_keys = set() + for metrics in metrics_list: + all_keys.update(metrics.keys()) + + self._print_header("COMPARISON") + + for key in sorted(all_keys): + parts = key.rsplit("_", 1) + if len(parts) == 2: + dataset, mode = parts[0], parts[1] + else: + dataset, mode = key, "unknown" + + print(f"\n{Colors.BOLD_CYAN}📊 {dataset.upper()} - {mode.upper()}{Colors.RESET}") + print("-" * 100) + + # Collect metrics from all runs + all_metric_names = set() + for metrics in metrics_list: + if key in metrics and "mean" in metrics[key]: + all_metric_names.update(metrics[key]["mean"].keys()) + + if not all_metric_names: + continue + + # Build comparison table + metric_width = max(15, max(len(m) for m in all_metric_names) + 2) + label_width = max(15, max(len(l) for l in labels) + 2) + + # Header + header = f"{'Metric':<{metric_width}}" + for label in labels: + header += f"{label:<{label_width}}" + print(header) + print("-" * len(strip_ansi(header))) + + # Collect values for highlighting + for metric_name in sorted(all_metric_names): + values = [] + for metrics in metrics_list: + if key in metrics and "mean" in metrics[key]: + val = metrics[key]["mean"].get(metric_name) + values.append(val if val is not None else float("nan")) + else: + values.append(float("nan")) + + # Find best/worst + valid_values = [v for v in values if not (v != v)] # Filter NaN + if valid_values: + lower_better = any( + lb in metric_name.lower() for lb in self.LOWER_IS_BETTER + ) + best_val = min(valid_values) if lower_better else max(valid_values) + worst_val = max(valid_values) if lower_better else min(valid_values) + else: + best_val = worst_val = None + + # Print row + row = f"{metric_name:<{metric_width}}" + for val in values: + if val != val: # NaN check + val_str = "N/A" + else: + val_str = f"{val:.4f}" + if self.use_color and len(valid_values) > 1: + lower_better = any( + lb in metric_name.lower() for lb in self.LOWER_IS_BETTER + ) + is_best = abs(val - best_val) < 1e-8 if best_val else False + is_worst = abs(val - worst_val) < 1e-8 if worst_val else False + val_str_padded = f"{val_str:<{label_width}}" + val_str = colorize_value( + val_str_padded, is_best, is_worst, lower_better + ) + row += val_str + continue + row += f"{val_str:<{label_width}}" + print(row) + + self._print_footer() + + def _print_header(self, title: str = "EVALUATION RESULTS") -> None: + """Print report header.""" + width = 100 + print() + print("=" * width) + print(f"{Colors.BOLD_CYAN}📊 DEPTH ANYTHING 3 {title}{Colors.RESET}") + print("=" * width) + + def _print_footer(self) -> None: + """Print report footer.""" + width = 100 + print() + print("=" * width) + print(f"{Colors.BOLD_GREEN}✅ Evaluation Complete{Colors.RESET}") + print("=" * width) + print() + + def _group_by_dataset(self, metrics: TDict[str, dict]) -> TDict[str, dict]: + """Group metrics by dataset.""" + grouped = {} + for key, data in metrics.items(): + if not isinstance(data, dict) or "mean" not in data: + continue + # Parse key format: "dataset_mode" (e.g., "dtu_recon_unposed") + parts = key.split("_", 1) + if len(parts) == 2: + dataset, mode = parts + if dataset not in grouped: + grouped[dataset] = {} + grouped[dataset][mode] = data + return grouped + + def _print_dataset_section(self, dataset: str, modes_data: TDict[str, dict]) -> None: + """Print metrics section for a single dataset.""" + print(f"\n{Colors.BOLD_MAGENTA}🔍 {dataset.upper()}{Colors.RESET}") + print("-" * 100) + + # Collect all unique metrics across all modes + all_metrics = set() + for mode_data in modes_data.values(): + all_metrics.update(mode_data["mean"].keys()) + all_metrics = sorted(list(all_metrics)) + + if not all_metrics: + print(" No metrics available") + return + + # Calculate column widths + metric_width = max(18, max(len(m) for m in all_metrics) + 2) + mode_width = 18 + modes = list(modes_data.keys()) + + # Print header + header = f"{'Metric':<{metric_width}}" + for mode in modes: + header += f"{mode.upper():<{mode_width}}" + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(header)) + + # Print each metric row + for metric in all_metrics: + row = f"{metric:<{metric_width}}" + + # Collect values for this metric across modes + values = [] + for mode in modes: + if metric in modes_data[mode]["mean"]: + values.append(modes_data[mode]["mean"][metric]) + else: + values.append(None) + + # Find best/worst values + valid_values = [v for v in values if v is not None] + if valid_values: + lower_better = any(lb in metric.lower() for lb in self.LOWER_IS_BETTER) + best_val = min(valid_values) if lower_better else max(valid_values) + worst_val = max(valid_values) if lower_better else min(valid_values) + else: + best_val = worst_val = None + + # Format each value + for val in values: + if val is None: + row += f"{'N/A':<{mode_width}}" + else: + val_str = f"{val:.4f}" + if self.use_color and len(valid_values) > 1: + is_best = abs(val - best_val) < 1e-8 if best_val else False + is_worst = abs(val - worst_val) < 1e-8 if worst_val else False + lower_better = any( + lb in metric.lower() for lb in self.LOWER_IS_BETTER + ) + # Pad before colorizing to maintain alignment + val_str_padded = f"{val_str:<{mode_width}}" + row += colorize_value( + val_str_padded, is_best, is_worst, lower_better + ) + else: + row += f"{val_str:<{mode_width}}" + print(row) + + # Show scene counts + scene_info = [] + for mode, mode_data in modes_data.items(): + scene_count = len([k for k in mode_data.keys() if k != "mean"]) + scene_info.append(f"{mode}: {scene_count} scenes") + print(f"\n{Colors.CYAN}📈 {' | '.join(scene_info)}{Colors.RESET}") + + def _print_summary(self, metrics: TDict[str, dict]) -> None: + """ + Print summary table with key metrics across all datasets. + + Format: One row per metric, datasets as columns. + Order: HiRoom, ETH3D, DTU, 7Scenes, ScanNet++, (DTU-64 for pose only) + """ + print(f"\n{Colors.BOLD_CYAN}{'=' * 120}{Colors.RESET}") + print(f"{Colors.BOLD_CYAN}📊 SUMMARY{Colors.RESET}") + print(f"{Colors.BOLD_CYAN}{'=' * 120}{Colors.RESET}") + + # Dataset display order and names + DATASET_ORDER = ["hiroom", "eth3d", "dtu", "7scenes", "scannetpp", "dtu64"] + DATASET_DISPLAY = { + "hiroom": "HiRoom", + "eth3d": "ETH3D", + "dtu": "DTU", + "7scenes": "7Scenes", + "scannetpp": "ScanNet++", + "dtu64": "DTU-64", + } + + # Collect all metrics into a structured dict + # metric_data[dataset][mode] = {"Auc_3": x, "Auc_30": x, "fscore": x, "overall": x} + metric_data = {} + for key, data in metrics.items(): + if not isinstance(data, dict) or "mean" not in data: + continue + parts = key.split("_", 1) + if len(parts) != 2: + continue + dataset, mode = parts + dataset_lower = dataset.lower() + if dataset_lower not in metric_data: + metric_data[dataset_lower] = {} + metric_data[dataset_lower][mode] = data["mean"] + + col_width = 12 + + def fmt_val(val): + """Format value or return N/A.""" + if val is None: + return "N/A" + return f"{val:.4f}" + + def get_metric(dataset, mode, metric_name): + """Get metric value or None.""" + if dataset not in metric_data: + return None + if mode not in metric_data[dataset]: + return None + return metric_data[dataset][mode].get(metric_name) + + # ============ POSE METRICS ============ + print(f"\n{Colors.BOLD_MAGENTA}🎯 POSE ESTIMATION{Colors.RESET}") + + # Pose: show all datasets except DTU (keep DTU-64 only) + # Order: HiRoom, ETH3D, DTU-64, 7Scenes, ScanNet++ + pose_datasets = ["hiroom", "eth3d", "dtu64", "7scenes", "scannetpp"] + + # Header: Avg first, then datasets + header = f"{'Metric':<15}{'Avg':<{col_width}}" + for ds in pose_datasets: + header += f"{DATASET_DISPLAY[ds]:<{col_width}}" + print("-" * len(strip_ansi(header))) + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(strip_ansi(header))) + + # Helper to get metric with fallback names + def get_pose_metric(dataset, metric_name): + """Get pose metric with fallback for different naming conventions.""" + # Try different naming conventions + names = { + "Auc3": ["Auc_3", "auc03", "auc_3", "AUC_3", "Auc3", "auc3"], + "Auc30": ["Auc_30", "auc30", "auc_30", "AUC_30", "Auc30"], + } + for name in names.get(metric_name, [metric_name]): + val = get_metric(dataset, "pose", name) + if val is not None: + return val + return None + + # Auc3 row + values = [] + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc3") + if val is not None: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Auc3':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc3") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # Auc30 row + values = [] + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc30") + if val is not None: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Auc30':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in pose_datasets: + val = get_pose_metric(ds, "Auc30") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # ============ RECON_UNPOSED METRICS ============ + print(f"\n{Colors.BOLD_MAGENTA}🏗️ RECON_UNPOSED (Pred Pose){Colors.RESET}") + + # For recon, exclude dtu64 from columns + recon_datasets = ["hiroom", "eth3d", "dtu", "7scenes", "scannetpp"] + avg_datasets = ["hiroom", "eth3d", "7scenes", "scannetpp"] # Exclude DTU from avg + + # Header: Avg first, then datasets + header = f"{'Metric':<15}{'Avg*':<{col_width}}" + for ds in recon_datasets: + header += f"{DATASET_DISPLAY[ds]:<{col_width}}" + print("-" * len(strip_ansi(header))) + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(strip_ansi(header))) + + # F-score row (only metric for avg) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "fscore") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'F-score':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "fscore") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # Overall row (avg over 4 datasets excluding DTU) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "overall") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Overall':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_unposed", "overall") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # ============ RECON_POSED METRICS ============ + print(f"\n{Colors.BOLD_MAGENTA}🏗️ RECON_POSED (GT Pose){Colors.RESET}") + + # Header: Avg first, then datasets + header = f"{'Metric':<15}{'Avg*':<{col_width}}" + for ds in recon_datasets: + header += f"{DATASET_DISPLAY[ds]:<{col_width}}" + print("-" * len(strip_ansi(header))) + print(f"{Colors.BOLD}{header}{Colors.RESET}") + print("-" * len(strip_ansi(header))) + + # F-score row (only metric for avg) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "fscore") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'F-score':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "fscore") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + # Overall row (avg over 4 datasets excluding DTU) + values = [] + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "overall") + if val is not None and ds in avg_datasets: + values.append(val) + avg = sum(values) / len(values) if values else None + row = f"{'Overall':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}" + for ds in recon_datasets: + val = get_metric(ds, "recon_posed", "overall") + row += f"{fmt_val(val):<{col_width}}" + print(row) + + print(f"\n{Colors.CYAN}* Avg F-score / Overall = average over HiRoom, ETH3D, 7Scenes, ScanNet++ (4 datasets){Colors.RESET}") + + +def load_metrics_from_dir(metric_dir: str) -> TDict[str, dict]: + """ + Load all metrics JSON files from a directory. + + Args: + metric_dir: Path to directory containing metric JSON files + + Returns: + Dictionary mapping filename (without .json) to metric data + """ + metrics = {} + if not os.path.exists(metric_dir): + return metrics + + for filename in os.listdir(metric_dir): + if filename.endswith(".json"): + filepath = os.path.join(metric_dir, filename) + try: + with open(filepath, encoding="utf-8") as f: + content = f.read() + # Handle trailing commas in JSON + content = re.sub(r",\s*([\]\}])", r"\1", content) + data = json.loads(content) + key = filename[:-5] + metrics[key] = data + except Exception as e: + print(f"Warning: Failed to load {filename}: {e}") + + return metrics + + +def main(): + """Command-line interface for metrics printing.""" + parser = argparse.ArgumentParser( + description="Print DepthAnything3 benchmark evaluation metrics." + ) + parser.add_argument( + "--input_dir", + type=str, + default="./eval_workspace/metric_results", + help="Directory containing metric JSON files (comma-separated for comparison)", + ) + parser.add_argument( + "--no_color", + action="store_true", + help="Disable colored output", + ) + parser.add_argument( + "--key", + type=str, + default=None, + help="Specific metric key to highlight", + ) + args = parser.parse_args() + + # Support multiple directories for comparison + input_dirs = [d.strip() for d in args.input_dir.split(",") if d.strip()] + + printer = MetricsPrinter(use_color=not args.no_color) + + if len(input_dirs) == 1: + # Single directory - simple print + metrics = load_metrics_from_dir(input_dirs[0]) + printer.print_results(metrics) + else: + # Multiple directories - comparison mode + metrics_list = [] + labels = [] + for d in input_dirs: + metrics = load_metrics_from_dir(d) + if metrics: + metrics_list.append(metrics) + labels.append(os.path.basename(d.rstrip("/"))) + + if metrics_list: + printer.print_comparison(metrics_list, labels) + else: + print("No metrics found in specified directories") + + +if __name__ == "__main__": + main() + diff --git a/src/nn/depth_anything_3/bench/registries.py b/src/nn/depth_anything_3/bench/registries.py new file mode 100644 index 0000000..0744661 --- /dev/null +++ b/src/nn/depth_anything_3/bench/registries.py @@ -0,0 +1,85 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Auto-loading registry system for benchmark datasets. + +This module provides registry classes that automatically discover and import +dataset implementations from the datasets subpackage on first access. +""" + +import importlib +import pkgutil +import threading + +from depth_anything_3.utils.registry import Registry + +__all__ = ["METRIC_REGISTRY", "MONO_REGISTRY", "MV_REGISTRY", "NVS_REGISTRY"] + +# ---- Lazy import: Only scan and import all datasets submodules on first registry access ---- +_loaded = False +_lock = threading.Lock() + + +def _import_all_datasets_once(): + """ + Scan and import all .py submodules under depth_anything_3.bench.datasets + (skip files/packages starting with underscore), to trigger @REGISTRY.register(...) in each module. + """ + global _loaded + if _loaded: + return + + with _lock: + if _loaded: + return + + pkg_name = "depth_anything_3.bench.datasets" + pkg = importlib.import_module(pkg_name) + pkg_paths = list(getattr(pkg, "__path__", [])) + + for finder, name, ispkg in pkgutil.walk_packages(pkg_paths, prefix=pkg_name + "."): + base = name.rsplit(".", 1)[-1] + if base.startswith("_"): + continue + try: + importlib.import_module(name) + except Exception as e: + print(f"[datasets auto-import] Failed to import {name}: {e}") + + _loaded = True + + +class AutoRegistry(Registry): + """Registry that ensures all datasets are auto-discovered and imported on first use.""" + + def get(self, name): + _import_all_datasets_once() + return super().get(name) + + def all(self): + _import_all_datasets_once() + return super().all() + + def has(self, name): + _import_all_datasets_once() + return name in self._map + + +# Four auto-lazy registry instances for different evaluation types +METRIC_REGISTRY = AutoRegistry() # For metric depth evaluation +MONO_REGISTRY = AutoRegistry() # For monocular depth evaluation +MV_REGISTRY = AutoRegistry() # For multi-view evaluation +NVS_REGISTRY = AutoRegistry() # For novel view synthesis evaluation + diff --git a/src/nn/depth_anything_3/bench/utils.py b/src/nn/depth_anything_3/bench/utils.py new file mode 100644 index 0000000..ca65477 --- /dev/null +++ b/src/nn/depth_anything_3/bench/utils.py @@ -0,0 +1,525 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utility functions for benchmark evaluation. + +Contains: +- Pose evaluation metrics (AUC) and helper functions +- 3D reconstruction evaluation metrics (Acc/Comp/F-score) +- Geometry utilities (quaternion conversion, etc.) +""" + +from typing import Dict as TDict, Optional, Tuple, Union + +import numpy as np +import open3d as o3d +import torch +from addict import Dict +from scipy.spatial import KDTree + +from depth_anything_3.utils.geometry import mat_to_quat + + +# ============================================================================= +# Geometry Utilities +# ============================================================================= + + +def quat2rotmat(qvec: list) -> np.ndarray: + """ + Convert quaternion (WXYZ order) to rotation matrix. + + Args: + qvec: Quaternion as [w, x, y, z] + + Returns: + 3x3 rotation matrix + """ + rotmat = np.array( + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ] + ) + rotmat = rotmat.reshape(3, 3) + return rotmat + + +# ============================================================================= +# 3D Reconstruction Evaluation +# ============================================================================= + + +def nn_correspondance(verts1: np.ndarray, verts2: np.ndarray) -> np.ndarray: + """ + Compute nearest neighbor distances from verts2 to verts1 using KDTree. + + Args: + verts1: Reference point cloud [N, 3] + verts2: Query point cloud [M, 3] + + Returns: + Distance array [M,] - distance from each point in verts2 to nearest in verts1 + """ + if len(verts1) == 0 or len(verts2) == 0: + return np.array([]) + + kdtree = KDTree(verts1) + distances, _ = kdtree.query(verts2) + return distances.reshape(-1) + + +def evaluate_3d_reconstruction( + pcd_pred: Union[o3d.geometry.PointCloud, np.ndarray], + pcd_trgt: Union[o3d.geometry.PointCloud, np.ndarray], + threshold: float = 0.05, + down_sample: Optional[float] = None, +) -> TDict[str, float]: + """ + Evaluate 3D reconstruction quality using standard metrics. + + This function computes: + - Accuracy: Mean distance from predicted points to GT surface + - Completeness: Mean distance from GT points to predicted surface + - Overall: Average of accuracy and completeness + - Precision: Fraction of predicted points within threshold of GT + - Recall: Fraction of GT points within threshold of prediction + - F-score: Harmonic mean of precision and recall + + Args: + pcd_pred: Predicted point cloud (Open3D or numpy array) + pcd_trgt: Ground truth point cloud (Open3D or numpy array) + threshold: Distance threshold for precision/recall (meters) + down_sample: Voxel size for downsampling (None to skip) + + Returns: + Dict with metrics: acc, comp, overall, precision, recall, fscore + """ + # Convert to Open3D if needed + if isinstance(pcd_pred, np.ndarray): + pcd_pred_o3d = o3d.geometry.PointCloud() + pcd_pred_o3d.points = o3d.utility.Vector3dVector(pcd_pred) + pcd_pred = pcd_pred_o3d + if isinstance(pcd_trgt, np.ndarray): + pcd_trgt_o3d = o3d.geometry.PointCloud() + pcd_trgt_o3d.points = o3d.utility.Vector3dVector(pcd_trgt) + pcd_trgt = pcd_trgt_o3d + + # Downsample if requested + if down_sample is not None and down_sample > 0: + pcd_pred = pcd_pred.voxel_down_sample(down_sample) + pcd_trgt = pcd_trgt.voxel_down_sample(down_sample) + + verts_pred = np.asarray(pcd_pred.points) + verts_trgt = np.asarray(pcd_trgt.points) + + # Handle empty point clouds + if len(verts_pred) == 0 or len(verts_trgt) == 0: + return { + "acc": float("inf"), + "comp": float("inf"), + "overall": float("inf"), + "precision": 0.0, + "recall": 0.0, + "fscore": 0.0, + } + + # Compute distances + dist_pred_to_gt = nn_correspondance(verts_trgt, verts_pred) # Accuracy + dist_gt_to_pred = nn_correspondance(verts_pred, verts_trgt) # Completeness + + # Compute metrics + accuracy = float(np.mean(dist_pred_to_gt)) + completeness = float(np.mean(dist_gt_to_pred)) + overall = (accuracy + completeness) / 2 + + precision = float(np.mean((dist_pred_to_gt < threshold).astype(float))) + recall = float(np.mean((dist_gt_to_pred < threshold).astype(float))) + + if precision + recall > 0: + fscore = 2 * precision * recall / (precision + recall) + else: + fscore = 0.0 + + return { + "acc": accuracy, + "comp": completeness, + "overall": overall, + "precision": precision, + "recall": recall, + "fscore": fscore, + } + + +def create_tsdf_volume( + voxel_length: float = 4.0 / 512.0, + sdf_trunc: float = 0.04, + color_type: str = "RGB8", +) -> o3d.pipelines.integration.ScalableTSDFVolume: + """ + Create a scalable TSDF volume for depth fusion. + + Args: + voxel_length: Size of each voxel + sdf_trunc: Truncation distance for SDF + color_type: Color integration type ("RGB8" or "Gray32") + + Returns: + Initialized ScalableTSDFVolume + """ + if color_type == "RGB8": + color_enum = o3d.pipelines.integration.TSDFVolumeColorType.RGB8 + else: + color_enum = o3d.pipelines.integration.TSDFVolumeColorType.Gray32 + + volume = o3d.pipelines.integration.ScalableTSDFVolume( + voxel_length=voxel_length, + sdf_trunc=sdf_trunc, + color_type=color_enum, + ) + return volume + + +def fuse_depth_to_tsdf( + volume: o3d.pipelines.integration.ScalableTSDFVolume, + depths: np.ndarray, + images: np.ndarray, + intrinsics: np.ndarray, + extrinsics: np.ndarray, + max_depth: float = 10.0, +) -> o3d.geometry.TriangleMesh: + """ + Fuse multiple depth maps into TSDF volume and extract mesh. + + Args: + volume: TSDF volume to integrate into + depths: Depth maps [N, H, W] + images: RGB images [N, H, W, 3] + intrinsics: Camera intrinsics [N, 3, 3] + extrinsics: Camera extrinsics (world-to-camera) [N, 4, 4] + max_depth: Maximum depth for truncation + + Returns: + Extracted triangle mesh + """ + for i in range(len(depths)): + depth = depths[i] + image = images[i] + ixt = intrinsics[i] + ext = extrinsics[i] + + h, w = depth.shape[:2] + + # Create RGBD image + depth_o3d = o3d.geometry.Image(depth.astype(np.float32)) + color_o3d = o3d.geometry.Image(image.astype(np.uint8)) + rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth( + color_o3d, + depth_o3d, + depth_trunc=max_depth, + convert_rgb_to_intensity=False, + depth_scale=1.0, + ) + + # Create camera intrinsics + ixt_o3d = o3d.camera.PinholeCameraIntrinsic( + w, h, ixt[0, 0], ixt[1, 1], ixt[0, 2], ixt[1, 2] + ) + + # Integrate into volume + volume.integrate(rgbd, ixt_o3d, ext) + + # Extract mesh + mesh = volume.extract_triangle_mesh() + return mesh + + +def sample_points_from_mesh( + mesh: o3d.geometry.TriangleMesh, + num_points: int = 1000000, +) -> o3d.geometry.PointCloud: + """ + Uniformly sample points from a triangle mesh. + + Args: + mesh: Input triangle mesh + num_points: Number of points to sample + + Returns: + Sampled point cloud + """ + try: + pcd = mesh.sample_points_uniformly(number_of_points=num_points) + # Clamp colors to valid range [0, 1] for Open3D PLY export + if pcd.has_colors(): + colors = np.asarray(pcd.colors) + colors = np.clip(colors, 0.0, 1.0) + pcd.colors = o3d.utility.Vector3dVector(colors) + except Exception: + # Fallback: create random points if mesh is invalid (with fixed seed for reproducibility) + rng = np.random.default_rng(seed=42) + points = rng.uniform(-1, 1, size=(num_points, 3)) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + return pcd + + +# ============================================================================= +# Pose Evaluation +# ============================================================================= + + +def build_pair_index(N: int, B: int = 1): + """ + Build indices for all possible pairs of frames. + + Args: + N: Number of frames + B: Batch size + + Returns: + i1, i2: Indices for all possible pairs + """ + i1_, i2_ = torch.combinations(torch.arange(N), 2, with_replacement=False).unbind(-1) + i1, i2 = ((i[None] + torch.arange(B)[:, None] * N).reshape(-1) for i in [i1_, i2_]) + return i1, i2 + + +def compute_pose(pred_se3: torch.Tensor, gt_se3: torch.Tensor) -> Dict: + """ + Compute pose estimation metrics between predicted and ground truth trajectories. + + Args: + pred_se3: Predicted SE(3) transformations [N, 4, 4] + gt_se3: Ground truth SE(3) transformations [N, 4, 4] + + Returns: + Dict with AUC metrics at different thresholds (auc30, auc15, auc05, auc03) + """ + pred_se3 = align_to_first_camera(pred_se3) + gt_se3 = align_to_first_camera(gt_se3) + + rel_rangle_deg, rel_tangle_deg = se3_to_relative_pose_error(pred_se3, gt_se3, len(pred_se3)) + rError = rel_rangle_deg.cpu().numpy() + tError = rel_tangle_deg.cpu().numpy() + + output = Dict() + output.auc30, _ = calculate_auc_np(rError, tError, max_threshold=30) + output.auc15, _ = calculate_auc_np(rError, tError, max_threshold=15) + output.auc05, _ = calculate_auc_np(rError, tError, max_threshold=5) + output.auc03, _ = calculate_auc_np(rError, tError, max_threshold=3) + return output + + +def align_to_first_camera(camera_poses: torch.Tensor) -> torch.Tensor: + """ + Align all camera poses to the first camera's coordinate frame. + + Args: + camera_poses: Camera poses as SE3 transformations [N, 4, 4] + + Returns: + Aligned camera poses [N, 4, 4] + """ + first_cam_extrinsic_inv = closed_form_inverse_se3(camera_poses[0][None]) + aligned_poses = torch.matmul(camera_poses, first_cam_extrinsic_inv) + return aligned_poses + + +def rotation_angle( + rot_gt: torch.Tensor, rot_pred: torch.Tensor, batch_size: int = None, eps: float = 1e-15 +) -> torch.Tensor: + """ + Calculate rotation angle error between ground truth and predicted rotations. + + Args: + rot_gt: Ground truth rotation matrices + rot_pred: Predicted rotation matrices + batch_size: Batch size for reshaping the result + eps: Small value to avoid numerical issues + + Returns: + Rotation angle error in degrees + """ + q_pred = mat_to_quat(rot_pred) + q_gt = mat_to_quat(rot_gt) + + loss_q = (1 - (q_pred * q_gt).sum(dim=1) ** 2).clamp(min=eps) + err_q = torch.arccos(1 - 2 * loss_q) + + rel_rangle_deg = err_q * 180 / np.pi + + if batch_size is not None: + rel_rangle_deg = rel_rangle_deg.reshape(batch_size, -1) + + return rel_rangle_deg + + +def translation_angle( + tvec_gt: torch.Tensor, + tvec_pred: torch.Tensor, + batch_size: int = None, + ambiguity: bool = True, +) -> torch.Tensor: + """ + Calculate translation angle error between ground truth and predicted translations. + + Args: + tvec_gt: Ground truth translation vectors + tvec_pred: Predicted translation vectors + batch_size: Batch size for reshaping the result + ambiguity: Whether to handle direction ambiguity + + Returns: + Translation angle error in degrees + """ + rel_tangle_deg = compare_translation_by_angle(tvec_gt, tvec_pred) + rel_tangle_deg = rel_tangle_deg * 180.0 / np.pi + + if ambiguity: + rel_tangle_deg = torch.min(rel_tangle_deg, (180 - rel_tangle_deg).abs()) + + if batch_size is not None: + rel_tangle_deg = rel_tangle_deg.reshape(batch_size, -1) + + return rel_tangle_deg + + +def compare_translation_by_angle( + t_gt: torch.Tensor, t: torch.Tensor, eps: float = 1e-15, default_err: float = 1e6 +) -> torch.Tensor: + """ + Normalize the translation vectors and compute the angle between them. + + Args: + t_gt: Ground truth translation vectors + t: Predicted translation vectors + eps: Small value to avoid division by zero + default_err: Default error value for invalid cases + + Returns: + Angular error between translation vectors in radians + """ + t_norm = torch.norm(t, dim=1, keepdim=True) + t = t / (t_norm + eps) + + t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True) + t_gt = t_gt / (t_gt_norm + eps) + + loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps) + err_t = torch.acos(torch.sqrt(1 - loss_t)) + + err_t[torch.isnan(err_t) | torch.isinf(err_t)] = default_err + return err_t + + +def calculate_auc_np( + r_error: np.ndarray, t_error: np.ndarray, max_threshold: int = 30 +) -> tuple: + """ + Calculate the Area Under the Curve (AUC) for the given error arrays. + + Args: + r_error: Rotation error values in degrees + t_error: Translation error values in degrees + max_threshold: Maximum threshold value for binning + + Returns: + Tuple of (AUC value, normalized histogram) + """ + error_matrix = np.concatenate((r_error[:, None], t_error[:, None]), axis=1) + max_errors = np.max(error_matrix, axis=1) + bins = np.arange(max_threshold + 1) + histogram, _ = np.histogram(max_errors, bins=bins) + num_pairs = float(len(max_errors)) + normalized_histogram = histogram.astype(float) / num_pairs + return np.mean(np.cumsum(normalized_histogram)), normalized_histogram + + +def se3_to_relative_pose_error( + pred_se3: torch.Tensor, gt_se3: torch.Tensor, num_frames: int +) -> tuple: + """ + Compute rotation and translation errors between predicted and ground truth poses. + + Args: + pred_se3: Predicted SE(3) transformations + gt_se3: Ground truth SE(3) transformations + num_frames: Number of frames + + Returns: + Tuple of (rotation angle errors, translation angle errors) in degrees + """ + pair_idx_i1, pair_idx_i2 = build_pair_index(num_frames) + + # Compute relative camera poses between pairs using closed-form inverse + relative_pose_gt = closed_form_inverse_se3(gt_se3[pair_idx_i1]).bmm(gt_se3[pair_idx_i2]) + relative_pose_pred = closed_form_inverse_se3(pred_se3[pair_idx_i1]).bmm(pred_se3[pair_idx_i2]) + + # Compute the difference in rotation and translation + rel_rangle_deg = rotation_angle(relative_pose_gt[:, :3, :3], relative_pose_pred[:, :3, :3]) + rel_tangle_deg = translation_angle(relative_pose_gt[:, :3, 3], relative_pose_pred[:, :3, 3]) + + return rel_rangle_deg, rel_tangle_deg + + +def closed_form_inverse_se3( + se3: torch.Tensor, R: torch.Tensor = None, T: torch.Tensor = None +) -> torch.Tensor: + """ + Compute the inverse of each 4x4 (or 3x4) SE3 matrix in a batch. + + Uses closed-form solution instead of torch.inverse() for numerical stability. + + Args: + se3: Nx4x4 or Nx3x4 tensor of SE3 matrices + R: Optional Nx3x3 rotation matrices + T: Optional Nx3x1 translation vectors + + Returns: + Inverted SE3 matrices with same shape as input + """ + is_numpy = isinstance(se3, np.ndarray) + + if se3.shape[-2:] != (4, 4) and se3.shape[-2:] != (3, 4): + raise ValueError(f"se3 must be of shape (N,4,4), got {se3.shape}.") + + if R is None: + R = se3[:, :3, :3] + if T is None: + T = se3[:, :3, 3:] + + if is_numpy: + R_transposed = np.transpose(R, (0, 2, 1)) + top_right = -np.matmul(R_transposed, T) + inverted_matrix = np.tile(np.eye(4), (len(R), 1, 1)) + else: + R_transposed = R.transpose(1, 2) + top_right = -torch.bmm(R_transposed, T) + inverted_matrix = torch.eye(4, 4)[None].repeat(len(R), 1, 1) + inverted_matrix = inverted_matrix.to(R.dtype).to(R.device) + + inverted_matrix[:, :3, :3] = R_transposed + inverted_matrix[:, :3, 3:] = top_right + + return inverted_matrix + diff --git a/src/nn/depth_anything_3/cfg.py b/src/nn/depth_anything_3/cfg.py new file mode 100644 index 0000000..4607ff8 --- /dev/null +++ b/src/nn/depth_anything_3/cfg.py @@ -0,0 +1,144 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Configuration utility functions +""" + +import importlib +from pathlib import Path +from typing import Any, Callable, List, Union +from omegaconf import DictConfig, ListConfig, OmegaConf + +try: + OmegaConf.register_new_resolver("eval", eval) +except Exception as e: + # if eval is not available, we can just pass + print(f"Error registering eval resolver: {e}") + + +def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]: + """ + Load a configuration. Will resolve inheritance. + Supports both file paths and module paths (e.g., depth_anything_3.configs.giant). + """ + # Check if path is a module path (contains dots but no slashes and doesn't end with .yaml) + if "." in path and "/" not in path and not path.endswith(".yaml"): + # It's a module path, load from package resources + path_parts = path.split(".")[1:] + config_path = Path(__file__).resolve().parent + for part in path_parts: + config_path = config_path.joinpath(part) + config_path = config_path.with_suffix(".yaml") + config = OmegaConf.load(str(config_path)) + else: + # It's a file path (absolute, relative, or with .yaml extension) + config = OmegaConf.load(path) + + if argv is not None: + config_argv = OmegaConf.from_dotlist(argv) + config = OmegaConf.merge(config, config_argv) + config = resolve_recursive(config, resolve_inheritance) + return config + + +def resolve_recursive( + config: Any, + resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]], +) -> Any: + config = resolver(config) + if isinstance(config, DictConfig): + for k in config.keys(): + v = config.get(k) + if isinstance(v, (DictConfig, ListConfig)): + config[k] = resolve_recursive(v, resolver) + if isinstance(config, ListConfig): + for i in range(len(config)): + v = config.get(i) + if isinstance(v, (DictConfig, ListConfig)): + config[i] = resolve_recursive(v, resolver) + return config + + +def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any: + """ + Recursively resolve inheritance if the config contains: + __inherit__: path/to/parent.yaml or a ListConfig of such paths. + """ + if isinstance(config, DictConfig): + inherit = config.pop("__inherit__", None) + + if inherit: + inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit] + + parent_config = None + for parent_path in inherit_list: + assert isinstance(parent_path, str) + parent_config = ( + load_config(parent_path) + if parent_config is None + else OmegaConf.merge(parent_config, load_config(parent_path)) + ) + + if len(config.keys()) > 0: + config = OmegaConf.merge(parent_config, config) + else: + config = parent_config + return config + + +def import_item(path: str, name: str) -> Any: + """ + Import a python item. Example: import_item("path.to.file", "MyClass") -> MyClass + """ + return getattr(importlib.import_module(path), name) + + +def create_object(config: DictConfig) -> Any: + """ + Create an object from config. + The config is expected to contains the following: + __object__: + path: path.to.module + name: MyClass + args: as_config | as_params (default to as_config) + """ + config = DictConfig(config) + item = import_item( + path=config.__object__.path, + name=config.__object__.name, + ) + args = config.__object__.get("args", "as_config") + if args == "as_config": + return item(config) + if args == "as_params": + config = OmegaConf.to_object(config) + config.pop("__object__") + return item(**config) + raise NotImplementedError(f"Unknown args type: {args}") + + +def create_dataset(path: str, *args, **kwargs) -> Any: + """ + Create a dataset. Requires the file to contain a "create_dataset" function. + """ + return import_item(path, "create_dataset")(*args, **kwargs) + + +def to_dict_recursive(config_obj): + if isinstance(config_obj, DictConfig): + return {k: to_dict_recursive(v) for k, v in config_obj.items()} + elif isinstance(config_obj, ListConfig): + return [to_dict_recursive(item) for item in config_obj] + return config_obj diff --git a/src/nn/depth_anything_3/cli.py b/src/nn/depth_anything_3/cli.py new file mode 100644 index 0000000..c946302 --- /dev/null +++ b/src/nn/depth_anything_3/cli.py @@ -0,0 +1,803 @@ +# flake8: noqa: E402 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Refactored Depth Anything 3 CLI +Clean, modular command-line interface +""" + +from __future__ import annotations + +import os +import typer + +from depth_anything_3.services import start_server +from depth_anything_3.services.gallery import gallery as gallery_main +from depth_anything_3.services.inference_service import run_inference +from depth_anything_3.services.input_handlers import ( + ColmapHandler, + ImageHandler, + ImagesHandler, + InputHandler, + VideoHandler, + parse_export_feat, +) +from depth_anything_3.utils.constants import ( + DEFAULT_EXPORT_DIR, + DEFAULT_GALLERY_DIR, + DEFAULT_GRADIO_DIR, + DEFAULT_MODEL, +) + +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + +app = typer.Typer(help="Depth Anything 3 - Video depth estimation CLI", add_completion=False) + + +# ============================================================================ +# Input type detection utilities +# ============================================================================ + +# Supported file extensions +IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif"} +VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".webm", ".m4v"} + + +def detect_input_type(input_path: str) -> str: + """ + Detect input type from path. + + Returns: + - "image": Single image file + - "images": Directory containing images + - "video": Video file + - "colmap": COLMAP directory structure + - "unknown": Cannot determine type + """ + if not os.path.exists(input_path): + return "unknown" + + # Check if it's a file + if os.path.isfile(input_path): + ext = os.path.splitext(input_path)[1].lower() + if ext in IMAGE_EXTENSIONS: + return "image" + elif ext in VIDEO_EXTENSIONS: + return "video" + return "unknown" + + # Check if it's a directory + if os.path.isdir(input_path): + # Check for COLMAP structure + images_dir = os.path.join(input_path, "images") + sparse_dir = os.path.join(input_path, "sparse") + + if os.path.isdir(images_dir) and os.path.isdir(sparse_dir): + return "colmap" + + # Check if directory contains image files + for item in os.listdir(input_path): + item_path = os.path.join(input_path, item) + if os.path.isfile(item_path): + ext = os.path.splitext(item)[1].lower() + if ext in IMAGE_EXTENSIONS: + return "images" + + return "unknown" + + return "unknown" + + +# ============================================================================ +# Common parameters and configuration +# ============================================================================ + +# ============================================================================ +# Inference commands +# ============================================================================ + + +@app.command() +def auto( + input_path: str = typer.Argument( + ..., help="Path to input (image, directory, video, or COLMAP)" + ), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS]Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Video-specific options + fps: float = typer.Option(1.0, help="[Video] Sampling FPS for frame extraction"), + # COLMAP-specific options + sparse_subdir: str = typer.Option( + "", help="[COLMAP] Sparse reconstruction subdirectory (e.g., '0' for sparse/0/)" + ), + align_to_input_ext_scale: bool = typer.Option( + True, help="[COLMAP] Align prediction to input extrinsics scale" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """ + Automatically detect input type and run appropriate processing. + + Supports: + - Single image file (.jpg, .png, etc.) + - Directory of images + - Video file (.mp4, .avi, etc.) + - COLMAP directory (with 'images' and 'sparse' subdirectories) + """ + # Detect input type + input_type = detect_input_type(input_path) + + if input_type == "unknown": + typer.echo(f"❌ Error: Cannot determine input type for: {input_path}", err=True) + typer.echo("Supported inputs:", err=True) + typer.echo(" - Single image file (.jpg, .png, etc.)", err=True) + typer.echo(" - Directory containing images", err=True) + typer.echo(" - Video file (.mp4, .avi, etc.)", err=True) + typer.echo(" - COLMAP directory (with 'images/' and 'sparse/' subdirectories)", err=True) + raise typer.Exit(1) + + # Display detected type + typer.echo(f"🔍 Detected input type: {input_type.upper()}") + typer.echo(f"📁 Input path: {input_path}") + typer.echo() + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Route to appropriate handler + if input_type == "image": + typer.echo("Processing single image...") + # Process input + image_files = ImageHandler.process(input_path) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + elif input_type == "images": + typer.echo("Processing directory of images...") + # Process input - use default extensions + image_files = ImagesHandler.process(input_path, "png,jpg,jpeg") + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + elif input_type == "video": + typer.echo(f"Processing video with FPS={fps}...") + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Process input + image_files = VideoHandler.process(input_path, export_dir, fps) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + elif input_type == "colmap": + typer.echo( + f"Processing COLMAP directory (sparse subdirectory: '{sparse_subdir or 'default'}')..." + ) + # Process input + image_files, extrinsics, intrinsics = ColmapHandler.process(input_path, sparse_subdir) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + typer.echo() + typer.echo("✅ Processing completed successfully!") + + +@app.command() +def image( + image_path: str = typer.Argument(..., help="Path to input image file"), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run camera pose and depth estimation on a single image.""" + # Process input + image_files = ImageHandler.process(image_path) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +@app.command() +def images( + images_dir: str = typer.Argument(..., help="Path to directory containing input images"), + image_extensions: str = typer.Option( + "png,jpg,jpeg", help="Comma-separated image file extensions to process" + ), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run camera pose and depth estimation on a directory of images.""" + # Process input + image_files = ImagesHandler.process(images_dir, image_extensions) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +@app.command() +def colmap( + colmap_dir: str = typer.Argument( + ..., help="Path to COLMAP directory containing 'images' and 'sparse' subdirectories" + ), + sparse_subdir: str = typer.Option( + "", help="Sparse reconstruction subdirectory (e.g., '0' for sparse/0/, empty for sparse/)" + ), + align_to_input_ext_scale: bool = typer.Option( + True, help="Align prediction to input extrinsics scale" + ), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run pose conditioned depth estimation on COLMAP data.""" + # Process input + image_files, extrinsics, intrinsics = ColmapHandler.process(colmap_dir, sparse_subdir) + + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +@app.command() +def video( + video_path: str = typer.Argument(..., help="Path to input video file"), + fps: float = typer.Option(1.0, help="Sampling FPS for frame extraction"), + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"), + export_format: str = typer.Option("glb", help="Export format"), + device: str = typer.Option("cuda", help="Device to use"), + use_backend: bool = typer.Option(False, help="Use backend service for inference"), + backend_url: str = typer.Option( + "http://localhost:8008", help="Backend URL (default: http://localhost:8008)" + ), + process_res: int = typer.Option(504, help="Processing resolution"), + process_res_method: str = typer.Option( + "upper_bound_resize", help="Processing resolution method" + ), + export_feat: str = typer.Option( + "", + help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').", + ), + auto_cleanup: bool = typer.Option( + False, help="Automatically clean export directory if it exists (no prompt)" + ), + # Pose estimation options + use_ray_pose: bool = typer.Option( + False, help="Use ray-based pose estimation instead of camera decoder" + ), + ref_view_strategy: str = typer.Option( + "saddle_balanced", + help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range", + ), + # GLB export options + conf_thresh_percentile: float = typer.Option( + 40.0, help="[GLB] Lower percentile for adaptive confidence threshold" + ), + num_max_points: int = typer.Option( + 1_000_000, help="[GLB] Maximum number of points in the point cloud" + ), + show_cameras: bool = typer.Option( + True, help="[GLB] Show camera wireframes in the exported scene" + ), + # Feat_vis export options + feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"), +): + """Run depth estimation on video by extracting frames and processing them.""" + # Handle export directory + export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup) + + # Process input + image_files = VideoHandler.process(video_path, export_dir, fps) + + # Parse export_feat parameter + export_feat_layers = parse_export_feat(export_feat) + + # Determine backend URL based on use_backend flag + final_backend_url = backend_url if use_backend else None + + # Run inference + run_inference( + image_paths=image_files, + export_dir=export_dir, + model_dir=model_dir, + device=device, + backend_url=final_backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + use_ray_pose=use_ray_pose, + reference_view_strategy=reference_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + + +# ============================================================================ +# Service management commands +# ============================================================================ + + +@app.command() +def backend( + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + device: str = typer.Option("cuda", help="Device to use"), + host: str = typer.Option("127.0.0.1", help="Host to bind to"), + port: int = typer.Option(8008, help="Port to bind to"), + gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery directory path (optional)"), +): + """Start model backend service with integrated gallery.""" + typer.echo("=" * 60) + typer.echo("🚀 Starting Depth Anything 3 Backend Server") + typer.echo("=" * 60) + typer.echo(f"Model directory: {model_dir}") + typer.echo(f"Device: {device}") + + # Check if gallery directory exists + if gallery_dir and os.path.exists(gallery_dir): + typer.echo(f"Gallery directory: {gallery_dir}") + else: + gallery_dir = None # Disable gallery if directory doesn't exist + + typer.echo() + typer.echo("📡 Server URLs (Ctrl/CMD+Click to open):") + typer.echo(f" 🏠 Home: http://{host}:{port}") + typer.echo(f" 📊 Dashboard: http://{host}:{port}/dashboard") + typer.echo(f" 📈 API Status: http://{host}:{port}/status") + + if gallery_dir: + typer.echo(f" 🎨 Gallery: http://{host}:{port}/gallery/") + + typer.echo("=" * 60) + + try: + start_server(model_dir, device, host, port, gallery_dir) + except KeyboardInterrupt: + typer.echo("\n👋 Backend server stopped.") + except Exception as e: + typer.echo(f"❌ Failed to start backend: {e}") + raise typer.Exit(1) + + +# ============================================================================ +# Application launch commands +# ============================================================================ + + +@app.command() +def gradio( + model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"), + workspace_dir: str = typer.Option(DEFAULT_GRADIO_DIR, help="Workspace directory path"), + gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery directory path"), + host: str = typer.Option("127.0.0.1", help="Host address to bind to"), + port: int = typer.Option(7860, help="Port number to bind to"), + share: bool = typer.Option(False, help="Create a public link for the app"), + debug: bool = typer.Option(False, help="Enable debug mode"), + cache_examples: bool = typer.Option( + False, help="Pre-cache all example scenes at startup for faster loading" + ), + cache_gs_tag: str = typer.Option( + "", + help="Tag to match scene names for high-res+3DGS caching (e.g., 'dl3dv'). Scenes containing this tag will use high_res and infer_gs=True; others will use low_res only.", + ), +): + """Launch Depth Anything 3 Gradio interactive web application""" + from depth_anything_3.app.gradio_app import DepthAnything3App + + # Create necessary directories + os.makedirs(workspace_dir, exist_ok=True) + os.makedirs(gallery_dir, exist_ok=True) + + typer.echo("Launching Depth Anything 3 Gradio application...") + typer.echo(f"Model directory: {model_dir}") + typer.echo(f"Workspace directory: {workspace_dir}") + typer.echo(f"Gallery directory: {gallery_dir}") + typer.echo(f"Host: {host}") + typer.echo(f"Port: {port}") + typer.echo(f"Share: {share}") + typer.echo(f"Debug mode: {debug}") + typer.echo(f"Cache examples: {cache_examples}") + if cache_examples: + if cache_gs_tag: + typer.echo( + f"Cache GS Tag: '{cache_gs_tag}' (scenes matching this tag will use high-res + 3DGS)" + ) + else: + typer.echo(f"Cache GS Tag: None (all scenes will use low-res only)") + + try: + # Initialize and launch application + app = DepthAnything3App( + model_dir=model_dir, workspace_dir=workspace_dir, gallery_dir=gallery_dir + ) + + # Pre-cache examples if requested + if cache_examples: + typer.echo("\n" + "=" * 60) + typer.echo("Pre-caching mode enabled") + if cache_gs_tag: + typer.echo(f"Scenes containing '{cache_gs_tag}' will use HIGH-RES + 3DGS") + typer.echo(f"Other scenes will use LOW-RES only") + else: + typer.echo(f"All scenes will use LOW-RES only") + typer.echo("=" * 60) + app.cache_examples( + show_cam=True, + filter_black_bg=False, + filter_white_bg=False, + save_percentage=20.0, + num_max_points=1000, + cache_gs_tag=cache_gs_tag, + gs_trj_mode="smooth", + gs_video_quality="low", + ) + + # Prepare launch arguments + launch_kwargs = {"share": share, "debug": debug} + + app.launch(host=host, port=port, **launch_kwargs) + + except KeyboardInterrupt: + typer.echo("\nGradio application stopped.") + except Exception as e: + typer.echo(f"Failed to launch Gradio application: {e}") + raise typer.Exit(1) + + +@app.command() +def gallery( + gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery root directory"), + host: str = typer.Option("127.0.0.1", help="Host address to bind to"), + port: int = typer.Option(8007, help="Port number to bind to"), + open_browser: bool = typer.Option(False, help="Open browser after launch"), +): + """Launch Depth Anything 3 Gallery server""" + + # Validate gallery directory + if not os.path.exists(gallery_dir): + raise typer.BadParameter(f"Gallery directory not found: {gallery_dir}") + + typer.echo("Launching Depth Anything 3 Gallery server...") + typer.echo(f"Gallery directory: {gallery_dir}") + typer.echo(f"Host: {host}") + typer.echo(f"Port: {port}") + typer.echo(f"Auto-open browser: {open_browser}") + + try: + # Set command line arguments + import sys + + sys.argv = ["gallery", "--dir", gallery_dir, "--host", host, "--port", str(port)] + if open_browser: + sys.argv.append("--open") + + # Launch gallery server + gallery_main() + + except KeyboardInterrupt: + typer.echo("\nGallery server stopped.") + except Exception as e: + typer.echo(f"Failed to launch Gallery server: {e}") + raise typer.Exit(1) + + +if __name__ == "__main__": + app() diff --git a/src/nn/depth_anything_3/configs/da3-base.yaml b/src/nn/depth_anything_3/configs/da3-base.yaml new file mode 100644 index 0000000..c52a7e5 --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3-base.yaml @@ -0,0 +1,45 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitb + out_layers: [5, 7, 9, 11] + alt_start: 4 + qknorm_start: 4 + rope_start: 4 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 1536 + output_dim: 2 + features: &head_features 128 + out_channels: &head_out_channels [96, 192, 384, 768] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 768 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 1536 diff --git a/src/nn/depth_anything_3/configs/da3-giant.yaml b/src/nn/depth_anything_3/configs/da3-giant.yaml new file mode 100644 index 0000000..f5a75c0 --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3-giant.yaml @@ -0,0 +1,71 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitg + out_layers: [19, 27, 33, 39] + alt_start: 13 + qknorm_start: 13 + rope_start: 13 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 3072 + output_dim: 2 + features: &head_features 256 + out_channels: &head_out_channels [256, 512, 1024, 1024] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 1536 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 3072 + + +gs_head: + __object__: + path: depth_anything_3.model.gsdpt + name: GSDPT + args: as_params + + dim_in: *head_dim_in + output_dim: 38 # should align with gs_adapter's setting, for gs params + features: *head_features + out_channels: *head_out_channels + + +gs_adapter: + __object__: + path: depth_anything_3.model.gs_adapter + name: GaussianAdapter + args: as_params + + sh_degree: 2 + pred_color: false # predict SH coefficient if false + pred_offset_depth: true + pred_offset_xy: true + gaussian_scale_min: 1e-5 + gaussian_scale_max: 30.0 diff --git a/src/nn/depth_anything_3/configs/da3-large.yaml b/src/nn/depth_anything_3/configs/da3-large.yaml new file mode 100644 index 0000000..4fa367c --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3-large.yaml @@ -0,0 +1,45 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitl + out_layers: [11, 15, 19, 23] + alt_start: 8 + qknorm_start: 8 + rope_start: 8 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 2048 + output_dim: 2 + features: &head_features 256 + out_channels: &head_out_channels [256, 512, 1024, 1024] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 1024 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 2048 diff --git a/src/nn/depth_anything_3/configs/da3-small.yaml b/src/nn/depth_anything_3/configs/da3-small.yaml new file mode 100644 index 0000000..1088743 --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3-small.yaml @@ -0,0 +1,45 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vits + out_layers: [5, 7, 9, 11] + alt_start: 4 + qknorm_start: 4 + rope_start: 4 + cat_token: True + +head: + __object__: + path: depth_anything_3.model.dualdpt + name: DualDPT + args: as_params + + dim_in: &head_dim_in 768 + output_dim: 2 + features: &head_features 64 + out_channels: &head_out_channels [48, 96, 192, 384] + + +cam_enc: + __object__: + path: depth_anything_3.model.cam_enc + name: CameraEnc + args: as_params + + dim_out: 384 + +cam_dec: + __object__: + path: depth_anything_3.model.cam_dec + name: CameraDec + args: as_params + + dim_in: 768 diff --git a/src/nn/depth_anything_3/configs/da3metric-large.yaml b/src/nn/depth_anything_3/configs/da3metric-large.yaml new file mode 100644 index 0000000..124635c --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3metric-large.yaml @@ -0,0 +1,28 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitl + out_layers: [4, 11, 17, 23] + alt_start: -1 # -1 means disable + qknorm_start: -1 + rope_start: -1 + cat_token: False + +head: + __object__: + path: depth_anything_3.model.dpt + name: DPT + args: as_params + + dim_in: 1024 + output_dim: 1 + features: 256 + out_channels: [256, 512, 1024, 1024] diff --git a/src/nn/depth_anything_3/configs/da3mono-large.yaml b/src/nn/depth_anything_3/configs/da3mono-large.yaml new file mode 100644 index 0000000..124635c --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3mono-large.yaml @@ -0,0 +1,28 @@ +__object__: + path: depth_anything_3.model.da3 + name: DepthAnything3Net + args: as_params + +net: + __object__: + path: depth_anything_3.model.dinov2.dinov2 + name: DinoV2 + args: as_params + + name: vitl + out_layers: [4, 11, 17, 23] + alt_start: -1 # -1 means disable + qknorm_start: -1 + rope_start: -1 + cat_token: False + +head: + __object__: + path: depth_anything_3.model.dpt + name: DPT + args: as_params + + dim_in: 1024 + output_dim: 1 + features: 256 + out_channels: [256, 512, 1024, 1024] diff --git a/src/nn/depth_anything_3/configs/da3nested-giant-large.yaml b/src/nn/depth_anything_3/configs/da3nested-giant-large.yaml new file mode 100644 index 0000000..595c122 --- /dev/null +++ b/src/nn/depth_anything_3/configs/da3nested-giant-large.yaml @@ -0,0 +1,10 @@ +__object__: + path: depth_anything_3.model.da3 + name: NestedDepthAnything3Net + args: as_params + +anyview: + __inherit__: depth_anything_3.configs.da3-giant + +metric: + __inherit__: depth_anything_3.configs.da3metric-large diff --git a/src/nn/depth_anything_3/model/__init__.py b/src/nn/depth_anything_3/model/__init__.py new file mode 100644 index 0000000..57a2a45 --- /dev/null +++ b/src/nn/depth_anything_3/model/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from depth_anything_3.model.da3 import DepthAnything3Net, NestedDepthAnything3Net + +__export__ = [ + NestedDepthAnything3Net, + DepthAnything3Net, +] diff --git a/src/nn/depth_anything_3/model/cam_dec.py b/src/nn/depth_anything_3/model/cam_dec.py new file mode 100644 index 0000000..3353b40 --- /dev/null +++ b/src/nn/depth_anything_3/model/cam_dec.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn + + +class CameraDec(nn.Module): + def __init__(self, dim_in=1536): + super().__init__() + output_dim = dim_in + self.backbone = nn.Sequential( + nn.Linear(output_dim, output_dim), + nn.ReLU(), + nn.Linear(output_dim, output_dim), + nn.ReLU(), + ) + self.fc_t = nn.Linear(output_dim, 3) + self.fc_qvec = nn.Linear(output_dim, 4) + self.fc_fov = nn.Sequential(nn.Linear(output_dim, 2), nn.ReLU()) + + def forward(self, feat, camera_encoding=None, *args, **kwargs): + B, N = feat.shape[:2] + feat = feat.reshape(B * N, -1) + feat = self.backbone(feat) + out_t = self.fc_t(feat.float()).reshape(B, N, 3) + if camera_encoding is None: + out_qvec = self.fc_qvec(feat.float()).reshape(B, N, 4) + out_fov = self.fc_fov(feat.float()).reshape(B, N, 2) + else: + out_qvec = camera_encoding[..., 3:7] + out_fov = camera_encoding[..., -2:] + pose_enc = torch.cat([out_t, out_qvec, out_fov], dim=-1) + return pose_enc diff --git a/src/nn/depth_anything_3/model/cam_enc.py b/src/nn/depth_anything_3/model/cam_enc.py new file mode 100644 index 0000000..bf28e70 --- /dev/null +++ b/src/nn/depth_anything_3/model/cam_enc.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch.nn as nn + +from depth_anything_3.model.utils.attention import Mlp +from depth_anything_3.model.utils.block import Block +from depth_anything_3.model.utils.transform import extri_intri_to_pose_encoding +from depth_anything_3.utils.geometry import affine_inverse + + +class CameraEnc(nn.Module): + """ + CameraHead predicts camera parameters from token representations using iterative refinement. + + It applies a series of transformer blocks (the "trunk") to dedicated camera tokens. + """ + + def __init__( + self, + dim_out: int = 1024, + dim_in: int = 9, + trunk_depth: int = 4, + target_dim: int = 9, + num_heads: int = 16, + mlp_ratio: int = 4, + init_values: float = 0.01, + **kwargs, + ): + super().__init__() + self.target_dim = target_dim + self.trunk_depth = trunk_depth + self.trunk = nn.Sequential( + *[ + Block( + dim=dim_out, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + init_values=init_values, + ) + for _ in range(trunk_depth) + ] + ) + self.token_norm = nn.LayerNorm(dim_out) + self.trunk_norm = nn.LayerNorm(dim_out) + self.pose_branch = Mlp( + in_features=dim_in, + hidden_features=dim_out // 2, + out_features=dim_out, + drop=0, + ) + + def forward( + self, + ext, + ixt, + image_size, + ) -> tuple: + c2ws = affine_inverse(ext) + pose_encoding = extri_intri_to_pose_encoding( + c2ws, + ixt, + image_size, + ) + pose_tokens = self.pose_branch(pose_encoding) + pose_tokens = self.token_norm(pose_tokens) + pose_tokens = self.trunk(pose_tokens) + pose_tokens = self.trunk_norm(pose_tokens) + return pose_tokens diff --git a/src/nn/depth_anything_3/model/da3.py b/src/nn/depth_anything_3/model/da3.py new file mode 100644 index 0000000..f9e5bd6 --- /dev/null +++ b/src/nn/depth_anything_3/model/da3.py @@ -0,0 +1,442 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import torch +import torch.nn as nn +from addict import Dict +from omegaconf import DictConfig, OmegaConf + +from depth_anything_3.cfg import create_object +from depth_anything_3.model.utils.transform import pose_encoding_to_extri_intri +from depth_anything_3.utils.alignment import ( + apply_metric_scaling, + compute_alignment_mask, + compute_sky_mask, + least_squares_scale_scalar, + sample_tensor_for_quantile, + set_sky_regions_to_max_depth, +) +from depth_anything_3.utils.geometry import affine_inverse, as_homogeneous, map_pdf_to_opacity +from depth_anything_3.utils.ray_utils import get_extrinsic_from_camray + + +def _wrap_cfg(cfg_obj): + return OmegaConf.create(cfg_obj) + + +class DepthAnything3Net(nn.Module): + """ + Depth Anything 3 network for depth estimation and camera pose estimation. + + This network consists of: + - Backbone: DinoV2 feature extractor + - Head: DPT or DualDPT for depth prediction + - Optional camera decoders for pose estimation + - Optional GSDPT for 3DGS prediction + + Args: + preset: Configuration preset containing network dimensions and settings + + Returns: + Dictionary containing: + - depth: Predicted depth map (B, H, W) + - depth_conf: Depth confidence map (B, H, W) + - extrinsics: Camera extrinsics (B, N, 4, 4) + - intrinsics: Camera intrinsics (B, N, 3, 3) + - gaussians: 3D Gaussian Splats (world space), type: model.gs_adapter.Gaussians + - aux: Auxiliary features for specified layers + """ + + # Patch size for feature extraction + PATCH_SIZE = 14 + + def __init__(self, net, head, cam_dec=None, cam_enc=None, gs_head=None, gs_adapter=None): + """ + Initialize DepthAnything3Net with given yaml-initialized configuration. + """ + super().__init__() + self.backbone = net if isinstance(net, nn.Module) else create_object(_wrap_cfg(net)) + self.head = head if isinstance(head, nn.Module) else create_object(_wrap_cfg(head)) + self.cam_dec, self.cam_enc = None, None + if cam_dec is not None: + self.cam_dec = ( + cam_dec if isinstance(cam_dec, nn.Module) else create_object(_wrap_cfg(cam_dec)) + ) + self.cam_enc = ( + cam_enc if isinstance(cam_enc, nn.Module) else create_object(_wrap_cfg(cam_enc)) + ) + self.gs_adapter, self.gs_head = None, None + if gs_head is not None and gs_adapter is not None: + self.gs_adapter = ( + gs_adapter + if isinstance(gs_adapter, nn.Module) + else create_object(_wrap_cfg(gs_adapter)) + ) + gs_out_dim = self.gs_adapter.d_in + 1 + if isinstance(gs_head, nn.Module): + assert ( + gs_head.out_dim == gs_out_dim + ), f"gs_head.out_dim should be {gs_out_dim}, got {gs_head.out_dim}" + self.gs_head = gs_head + else: + assert ( + gs_head["output_dim"] == gs_out_dim + ), f"gs_head output_dim should set to {gs_out_dim}, got {gs_head['output_dim']}" + self.gs_head = create_object(_wrap_cfg(gs_head)) + + def forward( + self, + x: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + export_feat_layers: list[int] | None = [], + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> Dict[str, torch.Tensor]: + """ + Forward pass through the network. + + Args: + x: Input images (B, N, 3, H, W) + extrinsics: Camera extrinsics (B, N, 4, 4) + intrinsics: Camera intrinsics (B, N, 3, 3) + feat_layers: List of layer indices to extract features from + infer_gs: Enable Gaussian Splatting branch + use_ray_pose: Use ray-based pose estimation + ref_view_strategy: Strategy for selecting reference view + + Returns: + Dictionary containing predictions and auxiliary features + """ + # Extract features using backbone + if extrinsics is not None: + with torch.autocast(device_type=x.device.type, enabled=False): + cam_token = self.cam_enc(extrinsics, intrinsics, x.shape[-2:]) + else: + cam_token = None + + feats, aux_feats = self.backbone( + x, cam_token=cam_token, export_feat_layers=export_feat_layers, ref_view_strategy=ref_view_strategy + ) + # feats = [[item for item in feat] for feat in feats] + H, W = x.shape[-2], x.shape[-1] + + # Process features through depth head + with torch.autocast(device_type=x.device.type, enabled=False): + output = self._process_depth_head(feats, H, W) + if use_ray_pose: + output = self._process_ray_pose_estimation(output, H, W) + else: + output = self._process_camera_estimation(feats, H, W, output) + if infer_gs: + output = self._process_gs_head(feats, H, W, output, x, extrinsics, intrinsics) + + output = self._process_mono_sky_estimation(output) + + # Extract auxiliary features if requested + output.aux = self._extract_auxiliary_features(aux_feats, export_feat_layers, H, W) + + return output + + def _process_mono_sky_estimation( + self, output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Process mono sky estimation.""" + if "sky" not in output: + return output + non_sky_mask = compute_sky_mask(output.sky, threshold=0.3) + if non_sky_mask.sum() <= 10: + return output + if (~non_sky_mask).sum() <= 10: + return output + + non_sky_depth = output.depth[non_sky_mask] + if non_sky_depth.numel() > 100000: + idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device) + sampled_depth = non_sky_depth[idx] + else: + sampled_depth = non_sky_depth + non_sky_max = torch.quantile(sampled_depth, 0.99) + + # Set sky regions to maximum depth and high confidence + output.depth, _ = set_sky_regions_to_max_depth( + output.depth, None, non_sky_mask, max_depth=non_sky_max + ) + return output + + def _process_ray_pose_estimation( + self, output: Dict[str, torch.Tensor], height: int, width: int + ) -> Dict[str, torch.Tensor]: + """Process ray pose estimation if ray pose decoder is available.""" + if "ray" in output and "ray_conf" in output: + pred_extrinsic, pred_focal_lengths, pred_principal_points = get_extrinsic_from_camray( + output.ray, + output.ray_conf, + output.ray.shape[-3], + output.ray.shape[-2], + ) + pred_extrinsic = affine_inverse(pred_extrinsic) # w2c -> c2w + pred_extrinsic = pred_extrinsic[:, :, :3, :] + pred_intrinsic = torch.eye(3, 3)[None, None].repeat(pred_extrinsic.shape[0], pred_extrinsic.shape[1], 1, 1).clone().to(pred_extrinsic.device) + pred_intrinsic[:, :, 0, 0] = pred_focal_lengths[:, :, 0] / 2 * width + pred_intrinsic[:, :, 1, 1] = pred_focal_lengths[:, :, 1] / 2 * height + pred_intrinsic[:, :, 0, 2] = pred_principal_points[:, :, 0] * width * 0.5 + pred_intrinsic[:, :, 1, 2] = pred_principal_points[:, :, 1] * height * 0.5 + del output.ray + del output.ray_conf + output.extrinsics = pred_extrinsic + output.intrinsics = pred_intrinsic + return output + + def _process_depth_head( + self, feats: list[torch.Tensor], H: int, W: int + ) -> Dict[str, torch.Tensor]: + """Process features through the depth prediction head.""" + return self.head(feats, H, W, patch_start_idx=0) + + def _process_camera_estimation( + self, feats: list[torch.Tensor], H: int, W: int, output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Process camera pose estimation if camera decoder is available.""" + if self.cam_dec is not None: + pose_enc = self.cam_dec(feats[-1][1]) + # Remove ray information as it's not needed for pose estimation + if "ray" in output: + del output.ray + if "ray_conf" in output: + del output.ray_conf + + # Convert pose encoding to extrinsics and intrinsics + c2w, ixt = pose_encoding_to_extri_intri(pose_enc, (H, W)) + output.extrinsics = affine_inverse(c2w) + output.intrinsics = ixt + + return output + + def _process_gs_head( + self, + feats: list[torch.Tensor], + H: int, + W: int, + output: Dict[str, torch.Tensor], + in_images: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + ) -> Dict[str, torch.Tensor]: + """Process 3DGS parameters estimation if 3DGS head is available.""" + if self.gs_head is None or self.gs_adapter is None: + return output + assert output.get("depth", None) is not None, "must provide MV depth for the GS head." + + # The depth is defined in the DA3 model's camera space, + # so even with provided GT camera poses, + # we instead use the predicted camera poses for better alignment. + ctx_extr = output.get("extrinsics", None) + ctx_intr = output.get("intrinsics", None) + assert ( + ctx_extr is not None and ctx_intr is not None + ), "must process camera info first if GT is not available" + + gt_extr = extrinsics + # homo the extr if needed + ctx_extr = as_homogeneous(ctx_extr) + if gt_extr is not None: + gt_extr = as_homogeneous(gt_extr) + + # forward through the gs_dpt head to get 'camera space' parameters + gs_outs = self.gs_head( + feats=feats, + H=H, + W=W, + patch_start_idx=0, + images=in_images, + ) + raw_gaussians = gs_outs.raw_gs + densities = gs_outs.raw_gs_conf + + # convert to 'world space' 3DGS parameters; ready to export and render + # gt_extr could be None, and will be used to align the pose scale if available + gs_world = self.gs_adapter( + extrinsics=ctx_extr, + intrinsics=ctx_intr, + depths=output.depth, + opacities=map_pdf_to_opacity(densities), + raw_gaussians=raw_gaussians, + image_shape=(H, W), + gt_extrinsics=gt_extr, + ) + output.gaussians = gs_world + + return output + + def _extract_auxiliary_features( + self, feats: list[torch.Tensor], feat_layers: list[int], H: int, W: int + ) -> Dict[str, torch.Tensor]: + """Extract auxiliary features from specified layers.""" + aux_features = Dict() + assert len(feats) == len(feat_layers) + for feat, feat_layer in zip(feats, feat_layers): + # Reshape features to spatial dimensions + feat_reshaped = feat.reshape( + [ + feat.shape[0], + feat.shape[1], + H // self.PATCH_SIZE, + W // self.PATCH_SIZE, + feat.shape[-1], + ] + ) + aux_features[f"feat_layer_{feat_layer}"] = feat_reshaped + + return aux_features + + +class NestedDepthAnything3Net(nn.Module): + """ + Nested Depth Anything 3 network with metric scaling capabilities. + + This network combines two DepthAnything3Net branches: + - Main branch: Standard depth estimation + - Metric branch: Metric depth estimation for scaling alignment + + The network performs depth alignment using least squares scaling + and handles sky region masking for improved depth estimation. + + Args: + preset: Configuration for the main depth estimation branch + second_preset: Configuration for the metric depth branch + """ + + def __init__(self, anyview: DictConfig, metric: DictConfig): + """ + Initialize NestedDepthAnything3Net with two branches. + + Args: + preset: Configuration for main depth estimation branch + second_preset: Configuration for metric depth branch + """ + super().__init__() + self.da3 = create_object(anyview) + self.da3_metric = create_object(metric) + + def forward( + self, + x: torch.Tensor, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + export_feat_layers: list[int] | None = [], + infer_gs: bool = False, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + ) -> Dict[str, torch.Tensor]: + """ + Forward pass through both branches with metric scaling alignment. + + Args: + x: Input images (B, N, 3, H, W) + extrinsics: Camera extrinsics (B, N, 4, 4) - unused + intrinsics: Camera intrinsics (B, N, 3, 3) - unused + feat_layers: List of layer indices to extract features from + infer_gs: Enable Gaussian Splatting branch + use_ray_pose: Use ray-based pose estimation + ref_view_strategy: Strategy for selecting reference view + + Returns: + Dictionary containing aligned depth predictions and camera parameters + """ + # Get predictions from both branches + output = self.da3( + x, extrinsics, intrinsics, export_feat_layers=export_feat_layers, infer_gs=infer_gs, use_ray_pose=use_ray_pose, ref_view_strategy=ref_view_strategy + ) + metric_output = self.da3_metric(x) + + # Apply metric scaling and alignment + output = self._apply_metric_scaling(output, metric_output) + output = self._apply_depth_alignment(output, metric_output) + output = self._handle_sky_regions(output, metric_output) + + return output + + def _apply_metric_scaling( + self, output: Dict[str, torch.Tensor], metric_output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Apply metric scaling to the metric depth output.""" + # Scale metric depth based on camera intrinsics + metric_output.depth = apply_metric_scaling( + metric_output.depth, + output.intrinsics, + ) + return output + + def _apply_depth_alignment( + self, output: Dict[str, torch.Tensor], metric_output: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """Apply depth alignment using least squares scaling.""" + # Compute non-sky mask + non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3) + + # Ensure we have enough non-sky pixels + assert non_sky_mask.sum() > 10, "Insufficient non-sky pixels for alignment" + + # Sample depth confidence for quantile computation + depth_conf_ns = output.depth_conf[non_sky_mask] + depth_conf_sampled = sample_tensor_for_quantile(depth_conf_ns, max_samples=100000) + median_conf = torch.quantile(depth_conf_sampled, 0.5) + + # Compute alignment mask + align_mask = compute_alignment_mask( + output.depth_conf, non_sky_mask, output.depth, metric_output.depth, median_conf + ) + + # Compute scale factor using least squares + valid_depth = output.depth[align_mask] + valid_metric_depth = metric_output.depth[align_mask] + scale_factor = least_squares_scale_scalar(valid_metric_depth, valid_depth) + + # Apply scaling to depth and extrinsics + output.depth *= scale_factor + output.extrinsics[:, :, :3, 3] *= scale_factor + output.is_metric = 1 + output.scale_factor = scale_factor.item() + + return output + + def _handle_sky_regions( + self, + output: Dict[str, torch.Tensor], + metric_output: Dict[str, torch.Tensor], + sky_depth_def: float = 200.0, + ) -> Dict[str, torch.Tensor]: + """Handle sky regions by setting them to maximum depth.""" + non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3) + + # Compute maximum depth for non-sky regions + # Use sampling to safely compute quantile on large tensors + non_sky_depth = output.depth[non_sky_mask] + if non_sky_depth.numel() > 100000: + idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device) + sampled_depth = non_sky_depth[idx] + else: + sampled_depth = non_sky_depth + non_sky_max = min(torch.quantile(sampled_depth, 0.99), sky_depth_def) + + # Set sky regions to maximum depth and high confidence + output.depth, output.depth_conf = set_sky_regions_to_max_depth( + output.depth, output.depth_conf, non_sky_mask, max_depth=non_sky_max + ) + + return output diff --git a/src/nn/depth_anything_3/model/dinov2/dinov2.py b/src/nn/depth_anything_3/model/dinov2/dinov2.py new file mode 100644 index 0000000..ee0d88b --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/dinov2.py @@ -0,0 +1,64 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + + +from typing import List +import torch.nn as nn + +from depth_anything_3.model.dinov2.vision_transformer import ( + vit_base, + vit_giant2, + vit_large, + vit_small, +) + + +class DinoV2(nn.Module): + def __init__( + self, + name: str, + out_layers: List[int], + alt_start: int = -1, + qknorm_start: int = -1, + rope_start: int = -1, + cat_token: bool = True, + **kwargs, + ): + super().__init__() + assert name in {"vits", "vitb", "vitl", "vitg"} + self.name = name + self.out_layers = out_layers + self.alt_start = alt_start + self.qknorm_start = qknorm_start + self.rope_start = rope_start + self.cat_token = cat_token + encoder_map = { + "vits": vit_small, + "vitb": vit_base, + "vitl": vit_large, + "vitg": vit_giant2, + } + encoder_fn = encoder_map[self.name] + ffn_layer = "swiglufused" if self.name == "vitg" else "mlp" + self.pretrained = encoder_fn( + img_size=518, + patch_size=14, + ffn_layer=ffn_layer, + alt_start=alt_start, + qknorm_start=qknorm_start, + rope_start=rope_start, + cat_token=cat_token, + ) + + def forward(self, x, **kwargs): + return self.pretrained.get_intermediate_layers( + x, + self.out_layers, + **kwargs, + ) diff --git a/src/nn/depth_anything_3/model/dinov2/layers/__init__.py b/src/nn/depth_anything_3/model/dinov2/layers/__init__.py new file mode 100644 index 0000000..97dfba9 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# from .attention import MemEffAttention +from .block import Block +from .layer_scale import LayerScale +from .mlp import Mlp +from .patch_embed import PatchEmbed +from .rope import PositionGetter, RotaryPositionEmbedding2D +from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused + +__all__ = [ + Mlp, + PatchEmbed, + SwiGLUFFN, + SwiGLUFFNFused, + Block, + # MemEffAttention, + LayerScale, + PositionGetter, + RotaryPositionEmbedding2D, +] diff --git a/src/nn/depth_anything_3/model/dinov2/layers/attention.py b/src/nn/depth_anything_3/model/dinov2/layers/attention.py new file mode 100644 index 0000000..096b9d4 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/attention.py @@ -0,0 +1,100 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import logging +import torch.nn.functional as F +from torch import Tensor, nn + +logger = logging.getLogger("dinov2") + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + qk_norm: bool = False, + fused_attn: bool = True, # use F.scaled_dot_product_attention or not + rope=None, + ) -> None: + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + self.fused_attn = fused_attn + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + self.rope = rope + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = qkv[0], qkv[1], qkv[2] + q, k = self.q_norm(q), self.k_norm(k) + if self.rope is not None and pos is not None: + q = self.rope(q, pos) + k = self.rope(k, pos) + if self.fused_attn: + x = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.attn_drop.p if self.training else 0.0, + attn_mask=( + (attn_mask)[:, None].repeat(1, self.num_heads, 1, 1) + if attn_mask is not None + else None + ), + ) + else: + q = q * self.scale + attn = q @ k.transpose(-2, -1) + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = attn @ v + + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + def _forward(self, x: Tensor) -> Tensor: + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + + q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] + attn = q @ k.transpose(-2, -1) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x diff --git a/src/nn/depth_anything_3/model/dinov2/layers/block.py b/src/nn/depth_anything_3/model/dinov2/layers/block.py new file mode 100644 index 0000000..731519b --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/block.py @@ -0,0 +1,143 @@ +# flake8: noqa: F821 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +import logging +from typing import Callable, Optional +import torch +from torch import Tensor, nn + +from .attention import Attention +from .drop_path import DropPath +from .layer_scale import LayerScale +from .mlp import Mlp + +logger = logging.getLogger("dinov2") +XFORMERS_AVAILABLE = True + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + qk_norm: bool = False, + rope=None, + ln_eps: float = 1e-6, + ) -> None: + super().__init__() + # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") + self.norm1 = norm_layer(dim, eps=ln_eps) + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + qk_norm=qk_norm, + rope=rope, + ) + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim, eps=ln_eps) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.sample_drop_ratio = drop_path + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + def attn_residual_func(x: Tensor, pos=None, attn_mask=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask)) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + if self.training and self.sample_drop_ratio > 0.1: + # the overhead is compensated only for a drop path rate larger than 0.1 + x = drop_add_residual_stochastic_depth( + x, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + pos=pos, + ) + x = drop_add_residual_stochastic_depth( + x, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + elif self.training and self.sample_drop_ratio > 0.0: + x = x + self.drop_path1(attn_residual_func(x, pos=pos, attn_mask=attn_mask)) + x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 + else: + x = x + attn_residual_func(x, pos=pos, attn_mask=attn_mask) + x = x + ffn_residual_func(x) + return x + + +def drop_add_residual_stochastic_depth( + x: Tensor, + residual_func: Callable[[Tensor], Tensor], + sample_drop_ratio: float = 0.0, + pos: Optional[Tensor] = None, +) -> Tensor: + # 1) extract subset using permutation + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + x_subset = x[brange] + + # 2) apply residual_func to get residual + if pos is not None: + # if necessary, apply rope to the subset + pos = pos[brange] + residual = residual_func(x_subset, pos=pos) + else: + residual = residual_func(x_subset) + + x_flat = x.flatten(1) + residual = residual.flatten(1) + + residual_scale_factor = b / sample_subset_size + + # 3) add the residual + x_plus_residual = torch.index_add( + x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor + ) + return x_plus_residual.view_as(x) + + +def get_branges_scales(x, sample_drop_ratio=0.0): + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + residual_scale_factor = b / sample_subset_size + return brange, residual_scale_factor diff --git a/src/nn/depth_anything_3/model/dinov2/layers/drop_path.py b/src/nn/depth_anything_3/model/dinov2/layers/drop_path.py new file mode 100644 index 0000000..1c2cc94 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/drop_path.py @@ -0,0 +1,35 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py + + +from torch import nn + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0: + random_tensor.div_(keep_prob) + output = x * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) diff --git a/src/nn/depth_anything_3/model/dinov2/layers/layer_scale.py b/src/nn/depth_anything_3/model/dinov2/layers/layer_scale.py new file mode 100644 index 0000000..898ee12 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/layer_scale.py @@ -0,0 +1,31 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 # noqa: E501 + +from typing import Union +import torch +from torch import Tensor, nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.dim = dim + self.inplace = inplace + self.init_values = init_values + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + def extra_repr(self) -> str: + return f"{self.dim}, init_values={self.init_values}, inplace={self.inplace}" diff --git a/src/nn/depth_anything_3/model/dinov2/layers/mlp.py b/src/nn/depth_anything_3/model/dinov2/layers/mlp.py new file mode 100644 index 0000000..78ad0d8 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/mlp.py @@ -0,0 +1,40 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py + + +from typing import Callable, Optional +from torch import Tensor, nn + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/src/nn/depth_anything_3/model/dinov2/layers/patch_embed.py b/src/nn/depth_anything_3/model/dinov2/layers/patch_embed.py new file mode 100644 index 0000000..64bf6be --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/patch_embed.py @@ -0,0 +1,94 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, Optional, Tuple, Union +import torch.nn as nn +from torch import Tensor + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert ( + H % patch_H == 0 + ), f"Input image height {H} is not a multiple of patch height {patch_H}" + assert ( + W % patch_W == 0 + ), f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = ( + Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + ) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops diff --git a/src/nn/depth_anything_3/model/dinov2/layers/rope.py b/src/nn/depth_anything_3/model/dinov2/layers/rope.py new file mode 100644 index 0000000..f75ba37 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/rope.py @@ -0,0 +1,200 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + + +# Implementation of 2D Rotary Position Embeddings (RoPE). + +# This module provides a clean implementation of 2D Rotary Position Embeddings, +# which extends the original RoPE concept to handle 2D spatial positions. + +# Inspired by: +# https://github.com/meta-llama/codellama/blob/main/llama/model.py +# https://github.com/naver-ai/rope-vit + + +from typing import Dict, Tuple +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class PositionGetter: + """Generates and caches 2D spatial positions for patches in a grid. + + This class efficiently manages the generation of spatial coordinates for patches + in a 2D grid, caching results to avoid redundant computations. + + Attributes: + position_cache: Dictionary storing precomputed position tensors for different + grid dimensions. + """ + + def __init__(self): + """Initializes the position generator with an empty cache.""" + self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {} + + def __call__( + self, batch_size: int, height: int, width: int, device: torch.device + ) -> torch.Tensor: + """Generates spatial positions for a batch of patches. + + Args: + batch_size: Number of samples in the batch. + height: Height of the grid in patches. + width: Width of the grid in patches. + device: Target device for the position tensor. + + Returns: + Tensor of shape (batch_size, height*width, 2) containing y,x coordinates + for each position in the grid, repeated for each batch item. + """ + if (height, width) not in self.position_cache: + y_coords = torch.arange(height, device=device) + x_coords = torch.arange(width, device=device) + positions = torch.cartesian_prod(y_coords, x_coords) + self.position_cache[height, width] = positions + + cached_positions = self.position_cache[height, width] + return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone() + + +class RotaryPositionEmbedding2D(nn.Module): + """2D Rotary Position Embedding implementation. + + This module applies rotary position embeddings to input tokens based on their + 2D spatial positions. It handles the position-dependent rotation of features + separately for vertical and horizontal dimensions. + + Args: + frequency: Base frequency for the position embeddings. Default: 100.0 + scaling_factor: Scaling factor for frequency computation. Default: 1.0 + + Attributes: + base_frequency: Base frequency for computing position embeddings. + scaling_factor: Factor to scale the computed frequencies. + frequency_cache: Cache for storing precomputed frequency components. + """ + + def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0): + """Initializes the 2D RoPE module.""" + super().__init__() + self.base_frequency = frequency + self.scaling_factor = scaling_factor + self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {} + + def _compute_frequency_components( + self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Computes frequency components for rotary embeddings. + + Args: + dim: Feature dimension (must be even). + seq_len: Maximum sequence length. + device: Target device for computations. + dtype: Data type for the computed tensors. + + Returns: + Tuple of (cosine, sine) tensors for frequency components. + """ + cache_key = (dim, seq_len, device, dtype) + if cache_key not in self.frequency_cache: + # Compute frequency bands + exponents = torch.arange(0, dim, 2, device=device).float() / dim + inv_freq = 1.0 / (self.base_frequency**exponents) + + # Generate position-dependent frequencies + positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + angles = torch.einsum("i,j->ij", positions, inv_freq) + + # Compute and cache frequency components + angles = angles.to(dtype) + angles = torch.cat((angles, angles), dim=-1) + cos_components = angles.cos().to(dtype) + sin_components = angles.sin().to(dtype) + self.frequency_cache[cache_key] = (cos_components, sin_components) + + return self.frequency_cache[cache_key] + + @staticmethod + def _rotate_features(x: torch.Tensor) -> torch.Tensor: + """Performs feature rotation by splitting and recombining feature dimensions. + + Args: + x: Input tensor to rotate. + + Returns: + Rotated feature tensor. + """ + feature_dim = x.shape[-1] + x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_1d_rope( + self, + tokens: torch.Tensor, + positions: torch.Tensor, + cos_comp: torch.Tensor, + sin_comp: torch.Tensor, + ) -> torch.Tensor: + """Applies 1D rotary position embeddings along one dimension. + + Args: + tokens: Input token features. + positions: Position indices. + cos_comp: Cosine components for rotation. + sin_comp: Sine components for rotation. + + Returns: + Tokens with applied rotary position embeddings. + """ + # Embed positions with frequency components + cos = F.embedding(positions, cos_comp)[:, None, :, :] + sin = F.embedding(positions, sin_comp)[:, None, :, :] + # Apply rotation + return (tokens * cos) + (self._rotate_features(tokens) * sin) + + def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Applies 2D rotary position embeddings to input tokens. + + Args: + tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim). + The feature dimension (dim) must be divisible by 4. + positions: Position tensor of shape (batch_size, n_tokens, 2) containing + the y and x coordinates for each token. + + Returns: + Tensor of same shape as input with applied 2D rotary position embeddings. + + Raises: + AssertionError: If input dimensions are invalid or positions are malformed. + """ + # Validate inputs + assert tokens.size(-1) % 2 == 0, "Feature dimension must be even" + assert ( + positions.ndim == 3 and positions.shape[-1] == 2 + ), "Positions must have shape (batch_size, n_tokens, 2)" + + # Compute feature dimension for each spatial direction + feature_dim = tokens.size(-1) // 2 + + # Get frequency components + max_position = int(positions.max()) + 1 + cos_comp, sin_comp = self._compute_frequency_components( + feature_dim, max_position, tokens.device, tokens.dtype + ) + + # Split features for vertical and horizontal processing + vertical_features, horizontal_features = tokens.chunk(2, dim=-1) + + # Apply RoPE separately for each dimension + vertical_features = self._apply_1d_rope( + vertical_features, positions[..., 0], cos_comp, sin_comp + ) + horizontal_features = self._apply_1d_rope( + horizontal_features, positions[..., 1], cos_comp, sin_comp + ) + + # Combine processed features + return torch.cat((vertical_features, horizontal_features), dim=-1) diff --git a/src/nn/depth_anything_3/model/dinov2/layers/swiglu_ffn.py b/src/nn/depth_anything_3/model/dinov2/layers/swiglu_ffn.py new file mode 100644 index 0000000..c8f58e5 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/layers/swiglu_ffn.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Callable, Optional +import torch.nn.functional as F +from torch import Tensor, nn + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +try: + from xformers.ops import SwiGLU + + XFORMERS_AVAILABLE = True +except ImportError: + SwiGLU = SwiGLUFFN + XFORMERS_AVAILABLE = False + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/src/nn/depth_anything_3/model/dinov2/vision_transformer.py b/src/nn/depth_anything_3/model/dinov2/vision_transformer.py new file mode 100644 index 0000000..f0f8e23 --- /dev/null +++ b/src/nn/depth_anything_3/model/dinov2/vision_transformer.py @@ -0,0 +1,456 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import math +from typing import Callable, List, Sequence, Tuple, Union +import numpy as np +import torch +import torch.nn as nn +import torch.utils.checkpoint +from einops import rearrange + +from depth_anything_3.utils.logger import logger + +from .layers import LayerScale # noqa: F401 +from .layers import Mlp # noqa: F401 +from .layers import ( # noqa: F401 + Block, + PatchEmbed, + PositionGetter, + RotaryPositionEmbedding2D, + SwiGLUFFNFused, +) +from depth_anything_3.model.reference_view_selector import ( + RefViewStrategy, + select_reference_view, + reorder_by_reference, + restore_original_order, +) +from depth_anything_3.utils.constants import THRESH_FOR_REF_SELECTION + +# logger = logging.getLogger("dinov2") + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def named_apply( + fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False +) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply( + fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True + ) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=1.0, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + alt_start=-1, + qknorm_start=-1, + rope_start=-1, + rope_freq=100, + plus_cam_token=False, + cat_token=True, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating + positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating + positional embeddings + """ + super().__init__() + self.patch_start_idx = 1 + norm_layer = nn.LayerNorm + self.num_features = self.embed_dim = ( + embed_dim # num_features for consistency with other models + ) + self.alt_start = alt_start + self.qknorm_start = qknorm_start + self.rope_start = rope_start + self.cat_token = cat_token + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + + self.patch_embed = embed_layer( + img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim + ) + num_patches = self.patch_embed.num_patches + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + if self.alt_start != -1: + self.camera_token = nn.Parameter(torch.randn(1, 2, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) + if num_register_tokens + else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + if ffn_layer == "mlp": + logger.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + logger.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + logger.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + if self.rope_start != -1: + self.rope = RotaryPositionEmbedding2D(frequency=rope_freq) if rope_freq > 0 else None + self.position_getter = PositionGetter() if self.rope is not None else None + else: + self.rope = None + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + qk_norm=i >= qknorm_start if qknorm_start != -1 else False, + rope=self.rope if i >= rope_start and rope_start != -1 else None, + ) + for i in range(depth) + ] + self.blocks = nn.ModuleList(blocks_list) + self.norm = norm_layer(embed_dim) + + def interpolate_pos_encoding(self, x, w, h): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + M = int(math.sqrt(N)) # Recover the number of patches in each dimension + assert N == M * M + kwargs = {} + if self.interpolate_offset: + # Historical kludge: add a small number to avoid floating point error in the + # interpolation, see https://github.com/facebookresearch/dino/issues/8 + # Note: still needed for backward-compatibility, the underlying operators are using + # both output size and scale factors + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + kwargs["scale_factor"] = (sx, sy) + else: + # Simply specify an output size instead of a scale factor + kwargs["size"] = (w0, h0) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), + mode="bicubic", + antialias=self.interpolate_antialias, + **kwargs, + ) + assert (w0, h0) == patch_pos_embed.shape[-2:] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) + + def prepare_cls_token(self, B, S): + cls_token = self.cls_token.expand(B, S, -1) + cls_token = cls_token.reshape(B * S, -1, self.embed_dim) + return cls_token + + def prepare_tokens_with_masks(self, x, masks=None, cls_token=None, **kwargs): + B, S, nc, w, h = x.shape + x = rearrange(x, "b s c h w -> (b s) c h w") + x = self.patch_embed(x) + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + cls_token = self.prepare_cls_token(B, S) + x = torch.cat((cls_token, x), dim=1) + x = x + self.interpolate_pos_encoding(x, w, h) + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + x = rearrange(x, "(b s) n c -> b s n c", b=B, s=S) + return x + + def _prepare_rope(self, B, S, H, W, device): + pos = None + pos_nodiff = None + if self.rope is not None: + pos = self.position_getter( + B * S, H // self.patch_size, W // self.patch_size, device=device + ) + pos = rearrange(pos, "(b s) n c -> b s n c", b=B) + pos_nodiff = torch.zeros_like(pos).to(pos.dtype) + if self.patch_start_idx > 0: + pos = pos + 1 + pos_special = torch.zeros(B * S, self.patch_start_idx, 2).to(device).to(pos.dtype) + pos_special = rearrange(pos_special, "(b s) n c -> b s n c", b=B) + pos = torch.cat([pos_special, pos], dim=2) + pos_nodiff = pos_nodiff + 1 + pos_nodiff = torch.cat([pos_special, pos_nodiff], dim=2) + return pos, pos_nodiff + + def _get_intermediate_layers_not_chunked(self, x, n=1, export_feat_layers=[], **kwargs): + B, S, _, H, W = x.shape + x = self.prepare_tokens_with_masks(x) + output, total_block_len, aux_output = [], len(self.blocks), [] + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + pos, pos_nodiff = self._prepare_rope(B, S, H, W, x.device) + + for i, blk in enumerate(self.blocks): + if i < self.rope_start or self.rope is None: + g_pos, l_pos = None, None + else: + g_pos = pos_nodiff + l_pos = pos + + if self.alt_start != -1 and (i == self.alt_start - 1) and x.shape[1] >= THRESH_FOR_REF_SELECTION and kwargs.get("cam_token", None) is None: + # Select reference view using configured strategy + strategy = kwargs.get("ref_view_strategy", "saddle_balanced") + logger.info(f"Selecting reference view using strategy: {strategy}") + b_idx = select_reference_view(x, strategy=strategy) + # Reorder views to place reference view first + x = reorder_by_reference(x, b_idx) + local_x = reorder_by_reference(local_x, b_idx) + + if self.alt_start != -1 and i == self.alt_start: + if kwargs.get("cam_token", None) is not None: + logger.info("Using camera conditions provided by the user") + cam_token = kwargs.get("cam_token") + else: + ref_token = self.camera_token[:, :1].expand(B, -1, -1) + src_token = self.camera_token[:, 1:].expand(B, S - 1, -1) + cam_token = torch.cat([ref_token, src_token], dim=1) + x[:, :, 0] = cam_token + + if self.alt_start != -1 and i >= self.alt_start and i % 2 == 1: + x = self.process_attention( + x, blk, "global", pos=g_pos, attn_mask=kwargs.get("attn_mask", None) + ) + else: + x = self.process_attention(x, blk, "local", pos=l_pos) + local_x = x + + if i in blocks_to_take: + out_x = torch.cat([local_x, x], dim=-1) if self.cat_token else x + # Restore original view order if reordering was applied + if x.shape[1] >= THRESH_FOR_REF_SELECTION and self.alt_start != -1 and 'b_idx' in locals(): + out_x = restore_original_order(out_x, b_idx) + output.append((out_x[:, :, 0], out_x)) + if i in export_feat_layers: + aux_output.append(x) + return output, aux_output + + def process_attention(self, x, block, attn_type="global", pos=None, attn_mask=None): + b, s, n = x.shape[:3] + if attn_type == "local": + x = rearrange(x, "b s n c -> (b s) n c") + if pos is not None: + pos = rearrange(pos, "b s n c -> (b s) n c") + elif attn_type == "global": + x = rearrange(x, "b s n c -> b (s n) c") + if pos is not None: + pos = rearrange(pos, "b s n c -> b (s n) c") + else: + raise ValueError(f"Invalid attention type: {attn_type}") + + x = block(x, pos=pos, attn_mask=attn_mask) + + if attn_type == "local": + x = rearrange(x, "(b s) n c -> b s n c", b=b, s=s) + elif attn_type == "global": + x = rearrange(x, "b (s n) c -> b s n c", b=b, s=s) + return x + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + export_feat_layers: List[int] = [], + **kwargs, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + outputs, aux_outputs = self._get_intermediate_layers_not_chunked( + x, n, export_feat_layers=export_feat_layers, **kwargs + ) + camera_tokens = [out[0] for out in outputs] + if outputs[0][1].shape[-1] == self.embed_dim: + outputs = [self.norm(out[1]) for out in outputs] + elif outputs[0][1].shape[-1] == (self.embed_dim * 2): + outputs = [ + torch.cat( + [out[1][..., : self.embed_dim], self.norm(out[1][..., self.embed_dim :])], + dim=-1, + ) + for out in outputs + ] + else: + raise ValueError(f"Invalid output shape: {outputs[0][1].shape}") + aux_outputs = [self.norm(out) for out in aux_outputs] + outputs = [out[..., 1 + self.num_register_tokens :, :] for out in outputs] + aux_outputs = [out[..., 1 + self.num_register_tokens :, :] for out in aux_outputs] + return tuple(zip(outputs, camera_tokens)), aux_outputs + + +def vit_small(patch_size=16, num_register_tokens=0, depth=12, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=depth, + num_heads=6, + mlp_ratio=4, + # block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, depth=12, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=depth, + num_heads=12, + mlp_ratio=4, + # block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, depth=24, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=depth, + num_heads=16, + mlp_ratio=4, + # block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, depth=40, **kwargs): + """ + Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 + """ + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=depth, + num_heads=24, + mlp_ratio=4, + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model diff --git a/src/nn/depth_anything_3/model/dpt.py b/src/nn/depth_anything_3/model/dpt.py new file mode 100644 index 0000000..337a8a9 --- /dev/null +++ b/src/nn/depth_anything_3/model/dpt.py @@ -0,0 +1,458 @@ +# flake8: noqa E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict as TyDict +from typing import List, Sequence, Tuple +import torch +import torch.nn as nn +from addict import Dict +from einops import rearrange + +from depth_anything_3.model.utils.head_utils import ( + Permute, + create_uv_grid, + custom_interpolate, + position_grid_to_embed, +) + + +class DPT(nn.Module): + """ + DPT for dense prediction (main head + optional sky head, sky always 1 channel). + + Returns: + - Main head: + * If output_dim>1: { head_name, f"{head_name}_conf" } + * If output_dim==1: { head_name } + - Sky head (if use_sky_head=True): { sky_name } # [B, S, 1, H/down_ratio, W/down_ratio] + """ + + def __init__( + self, + dim_in: int, + *, + patch_size: int = 14, + output_dim: int = 1, + activation: str = "exp", + conf_activation: str = "expp1", + features: int = 256, + out_channels: Sequence[int] = (256, 512, 1024, 1024), + pos_embed: bool = False, + down_ratio: int = 1, + head_name: str = "depth", + # ---- sky head (fixed 1 channel) ---- + use_sky_head: bool = True, + sky_name: str = "sky", + sky_activation: str = "relu", # 'sigmoid' / 'relu' / 'linear' + use_ln_for_heads: bool = False, # If needed, apply LayerNorm on intermediate features of both heads + norm_type: str = "idt", # use to match legacy GS-DPT head, "idt" / "layer" + fusion_block_inplace: bool = False, + ) -> None: + super().__init__() + + # -------------------- configuration -------------------- + self.patch_size = patch_size + self.activation = activation + self.conf_activation = conf_activation + self.pos_embed = pos_embed + self.down_ratio = down_ratio + + # Names + self.head_main = head_name + self.sky_name = sky_name + + # Main head: output dimension and confidence switch + self.out_dim = output_dim + self.has_conf = output_dim > 1 + + # Sky head parameters (always 1 channel) + self.use_sky_head = use_sky_head + self.sky_activation = sky_activation + + # Fixed 4 intermediate outputs + self.intermediate_layer_idx: Tuple[int, int, int, int] = (0, 1, 2, 3) + + # -------------------- token pre-norm + per-stage projection -------------------- + if norm_type == "layer": + self.norm = nn.LayerNorm(dim_in) + elif norm_type == "idt": + self.norm = nn.Identity() + else: + raise Exception(f"Unknown norm_type {norm_type}, should be 'layer' or 'idt'.") + self.projects = nn.ModuleList( + [nn.Conv2d(dim_in, oc, kernel_size=1, stride=1, padding=0) for oc in out_channels] + ) + + # -------------------- Spatial re-size (align to common scale before fusion) -------------------- + # Design consistent with original: relative to patch grid (x4, x2, x1, /2) + self.resize_layers = nn.ModuleList( + [ + nn.ConvTranspose2d( + out_channels[0], out_channels[0], kernel_size=4, stride=4, padding=0 + ), + nn.ConvTranspose2d( + out_channels[1], out_channels[1], kernel_size=2, stride=2, padding=0 + ), + nn.Identity(), + nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1), + ] + ) + + # -------------------- scratch: stage adapters + main fusion chain -------------------- + self.scratch = _make_scratch(list(out_channels), features, expand=False) + + # Main fusion chain + self.scratch.refinenet1 = _make_fusion_block(features, inplace=fusion_block_inplace) + self.scratch.refinenet2 = _make_fusion_block(features, inplace=fusion_block_inplace) + self.scratch.refinenet3 = _make_fusion_block(features, inplace=fusion_block_inplace) + self.scratch.refinenet4 = _make_fusion_block( + features, has_residual=False, inplace=fusion_block_inplace + ) + + # Heads (shared neck1; then split into two heads) + head_features_1 = features + head_features_2 = 32 + self.scratch.output_conv1 = nn.Conv2d( + head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1 + ) + + ln_seq = ( + [Permute((0, 2, 3, 1)), nn.LayerNorm(head_features_2), Permute((0, 3, 1, 2))] + if use_ln_for_heads + else [] + ) + + # Main head + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1), + *ln_seq, + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0), + ) + + # Sky head (fixed 1 channel) + if self.use_sky_head: + self.scratch.sky_output_conv2 = nn.Sequential( + nn.Conv2d( + head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1 + ), + *ln_seq, + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), + ) + + # ------------------------------------------------------------------------- + # Public forward (supports frame chunking to save memory) + # ------------------------------------------------------------------------- + def forward( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + chunk_size: int = 8, + **kwargs, + ) -> Dict: + """ + Args: + feats: List of 4 entries, each entry is a tensor like [B, S, T, C] (or the 0th element of tuple/list is that tensor). + H, W: Original image dimensions + patch_start_idx: Starting index of patch tokens in sequence (for cropping non-patch tokens) + chunk_size: Chunk size along time dimension S + + Returns: + Dict[str, Tensor] + """ + B, S, N, C = feats[0][0].shape + feats = [feat[0].reshape(B * S, N, C) for feat in feats] + + # update image info, used by the GS-DPT head + extra_kwargs = {} + if "images" in kwargs: + extra_kwargs.update({"images": rearrange(kwargs["images"], "B S ... -> (B S) ...")}) + + if chunk_size is None or chunk_size >= S: + out_dict = self._forward_impl(feats, H, W, patch_start_idx, **extra_kwargs) + out_dict = {k: v.view(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + + out_dicts: List[TyDict[str, torch.Tensor]] = [] + for s0 in range(0, S, chunk_size): + s1 = min(s0 + chunk_size, S) + kw = {} + if "images" in extra_kwargs: + kw.update({"images": extra_kwargs["images"][s0:s1]}) + out_dicts.append( + self._forward_impl([f[s0:s1] for f in feats], H, W, patch_start_idx, **kw) + ) + out_dict = {k: torch.cat([od[k] for od in out_dicts], dim=0) for k in out_dicts[0].keys()} + out_dict = {k: v.view(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + + # ------------------------------------------------------------------------- + # Internal forward (single chunk) + # ------------------------------------------------------------------------- + def _forward_impl( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + ) -> TyDict[str, torch.Tensor]: + B, _, C = feats[0].shape + ph, pw = H // self.patch_size, W // self.patch_size + resized_feats = [] + for stage_idx, take_idx in enumerate(self.intermediate_layer_idx): + x = feats[take_idx][:, patch_start_idx:] # [B*S, N_patch, C] + x = self.norm(x) + # permute -> contiguous before reshape to keep conv input contiguous + x = x.permute(0, 2, 1).contiguous().reshape(B, C, ph, pw) # [B*S, C, ph, pw] + + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + x = self.resize_layers[stage_idx](x) # Align scale + resized_feats.append(x) + + # 2) Fusion pyramid (main branch only) + fused = self._fuse(resized_feats) + + # 3) Upsample to target resolution, optionally add position encoding again + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused = self.scratch.output_conv1(fused) + fused = custom_interpolate(fused, (h_out, w_out), mode="bilinear", align_corners=True) + if self.pos_embed: + fused = self._add_pos_embed(fused, W, H) + + # 4) Shared neck1 + feat = fused + + # 5) Main head: logits -> activation + main_logits = self.scratch.output_conv2(feat) + outs: TyDict[str, torch.Tensor] = {} + if self.has_conf: + fmap = main_logits.permute(0, 2, 3, 1) + pred = self._apply_activation_single(fmap[..., :-1], self.activation) + conf = self._apply_activation_single(fmap[..., -1], self.conf_activation) + outs[self.head_main] = pred.squeeze(1) + outs[f"{self.head_main}_conf"] = conf.squeeze(1) + else: + outs[self.head_main] = self._apply_activation_single( + main_logits, self.activation + ).squeeze(1) + + # 6) Sky head (fixed 1 channel) + if self.use_sky_head: + sky_logits = self.scratch.sky_output_conv2(feat) + outs[self.sky_name] = self._apply_sky_activation(sky_logits).squeeze(1) + + return outs + + # ------------------------------------------------------------------------- + # Subroutines + # ------------------------------------------------------------------------- + def _fuse(self, feats: List[torch.Tensor]) -> torch.Tensor: + """ + 4-layer top-down fusion, returns finest scale features (after fusion, before neck1). + """ + l1, l2, l3, l4 = feats + + l1_rn = self.scratch.layer1_rn(l1) + l2_rn = self.scratch.layer2_rn(l2) + l3_rn = self.scratch.layer3_rn(l3) + l4_rn = self.scratch.layer4_rn(l4) + + # 4 -> 3 -> 2 -> 1 + out = self.scratch.refinenet4(l4_rn, size=l3_rn.shape[2:]) + out = self.scratch.refinenet3(out, l3_rn, size=l2_rn.shape[2:]) + out = self.scratch.refinenet2(out, l2_rn, size=l1_rn.shape[2:]) + out = self.scratch.refinenet1(out, l1_rn) + return out + + def _apply_activation_single( + self, x: torch.Tensor, activation: str = "linear" + ) -> torch.Tensor: + """ + Apply activation to single channel output, maintaining semantic consistency with value branch in multi-channel case. + Supports: exp / relu / sigmoid / softplus / tanh / linear / expp1 + """ + act = activation.lower() if isinstance(activation, str) else activation + if act == "exp": + return torch.exp(x) + if act == "expp1": + return torch.exp(x) + 1 + if act == "expm1": + return torch.expm1(x) + if act == "relu": + return torch.relu(x) + if act == "sigmoid": + return torch.sigmoid(x) + if act == "softplus": + return torch.nn.functional.softplus(x) + if act == "tanh": + return torch.tanh(x) + # Default linear + return x + + def _apply_sky_activation(self, x: torch.Tensor) -> torch.Tensor: + """ + Sky head activation (fixed 1 channel): + * 'sigmoid' -> Sigmoid probability map + * 'relu' -> ReLU positive domain output + * 'linear' -> Original value (logits) + """ + act = ( + self.sky_activation.lower() + if isinstance(self.sky_activation, str) + else self.sky_activation + ) + if act == "sigmoid": + return torch.sigmoid(x) + if act == "relu": + return torch.relu(x) + # 'linear' + return x + + def _add_pos_embed(self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor: + """Simple UV position encoding directly added to feature map.""" + pw, ph = x.shape[-1], x.shape[-2] + pe = create_uv_grid(pw, ph, aspect_ratio=W / H, dtype=x.dtype, device=x.device) + pe = position_grid_to_embed(pe, x.shape[1]) * ratio + pe = pe.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1) + return x + pe + + +# ----------------------------------------------------------------------------- +# Building blocks (preserved, consistent with original) +# ----------------------------------------------------------------------------- +def _make_fusion_block( + features: int, + size: Tuple[int, int] = None, + has_residual: bool = True, + groups: int = 1, + inplace: bool = False, +) -> nn.Module: + return FeatureFusionBlock( + features=features, + activation=nn.ReLU(inplace=inplace), + deconv=False, + bn=False, + expand=False, + align_corners=True, + size=size, + has_residual=has_residual, + groups=groups, + ) + + +def _make_scratch( + in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False +) -> nn.Module: + scratch = nn.Module() + # Optional expansion by stage + c1 = out_shape + c2 = out_shape * (2 if expand else 1) + c3 = out_shape * (4 if expand else 1) + c4 = out_shape * (8 if expand else 1) + + scratch.layer1_rn = nn.Conv2d(in_shape[0], c1, 3, 1, 1, bias=False, groups=groups) + scratch.layer2_rn = nn.Conv2d(in_shape[1], c2, 3, 1, 1, bias=False, groups=groups) + scratch.layer3_rn = nn.Conv2d(in_shape[2], c3, 3, 1, 1, bias=False, groups=groups) + scratch.layer4_rn = nn.Conv2d(in_shape[3], c4, 3, 1, 1, bias=False, groups=groups) + return scratch + + +class ResidualConvUnit(nn.Module): + """Lightweight residual convolution block for fusion""" + + def __init__(self, features: int, activation: nn.Module, bn: bool, groups: int = 1) -> None: + super().__init__() + self.bn = bn + self.groups = groups + self.conv1 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) + self.conv2 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) + self.norm1 = None + self.norm2 = None + self.activation = activation + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] + out = self.activation(x) + out = self.conv1(out) + if self.norm1 is not None: + out = self.norm1(out) + + out = self.activation(out) + out = self.conv2(out) + if self.norm2 is not None: + out = self.norm2(out) + + return self.skip_add.add(out, x) + + +class FeatureFusionBlock(nn.Module): + """Top-down fusion block: (optional) residual merge + upsampling + 1x1 contraction""" + + def __init__( + self, + features: int, + activation: nn.Module, + deconv: bool = False, + bn: bool = False, + expand: bool = False, + align_corners: bool = True, + size: Tuple[int, int] = None, + has_residual: bool = True, + groups: int = 1, + ) -> None: + super().__init__() + self.align_corners = align_corners + self.size = size + self.has_residual = has_residual + + self.resConfUnit1 = ( + ResidualConvUnit(features, activation, bn, groups=groups) if has_residual else None + ) + self.resConfUnit2 = ResidualConvUnit(features, activation, bn, groups=groups) + + out_features = (features // 2) if expand else features + self.out_conv = nn.Conv2d(features, out_features, 1, 1, 0, bias=True, groups=groups) + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, *xs: torch.Tensor, size: Tuple[int, int] = None) -> torch.Tensor: # type: ignore[override] + """ + xs: + - xs[0]: Top branch input + - xs[1]: Lateral input (can do residual addition with top branch) + """ + y = xs[0] + if self.has_residual and len(xs) > 1 and self.resConfUnit1 is not None: + y = self.skip_add.add(y, self.resConfUnit1(xs[1])) + + y = self.resConfUnit2(y) + + # Upsampling + if (size is None) and (self.size is None): + up_kwargs = {"scale_factor": 2} + elif size is None: + up_kwargs = {"size": self.size} + else: + up_kwargs = {"size": size} + + y = custom_interpolate(y, **up_kwargs, mode="bilinear", align_corners=self.align_corners) + y = self.out_conv(y) + return y diff --git a/src/nn/depth_anything_3/model/dualdpt.py b/src/nn/depth_anything_3/model/dualdpt.py new file mode 100644 index 0000000..229e1a9 --- /dev/null +++ b/src/nn/depth_anything_3/model/dualdpt.py @@ -0,0 +1,488 @@ +# flake8: noqa E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Sequence, Tuple +import torch +import torch.nn as nn +from addict import Dict + +from depth_anything_3.model.dpt import _make_fusion_block, _make_scratch +from depth_anything_3.model.utils.head_utils import ( + Permute, + create_uv_grid, + custom_interpolate, + position_grid_to_embed, +) + + +class DualDPT(nn.Module): + """ + Dual-head DPT for dense prediction with an always-on auxiliary head. + + Architectural notes: + - Sky/object branches are removed. + - `intermediate_layer_idx` is fixed to (0, 1, 2, 3). + - Auxiliary head has its **own** fusion blocks (no fusion_inplace / no sharing). + - Auxiliary head is internally multi-level; **only the final level** is returned. + - Returns a **dict** with keys from `head_names`, e.g.: + { main_name, f"{main_name}_conf", aux_name, f"{aux_name}_conf" } + - `feature_only` is fixed to False. + """ + + def __init__( + self, + dim_in: int, + *, + patch_size: int = 14, + output_dim: int = 2, + activation: str = "exp", + conf_activation: str = "expp1", + features: int = 256, + out_channels: Sequence[int] = (256, 512, 1024, 1024), + pos_embed: bool = True, + down_ratio: int = 1, + aux_pyramid_levels: int = 4, + aux_out1_conv_num: int = 5, + head_names: Tuple[str, str] = ("depth", "ray"), + ) -> None: + super().__init__() + + # -------------------- configuration -------------------- + self.patch_size = patch_size + self.activation = activation + self.conf_activation = conf_activation + self.pos_embed = pos_embed + self.down_ratio = down_ratio + + self.aux_levels = aux_pyramid_levels + self.aux_out1_conv_num = aux_out1_conv_num + + # names ONLY come from config (no hard-coded strings elsewhere) + self.head_main, self.head_aux = head_names + + # Always expect 4 scales; enforce intermediate idx = (0, 1, 2, 3) + self.intermediate_layer_idx: Tuple[int, int, int, int] = (0, 1, 2, 3) + + # -------------------- token pre-norm + per-stage projection -------------------- + self.norm = nn.LayerNorm(dim_in) + self.projects = nn.ModuleList( + [nn.Conv2d(dim_in, oc, kernel_size=1, stride=1, padding=0) for oc in out_channels] + ) + + # -------------------- spatial re-sizers (align to common scale before fusion) -------------------- + # design: stage strides (x4, x2, x1, /2) relative to patch grid to align to a common pivot scale + self.resize_layers = nn.ModuleList( + [ + nn.ConvTranspose2d( + out_channels[0], out_channels[0], kernel_size=4, stride=4, padding=0 + ), + nn.ConvTranspose2d( + out_channels[1], out_channels[1], kernel_size=2, stride=2, padding=0 + ), + nn.Identity(), + nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1), + ] + ) + + # -------------------- scratch: stage adapters + fusion (main & aux are separate) -------------------- + self.scratch = _make_scratch(list(out_channels), features, expand=False) + + # Main fusion chain (independent) + self.scratch.refinenet1 = _make_fusion_block(features) + self.scratch.refinenet2 = _make_fusion_block(features) + self.scratch.refinenet3 = _make_fusion_block(features) + self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False) + + # Primary head neck + head (independent) + head_features_1 = features + head_features_2 = 32 + self.scratch.output_conv1 = nn.Conv2d( + head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1 + ) + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0), + ) + + # Auxiliary fusion chain (completely separate; no sharing, i.e., "fusion_inplace=False") + self.scratch.refinenet1_aux = _make_fusion_block(features) + self.scratch.refinenet2_aux = _make_fusion_block(features) + self.scratch.refinenet3_aux = _make_fusion_block(features) + self.scratch.refinenet4_aux = _make_fusion_block(features, has_residual=False) + + # Aux pre-head per level (we will only *return final level*) + self.scratch.output_conv1_aux = nn.ModuleList( + [self._make_aux_out1_block(head_features_1) for _ in range(self.aux_levels)] + ) + + # Aux final projection per level + use_ln = True + ln_seq = ( + [Permute((0, 2, 3, 1)), nn.LayerNorm(head_features_2), Permute((0, 3, 1, 2))] + if use_ln + else [] + ) + self.scratch.output_conv2_aux = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d( + head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1 + ), + *ln_seq, + nn.ReLU(inplace=True), + nn.Conv2d(head_features_2, 7, kernel_size=1, stride=1, padding=0), + ) + for _ in range(self.aux_levels) + ] + ) + + # ------------------------------------------------------------------------- + # Public forward (supports frame chunking for memory) + # ------------------------------------------------------------------------- + + def forward( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + chunk_size: int = 8, + ) -> Dict[str, torch.Tensor]: + """ + Args: + aggregated_tokens_list: List of 4 tensors [B, S, T, C] from transformer. + images: [B, S, 3, H, W], in [0, 1]. + patch_start_idx: Patch-token start in the token sequence (to drop non-patch tokens). + frames_chunk_size: Optional chunking along S for memory. + + Returns: + Dict[str, Tensor] with keys based on `head_names`, e.g.: + self.head_main, f"{self.head_main}_conf", + self.head_aux, f"{self.head_aux}_conf" + Shapes: + main: [B, S, out_dim, H/down_ratio, W/down_ratio] + main_cf: [B, S, 1, H/down_ratio, W/down_ratio] + aux: [B, S, 7, H/down_ratio, W/down_ratio] + aux_cf: [B, S, 1, H/down_ratio, W/down_ratio] + """ + B, S, N, C = feats[0][0].shape + feats = [feat[0].reshape(B * S, N, C) for feat in feats] + if chunk_size is None or chunk_size >= S: + out_dict = self._forward_impl(feats, H, W, patch_start_idx) + out_dict = {k: v.reshape(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + out_dicts = [] + for s0 in range(0, B * S, chunk_size): + s1 = min(s0 + chunk_size, B * S) + out_dict = self._forward_impl( + [feat[s0:s1] for feat in feats], + H, + W, + patch_start_idx, + ) + out_dicts.append(out_dict) + out_dict = { + k: torch.cat([out_dict[k] for out_dict in out_dicts], dim=0) + for k in out_dicts[0].keys() + } + out_dict = {k: v.view(B, S, *v.shape[1:]) for k, v in out_dict.items()} + return Dict(out_dict) + + # ------------------------------------------------------------------------- + # Internal forward (single chunk) + # ------------------------------------------------------------------------- + + def _forward_impl( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + ) -> Dict[str, torch.Tensor]: + B, _, C = feats[0].shape + ph, pw = H // self.patch_size, W // self.patch_size + resized_feats = [] + for stage_idx, take_idx in enumerate(self.intermediate_layer_idx): + x = feats[take_idx][:, patch_start_idx:] + x = self.norm(x) + x = x.permute(0, 2, 1).reshape(B, C, ph, pw) # [B*S, C, ph, pw] + + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + x = self.resize_layers[stage_idx](x) # align scales + resized_feats.append(x) + + # 2) Fuse pyramid (main & aux are completely independent) + fused_main, fused_aux_pyr = self._fuse(resized_feats) + + # 3) Upsample to target resolution and (optional) add pos-embed again + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused_main = custom_interpolate( + fused_main, (h_out, w_out), mode="bilinear", align_corners=True + ) + if self.pos_embed: + fused_main = self._add_pos_embed(fused_main, W, H) + + # Primary head: conv1 -> conv2 -> activate + # fused_main = self.scratch.output_conv1(fused_main) + main_logits = self.scratch.output_conv2(fused_main) + fmap = main_logits.permute(0, 2, 3, 1) + main_pred = self._apply_activation_single(fmap[..., :-1], self.activation) + main_conf = self._apply_activation_single(fmap[..., -1], self.conf_activation) + + # Auxiliary head (multi-level inside) -> only last level returned (after activation) + last_aux = fused_aux_pyr[-1] + if self.pos_embed: + last_aux = self._add_pos_embed(last_aux, W, H) + # neck (per-level pre-conv) then final projection (only for last level) + # last_aux = self.scratch.output_conv1_aux[-1](last_aux) + last_aux_logits = self.scratch.output_conv2_aux[-1](last_aux) + fmap_last = last_aux_logits.permute(0, 2, 3, 1) + aux_pred = self._apply_activation_single(fmap_last[..., :-1], "linear") + aux_conf = self._apply_activation_single(fmap_last[..., -1], self.conf_activation) + return { + self.head_main: main_pred.squeeze(-1), + f"{self.head_main}_conf": main_conf, + self.head_aux: aux_pred, + f"{self.head_aux}_conf": aux_conf, + } + + # ------------------------------------------------------------------------- + # Subroutines + # ------------------------------------------------------------------------- + + def _fuse(self, feats: List[torch.Tensor]) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """ + Feature pyramid fusion. + Returns: + fused_main: Tensor at finest scale (after refinenet1) + aux_pyr: List of aux tensors at each level (pre out_conv1_aux) + """ + l1, l2, l3, l4 = feats + + l1_rn = self.scratch.layer1_rn(l1) + l2_rn = self.scratch.layer2_rn(l2) + l3_rn = self.scratch.layer3_rn(l3) + l4_rn = self.scratch.layer4_rn(l4) + + # level 4 -> 3 + out = self.scratch.refinenet4(l4_rn, size=l3_rn.shape[2:]) + aux_out = self.scratch.refinenet4_aux(l4_rn, size=l3_rn.shape[2:]) + aux_list: List[torch.Tensor] = [] + if self.aux_levels >= 4: + aux_list.append(aux_out) + + # level 3 -> 2 + out = self.scratch.refinenet3(out, l3_rn, size=l2_rn.shape[2:]) + aux_out = self.scratch.refinenet3_aux(aux_out, l3_rn, size=l2_rn.shape[2:]) + if self.aux_levels >= 3: + aux_list.append(aux_out) + + # level 2 -> 1 + out = self.scratch.refinenet2(out, l2_rn, size=l1_rn.shape[2:]) + aux_out = self.scratch.refinenet2_aux(aux_out, l2_rn, size=l1_rn.shape[2:]) + if self.aux_levels >= 2: + aux_list.append(aux_out) + + # level 1 (final) + out = self.scratch.refinenet1(out, l1_rn) + aux_out = self.scratch.refinenet1_aux(aux_out, l1_rn) + aux_list.append(aux_out) + + out = self.scratch.output_conv1(out) + aux_list = [self.scratch.output_conv1_aux[i](aux) for i, aux in enumerate(aux_list)] + + return out, aux_list + + def _add_pos_embed(self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor: + """Simple UV positional embedding added to feature maps.""" + pw, ph = x.shape[-1], x.shape[-2] + pe = create_uv_grid(pw, ph, aspect_ratio=W / H, dtype=x.dtype, device=x.device) + pe = position_grid_to_embed(pe, x.shape[1]) * ratio + pe = pe.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1) + return x + pe + + def _make_aux_out1_block(self, in_ch: int) -> nn.Sequential: + """Factory for the aux pre-head stack before the final 1x1 projection.""" + if self.aux_out1_conv_num == 5: + return nn.Sequential( + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + nn.Conv2d(in_ch // 2, in_ch, 3, 1, 1), + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + nn.Conv2d(in_ch // 2, in_ch, 3, 1, 1), + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + ) + if self.aux_out1_conv_num == 3: + return nn.Sequential( + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + nn.Conv2d(in_ch // 2, in_ch, 3, 1, 1), + nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1), + ) + if self.aux_out1_conv_num == 1: + return nn.Sequential(nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1)) + raise ValueError(f"aux_out1_conv_num {self.aux_out1_conv_num} not supported") + + def _apply_activation_single( + self, x: torch.Tensor, activation: str = "linear" + ) -> torch.Tensor: + """ + Apply activation to single channel output, maintaining semantic consistency with value branch in multi-channel case. + Supports: exp / relu / sigmoid / softplus / tanh / linear / expp1 + """ + act = activation.lower() if isinstance(activation, str) else activation + if act == "exp": + return torch.exp(x) + if act == "expm1": + return torch.expm1(x) + if act == "expp1": + return torch.exp(x) + 1 + if act == "relu": + return torch.relu(x) + if act == "sigmoid": + return torch.sigmoid(x) + if act == "softplus": + return torch.nn.functional.softplus(x) + if act == "tanh": + return torch.tanh(x) + # Default linear + return x + + +# # ----------------------------------------------------------------------------- +# # Building blocks (tidy) +# # ----------------------------------------------------------------------------- + + +# def _make_fusion_block( +# features: int, +# size: Tuple[int, int] = None, +# has_residual: bool = True, +# groups: int = 1, +# inplace: bool = False, # <- activation uses inplace=True by default; not related to "fusion_inplace" +# ) -> nn.Module: +# return FeatureFusionBlock( +# features=features, +# activation=nn.ReLU(inplace=inplace), +# deconv=False, +# bn=False, +# expand=False, +# align_corners=True, +# size=size, +# has_residual=has_residual, +# groups=groups, +# ) + + +# def _make_scratch( +# in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False +# ) -> nn.Module: +# scratch = nn.Module() +# # optionally expand widths by stage +# c1 = out_shape +# c2 = out_shape * (2 if expand else 1) +# c3 = out_shape * (4 if expand else 1) +# c4 = out_shape * (8 if expand else 1) + +# scratch.layer1_rn = nn.Conv2d(in_shape[0], c1, 3, 1, 1, bias=False, groups=groups) +# scratch.layer2_rn = nn.Conv2d(in_shape[1], c2, 3, 1, 1, bias=False, groups=groups) +# scratch.layer3_rn = nn.Conv2d(in_shape[2], c3, 3, 1, 1, bias=False, groups=groups) +# scratch.layer4_rn = nn.Conv2d(in_shape[3], c4, 3, 1, 1, bias=False, groups=groups) +# return scratch + + +# class ResidualConvUnit(nn.Module): +# """Lightweight residual conv block used within fusion.""" + +# def __init__(self, features: int, activation: nn.Module, bn: bool, groups: int = 1) -> None: +# super().__init__() +# self.bn = bn +# self.groups = groups +# self.conv1 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) +# self.conv2 = nn.Conv2d(features, features, 3, 1, 1, bias=True, groups=groups) +# self.norm1 = None +# self.norm2 = None +# self.activation = activation +# self.skip_add = nn.quantized.FloatFunctional() + +# def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] +# out = self.activation(x) +# out = self.conv1(out) +# if self.norm1 is not None: +# out = self.norm1(out) + +# out = self.activation(out) +# out = self.conv2(out) +# if self.norm2 is not None: +# out = self.norm2(out) + +# return self.skip_add.add(out, x) + + +# class FeatureFusionBlock(nn.Module): +# """Top-down fusion block: (optional) residual merge + upsample + 1x1 shrink.""" + +# def __init__( +# self, +# features: int, +# activation: nn.Module, +# deconv: bool = False, +# bn: bool = False, +# expand: bool = False, +# align_corners: bool = True, +# size: Tuple[int, int] = None, +# has_residual: bool = True, +# groups: int = 1, +# ) -> None: +# super().__init__() +# self.align_corners = align_corners +# self.size = size +# self.has_residual = has_residual + +# self.resConfUnit1 = ( +# ResidualConvUnit(features, activation, bn, groups=groups) if has_residual else None +# ) +# self.resConfUnit2 = ResidualConvUnit(features, activation, bn, groups=groups) + +# out_features = (features // 2) if expand else features +# self.out_conv = nn.Conv2d(features, out_features, 1, 1, 0, bias=True, groups=groups) +# self.skip_add = nn.quantized.FloatFunctional() + +# def forward(self, *xs: torch.Tensor, size: Tuple[int, int] = None) -> torch.Tensor: # type: ignore[override] +# """ +# xs: +# - xs[0]: top input +# - xs[1]: (optional) lateral (to be added with residual) +# """ +# y = xs[0] +# if self.has_residual and len(xs) > 1 and self.resConfUnit1 is not None: +# y = self.skip_add.add(y, self.resConfUnit1(xs[1])) + +# y = self.resConfUnit2(y) + +# # upsample +# if (size is None) and (self.size is None): +# up_kwargs = {"scale_factor": 2} +# elif size is None: +# up_kwargs = {"size": self.size} +# else: +# up_kwargs = {"size": size} + +# y = custom_interpolate(y, **up_kwargs, mode="bilinear", align_corners=self.align_corners) +# y = self.out_conv(y) +# return y diff --git a/src/nn/depth_anything_3/model/gs_adapter.py b/src/nn/depth_anything_3/model/gs_adapter.py new file mode 100644 index 0000000..1335ce3 --- /dev/null +++ b/src/nn/depth_anything_3/model/gs_adapter.py @@ -0,0 +1,200 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional +import torch +from einops import einsum, rearrange, repeat +from torch import nn + +from depth_anything_3.model.utils.transform import cam_quat_xyzw_to_world_quat_wxyz +from depth_anything_3.specs import Gaussians +from depth_anything_3.utils.geometry import affine_inverse, get_world_rays, sample_image_grid +from depth_anything_3.utils.pose_align import batch_align_poses_umeyama +from depth_anything_3.utils.sh_helpers import rotate_sh + + +class GaussianAdapter(nn.Module): + + def __init__( + self, + sh_degree: int = 0, + pred_color: bool = False, + pred_offset_depth: bool = False, + pred_offset_xy: bool = True, + gaussian_scale_min: float = 1e-5, + gaussian_scale_max: float = 30.0, + ): + super().__init__() + self.sh_degree = sh_degree + self.pred_color = pred_color + self.pred_offset_depth = pred_offset_depth + self.pred_offset_xy = pred_offset_xy + self.gaussian_scale_min = gaussian_scale_min + self.gaussian_scale_max = gaussian_scale_max + + # Create a mask for the spherical harmonics coefficients. This ensures that at + # initialization, the coefficients are biased towards having a large DC + # component and small view-dependent components. + if not pred_color: + self.register_buffer( + "sh_mask", + torch.ones((self.d_sh,), dtype=torch.float32), + persistent=False, + ) + for degree in range(1, sh_degree + 1): + self.sh_mask[degree**2 : (degree + 1) ** 2] = 0.1 * 0.25**degree + + def forward( + self, + extrinsics: torch.Tensor, # "*#batch 4 4" + intrinsics: torch.Tensor, # "*#batch 3 3" + depths: torch.Tensor, # "*#batch" + opacities: torch.Tensor, # "*#batch" | "*#batch _" + raw_gaussians: torch.Tensor, # "*#batch _" + image_shape: tuple[int, int], + eps: float = 1e-8, + gt_extrinsics: Optional[torch.Tensor] = None, # "*#batch 4 4" + **kwargs, + ) -> Gaussians: + device = extrinsics.device + dtype = raw_gaussians.dtype + H, W = image_shape + b, v = raw_gaussians.shape[:2] + + # get cam2worlds and intr_normed to adapt to 3DGS codebase + cam2worlds = affine_inverse(extrinsics) + intr_normed = intrinsics.clone().detach() + intr_normed[..., 0, :] /= W + intr_normed[..., 1, :] /= H + + # 1. compute 3DGS means + # 1.1) offset the predicted depth if needed + if self.pred_offset_depth: + gs_depths = depths + raw_gaussians[..., -1] + raw_gaussians = raw_gaussians[..., :-1] + else: + gs_depths = depths + # 1.2) align predicted poses with GT if needed + if gt_extrinsics is not None and not torch.equal(extrinsics, gt_extrinsics): + try: + _, _, pose_scales = batch_align_poses_umeyama( + gt_extrinsics.detach().float(), + extrinsics.detach().float(), + ) + except Exception: + pose_scales = torch.ones_like(extrinsics[:, 0, 0, 0]) + pose_scales = torch.clamp(pose_scales, min=1 / 3.0, max=3.0) + cam2worlds[:, :, :3, 3] = cam2worlds[:, :, :3, 3] * rearrange( + pose_scales, "b -> b () ()" + ) # [b, i, j] + gs_depths = gs_depths * rearrange(pose_scales, "b -> b () () ()") # [b, v, h, w] + # 1.3) casting xy in image space + xy_ray, _ = sample_image_grid((H, W), device) + xy_ray = xy_ray[None, None, ...].expand(b, v, -1, -1, -1) # b v h w xy + # offset xy if needed + if self.pred_offset_xy: + pixel_size = 1 / torch.tensor((W, H), dtype=xy_ray.dtype, device=device) + offset_xy = raw_gaussians[..., :2] + xy_ray = xy_ray + offset_xy * pixel_size + raw_gaussians = raw_gaussians[..., 2:] # skip the offset_xy + # 1.4) unproject depth + xy to world ray + origins, directions = get_world_rays( + xy_ray, + repeat(cam2worlds, "b v i j -> b v h w i j", h=H, w=W), + repeat(intr_normed, "b v i j -> b v h w i j", h=H, w=W), + ) + gs_means_world = origins + directions * gs_depths[..., None] + gs_means_world = rearrange(gs_means_world, "b v h w d -> b (v h w) d") + + # 2. compute other GS attributes + scales, rotations, sh = raw_gaussians.split((3, 4, 3 * self.d_sh), dim=-1) + + # 2.1) 3DGS scales + # make the scale invarient to resolution + scale_min = self.gaussian_scale_min + scale_max = self.gaussian_scale_max + scales = scale_min + (scale_max - scale_min) * scales.sigmoid() + pixel_size = 1 / torch.tensor((W, H), dtype=dtype, device=device) + multiplier = self.get_scale_multiplier(intr_normed, pixel_size) + gs_scales = scales * gs_depths[..., None] * multiplier[..., None, None, None] + gs_scales = rearrange(gs_scales, "b v h w d -> b (v h w) d") + + # 2.2) 3DGS quaternion (world space) + # due to historical issue, assume quaternion in order xyzw, not wxyz + # Normalize the quaternion features to yield a valid quaternion. + rotations = rotations / (rotations.norm(dim=-1, keepdim=True) + eps) + # rotate them to world space + cam_quat_xyzw = rearrange(rotations, "b v h w c -> b (v h w) c") + c2w_mat = repeat( + cam2worlds, + "b v i j -> b (v h w) i j", + h=H, + w=W, + ) + world_quat_wxyz = cam_quat_xyzw_to_world_quat_wxyz(cam_quat_xyzw, c2w_mat) + gs_rotations_world = world_quat_wxyz # b (v h w) c + + # 2.3) 3DGS color / SH coefficient (world space) + sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3) + if not self.pred_color: + sh = sh * self.sh_mask + + if self.pred_color or self.sh_degree == 0: + # predict pre-computed color or predict only DC band, no need to transform + gs_sh_world = sh + else: + gs_sh_world = rotate_sh(sh, cam2worlds[:, :, None, None, None, :3, :3]) + gs_sh_world = rearrange(gs_sh_world, "b v h w xyz d_sh -> b (v h w) xyz d_sh") + + # 2.4) 3DGS opacity + gs_opacities = rearrange(opacities, "b v h w ... -> b (v h w) ...") + + return Gaussians( + means=gs_means_world, + harmonics=gs_sh_world, + opacities=gs_opacities, + scales=gs_scales, + rotations=gs_rotations_world, + ) + + def get_scale_multiplier( + self, + intrinsics: torch.Tensor, # "*#batch 3 3" + pixel_size: torch.Tensor, # "*#batch 2" + multiplier: float = 0.1, + ) -> torch.Tensor: # " *batch" + xy_multipliers = multiplier * einsum( + intrinsics[..., :2, :2].float().inverse().to(intrinsics), + pixel_size, + "... i j, j -> ... i", + ) + return xy_multipliers.sum(dim=-1) + + @property + def d_sh(self) -> int: + return 1 if self.pred_color else (self.sh_degree + 1) ** 2 + + @property + def d_in(self) -> int: + # provided as reference to the gs_dpt output dim + raw_gs_dim = 0 + if self.pred_offset_xy: + raw_gs_dim += 2 + raw_gs_dim += 3 # scales + raw_gs_dim += 4 # quaternion + raw_gs_dim += 3 * self.d_sh # color + if self.pred_offset_depth: + raw_gs_dim += 1 + + return raw_gs_dim diff --git a/src/nn/depth_anything_3/model/gsdpt.py b/src/nn/depth_anything_3/model/gsdpt.py new file mode 100644 index 0000000..70448b7 --- /dev/null +++ b/src/nn/depth_anything_3/model/gsdpt.py @@ -0,0 +1,133 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict as TyDict +from typing import List, Sequence +import torch +import torch.nn as nn + +from depth_anything_3.model.dpt import DPT +from depth_anything_3.model.utils.head_utils import activate_head_gs, custom_interpolate + + +class GSDPT(DPT): + + def __init__( + self, + dim_in: int, + patch_size: int = 14, + output_dim: int = 4, + activation: str = "linear", + conf_activation: str = "sigmoid", + features: int = 256, + out_channels: Sequence[int] = (256, 512, 1024, 1024), + pos_embed: bool = True, + feature_only: bool = False, + down_ratio: int = 1, + conf_dim: int = 1, + norm_type: str = "idt", # use to match legacy GS-DPT head, "idt" / "layer" + fusion_block_inplace: bool = False, + ) -> None: + super().__init__( + dim_in=dim_in, + patch_size=patch_size, + output_dim=output_dim, + activation=activation, + conf_activation=conf_activation, + features=features, + out_channels=out_channels, + pos_embed=pos_embed, + down_ratio=down_ratio, + head_name="raw_gs", + use_sky_head=False, + norm_type=norm_type, + fusion_block_inplace=fusion_block_inplace, + ) + self.conf_dim = conf_dim + if conf_dim and conf_dim > 1: + assert ( + conf_activation == "linear" + ), "use linear prediction when using view-dependent opacity" + + merger_out_dim = features if feature_only else features // 2 + self.images_merger = nn.Sequential( + nn.Conv2d(3, merger_out_dim // 4, 3, 1, 1), # fewer channels first + nn.GELU(), + nn.Conv2d(merger_out_dim // 4, merger_out_dim // 2, 3, 1, 1), + nn.GELU(), + nn.Conv2d(merger_out_dim // 2, merger_out_dim, 3, 1, 1), + nn.GELU(), + ) + + # ------------------------------------------------------------------------- + # Internal forward (single chunk) + # ------------------------------------------------------------------------- + def _forward_impl( + self, + feats: List[torch.Tensor], + H: int, + W: int, + patch_start_idx: int, + images: torch.Tensor, + ) -> TyDict[str, torch.Tensor]: + B, _, C = feats[0].shape + ph, pw = H // self.patch_size, W // self.patch_size + resized_feats = [] + for stage_idx, take_idx in enumerate(self.intermediate_layer_idx): + x = feats[take_idx][:, patch_start_idx:] # [B*S, N_patch, C] + x = self.norm(x) + x = x.permute(0, 2, 1).reshape(B, C, ph, pw) # [B*S, C, ph, pw] + + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + x = self.resize_layers[stage_idx](x) # Align scale + resized_feats.append(x) + + # 2) Fusion pyramid (main branch only) + fused = self._fuse(resized_feats) + fused = self.scratch.output_conv1(fused) + + # 3) Upsample to target resolution, optionally add position encoding again + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused = custom_interpolate(fused, (h_out, w_out), mode="bilinear", align_corners=True) + + # inject the image information here + fused = fused + self.images_merger(images) + + if self.pos_embed: + fused = self._add_pos_embed(fused, W, H) + + # 4) Shared neck1 + # feat = self.scratch.output_conv1(fused) + feat = fused + + # 5) Main head: logits -> activate_head or single channel activation + main_logits = self.scratch.output_conv2(feat) + outs: TyDict[str, torch.Tensor] = {} + if self.has_conf: + pred, conf = activate_head_gs( + main_logits, + activation=self.activation, + conf_activation=self.conf_activation, + conf_dim=self.conf_dim, + ) + outs[self.head_main] = pred.squeeze(1) + outs[f"{self.head_main}_conf"] = conf.squeeze(1) + else: + outs[self.head_main] = self._apply_activation_single(main_logits).squeeze(1) + + return outs diff --git a/src/nn/depth_anything_3/model/reference_view_selector.py b/src/nn/depth_anything_3/model/reference_view_selector.py new file mode 100644 index 0000000..f406f5f --- /dev/null +++ b/src/nn/depth_anything_3/model/reference_view_selector.py @@ -0,0 +1,223 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Reference View Selection Strategies + +This module provides different strategies for selecting a reference view +from multiple input views in multi-view depth estimation. +""" + +import torch +from typing import Literal + + +RefViewStrategy = Literal["first", "middle", "saddle_balanced", "saddle_sim_range"] + + +def select_reference_view( + x: torch.Tensor, + strategy: RefViewStrategy = "saddle_balanced", +) -> torch.Tensor: + """ + Select a reference view from multiple views using the specified strategy. + + Args: + x: Input tensor of shape (B, S, N, C) where + B = batch size + S = number of views + N = number of tokens + C = channel dimension + strategy: Selection strategy, one of: + - "first": Always select the first view + - "middle": Select the middle view + - "saddle_balanced": Select view with balanced features across multiple metrics + - "saddle_sim_range": Select view with largest similarity range + + Returns: + b_idx: Tensor of shape (B,) containing the selected view index for each batch + """ + B, S, N, C = x.shape + + # For single view, no reordering needed + if S <= 1: + return torch.zeros(B, dtype=torch.long, device=x.device) + + # Simple position-based strategies + if strategy == "first": + return torch.zeros(B, dtype=torch.long, device=x.device) + + elif strategy == "middle": + return torch.full((B,), S // 2, dtype=torch.long, device=x.device) + + # Feature-based strategies require normalized class tokens + # Extract and normalize class tokens (first token of each view) + img_class_feat = x[:, :, 0] / x[:, :, 0].norm(dim=-1, keepdim=True) # B S C + + if strategy == "saddle_balanced": + # Select view with balanced features across multiple metrics + # Compute similarity matrix + sim = torch.matmul(img_class_feat, img_class_feat.transpose(1, 2)) # B S S + sim_no_diag = sim - torch.eye(S, device=sim.device).unsqueeze(0) + sim_score = sim_no_diag.sum(dim=-1) / (S - 1) # B S + + feat_norm = x[:, :, 0].norm(dim=-1) # B S + feat_var = img_class_feat.var(dim=-1) # B S + + # Normalize all metrics to [0, 1] + def normalize_metric(metric): + min_val = metric.min(dim=1, keepdim=True).values + max_val = metric.max(dim=1, keepdim=True).values + return (metric - min_val) / (max_val - min_val + 1e-8) + + sim_score_norm = normalize_metric(sim_score) + norm_norm = normalize_metric(feat_norm) + var_norm = normalize_metric(feat_var) + + # Select view closest to the median (0.5) across all metrics + balance_score = ( + (sim_score_norm - 0.5).abs() + + (norm_norm - 0.5).abs() + + (var_norm - 0.5).abs() + ) + b_idx = balance_score.argmin(dim=1) + + elif strategy == "saddle_sim_range": + # Select view with largest similarity range (max - min) + sim = torch.matmul(img_class_feat, img_class_feat.transpose(1, 2)) # B S S + sim_no_diag = sim - torch.eye(S, device=sim.device).unsqueeze(0) + + sim_max = sim_no_diag.max(dim=-1).values # B S + sim_min = sim_no_diag.min(dim=-1).values # B S + sim_range = sim_max - sim_min + b_idx = sim_range.argmax(dim=1) + + else: + raise ValueError( + f"Unknown reference view selection strategy: {strategy}. " + f"Must be one of: 'first', 'middle', 'saddle_balanced', 'saddle_sim_range'" + ) + + return b_idx + + +def reorder_by_reference( + x: torch.Tensor, + b_idx: torch.Tensor, +) -> torch.Tensor: + """ + Reorder views to place the selected reference view first. + + Args: + x: Input tensor of shape (B, S, N, C) + b_idx: Reference view indices of shape (B,) + + Returns: + Reordered tensor with reference view at position 0 + + Example: + If b_idx = [2] and S = 5 (views [0,1,2,3,4]), + result order is [2,0,1,3,4] (ref_idx first, then others in order) + """ + B, S = x.shape[0], x.shape[1] + + # For single view, no reordering needed + if S <= 1: + return x + + # Create position indices: (B, S) where each row is [0, 1, 2, ..., S-1] + positions = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1) # B S + + # For each position, determine which original index it should take + # Position 0 gets ref_idx + # Position 1 to ref_idx gets indices 0 to ref_idx-1 + # Position ref_idx+1 to S-1 gets indices ref_idx+1 to S-1 + + b_idx_expanded = b_idx.unsqueeze(1) # B 1 + + # Create the reordering indices + # For positions 1 to ref_idx: map to indices 0 to ref_idx-1 (shift by -1) + # For positions > ref_idx: keep the same + reorder_indices = positions.clone() + reorder_indices = torch.where( + (positions > 0) & (positions <= b_idx_expanded), + positions - 1, + positions + ) + # Set position 0 to ref_idx + reorder_indices[:, 0] = b_idx + + # Gather using advanced indexing + batch_indices = torch.arange(B, device=x.device).unsqueeze(1) # B 1 + x_reordered = x[batch_indices, reorder_indices] + + return x_reordered + + +def restore_original_order( + x: torch.Tensor, + b_idx: torch.Tensor, +) -> torch.Tensor: + """ + Restore original view order after processing. + + Args: + x: Reordered tensor of shape (B, S, ...) + b_idx: Original reference view indices of shape (B,) + + Returns: + Tensor with original view order restored + + Example: + If original order was [0, 1, 2, 3, 4] and b_idx=2, + reordered becomes [2, 0, 1, 3, 4] (reference at position 0), + restore should return [0, 1, 2, 3, 4] (original order). + """ + B, S = x.shape[0], x.shape[1] + + # For single view, no restoration needed + if S <= 1: + return x + + # Create target position indices: (B, S) where each row is [0, 1, 2, ..., S-1] + target_positions = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1) # B S + + # For each target position, determine which current position it comes from + # Target position 0 to ref_idx-1 <- Current position 1 to ref_idx (shift by +1) + # Target position ref_idx <- Current position 0 + # Target position ref_idx+1 to S-1 <- Current position ref_idx+1 to S-1 (no change) + + b_idx_expanded = b_idx.unsqueeze(1) # B 1 + + # Create the restore indices + restore_indices = torch.where( + target_positions < b_idx_expanded, + target_positions + 1, # Positions before ref_idx come from current position + 1 + target_positions # Positions after ref_idx stay the same + ) + # Target position = ref_idx comes from current position 0 + # Use scatter to set specific positions + restore_indices = torch.scatter( + restore_indices, + dim=1, + index=b_idx_expanded, + src=torch.zeros_like(b_idx_expanded) + ) + + # Gather using advanced indexing + batch_indices = torch.arange(B, device=x.device).unsqueeze(1) # B 1 + x_restored = x[batch_indices, restore_indices] + + return x_restored + diff --git a/src/nn/depth_anything_3/model/utils/attention.py b/src/nn/depth_anything_3/model/utils/attention.py new file mode 100644 index 0000000..49c07a8 --- /dev/null +++ b/src/nn/depth_anything_3/model/utils/attention.py @@ -0,0 +1,109 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 # noqa + +from typing import Callable, Optional, Union +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + qk_norm: bool = False, + rope=None, + ) -> None: + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + self.rope = rope + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + # Debug breakpoint removed for production + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + q = self.rope(q, pos) if self.rope is not None else q + k = self.rope(k, pos) if self.rope is not None else k + x = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.attn_drop.p if self.training else 0.0, + attn_mask=attn_mask, + ) + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/src/nn/depth_anything_3/model/utils/block.py b/src/nn/depth_anything_3/model/utils/block.py new file mode 100644 index 0000000..993fb4c --- /dev/null +++ b/src/nn/depth_anything_3/model/utils/block.py @@ -0,0 +1,81 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import Callable +from torch import Tensor, nn + +from .attention import Attention, LayerScale, Mlp + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + qk_norm: bool = False, + rope=None, + ) -> None: + super().__init__() + + self.norm1 = norm_layer(dim) + + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + qk_norm=qk_norm, + rope=rope, + ) + + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + + self.sample_drop_ratio = 0.0 # Equivalent to always having drop_path=0 + + def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: + def attn_residual_func(x: Tensor, pos=None, attn_mask=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask)) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + # drop_path is always 0, so always take the else branch + x = x + attn_residual_func(x, pos=pos, attn_mask=attn_mask) + x = x + ffn_residual_func(x) + return x diff --git a/src/nn/depth_anything_3/model/utils/gs_renderer.py b/src/nn/depth_anything_3/model/utils/gs_renderer.py new file mode 100644 index 0000000..91414d3 --- /dev/null +++ b/src/nn/depth_anything_3/model/utils/gs_renderer.py @@ -0,0 +1,340 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from math import isqrt +from typing import Literal, Optional +import torch +from einops import rearrange, repeat +from tqdm import tqdm + +from depth_anything_3.specs import Gaussians +from depth_anything_3.utils.camera_trj_helpers import ( + interpolate_extrinsics, + interpolate_intrinsics, + render_dolly_zoom_path, + render_stabilization_path, + render_wander_path, + render_wobble_inter_path, +) +from depth_anything_3.utils.geometry import affine_inverse, as_homogeneous, get_fov +from depth_anything_3.utils.logger import logger + +try: + from gsplat import rasterization +except ImportError: + logger.warn( + "Dependency `gsplat` is required for rendering 3DGS. " + "Install via: pip install git+https://github.com/nerfstudio-project/" + "gsplat.git@0b4dddf04cb687367602c01196913cde6a743d70" + ) + + +def render_3dgs( + extrinsics: torch.Tensor, # "batch_views 4 4", w2c + intrinsics: torch.Tensor, # "batch_views 3 3", normalized + image_shape: tuple[int, int], + gaussian: Gaussians, + background_color: Optional[torch.Tensor] = None, # "batch_views 3" + use_sh: bool = True, + num_view: int = 1, + color_mode: Literal["RGB+D", "RGB+ED"] = "RGB+D", + **kwargs, +) -> tuple[ + torch.Tensor, # "batch_views 3 height width" + torch.Tensor, # "batch_views height width" +]: + # extract gaussian params + gaussian_means = gaussian.means + gaussian_scales = gaussian.scales + gaussian_quats = gaussian.rotations + gaussian_opacities = gaussian.opacities + gaussian_sh_coefficients = gaussian.harmonics + b, _, _ = extrinsics.shape + + if background_color is None: + background_color = repeat(torch.tensor([0.0, 0.0, 0.0]), "c -> b c", b=b).to( + gaussian_sh_coefficients + ) + + if use_sh: + _, _, _, n = gaussian_sh_coefficients.shape + degree = isqrt(n) - 1 + shs = rearrange(gaussian_sh_coefficients, "b g xyz n -> b g n xyz").contiguous() + else: # use color + shs = ( + gaussian_sh_coefficients.squeeze(-1).sigmoid().contiguous() + ) # (b, g, c), normed to (0, 1) + + h, w = image_shape + + fov_x, fov_y = get_fov(intrinsics).unbind(dim=-1) + tan_fov_x = (0.5 * fov_x).tan() + tan_fov_y = (0.5 * fov_y).tan() + focal_length_x = w / (2 * tan_fov_x) + focal_length_y = h / (2 * tan_fov_y) + + view_matrix = extrinsics.float() + + all_images = [] + all_radii = [] + all_depths = [] + # render view in a batch based, each batch contains one scene + # assume the Gaussian parameters are originally repeated along the view dim + batch_scene = b // num_view + + def index_i_gs_attr(full_attr, idx): + # return rearrange(full_attr, "(b v) ... -> b v ...", v=num_view)[idx, 0] + return full_attr[idx] + + for i in range(batch_scene): + K = repeat( + torch.tensor( + [ + [0, 0, w / 2.0], + [0, 0, h / 2.0], + [0, 0, 1], + ] + ), + "i j -> v i j", + v=num_view, + ).to(gaussian_means) + K[:, 0, 0] = focal_length_x.reshape(batch_scene, num_view)[i] + K[:, 1, 1] = focal_length_y.reshape(batch_scene, num_view)[i] + + i_means = index_i_gs_attr(gaussian_means, i) # [N, 3] + i_scales = index_i_gs_attr(gaussian_scales, i) + i_quats = index_i_gs_attr(gaussian_quats, i) + i_opacities = index_i_gs_attr(gaussian_opacities, i) # [N,] + i_colors = index_i_gs_attr(shs, i) # [N, K, 3] + i_viewmats = rearrange(view_matrix, "(b v) ... -> b v ...", v=num_view)[i] # [v, 4, 4] + i_backgrounds = rearrange(background_color, "(b v) ... -> b v ...", v=num_view)[ + i + ] # [v, 3] + + render_colors, render_alphas, info = rasterization( + means=i_means, + quats=i_quats, # [N, 4] + scales=i_scales, # [N, 3] + opacities=i_opacities, + colors=i_colors, + viewmats=i_viewmats, # [v, 4, 4] + Ks=K, # [v, 3, 3] + backgrounds=i_backgrounds, + render_mode=color_mode, + width=w, + height=h, + packed=False, + sh_degree=degree if use_sh else None, + ) + depth = render_colors[..., -1].unbind(dim=0) + + image = rearrange(render_colors[..., :3], "v h w c -> v c h w").unbind(dim=0) + radii = info["radii"].unbind(dim=0) + try: + info["means2d"].retain_grad() # [1, N, 2] + except Exception: + pass + all_images.extend(image) + all_depths.extend(depth) + all_radii.extend(radii) + + return torch.stack(all_images), torch.stack(all_depths) + + +def run_renderer_in_chunk_w_trj_mode( + gaussians: Gaussians, + extrinsics: torch.Tensor, # world2cam, "batch view 4 4" | "batch view 3 4" + intrinsics: torch.Tensor, # unnormed intrinsics, "batch view 3 3" + image_shape: tuple[int, int], + chunk_size: Optional[int] = 8, + trj_mode: Literal[ + "original", + "smooth", + "interpolate", + "interpolate_smooth", + "wander", + "dolly_zoom", + "extend", + "wobble_inter", + ] = "smooth", + input_shape: Optional[tuple[int, int]] = None, + enable_tqdm: Optional[bool] = False, + **kwargs, +) -> tuple[ + torch.Tensor, # color, "batch view 3 height width" + torch.Tensor, # depth, "batch view height width" +]: + cam2world = affine_inverse(as_homogeneous(extrinsics)) + if input_shape is not None: + in_h, in_w = input_shape + else: + in_h, in_w = image_shape + intr_normed = intrinsics.clone().detach() + intr_normed[..., 0, :] /= in_w + intr_normed[..., 1, :] /= in_h + if extrinsics.shape[1] <= 1: + assert trj_mode in [ + "wander", + "dolly_zoom", + ], "Please set trj_mode to 'wander' or 'dolly_zoom' when n_views=1" + + def _smooth_trj_fn_batch(raw_c2ws, k_size=50): + try: + smooth_c2ws = torch.stack( + [render_stabilization_path(c2w_i, k_size) for c2w_i in raw_c2ws], + dim=0, + ) + except Exception as e: + print(f"[DEBUG] Path smoothing failed with error: {e}.") + smooth_c2ws = raw_c2ws + return smooth_c2ws + + # get rendered trj + if trj_mode == "original": + tgt_c2w = cam2world + tgt_intr = intr_normed + elif trj_mode == "smooth": + tgt_c2w = _smooth_trj_fn_batch(cam2world) + tgt_intr = intr_normed + elif trj_mode in ["interpolate", "interpolate_smooth", "extend"]: + inter_len = 8 + total_len = (cam2world.shape[1] - 1) * inter_len + if total_len > 24 * 18: # no more than 18s + inter_len = max(1, 24 * 10 // (cam2world.shape[1] - 1)) + if total_len < 24 * 2: # no less than 2s + inter_len = max(1, 24 * 2 // (cam2world.shape[1] - 1)) + + if inter_len > 2: + t = torch.linspace(0, 1, inter_len, dtype=torch.float32, device=cam2world.device) + t = (torch.cos(torch.pi * (t + 1)) + 1) / 2 + tgt_c2w_b = [] + tgt_intr_b = [] + for b_idx in range(cam2world.shape[0]): + tgt_c2w = [] + tgt_intr = [] + for cur_idx in range(cam2world.shape[1] - 1): + tgt_c2w.append( + interpolate_extrinsics( + cam2world[b_idx, cur_idx], cam2world[b_idx, cur_idx + 1], t + )[(0 if cur_idx == 0 else 1) :] + ) + tgt_intr.append( + interpolate_intrinsics( + intr_normed[b_idx, cur_idx], intr_normed[b_idx, cur_idx + 1], t + )[(0 if cur_idx == 0 else 1) :] + ) + tgt_c2w_b.append(torch.cat(tgt_c2w)) + tgt_intr_b.append(torch.cat(tgt_intr)) + tgt_c2w = torch.stack(tgt_c2w_b) # b v 4 4 + tgt_intr = torch.stack(tgt_intr_b) # b v 3 3 + else: + tgt_c2w = cam2world + tgt_intr = intr_normed + if trj_mode in ["interpolate_smooth", "extend"]: + tgt_c2w = _smooth_trj_fn_batch(tgt_c2w) + if trj_mode == "extend": + # apply dolly_zoom and wander in the middle frame + assert cam2world.shape[0] == 1, "extend only supports for batch_size=1 currently." + mid_idx = tgt_c2w.shape[1] // 2 + c2w_wd, intr_wd = render_wander_path( + tgt_c2w[0, mid_idx], + tgt_intr[0, mid_idx], + h=in_h, + w=in_w, + num_frames=max(36, min(60, mid_idx // 2)), + max_disp=24.0, + ) + c2w_dz, intr_dz = render_dolly_zoom_path( + tgt_c2w[0, mid_idx], + tgt_intr[0, mid_idx], + h=in_h, + w=in_w, + num_frames=max(36, min(60, mid_idx // 2)), + ) + tgt_c2w = torch.cat( + [ + tgt_c2w[:, :mid_idx], + c2w_wd.unsqueeze(0), + c2w_dz.unsqueeze(0), + tgt_c2w[:, mid_idx:], + ], + dim=1, + ) + tgt_intr = torch.cat( + [ + tgt_intr[:, :mid_idx], + intr_wd.unsqueeze(0), + intr_dz.unsqueeze(0), + tgt_intr[:, mid_idx:], + ], + dim=1, + ) + elif trj_mode in ["wander", "dolly_zoom"]: + if trj_mode == "wander": + render_fn = render_wander_path + extra_kwargs = {"max_disp": 24.0} + else: + render_fn = render_dolly_zoom_path + extra_kwargs = {"D_focus": 30.0, "max_disp": 2.0} + tgt_c2w = [] + tgt_intr = [] + for b_idx in range(cam2world.shape[0]): + c2w_i, intr_i = render_fn( + cam2world[b_idx, 0], intr_normed[b_idx, 0], h=in_h, w=in_w, **extra_kwargs + ) + tgt_c2w.append(c2w_i) + tgt_intr.append(intr_i) + tgt_c2w = torch.stack(tgt_c2w) + tgt_intr = torch.stack(tgt_intr) + elif trj_mode == "wobble_inter": + tgt_c2w, tgt_intr = render_wobble_inter_path( + cam2world=cam2world, + intr_normed=intr_normed, + inter_len=10, + n_skip=3, + ) + else: + raise Exception(f"trj mode [{trj_mode}] is not implemented.") + + _, v = tgt_c2w.shape[:2] + tgt_extr = affine_inverse(tgt_c2w) + if chunk_size is None: + chunk_size = v + chunk_size = min(v, chunk_size) + all_colors = [] + all_depths = [] + for chunk_idx in tqdm( + range(math.ceil(v / chunk_size)), + desc="Rendering novel views", + disable=(not enable_tqdm), + leave=False, + ): + s = int(chunk_idx * chunk_size) + e = int((chunk_idx + 1) * chunk_size) + cur_n_view = tgt_extr[:, s:e].shape[1] + color, depth = render_3dgs( + extrinsics=rearrange(tgt_extr[:, s:e], "b v ... -> (b v) ..."), # w2c + intrinsics=rearrange(tgt_intr[:, s:e], "b v ... -> (b v) ..."), # normed + image_shape=image_shape, + gaussian=gaussians, + num_view=cur_n_view, + **kwargs, + ) + all_colors.append(rearrange(color, "(b v) ... -> b v ...", v=cur_n_view)) + all_depths.append(rearrange(depth, "(b v) ... -> b v ...", v=cur_n_view)) + all_colors = torch.cat(all_colors, dim=1) + all_depths = torch.cat(all_depths, dim=1) + + return all_colors, all_depths diff --git a/src/nn/depth_anything_3/model/utils/head_utils.py b/src/nn/depth_anything_3/model/utils/head_utils.py new file mode 100644 index 0000000..c120958 --- /dev/null +++ b/src/nn/depth_anything_3/model/utils/head_utils.py @@ -0,0 +1,230 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Tuple, Union +import torch +import torch.nn as nn +import torch.nn.functional as F + +# ----------------------------------------------------------------------------- +# Activation functions +# ----------------------------------------------------------------------------- + + +def activate_head_gs(out, activation="norm_exp", conf_activation="expp1", conf_dim=None): + """ + Process network output to extract GS params and density values. + Density could be view-dependent as SH coefficient + + + Args: + out: Network output tensor (B, C, H, W) + activation: Activation type for 3D points + conf_activation: Activation type for confidence values + + Returns: + Tuple of (3D points tensor, confidence tensor) + """ + # Move channels from last dim to the 4th dimension => (B, H, W, C) + fmap = out.permute(0, 2, 3, 1) # B,H,W,C expected + + # Split into xyz (first C-1 channels) and confidence (last channel) + conf_dim = 1 if conf_dim is None else conf_dim + xyz = fmap[:, :, :, :-conf_dim] + conf = fmap[:, :, :, -1] if conf_dim == 1 else fmap[:, :, :, -conf_dim:] + + if activation == "norm_exp": + d = xyz.norm(dim=-1, keepdim=True).clamp(min=1e-8) + xyz_normed = xyz / d + pts3d = xyz_normed * torch.expm1(d) + elif activation == "norm": + pts3d = xyz / xyz.norm(dim=-1, keepdim=True) + elif activation == "exp": + pts3d = torch.exp(xyz) + elif activation == "relu": + pts3d = F.relu(xyz) + elif activation == "sigmoid": + pts3d = torch.sigmoid(xyz) + elif activation == "linear": + pts3d = xyz + else: + raise ValueError(f"Unknown activation: {activation}") + + if conf_activation == "expp1": + conf_out = 1 + conf.exp() + elif conf_activation == "expp0": + conf_out = conf.exp() + elif conf_activation == "sigmoid": + conf_out = torch.sigmoid(conf) + elif conf_activation == "linear": + conf_out = conf + else: + raise ValueError(f"Unknown conf_activation: {conf_activation}") + + return pts3d, conf_out + + +# ----------------------------------------------------------------------------- +# Other utilities +# ----------------------------------------------------------------------------- + + +class Permute(nn.Module): + """nn.Module wrapper around Tensor.permute for cleaner nn.Sequential usage.""" + + dims: Tuple[int, ...] + + def __init__(self, dims: Tuple[int, ...]) -> None: + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] + return x.permute(*self.dims) + + +def position_grid_to_embed( + pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100 +) -> torch.Tensor: + """ + Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC) + + Args: + pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates + embed_dim: Output channel dimension for embeddings + + Returns: + Tensor of shape (H, W, embed_dim) with positional embeddings + """ + H, W, grid_dim = pos_grid.shape + assert grid_dim == 2 + pos_flat = pos_grid.reshape(-1, grid_dim) # Flatten to (H*W, 2) + + # Process x and y coordinates separately + emb_x = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 0], omega_0=omega_0) # [1, H*W, D/2] + emb_y = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 1], omega_0=omega_0) # [1, H*W, D/2] + + # Combine and reshape + emb = torch.cat([emb_x, emb_y], dim=-1) # [1, H*W, D] + + return emb.view(H, W, embed_dim) # [H, W, D] + + +def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 100) -> torch.Tensor: + """ + This function generates a 1D positional embedding from a given grid using sine and cosine functions. # noqa + + Args: + - embed_dim: The embedding dimension. + - pos: The position to generate the embedding from. + + Returns: + - emb: The generated 1D positional embedding. + """ + assert embed_dim % 2 == 0 + omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=pos.device) + omega /= embed_dim / 2.0 + omega = 1.0 / omega_0**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = torch.sin(out) # (M, D/2) + emb_cos = torch.cos(out) # (M, D/2) + + emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D) + return emb.float() + + +# Inspired by https://github.com/microsoft/moge + + +def create_uv_grid( + width: int, + height: int, + aspect_ratio: float = None, + dtype: torch.dtype = None, + device: torch.device = None, +) -> torch.Tensor: + """ + Create a normalized UV grid of shape (width, height, 2). + + The grid spans horizontally and vertically according to an aspect ratio, + ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right + corner is at (x_span, y_span), normalized by the diagonal of the plane. + + Args: + width (int): Number of points horizontally. + height (int): Number of points vertically. + aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height. + dtype (torch.dtype, optional): Data type of the resulting tensor. + device (torch.device, optional): Device on which the tensor is created. + + Returns: + torch.Tensor: A (width, height, 2) tensor of UV coordinates. + """ + # Derive aspect ratio if not explicitly provided + if aspect_ratio is None: + aspect_ratio = float(width) / float(height) + + # Compute normalized spans for X and Y + diag_factor = (aspect_ratio**2 + 1.0) ** 0.5 + span_x = aspect_ratio / diag_factor + span_y = 1.0 / diag_factor + + # Establish the linspace boundaries + left_x = -span_x * (width - 1) / width + right_x = span_x * (width - 1) / width + top_y = -span_y * (height - 1) / height + bottom_y = span_y * (height - 1) / height + + # Generate 1D coordinates + x_coords = torch.linspace(left_x, right_x, steps=width, dtype=dtype, device=device) + y_coords = torch.linspace(top_y, bottom_y, steps=height, dtype=dtype, device=device) + + # Create 2D meshgrid (width x height) and stack into UV + uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy") + uv_grid = torch.stack((uu, vv), dim=-1) + + return uv_grid + + +# ----------------------------------------------------------------------------- +# Interpolation (safe interpolation, avoid INT_MAX overflow) +# ----------------------------------------------------------------------------- +def custom_interpolate( + x: torch.Tensor, + size: Union[Tuple[int, int], None] = None, + scale_factor: Union[float, None] = None, + mode: str = "bilinear", + align_corners: bool = True, +) -> torch.Tensor: + """ + Safe interpolation implementation to avoid INT_MAX overflow in torch.nn.functional.interpolate. + """ + if size is None: + assert scale_factor is not None, "Either size or scale_factor must be provided." + size = (int(x.shape[-2] * scale_factor), int(x.shape[-1] * scale_factor)) + + INT_MAX = 1610612736 + total = size[0] * size[1] * x.shape[0] * x.shape[1] + + if total > INT_MAX: + chunks = torch.chunk(x, chunks=(total // INT_MAX) + 1, dim=0) + outs = [ + nn.functional.interpolate(c, size=size, mode=mode, align_corners=align_corners) + for c in chunks + ] + return torch.cat(outs, dim=0).contiguous() + + return nn.functional.interpolate(x, size=size, mode=mode, align_corners=align_corners) diff --git a/src/nn/depth_anything_3/model/utils/transform.py b/src/nn/depth_anything_3/model/utils/transform.py new file mode 100644 index 0000000..8d732b0 --- /dev/null +++ b/src/nn/depth_anything_3/model/utils/transform.py @@ -0,0 +1,208 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn.functional as F + + +def extri_intri_to_pose_encoding( + extrinsics, + intrinsics, + image_size_hw=None, +): + """Convert camera extrinsics and intrinsics to a compact pose encoding.""" + + # extrinsics: BxSx3x4 + # intrinsics: BxSx3x3 + R = extrinsics[:, :, :3, :3] # BxSx3x3 + T = extrinsics[:, :, :3, 3] # BxSx3 + + quat = mat_to_quat(R) + # Note the order of h and w here + H, W = image_size_hw + fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1]) + fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0]) + pose_encoding = torch.cat([T, quat, fov_h[..., None], fov_w[..., None]], dim=-1).float() + + return pose_encoding + + +def pose_encoding_to_extri_intri( + pose_encoding, + image_size_hw=None, +): + """Convert a pose encoding back to camera extrinsics and intrinsics.""" + + T = pose_encoding[..., :3] + quat = pose_encoding[..., 3:7] + fov_h = pose_encoding[..., 7] + fov_w = pose_encoding[..., 8] + + R = quat_to_mat(quat) + extrinsics = torch.cat([R, T[..., None]], dim=-1) + + H, W = image_size_hw + fy = (H / 2.0) / torch.clamp(torch.tan(fov_h / 2.0), 1e-6) + fx = (W / 2.0) / torch.clamp(torch.tan(fov_w / 2.0), 1e-6) + intrinsics = torch.zeros(pose_encoding.shape[:2] + (3, 3), device=pose_encoding.device) + intrinsics[..., 0, 0] = fx + intrinsics[..., 1, 1] = fy + intrinsics[..., 0, 2] = W / 2 + intrinsics[..., 1, 2] = H / 2 + intrinsics[..., 2, 2] = 1.0 # Set the homogeneous coordinate to 1 + + return extrinsics, intrinsics + + +def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor: + """ + Quaternion Order: XYZW or say ijkr, scalar-last + + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + i, j, k, r = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + quat_by_rijk = torch.stack( + [ + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape( + batch_dim + (4,) + ) + + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def cam_quat_xyzw_to_world_quat_wxyz(cam_quat_xyzw, c2w): + # cam_quat_xyzw: (b, n, 4) in xyzw + # c2w: (b, n, 4, 4) + b, n = cam_quat_xyzw.shape[:2] + # 1. xyzw -> wxyz + cam_quat_wxyz = torch.cat( + [ + cam_quat_xyzw[..., 3:4], # w + cam_quat_xyzw[..., 0:1], # x + cam_quat_xyzw[..., 1:2], # y + cam_quat_xyzw[..., 2:3], # z + ], + dim=-1, + ) + # 2. Quaternion to matrix + cam_quat_wxyz_flat = cam_quat_wxyz.reshape(-1, 4) + rotmat_cam = quat_to_mat(cam_quat_wxyz_flat).reshape(b, n, 3, 3) + # 3. Transform to world space + rotmat_c2w = c2w[..., :3, :3] + rotmat_world = torch.matmul(rotmat_c2w, rotmat_cam) + # 4. Matrix to quaternion (wxyz) + rotmat_world_flat = rotmat_world.reshape(-1, 3, 3) + world_quat_wxyz_flat = mat_to_quat(rotmat_world_flat) + world_quat_wxyz = world_quat_wxyz_flat.reshape(b, n, 4) + return world_quat_wxyz diff --git a/src/nn/depth_anything_3/registry.py b/src/nn/depth_anything_3/registry.py new file mode 100644 index 0000000..9645071 --- /dev/null +++ b/src/nn/depth_anything_3/registry.py @@ -0,0 +1,50 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import OrderedDict +from pathlib import Path + + +def get_all_models() -> OrderedDict: + """ + Scans all YAML files in the configs directory and returns a sorted dictionary where: + - Keys are model names (YAML filenames without the .yaml extension) + - Values are absolute paths to the corresponding YAML files + """ + # Get path to the configs directory within the da3 package + # Works both in development and after pip installation + # configs_dir = files("depth_anything_3").joinpath("configs") + configs_dir = Path(__file__).resolve().parent / "configs" + + # Ensure path is a Path object for consistent cross-platform handling + configs_dir = Path(configs_dir) + + model_entries = [] + # Iterate through all items in the configs directory + for item in configs_dir.iterdir(): + # Filter for YAML files (excluding directories) + if item.is_file() and item.suffix == ".yaml": + # Extract model name (filename without .yaml extension) + model_name = item.stem + # Get absolute path (resolve() handles symlinks) + file_abs_path = str(item.resolve()) + model_entries.append((model_name, file_abs_path)) + + # Sort entries by model name and convert to OrderedDict + sorted_entries = sorted(model_entries, key=lambda x: x[0]) + return OrderedDict(sorted_entries) + + +# Global registry for external imports +MODEL_REGISTRY = get_all_models() diff --git a/src/nn/depth_anything_3/services/__init__.py b/src/nn/depth_anything_3/services/__init__.py new file mode 100644 index 0000000..07ed8b6 --- /dev/null +++ b/src/nn/depth_anything_3/services/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Services module for Depth Anything 3. +""" + +from depth_anything_3.services.backend import create_app, start_server + +__all__ = [ + start_server, + create_app, +] diff --git a/src/nn/depth_anything_3/services/backend.py b/src/nn/depth_anything_3/services/backend.py new file mode 100644 index 0000000..e279e64 --- /dev/null +++ b/src/nn/depth_anything_3/services/backend.py @@ -0,0 +1,1417 @@ +# flake8: noqa: E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model backend service for Depth Anything 3. +Provides HTTP API for model inference with persistent model loading. +""" + +import os +import posixpath +import time +import uuid + +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional +from urllib.parse import quote +import numpy as np + +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.responses import FileResponse, HTMLResponse +from pydantic import BaseModel + +from ..api import DepthAnything3 +from ..utils.memory import ( + get_gpu_memory_info, + cleanup_cuda_memory, + check_memory_availability, + estimate_memory_requirement, +) + + +class InferenceRequest(BaseModel): + """Request model for inference API.""" + + image_paths: List[str] + export_dir: Optional[str] = None + export_format: str = "mini_npz-glb" + extrinsics: Optional[List[List[List[float]]]] = None + intrinsics: Optional[List[List[List[float]]]] = None + process_res: int = 504 + process_res_method: str = "upper_bound_resize" + export_feat_layers: List[int] = [] + align_to_input_ext_scale: bool = True + # GLB export parameters + conf_thresh_percentile: float = 40.0 + num_max_points: int = 1_000_000 + show_cameras: bool = True + # Feat_vis export parameters + feat_vis_fps: int = 15 + + +class InferenceResponse(BaseModel): + """Response model for inference API.""" + + success: bool + message: str + task_id: Optional[str] = None + export_dir: Optional[str] = None + export_format: str = "mini_npz-glb" + processing_time: Optional[float] = None + + +class TaskStatus(BaseModel): + """Task status model.""" + + task_id: str + status: str # "pending", "running", "completed", "failed" + message: str + progress: Optional[float] = None # 0.0 to 1.0 + created_at: float + started_at: Optional[float] = None + completed_at: Optional[float] = None + export_dir: Optional[str] = None + request: Optional[InferenceRequest] = None # Store the original request + + # Essential task parameters + num_images: Optional[int] = None # Number of input images + export_format: Optional[str] = None # Export format + process_res_method: Optional[str] = None # Processing resolution method + video_path: Optional[str] = None # Source video path + + +class ModelBackend: + """Model backend service with persistent model loading.""" + + def __init__(self, model_dir: str, device: str = "cuda"): + self.model_dir = model_dir + self.device = device + self.model = None + self.model_loaded = False + self.load_time = None + self.load_start_time = None # Time when model loading started + self.load_completed_time = None # Time when model loading completed + self.last_used = None + + def load_model(self): + """Load model if not already loaded.""" + if self.model_loaded and self.model is not None: + self.last_used = time.time() + return self.model + + try: + print(f"Loading model from {self.model_dir}...") + self.load_start_time = time.time() + start_time = time.time() + + self.model = DepthAnything3.from_pretrained(self.model_dir).to(self.device) + self.model.eval() + + self.model_loaded = True + self.load_time = time.time() - start_time + self.load_completed_time = time.time() + self.last_used = time.time() + + print(f"Model loaded successfully in {self.load_time:.2f}s") + return self.model + + except Exception as e: + print(f"Failed to load model: {e}") + raise e + + def get_model(self): + """Get model, loading if necessary.""" + if not self.model_loaded: + return self.load_model() + self.last_used = time.time() + return self.model + + def get_status(self) -> Dict[str, Any]: + """Get backend status information.""" + # Calculate uptime from when model loading completed + uptime = 0 + if self.model_loaded and self.load_completed_time: + uptime = time.time() - self.load_completed_time + + return { + "model_loaded": self.model_loaded, + "model_dir": self.model_dir, + "device": self.device, + "load_time": self.load_time, + "last_used": self.last_used, + "uptime": uptime, + } + + +# Global backend instance +_backend: Optional[ModelBackend] = None +_app: Optional[FastAPI] = None +_tasks: Dict[str, TaskStatus] = {} +_executor = ThreadPoolExecutor(max_workers=1) # Restrict to single-task execution +_running_task_id: Optional[str] = None # Currently running task ID +_task_queue: List[str] = [] # Pending task queue + +# Task cleanup configuration +MAX_TASK_HISTORY = 100 # Maximum number of tasks to keep in memory +CLEANUP_INTERVAL = 300 # Cleanup interval in seconds (5 minutes) + + +def _process_next_task(): + """Process the next task in the queue.""" + global _task_queue, _running_task_id + + if not _task_queue or _running_task_id is not None: + return + + # Get next task from queue + task_id = _task_queue.pop(0) + + # Get task request from tasks dict (we need to store the request) + if task_id not in _tasks: + return + + # Submit task to executor + _executor.submit(_run_inference_task, task_id) + + +# get_gpu_memory_info imported from depth_anything_3.utils.memory + + +# cleanup_cuda_memory imported from depth_anything_3.utils.memory + + +# check_memory_availability imported from depth_anything_3.utils.memory + + +# estimate_memory_requirement imported from depth_anything_3.utils.memory + + +def _run_inference_task(task_id: str): + """Run inference task in background thread with OOM protection.""" + global _tasks, _backend, _running_task_id, _task_queue + + model = None + inference_started = False + start_time = time.time() + + try: + # Get task request + if task_id not in _tasks or _tasks[task_id].request is None: + print(f"[{task_id}] Task not found or request missing") + return + + request = _tasks[task_id].request + num_images = len(request.image_paths) + + # Set current running task + _running_task_id = task_id + + # Update task status to running + _tasks[task_id].status = "running" + _tasks[task_id].started_at = start_time + _tasks[task_id].message = f"[{task_id}] Starting inference on {num_images} frames..." + print(f"[{task_id}] Starting inference on {num_images} frames") + + # Pre-inference cleanup to ensure maximum available memory + print(f"[{task_id}] Pre-inference cleanup...") + cleanup_cuda_memory() + + # Check memory availability + estimated_memory = estimate_memory_requirement(num_images, request.process_res) + mem_available, mem_msg = check_memory_availability(estimated_memory) + print(f"[{task_id}] {mem_msg}") + + if not mem_available: + # Try aggressive cleanup + print(f"[{task_id}] Insufficient memory, attempting aggressive cleanup...") + cleanup_cuda_memory() + time.sleep(0.5) # Give system time to reclaim memory + + # Check again + mem_available, mem_msg = check_memory_availability(estimated_memory) + if not mem_available: + raise RuntimeError( + f"Insufficient GPU memory after cleanup. {mem_msg}\n" + f"Suggestions:\n" + f" 1. Reduce process_res (current: {request.process_res})\n" + f" 2. Process fewer images at once (current: {num_images})\n" + f" 3. Clear other GPU processes" + ) + + # Get model (with error handling) + print(f"[{task_id}] Loading model...") + _tasks[task_id].message = f"[{task_id}] Loading model..." + _tasks[task_id].progress = 0.1 + + try: + model = _backend.get_model() + except RuntimeError as e: + if "out of memory" in str(e).lower(): + cleanup_cuda_memory() + raise RuntimeError( + f"OOM during model loading: {str(e)}\n" + f"Try reducing the batch size or resolution." + ) + raise + + print(f"[{task_id}] Model loaded successfully") + _tasks[task_id].progress = 0.2 + + # Prepare inference parameters + inference_kwargs = { + "image": request.image_paths, + "export_format": request.export_format, + "process_res": request.process_res, + "process_res_method": request.process_res_method, + "export_feat_layers": request.export_feat_layers, + "align_to_input_ext_scale": request.align_to_input_ext_scale, + "conf_thresh_percentile": request.conf_thresh_percentile, + "num_max_points": request.num_max_points, + "show_cameras": request.show_cameras, + "feat_vis_fps": request.feat_vis_fps, + } + + if request.export_dir: + inference_kwargs["export_dir"] = request.export_dir + + if request.extrinsics: + inference_kwargs["extrinsics"] = np.array(request.extrinsics, dtype=np.float32) + + if request.intrinsics: + inference_kwargs["intrinsics"] = np.array(request.intrinsics, dtype=np.float32) + + # Run inference with timing + inference_start_time = time.time() + print(f"[{task_id}] Running model inference...") + _tasks[task_id].message = f"[{task_id}] Running model inference on {num_images} images..." + _tasks[task_id].progress = 0.3 + + inference_started = True + + try: + model.inference(**inference_kwargs) + inference_time = time.time() - inference_start_time + avg_time_per_image = inference_time / num_images if num_images > 0 else 0 + + print( + f"[{task_id}] Inference completed in {inference_time:.2f}s " + f"({avg_time_per_image:.2f}s per image)" + ) + + except RuntimeError as e: + if "out of memory" in str(e).lower(): + cleanup_cuda_memory() + raise RuntimeError( + f"OOM during inference: {str(e)}\n" + f"Settings: {num_images} images, resolution={request.process_res}\n" + f"Suggestions:\n" + f" 1. Reduce process_res to {int(request.process_res * 0.75)}\n" + f" 2. Process images in smaller batches\n" + f" 3. Use process_res_method='resize' instead of 'upper_bound_resize'" + ) + raise + + _tasks[task_id].progress = 0.9 + + # Post-inference cleanup + print(f"[{task_id}] Post-inference cleanup...") + cleanup_cuda_memory() + + # Calculate total processing time + total_time = time.time() - start_time + + # Update task status to completed + _tasks[task_id].status = "completed" + _tasks[task_id].completed_at = time.time() + _tasks[task_id].message = ( + f"[{task_id}] Completed in {total_time:.2f}s " f"({avg_time_per_image:.2f}s per image)" + ) + _tasks[task_id].progress = 1.0 + _tasks[task_id].export_dir = request.export_dir + + # Clear running state + _running_task_id = None + + # Process next task in queue + _process_next_task() + + print(f"[{task_id}] Task completed successfully") + print( + f"[{task_id}] Total time: {total_time:.2f}s, " + f"Inference time: {inference_time:.2f}s, " + f"Avg per image: {avg_time_per_image:.2f}s" + ) + + except Exception as e: + # Update task status to failed + error_msg = str(e) + total_time = time.time() - start_time + + print(f"[{task_id}] Task failed after {total_time:.2f}s: {error_msg}") + + # Always attempt cleanup on failure + cleanup_cuda_memory() + + _tasks[task_id].status = "failed" + _tasks[task_id].completed_at = time.time() + _tasks[task_id].message = f"[{task_id}] Failed after {total_time:.2f}s: {error_msg}" + + # Clear running state + _running_task_id = None + + # Process next task in queue + _process_next_task() + + finally: + # Final cleanup in finally block to ensure it always runs + # This is critical for releasing resources even if unexpected errors occur + try: + if inference_started: + print(f"[{task_id}] Final cleanup in finally block...") + cleanup_cuda_memory() + except Exception as e: + print(f"[{task_id}] Warning: Finally block cleanup failed: {e}") + + # Schedule cleanup after task completion + _schedule_task_cleanup() + + +def _cleanup_old_tasks(): + """Clean up old completed/failed tasks to prevent memory buildup.""" + global _tasks + + current_time = time.time() + tasks_to_remove = [] + + # Find tasks to remove - more aggressive cleanup + for task_id, task in _tasks.items(): + # Remove completed/failed tasks older than 10 minutes (instead of 1 hour) + if ( + task.status in ["completed", "failed"] + and task.completed_at + and current_time - task.completed_at > 600 + ): # 10 minutes + tasks_to_remove.append(task_id) + + # Remove old tasks + for task_id in tasks_to_remove: + del _tasks[task_id] + print(f"[CLEANUP] Removed old task: {task_id}") + + # If still too many tasks, remove oldest completed/failed tasks + if len(_tasks) > MAX_TASK_HISTORY: + completed_tasks = [ + (task_id, task) + for task_id, task in _tasks.items() + if task.status in ["completed", "failed"] + ] + completed_tasks.sort(key=lambda x: x[1].completed_at or 0) + + excess_count = len(_tasks) - MAX_TASK_HISTORY + for i in range(min(excess_count, len(completed_tasks))): + task_id = completed_tasks[i][0] + del _tasks[task_id] + print(f"[CLEANUP] Removed excess task: {task_id}") + + # Count active tasks (only pending and running) + active_count = sum(1 for task in _tasks.values() if task.status in ["pending", "running"]) + print( + "[CLEANUP] Task cleanup completed. " + f"Total tasks: {len(_tasks)}, Active tasks: {active_count}" + ) + + +def _schedule_task_cleanup(): + """Schedule task cleanup in background.""" + + def cleanup_worker(): + try: + time.sleep(2) # Small delay to ensure task status is updated + _cleanup_old_tasks() + except Exception as e: + print(f"[CLEANUP] Cleanup worker failed: {e}") + + # Run cleanup in background thread + _executor.submit(cleanup_worker) + + +# ============================================================================ +# Gallery utilities (extracted from gallery.py) +# ============================================================================ + +GALLERY_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + + +def _load_gallery_html() -> str: + """ + Load and modify gallery HTML to work under /gallery/ subdirectory. + Replaces API paths from root to /gallery/ prefix. + """ + from ..services.gallery import HTML_PAGE + + # Replace API paths to be under /gallery/ subdirectory + html = ( + HTML_PAGE.replace("fetch('/manifest.json'", "fetch('/gallery/manifest.json'") + .replace("fetch('/manifest/'+", "fetch('/gallery/manifest/'+") + .replace( + "if(location.pathname!=\"/\")history.replaceState(null,'','/'+location.search)", + "if(!location.pathname.startsWith(\"/gallery\"))history.replaceState(null,'','/gallery/'+location.search)", + ) + ) + + return html + + +def _gallery_url_join(*parts: str) -> str: + """Join URL parts safely.""" + norm = posixpath.join(*[p.replace("\\", "/") for p in parts]) + segs = [s for s in norm.split("/") if s not in ("", ".")] + return "/".join(quote(s) for s in segs) + + +def _is_plain_name(name: str) -> bool: + """Check if name is safe for use in paths.""" + return all(c not in name for c in ("/", "\\")) and name not in (".", "..") + + +def build_group_list(root_dir: str) -> dict: + """Build list of groups from gallery directory.""" + groups = [] + try: + for gname in sorted(os.listdir(root_dir)): + gpath = os.path.join(root_dir, gname) + if not os.path.isdir(gpath): + continue + has_scene = False + try: + for sname in os.listdir(gpath): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + if os.path.exists(os.path.join(spath, "scene.glb")) and os.path.exists( + os.path.join(spath, "scene.jpg") + ): + has_scene = True + break + except Exception: + pass + if has_scene: + groups.append({"id": gname, "title": gname}) + except Exception as e: + print(f"[warn] build_group_list failed: {e}") + return {"groups": groups} + + +def build_group_manifest(root_dir: str, group: str) -> dict: + """Build manifest for a specific group.""" + items = [] + gpath = os.path.join(root_dir, group) + try: + if not os.path.isdir(gpath): + return {"group": group, "items": []} + for sname in sorted(os.listdir(gpath)): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + glb_fs = os.path.join(spath, "scene.glb") + jpg_fs = os.path.join(spath, "scene.jpg") + if not (os.path.exists(glb_fs) and os.path.exists(jpg_fs)): + continue + depth_images = [] + dpath = os.path.join(spath, "depth_vis") + if os.path.isdir(dpath): + files = [ + f + for f in os.listdir(dpath) + if os.path.splitext(f)[1].lower() in GALLERY_IMAGE_EXTS + ] + for fn in sorted(files): + depth_images.append( + "/gallery/" + _gallery_url_join(group, sname, "depth_vis", fn) + ) + items.append( + { + "id": sname, + "title": sname, + "model": "/gallery/" + _gallery_url_join(group, sname, "scene.glb"), + "thumbnail": "/gallery/" + _gallery_url_join(group, sname, "scene.jpg"), + "depth_images": depth_images, + } + ) + except Exception as e: + print(f"[warn] build_group_manifest failed for {group}: {e}") + return {"group": group, "items": items} + + +def create_app(model_dir: str, device: str = "cuda", gallery_dir: Optional[str] = None) -> FastAPI: + """Create FastAPI application with model backend.""" + global _backend, _app + + _backend = ModelBackend(model_dir, device) + _app = FastAPI( + title="Depth Anything 3 Backend", + description="Model inference service for Depth Anything 3", + version="1.0.0", + ) + + # Store gallery directory globally for use in routes + _gallery_dir = gallery_dir + + @_app.get("/", response_class=HTMLResponse) + async def root(): + """Home page with navigation to dashboard and gallery.""" + html_content = ( + """ + + + + + + Depth Anything 3 Backend + + + +
+

Depth Anything 3

+

Model Backend Service

+
+ +

📊 Dashboard

+

Monitor backend status, model information, and inference tasks in real-time.

+ + """ + + ( + '' + "

🎨 Gallery

" + "

Browse 3D reconstructions and depth visualizations from processed scenes.

" + "" + if _gallery_dir and os.path.exists(_gallery_dir) + else "" + ) + + """ +
+
+

Depth Anything 3 Backend API

+
+
+ + + """ + ) + return HTMLResponse(html_content) + + @_app.get("/dashboard", response_class=HTMLResponse) + async def dashboard(): + """HTML dashboard for monitoring backend status and tasks.""" + if _backend is None: + return HTMLResponse("

Backend not initialized

", status_code=500) + + # Get backend status + status = _backend.get_status() + + # Safely format status values + if status["load_time"] is not None: + load_time_str = f"{status['load_time']:.2f}s" + else: + load_time_str = "Not loaded" + + if status["uptime"] is not None: + uptime_str = f"{status['uptime']:.2f}s" + else: + uptime_str = "Not running" + + # Get tasks information + active_tasks = [task for task in _tasks.values() if task.status in ["pending", "running"]] + completed_tasks = [ + task for task in _tasks.values() if task.status in ["completed", "failed"] + ] + + # Generate task HTML + active_tasks_html = "" + if active_tasks: + for task in active_tasks: + task_details = f""" +
+
+ {task.task_id} + {task.status} +
+
{task.message}
+
+ + Images: {task.num_images or 'N/A'} | + Format: {task.export_format or 'N/A'} | + Method: {task.process_res_method or 'N/A'} | + Export Dir: {task.export_dir or 'N/A'} + + {f'
Video: {task.video_path}' if task.video_path else ''} +
+
+ """ + active_tasks_html += task_details + else: + active_tasks_html = "

No active tasks

" + + completed_tasks_html = "" + if completed_tasks: + for task in completed_tasks[-10:]: + task_details = f""" +
+
+ {task.task_id} + {task.status} +
+
{task.message}
+
+ + Images: {task.num_images or 'N/A'} | + Format: {task.export_format or 'N/A'} | + Method: {task.process_res_method or 'N/A'} | + Export Dir: {task.export_dir or 'N/A'} + + {f'
Video: {task.video_path}' if task.video_path else ''} +
+
+ """ + completed_tasks_html += task_details + else: + completed_tasks_html = "

No completed tasks

" + + # Generate HTML + html_content = f""" + + + + + + Depth Anything 3 Backend Dashboard + + + +
+
+

Depth Anything 3 Backend Dashboard

+

Real-time monitoring of model status and inference tasks

+
+ +
+
+

Model Status

+
+ Status: + + {'Online' if status['model_loaded'] else 'Offline'} + +
+
+ Model Directory: + {status['model_dir']} +
+
+ Device: + {status['device']} +
+
+ Load Time: + {load_time_str} +
+
+ Uptime: + {uptime_str} +
+
+ +
+

Task Summary

+
+ Active Tasks: + {len(active_tasks)} +
+
+ Completed Tasks: + {len(completed_tasks)} +
+
+ Total Tasks: + {len(_tasks)} +
+
+
+ +
+

Active Tasks

+ + +
Last updated: {time.strftime('%Y-%m-%d %H:%M:%S')}
+ + {active_tasks_html} +
+ +
+

Recent Completed Tasks

+ {completed_tasks_html} +
+
+ + + + + """ + + return HTMLResponse(html_content) + + @_app.get("/status") + async def get_status(): + """Get backend status with GPU memory information.""" + if _backend is None: + raise HTTPException(status_code=500, detail="Backend not initialized") + + status = _backend.get_status() + + # Add GPU memory information + gpu_memory = get_gpu_memory_info() + if gpu_memory: + status["gpu_memory"] = { + "total_gb": round(gpu_memory["total_gb"], 2), + "allocated_gb": round(gpu_memory["allocated_gb"], 2), + "reserved_gb": round(gpu_memory["reserved_gb"], 2), + "free_gb": round(gpu_memory["free_gb"], 2), + "utilization_percent": round(gpu_memory["utilization"], 1), + } + else: + status["gpu_memory"] = None + + return status + + @_app.post("/inference", response_model=InferenceResponse) + async def run_inference(request: InferenceRequest): + """Submit inference task and return task ID.""" + global _running_task_id + + if _backend is None: + raise HTTPException(status_code=500, detail="Backend not initialized") + + # Generate unique task ID + task_id = str(uuid.uuid4()) + + # Create task status + if _running_task_id is not None: + status_msg = f"[{task_id}] Task queued (waiting for {_running_task_id} to complete)" + else: + status_msg = f"[{task_id}] Task submitted" + + _tasks[task_id] = TaskStatus( + task_id=task_id, + status="pending", + message=status_msg, + created_at=time.time(), + export_dir=request.export_dir, + request=request, + # Record essential parameters + num_images=len(request.image_paths), + export_format=request.export_format, + process_res_method=request.process_res_method, + video_path=( + request.image_paths[0] if request.image_paths else None + ), # Use first image path as video reference + ) + + # Add task to queue + _task_queue.append(task_id) + + # If no task is running, start processing the queue + if _running_task_id is None: + _process_next_task() + + return InferenceResponse( + success=True, + message="Task submitted successfully", + task_id=task_id, + export_dir=request.export_dir, + export_format=request.export_format, + ) + + @_app.get("/task/{task_id}", response_model=TaskStatus) + async def get_task_status(task_id: str): + """Get task status by task ID.""" + if task_id not in _tasks: + raise HTTPException(status_code=404, detail="Task not found") + + return _tasks[task_id] + + @_app.get("/gpu-memory") + async def get_gpu_memory(): + """Get detailed GPU memory information.""" + gpu_memory = get_gpu_memory_info() + if gpu_memory is None: + return { + "available": False, + "message": "CUDA not available or memory info cannot be retrieved", + } + + return { + "available": True, + "total_gb": round(gpu_memory["total_gb"], 2), + "allocated_gb": round(gpu_memory["allocated_gb"], 2), + "reserved_gb": round(gpu_memory["reserved_gb"], 2), + "free_gb": round(gpu_memory["free_gb"], 2), + "utilization_percent": round(gpu_memory["utilization"], 1), + "status": ( + "healthy" + if gpu_memory["utilization"] < 80 + else "warning" if gpu_memory["utilization"] < 95 else "critical" + ), + } + + @_app.get("/tasks") + async def list_tasks(): + """List all tasks.""" + # Separate active and completed tasks + active_tasks = [task for task in _tasks.values() if task.status in ["pending", "running"]] + completed_tasks = [ + task for task in _tasks.values() if task.status in ["completed", "failed"] + ] + + return { + "tasks": list(_tasks.values()), + "active_tasks": active_tasks, + "completed_tasks": completed_tasks, + "active_count": len(active_tasks), + "total_count": len(_tasks), + } + + @_app.post("/cleanup") + async def manual_cleanup(): + """Manually trigger task cleanup.""" + try: + _cleanup_old_tasks() + return {"message": "Cleanup completed", "active_tasks": len(_tasks)} + except Exception as e: + raise HTTPException(status_code=500, detail=f"Cleanup failed: {str(e)}") + + @_app.delete("/task/{task_id}") + async def delete_task(task_id: str): + """Delete a specific task.""" + if task_id not in _tasks: + raise HTTPException(status_code=404, detail="Task not found") + + # Only allow deletion of completed/failed tasks + if _tasks[task_id].status not in ["completed", "failed"]: + raise HTTPException(status_code=400, detail="Cannot delete running or pending tasks") + + del _tasks[task_id] + return {"message": f"Task {task_id} deleted successfully"} + + @_app.post("/reload") + async def reload_model(): + """Reload the model.""" + if _backend is None: + raise HTTPException(status_code=500, detail="Backend not initialized") + + try: + _backend.model = None + _backend.model_loaded = False + _backend.load_model() + return {"message": "Model reloaded successfully"} + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to reload model: {str(e)}") + + # ============================================================================ + # Gallery routes + # ============================================================================ + + if _gallery_dir and os.path.exists(_gallery_dir): + # Load gallery HTML page (with modified paths for /gallery/ subdirectory) + _gallery_html = _load_gallery_html() + + @_app.get("/gallery/", response_class=HTMLResponse) + @_app.get("/gallery", response_class=HTMLResponse) + async def gallery_home(): + """Gallery home page.""" + return HTMLResponse(_gallery_html) + + @_app.get("/gallery/manifest.json") + async def gallery_manifest(): + """Get gallery group list.""" + try: + return build_group_list(_gallery_dir) + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to build group list: {str(e)}" + ) + + @_app.get("/gallery/manifest/{group}.json") + async def gallery_group_manifest(group: str): + """Get manifest for a specific group.""" + if not _is_plain_name(group): + raise HTTPException(status_code=400, detail="Invalid group name") + try: + return build_group_manifest(_gallery_dir, group) + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to build group manifest: {str(e)}" + ) + + @_app.get("/gallery/{path:path}") + async def gallery_files(path: str): + """Serve gallery static files (GLB, JPG, etc.).""" + # Security check: prevent directory traversal + path_parts = path.split("/") + if any(not _is_plain_name(part) for part in path_parts if part): + raise HTTPException(status_code=400, detail="Invalid path") + + file_path = os.path.join(_gallery_dir, *path_parts) + + # Ensure the file is within gallery directory + real_file_path = os.path.realpath(file_path) + real_gallery_dir = os.path.realpath(_gallery_dir) + if not real_file_path.startswith(real_gallery_dir): + raise HTTPException(status_code=403, detail="Access denied") + + if not os.path.exists(file_path) or not os.path.isfile(file_path): + raise HTTPException(status_code=404, detail="File not found") + + return FileResponse(file_path) + + return _app + + +def start_server( + model_dir: str, + device: str = "cuda", + host: str = "127.0.0.1", + port: int = 8000, + gallery_dir: Optional[str] = None, +): + """Start the backend server.""" + app = create_app(model_dir, device, gallery_dir) + + print("Starting Depth Anything 3 Backend...") + print(f"Model directory: {model_dir}") + print(f"Device: {device}") + print(f"Server: http://{host}:{port}") + print(f"Dashboard: http://{host}:{port}/dashboard") + print(f"API Status: http://{host}:{port}/status") + + if gallery_dir and os.path.exists(gallery_dir): + print(f"Gallery: http://{host}:{port}/gallery/") + + print("=" * 60) + print("Backend is running! You can now:") + print(f" • Open home page: http://{host}:{port}") + print(f" • Open dashboard: http://{host}:{port}/dashboard") + print(f" • Check API status: http://{host}:{port}/status") + + if gallery_dir and os.path.exists(gallery_dir): + print(f" • Browse gallery: http://{host}:{port}/gallery/") + + print(" • Submit inference tasks via API") + print("=" * 60) + + uvicorn.run(app, host=host, port=port, log_level="info") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Depth Anything 3 Backend Server") + parser.add_argument("--model-dir", required=True, help="Model directory path") + parser.add_argument("--device", default="cuda", help="Device to use") + parser.add_argument("--host", default="127.0.0.1", help="Host to bind to") + parser.add_argument("--port", type=int, default=8000, help="Port to bind to") + parser.add_argument("--gallery-dir", help="Gallery directory path (optional)") + + args = parser.parse_args() + start_server(args.model_dir, args.device, args.host, args.port, args.gallery_dir) diff --git a/src/nn/depth_anything_3/services/gallery.py b/src/nn/depth_anything_3/services/gallery.py new file mode 100644 index 0000000..f72bb5e --- /dev/null +++ b/src/nn/depth_anything_3/services/gallery.py @@ -0,0 +1,806 @@ +#!/usr/bin/env python3 +# flake8: noqa: E501 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Depth Anything 3 Gallery Server (two-level, single-file) +Now supports paginated depth preview (4 per page). +""" + +import argparse +import json +import mimetypes +import os +import posixpath +import sys +from functools import partial +from http import HTTPStatus +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import quote, unquote + +# ------------------------------ Embedded HTML ------------------------------ # + +HTML_PAGE = r""" + + + + Depth Anything 3 Gallery + + + + + + + +
+
+ +

Depth Anything 3 Gallery

+ + +
+
Level 1 shows groups only; click a group to browse scenes and previews.
+
+ +
+ +
+

+ 🎯 Depth Anything 3 Gallery +

+

+ Explore 3D reconstructions and depth visualizations from Depth Anything 3. + Browse through groups of scenes, preview 3D models, and examine depth maps interactively. +

+
+ +
+
    + +
    + + +
    + + + +
    Depth Anything 3 Gallery. Copyright 2025 Depth Anything 3 authors.
    + + + + +""" + +# ------------------------------ Utilities ------------------------------ # + +IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + + +def _url_join(*parts: str) -> str: + norm = posixpath.join(*[p.replace("\\", "/") for p in parts]) + segs = [s for s in norm.split("/") if s not in ("", ".")] + return "/".join(quote(s) for s in segs) + + +def _is_plain_name(name: str) -> bool: + return all(c not in name for c in ("/", "\\")) and name not in (".", "..") + + +def build_group_list(root_dir: str) -> dict: + groups = [] + try: + for gname in sorted(os.listdir(root_dir)): + gpath = os.path.join(root_dir, gname) + if not os.path.isdir(gpath): + continue + has_scene = False + try: + for sname in os.listdir(gpath): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + if os.path.exists(os.path.join(spath, "scene.glb")) and os.path.exists( + os.path.join(spath, "scene.jpg") + ): + has_scene = True + break + except Exception: + pass + if has_scene: + groups.append({"id": gname, "title": gname}) + except Exception as e: + print(f"[warn] build_group_list failed: {e}", file=sys.stderr) + return {"groups": groups} + + +def build_group_manifest(root_dir: str, group: str) -> dict: + items = [] + gpath = os.path.join(root_dir, group) + try: + if not os.path.isdir(gpath): + return {"group": group, "items": []} + for sname in sorted(os.listdir(gpath)): + spath = os.path.join(gpath, sname) + if not os.path.isdir(spath): + continue + glb_fs = os.path.join(spath, "scene.glb") + jpg_fs = os.path.join(spath, "scene.jpg") + if not (os.path.exists(glb_fs) and os.path.exists(jpg_fs)): + continue + depth_images = [] + dpath = os.path.join(spath, "depth_vis") + if os.path.isdir(dpath): + files = [ + f for f in os.listdir(dpath) if os.path.splitext(f)[1].lower() in IMAGE_EXTS + ] + for fn in sorted(files): + depth_images.append("/" + _url_join(group, sname, "depth_vis", fn)) + items.append( + { + "id": sname, + "title": sname, + "model": "/" + _url_join(group, sname, "scene.glb"), + "thumbnail": "/" + _url_join(group, sname, "scene.jpg"), + "depth_images": depth_images, + } + ) + except Exception as e: + print(f"[warn] build_group_manifest failed for {group}: {e}", file=sys.stderr) + return {"group": group, "items": items} + + +class GalleryHandler(SimpleHTTPRequestHandler): + def __init__(self, *args, directory=None, **kwargs): + super().__init__(*args, directory=directory, **kwargs) + + def do_GET(self): + if self.path in ("/", "/index.html") or self.path.startswith("/?"): + content = HTML_PAGE.encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(content))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(content) + return + if self.path == "/manifest.json": + data = json.dumps( + build_group_list(self.directory), ensure_ascii=False, indent=2 + ).encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(data))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(data) + return + if self.path.startswith("/manifest/") and self.path.endswith(".json"): + group_enc = self.path[len("/manifest/") : -len(".json")] + try: + group = unquote(group_enc) + except Exception: + group = group_enc + if not _is_plain_name(group): + self.send_error(HTTPStatus.BAD_REQUEST, "Invalid group name") + return + data = json.dumps( + build_group_manifest(self.directory, group), ensure_ascii=False, indent=2 + ).encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(data))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(data) + return + if self.path == "/favicon.ico": + self.send_response(HTTPStatus.NO_CONTENT) + self.end_headers() + return + return super().do_GET() + + def list_directory(self, path): + self.send_error(HTTPStatus.NOT_FOUND, "Directory listing disabled") + return None + + +def gallery(): + parser = argparse.ArgumentParser( + description="Depth Anything 3 Gallery Server (two-level, with pagination)" + ) + parser.add_argument( + "-d", "--dir", required=True, help="Gallery root directory (two-level: group/scene)" + ) + parser.add_argument("-p", "--port", type=int, default=8000, help="Port (default 8000)") + parser.add_argument("--host", default="127.0.0.1", help="Host address (default 127.0.0.1)") + parser.add_argument("--open", action="store_true", help="Open browser after launch") + args = parser.parse_args() + + root_dir = os.path.abspath(args.dir) + if not os.path.isdir(root_dir): + print(f"[error] Directory not found: {root_dir}", file=sys.stderr) + sys.exit(1) + + Handler = partial(GalleryHandler, directory=root_dir) + server = ThreadingHTTPServer((args.host, args.port), Handler) + + addr = f"http://{args.host}:{args.port}/" + print(f"[info] Serving gallery from: {root_dir}") + print(f"[info] Open: {addr}") + + if args.open: + try: + import webbrowser + + webbrowser.open(addr) + except Exception as e: + print(f"[warn] Failed to open browser: {e}", file=sys.stderr) + + try: + server.serve_forever() + except KeyboardInterrupt: + print("\n[info] Shutting down...") + finally: + server.server_close() + + +def main(): + """Main entry point for gallery server.""" + mimetypes.add_type("model/gltf-binary", ".glb") + gallery() + + +if __name__ == "__main__": + main() diff --git a/src/nn/depth_anything_3/services/inference_service.py b/src/nn/depth_anything_3/services/inference_service.py new file mode 100644 index 0000000..2ec8dbc --- /dev/null +++ b/src/nn/depth_anything_3/services/inference_service.py @@ -0,0 +1,239 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unified Inference Service +Provides unified interface for local and remote inference +""" + +from typing import Any, Dict, List, Optional, Union +import numpy as np +import requests +import typer + +from ..api import DepthAnything3 + + +class InferenceService: + """Unified inference service class""" + + def __init__(self, model_dir: str, device: str = "cuda"): + self.model_dir = model_dir + self.device = device + self.model = None + + def load_model(self): + """Load model""" + if self.model is None: + typer.echo(f"Loading model from {self.model_dir}...") + self.model = DepthAnything3.from_pretrained(self.model_dir).to(self.device) + return self.model + + def run_local_inference( + self, + image_paths: List[str], + export_dir: str, + export_format: str = "mini_npz-glb", + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_feat_layers: List[int] = None, + extrinsics: Optional[np.ndarray] = None, + intrinsics: Optional[np.ndarray] = None, + align_to_input_ext_scale: bool = True, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + feat_vis_fps: int = 15, + ) -> Any: + """Run local inference""" + if export_feat_layers is None: + export_feat_layers = [] + + model = self.load_model() + + # Prepare inference parameters + inference_kwargs = { + "image": image_paths, + "export_dir": export_dir, + "export_format": export_format, + "process_res": process_res, + "process_res_method": process_res_method, + "export_feat_layers": export_feat_layers, + "align_to_input_ext_scale": align_to_input_ext_scale, + "use_ray_pose": use_ray_pose, + "ref_view_strategy": ref_view_strategy, + "conf_thresh_percentile": conf_thresh_percentile, + "num_max_points": num_max_points, + "show_cameras": show_cameras, + "feat_vis_fps": feat_vis_fps, + } + + # Add pose data (if exists) + if extrinsics is not None: + inference_kwargs["extrinsics"] = extrinsics + if intrinsics is not None: + inference_kwargs["intrinsics"] = intrinsics + + # Run inference + typer.echo(f"Running inference on {len(image_paths)} images...") + prediction = model.inference(**inference_kwargs) + + typer.echo(f"Results saved to {export_dir}") + typer.echo(f"Export format: {export_format}") + + return prediction + + def run_backend_inference( + self, + image_paths: List[str], + export_dir: str, + backend_url: str, + export_format: str = "mini_npz-glb", + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_feat_layers: List[int] = None, + extrinsics: Optional[np.ndarray] = None, + intrinsics: Optional[np.ndarray] = None, + align_to_input_ext_scale: bool = True, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + feat_vis_fps: int = 15, + ) -> Dict[str, Any]: + """Run backend inference""" + if export_feat_layers is None: + export_feat_layers = [] + + # Check backend status + if not self._check_backend_status(backend_url): + raise typer.BadParameter(f"Backend service is not running at {backend_url}") + + # Prepare payload + payload = { + "image_paths": image_paths, + "export_dir": export_dir, + "export_format": export_format, + "process_res": process_res, + "process_res_method": process_res_method, + "export_feat_layers": export_feat_layers, + "align_to_input_ext_scale": align_to_input_ext_scale, + "use_ray_pose": use_ray_pose, + "ref_view_strategy": ref_view_strategy, + "conf_thresh_percentile": conf_thresh_percentile, + "num_max_points": num_max_points, + "show_cameras": show_cameras, + "feat_vis_fps": feat_vis_fps, + } + + # Add pose data (if exists) + if extrinsics is not None: + payload["extrinsics"] = [ext.astype(np.float64).tolist() for ext in extrinsics] + if intrinsics is not None: + payload["intrinsics"] = [intr.astype(np.float64).tolist() for intr in intrinsics] + + # Submit task + typer.echo("Submitting inference task to backend...") + try: + response = requests.post(f"{backend_url}/inference", json=payload, timeout=30) + response.raise_for_status() + result = response.json() + + if result["success"]: + task_id = result["task_id"] + typer.echo("Task submitted successfully!") + typer.echo(f"Task ID: {task_id}") + typer.echo(f"Results will be saved to: {export_dir}") + typer.echo(f"Check backend logs for progress updates with task ID: {task_id}") + return result + else: + raise typer.BadParameter( + f"Backend inference submission failed: {result['message']}" + ) + except requests.exceptions.RequestException as e: + raise typer.BadParameter(f"Backend inference submission failed: {e}") + + def _check_backend_status(self, backend_url: str) -> bool: + """Check backend status""" + try: + response = requests.get(f"{backend_url}/status", timeout=5) + return response.status_code == 200 + except Exception: + return False + + +def run_inference( + image_paths: List[str], + export_dir: str, + model_dir: str, + device: str = "cuda", + backend_url: Optional[str] = None, + export_format: str = "mini_npz-glb", + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + export_feat_layers: List[int] = None, + extrinsics: Optional[np.ndarray] = None, + intrinsics: Optional[np.ndarray] = None, + align_to_input_ext_scale: bool = True, + use_ray_pose: bool = False, + ref_view_strategy: str = "saddle_balanced", + conf_thresh_percentile: float = 40.0, + num_max_points: int = 1_000_000, + show_cameras: bool = True, + feat_vis_fps: int = 15, +) -> Union[Any, Dict[str, Any]]: + """Unified inference interface""" + + service = InferenceService(model_dir, device) + + if backend_url: + return service.run_backend_inference( + image_paths=image_paths, + export_dir=export_dir, + backend_url=backend_url, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) + else: + return service.run_local_inference( + image_paths=image_paths, + export_dir=export_dir, + export_format=export_format, + process_res=process_res, + process_res_method=process_res_method, + export_feat_layers=export_feat_layers, + extrinsics=extrinsics, + intrinsics=intrinsics, + align_to_input_ext_scale=align_to_input_ext_scale, + use_ray_pose=use_ray_pose, + ref_view_strategy=ref_view_strategy, + conf_thresh_percentile=conf_thresh_percentile, + num_max_points=num_max_points, + show_cameras=show_cameras, + feat_vis_fps=feat_vis_fps, + ) diff --git a/src/nn/depth_anything_3/services/input_handlers.py b/src/nn/depth_anything_3/services/input_handlers.py new file mode 100644 index 0000000..dc0536b --- /dev/null +++ b/src/nn/depth_anything_3/services/input_handlers.py @@ -0,0 +1,266 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Input Processing Service +Handles different types of inputs (image, images, colmap, video) +""" + +import glob +import os +from typing import List, Tuple +import cv2 +import numpy as np +import typer + +from ..utils.read_write_model import read_model + + +class InputHandler: + """Base input handler class""" + + @staticmethod + def validate_path(path: str, path_type: str = "file") -> str: + """Validate path""" + if not os.path.exists(path): + raise typer.BadParameter(f"{path_type} not found: {path}") + return path + + @staticmethod + def handle_export_dir(export_dir: str, auto_cleanup: bool = False) -> str: + """Handle export directory""" + if os.path.exists(export_dir): + if auto_cleanup: + typer.echo(f"Auto-cleaning existing export directory: {export_dir}") + import shutil + + shutil.rmtree(export_dir) + os.makedirs(export_dir, exist_ok=True) + else: + typer.echo(f"Export directory '{export_dir}' already exists.") + if typer.confirm("Do you want to clean it and continue?"): + import shutil + + shutil.rmtree(export_dir) + os.makedirs(export_dir, exist_ok=True) + typer.echo(f"Cleaned export directory: {export_dir}") + else: + typer.echo("Operation cancelled.") + raise typer.Exit(0) + else: + os.makedirs(export_dir, exist_ok=True) + return export_dir + + +class ImageHandler(InputHandler): + """Single image handler""" + + @staticmethod + def process(image_path: str) -> List[str]: + """Process single image""" + InputHandler.validate_path(image_path, "Image file") + return [image_path] + + +class ImagesHandler(InputHandler): + """Image directory handler""" + + @staticmethod + def process(images_dir: str, image_extensions: str = "png,jpg,jpeg") -> List[str]: + """Process image directory""" + InputHandler.validate_path(images_dir, "Images directory") + + # Parse extensions + extensions = [ext.strip().lower() for ext in image_extensions.split(",")] + extensions = [ext if ext.startswith(".") else f".{ext}" for ext in extensions] + + # Find image files + image_files = [] + for ext in extensions: + pattern = f"*{ext}" + image_files.extend(glob.glob(os.path.join(images_dir, pattern))) + image_files.extend(glob.glob(os.path.join(images_dir, pattern.upper()))) + + image_files = sorted(list(set(image_files))) # Remove duplicates and sort + + if not image_files: + raise typer.BadParameter( + f"No image files found in {images_dir} with extensions: {extensions}" + ) + + typer.echo(f"Found {len(image_files)} images to process") + return image_files + + +class ColmapHandler(InputHandler): + """COLMAP data handler""" + + @staticmethod + def process( + colmap_dir: str, sparse_subdir: str = "" + ) -> Tuple[List[str], np.ndarray, np.ndarray]: + """Process COLMAP data""" + InputHandler.validate_path(colmap_dir, "COLMAP directory") + + # Build paths + images_dir = os.path.join(colmap_dir, "images") + if sparse_subdir: + sparse_dir = os.path.join(colmap_dir, "sparse", sparse_subdir) + else: + sparse_dir = os.path.join(colmap_dir, "sparse") + + InputHandler.validate_path(images_dir, "Images directory") + InputHandler.validate_path(sparse_dir, "Sparse reconstruction directory") + + # Load COLMAP data + typer.echo("Loading COLMAP reconstruction data...") + try: + cameras, images, points3D = read_model(sparse_dir) + + typer.echo( + f"Loaded COLMAP data: {len(cameras)} cameras, {len(images)} images, " + f"{len(points3D)} 3D points." + ) + + # Get image files and pose data + image_files = [] + extrinsics = [] + intrinsics = [] + + for image_id, image_data in images.items(): + image_name = image_data.name + image_path = os.path.join(images_dir, image_name) + + if os.path.exists(image_path): + image_files.append(image_path) + + # Get camera parameters + camera = cameras[image_data.camera_id] + + # Convert quaternion to rotation matrix + R = image_data.qvec2rotmat() + t = image_data.tvec + + # Create extrinsic matrix (world to camera) + extrinsic = np.eye(4) + extrinsic[:3, :3] = R + extrinsic[:3, 3] = t + extrinsics.append(extrinsic) + + # Create intrinsic matrix + if camera.model == "PINHOLE": + fx, fy, cx, cy = camera.params + elif camera.model == "SIMPLE_PINHOLE": + f, cx, cy = camera.params + fx = fy = f + else: + # For other models, use basic pinhole approximation + fx = fy = camera.params[0] if len(camera.params) > 0 else 1000 + cx = camera.width / 2 + cy = camera.height / 2 + + intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) + intrinsics.append(intrinsic) + + if not image_files: + raise typer.BadParameter("No valid images found in COLMAP data") + + typer.echo(f"Found {len(image_files)} valid images with pose data") + + return image_files, np.array(extrinsics), np.array(intrinsics) + + except Exception as e: + raise typer.BadParameter(f"Failed to load COLMAP data: {e}") + + +class VideoHandler(InputHandler): + """Video handler""" + + @staticmethod + def process(video_path: str, output_dir: str, fps: float = 1.0) -> List[str]: + """Process video, extract frames""" + InputHandler.validate_path(video_path, "Video file") + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise typer.BadParameter(f"Cannot open video: {video_path}") + + # Get video properties + video_fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + duration = total_frames / video_fps + + # Calculate frame interval (ensure at least 1) + frame_interval = max(1, int(video_fps / fps)) + actual_fps = video_fps / frame_interval + + typer.echo(f"Video FPS: {video_fps:.2f}, Duration: {duration:.2f}s") + + # Warn if requested FPS is higher than video FPS + if fps > video_fps: + typer.echo( + f"⚠️ Warning: Requested sampling FPS ({fps:.2f}) exceeds video FPS ({video_fps:.2f})", # noqa: E501 + err=True, + ) + typer.echo( + f"⚠️ Using maximum available FPS: {actual_fps:.2f} (extracting every frame)", + err=True, + ) + + typer.echo(f"Extracting frames at {actual_fps:.2f} FPS (every {frame_interval} frame(s))") + + # Create output directory + frames_dir = os.path.join(output_dir, "input_images") + os.makedirs(frames_dir, exist_ok=True) + + frame_count = 0 + saved_count = 0 + + while True: + ret, frame = cap.read() + if not ret: + break + + if frame_count % frame_interval == 0: + frame_path = os.path.join(frames_dir, f"{saved_count:06d}.png") + cv2.imwrite(frame_path, frame) + saved_count += 1 + + frame_count += 1 + + cap.release() + typer.echo(f"Extracted {saved_count} frames to {frames_dir}") + + # Get frame file list + frame_files = sorted( + [f for f in os.listdir(frames_dir) if f.endswith((".png", ".jpg", ".jpeg"))] + ) + if not frame_files: + raise typer.BadParameter("No frames extracted from video") + + return [os.path.join(frames_dir, f) for f in frame_files] + + +def parse_export_feat(export_feat_str: str) -> List[int]: + """Parse export_feat parameter""" + if not export_feat_str: + return [] + + try: + return [int(x.strip()) for x in export_feat_str.split(",") if x.strip()] + except ValueError: + raise typer.BadParameter( + f"Invalid export_feat format: {export_feat_str}. " + "Use comma-separated integers like '0,1,2'" + ) diff --git a/src/nn/depth_anything_3/specs.py b/src/nn/depth_anything_3/specs.py new file mode 100644 index 0000000..fe5b302 --- /dev/null +++ b/src/nn/depth_anything_3/specs.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional +import numpy as np +import torch + + +@dataclass +class Gaussians: + """3DGS parameters, all in world space""" + + means: torch.Tensor # world points, "batch gaussian dim" + scales: torch.Tensor # scales_std, "batch gaussian 3" + rotations: torch.Tensor # world_quat_wxyz, "batch gaussian 4" + harmonics: torch.Tensor # world SH, "batch gaussian 3 d_sh" + opacities: torch.Tensor # opacity | opacity SH, "batch gaussian" | "batch gaussian 1 d_sh" + + +@dataclass +class Prediction: + depth: np.ndarray # N, H, W + is_metric: int + sky: np.ndarray | None = None # N, H, W + conf: np.ndarray | None = None # N, H, W + extrinsics: np.ndarray | None = None # N, 4, 4 + intrinsics: np.ndarray | None = None # N, 3, 3 + processed_images: np.ndarray | None = None # N, H, W, 3 - processed images for visualization + gaussians: Gaussians | None = None # 3D gaussians + aux: dict[str, Any] = None # + scale_factor: Optional[float] = None # metric scale diff --git a/src/nn/depth_anything_3/utils/alignment.py b/src/nn/depth_anything_3/utils/alignment.py new file mode 100644 index 0000000..ceb8983 --- /dev/null +++ b/src/nn/depth_anything_3/utils/alignment.py @@ -0,0 +1,163 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Alignment utilities for depth estimation and metric scaling. +""" + +from typing import Tuple +import torch + + +def least_squares_scale_scalar( + a: torch.Tensor, b: torch.Tensor, eps: float = 1e-12 +) -> torch.Tensor: + """ + Compute least squares scale factor s such that a ≈ s * b. + + Args: + a: First tensor + b: Second tensor + eps: Small epsilon for numerical stability + + Returns: + Scalar tensor containing the scale factor + + Raises: + ValueError: If tensors have mismatched shapes or devices + TypeError: If tensors are not floating point + """ + if a.shape != b.shape: + raise ValueError(f"Shape mismatch: {a.shape} vs {b.shape}") + if a.device != b.device: + raise ValueError(f"Device mismatch: {a.device} vs {b.device}") + if not a.is_floating_point() or not b.is_floating_point(): + raise TypeError("Tensors must be floating point type") + + # Compute dot products for least squares solution + num = torch.dot(a.reshape(-1), b.reshape(-1)) + den = torch.dot(b.reshape(-1), b.reshape(-1)).clamp_min(eps) + return num / den + + +def compute_sky_mask(sky_prediction: torch.Tensor, threshold: float = 0.3) -> torch.Tensor: + """ + Compute non-sky mask from sky prediction. + + Args: + sky_prediction: Sky prediction tensor + threshold: Threshold for sky classification + + Returns: + Boolean mask where True indicates non-sky regions + """ + return sky_prediction < threshold + + +def compute_alignment_mask( + depth_conf: torch.Tensor, + non_sky_mask: torch.Tensor, + depth: torch.Tensor, + metric_depth: torch.Tensor, + median_conf: torch.Tensor, + min_depth_threshold: float = 1e-3, + min_metric_depth_threshold: float = 1e-2, +) -> torch.Tensor: + """ + Compute mask for depth alignment based on confidence and depth thresholds. + + Args: + depth_conf: Depth confidence tensor + non_sky_mask: Non-sky region mask + depth: Predicted depth tensor + metric_depth: Metric depth tensor + median_conf: Median confidence threshold + min_depth_threshold: Minimum depth threshold + min_metric_depth_threshold: Minimum metric depth threshold + + Returns: + Boolean mask for valid alignment regions + """ + return ( + (depth_conf >= median_conf) + & non_sky_mask + & (metric_depth > min_metric_depth_threshold) + & (depth > min_depth_threshold) + ) + + +def sample_tensor_for_quantile(tensor: torch.Tensor, max_samples: int = 100000) -> torch.Tensor: + """ + Sample tensor elements for quantile computation to reduce memory usage. + + Args: + tensor: Input tensor to sample + max_samples: Maximum number of samples to take + + Returns: + Sampled tensor + """ + if tensor.numel() <= max_samples: + return tensor + + idx = torch.randperm(tensor.numel(), device=tensor.device)[:max_samples] + return tensor.flatten()[idx] + + +def apply_metric_scaling( + depth: torch.Tensor, intrinsics: torch.Tensor, scale_factor: float = 300.0 +) -> torch.Tensor: + """ + Apply metric scaling to depth based on camera intrinsics. + + Args: + depth: Input depth tensor + intrinsics: Camera intrinsics tensor + scale_factor: Scaling factor for metric conversion + + Returns: + Scaled depth tensor + """ + focal_length = (intrinsics[:, :, 0, 0] + intrinsics[:, :, 1, 1]) / 2 + return depth * (focal_length[:, :, None, None] / scale_factor) + + +def set_sky_regions_to_max_depth( + depth: torch.Tensor, + depth_conf: torch.Tensor, + non_sky_mask: torch.Tensor, + max_depth: float = 200.0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Set sky regions to maximum depth and high confidence. + + Args: + depth: Depth tensor + depth_conf: Depth confidence tensor + non_sky_mask: Non-sky region mask + max_depth: Maximum depth value for sky regions + + Returns: + Tuple of (updated_depth, updated_depth_conf) + """ + depth = depth.clone() + + # Set sky regions to max depth and high confidence + depth[~non_sky_mask] = max_depth + if depth_conf is not None: + depth_conf = depth_conf.clone() + depth_conf[~non_sky_mask] = 1.0 + return depth, depth_conf + else: + return depth, None diff --git a/src/nn/depth_anything_3/utils/api_helpers.py b/src/nn/depth_anything_3/utils/api_helpers.py new file mode 100644 index 0000000..b327331 --- /dev/null +++ b/src/nn/depth_anything_3/utils/api_helpers.py @@ -0,0 +1,58 @@ +import argparse + + +def parse_scalar(s): + if not isinstance(s, str): + return s + t = s.strip() + l = t.lower() + if l == "true": + return True + if l == "false": + return False + if l in ("none", "null"): + return None + try: + return int(t, 10) + except Exception: + pass + try: + return float(t) + except Exception: + return s + + +def fn_kv_csv(s: str) -> dict[str, dict[str, object]]: + """ + Parse a string of comma-separated triplets: fn:key:value + + Returns: + dict[fn_name] -> dict[key] = parsed_value + + Example: + "fn1:width:1920,fn1:height:1080,fn2:quality:0.8" + -> {"fn1": {"width": 1920, "height": 1080}, "fn2": {"quality": 0.8}} + """ + result: dict[str, dict[str, object]] = {} + if not s: + return result + + for item in s.split(","): + if not item: + continue + parts = item.split(":", 2) # allow value to contain ":" beyond first two separators + if len(parts) < 3: + raise argparse.ArgumentTypeError(f"Bad item '{item}', expected FN:KEY:VALUE") + fn, key, raw_val = parts[0], parts[1], parts[2] + # If you need to allow colons in values, join leftover parts: + # fn, key, raw_val = parts[0], parts[1], ":".join(parts[2:]) + + if not fn: + raise argparse.ArgumentTypeError(f"Bad item '{item}': empty function name") + if not key: + raise argparse.ArgumentTypeError(f"Bad item '{item}': empty key") + + val = parse_scalar(raw_val) + bucket = result.setdefault(fn, {}) + bucket[key] = val + return result diff --git a/src/nn/depth_anything_3/utils/camera_trj_helpers.py b/src/nn/depth_anything_3/utils/camera_trj_helpers.py new file mode 100644 index 0000000..83624f3 --- /dev/null +++ b/src/nn/depth_anything_3/utils/camera_trj_helpers.py @@ -0,0 +1,479 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from einops import einsum, rearrange, reduce + +try: + from scipy.spatial.transform import Rotation as R +except ImportError: + from depth_anything_3.utils.logger import logger + + logger.warn("Dependency 'scipy' not found. Required for interpolating camera trajectory.") + +from depth_anything_3.utils.geometry import as_homogeneous + + +@torch.no_grad() +def render_stabilization_path(poses, k_size=45): + """Rendering stabilized camera path. + poses: [batch, 4, 4] or [batch, 3, 4], + return: + smooth path: [batch 4 4]""" + num_frames = poses.shape[0] + device = poses.device + dtype = poses.dtype + + # Early exit for trivial cases + if num_frames <= 1: + return as_homogeneous(poses) + + # Make k_size safe: positive odd and not larger than num_frames + # 1) Ensure odd + if k_size < 1: + k_size = 1 + if k_size % 2 == 0: + k_size += 1 + # 2) Cap to num_frames (keep odd) + max_odd = num_frames if (num_frames % 2 == 1) else (num_frames - 1) + if max_odd < 1: + max_odd = 1 # covers num_frames == 0 theoretically + k_size = min(k_size, max_odd) + # 3) enforce a minimum of 3 when possible (for better smoothing) + if num_frames >= 3 and k_size < 3: + k_size = 3 + + input_poses = [] + for i in range(num_frames): + input_poses.append( + torch.cat([poses[i, :3, 0:1], poses[i, :3, 1:2], poses[i, :3, 3:4]], dim=-1) + ) + input_poses = torch.stack(input_poses) # (num_frames, 3, 3) + + # Prepare Gaussian kernel + gaussian_kernel = cv2.getGaussianKernel(ksize=k_size, sigma=-1).astype(np.float32).squeeze() + gaussian_kernel = torch.tensor(gaussian_kernel, dtype=dtype, device=device).view(1, 1, -1) + pad = k_size // 2 + + output_vectors = [] + for idx in range(3): # For r1, r2, t + vec = ( + input_poses[:, :, idx].T.unsqueeze(0).unsqueeze(0) + ) # (1, 1, 3, num_frames) -> (1, 1, 3, num_frames) + # But actually, we want (batch=3, channel=1, width=num_frames) + # So: + vec = input_poses[:, :, idx].T.unsqueeze(1) # (3, 1, num_frames) + vec_padded = F.pad(vec, (pad, pad), mode="reflect") + filtered = F.conv1d(vec_padded, gaussian_kernel) + output_vectors.append(filtered.squeeze(1).T) # (num_frames, 3) + + output_r1, output_r2, output_t = output_vectors # Each is (num_frames, 3) + + # Normalize r1 and r2 + output_r1 = output_r1 / output_r1.norm(dim=-1, keepdim=True) + output_r2 = output_r2 / output_r2.norm(dim=-1, keepdim=True) + + output_poses = [] + for i in range(num_frames): + output_r3 = torch.linalg.cross(output_r1[i], output_r2[i]) + render_pose = torch.cat( + [ + output_r1[i].unsqueeze(-1), + output_r2[i].unsqueeze(-1), + output_r3.unsqueeze(-1), + output_t[i].unsqueeze(-1), + ], + dim=-1, + ) + output_poses.append(render_pose[:3, :]) + output_poses = as_homogeneous(torch.stack(output_poses, dim=0)) + + return output_poses + + +@torch.no_grad() +def render_wander_path( + cam2world: torch.Tensor, + intrinsic: torch.Tensor, + h: int, + w: int, + num_frames: int = 120, + max_disp: float = 48.0, +): + device, dtype = cam2world.device, cam2world.dtype + fx = intrinsic[0, 0] * w + r = max_disp / fx + th = torch.linspace(0, 2.0 * torch.pi, steps=num_frames, device=device, dtype=dtype) + x = r * torch.sin(th) + yz = r * torch.cos(th) / 3.0 + T = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(num_frames, 1, 1) + T[:, :3, 3] = torch.stack([x, yz, yz], dim=-1) * -1.0 + c2ws = cam2world.unsqueeze(0) @ T + # Start at reference pose and end back at reference pose + c2ws = torch.cat([cam2world.unsqueeze(0), c2ws, cam2world.unsqueeze(0)], dim=0) + Ks = intrinsic.unsqueeze(0).repeat(c2ws.shape[0], 1, 1) + return c2ws, Ks + + +@torch.no_grad() +def render_dolly_zoom_path( + cam2world: torch.Tensor, + intrinsic: torch.Tensor, + h: int, + w: int, + num_frames: int = 120, + max_disp: float = 0.1, + D_focus: float = 10.0, +): + device, dtype = cam2world.device, cam2world.dtype + fx0, fy0 = intrinsic[0, 0] * w, intrinsic[1, 1] * h + t = torch.linspace(0.0, 2.0, steps=num_frames, device=device, dtype=dtype) + z = 0.5 * (1.0 - torch.cos(torch.pi * t)) * max_disp + T = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(num_frames, 1, 1) + T[:, 2, 3] = -z + c2ws = cam2world.unsqueeze(0) @ T + Df = torch.as_tensor(D_focus, device=device, dtype=dtype) + scale = (Df / (Df + z)).clamp(min=1e-6) + Ks = intrinsic.unsqueeze(0).repeat(num_frames, 1, 1) + Ks[:, 0, 0] = (fx0 * scale) / w + Ks[:, 1, 1] = (fy0 * scale) / h + return c2ws, Ks + + +@torch.no_grad() +def interpolate_intrinsics( + initial: torch.Tensor, # "*#batch 3 3" + final: torch.Tensor, # "*#batch 3 3" + t: torch.Tensor, # " time_step" +) -> torch.Tensor: # "*batch time_step 3 3" + initial = rearrange(initial, "... i j -> ... () i j") + final = rearrange(final, "... i j -> ... () i j") + t = rearrange(t, "t -> t () ()") + return initial + (final - initial) * t + + +def intersect_rays( + a_origins: torch.Tensor, # "*#batch dim" + a_directions: torch.Tensor, # "*#batch dim" + b_origins: torch.Tensor, # "*#batch dim" + b_directions: torch.Tensor, # "*#batch dim" +) -> torch.Tensor: # "*batch dim" + """Compute the least-squares intersection of rays. Uses the math from here: + https://math.stackexchange.com/a/1762491/286022 + """ + + # Broadcast and stack the tensors. + a_origins, a_directions, b_origins, b_directions = torch.broadcast_tensors( + a_origins, a_directions, b_origins, b_directions + ) + origins = torch.stack((a_origins, b_origins), dim=-2) + directions = torch.stack((a_directions, b_directions), dim=-2) + + # Compute n_i * n_i^T - eye(3) from the equation. + n = einsum(directions, directions, "... n i, ... n j -> ... n i j") + n = n - torch.eye(3, dtype=origins.dtype, device=origins.device) + + # Compute the left-hand side of the equation. + lhs = reduce(n, "... n i j -> ... i j", "sum") + + # Compute the right-hand side of the equation. + rhs = einsum(n, origins, "... n i j, ... n j -> ... n i") + rhs = reduce(rhs, "... n i -> ... i", "sum") + + # Left-matrix-multiply both sides by the inverse of lhs to find p. + return torch.linalg.lstsq(lhs, rhs).solution + + +def normalize(a: torch.Tensor) -> torch.Tensor: # "*#batch dim" -> "*#batch dim" + return a / a.norm(dim=-1, keepdim=True) + + +def generate_coordinate_frame( + y: torch.Tensor, # "*#batch 3" + z: torch.Tensor, # "*#batch 3" +) -> torch.Tensor: # "*batch 3 3" + """Generate a coordinate frame given perpendicular, unit-length Y and Z vectors.""" + y, z = torch.broadcast_tensors(y, z) + return torch.stack([y.cross(z, dim=-1), y, z], dim=-1) + + +def generate_rotation_coordinate_frame( + a: torch.Tensor, # "*#batch 3" + b: torch.Tensor, # "*#batch 3" + eps: float = 1e-4, +) -> torch.Tensor: # "*batch 3 3" + """Generate a coordinate frame where the Y direction is normal to the plane defined + by unit vectors a and b. The other axes are arbitrary.""" + device = a.device + + # Replace every entry in b that's parallel to the corresponding entry in a with an + # arbitrary vector. + b = b.detach().clone() + parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps + b[parallel] = torch.tensor([0, 0, 1], dtype=b.dtype, device=device) + parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps + b[parallel] = torch.tensor([0, 1, 0], dtype=b.dtype, device=device) + + # Generate the coordinate frame. The initial cross product defines the plane. + return generate_coordinate_frame(normalize(torch.linalg.cross(a, b)), a) + + +def matrix_to_euler( + rotations: torch.Tensor, # "*batch 3 3" + pattern: str, +) -> torch.Tensor: # "*batch 3" + *batch, _, _ = rotations.shape + rotations = rotations.reshape(-1, 3, 3) + angles_np = R.from_matrix(rotations.detach().cpu().numpy()).as_euler(pattern) + rotations = torch.tensor(angles_np, dtype=rotations.dtype, device=rotations.device) + return rotations.reshape(*batch, 3) + + +def euler_to_matrix( + rotations: torch.Tensor, # "*batch 3" + pattern: str, +) -> torch.Tensor: # "*batch 3 3" + *batch, _ = rotations.shape + rotations = rotations.reshape(-1, 3) + matrix_np = R.from_euler(pattern, rotations.detach().cpu().numpy()).as_matrix() + rotations = torch.tensor(matrix_np, dtype=rotations.dtype, device=rotations.device) + return rotations.reshape(*batch, 3, 3) + + +def extrinsics_to_pivot_parameters( + extrinsics: torch.Tensor, # "*#batch 4 4" + pivot_coordinate_frame: torch.Tensor, # "*#batch 3 3" + pivot_point: torch.Tensor, # "*#batch 3" +) -> torch.Tensor: # "*batch 5" + """Convert the extrinsics to a representation with 5 degrees of freedom: + 1. Distance from pivot point in the "X" (look cross pivot axis) direction. + 2. Distance from pivot point in the "Y" (pivot axis) direction. + 3. Distance from pivot point in the Z (look) direction + 4. Angle in plane + 5. Twist (rotation not in plane) + """ + + # The pivot coordinate frame's Z axis is normal to the plane. + pivot_axis = pivot_coordinate_frame[..., :, 1] + + # Compute the translation elements of the pivot parametrization. + translation_frame = generate_coordinate_frame(pivot_axis, extrinsics[..., :3, 2]) + origin = extrinsics[..., :3, 3] + delta = pivot_point - origin + translation = einsum(translation_frame, delta, "... i j, ... i -> ... j") + + # Add the rotation elements of the pivot parametrization. + inverted = pivot_coordinate_frame.inverse() @ extrinsics[..., :3, :3] + y, _, z = matrix_to_euler(inverted, "YXZ").unbind(dim=-1) + + return torch.cat([translation, y[..., None], z[..., None]], dim=-1) + + +def pivot_parameters_to_extrinsics( + parameters: torch.Tensor, # "*#batch 5" + pivot_coordinate_frame: torch.Tensor, # "*#batch 3 3" + pivot_point: torch.Tensor, # "*#batch 3" +) -> torch.Tensor: # "*batch 4 4" + translation, y, z = parameters.split((3, 1, 1), dim=-1) + + euler = torch.cat((y, torch.zeros_like(y), z), dim=-1) + rotation = pivot_coordinate_frame @ euler_to_matrix(euler, "YXZ") + + # The pivot coordinate frame's Z axis is normal to the plane. + pivot_axis = pivot_coordinate_frame[..., :, 1] + + translation_frame = generate_coordinate_frame(pivot_axis, rotation[..., :3, 2]) + delta = einsum(translation_frame, translation, "... i j, ... j -> ... i") + origin = pivot_point - delta + + *batch, _ = origin.shape + extrinsics = torch.eye(4, dtype=parameters.dtype, device=parameters.device) + extrinsics = extrinsics.broadcast_to((*batch, 4, 4)).clone() + extrinsics[..., 3, 3] = 1 + extrinsics[..., :3, :3] = rotation + extrinsics[..., :3, 3] = origin + return extrinsics + + +def interpolate_circular( + a: torch.Tensor, # "*#batch" + b: torch.Tensor, # "*#batch" + t: torch.Tensor, # "*#batch" +) -> torch.Tensor: # " *batch" + a, b, t = torch.broadcast_tensors(a, b, t) + + tau = 2 * torch.pi + a = a % tau + b = b % tau + + # Consider piecewise edge cases. + d = (b - a).abs() + a_left = a - tau + d_left = (b - a_left).abs() + a_right = a + tau + d_right = (b - a_right).abs() + use_d = (d < d_left) & (d < d_right) + use_d_left = (d_left < d_right) & (~use_d) + use_d_right = (~use_d) & (~use_d_left) + + result = a + (b - a) * t + result[use_d_left] = (a_left + (b - a_left) * t)[use_d_left] + result[use_d_right] = (a_right + (b - a_right) * t)[use_d_right] + + return result + + +def interpolate_pivot_parameters( + initial: torch.Tensor, # "*#batch 5" + final: torch.Tensor, # "*#batch 5" + t: torch.Tensor, # " time_step" +) -> torch.Tensor: # "*batch time_step 5" + initial = rearrange(initial, "... d -> ... () d") + final = rearrange(final, "... d -> ... () d") + t = rearrange(t, "t -> t ()") + ti, ri = initial.split((3, 2), dim=-1) + tf, rf = final.split((3, 2), dim=-1) + + t_lerp = ti + (tf - ti) * t + r_lerp = interpolate_circular(ri, rf, t) + + return torch.cat((t_lerp, r_lerp), dim=-1) + + +@torch.no_grad() +def interpolate_extrinsics( + initial: torch.Tensor, # "*#batch 4 4" + final: torch.Tensor, # "*#batch 4 4" + t: torch.Tensor, # " time_step" + eps: float = 1e-4, +) -> torch.Tensor: # "*batch time_step 4 4" + """Interpolate extrinsics by rotating around their "focus point," which is the + least-squares intersection between the look vectors of the initial and final + extrinsics. + """ + + initial = initial.type(torch.float64) + final = final.type(torch.float64) + t = t.type(torch.float64) + + # Based on the dot product between the look vectors, pick from one of two cases: + # 1. Look vectors are parallel: interpolate about their origins' midpoint. + # 3. Look vectors aren't parallel: interpolate about their focus point. + initial_look = initial[..., :3, 2] + final_look = final[..., :3, 2] + dot_products = einsum(initial_look, final_look, "... i, ... i -> ...") + parallel_mask = (dot_products.abs() - 1).abs() < eps + + # Pick focus points. + initial_origin = initial[..., :3, 3] + final_origin = final[..., :3, 3] + pivot_point = 0.5 * (initial_origin + final_origin) + pivot_point[~parallel_mask] = intersect_rays( + initial_origin[~parallel_mask], + initial_look[~parallel_mask], + final_origin[~parallel_mask], + final_look[~parallel_mask], + ) + + # Convert to pivot parameters. + pivot_frame = generate_rotation_coordinate_frame(initial_look, final_look, eps=eps) + initial_params = extrinsics_to_pivot_parameters(initial, pivot_frame, pivot_point) + final_params = extrinsics_to_pivot_parameters(final, pivot_frame, pivot_point) + + # Interpolate the pivot parameters. + interpolated_params = interpolate_pivot_parameters(initial_params, final_params, t) + + # Convert back. + return pivot_parameters_to_extrinsics( + interpolated_params.type(torch.float32), + rearrange(pivot_frame, "... i j -> ... () i j").type(torch.float32), + rearrange(pivot_point, "... xyz -> ... () xyz").type(torch.float32), + ) + + +@torch.no_grad() +def generate_wobble_transformation( + radius: torch.Tensor, # "*#batch" + t: torch.Tensor, # " time_step" + num_rotations: int = 1, + scale_radius_with_t: bool = True, +) -> torch.Tensor: # "*batch time_step 4 4"]: + # Generate a translation in the image plane. + tf = torch.eye(4, dtype=torch.float32, device=t.device) + tf = tf.broadcast_to((*radius.shape, t.shape[0], 4, 4)).clone() + radius = radius[..., None] + if scale_radius_with_t: + radius = radius * t + tf[..., 0, 3] = torch.sin(2 * torch.pi * num_rotations * t) * radius + tf[..., 1, 3] = -torch.cos(2 * torch.pi * num_rotations * t) * radius + return tf + + +@torch.no_grad() +def render_wobble_inter_path( + cam2world: torch.Tensor, intr_normed: torch.Tensor, inter_len: int, n_skip: int = 3 +): + """ + cam2world: [batch, 4, 4], + intr_normed: [batch, 3, 3] + """ + frame_per_round = n_skip * inter_len + num_rotations = 1 + + t = torch.linspace(0, 1, frame_per_round, dtype=torch.float32, device=cam2world.device) + # t = (torch.cos(torch.pi * (t + 1)) + 1) / 2 + tgt_c2w_b = [] + tgt_intr_b = [] + for b_idx in range(cam2world.shape[0]): + tgt_c2w = [] + tgt_intr = [] + for cur_idx in range(0, cam2world.shape[1] - n_skip, n_skip): + origin_a = cam2world[b_idx, cur_idx, :3, 3] + origin_b = cam2world[b_idx, cur_idx + n_skip, :3, 3] + delta = (origin_a - origin_b).norm(dim=-1) + if cur_idx == 0: + delta_prev = delta + else: + delta = (delta_prev + delta) / 2 + delta_prev = delta + tf = generate_wobble_transformation( + radius=delta * 0.5, + t=t, + num_rotations=num_rotations, + scale_radius_with_t=False, + ) + cur_extrs = ( + interpolate_extrinsics( + cam2world[b_idx, cur_idx], + cam2world[b_idx, cur_idx + n_skip], + t, + ) + @ tf + ) + tgt_c2w.append(cur_extrs[(0 if cur_idx == 0 else 1) :]) + tgt_intr.append( + interpolate_intrinsics( + intr_normed[b_idx, cur_idx], + intr_normed[b_idx, cur_idx + n_skip], + t, + )[(0 if cur_idx == 0 else 1) :] + ) + tgt_c2w_b.append(torch.cat(tgt_c2w)) + tgt_intr_b.append(torch.cat(tgt_intr)) + tgt_c2w = torch.stack(tgt_c2w_b) # b v 4 4 + tgt_intr = torch.stack(tgt_intr_b) # b v 3 3 + return tgt_c2w, tgt_intr diff --git a/src/nn/depth_anything_3/utils/constants.py b/src/nn/depth_anything_3/utils/constants.py new file mode 100644 index 0000000..25c330e --- /dev/null +++ b/src/nn/depth_anything_3/utils/constants.py @@ -0,0 +1,270 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +DEFAULT_MODEL = "depth-anything/DA3NESTED-GIANT-LARGE-1.1" +DEFAULT_EXPORT_DIR = "workspace/gallery/scene" +DEFAULT_GALLERY_DIR = "workspace/gallery" +DEFAULT_GRADIO_DIR = "workspace/gradio" +THRESH_FOR_REF_SELECTION = 3 + +# ============================================================================= +# Benchmark Evaluation Constants +# ============================================================================= + +# Default evaluation workspace directory +DEFAULT_EVAL_WORKSPACE = "workspace/evaluation" + +# Default reference view selection strategy for evaluation +# Use "first" for consistent and reproducible evaluation results +# Other options: "saddle_balanced", "auto", "mid" +EVAL_REF_VIEW_STRATEGY = "first" + +# ----------------------------------------------------------------------------- +# DTU Dataset Configuration +# Reference: https://roboimagedata.compute.dtu.dk/ +# Note: DepthAnything3 was never trained on any images from DTU. +# ----------------------------------------------------------------------------- + +# Root directory for DTU evaluation data (MVSNet format) +# Download from: https://drive.google.com/file/d/1rX0EXlUL4prRxrRu2DgLJv2j7-tpUD4D/view +DTU_EVAL_DATA_ROOT = "workspace/benchmark_dataset/dtu" + +# List of DTU evaluation scenes +DTU_SCENES = [ + "scan1", + "scan4", + "scan9", + "scan10", + "scan11", + "scan12", + "scan13", + "scan15", + "scan23", + "scan24", + "scan29", + "scan32", + "scan33", + "scan34", + "scan48", + "scan49", + "scan62", + "scan75", + "scan77", + "scan110", + "scan114", + "scan118", +] + +# Point cloud fusion hyperparameters +DTU_DIST_THRESH = 0.2 # Distance threshold for geometric consistency (mm) +DTU_NUM_CONSIST = 4 # Minimum number of consistent views for a point +DTU_MAX_POINTS = 4_000_000 # Maximum points in fused point cloud + +# 3D reconstruction evaluation hyperparameters +DTU_DOWN_DENSE = 0.2 # Downsample density for evaluation (mm) +DTU_PATCH_SIZE = 60 # Patch size for boundary handling +DTU_MAX_DIST = 20 # Outlier threshold for accuracy/completeness (mm) + +# ----------------------------------------------------------------------------- +# DTU-64 Dataset Configuration (Pose Evaluation Only) +# This is a subset of DTU with 64 images per scene for pose evaluation. +# Note: This dataset is ONLY for pose evaluation, not 3D reconstruction. +# ----------------------------------------------------------------------------- + +# Root directory for DTU-64 evaluation data +DTU64_EVAL_DATA_ROOT = "workspace/benchmark_dataset/dtu64" +DTU64_CAMERA_ROOT = "workspace/benchmark_dataset/dtu64/Cameras" + +# List of DTU-64 evaluation scenes (13 scenes) +DTU64_SCENES = [ + "scan105", + "scan114", + "scan118", + "scan122", + "scan24", + "scan37", + "scan40", + "scan55", + "scan63", + "scan65", + "scan69", + "scan83", + "scan97", +] + +# ----------------------------------------------------------------------------- +# ETH3D Dataset Configuration +# Reference: https://www.eth3d.net/ +# High-resolution multi-view stereo benchmark with laser-scanned ground truth. +# Note: DepthAnything3 was never trained on any images from ETH3D. +# ----------------------------------------------------------------------------- + +# Root directory for ETH3D evaluation data +ETH3D_EVAL_DATA_ROOT = "workspace/benchmark_dataset/eth3d" + +# List of ETH3D evaluation scenes (indoor and outdoor) +ETH3D_SCENES = [ + "courtyard", + "electro", + "kicker", + "pipes", + "relief", + # "terrace", # Excluded: known issues + "delivery_area", + "facade", + # "meadow", # Excluded: known issues + "office", + "playground", + "relief_2", + "terrains", +] + +# Images to filter out (known problematic views per scene) +ETH3D_FILTER_KEYS = { + "delivery_area": ["711.JPG", "712.JPG", "713.JPG", "714.JPG"], + "electro": ["9289.JPG", "9290.JPG", "9291.JPG", "9292.JPG", "9293.JPG", "9298.JPG"], + "playground": ["587.JPG", "588.JPG", "589.JPG", "590.JPG", "591.JPG", "592.JPG"], + "relief": [ + "427.JPG", "428.JPG", "429.JPG", "430.JPG", "431.JPG", "432.JPG", + "433.JPG", "434.JPG", "435.JPG", "436.JPG", "437.JPG", "438.JPG", + ], + "relief_2": [ + "458.JPG", "459.JPG", "460.JPG", "461.JPG", "462.JPG", "463.JPG", + "464.JPG", "465.JPG", "466.JPG", "467.JPG", "468.JPG", + ], +} + +# TSDF fusion hyperparameters (scaled for outdoor scenes) +ETH3D_VOXEL_LENGTH = 4.0 / 512.0 * 5 # Voxel size for TSDF (meters) +ETH3D_SDF_TRUNC = 0.04 * 5 # SDF truncation distance (meters) +ETH3D_MAX_DEPTH = 100000.0 # Maximum depth for integration (effectively no truncation) + +# Point cloud sampling +ETH3D_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +ETH3D_EVAL_THRESHOLD = 0.05 * 5 # Distance threshold for precision/recall (meters) +ETH3D_DOWN_SAMPLE = 4.0 / 512.0 * 5 # Voxel size for evaluation downsampling (meters) + + +# ============================================================================== +# 7Scenes Dataset Configuration +# ============================================================================== +# Reference: https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/ +# Note: Indoor RGB-D dataset with ground truth poses and meshes. + +# Root directory for 7Scenes evaluation data +SEVENSCENES_EVAL_DATA_ROOT = "workspace/benchmark_dataset/7scenes" + +# List of 7Scenes evaluation scenes +SEVENSCENES_SCENES = [ + "chess", + "fire", + "heads", + "office", + "pumpkin", + "redkitchen", + "stairs", +] + +# Fixed camera intrinsics for 7Scenes (all images share same intrinsics) +SEVENSCENES_FX = 585.0 +SEVENSCENES_FY = 585.0 +SEVENSCENES_CX = 320.0 +SEVENSCENES_CY = 240.0 + +# TSDF fusion hyperparameters (indoor scenes, smaller voxels) +SEVENSCENES_VOXEL_LENGTH = 4.0 / 512.0 # Voxel size for TSDF (meters) +SEVENSCENES_SDF_TRUNC = 0.04 # SDF truncation distance (meters) +SEVENSCENES_MAX_DEPTH = 1000000.0 # Maximum depth for integration (no truncation) + +# Point cloud sampling +SEVENSCENES_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +SEVENSCENES_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) +SEVENSCENES_DOWN_SAMPLE = 4.0 / 512.0 # Voxel size for evaluation downsampling (meters) + + +# ============================================================================== +# ScanNet++ Dataset Configuration +# ============================================================================== +# Reference: https://kaldir.vc.in.tum.de/scannetpp/ +# Note: High-quality indoor RGB-D dataset with iPhone and DSLR images. + +# Root directory for ScanNet++ evaluation data +SCANNETPP_EVAL_DATA_ROOT = "workspace/benchmark_dataset/scannetpp" + +# List of ScanNet++ evaluation scenes +SCANNETPP_SCENES = [ + "09c1414f1b", + "1ada7a0617", + "40aec5fffa", + "3e8bba0176", + "acd95847c5", + "578511c8a9", + "5f99900f09", + "c4c04e6d6c", + "f3d64c30f8", + "7bc286c1b6", + "c5439f4607", + "286b55a2bf", + "fb5a96b1a2", + "7831862f02", + "38d58a7a31", + "bde1e479ad", + "9071e139d9", + "21d970d8de", + "bcd2436daf", + "cc5237fd77", +] + +# Input resolution for ScanNet++ (after undistortion and resize) +SCANNETPP_INPUT_H = 768 +SCANNETPP_INPUT_W = 1024 + +# TSDF fusion hyperparameters (indoor scenes) +SCANNETPP_VOXEL_LENGTH = 0.02 # Voxel size for TSDF (meters) +SCANNETPP_SDF_TRUNC = 0.15 # SDF truncation distance (meters) +SCANNETPP_MAX_DEPTH = 5.0 # Maximum depth for integration (meters) + +# Point cloud sampling +SCANNETPP_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +SCANNETPP_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) +SCANNETPP_DOWN_SAMPLE = 0.02 # Voxel size for evaluation downsampling (meters) + + +# ============================================================================== +# HiRoom Dataset Configuration +# ============================================================================== +# Note: Indoor RGB-D dataset. + +# Root directory for HiRoom evaluation data +HIROOM_EVAL_DATA_ROOT = "workspace/benchmark_dataset/hiroom/data" +HIROOM_GT_ROOT_PATH = "workspace/benchmark_dataset/hiroom/fused_pcd" +HIROOM_SCENE_LIST_PATH = "workspace/benchmark_dataset/hiroom/selected_scene_list_val.txt" + +# TSDF fusion hyperparameters (indoor scenes) +HIROOM_VOXEL_LENGTH = 4.0 / 512.0 # Voxel size for TSDF (meters) +HIROOM_SDF_TRUNC = 0.04 # SDF truncation distance (meters) +HIROOM_MAX_DEPTH = 10000.0 # Maximum depth for integration (no truncation) + +# Point cloud sampling +HIROOM_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh + +# 3D reconstruction evaluation hyperparameters +HIROOM_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) +HIROOM_DOWN_SAMPLE = 4.0 / 512.0 # Voxel size for evaluation downsampling (meters) diff --git a/src/nn/depth_anything_3/utils/export/__init__.py b/src/nn/depth_anything_3/utils/export/__init__.py new file mode 100644 index 0000000..e3e4c65 --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/__init__.py @@ -0,0 +1,59 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.export.gs import export_to_gs_ply, export_to_gs_video + +from .colmap import export_to_colmap +from .depth_vis import export_to_depth_vis +from .feat_vis import export_to_feat_vis +from .glb import export_to_glb +from .npz import export_to_mini_npz, export_to_npz + + +def export( + prediction: Prediction, + export_format: str, + export_dir: str, + **kwargs, +): + if "-" in export_format: + export_formats = export_format.split("-") + for export_format in export_formats: + export(prediction, export_format, export_dir, **kwargs) + return # Prevent falling through to single-format handling + + if export_format == "glb": + export_to_glb(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "mini_npz": + export_to_mini_npz(prediction, export_dir) + elif export_format == "npz": + export_to_npz(prediction, export_dir) + elif export_format == "feat_vis": + export_to_feat_vis(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "depth_vis": + export_to_depth_vis(prediction, export_dir) + elif export_format == "gs_ply": + export_to_gs_ply(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "gs_video": + export_to_gs_video(prediction, export_dir, **kwargs.get(export_format, {})) + elif export_format == "colmap": + export_to_colmap(prediction, export_dir, **kwargs.get(export_format, {})) + else: + raise ValueError(f"Unsupported export format: {export_format}") + + +__all__ = [ + export, +] diff --git a/src/nn/depth_anything_3/utils/export/colmap.py b/src/nn/depth_anything_3/utils/export/colmap.py new file mode 100644 index 0000000..e57964c --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/colmap.py @@ -0,0 +1,150 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import pycolmap +import cv2 as cv +import numpy as np + +from PIL import Image + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.logger import logger + +from .glb import _depths_to_world_points_with_colors + + +def export_to_colmap( + prediction: Prediction, + export_dir: str, + image_paths: list[str], + conf_thresh_percentile: float = 40.0, + process_res_method: str = "upper_bound_resize", +) -> None: + # 1. Data preparation + conf_thresh = np.percentile(prediction.conf, conf_thresh_percentile) + points, colors = _depths_to_world_points_with_colors( + prediction.depth, + prediction.intrinsics, + prediction.extrinsics, # w2c + prediction.processed_images, + prediction.conf, + conf_thresh, + ) + num_points = len(points) + logger.info(f"Exporting to COLMAP with {num_points} points") + num_frames = len(prediction.processed_images) + h, w = prediction.processed_images.shape[1:3] + points_xyf = _create_xyf(num_frames, h, w) + points_xyf = points_xyf[prediction.conf >= conf_thresh] + + # 2. Set Reconstruction + reconstruction = pycolmap.Reconstruction() + + point3d_ids = [] + for vidx in range(num_points): + point3d_id = reconstruction.add_point3D(points[vidx], pycolmap.Track(), colors[vidx]) + point3d_ids.append(point3d_id) + + for fidx in range(num_frames): + orig_w, orig_h = Image.open(image_paths[fidx]).size + + intrinsic = prediction.intrinsics[fidx] + if process_res_method.endswith("resize"): + intrinsic[:1] *= orig_w / w + intrinsic[1:2] *= orig_h / h + elif process_res_method == "crop": + raise NotImplementedError("COLMAP export for crop method is not implemented") + else: + raise ValueError(f"Unknown process_res_method: {process_res_method}") + + pycolmap_intri = np.array( + [intrinsic[0, 0], intrinsic[1, 1], intrinsic[0, 2], intrinsic[1, 2]] + ) + + extrinsic = prediction.extrinsics[fidx] + cam_from_world = pycolmap.Rigid3d(pycolmap.Rotation3d(extrinsic[:3, :3]), extrinsic[:3, 3]) + + # set and add camera + camera = pycolmap.Camera() + camera.camera_id = fidx + 1 + camera.model = pycolmap.CameraModelId.PINHOLE + camera.width = orig_w + camera.height = orig_h + camera.params = pycolmap_intri + reconstruction.add_camera(camera) + + # set and add rig (from camera) + rig = pycolmap.Rig() + rig.rig_id = camera.camera_id + rig.add_ref_sensor(camera.sensor_id) + reconstruction.add_rig(rig) + + # set image + image = pycolmap.Image() + image.image_id = fidx + 1 + image.camera_id = camera.camera_id + + # set and add frame (from image) + frame = pycolmap.Frame() + frame.frame_id = image.image_id + frame.rig_id = camera.camera_id + frame.add_data_id(image.data_id) + frame.rig_from_world = cam_from_world + reconstruction.add_frame(frame) + + # set point2d and update track + point2d_list = [] + points_in_frame = points_xyf[:, 2].astype(np.int32) == fidx + for vidx in np.where(points_in_frame)[0]: + point2d = points_xyf[vidx][:2] + point2d[0] *= orig_w / w + point2d[1] *= orig_h / h + point3d_id = point3d_ids[vidx] + point2d_list.append(pycolmap.Point2D(point2d, point3d_id)) + reconstruction.point3D(point3d_id).track.add_element( + image.image_id, len(point2d_list) - 1 + ) + + # set and add image + image.frame_id = image.image_id + image.name = os.path.basename(image_paths[fidx]) + image.points2D = pycolmap.Point2DList(point2d_list) + reconstruction.add_image(image) + + # 3. Export + reconstruction.write(export_dir) + + +def _create_xyf(num_frames, height, width): + """ + Creates a grid of pixel coordinates and frame indices (fidx) for all frames. + """ + # Create coordinate grids for a single frame + y_grid, x_grid = np.indices((height, width), dtype=np.int32) + x_grid = x_grid[np.newaxis, :, :] + y_grid = y_grid[np.newaxis, :, :] + + # Broadcast to all frames + x_coords = np.broadcast_to(x_grid, (num_frames, height, width)) + y_coords = np.broadcast_to(y_grid, (num_frames, height, width)) + + # Create frame indices and broadcast + f_idx = np.arange(num_frames, dtype=np.int32)[:, np.newaxis, np.newaxis] + f_coords = np.broadcast_to(f_idx, (num_frames, height, width)) + + # Stack coordinates and frame indices + points_xyf = np.stack((x_coords, y_coords, f_coords), axis=-1) + + return points_xyf diff --git a/src/nn/depth_anything_3/utils/export/depth_vis.py b/src/nn/depth_anything_3/utils/export/depth_vis.py new file mode 100644 index 0000000..8accc04 --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/depth_vis.py @@ -0,0 +1,41 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import imageio +import numpy as np + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.visualize import visualize_depth + + +def export_to_depth_vis( + prediction: Prediction, + export_dir: str, +): + # Use prediction.processed_images, which is already processed image data + if prediction.processed_images is None: + raise ValueError("prediction.processed_images is required but not available") + + images_u8 = prediction.processed_images # (N,H,W,3) uint8 + + os.makedirs(os.path.join(export_dir, "depth_vis"), exist_ok=True) + for idx in range(prediction.depth.shape[0]): + depth_vis = visualize_depth(prediction.depth[idx]) + image_vis = images_u8[idx] + depth_vis = depth_vis.astype(np.uint8) + image_vis = image_vis.astype(np.uint8) + vis_image = np.concatenate([image_vis, depth_vis], axis=1) + save_path = os.path.join(export_dir, f"depth_vis/{idx:04d}.jpg") + imageio.imwrite(save_path, vis_image, quality=95) diff --git a/src/nn/depth_anything_3/utils/export/feat_vis.py b/src/nn/depth_anything_3/utils/export/feat_vis.py new file mode 100644 index 0000000..e3dc780 --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/feat_vis.py @@ -0,0 +1,65 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import cv2 +import imageio +import numpy as np +from tqdm.auto import tqdm + +from depth_anything_3.utils.parallel_utils import async_call +from depth_anything_3.utils.pca_utils import PCARGBVisualizer + + +@async_call +def export_to_feat_vis( + prediction, + export_dir, + fps=15, +): + """Export feature visualization with PCA. + + Args: + prediction: Model prediction containing feature maps + export_dir: Directory to export results + fps: Frame rate for output video (default: 15) + """ + out_dir = os.path.join(export_dir, "feat_vis") + os.makedirs(out_dir, exist_ok=True) + + images = prediction.processed_images + for k, v in prediction.aux.items(): + if not k.startswith("feat_layer_"): + continue + os.makedirs(os.path.join(out_dir, k), exist_ok=True) + viz = PCARGBVisualizer(basis_mode="fixed", percentile_mode="global", clip_percent=10.0) + viz.fit_reference(v) + feats_vis = viz.transform_video(v) + for idx in tqdm(range(len(feats_vis))): + img = images[idx] + feat_vis = (feats_vis[idx] * 255).astype(np.uint8) + feat_vis = cv2.resize( + feat_vis, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST + ) + save_path = os.path.join(out_dir, f"{k}/{idx:06d}.jpg") + save = np.concatenate([img, feat_vis], axis=1) + imageio.imwrite(save_path, save, quality=95) + cmd = ( + "ffmpeg -loglevel error -hide_banner -y " + f"-framerate {fps} -start_number 0 " + f"-i {out_dir}/{k}/%06d.jpg " + f"-c:v libx264 -pix_fmt yuv420p " + f"{out_dir}/{k}.mp4" + ) + os.system(cmd) diff --git a/src/nn/depth_anything_3/utils/export/glb.py b/src/nn/depth_anything_3/utils/export/glb.py new file mode 100644 index 0000000..ece1379 --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/glb.py @@ -0,0 +1,432 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import os +import numpy as np +import trimesh + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.logger import logger + +from .depth_vis import export_to_depth_vis + + +def set_sky_depth(prediction: Prediction, sky_mask: np.ndarray, sky_depth_def: float = 98.0): + non_sky_mask = ~sky_mask + valid_depth = prediction.depth[non_sky_mask] + if valid_depth.size > 0: + max_depth = np.percentile(valid_depth, sky_depth_def) + prediction.depth[sky_mask] = max_depth + + +def get_conf_thresh( + prediction: Prediction, + sky_mask: np.ndarray, + conf_thresh: float, + conf_thresh_percentile: float = 10.0, + ensure_thresh_percentile: float = 90.0, +): + if sky_mask is not None and (~sky_mask).sum() > 10: + conf_pixels = prediction.conf[~sky_mask] + else: + conf_pixels = prediction.conf + lower = np.percentile(conf_pixels, conf_thresh_percentile) + upper = np.percentile(conf_pixels, ensure_thresh_percentile) + conf_thresh = min(max(conf_thresh, lower), upper) + return conf_thresh + + +def export_to_glb( + prediction: Prediction, + export_dir: str, + num_max_points: int = 1_000_000, + conf_thresh: float = 1.05, + filter_black_bg: bool = False, + filter_white_bg: bool = False, + conf_thresh_percentile: float = 40.0, + ensure_thresh_percentile: float = 90.0, + sky_depth_def: float = 98.0, + show_cameras: bool = True, + camera_size: float = 0.03, + export_depth_vis: bool = True, +) -> str: + """Generate a 3D point cloud and camera wireframes and export them as a ``.glb`` file. + + The function builds a point cloud from the predicted depth maps, aligns it to the + first camera in glTF coordinates (X-right, Y-up, Z-backward), optionally draws + camera wireframes, and writes the result to ``scene.glb``. Auxiliary assets such as + depth visualizations can also be generated alongside the main export. + + Args: + prediction: Model prediction containing depth, confidence, intrinsics, extrinsics, + and pre-processed images. + export_dir: Output directory where the glTF assets will be written. + num_max_points: Maximum number of points retained after downsampling. + conf_thresh: Base confidence threshold used before percentile adjustments. + filter_black_bg: Mark near-black background pixels for removal during confidence filtering. + filter_white_bg: Mark near-white background pixels for removal during confidence filtering. + conf_thresh_percentile: Lower percentile used when adapting the confidence threshold. + ensure_thresh_percentile: Upper percentile clamp for the adaptive threshold. + sky_depth_def: Percentile used to fill sky pixels with plausible depth values. + show_cameras: Whether to render camera wireframes in the exported scene. + camera_size: Relative camera wireframe scale as a fraction of the scene diagonal. + export_depth_vis: Whether to export raster depth visualisations alongside the glTF. + + Returns: + Path to the exported ``scene.glb`` file. + """ + # 1) Use prediction.processed_images, which is already processed image data + assert ( + prediction.processed_images is not None + ), "Export to GLB: prediction.processed_images is required but not available" + assert ( + prediction.depth is not None + ), "Export to GLB: prediction.depth is required but not available" + assert ( + prediction.intrinsics is not None + ), "Export to GLB: prediction.intrinsics is required but not available" + assert ( + prediction.extrinsics is not None + ), "Export to GLB: prediction.extrinsics is required but not available" + assert ( + prediction.conf is not None + ), "Export to GLB: prediction.conf is required but not available" + logger.info(f"conf_thresh_percentile: {conf_thresh_percentile}") + logger.info(f"num max points: {num_max_points}") + logger.info(f"Exporting to GLB with num_max_points: {num_max_points}") + if prediction.processed_images is None: + raise ValueError("prediction.processed_images is required but not available") + + images_u8 = prediction.processed_images # (N,H,W,3) uint8 + + # 2) Sky processing (if sky_mask is provided) + if getattr(prediction, "sky_mask", None) is not None: + set_sky_depth(prediction, prediction.sky_mask, sky_depth_def) + + # 3) Confidence threshold (if no conf, then no filtering) + if filter_black_bg: + prediction.conf[(prediction.processed_images < 16).all(axis=-1)] = 1.0 + if filter_white_bg: + prediction.conf[(prediction.processed_images >= 240).all(axis=-1)] = 1.0 + conf_thr = get_conf_thresh( + prediction, + getattr(prediction, "sky_mask", None), + conf_thresh, + conf_thresh_percentile, + ensure_thresh_percentile, + ) + + # 4) Back-project to world coordinates and get colors (world frame) + points, colors = _depths_to_world_points_with_colors( + prediction.depth, + prediction.intrinsics, + prediction.extrinsics, # w2c + images_u8, + prediction.conf, + conf_thr, + ) + + # 5) Based on first camera orientation + glTF axis system, center by point cloud, + # construct alignment transform, and apply to point cloud + A = _compute_alignment_transform_first_cam_glTF_center_by_points( + prediction.extrinsics[0], points + ) # (4,4) + + if points.shape[0] > 0: + points = trimesh.transform_points(points, A) + + # 6) Clean + downsample + points, colors = _filter_and_downsample(points, colors, num_max_points) + + # 7) Assemble scene (add point cloud first) + scene = trimesh.Scene() + if scene.metadata is None: + scene.metadata = {} + scene.metadata["hf_alignment"] = A # For camera wireframes and external reuse + + if points.shape[0] > 0: + pc = trimesh.points.PointCloud(vertices=points, colors=colors) + scene.add_geometry(pc) + + # 8) Draw cameras (wireframe pyramids), using the same transform A + if show_cameras and prediction.intrinsics is not None and prediction.extrinsics is not None: + scene_scale = _estimate_scene_scale(points, fallback=1.0) + H, W = prediction.depth.shape[1:] + _add_cameras_to_scene( + scene=scene, + K=prediction.intrinsics, + ext_w2c=prediction.extrinsics, + image_sizes=[(H, W)] * prediction.depth.shape[0], + scale=scene_scale * camera_size, + ) + + # 9) Export + os.makedirs(export_dir, exist_ok=True) + out_path = os.path.join(export_dir, "scene.glb") + scene.export(out_path) + + if export_depth_vis: + export_to_depth_vis(prediction, export_dir) + os.system(f"cp -r {export_dir}/depth_vis/0000.jpg {export_dir}/scene.jpg") + return out_path + + +# ========================= +# utilities +# ========================= + + +def _as_homogeneous44(ext: np.ndarray) -> np.ndarray: + """ + Accept (4,4) or (3,4) extrinsic parameters, return (4,4) homogeneous matrix. + """ + if ext.shape == (4, 4): + return ext + if ext.shape == (3, 4): + H = np.eye(4, dtype=ext.dtype) + H[:3, :4] = ext + return H + raise ValueError(f"extrinsic must be (4,4) or (3,4), got {ext.shape}") + + +def _depths_to_world_points_with_colors( + depth: np.ndarray, + K: np.ndarray, + ext_w2c: np.ndarray, + images_u8: np.ndarray, + conf: np.ndarray | None, + conf_thr: float, +) -> tuple[np.ndarray, np.ndarray]: + """ + For each frame, transform (u,v,1) through K^{-1} to get rays, + multiply by depth to camera frame, then use (w2c)^{-1} to transform to world frame. + Simultaneously extract colors. + """ + N, H, W = depth.shape + us, vs = np.meshgrid(np.arange(W), np.arange(H)) + ones = np.ones_like(us) + pix = np.stack([us, vs, ones], axis=-1).reshape(-1, 3) # (H*W,3) + + pts_all, col_all = [], [] + + for i in range(N): + d = depth[i] # (H,W) + valid = np.isfinite(d) & (d > 0) + if conf is not None: + valid &= conf[i] >= conf_thr + if not np.any(valid): + continue + + d_flat = d.reshape(-1) + vidx = np.flatnonzero(valid.reshape(-1)) + + K_inv = np.linalg.inv(K[i]) # (3,3) + c2w = np.linalg.inv(_as_homogeneous44(ext_w2c[i])) # (4,4) + + rays = K_inv @ pix[vidx].T # (3,M) + Xc = rays * d_flat[vidx][None, :] # (3,M) + Xc_h = np.vstack([Xc, np.ones((1, Xc.shape[1]))]) + Xw = (c2w @ Xc_h)[:3].T.astype(np.float32) # (M,3) + + cols = images_u8[i].reshape(-1, 3)[vidx].astype(np.uint8) # (M,3) + + pts_all.append(Xw) + col_all.append(cols) + + if len(pts_all) == 0: + return np.zeros((0, 3), dtype=np.float32), np.zeros((0, 3), dtype=np.uint8) + + return np.concatenate(pts_all, 0), np.concatenate(col_all, 0) + + +def _filter_and_downsample(points: np.ndarray, colors: np.ndarray, num_max: int): + if points.shape[0] == 0: + return points, colors + finite = np.isfinite(points).all(axis=1) + points, colors = points[finite], colors[finite] + if points.shape[0] > num_max: + idx = np.random.choice(points.shape[0], num_max, replace=False) + points, colors = points[idx], colors[idx] + return points, colors + + +def _estimate_scene_scale(points: np.ndarray, fallback: float = 1.0) -> float: + if points.shape[0] < 2: + return fallback + lo = np.percentile(points, 5, axis=0) + hi = np.percentile(points, 95, axis=0) + diag = np.linalg.norm(hi - lo) + return float(diag if np.isfinite(diag) and diag > 0 else fallback) + + +def _compute_alignment_transform_first_cam_glTF_center_by_points( + ext_w2c0: np.ndarray, + points_world: np.ndarray, +) -> np.ndarray: + """Computes the transformation matrix to align the scene with glTF standards. + + This function calculates a 4x4 homogeneous matrix that centers the scene's + point cloud and transforms its coordinate system from the computer vision (CV) + standard to the glTF standard. + + The transformation process involves three main steps: + 1. **Initial Alignment**: Orients the world coordinate system to match the + first camera's view (x-right, y-down, z-forward). + 2. **Coordinate System Conversion**: Converts the CV camera frame to the + glTF frame (x-right, y-up, z-backward) by flipping the Y and Z axes. + 3. **Centering**: Translates the entire scene so that the median of the + point cloud becomes the new origin (0,0,0). + + Returns: + A 4x4 homogeneous transformation matrix (torch.Tensor or np.ndarray) + that applies these transformations. A: X' = A @ [X;1] + """ + + w2c0 = _as_homogeneous44(ext_w2c0).astype(np.float64) + + # CV -> glTF axis transformation + M = np.eye(4, dtype=np.float64) + M[1, 1] = -1.0 # flip Y + M[2, 2] = -1.0 # flip Z + + # Don't center first + A_no_center = M @ w2c0 + + # Calculate point cloud center in new coordinate system (use median to resist outliers) + if points_world.shape[0] > 0: + pts_tmp = trimesh.transform_points(points_world, A_no_center) + center = np.median(pts_tmp, axis=0) + else: + center = np.zeros(3, dtype=np.float64) + + T_center = np.eye(4, dtype=np.float64) + T_center[:3, 3] = -center + + A = T_center @ A_no_center + return A + + +def _add_cameras_to_scene( + scene: trimesh.Scene, + K: np.ndarray, + ext_w2c: np.ndarray, + image_sizes: list[tuple[int, int]], + scale: float, +) -> None: + """Draws camera frustums to visualize their position and orientation. + + This function renders each camera as a wireframe pyramid, originating from + the camera's center and extending to the corners of its imaging plane. + + It reads the 'hf_alignment' metadata from the scene to ensure the + wireframes are correctly aligned with the 3D point cloud. + """ + N = K.shape[0] + if N == 0: + return + + # Alignment matrix consistent with point cloud (use identity matrix if missing) + A = None + try: + A = scene.metadata.get("hf_alignment", None) if scene.metadata else None + except Exception: + A = None + if A is None: + A = np.eye(4, dtype=np.float64) + + for i in range(N): + H, W = image_sizes[i] + segs = _camera_frustum_lines(K[i], ext_w2c[i], W, H, scale) # (8,2,3) world frame + # Apply unified transformation + segs = trimesh.transform_points(segs.reshape(-1, 3), A).reshape(-1, 2, 3) + path = trimesh.load_path(segs) + color = _index_color_rgb(i, N) + if hasattr(path, "colors"): + path.colors = np.tile(color, (len(path.entities), 1)) + scene.add_geometry(path) + + +def _camera_frustum_lines( + K: np.ndarray, ext_w2c: np.ndarray, W: int, H: int, scale: float +) -> np.ndarray: + corners = np.array( + [ + [0, 0, 1.0], + [W - 1, 0, 1.0], + [W - 1, H - 1, 1.0], + [0, H - 1, 1.0], + ], + dtype=float, + ) # (4,3) + + K_inv = np.linalg.inv(K) + c2w = np.linalg.inv(_as_homogeneous44(ext_w2c)) + + # camera center in world + Cw = (c2w @ np.array([0, 0, 0, 1.0]))[:3] + + # rays -> z=1 plane points (camera frame) + rays = (K_inv @ corners.T).T + z = rays[:, 2:3] + z[z == 0] = 1.0 + plane_cam = (rays / z) * scale # (4,3) + + # to world + plane_w = [] + for p in plane_cam: + pw = (c2w @ np.array([p[0], p[1], p[2], 1.0]))[:3] + plane_w.append(pw) + plane_w = np.stack(plane_w, 0) # (4,3) + + segs = [] + # center to corners + for k in range(4): + segs.append(np.stack([Cw, plane_w[k]], 0)) + # rectangle edges + order = [0, 1, 2, 3, 0] + for a, b in zip(order[:-1], order[1:]): + segs.append(np.stack([plane_w[a], plane_w[b]], 0)) + + return np.stack(segs, 0) # (8,2,3) + + +def _index_color_rgb(i: int, n: int) -> np.ndarray: + h = (i + 0.5) / max(n, 1) + s, v = 0.85, 0.95 + r, g, b = _hsv_to_rgb(h, s, v) + return (np.array([r, g, b]) * 255).astype(np.uint8) + + +def _hsv_to_rgb(h: float, s: float, v: float) -> tuple[float, float, float]: + i = int(h * 6.0) + f = h * 6.0 - i + p = v * (1.0 - s) + q = v * (1.0 - f * s) + t = v * (1.0 - (1.0 - f) * s) + i = i % 6 + if i == 0: + r, g, b = v, t, p + elif i == 1: + r, g, b = q, v, p + elif i == 2: + r, g, b = p, v, t + elif i == 3: + r, g, b = p, q, v + elif i == 4: + r, g, b = t, p, v + else: + r, g, b = v, p, q + return r, g, b diff --git a/src/nn/depth_anything_3/utils/export/gs.py b/src/nn/depth_anything_3/utils/export/gs.py new file mode 100644 index 0000000..90077cf --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/gs.py @@ -0,0 +1,154 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from typing import Literal, Optional +import moviepy.editor as mpy +import torch + +from depth_anything_3.model.utils.gs_renderer import run_renderer_in_chunk_w_trj_mode +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.gsply_helpers import save_gaussian_ply +from depth_anything_3.utils.layout_helpers import hcat, vcat +from depth_anything_3.utils.visualize import vis_depth_map_tensor + +VIDEO_QUALITY_MAP = { + "low": {"crf": "28", "preset": "veryfast"}, + "medium": {"crf": "23", "preset": "medium"}, + "high": {"crf": "18", "preset": "slow"}, +} + + +def export_to_gs_ply( + prediction: Prediction, + export_dir: str, + gs_views_interval: Optional[ + int + ] = 1, # export GS every N views, useful for extremely dense inputs +): + gs_world = prediction.gaussians + pred_depth = torch.from_numpy(prediction.depth).unsqueeze(-1).to(gs_world.means) # v h w 1 + idx = 0 + os.makedirs(os.path.join(export_dir, "gs_ply"), exist_ok=True) + save_path = os.path.join(export_dir, f"gs_ply/{idx:04d}.ply") + if gs_views_interval is None: # select around 12 views in total + gs_views_interval = max(pred_depth.shape[0] // 12, 1) + save_gaussian_ply( + gaussians=gs_world, + save_path=save_path, + ctx_depth=pred_depth, + shift_and_scale=False, + save_sh_dc_only=True, + gs_views_interval=gs_views_interval, + inv_opacity=True, + prune_by_depth_percent=0.9, + prune_border_gs=True, + match_3dgs_mcmc_dev=False, + ) + + +def export_to_gs_video( + prediction: Prediction, + export_dir: str, + extrinsics: Optional[torch.Tensor] = None, # render views' world2cam, "b v 4 4" + intrinsics: Optional[torch.Tensor] = None, # render views' unnormed intrinsics, "b v 3 3" + out_image_hw: Optional[tuple[int, int]] = None, # render views' resolution, (h, w) + chunk_size: Optional[int] = 4, + trj_mode: Literal[ + "original", + "smooth", + "interpolate", + "interpolate_smooth", + "wander", + "dolly_zoom", + "extend", + "wobble_inter", + ] = "extend", + color_mode: Literal["RGB+D", "RGB+ED"] = "RGB+ED", + vis_depth: Optional[Literal["hcat", "vcat"]] = "hcat", + enable_tqdm: Optional[bool] = True, + output_name: Optional[str] = None, + video_quality: Literal["low", "medium", "high"] = "high", +) -> None: + gs_world = prediction.gaussians + # if target poses are not provided, render the (smooth/interpolate) input poses + if extrinsics is not None: + tgt_extrs = extrinsics + else: + tgt_extrs = torch.from_numpy(prediction.extrinsics).unsqueeze(0).to(gs_world.means) + if prediction.is_metric: + scale_factor = prediction.scale_factor + if scale_factor is not None: + tgt_extrs[:, :, :3, 3] /= scale_factor + tgt_intrs = ( + intrinsics + if intrinsics is not None + else torch.from_numpy(prediction.intrinsics).unsqueeze(0).to(gs_world.means) + ) + # if render resolution is not provided, render the input ones + if out_image_hw is not None: + H, W = out_image_hw + else: + H, W = prediction.depth.shape[-2:] + # if single views, render wander trj + if tgt_extrs.shape[1] <= 1: + trj_mode = "wander" + # trj_mode = "dolly_zoom" + + color, depth = run_renderer_in_chunk_w_trj_mode( + gaussians=gs_world, + extrinsics=tgt_extrs, + intrinsics=tgt_intrs, + image_shape=(H, W), + chunk_size=chunk_size, + trj_mode=trj_mode, + use_sh=True, + color_mode=color_mode, + enable_tqdm=enable_tqdm, + ) + + # save as video + ffmpeg_params = [ + "-crf", + VIDEO_QUALITY_MAP[video_quality]["crf"], + "-preset", + VIDEO_QUALITY_MAP[video_quality]["preset"], + "-pix_fmt", + "yuv420p", + ] # best compatibility + + os.makedirs(os.path.join(export_dir, "gs_video"), exist_ok=True) + for idx in range(color.shape[0]): + video_i = color[idx] + if vis_depth is not None: + depth_i = vis_depth_map_tensor(depth[0]) + cat_fn = hcat if vis_depth == "hcat" else vcat + video_i = torch.stack([cat_fn(c, d) for c, d in zip(video_i, depth_i)]) + frames = list( + (video_i.clamp(0, 1) * 255).byte().permute(0, 2, 3, 1).cpu().numpy() + ) # T x H x W x C, uint8, numpy() + + fps = 24 + clip = mpy.ImageSequenceClip(frames, fps=fps) + output_name = f"{idx:04d}_{trj_mode}" if output_name is None else output_name + save_path = os.path.join(export_dir, f"gs_video/{output_name}.mp4") + # clip.write_videofile(save_path, codec="libx264", audio=False, bitrate="4000k") + clip.write_videofile( + save_path, + codec="libx264", + audio=False, + fps=fps, + ffmpeg_params=ffmpeg_params, + ) + return diff --git a/src/nn/depth_anything_3/utils/export/npz.py b/src/nn/depth_anything_3/utils/export/npz.py new file mode 100644 index 0000000..75df4f5 --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/npz.py @@ -0,0 +1,73 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import numpy as np + +from depth_anything_3.specs import Prediction +from depth_anything_3.utils.parallel_utils import async_call + + +@async_call +def export_to_npz( + prediction: Prediction, + export_dir: str, +): + output_file = os.path.join(export_dir, "exports", "npz", "results.npz") + os.makedirs(os.path.dirname(output_file), exist_ok=True) + + # Use prediction.processed_images, which is already processed image data + if prediction.processed_images is None: + raise ValueError("prediction.processed_images is required but not available") + + image = prediction.processed_images # (N,H,W,3) uint8 + + # Build save dict with only non-None values + save_dict = { + "image": image, + "depth": np.round(prediction.depth, 6), + } + + if prediction.conf is not None: + save_dict["conf"] = np.round(prediction.conf, 2) + if prediction.extrinsics is not None: + save_dict["extrinsics"] = prediction.extrinsics + if prediction.intrinsics is not None: + save_dict["intrinsics"] = prediction.intrinsics + + # aux = {k: np.round(v, 4) for k, v in prediction.aux.items()} + np.savez_compressed(output_file, **save_dict) + + +@async_call +def export_to_mini_npz( + prediction: Prediction, + export_dir: str, +): + output_file = os.path.join(export_dir, "exports", "mini_npz", "results.npz") + os.makedirs(os.path.dirname(output_file), exist_ok=True) + + # Build save dict with only non-None values + save_dict = { + "depth": np.round(prediction.depth, 8), + } + + if prediction.conf is not None: + save_dict["conf"] = np.round(prediction.conf, 2) + if prediction.extrinsics is not None: + save_dict["extrinsics"] = prediction.extrinsics + if prediction.intrinsics is not None: + save_dict["intrinsics"] = prediction.intrinsics + + np.savez_compressed(output_file, **save_dict) diff --git a/src/nn/depth_anything_3/utils/export/utils.py b/src/nn/depth_anything_3/utils/export/utils.py new file mode 100644 index 0000000..81f45fb --- /dev/null +++ b/src/nn/depth_anything_3/utils/export/utils.py @@ -0,0 +1,30 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + + +def _denorm_and_to_uint8(image_tensor: torch.Tensor) -> np.ndarray: + """Denormalize to [0,255] and output (N, H, W, 3) uint8.""" + resnet_mean = torch.tensor( + [0.485, 0.456, 0.406], dtype=image_tensor.dtype, device=image_tensor.device + ) + resnet_std = torch.tensor( + [0.229, 0.224, 0.225], dtype=image_tensor.dtype, device=image_tensor.device + ) + img = image_tensor * resnet_std[None, :, None, None] + resnet_mean[None, :, None, None] + img = torch.clamp(img, 0.0, 1.0) + img = (img.permute(0, 2, 3, 1).cpu().numpy() * 255.0).round().astype(np.uint8) # (N,H,W,3) + return img diff --git a/src/nn/depth_anything_3/utils/geometry.py b/src/nn/depth_anything_3/utils/geometry.py new file mode 100644 index 0000000..41a4219 --- /dev/null +++ b/src/nn/depth_anything_3/utils/geometry.py @@ -0,0 +1,498 @@ +# flake8: noqa: F722 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from types import SimpleNamespace +from typing import Optional +import numpy as np +import torch +import torch.nn.functional as F +from einops import einsum + + +def as_homogeneous(ext): + """ + Accept (..., 3,4) or (..., 4,4) extrinsics, return (...,4,4) homogeneous matrix. + Supports torch.Tensor or np.ndarray. + """ + if isinstance(ext, torch.Tensor): + # If already in homogeneous form + if ext.shape[-2:] == (4, 4): + return ext + elif ext.shape[-2:] == (3, 4): + # Create a new homogeneous matrix + ones = torch.zeros_like(ext[..., :1, :4]) + ones[..., 0, 3] = 1.0 + return torch.cat([ext, ones], dim=-2) + else: + raise ValueError(f"Invalid shape for torch.Tensor: {ext.shape}") + + elif isinstance(ext, np.ndarray): + if ext.shape[-2:] == (4, 4): + return ext + elif ext.shape[-2:] == (3, 4): + ones = np.zeros_like(ext[..., :1, :4]) + ones[..., 0, 3] = 1.0 + return np.concatenate([ext, ones], axis=-2) + else: + raise ValueError(f"Invalid shape for np.ndarray: {ext.shape}") + + else: + raise TypeError("Input must be a torch.Tensor or np.ndarray.") + + +@torch.jit.script +def affine_inverse(A: torch.Tensor): + R = A[..., :3, :3] # ..., 3, 3 + T = A[..., :3, 3:] # ..., 3, 1 + P = A[..., 3:, :] # ..., 1, 4 + return torch.cat([torch.cat([R.mT, -R.mT @ T], dim=-1), P], dim=-2) + + +def transpose_last_two_axes(arr): + """ + for np < 2 + """ + if arr.ndim < 2: + return arr + axes = list(range(arr.ndim)) + # swap the last two + axes[-2], axes[-1] = axes[-1], axes[-2] + return arr.transpose(axes) + + +def affine_inverse_np(A: np.ndarray): + R = A[..., :3, :3] + T = A[..., :3, 3:] + P = A[..., 3:, :] + return np.concatenate( + [ + np.concatenate([transpose_last_two_axes(R), -transpose_last_two_axes(R) @ T], axis=-1), + P, + ], + axis=-2, + ) + + +def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor: + """ + Quaternion Order: XYZW or say ijkr, scalar-last + + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + i, j, k, r = torch.unbind(quaternions, -1) + # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape( + batch_dim + (4,) + ) + + # Convert from rijk to ijkr + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def sample_image_grid( + shape: tuple[int, ...], + device: torch.device = torch.device("cpu"), +) -> tuple[ + torch.Tensor, # float coordinates (xy indexing), "*shape dim" + torch.Tensor, # integer indices (ij indexing), "*shape dim" +]: + """Get normalized (range 0 to 1) coordinates and integer indices for an image.""" + + # Each entry is a pixel-wise integer coordinate. In the 2D case, each entry is a + # (row, col) coordinate. + indices = [torch.arange(length, device=device) for length in shape] + stacked_indices = torch.stack(torch.meshgrid(*indices, indexing="ij"), dim=-1) + + # Each entry is a floating-point coordinate in the range (0, 1). In the 2D case, + # each entry is an (x, y) coordinate. + coordinates = [(idx + 0.5) / length for idx, length in zip(indices, shape)] + coordinates = reversed(coordinates) + coordinates = torch.stack(torch.meshgrid(*coordinates, indexing="xy"), dim=-1) + + return coordinates, stacked_indices + + +def homogenize_points(points: torch.Tensor) -> torch.Tensor: # "*batch dim" # "*batch dim+1" + """Convert batched points (xyz) to (xyz1).""" + return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1) + + +def homogenize_vectors(vectors: torch.Tensor) -> torch.Tensor: # "*batch dim" # "*batch dim+1" + """Convert batched vectors (xyz) to (xyz0).""" + return torch.cat([vectors, torch.zeros_like(vectors[..., :1])], dim=-1) + + +def transform_rigid( + homogeneous_coordinates: torch.Tensor, # "*#batch dim" + transformation: torch.Tensor, # "*#batch dim dim" +) -> torch.Tensor: # "*batch dim" + """Apply a rigid-body transformation to points or vectors.""" + return einsum( + transformation, + homogeneous_coordinates.to(transformation.dtype), + "... i j, ... j -> ... i", + ) + + +def transform_cam2world( + homogeneous_coordinates: torch.Tensor, # "*#batch dim" + extrinsics: torch.Tensor, # "*#batch dim dim" +) -> torch.Tensor: # "*batch dim" + """Transform points from 3D camera coordinates to 3D world coordinates.""" + return transform_rigid(homogeneous_coordinates, extrinsics) + + +def unproject( + coordinates: torch.Tensor, # "*#batch dim" + z: torch.Tensor, # "*#batch" + intrinsics: torch.Tensor, # "*#batch dim+1 dim+1" +) -> torch.Tensor: # "*batch dim+1" + """Unproject 2D camera coordinates with the given Z values.""" + + # Apply the inverse intrinsics to the coordinates. + coordinates = homogenize_points(coordinates) + ray_directions = einsum( + intrinsics.float().inverse().to(intrinsics), + coordinates.to(intrinsics.dtype), + "... i j, ... j -> ... i", + ) + + # Apply the supplied depth values. + return ray_directions * z[..., None] + + +def get_world_rays( + coordinates: torch.Tensor, # "*#batch dim" + extrinsics: torch.Tensor, # "*#batch dim+2 dim+2" + intrinsics: torch.Tensor, # "*#batch dim+1 dim+1" +) -> tuple[ + torch.Tensor, # origins, "*batch dim+1" + torch.Tensor, # directions, "*batch dim+1" +]: + # Get camera-space ray directions. + directions = unproject( + coordinates, + torch.ones_like(coordinates[..., 0]), + intrinsics, + ) + directions = directions / directions.norm(dim=-1, keepdim=True) + + # Transform ray directions to world coordinates. + directions = homogenize_vectors(directions) + directions = transform_cam2world(directions, extrinsics)[..., :-1] + + # Tile the ray origins to have the same shape as the ray directions. + origins = extrinsics[..., :-1, -1].broadcast_to(directions.shape) + + return origins, directions + + +def get_fov(intrinsics: torch.Tensor) -> torch.Tensor: # "batch 3 3" -> "batch 2" + intrinsics_inv = intrinsics.float().inverse().to(intrinsics) + + def process_vector(vector): + vector = torch.tensor(vector, dtype=intrinsics.dtype, device=intrinsics.device) + vector = einsum(intrinsics_inv, vector, "b i j, j -> b i") + return vector / vector.norm(dim=-1, keepdim=True) + + left = process_vector([0, 0.5, 1]) + right = process_vector([1, 0.5, 1]) + top = process_vector([0.5, 0, 1]) + bottom = process_vector([0.5, 1, 1]) + fov_x = (left * right).sum(dim=-1).acos() + fov_y = (top * bottom).sum(dim=-1).acos() + return torch.stack((fov_x, fov_y), dim=-1) + + +def map_pdf_to_opacity( + pdf: torch.Tensor, # " *batch" + global_step: int = 0, + opacity_mapping: Optional[dict] = None, +) -> torch.Tensor: # " *batch" + # https://www.desmos.com/calculator/opvwti3ba9 + + # Figure out the exponent. + if opacity_mapping is not None: + cfg = SimpleNamespace(**opacity_mapping) + x = cfg.initial + min(global_step / cfg.warm_up, 1) * (cfg.final - cfg.initial) + else: + x = 0.0 + exponent = 2**x + + # Map the probability density to an opacity. + return 0.5 * (1 - (1 - pdf) ** exponent + pdf ** (1 / exponent)) + +def normalize_homogenous_points(points): + """Normalize the point vectors""" + return points / points[..., -1:] + +def inverse_intrinsic_matrix(ixts): + """ """ + return torch.inverse(ixts) + +def pixel_space_to_camera_space(pixel_space_points, depth, intrinsics): + """ + Convert pixel space points to camera space points. + + Args: + pixel_space_points (torch.Tensor): Pixel space points with shape (h, w, 2) + depth (torch.Tensor): Depth map with shape (b, v, h, w, 1) + intrinsics (torch.Tensor): Camera intrinsics with shape (b, v, 3, 3) + + Returns: + torch.Tensor: Camera space points with shape (b, v, h, w, 3). + """ + pixel_space_points = homogenize_points(pixel_space_points) + # camera_space_points = torch.einsum( + # "b v i j , h w j -> b v h w i", intrinsics.inverse(), pixel_space_points + # ) + camera_space_points = torch.einsum( + "b v i j , h w j -> b v h w i", inverse_intrinsic_matrix(intrinsics), pixel_space_points + ) + camera_space_points = camera_space_points * depth + return camera_space_points + + +def camera_space_to_world_space(camera_space_points, c2w): + """ + Convert camera space points to world space points. + + Args: + camera_space_points (torch.Tensor): Camera space points with shape (b, v, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v, 4, 4) + + Returns: + torch.Tensor: World space points with shape (b, v, h, w, 3). + """ + camera_space_points = homogenize_points(camera_space_points) + world_space_points = torch.einsum("b v i j , b v h w j -> b v h w i", c2w, camera_space_points) + return world_space_points[..., :3] + + +def camera_space_to_pixel_space(camera_space_points, intrinsics): + """ + Convert camera space points to pixel space points. + + Args: + camera_space_points (torch.Tensor): Camera space points with shape (b, v1, v2, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v2, 3, 3) + + Returns: + torch.Tensor: World space points with shape (b, v1, v2, h, w, 2). + """ + camera_space_points = normalize_homogenous_points(camera_space_points) + pixel_space_points = torch.einsum( + "b u i j , b v u h w j -> b v u h w i", intrinsics, camera_space_points + ) + return pixel_space_points[..., :2] + + +def world_space_to_camera_space(world_space_points, c2w): + """ + Convert world space points to pixel space points. + + Args: + world_space_points (torch.Tensor): World space points with shape (b, v1, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v2, 4, 4) + + Returns: + torch.Tensor: Camera space points with shape (b, v1, v2, h, w, 3). + """ + world_space_points = homogenize_points(world_space_points) + camera_space_points = torch.einsum( + "b u i j , b v h w j -> b v u h w i", c2w.inverse(), world_space_points + ) + return camera_space_points[..., :3] + + +def unproject_depth( + depth, intrinsics, c2w=None, ixt_normalized=False, num_patches_x=None, num_patches_y=None +): + """ + Turn the depth map into a 3D point cloud in world space + + Args: + depth: (b, v, h, w, 1) + intrinsics: (b, v, 3, 3) + c2w: (b, v, 4, 4) + + Returns: + torch.Tensor: World space points with shape (b, v, h, w, 3). + """ + if c2w is None: + c2w = torch.eye(4, device=depth.device, dtype=depth.dtype) + c2w = c2w[None, None].repeat(depth.shape[0], depth.shape[1], 1, 1) + + if not ixt_normalized: + # Compute indices of pixels + h, w = depth.shape[-3], depth.shape[-2] + x_grid, y_grid = torch.meshgrid( + torch.arange(w, device=depth.device, dtype=depth.dtype), + torch.arange(h, device=depth.device, dtype=depth.dtype), + indexing="xy", + ) # (h, w), (h, w) + else: + # ixt_normalized: h=w=2.0. cx, cy, fx, fy are normalized according to h=w=2.0 + assert num_patches_x is not None and num_patches_y is not None + dx = 1 / num_patches_x + dy = 1 / num_patches_y + max_y = 1 - dy + min_y = -max_y + max_x = 1 - dx + min_x = -max_x + + grid_shift = 1.0 + y_grid, x_grid = torch.meshgrid( + torch.linspace( + min_y + grid_shift, + max_y + grid_shift, + num_patches_y, + dtype=torch.float32, + device=depth.device, + ), + torch.linspace( + min_x + grid_shift, + max_x + grid_shift, + num_patches_x, + dtype=torch.float32, + device=depth.device, + ), + indexing="ij", + ) + + # Compute coordinates of pixels in camera space + pixel_space_points = torch.stack((x_grid, y_grid), dim=-1) # (..., h, w, 2) + camera_points = pixel_space_to_camera_space( + pixel_space_points, depth, intrinsics + ) # (..., h, w, 3) + + # Convert points to world space + world_points = camera_space_to_world_space(camera_points, c2w) # (..., h, w, 3) + + return world_points \ No newline at end of file diff --git a/src/nn/depth_anything_3/utils/gsply_helpers.py b/src/nn/depth_anything_3/utils/gsply_helpers.py new file mode 100644 index 0000000..5733009 --- /dev/null +++ b/src/nn/depth_anything_3/utils/gsply_helpers.py @@ -0,0 +1,173 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from pathlib import Path +from typing import Optional +import numpy as np +import torch +from einops import rearrange, repeat +from plyfile import PlyData, PlyElement +from torch import Tensor + +from depth_anything_3.specs import Gaussians + + +def construct_list_of_attributes(num_rest: int) -> list[str]: + attributes = ["x", "y", "z", "nx", "ny", "nz"] + for i in range(3): + attributes.append(f"f_dc_{i}") + for i in range(num_rest): + attributes.append(f"f_rest_{i}") + attributes.append("opacity") + for i in range(3): + attributes.append(f"scale_{i}") + for i in range(4): + attributes.append(f"rot_{i}") + return attributes + + +def export_ply( + means: Tensor, # "gaussian 3" + scales: Tensor, # "gaussian 3" + rotations: Tensor, # "gaussian 4" + harmonics: Tensor, # "gaussian 3 d_sh" + opacities: Tensor, # "gaussian" + path: Path, + shift_and_scale: bool = False, + save_sh_dc_only: bool = True, + match_3dgs_mcmc_dev: Optional[bool] = False, +): + if shift_and_scale: + # Shift the scene so that the median Gaussian is at the origin. + means = means - means.median(dim=0).values + + # Rescale the scene so that most Gaussians are within range [-1, 1]. + scale_factor = means.abs().quantile(0.95, dim=0).max() + means = means / scale_factor + scales = scales / scale_factor + + rotations = rotations.detach().cpu().numpy() + + # Since current model use SH_degree = 4, + # which require large memory to store, we can only save the DC band to save memory. + f_dc = harmonics[..., 0] + f_rest = harmonics[..., 1:].flatten(start_dim=1) + + if match_3dgs_mcmc_dev: + sh_degree = 3 + n_rest = 3 * (sh_degree + 1) ** 2 - 3 + f_rest = repeat( + torch.zeros_like(harmonics[..., :1]), "... i -> ... (n i)", n=(n_rest // 3) + ).flatten(start_dim=1) + dtype_full = [ + (attribute, "f4") + for attribute in construct_list_of_attributes(num_rest=n_rest) + if attribute not in ("nx", "ny", "nz") + ] + else: + dtype_full = [ + (attribute, "f4") + for attribute in construct_list_of_attributes( + 0 if save_sh_dc_only else f_rest.shape[1] + ) + ] + elements = np.empty(means.shape[0], dtype=dtype_full) + attributes = [ + means.detach().cpu().numpy(), + torch.zeros_like(means).detach().cpu().numpy(), + f_dc.detach().cpu().contiguous().numpy(), + f_rest.detach().cpu().contiguous().numpy(), + opacities[..., None].detach().cpu().numpy(), + scales.log().detach().cpu().numpy(), + rotations, + ] + if match_3dgs_mcmc_dev: + attributes.pop(1) # dummy normal is not needed + elif save_sh_dc_only: + attributes.pop(3) # remove f_rest from attributes + + attributes = np.concatenate(attributes, axis=1) + elements[:] = list(map(tuple, attributes)) + path.parent.mkdir(exist_ok=True, parents=True) + PlyData([PlyElement.describe(elements, "vertex")]).write(path) + + +def inverse_sigmoid(x): + return torch.log(x / (1 - x)) + + +def save_gaussian_ply( + gaussians: Gaussians, + save_path: str, + ctx_depth: torch.Tensor, # depth of input views; for getting shape and filtering, "v h w 1" + shift_and_scale: bool = False, + save_sh_dc_only: bool = True, + gs_views_interval: int = 1, + inv_opacity: Optional[bool] = True, + prune_by_depth_percent: Optional[float] = 1.0, + prune_border_gs: Optional[bool] = True, + match_3dgs_mcmc_dev: Optional[bool] = False, +): + b = gaussians.means.shape[0] + assert b == 1, "must set batch_size=1 when exporting 3D gaussians" + src_v, out_h, out_w, _ = ctx_depth.shape + + # extract gs params + world_means = gaussians.means + world_shs = gaussians.harmonics + world_rotations = gaussians.rotations + gs_scales = gaussians.scales + gs_opacities = inverse_sigmoid(gaussians.opacities) if inv_opacity else gaussians.opacities + + # Create a mask to filter the Gaussians. + + # TODO: prune the sky region here + + # throw away Gaussians at the borders, since they're generally of lower quality. + if prune_border_gs: + mask = torch.zeros_like(ctx_depth, dtype=torch.bool) + gstrim_h = int(8 / 256 * out_h) + gstrim_w = int(8 / 256 * out_w) + mask[:, gstrim_h:-gstrim_h, gstrim_w:-gstrim_w, :] = 1 + else: + mask = torch.ones_like(ctx_depth, dtype=torch.bool) + + # trim the far away point based on depth; + if prune_by_depth_percent is not None and prune_by_depth_percent < 1: + in_depths = ctx_depth + d_percentile = torch.quantile( + in_depths.view(in_depths.shape[0], -1), q=prune_by_depth_percent, dim=1 + ).view(-1, 1, 1) + d_mask = (in_depths[..., 0] <= d_percentile).unsqueeze(-1) + mask = mask & d_mask + mask = mask.squeeze(-1) # v h w + + # helper fn, must place after mask + def trim_select_reshape(element): + selected_element = rearrange( + element[0], "(v h w) ... -> v h w ...", v=src_v, h=out_h, w=out_w + ) + selected_element = selected_element[::gs_views_interval][mask[::gs_views_interval]] + return selected_element + + export_ply( + means=trim_select_reshape(world_means), + scales=trim_select_reshape(gs_scales), + rotations=trim_select_reshape(world_rotations), + harmonics=trim_select_reshape(world_shs), + opacities=trim_select_reshape(gs_opacities), + path=Path(save_path), + shift_and_scale=shift_and_scale, + save_sh_dc_only=save_sh_dc_only, + match_3dgs_mcmc_dev=match_3dgs_mcmc_dev, + ) diff --git a/src/nn/depth_anything_3/utils/io/input_processor.py b/src/nn/depth_anything_3/utils/io/input_processor.py new file mode 100644 index 0000000..fa60194 --- /dev/null +++ b/src/nn/depth_anything_3/utils/io/input_processor.py @@ -0,0 +1,501 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Input processor for Depth Anything 3 (parallelized). + +This version removes the square center-crop step for "*crop" methods (same as your note). +In addition, it parallelizes per-image preprocessing using the provided `parallel_execution`. +""" + +from __future__ import annotations + +from typing import Sequence +import cv2 +import numpy as np +import torch +import torchvision.transforms as T +from PIL import Image + +from depth_anything_3.utils.logger import logger +from depth_anything_3.utils.parallel_utils import parallel_execution + + +class InputProcessor: + """Prepares a batch of images for model inference. + This processor converts a list of image file paths into a single, model-ready + tensor. The processing pipeline is executed in parallel across multiple workers + for efficiency. + + Pipeline: + 1) Load image and convert to RGB + 2) Boundary resize (upper/lower bound, preserving aspect ratio) + 3) Enforce divisibility by PATCH_SIZE: + - "*resize" methods: each dimension is rounded to nearest multiple + (may up/downscale a few px) + - "*crop" methods: each dimension is floored to nearest multiple via center crop + 4) Convert to tensor and apply ImageNet normalization + 5) Stack into (1, N, 3, H, W) + + Parallelization: + - Each image is processed independently in a worker. + - Order of outputs matches the input order. + """ + + NORMALIZE = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + PATCH_SIZE = 14 + + def __init__(self): + pass + + # ----------------------------- + # Public API + # ----------------------------- + def __call__( + self, + image: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + process_res: int = 504, + process_res_method: str = "upper_bound_resize", + *, + num_workers: int = 8, + print_progress: bool = False, + sequential: bool | None = None, + desc: str | None = "Preprocess", + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Returns: + (tensor, extrinsics_list, intrinsics_list) + tensor shape: (1, N, 3, H, W) + """ + sequential = self._resolve_sequential(sequential, num_workers) + exts_list, ixts_list = self._validate_and_pack_meta(image, extrinsics, intrinsics) + + results = self._run_parallel( + image=image, + exts_list=exts_list, + ixts_list=ixts_list, + process_res=process_res, + process_res_method=process_res_method, + num_workers=num_workers, + print_progress=print_progress, + sequential=sequential, + desc=desc, + ) + + proc_imgs, out_sizes, out_ixts, out_exts = self._unpack_results(results) + proc_imgs, out_sizes, out_ixts = self._unify_batch_shapes(proc_imgs, out_sizes, out_ixts) + + batch_tensor = self._stack_batch(proc_imgs) + out_exts = ( + torch.from_numpy(np.asarray(out_exts)).float() + if out_exts is not None and out_exts[0] is not None + else None + ) + out_ixts = ( + torch.from_numpy(np.asarray(out_ixts)).float() + if out_ixts is not None and out_ixts[0] is not None + else None + ) + return (batch_tensor, out_exts, out_ixts) + + # ----------------------------- + # __call__ helpers + # ----------------------------- + def _resolve_sequential(self, sequential: bool | None, num_workers: int) -> bool: + return (num_workers <= 1) if sequential is None else sequential + + def _validate_and_pack_meta( + self, + images: list[np.ndarray | Image.Image | str], + extrinsics: np.ndarray | None, + intrinsics: np.ndarray | None, + ) -> tuple[list[np.ndarray | None] | None, list[np.ndarray | None] | None]: + if extrinsics is not None and len(extrinsics) != len(images): + raise ValueError("Length of extrinsics must match images when provided.") + if intrinsics is not None and len(intrinsics) != len(images): + raise ValueError("Length of intrinsics must match images when provided.") + exts_list = [e for e in extrinsics] if extrinsics is not None else None + ixts_list = [k for k in intrinsics] if intrinsics is not None else None + return exts_list, ixts_list + + def _run_parallel( + self, + *, + image: list[np.ndarray | Image.Image | str], + exts_list: list[np.ndarray | None] | None, + ixts_list: list[np.ndarray | None] | None, + process_res: int, + process_res_method: str, + num_workers: int, + print_progress: bool, + sequential: bool, + desc: str | None, + ): + results = parallel_execution( + image, + exts_list, + ixts_list, + action=self._process_one, # (img, extrinsic, intrinsic, ...) + num_processes=num_workers, + print_progress=print_progress, + sequential=sequential, + desc=desc, + process_res=process_res, + process_res_method=process_res_method, + ) + if not results: + raise RuntimeError( + "No preprocessing results returned. Check inputs and parallel_execution." + ) + return results + + def _unpack_results(self, results): + """ + results: List[Tuple[torch.Tensor, Tuple[H, W], Optional[np.ndarray], Optional[np.ndarray]]] + -> processed_images, out_sizes, out_intrinsics, out_extrinsics + """ + try: + processed_images, out_sizes, out_intrinsics, out_extrinsics = zip(*results) + except Exception as e: + raise RuntimeError( + "Unexpected results structure from parallel_execution: " + f"{type(results)} / sample: {results[0]}" + ) from e + + return list(processed_images), list(out_sizes), list(out_intrinsics), list(out_extrinsics) + + def _unify_batch_shapes( + self, + processed_images: list[torch.Tensor], + out_sizes: list[tuple[int, int]], + out_intrinsics: list[np.ndarray | None], + ) -> tuple[list[torch.Tensor], list[tuple[int, int]], list[np.ndarray | None]]: + """Center-crop all tensors to the smallest H, W; adjust intrinsics' cx, cy accordingly.""" + if len(set(out_sizes)) <= 1: + return processed_images, out_sizes, out_intrinsics + + min_h = min(h for h, _ in out_sizes) + min_w = min(w for _, w in out_sizes) + logger.warn( + f"Images in batch have different sizes {out_sizes}; " + f"center-cropping all to smallest ({min_h},{min_w})" + ) + + center_crop = T.CenterCrop((min_h, min_w)) + new_imgs, new_sizes, new_ixts = [], [], [] + for img_t, (H, W), K in zip(processed_images, out_sizes, out_intrinsics): + crop_top = max(0, (H - min_h) // 2) + crop_left = max(0, (W - min_w) // 2) + new_imgs.append(center_crop(img_t)) + new_sizes.append((min_h, min_w)) + if K is None: + new_ixts.append(None) + else: + K_adj = K.copy() + K_adj[0, 2] -= crop_left + K_adj[1, 2] -= crop_top + new_ixts.append(K_adj) + return new_imgs, new_sizes, new_ixts + + def _stack_batch(self, processed_images: list[torch.Tensor]) -> torch.Tensor: + return torch.stack(processed_images) + + # ----------------------------- + # Per-item worker + # ----------------------------- + def _process_one( + self, + img: np.ndarray | Image.Image | str, + extrinsic: np.ndarray | None = None, + intrinsic: np.ndarray | None = None, + *, + process_res: int, + process_res_method: str, + ) -> tuple[torch.Tensor, tuple[int, int], np.ndarray | None, np.ndarray | None]: + # Load & remember original size + pil_img = self._load_image(img) + orig_w, orig_h = pil_img.size + + # Boundary resize + pil_img = self._resize_image(pil_img, process_res, process_res_method) + w, h = pil_img.size + intrinsic = self._resize_ixt(intrinsic, orig_w, orig_h, w, h) + + # Enforce divisibility by PATCH_SIZE + if process_res_method.endswith("resize"): + pil_img = self._make_divisible_by_resize(pil_img, self.PATCH_SIZE) + new_w, new_h = pil_img.size + intrinsic = self._resize_ixt(intrinsic, w, h, new_w, new_h) + w, h = new_w, new_h + elif process_res_method.endswith("crop"): + pil_img = self._make_divisible_by_crop(pil_img, self.PATCH_SIZE) + new_w, new_h = pil_img.size + intrinsic = self._crop_ixt(intrinsic, w, h, new_w, new_h) + w, h = new_w, new_h + else: + raise ValueError(f"Unsupported process_res_method: {process_res_method}") + + # Convert to tensor & normalize + img_tensor = self._normalize_image(pil_img) + _, H, W = img_tensor.shape + assert (W, H) == (w, h), "Tensor size mismatch with PIL image size after processing." + + # Return: (img_tensor, (H, W), intrinsic, extrinsic) + return img_tensor, (H, W), intrinsic, extrinsic + + # ----------------------------- + # Intrinsics transforms + # ----------------------------- + def _resize_ixt( + self, + intrinsic: np.ndarray | None, + orig_w: int, + orig_h: int, + w: int, + h: int, + ) -> np.ndarray | None: + if intrinsic is None: + return None + K = intrinsic.copy() + # scale fx, cx by w ratio; fy, cy by h ratio + K[:1] *= w / float(orig_w) + K[1:2] *= h / float(orig_h) + return K + + def _crop_ixt( + self, + intrinsic: np.ndarray | None, + orig_w: int, + orig_h: int, + w: int, + h: int, + ) -> np.ndarray | None: + if intrinsic is None: + return None + K = intrinsic.copy() + crop_h = (orig_h - h) // 2 + crop_w = (orig_w - w) // 2 + K[0, 2] -= crop_w + K[1, 2] -= crop_h + return K + + # ----------------------------- + # I/O & normalization + # ----------------------------- + def _load_image(self, img: np.ndarray | Image.Image | str) -> Image.Image: + if isinstance(img, str): + return Image.open(img).convert("RGB") + elif isinstance(img, np.ndarray): + # Assume HxWxC uint8/RGB + return Image.fromarray(img).convert("RGB") + elif isinstance(img, Image.Image): + return img.convert("RGB") + else: + raise ValueError(f"Unsupported image type: {type(img)}") + + def _normalize_image(self, img: Image.Image) -> torch.Tensor: + img_tensor = T.ToTensor()(img) + return self.NORMALIZE(img_tensor) + + # ----------------------------- + # Boundary resizing + # ----------------------------- + def _resize_image(self, img: Image.Image, target_size: int, method: str) -> Image.Image: + if method in ("upper_bound_resize", "upper_bound_crop"): + return self._resize_longest_side(img, target_size) + elif method in ("lower_bound_resize", "lower_bound_crop"): + return self._resize_shortest_side(img, target_size) + else: + raise ValueError(f"Unsupported resize method: {method}") + + def _resize_longest_side(self, img: Image.Image, target_size: int) -> Image.Image: + w, h = img.size + longest = max(w, h) + if longest == target_size: + return img + scale = target_size / float(longest) + new_w = max(1, int(round(w * scale))) + new_h = max(1, int(round(h * scale))) + interpolation = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA + arr = cv2.resize(np.asarray(img), (new_w, new_h), interpolation=interpolation) + return Image.fromarray(arr) + + def _resize_shortest_side(self, img: Image.Image, target_size: int) -> Image.Image: + w, h = img.size + shortest = min(w, h) + if shortest == target_size: + return img + scale = target_size / float(shortest) + new_w = max(1, int(round(w * scale))) + new_h = max(1, int(round(h * scale))) + interpolation = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA + arr = cv2.resize(np.asarray(img), (new_w, new_h), interpolation=interpolation) + return Image.fromarray(arr) + + # ----------------------------- + # Make divisible by PATCH_SIZE + # ----------------------------- + def _make_divisible_by_crop(self, img: Image.Image, patch: int) -> Image.Image: + """ + Floor each dimension to the nearest multiple of PATCH_SIZE via center crop. + Example: 504x377 -> 504x364 + """ + w, h = img.size + new_w = (w // patch) * patch + new_h = (h // patch) * patch + if new_w == w and new_h == h: + return img + left = (w - new_w) // 2 + top = (h - new_h) // 2 + return img.crop((left, top, left + new_w, top + new_h)) + + def _make_divisible_by_resize(self, img: Image.Image, patch: int) -> Image.Image: + """ + Round each dimension to nearest multiple of PATCH_SIZE via small resize. + """ + w, h = img.size + + def nearest_multiple(x: int, p: int) -> int: + down = (x // p) * p + up = down + p + return up if abs(up - x) <= abs(x - down) else down + + new_w = max(1, nearest_multiple(w, patch)) + new_h = max(1, nearest_multiple(h, patch)) + if new_w == w and new_h == h: + return img + upscale = (new_w > w) or (new_h > h) + interpolation = cv2.INTER_CUBIC if upscale else cv2.INTER_AREA + arr = cv2.resize(np.asarray(img), (new_w, new_h), interpolation=interpolation) + return Image.fromarray(arr) + + +# Backward compatibility alias +InputAdapter = InputProcessor + + +# =========================== +# Minimal test runner (parallel execution) +# =========================== +if __name__ == "__main__": + """ + Minimal test suite: + - Creates pairs of images so batch shapes match. + - Tests all four process_res_methods. + - Prints fx fy cx cy IN->OUT per image. + - Includes cases with K/E provided and with None. + """ + + def fmt_k_line(K: np.ndarray | None) -> str: + if K is None: + return "None" + fx, fy, cx, cy = float(K[0, 0]), float(K[1, 1]), float(K[0, 2]), float(K[1, 2]) + return f"fx={fx:.3f} fy={fy:.3f} cx={cx:.3f} cy={cy:.3f}" + + def show_result( + tag: str, + tensor: torch.Tensor, + Ks_in: Sequence[np.ndarray | None] | None = None, + Ks_out: Sequence[np.ndarray | None] | None = None, + ): + B, N, C, H, W = tensor.shape + print(f"[{tag}] shape={tuple(tensor.shape)} HxW=({H},{W}) div14=({H%14==0},{W%14==0})") + assert H % 14 == 0 and W % 14 == 0, f"{tag}: output size not divisible by 14!" + if Ks_in is not None or Ks_out is not None: + Ks_in = Ks_in or [None] * N + Ks_out = Ks_out or [None] * N + for i in range(N): + print(f" K[{i}]: {fmt_k_line(Ks_in[i])} -> {fmt_k_line(Ks_out[i])}") + + proc = InputProcessor() + process_res = 504 + methods = ["upper_bound_resize", "upper_bound_crop", "lower_bound_resize", "lower_bound_crop"] + + # Example sizes (two orientations) + small_sizes = [(680, 1208), (1208, 680)] + large_sizes = [(1208, 680), (680, 1208)] + + def make_K(w, h, fx=1200.0, fy=1100.0): + cx, cy = w / 2.0, h / 2.0 + K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32) + return K + + def run_suite(suite_name: str, sizes: list[tuple[int, int]]): + print(f"\n===== {suite_name} =====") + for w, h in sizes: + img = Image.new("RGB", (w, h), color=(123, 222, 100)) + batch_imgs = [img, img] + + # intrinsics / extrinsics examples + Ks_in = [make_K(w, h), make_K(w, h)] + Es_in = [np.eye(4, dtype=np.float32), np.eye(4, dtype=np.float32)] + + for m in methods: + tensor, Es_out, Ks_out = proc( + image=batch_imgs, + process_res=process_res, + process_res_method=m, + num_workers=8, + print_progress=False, + intrinsics=Ks_in, # test with non-None + extrinsics=Es_in, + ) + show_result(f"{suite_name} size=({w},{h}) | {m}", tensor, Ks_in, Ks_out) + + # Also test None path + tensor2, Es_out2, Ks_out2 = proc( + image=batch_imgs, + process_res=process_res, + process_res_method="upper_bound_resize", + num_workers=8, + intrinsics=None, + extrinsics=None, + ) + show_result( + f"{suite_name} size=({w},{h}) | upper_bound_resize | no K/E", + tensor2, + None, + Ks_out2, + ) + + run_suite("SMALL", small_sizes) + run_suite("LARGE", large_sizes) + + # Extra sanity for 504x376 + print("\n===== EXTRA sanity for 504x376 =====") + img_example = Image.new("RGB", (504, 376), color=(10, 20, 30)) + Ks_in_extra = [make_K(504, 376, fx=900.0, fy=900.0), make_K(504, 376, fx=900.0, fy=900.0)] + + out_r, _, Ks_out_r = proc( + image=[img_example, img_example], + process_res=504, + process_res_method="upper_bound_resize", + num_workers=8, + intrinsics=Ks_in_extra, + ) + out_c, _, Ks_out_c = proc( + image=[img_example, img_example], + process_res=504, + process_res_method="upper_bound_crop", + num_workers=8, + intrinsics=Ks_in_extra, + ) + _, _, _, Hr, Wr = out_r.shape + _, _, _, Hc, Wc = out_c.shape + print(f"upper_bound_resize -> ({Hr},{Wr}) (rounded to nearest multiple of 14)") + show_result("Ks after upper_bound_resize", out_r, Ks_in_extra, Ks_out_r) + print(f"upper_bound_crop -> ({Hc},{Wc}) (floored to multiple of 14)") + show_result("Ks after upper_bound_crop", out_c, Ks_in_extra, Ks_out_c) diff --git a/src/nn/depth_anything_3/utils/io/output_processor.py b/src/nn/depth_anything_3/utils/io/output_processor.py new file mode 100644 index 0000000..c317eb9 --- /dev/null +++ b/src/nn/depth_anything_3/utils/io/output_processor.py @@ -0,0 +1,172 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Output processor for Depth Anything 3. + +This module handles model output processing, including tensor-to-numpy conversion, +batch dimension removal, and Prediction object creation. +""" + +from __future__ import annotations + +import numpy as np +import torch +from addict import Dict as AddictDict + +from depth_anything_3.specs import Prediction + + +class OutputProcessor: + """ + Output processor for converting model outputs to Prediction objects. + + Handles tensor-to-numpy conversion, batch dimension removal, + and creates structured Prediction objects with proper data types. + """ + + def __init__(self) -> None: + """Initialize the output processor.""" + + def __call__(self, model_output: dict[str, torch.Tensor]) -> Prediction: + """ + Convert model output to Prediction object. + + Args: + model_output: Model output dictionary containing depth, conf, extrinsics, intrinsics + Expected shapes: depth (B, N, 1, H, W), conf (B, N, 1, H, W), + extrinsics (B, N, 4, 4), intrinsics (B, N, 3, 3) + + Returns: + Prediction: Object containing depth estimation results with shapes: + depth (N, H, W), conf (N, H, W), extrinsics (N, 4, 4), intrinsics (N, 3, 3) + """ + # Extract data from batch dimension (B=1, N=number of images) + depth = self._extract_depth(model_output) + conf = self._extract_conf(model_output) + extrinsics = self._extract_extrinsics(model_output) + intrinsics = self._extract_intrinsics(model_output) + sky = self._extract_sky(model_output) + aux = self._extract_aux(model_output) + gaussians = model_output.get("gaussians", None) + scale_factor = model_output.get("scale_factor", None) + + return Prediction( + depth=depth, + sky=sky, + conf=conf, + extrinsics=extrinsics, + intrinsics=intrinsics, + is_metric=getattr(model_output, "is_metric", 0), + gaussians=gaussians, + aux=aux, + scale_factor=scale_factor, + ) + + def _extract_depth(self, model_output: dict[str, torch.Tensor]) -> np.ndarray: + """ + Extract depth tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Depth array with shape (N, H, W) + """ + depth = model_output["depth"].squeeze(0).squeeze(-1).cpu().numpy() # (N, H, W) + return depth + + def _extract_conf(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract confidence tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Confidence array with shape (N, H, W) or None + """ + conf = model_output.get("depth_conf", None) + if conf is not None: + conf = conf.squeeze(0).cpu().numpy() # (N, H, W) + return conf + + def _extract_extrinsics(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract extrinsics tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Extrinsics array with shape (N, 4, 4) or None + """ + extrinsics = model_output.get("extrinsics", None) + if extrinsics is not None: + extrinsics = extrinsics.squeeze(0).cpu().numpy() # (N, 4, 4) + return extrinsics + + def _extract_intrinsics(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract intrinsics tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Intrinsics array with shape (N, 3, 3) or None + """ + intrinsics = model_output.get("intrinsics", None) + if intrinsics is not None: + intrinsics = intrinsics.squeeze(0).cpu().numpy() # (N, 3, 3) + return intrinsics + + def _extract_sky(self, model_output: dict[str, torch.Tensor]) -> np.ndarray | None: + """ + Extract sky tensor from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Sky mask array with shape (N, H, W) or None + """ + sky = model_output.get("sky", None) + if sky is not None: + sky = sky.squeeze(0).cpu().numpy() >= 0.5 # (N, H, W) + return sky + + def _extract_aux(self, model_output: dict[str, torch.Tensor]) -> AddictDict: + """ + Extract auxiliary data from model output and convert to numpy. + + Args: + model_output: Model output dictionary + + Returns: + Dictionary containing auxiliary data + """ + aux = model_output.get("aux", None) + ret = AddictDict() + if aux is not None: + for k in aux.keys(): + if isinstance(aux[k], torch.Tensor): + ret[k] = aux[k].squeeze(0).cpu().numpy() + else: + ret[k] = aux[k] + return ret + + +# Backward compatibility alias +OutputAdapter = OutputProcessor diff --git a/src/nn/depth_anything_3/utils/layout_helpers.py b/src/nn/depth_anything_3/utils/layout_helpers.py new file mode 100644 index 0000000..189c170 --- /dev/null +++ b/src/nn/depth_anything_3/utils/layout_helpers.py @@ -0,0 +1,216 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""This file contains useful layout utilities for images. They are: + +- add_border: Add a border to an image. +- cat/hcat/vcat: Join images by arranging them in a line. If the images have different + sizes, they are aligned as specified (start, end, center). Allows you to specify a gap + between images. + +Images are assumed to be float32 tensors with shape (channel, height, width). +""" + +from typing import Any, Generator, Iterable, Literal, Union +import torch +from torch import Tensor + +Alignment = Literal["start", "center", "end"] +Axis = Literal["horizontal", "vertical"] +Color = Union[ + int, + float, + Iterable[int], + Iterable[float], + Tensor, + Tensor, +] + + +def _sanitize_color(color: Color) -> Tensor: # "#channel" + # Convert tensor to list (or individual item). + if isinstance(color, torch.Tensor): + color = color.tolist() + + # Turn iterators and individual items into lists. + if isinstance(color, Iterable): + color = list(color) + else: + color = [color] + + return torch.tensor(color, dtype=torch.float32) + + +def _intersperse(iterable: Iterable, delimiter: Any) -> Generator[Any, None, None]: + it = iter(iterable) + yield next(it) + for item in it: + yield delimiter + yield item + + +def _get_main_dim(main_axis: Axis) -> int: + return { + "horizontal": 2, + "vertical": 1, + }[main_axis] + + +def _get_cross_dim(main_axis: Axis) -> int: + return { + "horizontal": 1, + "vertical": 2, + }[main_axis] + + +def _compute_offset(base: int, overlay: int, align: Alignment) -> slice: + assert base >= overlay + offset = { + "start": 0, + "center": (base - overlay) // 2, + "end": base - overlay, + }[align] + return slice(offset, offset + overlay) + + +def overlay( + base: Tensor, # "channel base_height base_width" + overlay: Tensor, # "channel overlay_height overlay_width" + main_axis: Axis, + main_axis_alignment: Alignment, + cross_axis_alignment: Alignment, +) -> Tensor: # "channel base_height base_width" + # The overlay must be smaller than the base. + _, base_height, base_width = base.shape + _, overlay_height, overlay_width = overlay.shape + assert base_height >= overlay_height and base_width >= overlay_width + + # Compute spacing on the main dimension. + main_dim = _get_main_dim(main_axis) + main_slice = _compute_offset( + base.shape[main_dim], overlay.shape[main_dim], main_axis_alignment + ) + + # Compute spacing on the cross dimension. + cross_dim = _get_cross_dim(main_axis) + cross_slice = _compute_offset( + base.shape[cross_dim], overlay.shape[cross_dim], cross_axis_alignment + ) + + # Combine the slices and paste the overlay onto the base accordingly. + selector = [..., None, None] + selector[main_dim] = main_slice + selector[cross_dim] = cross_slice + result = base.clone() + result[selector] = overlay + return result + + +def cat( + main_axis: Axis, + *images: Iterable[Tensor], # "channel _ _" + align: Alignment = "center", + gap: int = 8, + gap_color: Color = 1, +) -> Tensor: # "channel height width" + """Arrange images in a line. The interface resembles a CSS div with flexbox.""" + device = images[0].device + gap_color = _sanitize_color(gap_color).to(device) + + # Find the maximum image side length in the cross axis dimension. + cross_dim = _get_cross_dim(main_axis) + cross_axis_length = max(image.shape[cross_dim] for image in images) + + # Pad the images. + padded_images = [] + for image in images: + # Create an empty image with the correct size. + padded_shape = list(image.shape) + padded_shape[cross_dim] = cross_axis_length + base = torch.ones(padded_shape, dtype=torch.float32, device=device) + base = base * gap_color[:, None, None] + padded_images.append(overlay(base, image, main_axis, "start", align)) + + # Intersperse separators if necessary. + if gap > 0: + # Generate a separator. + c, _, _ = images[0].shape + separator_size = [gap, gap] + separator_size[cross_dim - 1] = cross_axis_length + separator = torch.ones((c, *separator_size), dtype=torch.float32, device=device) + separator = separator * gap_color[:, None, None] + + # Intersperse the separator between the images. + padded_images = list(_intersperse(padded_images, separator)) + + return torch.cat(padded_images, dim=_get_main_dim(main_axis)) + + +def hcat( + *images: Iterable[Tensor], # "channel _ _" + align: Literal["start", "center", "end", "top", "bottom"] = "start", + gap: int = 8, + gap_color: Color = 1, +): + """Shorthand for a horizontal linear concatenation.""" + return cat( + "horizontal", + *images, + align={ + "start": "start", + "center": "center", + "end": "end", + "top": "start", + "bottom": "end", + }[align], + gap=gap, + gap_color=gap_color, + ) + + +def vcat( + *images: Iterable[Tensor], # "channel _ _" + align: Literal["start", "center", "end", "left", "right"] = "start", + gap: int = 8, + gap_color: Color = 1, +): + """Shorthand for a horizontal linear concatenation.""" + return cat( + "vertical", + *images, + align={ + "start": "start", + "center": "center", + "end": "end", + "left": "start", + "right": "end", + }[align], + gap=gap, + gap_color=gap_color, + ) + + +def add_border( + image: Tensor, # "channel height width" + border: int = 8, + color: Color = 1, +) -> Tensor: # "channel new_height new_width" + color = _sanitize_color(color).to(image) + c, h, w = image.shape + result = torch.empty( + (c, h + 2 * border, w + 2 * border), dtype=torch.float32, device=image.device + ) + result[:] = color[:, None, None] + result[:, border : h + border, border : w + border] = image + return result diff --git a/src/nn/depth_anything_3/utils/logger.py b/src/nn/depth_anything_3/utils/logger.py new file mode 100644 index 0000000..0eb4f60 --- /dev/null +++ b/src/nn/depth_anything_3/utils/logger.py @@ -0,0 +1,82 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys + + +class Color: + RED = "\033[91m" + YELLOW = "\033[93m" + WHITE = "\033[97m" + GREEN = "\033[92m" + RESET = "\033[0m" + + +LOG_LEVELS = {"ERROR": 0, "WARN": 1, "INFO": 2, "DEBUG": 3} + +COLOR_MAP = {"ERROR": Color.RED, "WARN": Color.YELLOW, "INFO": Color.WHITE, "DEBUG": Color.GREEN} + + +def get_env_log_level(): + level = os.environ.get("DA3_LOG_LEVEL", "INFO").upper() + return LOG_LEVELS.get(level, LOG_LEVELS["INFO"]) + + +class Logger: + def __init__(self): + self.level = get_env_log_level() + + def log(self, level_str, *args, **kwargs): + level_key = level_str.split(":")[0].strip() + level_val = LOG_LEVELS.get(level_key) + if level_val is None: + raise ValueError(f"Unknown log level: {level_str}") + if self.level >= level_val: + color = COLOR_MAP[level_key] + msg = " ".join(str(arg) for arg in args) + + # Align log level output in square brackets + # ERROR and DEBUG are 5 characters, INFO and WARN have an extra space for alignment + tag = level_key + if tag in ("INFO", "WARN"): + tag += " " + print( + f"{color}[{tag}] {msg}{Color.RESET}", + file=sys.stderr if level_key == "ERROR" else sys.stdout, + **kwargs, + ) + + def error(self, *args, **kwargs): + self.log("ERROR:", *args, **kwargs) + + def warn(self, *args, **kwargs): + self.log("WARN:", *args, **kwargs) + + def info(self, *args, **kwargs): + self.log("INFO:", *args, **kwargs) + + def debug(self, *args, **kwargs): + self.log("DEBUG:", *args, **kwargs) + + +logger = Logger() + +__all__ = ["logger"] + +if __name__ == "__main__": + logger.info("This is an info message") + logger.warn("This is a warning message") + logger.error("This is an error message") + logger.debug("This is a debug message") diff --git a/src/nn/depth_anything_3/utils/memory.py b/src/nn/depth_anything_3/utils/memory.py new file mode 100644 index 0000000..682dad7 --- /dev/null +++ b/src/nn/depth_anything_3/utils/memory.py @@ -0,0 +1,127 @@ +""" +GPU memory utility helpers. + +Shared cleanup and memory checking logic used by both the backend API and +the Gradio UI to keep memory-management behavior consistent. +""" +from __future__ import annotations + +import gc + +from typing import Any, Dict, Optional + +import torch + + +def get_gpu_memory_info() -> Optional[Dict[str, Any]]: + """Return a snapshot of current GPU memory usage or None if CUDA not available. + + Keys in returned dict: total_gb, allocated_gb, reserved_gb, free_gb, utilization + """ + if not torch.cuda.is_available(): + return None + + try: + device = torch.cuda.current_device() + total_memory = torch.cuda.get_device_properties(device).total_memory + allocated_memory = torch.cuda.memory_allocated(device) + reserved_memory = torch.cuda.memory_reserved(device) + free_memory = total_memory - reserved_memory + + return { + "total_gb": total_memory / 1024 ** 3, + "allocated_gb": allocated_memory / 1024 ** 3, + "reserved_gb": reserved_memory / 1024 ** 3, + "free_gb": free_memory / 1024 ** 3, + "utilization": (reserved_memory / total_memory) * 100, + } + except Exception: + return None + + +def cleanup_cuda_memory() -> None: + """Perform a robust GPU cleanup sequence. + + This includes synchronizing, emptying caches, collecting IPC handles and + running the Python garbage collector. Use this instead of a raw + ``torch.cuda.empty_cache()`` where you need reliable freeing of GPU memory + between model loads or in error handling paths. + """ + try: + if torch.cuda.is_available(): + mem_before = get_gpu_memory_info() + + torch.cuda.synchronize() + torch.cuda.empty_cache() + # Collect cross-process cuda resources + try: + torch.cuda.ipc_collect() + except Exception: + # Older PyTorch versions or non-cuda devices may not support + # ipc_collect (no-op if not available) + pass + gc.collect() + + mem_after = get_gpu_memory_info() + if mem_before and mem_after: + freed = mem_before["reserved_gb"] - mem_after["reserved_gb"] + print( + f"CUDA cleanup: freed {freed:.2f}GB, " + f"available: {mem_after['free_gb']:.2f}GB/{mem_after['total_gb']:.2f}GB" + ) + else: + print("CUDA memory cleanup completed") + except Exception as e: + print(f"Warning: CUDA cleanup failed: {e}") + + +def check_memory_availability(required_gb: float = 2.0) -> tuple[bool, str]: + """Return whether at least ``required_gb`` seems available on the current GPU. + + The returned tuple is (is_available, message) with a human-friendly message. + """ + try: + if not torch.cuda.is_available(): + return False, "CUDA is not available" + + mem_info = get_gpu_memory_info() + if mem_info is None: + return True, "Cannot check memory, proceeding anyway" + + if mem_info["free_gb"] < required_gb: + return ( + False, + ( + f"Insufficient GPU memory: {mem_info['free_gb']:.2f}GB available, " + f"{required_gb:.2f}GB required. Total: {mem_info['total_gb']:.2f}GB, " + f"Used: {mem_info['reserved_gb']:.2f}GB ({mem_info['utilization']:.1f}%)" + ), + ) + + return ( + True, + ( + f"Memory check passed: {mem_info['free_gb']:.2f}GB available, " + f"{required_gb:.2f}GB required" + ), + ) + except Exception as e: + return True, f"Memory check failed: {e}, proceeding anyway" +def estimate_memory_requirement(num_images: int, process_res: int) -> float: + """Heuristic estimate for memory usage (GB) based on image count and resolution. + + This mirrors the simple policy used by the backend service so other code + (e.g., Gradio UI) can make consistent decisions when checking available + memory before loading a model or running inference. + + Args: + num_images: Number of images to process. + process_res: Processing resolution. + + Returns: + Estimated memory requirement in GB. + """ + base_memory = 2.0 + per_image_memory = (process_res / 504) ** 2 * 0.5 + total_memory = base_memory + (num_images * per_image_memory * 0.1) + return total_memory diff --git a/src/nn/depth_anything_3/utils/model_loading.py b/src/nn/depth_anything_3/utils/model_loading.py new file mode 100644 index 0000000..b6d43b5 --- /dev/null +++ b/src/nn/depth_anything_3/utils/model_loading.py @@ -0,0 +1,149 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Model loading and state dict conversion utilities. +""" + +from typing import Dict, Tuple +import torch + +from depth_anything_3.utils.logger import logger + + +def convert_general_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Convert general model state dict to match current model architecture. + + Args: + state_dict: Original state dictionary + + Returns: + Converted state dictionary + """ + # Replace module prefixes + state_dict = {k.replace("module.", "model."): v for k, v in state_dict.items()} + state_dict = {k.replace(".net.", ".backbone."): v for k, v in state_dict.items()} + + # Remove camera token if present + if "model.backbone.pretrained.camera_token" in state_dict: + del state_dict["model.backbone.pretrained.camera_token"] + + # Replace camera token naming + state_dict = { + k.replace(".camera_token_extra", ".camera_token"): v for k, v in state_dict.items() + } + + # Replace head naming + state_dict = { + k.replace("model.all_heads.camera_cond_head", "model.cam_enc"): v + for k, v in state_dict.items() + } + state_dict = { + k.replace("model.all_heads.camera_head", "model.cam_dec"): v for k, v in state_dict.items() + } + state_dict = {k.replace(".more_mlps.", ".backbone."): v for k, v in state_dict.items()} + state_dict = {k.replace(".fc_rot.", ".fc_qvec."): v for k, v in state_dict.items()} + state_dict = { + k.replace("model.all_heads.head", "model.head"): v for k, v in state_dict.items() + } + + # Replace output naming + state_dict = { + k.replace("output_conv2_additional.sky_mask", "sky_output_conv2"): v + for k, v in state_dict.items() + } + state_dict = {k.replace("_ray.", "_aux."): v for k, v in state_dict.items()} + + # Update GS-DPT head naming and value + state_dict = {k.replace("gaussian_param_head.", "gs_head."): v for k, v in state_dict.items()} + + return state_dict + + +def convert_metric_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Convert metric model state dict to match current model architecture. + + Args: + state_dict: Original metric state dictionary + + Returns: + Converted state dictionary + """ + # Add module prefix for metric models + state_dict = {"module." + k: v for k, v in state_dict.items()} + return convert_general_state_dict(state_dict) + + +def load_pretrained_weights(model, model_path: str, is_metric: bool = False) -> Tuple[list, list]: + """ + Load pretrained weights for a single model. + + Args: + model: Model instance to load weights into + model_path: Path to the pretrained weights + is_metric: Whether this is a metric model + + Returns: + Tuple of (missed_keys, unexpected_keys) + """ + state_dict = torch.load(model_path, map_location="cpu") + + if is_metric: + state_dict = convert_metric_state_dict(state_dict) + else: + state_dict = convert_general_state_dict(state_dict) + + missed, unexpected = model.load_state_dict(state_dict, strict=False) + logger.info("Missed keys:", missed) + logger.info("Unexpected keys:", unexpected) + + return missed, unexpected + + +def load_pretrained_nested_weights( + model, main_model_path: str, metric_model_path: str +) -> Tuple[list, list]: + """ + Load pretrained weights for a nested model with both main and metric branches. + + Args: + model: Nested model instance + main_model_path: Path to main model weights + metric_model_path: Path to metric model weights + + Returns: + Tuple of (missed_keys, unexpected_keys) + """ + # Load main model weights + state_dict0 = torch.load(main_model_path, map_location="cpu") + state_dict0 = convert_general_state_dict(state_dict0) + state_dict0 = {k.replace("model.", "model.da3."): v for k, v in state_dict0.items()} + + # Load metric model weights + state_dict1 = torch.load(metric_model_path, map_location="cpu") + state_dict1 = convert_metric_state_dict(state_dict1) + state_dict1 = {k.replace("model.", "model.da3_metric."): v for k, v in state_dict1.items()} + + # Combine state dictionaries + combined_state_dict = state_dict0.copy() + combined_state_dict.update(state_dict1) + + missed, unexpected = model.load_state_dict(combined_state_dict, strict=False) + + print("Missed keys:", missed) + print("Unexpected keys:", unexpected) + + return missed, unexpected diff --git a/src/nn/depth_anything_3/utils/parallel_utils.py b/src/nn/depth_anything_3/utils/parallel_utils.py new file mode 100644 index 0000000..9ff108e --- /dev/null +++ b/src/nn/depth_anything_3/utils/parallel_utils.py @@ -0,0 +1,133 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os +from functools import wraps +from multiprocessing.pool import ThreadPool +from threading import Thread +from typing import Callable, Dict, List +import imageio +from tqdm import tqdm + + +def async_call_func(func): + @wraps(func) + async def wrapper(*args, **kwargs): + loop = asyncio.get_event_loop() + # Use run_in_executor to run the blocking function in a separate thread + return await loop.run_in_executor(None, func, *args, **kwargs) + + return wrapper + + +slice_func = lambda chunk_index, chunk_dim, chunk_size: [slice(None)] * chunk_dim + [ + slice(chunk_index, chunk_index + chunk_size) +] + + +def async_call(fn): + def wrapper(*args, **kwargs): + Thread(target=fn, args=args, kwargs=kwargs).start() + + return wrapper + + +def _save_image_impl(save_img, save_path): + """Common implementation for saving images synchronously or asynchronously""" + os.makedirs(os.path.dirname(save_path), exist_ok=True) + imageio.imwrite(save_path, save_img) + + +@async_call +def save_image_async(save_img, save_path): + """Save image asynchronously""" + _save_image_impl(save_img, save_path) + + +def save_image(save_img, save_path): + """Save image synchronously""" + _save_image_impl(save_img, save_path) + + +def parallel_execution( + *args, + action: Callable, + num_processes=32, + print_progress=False, + sequential=False, + async_return=False, + desc=None, + **kwargs, +): + # Partially copy from EasyVolumetricVideo (parallel_execution) + # NOTE: we expect first arg / or kwargs to be distributed + # NOTE: print_progress arg is reserved. + # `*args` packs all positional arguments passed to the function into a tuple + args = list(args) + + def get_length(args: List, kwargs: Dict): + for a in args: + if isinstance(a, list): + return len(a) + for v in kwargs.values(): + if isinstance(v, list): + return len(v) + raise NotImplementedError + + def get_action_args(length: int, args: List, kwargs: Dict, i: int): + action_args = [ + (arg[i] if isinstance(arg, list) and len(arg) == length else arg) for arg in args + ] + # TODO: Support all types of iterable + action_kwargs = { + key: ( + kwargs[key][i] + if isinstance(kwargs[key], list) and len(kwargs[key]) == length + else kwargs[key] + ) + for key in kwargs + } + return action_args, action_kwargs + + if not sequential: + # Create ThreadPool + pool = ThreadPool(processes=num_processes) + + # Spawn threads + results = [] + asyncs = [] + length = get_length(args, kwargs) + for i in range(length): + action_args, action_kwargs = get_action_args(length, args, kwargs, i) + async_result = pool.apply_async(action, action_args, action_kwargs) + asyncs.append(async_result) + + # Join threads and get return values + if not async_return: + for async_result in tqdm(asyncs, desc=desc, disable=not print_progress): + results.append(async_result.get()) # will sync the corresponding thread + pool.close() + pool.join() + return results + else: + return pool + else: + results = [] + length = get_length(args, kwargs) + for i in tqdm(range(length), desc=desc, disable=not print_progress): + action_args, action_kwargs = get_action_args(length, args, kwargs, i) + async_result = action(*action_args, **action_kwargs) + results.append(async_result) + return results diff --git a/src/nn/depth_anything_3/utils/pca_utils.py b/src/nn/depth_anything_3/utils/pca_utils.py new file mode 100644 index 0000000..2b9eee2 --- /dev/null +++ b/src/nn/depth_anything_3/utils/pca_utils.py @@ -0,0 +1,284 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +PCA utilities for feature visualization and dimensionality reduction (video-friendly). +- Support frame-by-frame: transform_frame / transform_video +- Support one-time global PCA fitting and reuse (mean, V3) for stable colors +- Support Procrustes alignment (solving principal component order/sign/rotation jumps) +- Support global fixed or temporal EMA for percentiles (time dimension only, no spatial) +""" + +import numpy as np +import torch + + +def pca_to_rgb_4d_bf16_percentile( + x_np: np.ndarray, + device=None, + q_oversample: int = 6, + clip_percent: float = 10.0, # Percentage to clip from top and bottom (0~49.9) + return_uint8: bool = False, + enable_autocast_bf16: bool = True, +): + """ + Reduce numpy array of shape (49, 27, 36, 3072) to 3D via PCA and visualize as (49, 27, 36, 3). + - PCA uses torch.pca_lowrank (randomized SVD), defaults to GPU. + - Uses CUDA bf16 autocast in computation (if available), + then per-channel percentile clipping and normalization. + - Default removes 5% outliers from top and bottom (adjustable via clip_percent) to + improve visualization contrast. + + Parameters + ---------- + x_np : np.ndarray + Shape must be (49, 27, 36, 3072). dtype recommended float32/float64. + device : str | None + Specify 'cuda' or 'cpu'. Auto-select if None (prefer cuda). + q_oversample : int + Oversampling q for pca_lowrank, must be >= 3. + Slightly larger than target dim (3) is more stable, default 6. + clip_percent : float + Percentage to clip from top and bottom (0~49.9), + e.g. 5.0 means clip lowest 5% and highest 5% per channel. + return_uint8 : bool + True returns uint8(0~255), otherwise returns float32(0~1). + enable_autocast_bf16 : bool + Enable bf16 autocast on CUDA. + + Returns + ------- + np.ndarray + Array of shape (49, 27, 36, 3), float32[0,1] or uint8[0,255]. + """ + assert ( + x_np.ndim == 4 + ) # and x_np.shape[-1] == 3072, f"expect (49,27,36,3072), got {x_np.shape}" + B1, B2, B3, D = x_np.shape + N = B1 * B2 * B3 + + # Device selection + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Convert input to torch, unified float32 + X = torch.from_numpy(x_np.reshape(N, D)).to(device=device, dtype=torch.float32) + + # Parameter and safety checks + k = 3 + q = max(int(q_oversample), k) + clip_percent = float(clip_percent) + if not (0.0 <= clip_percent < 50.0): + raise ValueError( + "clip_percent must be in [0, 50), e.g. 5.0 means clip 5% from top and bottom" + ) + low = clip_percent / 100.0 + high = 1.0 - low + + with torch.no_grad(): + # Zero mean + mean = X.mean(dim=0, keepdim=True) + Xc = X - mean + + # Main computation: PCA + projection, try to use bf16 + # (auto-fallback if operator not supported) + device.startswith("cuda") and enable_autocast_bf16 + U, S, V = torch.pca_lowrank(Xc, q=q, center=False) # V: (D, q) + V3 = V[:, :k] # (3072, 3) + PCs = Xc @ V3 # (N, 3) + + # === Per-channel percentile clipping and normalization to [0,1] === + # Vectorized one-time calculation of low/high percentiles for each channel + qs = torch.tensor([low, high], device=PCs.device, dtype=PCs.dtype) + qvals = torch.quantile(PCs, q=qs, dim=0) # Shape (2, 3) + lo = qvals[0] # (3,) + hi = qvals[1] # (3,) + + # Avoid degenerate case where hi==lo + denom = torch.clamp(hi - lo, min=1e-8) + + # Broadcast clipping + normalization + PCs = torch.clamp(PCs, lo, hi) + PCs = (PCs - lo) / denom # (N, 3) in [0,1] + + # Restore 4D + PCs = PCs.reshape(B1, B2, B3, k) + + # Output + if return_uint8: + out = (PCs * 255.0).round().clamp(0, 255).to(torch.uint8).cpu().numpy() + else: + out = PCs.clamp(0, 1).to(torch.float32).cpu().numpy() + + return out + + +class PCARGBVisualizer: + """ + Stable PCA→RGB for video features shaped (T, H, W, D) or a single frame (H, W, D). + - Global mean/V3 reference for stable colors + - Per-frame PCA with Procrustes alignment to V3_ref (basis_mode='procrustes') + - Percentile normalization with global or EMA stats (time-only, no spatial smoothing) + """ + + def __init__( + self, + device=None, + q_oversample: int = 16, + clip_percent: float = 10.0, + return_uint8: bool = False, + enable_autocast_bf16: bool = True, + basis_mode: str = "procrustes", # 'fixed' | 'procrustes' + percentile_mode: str = "ema", # 'global' | 'ema' + ema_alpha: float = 0.1, + denom_eps: float = 1e-4, + ): + assert 0.0 <= clip_percent < 50.0 + assert basis_mode in ("fixed", "procrustes") + assert percentile_mode in ("global", "ema") + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.q = max(int(q_oversample), 6) + self.clip_percent = float(clip_percent) + self.return_uint8 = return_uint8 + self.enable_autocast_bf16 = enable_autocast_bf16 + self.basis_mode = basis_mode + self.percentile_mode = percentile_mode + self.ema_alpha = float(ema_alpha) + self.denom_eps = float(denom_eps) + + # reference state + self.mean_ref = None # (1, D) + self.V3_ref = None # (D, 3) + self.lo_ref = None # (3,) + self.hi_ref = None # (3,) + + @torch.no_grad() + def fit_reference(self, frames): + """ + Fit global mean/V3 and initialize percentiles from a reference set. + frames: ndarray (T,H,W,D) or list of (H,W,D) + """ + if isinstance(frames, np.ndarray): + if frames.ndim != 4: + raise ValueError("fit_reference expects (T,H,W,D) ndarray.") + T, H, W, D = frames.shape + X = torch.from_numpy(frames.reshape(T * H * W, D)) + else: # list of (H,W,D) + xs = [torch.from_numpy(x.reshape(-1, x.shape[-1])) for x in frames] + D = xs[0].shape[-1] + X = torch.cat(xs, dim=0) + + X = X.to(self.device, dtype=torch.float32) + X = torch.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6) + + mean = X.mean(0, keepdim=True) + Xc = X - mean + + U, S, V = torch.pca_lowrank(Xc, q=max(self.q, 8), center=False) + V3 = V[:, :3] # (D,3) + + PCs = Xc @ V3 + low = self.clip_percent / 100.0 + high = 1.0 - low + qs = torch.tensor([low, high], device=PCs.device, dtype=PCs.dtype) + qvals = torch.quantile(PCs, q=qs, dim=0) + lo, hi = qvals[0], qvals[1] + + self.mean_ref = mean + self.V3_ref = V3 + if self.percentile_mode == "global": + self.lo_ref, self.hi_ref = lo, hi + else: + self.lo_ref = lo.clone() + self.hi_ref = hi.clone() + + @torch.no_grad() + def _project_with_stable_colors(self, X: torch.Tensor) -> torch.Tensor: + """ + X: (N,D) where N = H*W + Returns PCs_raw: (N,3) using stable basis (fixed or Procrustes-aligned) + """ + assert self.mean_ref is not None and self.V3_ref is not None, "Call fit_reference() first." + X = torch.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6) + Xc = X - self.mean_ref + + if self.basis_mode == "fixed": + V3_used = self.V3_ref + else: + U, S, V = torch.pca_lowrank(Xc, q=max(self.q, 6), center=False) + V3 = V[:, :3] # (D,3) + M = V3.T @ self.V3_ref + Uo, So, Vh = torch.linalg.svd(M) + R = Uo @ Vh + V3_used = V3 @ R + # Optional polarity fix via anchor + a = self.V3_ref.mean(0, keepdim=True) + sign = torch.sign((V3_used * a).sum(0, keepdim=True)).clamp(min=-1) + V3_used = V3_used * sign + + return Xc @ V3_used + + @torch.no_grad() + def _normalize_rgb(self, PCs_raw: torch.Tensor) -> torch.Tensor: + assert self.lo_ref is not None and self.hi_ref is not None + if self.percentile_mode == "global": + lo, hi = self.lo_ref, self.hi_ref + else: + low = self.clip_percent / 100.0 + high = 1.0 - low + qs = torch.tensor([low, high], device=PCs_raw.device, dtype=PCs_raw.dtype) + qvals = torch.quantile(PCs_raw, q=qs, dim=0) + lo_now, hi_now = qvals[0], qvals[1] + a = self.ema_alpha + self.lo_ref = (1 - a) * self.lo_ref + a * lo_now + self.hi_ref = (1 - a) * self.hi_ref + a * hi_now + lo, hi = self.lo_ref, self.hi_ref + + denom = torch.clamp(hi - lo, min=self.denom_eps) + PCs = torch.clamp(PCs_raw, lo, hi) + PCs = (PCs - lo) / denom + return PCs.clamp_(0, 1) + + @torch.no_grad() + def transform_frame(self, frame: np.ndarray) -> np.ndarray: + """ + frame: (H,W,D) -> (H,W,3) + """ + if frame.ndim != 3: + raise ValueError("transform_frame expects (H,W,D).") + H, W, D = frame.shape + X = torch.from_numpy(frame.reshape(H * W, D)).to(self.device, dtype=torch.float32) + PCs_raw = self._project_with_stable_colors(X) + PCs = self._normalize_rgb(PCs_raw).reshape(H, W, 3) + if self.return_uint8: + return (PCs * 255.0).round().clamp(0, 255).to(torch.uint8).cpu().numpy() + return PCs.to(torch.float32).cpu().numpy() + + @torch.no_grad() + def transform_video(self, frames) -> np.ndarray: + """ + frames: (T,H,W,D) or list of (H,W,D) + returns: (T,H,W,3) + """ + outs = [] + if isinstance(frames, np.ndarray): + if frames.ndim != 4: + raise ValueError("transform_video expects (T,H,W,D).") + T, H, W, D = frames.shape + for t in range(T): + outs.append(self.transform_frame(frames[t])) + else: + for f in frames: + outs.append(self.transform_frame(f)) + return np.stack(outs, axis=0) diff --git a/src/nn/depth_anything_3/utils/pose_align.py b/src/nn/depth_anything_3/utils/pose_align.py new file mode 100644 index 0000000..695d07f --- /dev/null +++ b/src/nn/depth_anything_3/utils/pose_align.py @@ -0,0 +1,347 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List +import numpy as np +import torch +from evo.core.trajectory import PosePath3D + +from depth_anything_3.utils.geometry import affine_inverse, affine_inverse_np + + +def batch_apply_alignment_to_enc( + rots: torch.Tensor, trans: torch.Tensor, scales: torch.Tensor, enc_list: List[torch.Tensor] +): + pass + + +def batch_apply_alignment_to_ext( + rots: torch.Tensor, trans: torch.Tensor, scales: torch.Tensor, ext: torch.Tensor +): + device, _ = ext.device, ext.dtype + if ext.shape[-2:] == (3, 4): + pad = torch.zeros((*ext.shape[:-2], 4, 4), dtype=ext.dtype, device=device) + pad[..., :3, :4] = ext + pad[..., 3, 3] = 1.0 + ext = pad + pose_est = affine_inverse(ext) + pose_new_align_rot = rots[:, None] @ pose_est[..., :3, :3] + pose_new_align_trans = ( + scales[:, None, None] * (rots[:, None] @ pose_est[..., :3, 3:])[..., 0] + trans[:, None] + ) + pose_new_align = torch.zeros_like(ext) + pose_new_align[..., :3, :3] = pose_new_align_rot + pose_new_align[..., :3, 3] = pose_new_align_trans + pose_new_align[..., 3, 3] = 1.0 + return affine_inverse(pose_new_align)[:, :3] + + +def batch_align_poses_umeyama(ext_ref: torch.Tensor, ext_est: torch.Tensor): + device, dtype = ext_ref.device, ext_ref.dtype + assert ext_ref.dtype in [torch.float32, torch.float64] + assert ext_est.dtype in [torch.float32, torch.float64] + assert ext_ref.requires_grad is False + assert ext_est.requires_grad is False + rots, trans, scales = [], [], [] + for b in range(ext_ref.shape[0]): + r, t, s = align_poses_umeyama(ext_ref[b].cpu().numpy(), ext_est[b].cpu().numpy()) + rots.append(torch.from_numpy(r).to(device=device, dtype=dtype)) + trans.append(torch.from_numpy(t).to(device=device, dtype=dtype)) + scales.append(torch.tensor(s, device=device, dtype=dtype)) + return torch.stack(rots), torch.stack(trans), torch.stack(scales) + + +# Dependencies: affine_inverse_np, PosePath3D (maintain consistency with your existing project) + + +def _to44(ext): + if ext.shape[1] == 3: + out = np.eye(4)[None].repeat(len(ext), 0) + out[:, :3, :4] = ext + return out + return ext + + +def _poses_from_ext(ext_ref, ext_est): + ext_ref = _to44(ext_ref) + ext_est = _to44(ext_est) + pose_ref = affine_inverse_np(ext_ref) + pose_est = affine_inverse_np(ext_est) + return pose_ref, pose_est + + +def _umeyama_sim3_from_paths(pose_ref, pose_est): + path_ref = PosePath3D(poses_se3=pose_ref.copy()) + path_est = PosePath3D(poses_se3=pose_est.copy()) + r, t, s = path_est.align(path_ref, correct_scale=True) + pose_est_aligned = np.stack(path_est.poses_se3) + return r, t, s, pose_est_aligned + + +def _apply_sim3_to_poses(poses, r, t, s): + out = poses.copy() + Ri = poses[:, :3, :3] + ti = poses[:, :3, 3] + out[:, :3, :3] = r @ Ri + out[:, :3, 3] = (r @ (s * ti.T)).T + t + return out + + +def _median_nn_thresh(pose_ref, pose_est_aligned): + P_ref = pose_ref[:, :3, 3] + P_est = pose_est_aligned[:, :3, 3] + dists = [] + for p in P_est: + dd = np.linalg.norm(P_ref - p[None, :], axis=1) + dists.append(dd.min()) + return float(np.median(dists)) if dists else 0.0 + + +def _ransac_align_sim3( + pose_ref, pose_est, sub_n=None, inlier_thresh=None, max_iters=10, random_state=None +): + rng = np.random.default_rng(random_state) + N = pose_ref.shape[0] + idx_all = np.arange(N) + if sub_n is None: + sub_n = max(3, (N + 1) // 2) + else: + sub_n = max(3, min(sub_n, N)) + + # Pre-alignment + default threshold + r0, t0, s0, pose_est0 = _umeyama_sim3_from_paths(pose_ref, pose_est) + if inlier_thresh is None: + inlier_thresh = _median_nn_thresh(pose_ref, pose_est0) + + P_ref_all = pose_ref[:, :3, 3] + + best_model = (r0, t0, s0) + best_inliers = None + best_score = (-1, np.inf) # (num_inliers, mean_err) + + for _ in range(max_iters): + sample = rng.choice(idx_all, size=sub_n, replace=False) + try: + r, t, s, _ = _umeyama_sim3_from_paths(pose_ref[sample], pose_est[sample]) + except Exception: + continue + pose_h = _apply_sim3_to_poses(pose_est, r, t, s) + P_h = pose_h[:, :3, 3] + errs = np.linalg.norm(P_h - P_ref_all, axis=1) # Match by same index + inliers = errs <= inlier_thresh + k = int(inliers.sum()) + mean_err = float(errs[inliers].mean()) if k > 0 else np.inf + if (k > best_score[0]) or (k == best_score[0] and mean_err < best_score[1]): + best_score = (k, mean_err) + best_model = (r, t, s) + best_inliers = inliers + + # Fit again with best inliers + if best_inliers is not None and best_inliers.sum() >= 3: + r, t, s, _ = _umeyama_sim3_from_paths(pose_ref[best_inliers], pose_est[best_inliers]) + else: + r, t, s = best_model + return r, t, s + + +def align_poses_umeyama( + ext_ref: np.ndarray, + ext_est: np.ndarray, + return_aligned=False, + ransac=False, + sub_n=None, + inlier_thresh=None, + ransac_max_iters=10, + random_state=None, +): + """ + Align estimated trajectory to reference using Umeyama Sim(3). + Default no RANSAC; if ransac=True, use RANSAC (max iterations default 10). + - sub_n defaults to half the number of frames (rounded up, at least 3) + - inlier_thresh defaults to median of "distance from each estimated pose to + nearest reference pose after pre-alignment" + Returns rotation (3x3), translation (3,), scale; optionally returns aligned extrinsics (4x4). + """ + pose_ref, pose_est = _poses_from_ext(ext_ref, ext_est) + + if not ransac: + r, t, s, pose_est_aligned = _umeyama_sim3_from_paths(pose_ref, pose_est) + else: + r, t, s = _ransac_align_sim3( + pose_ref, + pose_est, + sub_n=sub_n, + inlier_thresh=inlier_thresh, + max_iters=ransac_max_iters, + random_state=random_state, + ) + pose_est_aligned = _apply_sim3_to_poses(pose_est, r, t, s) + + if return_aligned: + ext_est_aligned = affine_inverse_np(pose_est_aligned) + return r, t, s, ext_est_aligned + return r, t, s + + +# def align_poses_umeyama(ext_ref: np.ndarray, ext_est: np.ndarray, return_aligned=False): +# """ +# Align estimated trajectory to reference trajectory using Umeyama Sim(3) +# alignment (via evo PosePath3D). # noqa +# Returns rotation, translation, and scale. +# """ +# # If input extrinsics are 3x4, convert to 4x4 by padding +# if ext_ref.shape[1] == 3: +# ext_ref_ = np.eye(4)[None].repeat(len(ext_ref), 0) +# ext_ref_[:, :3] = ext_ref +# ext_ref = ext_ref_ +# if ext_est.shape[1] == 3: +# ext_est_ = np.eye(4)[None].repeat(len(ext_est), 0) +# ext_est_[:, :3] = ext_est +# ext_est = ext_est_ + +# # Convert to camera poses (inverse extrinsics) +# pose_ref = affine_inverse_np(ext_ref) +# pose_est = affine_inverse_np(ext_est) + +# # Create evo PosePath3D objects +# path_ref = PosePath3D(poses_se3=pose_ref) +# path_est = PosePath3D(poses_se3=pose_est) +# r, t, s = path_est.align(path_ref, correct_scale=True) +# if return_aligned: +# return r, t, s, affine_inverse_np(np.stack(path_est.poses_se3)) +# else: +# return r, t, s + + +def apply_umeyama_alignment_to_ext( + rot: np.ndarray, # (3,3) + trans: np.ndarray, # (3,) or (1,3) + scale: float, + ext_est: np.ndarray, # (...,4,4) or (...,3,4) +) -> np.ndarray: + """ + Apply Sim(3) (R, t, s) to a batch of world-to-camera extrinsics ext_est. + Returns the aligned extrinsics, with the same shape as input. + """ + + # Allow 3x4 extrinsics: pad to 4x4 + if ext_est.shape[-2:] == (3, 4): + pad = np.zeros((*ext_est.shape[:-2], 4, 4), dtype=ext_est.dtype) + pad[..., :3, :4] = ext_est + pad[..., 3, 3] = 1.0 + ext_est = pad + + # Convert world-to-camera to camera-to-world + pose_est = affine_inverse_np(ext_est) # (...,4,4) + R_e = pose_est[..., :3, :3] # (...,3,3) + t_e = pose_est[..., :3, 3] # (...,3) + + # Apply Sim(3) transformation + R_a = np.einsum("ij,...jk->...ik", rot, R_e) # (...,3,3) + t_a = scale * np.einsum("ij,...j->...i", rot, t_e) + trans # (...,3) + + # Assemble the transformed pose + pose_a = np.zeros_like(pose_est) + pose_a[..., :3, :3] = R_a + pose_a[..., :3, 3] = t_a + pose_a[..., 3, 3] = 1.0 + + # Convert back to world-to-camera + return affine_inverse_np(pose_a) + + +def transform_points_sim3(points, rot, trans, scale, inverse=False): + """ + Sim(3) transform point cloud + points: (N, 3) + rot: (3, 3) + trans: (3,) or (1, 3) + scale: float + inverse: Whether to do inverse transform (ref->est) + Returns: (N, 3) + """ + if not inverse: + # Forward: est -> ref + return scale * (points @ rot.T) + trans + else: + # Inverse: ref -> est + return ((points - trans) @ rot) / scale + + +def _rand_rot(): + u1, u2, u3 = np.random.rand(3) + q = np.array( + [ + np.sqrt(1 - u1) * np.sin(2 * np.math.pi * u2), + np.sqrt(1 - u1) * np.cos(2 * np.math.pi * u2), + np.sqrt(u1) * np.sin(2 * np.math.pi * u3), + np.sqrt(u1) * np.cos(2 * np.math.pi * u3), + ] + ) + w, x, y, z = q + return np.array( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ] + ) + + +def _rand_pose(): + R, t = _rand_rot(), np.random.randn(3) + P = np.eye(4) + P[:3, :3] = R + P[:3, 3] = t + return P + + +if __name__ == "__main__": + np.random.seed(42) + # 1. Randomly generate reference trajectory and Sim(3) + N = 8 + pose_ref = np.stack([_rand_pose() for _ in range(N)]) # (N,4,4) cam→world + rot_gt = _rand_rot() + scale_gt = 2.3 + trans_gt = np.random.randn(3) + # 2. Generate estimated trajectory (apply Sim(3)) + pose_est = np.zeros_like(pose_ref) + for i in range(N): + R = pose_ref[i][:3, :3] + t = pose_ref[i][:3, 3] + pose_est[i][:3, :3] = rot_gt @ R + pose_est[i][:3, 3] = scale_gt * (rot_gt @ t) + trans_gt + pose_est[i][3, 3] = 1.0 + # 3. Get extrinsics (world->cam) + ext_ref = affine_inverse_np(pose_ref) + ext_est = affine_inverse_np(pose_est) + # 4. Use umeyama alignment, estimate Sim(3) + r_est, t_est, s_est = align_poses_umeyama(ext_ref, ext_est) + print("GT scale:", scale_gt, "Estimated:", s_est) + print("GT trans:", trans_gt, "Estimated:", t_est) + print("GT rot:\n", rot_gt, "\nEstimated:\n", r_est) + # 5. Random point cloud, in ref frame + num_points = 100 + points_ref = np.random.randn(num_points, 3) + # 6. Use GT Sim(3) inverse transform to est frame + points_est = transform_points_sim3(points_ref, rot_gt, trans_gt, scale_gt, inverse=True) + # 7. Use estimated Sim(3) forward transform back to ref frame + points_ref_recovered = transform_points_sim3(points_est, r_est, t_est, s_est, inverse=False) + # 8. Check error + err = np.abs(points_ref_recovered - points_ref) + print("Point cloud sim3 transform error (mean abs):", err.mean()) + print("Point cloud sim3 transform error (max abs):", err.max()) + assert err.mean() < 1e-6, "Mean sim3 transform error too large!" + assert err.max() < 1e-5, "Max sim3 transform error too large!" + print("Sim(3) point cloud transform & alignment test passed!") diff --git a/src/nn/depth_anything_3/utils/ray_utils.py b/src/nn/depth_anything_3/utils/ray_utils.py new file mode 100644 index 0000000..6244dc4 --- /dev/null +++ b/src/nn/depth_anything_3/utils/ray_utils.py @@ -0,0 +1,523 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from einops import repeat +from .geometry import unproject_depth + + +def compute_optimal_rotation_intrinsics_batch( + rays_origin, rays_target, z_threshold=1e-4, reproj_threshold=0.2, weights=None, + n_sample = None, + n_iter=100, + num_sample_for_ransac=8, + rand_sample_iters_idx=None, +): + """ + Args: + rays_origin (torch.Tensor): (B, N, 3) + rays_target (torch.Tensor): (B, N, 3) + z_threshold (float): Threshold for z value to be considered valid. + + Returns: + R (torch.tensor): (3, 3) + focal_length (torch.tensor): (2,) + principal_point (torch.tensor): (2,) + """ + device = rays_origin.device + B, N, _ = rays_origin.shape + z_mask = torch.logical_and( + torch.abs(rays_target[:, :, 2]) > z_threshold, torch.abs(rays_origin[:, :, 2]) > z_threshold + ) # (B, N, 1) + rays_origin = rays_origin.clone() + rays_target = rays_target.clone() + rays_origin[:, :, 0][z_mask] /= rays_origin[:, :, 2][z_mask] + rays_origin[:, :, 1][z_mask] /= rays_origin[:, :, 2][z_mask] + rays_target[:, :, 0][z_mask] /= rays_target[:, :, 2][z_mask] + rays_target[:, :, 1][z_mask] /= rays_target[:, :, 2][z_mask] + + rays_origin = rays_origin[:, :, :2] + rays_target = rays_target[:, :, :2] + assert weights is not None, "weights must be provided" + weights[~z_mask] = 0 + + A_list = [] + max_chunk_size = 2 + for i in range(0, rays_origin.shape[0], max_chunk_size): + A = ransac_find_homography_weighted_fast_batch( + rays_origin[i:i+max_chunk_size], + rays_target[i:i+max_chunk_size], + weights[i:i+max_chunk_size], + n_iter=n_iter, + n_sample = n_sample, + num_sample_for_ransac=num_sample_for_ransac, + reproj_threshold=reproj_threshold, + rand_sample_iters_idx=rand_sample_iters_idx, + max_inlier_num=8000, + ) + A = A.to(device) + A_need_inv_mask = torch.linalg.det(A) < 0 + A[A_need_inv_mask] = -A[A_need_inv_mask] + A_list.append(A) + + A = torch.cat(A_list, dim=0) + + R_list = [] + f_list = [] + pp_list = [] + for i in range(A.shape[0]): + R, L = ql_decomposition(A[i]) + L = L / L[2][2] + + f = torch.stack((L[0][0], L[1][1])) + pp = torch.stack((L[2][0], L[2][1])) + R_list.append(R) + f_list.append(f) + pp_list.append(pp) + + R = torch.stack(R_list) + f = torch.stack(f_list) + pp = torch.stack(pp_list) + + return R, f, pp + + +# https://www.reddit.com/r/learnmath/comments/v1crd7/linear_algebra_qr_to_ql_decomposition/ +def ql_decomposition(A): + P = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=A.device).float() + A_tilde = torch.matmul(A, P) + Q_tilde, R_tilde = torch.linalg.qr(A_tilde) + Q = torch.matmul(Q_tilde, P) + L = torch.matmul(torch.matmul(P, R_tilde), P) + d = torch.diag(L) + Q[:, 0] *= torch.sign(d[0]) + Q[:, 1] *= torch.sign(d[1]) + Q[:, 2] *= torch.sign(d[2]) + L[0] *= torch.sign(d[0]) + L[1] *= torch.sign(d[1]) + L[2] *= torch.sign(d[2]) + return Q, L + +def find_homography_least_squares_weighted_torch(src_pts, dst_pts, confident_weight): + """ + src_pts: (N,2) source points (torch.Tensor, float32/float64) + dst_pts: (N,2) target points (torch.Tensor, float32/float64) + confident_weight: (N,) weights (torch.Tensor) + Returns: (3,3) homography matrix H (torch.Tensor) + """ + assert src_pts.shape == dst_pts.shape + N = src_pts.shape[0] + if N < 4: + raise ValueError("At least 4 points are required to compute homography.") + assert confident_weight.shape == (N,) + + w = confident_weight.sqrt().unsqueeze(1) # (N,1) + + x = src_pts[:, 0:1] # (N,1) + y = src_pts[:, 1:2] # (N,1) + u = dst_pts[:, 0:1] + v = dst_pts[:, 1:2] + + zeros = torch.zeros_like(x) + + # Construct A matrix (2N, 9) + A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=1) + A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=1) + A = torch.cat([A1, A2], dim=0) # (2N, 9) + + # SVD + # Note: torch.linalg.svd returns U, S, Vh, where Vh is the transpose of V + _, _, Vh = torch.linalg.svd(A) + H = Vh[-1].reshape(3, 3) + H = H / H[-1, -1] + return H + + +def ransac_find_homography_weighted( + src_pts, + dst_pts, + confident_weight, + n_iter=100, + sample_ratio=0.2, + reproj_threshold=3.0, + num_sample_for_ransac=16, + random_seed=None, +): + """ + RANSAC version of weighted Homography estimation. + Sample 4 points from the top 50% weighted points each time. + reproj_threshold: points with reprojection error less than this value are inliers + Returns: best_H + """ + if random_seed is not None: + torch.manual_seed(random_seed) + N = src_pts.shape[0] + assert N >= 4 + # 1. Select top 50% weighted points + sorted_idx = torch.argsort(confident_weight, descending=True) + n_sample = max(num_sample_for_ransac, int(N * sample_ratio)) + candidate_idx = sorted_idx[:n_sample] + best_inlier_mask = None + best_score = 0 + for _ in range(n_iter): + # 2. Randomly sample 4 points + idx = candidate_idx[torch.randperm(n_sample)[:num_sample_for_ransac]] + # 3. Compute Homography + try: + H = find_homography_least_squares_weighted_torch( + src_pts[idx], dst_pts[idx], confident_weight[idx] + ) + except Exception: + H = torch.eye(3, dtype=src_pts.dtype, device=src_pts.device) + # 4. Compute reprojection error for all points + src_homo = torch.cat( + [src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1 + ) + proj = (H @ src_homo.T).T + proj = proj[:, :2] / proj[:, 2:3] + error = ((proj - dst_pts) ** 2).sum(dim=1).sqrt() # Euclidean distance + inlier_mask = error < reproj_threshold + total_score = (inlier_mask * confident_weight).sum().item() + n_inlier = inlier_mask.sum().item() + if n_inlier < 4: + continue # At least 4 inliers required for fitting + + if total_score > best_score: + best_score = total_score + best_inlier_mask = inlier_mask + + # 5. Refit Homography using inliers + H_inlier = find_homography_least_squares_weighted_torch( + src_pts[best_inlier_mask], dst_pts[best_inlier_mask], confident_weight[best_inlier_mask] + ) + + return H_inlier + + +def find_homography_least_squares_weighted_torch_batch( + src_pts_batch, dst_pts_batch, confident_weight_batch +): + """ + Batch version of weighted least squares Homography + src_pts_batch: (B, K, 2) + dst_pts_batch: (B, K, 2) + confident_weight_batch: (B, K) + Returns: (B, 3, 3) + """ + B, K, _ = src_pts_batch.shape + w = confident_weight_batch.sqrt().unsqueeze(2) # (B,K,1) + x = src_pts_batch[:, :, 0:1] + y = src_pts_batch[:, :, 1:2] + u = dst_pts_batch[:, :, 0:1] + v = dst_pts_batch[:, :, 1:2] + zeros = torch.zeros_like(x) + A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=2) + A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=2) + A = torch.cat([A1, A2], dim=1) # (B, 2K, 9) + # SVD: torch.linalg.svd supports batch + _, _, Vh = torch.linalg.svd(A) + H = Vh[:, -1].reshape(B, 3, 3) + H = H / H[:, 2:3, 2:3] + return H + + +def ransac_find_homography_weighted_fast( + src_pts, + dst_pts, + confident_weight, + n_sample, + n_iter=100, + reproj_threshold=3.0, + num_sample_for_ransac=8, + random_seed=None, + rand_sample_iters_idx=None, +): + """ + Batch version of RANSAC weighted Homography estimation. + Returns: H_inlier + """ + if random_seed is not None: + torch.manual_seed(random_seed) + N = src_pts.shape[0] + device = src_pts.device + assert N >= 4 + # 1. Select top weighted points by sample_ratio + sorted_idx = torch.argsort(confident_weight, descending=True) + candidate_idx = sorted_idx[:n_sample] # (n_sample,) + if rand_sample_iters_idx is None: + rand_sample_iters_idx = torch.stack( + [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], + dim=0, + ) # (n_iter, num_sample_for_ransac) + # 2. Generate all sampling groups at once + # shape: (n_iter, num_sample_for_ransac) + rand_idx = candidate_idx[rand_sample_iters_idx] # (n_iter, num_sample_for_ransac) + # 3. Construct batch input + src_pts_batch = src_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2) + dst_pts_batch = dst_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2) + confident_weight_batch = confident_weight[rand_idx] # (n_iter, num_sample_for_ransac) + # 4. Batch fit Homography + H_batch = find_homography_least_squares_weighted_torch_batch( + src_pts_batch, dst_pts_batch, confident_weight_batch + ) # (n_iter, 3, 3) + # 5. Batch evaluate inliers for all H + src_homo = torch.cat( + [src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1 + ) # (N,3) + src_homo_expand = src_homo.unsqueeze(0).expand(n_iter, N, 3) # (n_iter, N, 3) + dst_pts_expand = dst_pts.unsqueeze(0).expand(n_iter, N, 2) # (n_iter, N, 2) + confident_weight_expand = confident_weight.unsqueeze(0).expand(n_iter, N) # (n_iter, N) + # H_batch: (n_iter, 3, 3) + proj = torch.bmm(src_homo_expand, H_batch.transpose(1, 2)) # (n_iter, N, 3) + proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (n_iter, N, 2) + error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=2).sqrt() # (n_iter, N) + inlier_mask = error < reproj_threshold # (n_iter, N) + total_score = (inlier_mask * confident_weight_expand).sum(dim=1) # (n_iter,) + # 6. Select the sampling group with the highest score + best_idx = torch.argmax(total_score) + best_inlier_mask = inlier_mask[best_idx] # (N,) + inlier_src_pts = src_pts[best_inlier_mask] + inlier_dst_pts = dst_pts[best_inlier_mask] + inlier_confident_weight = confident_weight[best_inlier_mask] + + max_inlier_num = 10000 + sorted_idx = torch.argsort(inlier_confident_weight, descending=True) + + # method 1: sort according to confident_weight, and only keep max_inlier_num pts + # sorted_idx = sorted_idx[:max_inlier_num] + + # method 2: random choose max_inlier_num pts + sorted_idx = sorted_idx[torch.randperm(len(sorted_idx))[:max_inlier_num]] + + inlier_src_pts = inlier_src_pts[sorted_idx] + inlier_dst_pts = inlier_dst_pts[sorted_idx] + inlier_confident_weight = inlier_confident_weight[sorted_idx] + # 7. Refit Homography using inliers + H_inlier = find_homography_least_squares_weighted_torch( + inlier_src_pts, inlier_dst_pts, inlier_confident_weight + ) + return H_inlier + + +def ransac_find_homography_weighted_fast_batch( + src_pts, # (B, N, 3) + dst_pts, # (B, N, 2) + confident_weight, # (B, N) + n_sample, + n_iter=100, + reproj_threshold=3.0, + num_sample_for_ransac=8, + max_inlier_num=10000, + random_seed=None, + rand_sample_iters_idx=None, +): + """ + Batch version of RANSAC weighted Homography estimation (supports batch). + Input: + src_pts: (B, N, 2) + dst_pts: (B, N, 2) + confident_weight: (B, N) + Returns: + H_inlier: (B, 3, 3) + """ + if random_seed is not None: + torch.manual_seed(random_seed) + B, N, _ = src_pts.shape + assert N >= 4 + + device = src_pts.device + + # 1. Select top weighted points by sample_ratio + sorted_idx = torch.argsort(confident_weight, descending=True, dim=1) # (B, N) + candidate_idx = sorted_idx[:, :n_sample] # (B, n_sample) + + # 2. Generate all sampling groups at once + # rand_idx: (B, n_iter, num_sample_for_ransac) + if rand_sample_iters_idx is None: + rand_sample_iters_idx = torch.stack( + [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], + dim=0, + ) # (n_iter, num_sample_for_ransac) + + rand_idx = candidate_idx[:, rand_sample_iters_idx] # (B, n_iter, num_sample_for_ransac) + + # 3. Construct batch input + # Indexing method below: (B, n_iter, num_sample_for_ransac, ...) + b_idx = torch.arange(B, device=device).view(B, 1, 1).expand(B, n_iter, num_sample_for_ransac) + src_pts_batch = src_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2) + dst_pts_batch = dst_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2) + confident_weight_batch = confident_weight[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac) + + # 4. Batch fit Homography + # Need to implement batch version that supports (B, n_iter, num_sample_for_ransac, ...) input + # Output H_batch: (B, n_iter, 3, 3) + cB, cN = src_pts_batch.shape[:2] + H_batch = find_homography_least_squares_weighted_torch_batch( + src_pts_batch.flatten(0, 1), dst_pts_batch.flatten(0, 1), confident_weight_batch.flatten(0, 1) + ) # (B, n_iter, 3, 3) + H_batch = H_batch.unflatten(0, (cB, cN)) + + # 5. Batch evaluate inliers for all H + src_homo = torch.cat( + [src_pts, torch.ones(B, N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=2 + ) # (B, N, 3) + src_homo_expand = src_homo.unsqueeze(1).expand(B, n_iter, N, 3) # (B, n_iter, N, 3) + dst_pts_expand = dst_pts.unsqueeze(1).expand(B, n_iter, N, 2) # (B, n_iter, N, 2) + confident_weight_expand = confident_weight.unsqueeze(1).expand(B, n_iter, N) # (B, n_iter, N) + + # H_batch: (B, n_iter, 3, 3) + # Need to reshape H_batch to (B*n_iter, 3, 3), src_homo_expand to (B*n_iter, N, 3) + H_batch_flat = H_batch.reshape(-1, 3, 3) + src_homo_expand_flat = src_homo_expand.reshape(-1, N, 3) + proj = torch.bmm(src_homo_expand_flat, H_batch_flat.transpose(1, 2)) # (B*n_iter, N, 3) + proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (B*n_iter, N, 2) + proj_xy = proj_xy.reshape(B, n_iter, N, 2) + error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=3).sqrt() # (B, n_iter, N) + inlier_mask = error < reproj_threshold # (B, n_iter, N) + total_score = (inlier_mask * confident_weight_expand).sum(dim=2) # (B, n_iter) + + # 6. Select the sampling group with the highest score + best_idx = torch.argmax(total_score, dim=1) # (B,) + best_inlier_mask = inlier_mask[torch.arange(B, device=device), best_idx] # (B, N) + + # 7. Refit Homography using inliers + H_inlier_list = [] + for b in range(B): + mask = best_inlier_mask[b] + inlier_src_pts = src_pts[b][mask] # (?, 3) + inlier_dst_pts = dst_pts[b][mask] # (?, 2) + inlier_confident_weight = confident_weight[b][mask] # (?) + + sorted_idx = torch.argsort(inlier_confident_weight, descending=True) + # # method 1: sort according to confident_weight, and only keep max_inlier_num pts + # sorted_idx = sorted_idx[:max_inlier_num] + # method 2: random choose max_inlier_num pts + if len(sorted_idx) > max_inlier_num: + # random choose from first 95% confident pts + keep_len = max(int(len(sorted_idx) * 0.95), max_inlier_num) + sorted_idx = sorted_idx[:keep_len] + perm = torch.randperm(len(sorted_idx), device=device)[:max_inlier_num] + sorted_idx = sorted_idx[perm] + inlier_src_pts = inlier_src_pts[sorted_idx] + inlier_dst_pts = inlier_dst_pts[sorted_idx] + inlier_confident_weight = inlier_confident_weight[sorted_idx] + + H_inlier = find_homography_least_squares_weighted_torch( + inlier_src_pts, inlier_dst_pts, inlier_confident_weight + ) # (3, 3) + H_inlier_list.append(H_inlier) + H_inlier = torch.stack(H_inlier_list, dim=0) # (B, 3, 3) + return H_inlier + +def get_params_for_ransac(N, device): + n_iter=100 + sample_ratio=0.3 + num_sample_for_ransac=8 + n_sample = max(num_sample_for_ransac, int(N * sample_ratio)) + rand_sample_iters_idx = torch.stack( + [torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)], + dim=0, + ) # (n_iter, num_sample_for_ransac) + return n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx + + +def camray_to_caminfo(camray, confidence=None, reproj_threshold=0.2, training=False): + """ + Args: + camray: (B, S, num_patches_y, num_patches_x, 6) + confidence: (B, S, num_patches_y, num_patches_x) + Returns: + R: (B, S, 3, 3) + T: (B, S, 3) + focal_lengths: (B, S, 2) + principal_points: (B, S, 2) + """ + if confidence is None: + confidence = torch.ones_like(camray[:, :, :, :, 0]) + B, S, num_patches_y, num_patches_x, _ = camray.shape + # identity K, assume imw=imh=2.0 + I_K = torch.eye(3, dtype=camray.dtype, device=camray.device) + I_K[0, 2] = 1.0 + I_K[1, 2] = 1.0 + # repeat I_K to match camray + I_K = I_K.unsqueeze(0).unsqueeze(0).expand(B, S, -1, -1) + + cam_plane_depth = torch.ones( + B, S, num_patches_y, num_patches_x, 1, dtype=camray.dtype, device=camray.device + ) + I_cam_plane_unproj = unproject_depth( + cam_plane_depth, + I_K, + c2w=None, + ixt_normalized=True, + num_patches_x=num_patches_x, + num_patches_y=num_patches_y, + ) # (B, S, num_patches_y, num_patches_x, 3) + + camray = camray.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x, 6) + I_cam_plane_unproj = I_cam_plane_unproj.flatten(0, 1).flatten( + 1, 2 + ) # (B*S, num_patches_y*num_patches_x, 3) + confidence = confidence.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x) + + # Compute optimal rotation to align rays + N = camray.shape[-2] + device = camray.device + n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx = get_params_for_ransac(N, device) + + # Use batch processing (confidence is guaranteed to be not None at this point) + if training: + camray = camray.clone().detach() + I_cam_plane_unproj = I_cam_plane_unproj.clone().detach() + confidence = confidence.clone().detach() + R, focal_lengths, principal_points = compute_optimal_rotation_intrinsics_batch( + I_cam_plane_unproj, + camray[:, :, :3], + reproj_threshold=reproj_threshold, + weights=confidence, + n_sample = n_sample, + n_iter=n_iter, + num_sample_for_ransac=num_sample_for_ransac, + rand_sample_iters_idx=rand_sample_iters_idx, + ) + + T = torch.sum(camray[:, :, 3:] * confidence.unsqueeze(-1), dim=1) / torch.sum( + confidence, dim=-1, keepdim=True + ) + + R = R.reshape(B, S, 3, 3) + T = T.reshape(B, S, 3) + focal_lengths = focal_lengths.reshape(B, S, 2) + principal_points = principal_points.reshape(B, S, 2) + + return R, T, 1.0 / focal_lengths, principal_points + 1.0 + +def get_extrinsic_from_camray(camray, conf, patch_size_y, patch_size_x, training=False): + pred_R, pred_T, pred_focal_lengths, pred_principal_points = camray_to_caminfo( + camray, confidence=conf.squeeze(-1), training=training + ) + + pred_extrinsic = torch.cat( + [ + torch.cat([pred_R, pred_T.unsqueeze(-1)], dim=-1), + repeat( + torch.tensor([0, 0, 0, 1], dtype=pred_R.dtype, device=pred_R.device), + "c -> b s 1 c", + b=pred_R.shape[0], + s=pred_R.shape[1], + ), + ], + dim=-2, + ) # B, S, 4, 4 + return pred_extrinsic, pred_focal_lengths, pred_principal_points \ No newline at end of file diff --git a/src/nn/depth_anything_3/utils/read_write_model.py b/src/nn/depth_anything_3/utils/read_write_model.py new file mode 100644 index 0000000..4b4cf19 --- /dev/null +++ b/src/nn/depth_anything_3/utils/read_write_model.py @@ -0,0 +1,585 @@ +# Copyright (c), ETH Zurich and UNC Chapel Hill. +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# All rights reserved. +# +# This file has been modified by ByteDance Ltd. and/or its affiliates. on 11/05/2025 +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of +# its contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + + +import argparse +import collections +import os +import struct +import numpy as np + +CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"]) +Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) +BaseImage = collections.namedtuple( + "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"] +) +Point3D = collections.namedtuple( + "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"] +) + + +class Image(BaseImage): + def qvec2rotmat(self): + return qvec2rotmat(self.qvec) + + +CAMERA_MODELS = { + CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), + CameraModel(model_id=1, model_name="PINHOLE", num_params=4), + CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), + CameraModel(model_id=3, model_name="RADIAL", num_params=5), + CameraModel(model_id=4, model_name="OPENCV", num_params=8), + CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), + CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), + CameraModel(model_id=7, model_name="FOV", num_params=5), + CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), + CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), + CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), +} +CAMERA_MODEL_IDS = {camera_model.model_id: camera_model for camera_model in CAMERA_MODELS} +CAMERA_MODEL_NAMES = {camera_model.model_name: camera_model for camera_model in CAMERA_MODELS} + + +def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): + """Read and unpack the next bytes from a binary file. + :param fid: + :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. + :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. + :param endian_character: Any of {@, =, <, >, !} + :return: Tuple of read and unpacked values. + """ + data = fid.read(num_bytes) + return struct.unpack(endian_character + format_char_sequence, data) + + +def write_next_bytes(fid, data, format_char_sequence, endian_character="<"): + """pack and write to a binary file. + :param fid: + :param data: data to send, if multiple elements are sent at the same time, + they should be encapsuled either in a list or a tuple + :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. + should be the same length as the data list or tuple + :param endian_character: Any of {@, =, <, >, !} + """ + if isinstance(data, (list, tuple)): + bytes = struct.pack(endian_character + format_char_sequence, *data) + else: + bytes = struct.pack(endian_character + format_char_sequence, data) + fid.write(bytes) + + +def read_cameras_text(path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasText(const std::string& path) + void Reconstruction::ReadCamerasText(const std::string& path) + """ + cameras = {} + with open(path) as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + camera_id = int(elems[0]) + model = elems[1] + width = int(elems[2]) + height = int(elems[3]) + params = np.array(tuple(map(float, elems[4:]))) + cameras[camera_id] = Camera( + id=camera_id, + model=model, + width=width, + height=height, + params=params, + ) + return cameras + + +def read_cameras_binary(path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasBinary(const std::string& path) + void Reconstruction::ReadCamerasBinary(const std::string& path) + """ + cameras = {} + with open(path_to_model_file, "rb") as fid: + num_cameras = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_cameras): + camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ") + camera_id = camera_properties[0] + model_id = camera_properties[1] + model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name + width = camera_properties[2] + height = camera_properties[3] + num_params = CAMERA_MODEL_IDS[model_id].num_params + params = read_next_bytes( + fid, + num_bytes=8 * num_params, + format_char_sequence="d" * num_params, + ) + cameras[camera_id] = Camera( + id=camera_id, + model=model_name, + width=width, + height=height, + params=np.array(params), + ) + assert len(cameras) == num_cameras + return cameras + + +def write_cameras_text(cameras, path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasText(const std::string& path) + void Reconstruction::ReadCamerasText(const std::string& path) + """ + HEADER = ( + "# Camera list with one line of data per camera:\n" + + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" + + f"# Number of cameras: {len(cameras)}\n" + ) + with open(path, "w") as fid: + fid.write(HEADER) + for _, cam in cameras.items(): + to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params] + line = " ".join([str(elem) for elem in to_write]) + fid.write(line + "\n") + + +def write_cameras_binary(cameras, path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::WriteCamerasBinary(const std::string& path) + void Reconstruction::ReadCamerasBinary(const std::string& path) + """ + with open(path_to_model_file, "wb") as fid: + write_next_bytes(fid, len(cameras), "Q") + for _, cam in cameras.items(): + model_id = CAMERA_MODEL_NAMES[cam.model].model_id + camera_properties = [cam.id, model_id, cam.width, cam.height] + write_next_bytes(fid, camera_properties, "iiQQ") + for p in cam.params: + write_next_bytes(fid, float(p), "d") + return cameras + + +def read_images_text(path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesText(const std::string& path) + void Reconstruction::WriteImagesText(const std::string& path) + """ + images = {} + with open(path) as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + image_id = int(elems[0]) + qvec = np.array(tuple(map(float, elems[1:5]))) + tvec = np.array(tuple(map(float, elems[5:8]))) + camera_id = int(elems[8]) + image_name = elems[9] + elems = fid.readline().split() + xys = np.column_stack( + [ + tuple(map(float, elems[0::3])), + tuple(map(float, elems[1::3])), + ] + ) + point3D_ids = np.array(tuple(map(int, elems[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def read_images_binary(path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesBinary(const std::string& path) + void Reconstruction::WriteImagesBinary(const std::string& path) + """ + images = {} + with open(path_to_model_file, "rb") as fid: + num_reg_images = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_reg_images): + binary_image_properties = read_next_bytes( + fid, num_bytes=64, format_char_sequence="idddddddi" + ) + image_id = binary_image_properties[0] + qvec = np.array(binary_image_properties[1:5]) + tvec = np.array(binary_image_properties[5:8]) + camera_id = binary_image_properties[8] + binary_image_name = b"" + current_char = read_next_bytes(fid, 1, "c")[0] + while current_char != b"\x00": # look for the ASCII 0 entry + binary_image_name += current_char + current_char = read_next_bytes(fid, 1, "c")[0] + image_name = binary_image_name.decode("utf-8") + num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] + x_y_id_s = read_next_bytes( + fid, + num_bytes=24 * num_points2D, + format_char_sequence="ddq" * num_points2D, + ) + xys = np.column_stack( + [ + tuple(map(float, x_y_id_s[0::3])), + tuple(map(float, x_y_id_s[1::3])), + ] + ) + point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def write_images_text(images, path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesText(const std::string& path) + void Reconstruction::WriteImagesText(const std::string& path) + """ + if len(images) == 0: + mean_observations = 0 + else: + mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images) + HEADER = ( + "# Image list with two lines of data per image:\n" + + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + + "# POINTS2D[] as (X, Y, POINT3D_ID)\n" + + "# Number of images: {}, mean observations per image: {}\n".format( + len(images), mean_observations + ) + ) + + with open(path, "w") as fid: + fid.write(HEADER) + for _, img in images.items(): + image_header = [ + img.id, + *img.qvec, + *img.tvec, + img.camera_id, + img.name, + ] + first_line = " ".join(map(str, image_header)) + fid.write(first_line + "\n") + + points_strings = [] + for xy, point3D_id in zip(img.xys, img.point3D_ids): + points_strings.append(" ".join(map(str, [*xy, point3D_id]))) + fid.write(" ".join(points_strings) + "\n") + + +def write_images_binary(images, path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadImagesBinary(const std::string& path) + void Reconstruction::WriteImagesBinary(const std::string& path) + """ + with open(path_to_model_file, "wb") as fid: + write_next_bytes(fid, len(images), "Q") + for _, img in images.items(): + write_next_bytes(fid, img.id, "i") + write_next_bytes(fid, img.qvec.tolist(), "dddd") + write_next_bytes(fid, img.tvec.tolist(), "ddd") + write_next_bytes(fid, img.camera_id, "i") + for char in img.name: + write_next_bytes(fid, char.encode("utf-8"), "c") + write_next_bytes(fid, b"\x00", "c") + write_next_bytes(fid, len(img.point3D_ids), "Q") + for xy, p3d_id in zip(img.xys, img.point3D_ids): + write_next_bytes(fid, [*xy, p3d_id], "ddq") + + +def read_points3D_text(path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DText(const std::string& path) + void Reconstruction::WritePoints3DText(const std::string& path) + """ + points3D = {} + with open(path) as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + point3D_id = int(elems[0]) + xyz = np.array(tuple(map(float, elems[1:4]))) + rgb = np.array(tuple(map(int, elems[4:7]))) + error = float(elems[7]) + image_ids = np.array(tuple(map(int, elems[8::2]))) + point2D_idxs = np.array(tuple(map(int, elems[9::2]))) + points3D[point3D_id] = Point3D( + id=point3D_id, + xyz=xyz, + rgb=rgb, + error=error, + image_ids=image_ids, + point2D_idxs=point2D_idxs, + ) + return points3D + + +def read_points3D_binary(path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DBinary(const std::string& path) + void Reconstruction::WritePoints3DBinary(const std::string& path) + """ + points3D = {} + with open(path_to_model_file, "rb") as fid: + num_points = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_points): + binary_point_line_properties = read_next_bytes( + fid, num_bytes=43, format_char_sequence="QdddBBBd" + ) + point3D_id = binary_point_line_properties[0] + xyz = np.array(binary_point_line_properties[1:4]) + rgb = np.array(binary_point_line_properties[4:7]) + error = np.array(binary_point_line_properties[7]) + track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] + track_elems = read_next_bytes( + fid, + num_bytes=8 * track_length, + format_char_sequence="ii" * track_length, + ) + image_ids = np.array(tuple(map(int, track_elems[0::2]))) + point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) + points3D[point3D_id] = Point3D( + id=point3D_id, + xyz=xyz, + rgb=rgb, + error=error, + image_ids=image_ids, + point2D_idxs=point2D_idxs, + ) + return points3D + + +def write_points3D_text(points3D, path): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DText(const std::string& path) + void Reconstruction::WritePoints3DText(const std::string& path) + """ + if len(points3D) == 0: + mean_track_length = 0 + else: + mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items())) / len(points3D) + HEADER = ( + "# 3D point list with one line of data per point:\n" + + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" + + "# Number of points: {}, mean track length: {}\n".format( + len(points3D), mean_track_length + ) + ) + + with open(path, "w") as fid: + fid.write(HEADER) + for _, pt in points3D.items(): + point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error] + fid.write(" ".join(map(str, point_header)) + " ") + track_strings = [] + for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs): + track_strings.append(" ".join(map(str, [image_id, point2D]))) + fid.write(" ".join(track_strings) + "\n") + + +def write_points3D_binary(points3D, path_to_model_file): + """ + see: src/colmap/scene/reconstruction.cc + void Reconstruction::ReadPoints3DBinary(const std::string& path) + void Reconstruction::WritePoints3DBinary(const std::string& path) + """ + with open(path_to_model_file, "wb") as fid: + write_next_bytes(fid, len(points3D), "Q") + for _, pt in points3D.items(): + write_next_bytes(fid, pt.id, "Q") + write_next_bytes(fid, pt.xyz.tolist(), "ddd") + write_next_bytes(fid, pt.rgb.tolist(), "BBB") + write_next_bytes(fid, pt.error, "d") + track_length = pt.image_ids.shape[0] + write_next_bytes(fid, track_length, "Q") + for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs): + write_next_bytes(fid, [image_id, point2D_id], "ii") + + +def detect_model_format(path, ext): + if ( + os.path.isfile(os.path.join(path, "cameras" + ext)) + and os.path.isfile(os.path.join(path, "images" + ext)) + and os.path.isfile(os.path.join(path, "points3D" + ext)) + ): + print("Detected model format: '" + ext + "'") + return True + + return False + + +def read_model(path, ext=""): + # try to detect the extension automatically + if ext == "": + if detect_model_format(path, ".bin"): + ext = ".bin" + elif detect_model_format(path, ".txt"): + ext = ".txt" + else: + print("Provide model format: '.bin' or '.txt'") + return + + if ext == ".txt": + cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) + images = read_images_text(os.path.join(path, "images" + ext)) + points3D = read_points3D_text(os.path.join(path, "points3D") + ext) + else: + cameras = read_cameras_binary(os.path.join(path, "cameras" + ext)) + images = read_images_binary(os.path.join(path, "images" + ext)) + points3D = read_points3D_binary(os.path.join(path, "points3D") + ext) + return cameras, images, points3D + + +def write_model(cameras, images, points3D, path, ext=".bin"): + if ext == ".txt": + write_cameras_text(cameras, os.path.join(path, "cameras" + ext)) + write_images_text(images, os.path.join(path, "images" + ext)) + write_points3D_text(points3D, os.path.join(path, "points3D") + ext) + else: + write_cameras_binary(cameras, os.path.join(path, "cameras" + ext)) + write_images_binary(images, os.path.join(path, "images" + ext)) + write_points3D_binary(points3D, os.path.join(path, "points3D") + ext) + return cameras, images, points3D + + +def qvec2rotmat(qvec): + return np.array( + [ + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + ], + [ + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + ], + [ + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ], + ] + ) + + +def rotmat2qvec(R): + Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat + K = ( + np.array( + [ + [Rxx - Ryy - Rzz, 0, 0, 0], + [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], + [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], + [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz], + ] + ) + / 3.0 + ) + eigvals, eigvecs = np.linalg.eigh(K) + qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] + if qvec[0] < 0: + qvec *= -1 + return qvec + + +def main(): + parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models") + parser.add_argument("--input_model", help="path to input model folder") + parser.add_argument( + "--input_format", + choices=[".bin", ".txt"], + help="input model format", + default="", + ) + parser.add_argument("--output_model", help="path to output model folder") + parser.add_argument( + "--output_format", + choices=[".bin", ".txt"], + help="output model format", + default=".txt", + ) + args = parser.parse_args() + + cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format) + + print("num_cameras:", len(cameras)) + print("num_images:", len(images)) + print("num_points3D:", len(points3D)) + + if args.output_model is not None: + write_model( + cameras, + images, + points3D, + path=args.output_model, + ext=args.output_format, + ) + + +if __name__ == "__main__": + main() diff --git a/src/nn/depth_anything_3/utils/registry.py b/src/nn/depth_anything_3/utils/registry.py new file mode 100644 index 0000000..7db16d5 --- /dev/null +++ b/src/nn/depth_anything_3/utils/registry.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any +from addict import Dict + + +class Registry(Dict[str, Any]): + def __init__(self): + super().__init__() + self._map = Dict({}) + + def register(self, name=None): + def decorator(cls): + key = name or cls.__name__ + self._map[key] = cls + return cls + + return decorator + + def get(self, name): + return self._map[name] + + def all(self): + return self._map diff --git a/src/nn/depth_anything_3/utils/sh_helpers.py b/src/nn/depth_anything_3/utils/sh_helpers.py new file mode 100644 index 0000000..1a1a4ca --- /dev/null +++ b/src/nn/depth_anything_3/utils/sh_helpers.py @@ -0,0 +1,88 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from math import isqrt +import torch +from einops import einsum + +try: + from e3nn.o3 import matrix_to_angles, wigner_D +except ImportError: + from depth_anything_3.utils.logger import logger + + logger.warn("Dependency 'e3nn' not found. Required for rotating the camera space SH coeff") + + +def project_to_so3_strict(M: torch.Tensor) -> torch.Tensor: + if M.shape[-2:] != (3, 3): + raise ValueError("Input must be a batch of 3x3 matrices (i.e., shape [..., 3, 3]).") + + # 1. Compute SVD + U, S, Vh = torch.linalg.svd(M) + V = Vh.mH + + # 2. Handle reflection case (det = -1) + det_U = torch.det(U) + det_V = torch.det(V) + is_reflection = (det_U * det_V) < 0 + correction_sign = torch.where( + is_reflection[..., None], + torch.tensor([1, 1, -1.0], device=M.device, dtype=M.dtype), + torch.tensor([1, 1, 1.0], device=M.device, dtype=M.dtype), + ) + correction_matrix = torch.diag_embed(correction_sign) + U_corrected = U @ correction_matrix + R_so3_initial = U_corrected @ V.transpose(-2, -1) + + # 3. Explicitly ensure determinant is 1 (or extremely close) + current_det = torch.det(R_so3_initial) + det_correction_factor = torch.pow(current_det, -1 / 3)[..., None, None] + R_so3_final = R_so3_initial * det_correction_factor + + return R_so3_final + + +def rotate_sh( + sh_coefficients: torch.Tensor, # "*#batch n" + rotations: torch.Tensor, # "*#batch 3 3" +) -> torch.Tensor: # "*batch n" + # https://github.com/graphdeco-inria/gaussian-splatting/issues/176#issuecomment-2452412653 + device = sh_coefficients.device + dtype = sh_coefficients.dtype + + *_, n = sh_coefficients.shape + + with torch.autocast(device_type=rotations.device.type, enabled=False): + rotations_float32 = rotations.to(torch.float32) + + # switch axes: yzx -> xyz + P = torch.tensor([[0, 0, 1], [1, 0, 0], [0, 1, 0]]).unsqueeze(0).to(rotations_float32) + permuted_rotations = torch.linalg.inv(P) @ rotations_float32 @ P + + # ensure rotation has det == 1 in float32 type + permuted_rotations_so3 = project_to_so3_strict(permuted_rotations) + + alpha, beta, gamma = matrix_to_angles(permuted_rotations_so3) + result = [] + for degree in range(isqrt(n)): + with torch.device(device): + sh_rotations = wigner_D(degree, alpha, -beta, gamma).type(dtype) + sh_rotated = einsum( + sh_rotations, + sh_coefficients[..., degree**2 : (degree + 1) ** 2], + "... i j, ... j -> ... i", + ) + result.append(sh_rotated) + + return torch.cat(result, dim=-1) diff --git a/src/nn/depth_anything_3/utils/visualize.py b/src/nn/depth_anything_3/utils/visualize.py new file mode 100644 index 0000000..8fd32bd --- /dev/null +++ b/src/nn/depth_anything_3/utils/visualize.py @@ -0,0 +1,120 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import matplotlib +import numpy as np +import torch +from einops import rearrange + +from depth_anything_3.utils.logger import logger + + +def visualize_depth( + depth: np.ndarray, + depth_min=None, + depth_max=None, + percentile=2, + ret_minmax=False, + ret_type=np.uint8, + cmap="Spectral", +): + """ + Visualize a depth map using a colormap. + + Args: + depth: Input depth map array + depth_min: Minimum depth value for normalization. If None, uses percentile + depth_max: Maximum depth value for normalization. If None, uses percentile + percentile: Percentile for min/max computation if not provided + ret_minmax: Whether to return min/max depth values + ret_type: Return array type (uint8 or float) + cmap: Matplotlib colormap name to use + + Returns: + Colored depth visualization as numpy array + If ret_minmax=True, also returns depth_min and depth_max + """ + depth = depth.copy() + depth.copy() + valid_mask = depth > 0 + depth[valid_mask] = 1 / depth[valid_mask] + if depth_min is None: + if valid_mask.sum() <= 10: + depth_min = 0 + else: + depth_min = np.percentile(depth[valid_mask], percentile) + if depth_max is None: + if valid_mask.sum() <= 10: + depth_max = 0 + else: + depth_max = np.percentile(depth[valid_mask], 100 - percentile) + if depth_min == depth_max: + depth_min = depth_min - 1e-6 + depth_max = depth_max + 1e-6 + cm = matplotlib.colormaps[cmap] + depth = ((depth - depth_min) / (depth_max - depth_min)).clip(0, 1) + depth = 1 - depth + img_colored_np = cm(depth[None], bytes=False)[:, :, :, 0:3] # value from 0 to 1 + if ret_type == np.uint8: + img_colored_np = (img_colored_np[0] * 255.0).astype(np.uint8) + elif ret_type == np.float32 or ret_type == np.float64: + img_colored_np = img_colored_np[0] + else: + raise ValueError(f"Invalid return type: {ret_type}") + if ret_minmax: + return img_colored_np, depth_min, depth_max + else: + return img_colored_np + + +# GS video rendering visulization function, since it operates in Tensor space... + + +def vis_depth_map_tensor( + result: torch.Tensor, # "*batch height width" + color_map: str = "Spectral", +) -> torch.Tensor: # "*batch 3 height with" + """ + Color-map the depth map. + """ + far = result.reshape(-1)[:16_000_000].float().quantile(0.99).log().to(result) + try: + near = result[result > 0][:16_000_000].float().quantile(0.01).log().to(result) + except (RuntimeError, ValueError) as e: + logger.error(f"No valid depth values found. Reason: {e}") + near = torch.zeros_like(far) + result = result.log() + result = (result - near) / (far - near) + return apply_color_map_to_image(result, color_map) + + +def apply_color_map( + x: torch.Tensor, # " *batch" + color_map: str = "inferno", +) -> torch.Tensor: # "*batch 3" + cmap = matplotlib.cm.get_cmap(color_map) + + # Convert to NumPy so that Matplotlib color maps can be used. + mapped = cmap(x.float().detach().clip(min=0, max=1).cpu().numpy())[..., :3] + + # Convert back to the original format. + return torch.tensor(mapped, device=x.device, dtype=torch.float32) + + +def apply_color_map_to_image( + image: torch.Tensor, # "*batch height width" + color_map: str = "inferno", +) -> torch.Tensor: # "*batch 3 height with" + image = apply_color_map(image, color_map) + return rearrange(image, "... h w c -> ... c h w") diff --git a/src/nn/segearth_ov3/pamr.py b/src/nn/segearth_ov3/pamr.py new file mode 100644 index 0000000..2e637d1 --- /dev/null +++ b/src/nn/segearth_ov3/pamr.py @@ -0,0 +1,135 @@ +# Copyright 2020 TU Darmstadt +# Licnese: Apache 2.0 License. +# https://github.com/visinf/1-stage-wseg/blob/master/models/mods/pamr.py +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# Helper modules +class LocalAffinity(nn.Module): + def __init__(self, dilations=[1]): + super(LocalAffinity, self).__init__() + self.dilations = dilations + weight = self._init_aff() + self.register_buffer('kernel', weight) + + def _init_aff(self): + # initialising the shift kernel + weight = torch.zeros(8, 1, 3, 3) + + for i in range(weight.size(0)): + weight[i, 0, 1, 1] = 1 + + weight[0, 0, 0, 0] = -1 + weight[1, 0, 0, 1] = -1 + weight[2, 0, 0, 2] = -1 + + weight[3, 0, 1, 0] = -1 + weight[4, 0, 1, 2] = -1 + + weight[5, 0, 2, 0] = -1 + weight[6, 0, 2, 1] = -1 + weight[7, 0, 2, 2] = -1 + + self.weight_check = weight.clone() + return weight + + def forward(self, x): + # self.weight_check = self.weight_check.type_as(x) + # assert torch.all(self.weight_check.eq(self.kernel)) + + B, K, H, W = x.size() + x = x.view(B * K, 1, H, W) + + x_aff = torch.empty((K * B, len(self.dilations) * self.kernel.size(0), H, W), device=x.device) + for i, d in enumerate(self.dilations): + x_pad = F.pad(x, [d] * 4, mode='replicate') + x_aff[:, self.kernel.size(0) * i:(i + 1) * self.kernel.size(0), :, :] = F.conv2d(x_pad, self.kernel, + dilation=d) + return x_aff.view(B, K, -1, H, W) + + +class LocalAffinityCopy(LocalAffinity): + def _init_aff(self): + # initialising the shift kernel + weight = torch.zeros(8, 1, 3, 3) + + weight[0, 0, 0, 0] = 1 + weight[1, 0, 0, 1] = 1 + weight[2, 0, 0, 2] = 1 + + weight[3, 0, 1, 0] = 1 + weight[4, 0, 1, 2] = 1 + + weight[5, 0, 2, 0] = 1 + weight[6, 0, 2, 1] = 1 + weight[7, 0, 2, 2] = 1 + + self.weight_check = weight.clone() + return weight + + +class LocalStDev(LocalAffinity): + + def _init_aff(self): + weight = torch.zeros(9, 1, 3, 3) + weight.zero_() + + weight[0, 0, 0, 0] = 1 + weight[1, 0, 0, 1] = 1 + weight[2, 0, 0, 2] = 1 + + weight[3, 0, 1, 0] = 1 + weight[4, 0, 1, 1] = 1 + weight[5, 0, 1, 2] = 1 + + weight[6, 0, 2, 0] = 1 + weight[7, 0, 2, 1] = 1 + weight[8, 0, 2, 2] = 1 + + self.weight_check = weight.clone() + return weight + + def forward(self, x): + # returns (B,K,P,H,W), where P is the number of locations + x = super(LocalStDev, self).forward(x) + return x.std(2, keepdim=True) + + +class LocalAffinityAbs(LocalAffinity): + def forward(self, x): + x = super(LocalAffinityAbs, self).forward(x) + return torch.abs(x) + + +# PAMR module +class PAMR(nn.Module): + def __init__(self, num_iter=1, dilations=[1]): + super(PAMR, self).__init__() + + self.num_iter = num_iter + self.aff_x = LocalAffinityAbs(dilations) + self.aff_m = LocalAffinityCopy(dilations) + self.aff_std = LocalStDev(dilations) + + def forward(self, x, mask): + mask = F.interpolate(mask, size=x.size()[-2:], mode="bilinear", align_corners=True) + + # x: [BxKxHxW] + # mask: [BxCxHxW] + B, K, H, W = x.size() + _, C, _, _ = mask.size() + + x_std = self.aff_std(x) + + x = -self.aff_x(x) / (1e-8 + 0.1 * x_std) + x = x.mean(1, keepdim=True) + x = F.softmax(x, 2) + + for _ in range(self.num_iter): + m = self.aff_m(mask) # [BxCxPxHxW] + mask = (m * x).sum(2) + + # xvals: [BxCxHxW] + return mask diff --git a/src/nn/segearth_ov3/sam3/__init__.py b/src/nn/segearth_ov3/sam3/__init__.py new file mode 100644 index 0000000..14270a6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from .model_builder import build_sam3_image_model + +__version__ = "0.1.0" + +__all__ = ["build_sam3_image_model"] diff --git a/src/nn/segearth_ov3/sam3/agent/__init__.py b/src/nn/segearth_ov3/sam3/agent/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/agent/agent_core.py b/src/nn/segearth_ov3/sam3/agent/agent_core.py new file mode 100644 index 0000000..f0016c7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/agent_core.py @@ -0,0 +1,563 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import copy +import json +import os + +import cv2 +from PIL import Image + +from .client_llm import send_generate_request +from .client_sam3 import call_sam_service +from .viz import visualize + + +def save_debug_messages(messages_list, debug, debug_folder_path, debug_jsonl_path): + """Save messages to debug jsonl file if debug is enabled""" + if debug and debug_jsonl_path: + # Ensure the debug directory exists before writing + os.makedirs(debug_folder_path, exist_ok=True) + with open(debug_jsonl_path, "w") as f: + for msg in messages_list: + f.write(json.dumps(msg, indent=4) + "\n") + + +def cleanup_debug_files(debug, debug_folder_path, debug_jsonl_path): + """Clean up debug files when function successfully returns""" + if debug and debug_folder_path: + try: + if os.path.exists(debug_jsonl_path): + os.remove(debug_jsonl_path) + if os.path.exists(debug_folder_path): + os.rmdir(debug_folder_path) + except Exception as e: + print(f"Warning: Could not clean up debug files: {e}") + + +def count_images(messages): + """Count the total number of images present in the messages history.""" + total = 0 + for message in messages: + # Check if message has content (should be a list) + if "content" in message and isinstance(message["content"], list): + # Iterate through each content item + for content_item in message["content"]: + # Check if content item is a dict with type "image" + if ( + isinstance(content_item, dict) + and content_item.get("type") == "image" + ): + total += 1 + return total + + +def _prune_messages_for_next_round( + messages_list, + used_text_prompts, + latest_sam3_text_prompt, + img_path, + initial_text_prompt, +): + """Return a new messages list that contains only: + 1) messages[:2] (with optional warning text added to the second message's content) + 2) the latest assistant message (and everything after it) that contains a segment_phrase tool call + """ + # There should not be more than 10 messages in the conversation history + assert len(messages_list) < 10 + + # Part 1: always keep the first two message JSONs + part1 = copy.deepcopy(messages_list[:2]) + + # Part 2: search backwards for the latest assistant message containing a segment_phrase tool call + part2_start_idx = None + for idx in range(len(messages_list) - 1, 1, -1): + msg = messages_list[idx] + # We only consider assistant messages with a "content" list + if msg.get("role") != "assistant" or "content" not in msg: + continue + # Look for any content element that is a text containing the segment_phrase tool call + for content in msg["content"]: + if ( + isinstance(content, dict) + and content.get("type") == "text" + and "" in content.get("text", "") + and "segment_phrase" in content.get("text", "") + ): + part2_start_idx = idx + break + if part2_start_idx is not None: + break + + part2 = messages_list[part2_start_idx:] if part2_start_idx is not None else [] + + # Part 3: decide whether to add warning text to the second message in part1 + previously_used = ( + [p for p in used_text_prompts if p != latest_sam3_text_prompt] + if latest_sam3_text_prompt + else list(used_text_prompts) + ) + if part2 and len(previously_used) > 0: + warning_text = f'Note that we have previously called the segment_phrase tool with each "text_prompt" in this list: {list(previously_used)}, but none of the generated results were satisfactory. So make sure that you do not use any of these phrases as the "text_prompt" to call the segment_phrase tool again.' + # Replace the second message entirely to keep exactly 2 content items + part1[1] = { + "role": "user", + "content": [ + {"type": "image", "image": img_path}, + { + "type": "text", + "text": f"The above image is the raw input image. The initial user input query is: '{initial_text_prompt}'." + + " " + + warning_text, + }, + ], + } + assert len(part1[1]["content"]) == 2 + + # Build the new messages list: part1 (with optional warning), then part2 + new_messages = list(part1) + new_messages.extend(part2) + return new_messages + + +def agent_inference( + img_path: str, + initial_text_prompt: str, + debug: bool = False, + send_generate_request=send_generate_request, + call_sam_service=call_sam_service, + max_generations: int = 100, + output_dir="../../sam3_agent_out", +): + """ + Given a text prompt and an image, this tool will perform all aspects of agentic problem solving, + while saving sam3 and MLLM outputs to their respective directories. + + Args: + img_path: Path to the input image + initial_text_prompt: Initial text prompt from the user + debug: Whether to enable debug mode + max_generations: Maximum number of send_generate_request calls allowed (default: 100) + """ + # setup dir + sam_output_dir = os.path.join(output_dir, "sam_out") + error_save_dir = os.path.join(output_dir, "none_out") + debug_save_dir = os.path.join(output_dir, "agent_debug_out") + os.makedirs(sam_output_dir, exist_ok=True) + os.makedirs(error_save_dir, exist_ok=True) + os.makedirs(debug_save_dir, exist_ok=True) + current_dir = os.path.dirname(os.path.abspath(__file__)) + MLLM_SYSTEM_PROMPT_PATH = os.path.join( + current_dir, "system_prompts/system_prompt.txt" + ) + ITERATIVE_CHECKING_SYSTEM_PROMPT_PATH = os.path.join( + current_dir, "system_prompts/system_prompt_iterative_checking.txt" + ) + # init variables + PATH_TO_LATEST_OUTPUT_JSON = "" + LATEST_SAM3_TEXT_PROMPT = "" + USED_TEXT_PROMPTS = ( + set() + ) # Track all previously used text prompts for segment_phrase + generation_count = 0 # Counter for number of send_generate_request calls + + # debug setup + debug_folder_path = None + debug_jsonl_path = None + if debug: + debug_folder_path = os.path.join( + debug_save_dir, f"{img_path.rsplit('/', 1)[-1].rsplit('.', 1)[0]}" + ) + debug_jsonl_path = os.path.join(debug_folder_path, "debug_history.json") + os.makedirs(debug_folder_path, exist_ok=True) + + # The helper functions are now defined outside the agent_inference function + with open(MLLM_SYSTEM_PROMPT_PATH, "r") as f: + system_prompt = f.read().strip() + with open(ITERATIVE_CHECKING_SYSTEM_PROMPT_PATH, "r") as f: + iterative_checking_system_prompt = f.read().strip() + + # Construct the initial message list + messages = [ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": [ + {"type": "image", "image": img_path}, + { + "type": "text", + "text": f"The above image is the raw input image. The initial user input query is: '{initial_text_prompt}'.", + }, + ], + }, + ] + print(f"> Text prompt: {initial_text_prompt}") + print(f"> Image path: {img_path}") + + print("\n\n") + print("-" * 30 + f" Round {str(generation_count + 1)}" + "-" * 30) + print("\n\n") + generated_text = send_generate_request(messages) + print(f"\n>>> MLLM Response [start]\n{generated_text}\n<<< MLLM Response [end]\n") + while generated_text is not None: + save_debug_messages(messages, debug, debug_folder_path, debug_jsonl_path) + assert ( + "" in generated_text, + f"Generated text does not contain tag: {generated_text}", + ) + generated_text = generated_text.split("", 1)[0] + "" + tool_call_json_str = ( + generated_text.split("")[-1] + .split("")[0] + .strip() + .replace(r"}}}", r"}}") # remove extra } if any + ) + try: + tool_call = json.loads(tool_call_json_str) + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON in tool call: {tool_call_json_str}") + + if PATH_TO_LATEST_OUTPUT_JSON == "": + # The first tool call must be segment_phrase or report_no_mask + assert ( + tool_call["name"] == "segment_phrase" + or tool_call["name"] == "report_no_mask" + ) + + if tool_call["name"] == "segment_phrase": + print("🔍 Calling segment_phrase tool...") + assert list(tool_call["parameters"].keys()) == ["text_prompt"] + + # Check if this text_prompt has been used before + current_text_prompt = tool_call["parameters"]["text_prompt"] + if current_text_prompt in USED_TEXT_PROMPTS: + print( + f"❌ Text prompt '{current_text_prompt}' has been used before. Requesting a different prompt." + ) + duplicate_prompt_message = f"You have previously used '{current_text_prompt}' as your text_prompt to call the segment_phrase tool. You may not use it again. Please call the segment_phrase tool again with a different, perhaps more general, or more creative simple noun phrase prompt, while adhering to all the rules stated in the system prompt. You must also never use any of the following text_prompt(s): {str(list(USED_TEXT_PROMPTS))}." + messages.append( + { + "role": "assistant", + "content": [{"type": "text", "text": generated_text}], + } + ) + messages.append( + { + "role": "user", + "content": [{"type": "text", "text": duplicate_prompt_message}], + } + ) + else: + # Add the text_prompt to the set of used prompts + USED_TEXT_PROMPTS.add(current_text_prompt) + LATEST_SAM3_TEXT_PROMPT = current_text_prompt + PATH_TO_LATEST_OUTPUT_JSON = call_sam_service( + image_path=img_path, + text_prompt=current_text_prompt, + output_folder_path=sam_output_dir, + ) + sam3_outputs = json.load(open(PATH_TO_LATEST_OUTPUT_JSON, "r")) + sam3_output_image_path = sam3_outputs["output_image_path"] + num_masks = len(sam3_outputs["pred_boxes"]) + + messages.append( + { + "role": "assistant", + "content": [{"type": "text", "text": generated_text}], + } + ) + if num_masks == 0: + print("❌ No masks generated by SAM3, reporting no mask to Qwen.") + sam3_output_text_message = f"The segment_phrase tool did not generate any masks for the text_prompt '{current_text_prompt}'. Now, please call the segment_phrase tool again with a different, perhaps more general, or more creative simple noun phrase text_prompt, while adhering to all the rules stated in the system prompt. Please be reminded that the original user query was '{initial_text_prompt}'." + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": sam3_output_text_message} + ], + } + ) + else: + sam3_output_text_message = rf"The segment_phrase tool generated {num_masks} available masks. All {num_masks} available masks are rendered in this image below, now you must analyze the {num_masks} available mask(s) carefully, compare them against the raw input image and the original user query, and determine your next action. Please be reminded that the original user query was '{initial_text_prompt}'." + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": sam3_output_text_message}, + {"type": "image", "image": sam3_output_image_path}, + ], + } + ) + print("\n\n>>> sam3_output_text_message:\n", sam3_output_text_message) + + elif tool_call["name"] == "examine_each_mask": + print("🔍 Calling examine_each_mask tool...") + assert LATEST_SAM3_TEXT_PROMPT != "" + + # Make sure that the last message is a image + assert ( + messages[-1]["content"][1]["type"] == "image" + ), "Second content element should be an image" + messages.pop() # Remove the last user message + # Add simplified replacement message + simplified_message = { + "role": "user", + "content": [ + { + "type": "text", + "text": "The segment_phrase tool generated several masks. Now you must analyze the mask(s) carefully, compare them against the raw input image and the original user query, and determine your next action.", + } + ], + } + messages.append(simplified_message) + + current_outputs = json.load(open(PATH_TO_LATEST_OUTPUT_JSON, "r")) + num_masks = len(current_outputs["pred_masks"]) + masks_to_keep = [] + + # MLLM check the mask one by one + for i in range(num_masks): + print(f"🔍 Checking mask {i+1}/{num_masks}...") + image_w_mask_i, image_w_zoomed_in_mask_i = visualize(current_outputs, i) + + image_w_zoomed_in_mask_i_path = os.path.join( + sam_output_dir, rf"{LATEST_SAM3_TEXT_PROMPT}.png".replace("/", "_") + ).replace(".png", f"_zoom_in_mask_{i + 1}.png") + image_w_mask_i_path = os.path.join( + sam_output_dir, rf"{LATEST_SAM3_TEXT_PROMPT}.png".replace("/", "_") + ).replace(".png", f"_selected_mask_{i + 1}.png") + image_w_zoomed_in_mask_i.save(image_w_zoomed_in_mask_i_path) + image_w_mask_i.save(image_w_mask_i_path) + + iterative_checking_messages = [ + {"role": "system", "content": iterative_checking_system_prompt}, + { + "role": "user", + "content": [ + {"type": "text", "text": f"The raw input image: "}, + {"type": "image", "image": img_path}, + { + "type": "text", + "text": f"The initial user input query is: '{initial_text_prompt}'", + }, + { + "type": "text", + "text": f"Image with the predicted segmentation mask rendered on it: ", + }, + {"type": "image", "image": image_w_mask_i_path}, + { + "type": "text", + "text": f"Image with the zoomed-in mask: ", + }, + {"type": "image", "image": image_w_zoomed_in_mask_i_path}, + ], + }, + ] + checking_generated_text = send_generate_request( + iterative_checking_messages + ) + + # Process the generated text to determine if the mask should be kept or rejected + if checking_generated_text is None: + raise ValueError( + "Generated text is None, which is unexpected. Please check the Qwen server and the input parameters." + ) + print(f"Generated text for mask {i+1}: {checking_generated_text}") + verdict = ( + checking_generated_text.split("")[-1] + .split("")[0] + .strip() + ) + if "Accept" in verdict: + assert not "Reject" in verdict + print(f"Mask {i+1} accepted, keeping it in the outputs.") + masks_to_keep.append(i) + elif "Reject" in verdict: + assert not "Accept" in verdict + print(f"Mask {i+1} rejected, removing it from the outputs.") + else: + raise ValueError( + f"Unexpected verdict in generated text: {checking_generated_text}. Expected 'Accept' or 'Reject'." + ) + + updated_outputs = { + "original_image_path": current_outputs["original_image_path"], + "orig_img_h": current_outputs["orig_img_h"], + "orig_img_w": current_outputs["orig_img_w"], + "pred_boxes": [current_outputs["pred_boxes"][i] for i in masks_to_keep], + "pred_scores": [ + current_outputs["pred_scores"][i] for i in masks_to_keep + ], + "pred_masks": [current_outputs["pred_masks"][i] for i in masks_to_keep], + } + + image_w_check_masks = visualize(updated_outputs) + image_w_check_masks_path = os.path.join( + sam_output_dir, rf"{LATEST_SAM3_TEXT_PROMPT}.png" + ).replace( + ".png", + f"_selected_masks_{'-'.join(map(str, [i+1 for i in masks_to_keep]))}.png".replace( + "/", "_" + ), + ) + image_w_check_masks.save(image_w_check_masks_path) + # save the updated json outputs and append to message history + messages.append( + { + "role": "assistant", + "content": [{"type": "text", "text": generated_text}], + } + ) + if len(masks_to_keep) == 0: + messages.append( + { + "role": "user", + "content": [ + { + "type": "text", + "text": f"The original user query was: '{initial_text_prompt}'. The examine_each_mask tool examined and rejected all of the masks generated by the segment_phrase tool. Now, please call the segment_phrase tool again with a different, perhaps more general, or more creative simple noun phrase text_prompt, while adhering to all the rules stated in the system prompt.", + } + ], + } + ) + else: + messages.append( + { + "role": "user", + "content": [ + { + "type": "text", + "text": f"The original user query was: '{initial_text_prompt}'. After calling the examine_each_mask tool on the available masks, the number of available masks is now {len(masks_to_keep)}. All {len(masks_to_keep)} available masks are rendered in this image below, now you must analyze the {len(masks_to_keep)} available mask(s) carefully, compare them against the raw input image and the original user query, and determine your next action.", + }, + {"type": "image", "image": image_w_check_masks_path}, + ], + } + ) + + # Create a new filename based on the original path to avoid filename length issues + base_path = PATH_TO_LATEST_OUTPUT_JSON + # Remove any existing "masks_" suffix to avoid duplication + if "masks_" in base_path: + base_path = base_path.split("masks_")[0] + ".json" + # Create new filename with current masks; use a clearer suffix when empty + if len(masks_to_keep) == 0: + PATH_TO_LATEST_OUTPUT_JSON = base_path.replace( + ".json", "masks_none.json" + ) + else: + PATH_TO_LATEST_OUTPUT_JSON = base_path.replace( + ".json", f"masks_{'_'.join(map(str, masks_to_keep))}.json" + ) + json.dump(updated_outputs, open(PATH_TO_LATEST_OUTPUT_JSON, "w"), indent=4) + + elif tool_call["name"] == "select_masks_and_return": + print("🔍 Calling select_masks_and_return tool...") + current_outputs = json.load(open(PATH_TO_LATEST_OUTPUT_JSON, "r")) + + assert list(tool_call["parameters"].keys()) == ["final_answer_masks"] + masks_to_keep = tool_call["parameters"]["final_answer_masks"] + + # Keep only valid mask indices, remove duplicates, and preserve deterministic ascending order + available_masks = set(range(1, len(current_outputs["pred_masks"]) + 1)) + masks_to_keep = sorted({i for i in masks_to_keep if i in available_masks}) + # Change this to a update message telling the model to try again along with information about errors made. + + final_outputs = { + "original_image_path": current_outputs["original_image_path"], + "orig_img_h": current_outputs["orig_img_h"], + "orig_img_w": current_outputs["orig_img_w"], + "pred_boxes": [ + current_outputs["pred_boxes"][i - 1] for i in masks_to_keep + ], + "pred_scores": [ + current_outputs["pred_scores"][i - 1] for i in masks_to_keep + ], + "pred_masks": [ + current_outputs["pred_masks"][i - 1] for i in masks_to_keep + ], + } + + rendered_final_output = visualize(final_outputs) + messages.append( + { + "role": "assistant", + "content": [{"type": "text", "text": generated_text}], + } + ) + + # Clean up debug files before successful return + cleanup_debug_files(debug, debug_folder_path, debug_jsonl_path) + return messages, final_outputs, rendered_final_output + + elif tool_call["name"] == "report_no_mask": + print("🔍 Calling report_no_mask tool...") + height, width = cv2.imread(img_path).shape[:2] + final_outputs = { + "original_image_path": img_path, + "orig_img_h": height, + "orig_img_w": width, + "pred_boxes": [], + "pred_scores": [], + "pred_masks": [], + } + rendered_final_output = Image.open(img_path) + messages.append( + { + "role": "assistant", + "content": [{"type": "text", "text": generated_text}], + } + ) + return messages, final_outputs, rendered_final_output + + else: + raise ValueError(f"Unknown tool call: {tool_call['name']}") + + # sometimes the MLLM don't know when to stop, and generates multiple tool calls in one round, so we need to split the generated text by and only keep the first one + + for message in messages: + if message["role"] == "assistant" and "content" in message: + for content in message["content"]: + if ( + isinstance(content, dict) + and content.get("type") == "text" + and "text" in content + ): + content["text"] = ( + content["text"].split("", 1)[0] + "\n\n" + ) + # Prune the messages history before the next MLLM generation round according to the 3-part rules. + # This keeps history compact and ensures the model sees only the allowed parts. + messages = _prune_messages_for_next_round( + messages, + USED_TEXT_PROMPTS, + LATEST_SAM3_TEXT_PROMPT, + img_path, + initial_text_prompt, + ) + # make sure there can never be more than 2 images in the context + assert count_images(messages) <= 2 + generation_count += 1 + if generation_count > max_generations: + raise ValueError( + f"Exceeded maximum number of allowed generation requests ({max_generations})" + ) + + print("\n\n") + print("-" * 30 + f" Round {str(generation_count + 1)}" + "-" * 30) + print("\n\n") + generated_text = send_generate_request(messages) + print( + f"\n>>> MLLM Response [start]\n{generated_text}\n<<< MLLM Response [end]\n" + ) + + print("\n\n>>> SAM 3 Agent execution ended.\n\n") + + error_save_path = os.path.join( + error_save_dir, + f"{img_path.rsplit('/', 1)[-1].rsplit('.', 1)[0]}_error_history.json", + ) + with open(error_save_path, "w") as f: + json.dump(messages, f, indent=4) + print("Saved messages history that caused error to:", error_save_path) + raise ValueError( + rf"Generated text is None, which is unexpected. Please check the Qwen server and the input parameters for image path: {img_path} and initial text prompt: {initial_text_prompt}." + ) diff --git a/src/nn/segearth_ov3/sam3/agent/client_llm.py b/src/nn/segearth_ov3/sam3/agent/client_llm.py new file mode 100644 index 0000000..85e513c --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/client_llm.py @@ -0,0 +1,205 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import base64 +import os +from typing import Any, Optional + +from openai import OpenAI + + +def get_image_base64_and_mime(image_path): + """Convert image file to base64 string and get MIME type""" + try: + # Get MIME type based on file extension + ext = os.path.splitext(image_path)[1].lower() + mime_types = { + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", + } + mime_type = mime_types.get(ext, "image/jpeg") # Default to JPEG + + # Convert image to base64 + with open(image_path, "rb") as image_file: + base64_data = base64.b64encode(image_file.read()).decode("utf-8") + return base64_data, mime_type + except Exception as e: + print(f"Error converting image to base64: {e}") + return None, None + + +def send_generate_request( + messages, + server_url=None, + model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", + api_key=None, + max_tokens=4096, +): + """ + Sends a request to the OpenAI-compatible API endpoint using the OpenAI client library. + + Args: + server_url (str): The base URL of the server, e.g. "http://127.0.0.1:8000" + messages (list): A list of message dicts, each containing role and content. + model (str): The model to use for generation (default: "llama-4") + max_tokens (int): Maximum number of tokens to generate (default: 4096) + + Returns: + str: The generated response text from the server. + """ + # Process messages to convert image paths to base64 + processed_messages = [] + for message in messages: + processed_message = message.copy() + if message["role"] == "user" and "content" in message: + processed_content = [] + for c in message["content"]: + if isinstance(c, dict) and c.get("type") == "image": + # Convert image path to base64 format + image_path = c["image"] + + print("image_path", image_path) + new_image_path = image_path.replace( + "?", "%3F" + ) # Escape ? in the path + + # Read the image file and convert to base64 + try: + base64_image, mime_type = get_image_base64_and_mime( + new_image_path + ) + if base64_image is None: + print( + f"Warning: Could not convert image to base64: {new_image_path}" + ) + continue + + # Create the proper image_url structure with base64 data + processed_content.append( + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{base64_image}", + "detail": "high", + }, + } + ) + + except FileNotFoundError: + print(f"Warning: Image file not found: {new_image_path}") + continue + except Exception as e: + print(f"Warning: Error processing image {new_image_path}: {e}") + continue + else: + processed_content.append(c) + + processed_message["content"] = processed_content + processed_messages.append(processed_message) + + # Create OpenAI client with custom base URL + client = OpenAI(api_key=api_key, base_url=server_url) + + try: + print(f"🔍 Calling model {model}...") + response = client.chat.completions.create( + model=model, + messages=processed_messages, + max_completion_tokens=max_tokens, + n=1, + ) + # print(f"Received response: {response.choices[0].message}") + + # Extract the response content + if response.choices and len(response.choices) > 0: + return response.choices[0].message.content + else: + print(f"Unexpected response format: {response}") + return None + + except Exception as e: + print(f"Request failed: {e}") + return None + + +def send_direct_request( + llm: Any, + messages: list[dict[str, Any]], + sampling_params: Any, +) -> Optional[str]: + """ + Run inference on a vLLM model instance directly without using a server. + + Args: + llm: Initialized vLLM LLM instance (passed from external initialization) + messages: List of message dicts with role and content (OpenAI format) + sampling_params: vLLM SamplingParams instance (initialized externally) + + Returns: + str: Generated response text, or None if inference fails + """ + try: + # Process messages to handle images (convert to base64 if needed) + processed_messages = [] + for message in messages: + processed_message = message.copy() + if message["role"] == "user" and "content" in message: + processed_content = [] + for c in message["content"]: + if isinstance(c, dict) and c.get("type") == "image": + # Convert image path to base64 format + image_path = c["image"] + new_image_path = image_path.replace("?", "%3F") + + try: + base64_image, mime_type = get_image_base64_and_mime( + new_image_path + ) + if base64_image is None: + print( + f"Warning: Could not convert image: {new_image_path}" + ) + continue + + # vLLM expects image_url format + processed_content.append( + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{base64_image}" + }, + } + ) + except Exception as e: + print( + f"Warning: Error processing image {new_image_path}: {e}" + ) + continue + else: + processed_content.append(c) + + processed_message["content"] = processed_content + processed_messages.append(processed_message) + + print("🔍 Running direct inference with vLLM...") + + # Run inference using vLLM's chat interface + outputs = llm.chat( + messages=processed_messages, + sampling_params=sampling_params, + ) + + # Extract the generated text from the first output + if outputs and len(outputs) > 0: + generated_text = outputs[0].outputs[0].text + return generated_text + else: + print(f"Unexpected output format: {outputs}") + return None + + except Exception as e: + print(f"Direct inference failed: {e}") + return None diff --git a/src/nn/segearth_ov3/sam3/agent/client_sam3.py b/src/nn/segearth_ov3/sam3/agent/client_sam3.py new file mode 100755 index 0000000..d2f64b7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/client_sam3.py @@ -0,0 +1,138 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import json +import os + +import torch +from PIL import Image + +from sam3.model.box_ops import box_xyxy_to_xywh +from sam3.train.masks_ops import rle_encode + +from .helpers.mask_overlap_removal import remove_overlapping_masks +from .viz import visualize + + +def sam3_inference(processor, image_path, text_prompt): + """Run SAM 3 image inference with text prompts and format the outputs""" + image = Image.open(image_path) + orig_img_w, orig_img_h = image.size + + # model inference + inference_state = processor.set_image(image) + inference_state = processor.set_text_prompt( + state=inference_state, prompt=text_prompt + ) + + # format and assemble outputs + pred_boxes_xyxy = torch.stack( + [ + inference_state["boxes"][:, 0] / orig_img_w, + inference_state["boxes"][:, 1] / orig_img_h, + inference_state["boxes"][:, 2] / orig_img_w, + inference_state["boxes"][:, 3] / orig_img_h, + ], + dim=-1, + ) # normalized in range [0, 1] + pred_boxes_xywh = box_xyxy_to_xywh(pred_boxes_xyxy).tolist() + pred_masks = rle_encode(inference_state["masks"].squeeze(1)) + pred_masks = [m["counts"] for m in pred_masks] + outputs = { + "orig_img_h": orig_img_h, + "orig_img_w": orig_img_w, + "pred_boxes": pred_boxes_xywh, + "pred_masks": pred_masks, + "pred_scores": inference_state["scores"].tolist(), + } + return outputs + + +def call_sam_service( + sam3_processor, + image_path: str, + text_prompt: str, + output_folder_path: str = "sam3_output", +): + """ + Loads an image, sends it with a text prompt to the service, + saves the results, and renders the visualization. + """ + print(f"📞 Loading image '{image_path}' and sending with prompt '{text_prompt}'...") + + text_prompt_for_save_path = ( + text_prompt.replace("/", "_") if "/" in text_prompt else text_prompt + ) + + os.makedirs( + os.path.join(output_folder_path, image_path.replace("/", "-")), exist_ok=True + ) + output_json_path = os.path.join( + output_folder_path, + image_path.replace("/", "-"), + rf"{text_prompt_for_save_path}.json", + ) + output_image_path = os.path.join( + output_folder_path, + image_path.replace("/", "-"), + rf"{text_prompt_for_save_path}.png", + ) + + try: + # Send the image and text prompt as a multipart/form-data request + serialized_response = sam3_inference(sam3_processor, image_path, text_prompt) + + # 1. Prepare the response dictionary + serialized_response = remove_overlapping_masks(serialized_response) + serialized_response = { + "original_image_path": image_path, + "output_image_path": output_image_path, + **serialized_response, + } + + # 2. Reorder predictions by scores (highest to lowest) if scores are available + if "pred_scores" in serialized_response and serialized_response["pred_scores"]: + # Create indices sorted by scores in descending order + score_indices = sorted( + range(len(serialized_response["pred_scores"])), + key=lambda i: serialized_response["pred_scores"][i], + reverse=True, + ) + + # Reorder all three lists based on the sorted indices + serialized_response["pred_scores"] = [ + serialized_response["pred_scores"][i] for i in score_indices + ] + serialized_response["pred_boxes"] = [ + serialized_response["pred_boxes"][i] for i in score_indices + ] + serialized_response["pred_masks"] = [ + serialized_response["pred_masks"][i] for i in score_indices + ] + + # 3. Remove any invalid RLE masks that is too short (shorter than 5 characters) + valid_masks = [] + valid_boxes = [] + valid_scores = [] + for i, rle in enumerate(serialized_response["pred_masks"]): + if len(rle) > 4: + valid_masks.append(rle) + valid_boxes.append(serialized_response["pred_boxes"][i]) + valid_scores.append(serialized_response["pred_scores"][i]) + serialized_response["pred_masks"] = valid_masks + serialized_response["pred_boxes"] = valid_boxes + serialized_response["pred_scores"] = valid_scores + + with open(output_json_path, "w") as f: + json.dump(serialized_response, f, indent=4) + print(f"✅ Raw JSON response saved to '{output_json_path}'") + + # 4. Render and save visualizations on the image and save it in the SAM3 output folder + print("🔍 Rendering visualizations on the image ...") + viz_image = visualize(serialized_response) + os.makedirs(os.path.dirname(output_image_path), exist_ok=True) + viz_image.save(output_image_path) + print("✅ Saved visualization at:", output_image_path) + except Exception as e: + print(f"❌ Error calling service: {e}") + + return output_json_path diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/__init__.py b/src/nn/segearth_ov3/sam3/agent/helpers/__init__.py new file mode 100755 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/boxes.py b/src/nn/segearth_ov3/sam3/agent/helpers/boxes.py new file mode 100755 index 0000000..6cfef58 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/boxes.py @@ -0,0 +1,438 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math +from enum import IntEnum, unique +from typing import List, Tuple, Union + +import numpy as np +import torch +from torch import device + +_RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray] + + +@unique +class BoxMode(IntEnum): + """ + Enum of different ways to represent a box. + """ + + XYXY_ABS = 0 + """ + (x0, y0, x1, y1) in absolute floating points coordinates. + The coordinates in range [0, width or height]. + """ + XYWH_ABS = 1 + """ + (x0, y0, w, h) in absolute floating points coordinates. + """ + XYXY_REL = 2 + """ + Not yet supported! + (x0, y0, x1, y1) in range [0, 1]. They are relative to the size of the image. + """ + XYWH_REL = 3 + """ + Not yet supported! + (x0, y0, w, h) in range [0, 1]. They are relative to the size of the image. + """ + XYWHA_ABS = 4 + """ + (xc, yc, w, h, a) in absolute floating points coordinates. + (xc, yc) is the center of the rotated box, and the angle a is in degrees ccw. + """ + + @staticmethod + def convert( + box: _RawBoxType, from_mode: "BoxMode", to_mode: "BoxMode" + ) -> _RawBoxType: + """ + Args: + box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 5 + from_mode, to_mode (BoxMode) + + Returns: + The converted box of the same type. + """ + if from_mode == to_mode: + return box + + original_type = type(box) + is_numpy = isinstance(box, np.ndarray) + single_box = isinstance(box, (list, tuple)) + if single_box: + assert len(box) == 4 or len(box) == 5, ( + "BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor," + " where k == 4 or 5" + ) + arr = torch.tensor(box)[None, :] + else: + # avoid modifying the input box + if is_numpy: + arr = torch.from_numpy(np.asarray(box)).clone() + else: + arr = box.clone() + + assert to_mode not in [ + BoxMode.XYXY_REL, + BoxMode.XYWH_REL, + ] and from_mode not in [ + BoxMode.XYXY_REL, + BoxMode.XYWH_REL, + ], "Relative mode not yet supported!" + + if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS: + assert ( + arr.shape[-1] == 5 + ), "The last dimension of input shape must be 5 for XYWHA format" + original_dtype = arr.dtype + arr = arr.double() + + w = arr[:, 2] + h = arr[:, 3] + a = arr[:, 4] + c = torch.abs(torch.cos(a * math.pi / 180.0)) + s = torch.abs(torch.sin(a * math.pi / 180.0)) + # This basically computes the horizontal bounding rectangle of the rotated box + new_w = c * w + s * h + new_h = c * h + s * w + + # convert center to top-left corner + arr[:, 0] -= new_w / 2.0 + arr[:, 1] -= new_h / 2.0 + # bottom-right corner + arr[:, 2] = arr[:, 0] + new_w + arr[:, 3] = arr[:, 1] + new_h + + arr = arr[:, :4].to(dtype=original_dtype) + elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS: + original_dtype = arr.dtype + arr = arr.double() + arr[:, 0] += arr[:, 2] / 2.0 + arr[:, 1] += arr[:, 3] / 2.0 + angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype) + arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype) + else: + if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS: + arr[:, 2] += arr[:, 0] + arr[:, 3] += arr[:, 1] + elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS: + arr[:, 2] -= arr[:, 0] + arr[:, 3] -= arr[:, 1] + else: + raise NotImplementedError( + "Conversion from BoxMode {} to {} is not supported yet".format( + from_mode, to_mode + ) + ) + + if single_box: + return original_type(arr.flatten().tolist()) + if is_numpy: + return arr.numpy() + else: + return arr + + +class Boxes: + """ + This structure stores a list of boxes as a Nx4 torch.Tensor. + It supports some common methods about boxes + (`area`, `clip`, `nonempty`, etc), + and also behaves like a Tensor + (support indexing, `to(device)`, `.device`, and iteration over all boxes) + + Attributes: + tensor (torch.Tensor): float matrix of Nx4. Each row is (x1, y1, x2, y2). + """ + + def __init__(self, tensor: torch.Tensor): + """ + Args: + tensor (Tensor[float]): a Nx4 matrix. Each row is (x1, y1, x2, y2). + """ + if not isinstance(tensor, torch.Tensor): + tensor = torch.as_tensor( + tensor, dtype=torch.float32, device=torch.device("cpu") + ) + else: + tensor = tensor.to(torch.float32) + if tensor.numel() == 0: + # Use reshape, so we don't end up creating a new tensor that does not depend on + # the inputs (and consequently confuses jit) + tensor = tensor.reshape((-1, 4)).to(dtype=torch.float32) + assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size() + + self.tensor = tensor + + def clone(self) -> "Boxes": + """ + Clone the Boxes. + + Returns: + Boxes + """ + return Boxes(self.tensor.clone()) + + def to(self, device: torch.device): + # Boxes are assumed float32 and does not support to(dtype) + return Boxes(self.tensor.to(device=device)) + + def area(self) -> torch.Tensor: + """ + Computes the area of all the boxes. + + Returns: + torch.Tensor: a vector with areas of each box. + """ + box = self.tensor + area = (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1]) + return area + + def clip(self, box_size: Tuple[int, int]) -> None: + """ + Clip (in place) the boxes by limiting x coordinates to the range [0, width] + and y coordinates to the range [0, height]. + + Args: + box_size (height, width): The clipping box's size. + """ + assert torch.isfinite(self.tensor).all(), "Box tensor contains infinite or NaN!" + h, w = box_size + x1 = self.tensor[:, 0].clamp(min=0, max=w) + y1 = self.tensor[:, 1].clamp(min=0, max=h) + x2 = self.tensor[:, 2].clamp(min=0, max=w) + y2 = self.tensor[:, 3].clamp(min=0, max=h) + self.tensor = torch.stack((x1, y1, x2, y2), dim=-1) + + def nonempty(self, threshold: float = 0.0) -> torch.Tensor: + """ + Find boxes that are non-empty. + A box is considered empty, if either of its side is no larger than threshold. + + Returns: + Tensor: + a binary vector which represents whether each box is empty + (False) or non-empty (True). + """ + box = self.tensor + widths = box[:, 2] - box[:, 0] + heights = box[:, 3] - box[:, 1] + keep = (widths > threshold) & (heights > threshold) + return keep + + def __getitem__(self, item) -> "Boxes": + """ + Args: + item: int, slice, or a BoolTensor + + Returns: + Boxes: Create a new :class:`Boxes` by indexing. + + The following usage are allowed: + + 1. `new_boxes = boxes[3]`: return a `Boxes` which contains only one box. + 2. `new_boxes = boxes[2:10]`: return a slice of boxes. + 3. `new_boxes = boxes[vector]`, where vector is a torch.BoolTensor + with `length = len(boxes)`. Nonzero elements in the vector will be selected. + + Note that the returned Boxes might share storage with this Boxes, + subject to Pytorch's indexing semantics. + """ + if isinstance(item, int): + return Boxes(self.tensor[item].view(1, -1)) + b = self.tensor[item] + assert ( + b.dim() == 2 + ), "Indexing on Boxes with {} failed to return a matrix!".format(item) + return Boxes(b) + + def __len__(self) -> int: + return self.tensor.shape[0] + + def __repr__(self) -> str: + return "Boxes(" + str(self.tensor) + ")" + + def inside_box( + self, box_size: Tuple[int, int], boundary_threshold: int = 0 + ) -> torch.Tensor: + """ + Args: + box_size (height, width): Size of the reference box. + boundary_threshold (int): Boxes that extend beyond the reference box + boundary by more than boundary_threshold are considered "outside". + + Returns: + a binary vector, indicating whether each box is inside the reference box. + """ + height, width = box_size + inds_inside = ( + (self.tensor[..., 0] >= -boundary_threshold) + & (self.tensor[..., 1] >= -boundary_threshold) + & (self.tensor[..., 2] < width + boundary_threshold) + & (self.tensor[..., 3] < height + boundary_threshold) + ) + return inds_inside + + def get_centers(self) -> torch.Tensor: + """ + Returns: + The box centers in a Nx2 array of (x, y). + """ + return (self.tensor[:, :2] + self.tensor[:, 2:]) / 2 + + def scale(self, scale_x: float, scale_y: float) -> None: + """ + Scale the box with horizontal and vertical scaling factors + """ + self.tensor[:, 0::2] *= scale_x + self.tensor[:, 1::2] *= scale_y + + @classmethod + def cat(cls, boxes_list: List["Boxes"]) -> "Boxes": + """ + Concatenates a list of Boxes into a single Boxes + + Arguments: + boxes_list (list[Boxes]) + + Returns: + Boxes: the concatenated Boxes + """ + assert isinstance(boxes_list, (list, tuple)) + if len(boxes_list) == 0: + return cls(torch.empty(0)) + assert all([isinstance(box, Boxes) for box in boxes_list]) + + # use torch.cat (v.s. layers.cat) so the returned boxes never share storage with input + cat_boxes = cls(torch.cat([b.tensor for b in boxes_list], dim=0)) + return cat_boxes + + @property + def device(self) -> device: + return self.tensor.device + + # type "Iterator[torch.Tensor]", yield, and iter() not supported by torchscript + # https://github.com/pytorch/pytorch/issues/18627 + @torch.jit.unused + def __iter__(self): + """ + Yield a box as a Tensor of shape (4,) at a time. + """ + yield from self.tensor + + +def pairwise_intersection(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor: + """ + Given two lists of boxes of size N and M, + compute the intersection area between __all__ N x M pairs of boxes. + The box order must be (xmin, ymin, xmax, ymax) + + Args: + boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively. + + Returns: + Tensor: intersection, sized [N,M]. + """ + boxes1, boxes2 = boxes1.tensor, boxes2.tensor + width_height = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) - torch.max( + boxes1[:, None, :2], boxes2[:, :2] + ) # [N,M,2] + + width_height.clamp_(min=0) # [N,M,2] + intersection = width_height.prod(dim=2) # [N,M] + return intersection + + +# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py +# with slight modifications +def pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor: + """ + Given two lists of boxes of size N and M, compute the IoU + (intersection over union) between **all** N x M pairs of boxes. + The box order must be (xmin, ymin, xmax, ymax). + + Args: + boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively. + + Returns: + Tensor: IoU, sized [N,M]. + """ + area1 = boxes1.area() # [N] + area2 = boxes2.area() # [M] + inter = pairwise_intersection(boxes1, boxes2) + + # handle empty boxes + iou = torch.where( + inter > 0, + inter / (area1[:, None] + area2 - inter), + torch.zeros(1, dtype=inter.dtype, device=inter.device), + ) + return iou + + +def pairwise_ioa(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor: + """ + Similar to :func:`pariwise_iou` but compute the IoA (intersection over boxes2 area). + + Args: + boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively. + + Returns: + Tensor: IoA, sized [N,M]. + """ + area2 = boxes2.area() # [M] + inter = pairwise_intersection(boxes1, boxes2) + + # handle empty boxes + ioa = torch.where( + inter > 0, inter / area2, torch.zeros(1, dtype=inter.dtype, device=inter.device) + ) + return ioa + + +def pairwise_point_box_distance(points: torch.Tensor, boxes: Boxes): + """ + Pairwise distance between N points and M boxes. The distance between a + point and a box is represented by the distance from the point to 4 edges + of the box. Distances are all positive when the point is inside the box. + + Args: + points: Nx2 coordinates. Each row is (x, y) + boxes: M boxes + + Returns: + Tensor: distances of size (N, M, 4). The 4 values are distances from + the point to the left, top, right, bottom of the box. + """ + x, y = points.unsqueeze(dim=2).unbind(dim=1) # (N, 1) + x0, y0, x1, y1 = boxes.tensor.unsqueeze(dim=0).unbind(dim=2) # (1, M) + return torch.stack([x - x0, y - y0, x1 - x, y1 - y], dim=2) + + +def matched_pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor: + """ + Compute pairwise intersection over union (IOU) of two sets of matched + boxes that have the same number of boxes. + Similar to :func:`pairwise_iou`, but computes only diagonal elements of the matrix. + + Args: + boxes1 (Boxes): bounding boxes, sized [N,4]. + boxes2 (Boxes): same length as boxes1 + Returns: + Tensor: iou, sized [N]. + """ + assert len(boxes1) == len(boxes2), ( + "boxlists should have the same" "number of entries, got {}, {}".format( + len(boxes1), len(boxes2) + ) + ) + area1 = boxes1.area() # [N] + area2 = boxes2.area() # [N] + box1, box2 = boxes1.tensor, boxes2.tensor + lt = torch.max(box1[:, :2], box2[:, :2]) # [N,2] + rb = torch.min(box1[:, 2:], box2[:, 2:]) # [N,2] + wh = (rb - lt).clamp(min=0) # [N,2] + inter = wh[:, 0] * wh[:, 1] # [N] + iou = inter / (area1 + area2 - inter) # [N] + return iou diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/color_map.py b/src/nn/segearth_ov3/sam3/agent/helpers/color_map.py new file mode 100644 index 0000000..1ea1b29 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/color_map.py @@ -0,0 +1,150 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +An awesome colormap for really neat visualizations. +Copied from Detectron, and removed gray colors. +""" + +import random + +import numpy as np + +__all__ = ["colormap", "random_color", "random_colors"] + + +# A list of 25 bright and sharp colors for segmentation masks, +# generated from the edges of the sRGB color space for maximum intensity. +_COLORS = ( + np.array( + [ + # The original 8 sharp colors + 1.000, + 1.000, + 0.000, # 1. Yellow + 0.000, + 1.000, + 0.000, # 2. Lime + 0.000, + 1.000, + 1.000, # 3. Cyan + 1.000, + 0.000, + 1.000, # 4. Magenta + 1.000, + 0.000, + 0.000, # 5. Red + 1.000, + 0.498, + 0.000, # 6. Orange + 0.498, + 1.000, + 0.000, # 7. Chartreuse + 0.000, + 1.000, + 0.498, # 8. Spring Green + 1.000, + 0.000, + 0.498, # 9. Rose + 0.498, + 0.000, + 1.000, # 10. Violet + 0.753, + 1.000, + 0.000, # 11. Electric Lime + 1.000, + 0.753, + 0.000, # 12. Vivid Orange + 0.000, + 1.000, + 0.753, # 13. Turquoise + 0.753, + 0.000, + 1.000, # 14. Bright Violet + 1.000, + 0.000, + 0.753, # 15. Bright Pink + 1.000, + 0.251, + 0.000, # 16. Fiery Orange + 0.251, + 1.000, + 0.000, # 17. Bright Chartreuse + 0.000, + 1.000, + 0.251, # 18. Malachite Green + 0.251, + 0.000, + 1.000, # 19. Deep Violet + 1.000, + 0.000, + 0.251, # 20. Hot Pink + ] + ) + .astype(np.float32) + .reshape(-1, 3) +) + + +def colormap(rgb=False, maximum=255): + """ + Args: + rgb (bool): whether to return RGB colors or BGR colors. + maximum (int): either 255 or 1 + + Returns: + ndarray: a float32 array of Nx3 colors, in range [0, 255] or [0, 1] + """ + assert maximum in [255, 1], maximum + c = _COLORS * maximum + if not rgb: + c = c[:, ::-1] + return c + + +def random_color(rgb=False, maximum=255): + """ + Args: + rgb (bool): whether to return RGB colors or BGR colors. + maximum (int): either 255 or 1 + + Returns: + ndarray: a vector of 3 numbers + """ + idx = np.random.randint(0, len(_COLORS)) + ret = _COLORS[idx] * maximum + if not rgb: + ret = ret[::-1] + return ret + + +def random_colors(N, rgb=False, maximum=255): + """ + Args: + N (int): number of unique colors needed + rgb (bool): whether to return RGB colors or BGR colors. + maximum (int): either 255 or 1 + + Returns: + ndarray: a list of random_color + """ + indices = random.sample(range(len(_COLORS)), N) + ret = [_COLORS[i] * maximum for i in indices] + if not rgb: + ret = [x[::-1] for x in ret] + return ret + + +if __name__ == "__main__": + import cv2 + + size = 100 + H, W = 10, 10 + canvas = np.random.rand(H * size, W * size, 3).astype("float32") + for h in range(H): + for w in range(W): + idx = h * W + w + if idx >= len(_COLORS): + break + canvas[h * size : (h + 1) * size, w * size : (w + 1) * size] = _COLORS[idx] + cv2.imshow("a", canvas) + cv2.waitKey(0) diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/keypoints.py b/src/nn/segearth_ov3/sam3/agent/helpers/keypoints.py new file mode 100755 index 0000000..040810f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/keypoints.py @@ -0,0 +1,244 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Any, List, Tuple, Union + +import numpy as np +import torch +from torch.nn import functional as F + + +class Keypoints: + """ + Stores keypoint **annotation** data. GT Instances have a `gt_keypoints` property + containing the x,y location and visibility flag of each keypoint. This tensor has shape + (N, K, 3) where N is the number of instances and K is the number of keypoints per instance. + + The visibility flag follows the COCO format and must be one of three integers: + + * v=0: not labeled (in which case x=y=0) + * v=1: labeled but not visible + * v=2: labeled and visible + """ + + def __init__(self, keypoints: Union[torch.Tensor, np.ndarray, List[List[float]]]): + """ + Arguments: + keypoints: A Tensor, numpy array, or list of the x, y, and visibility of each keypoint. + The shape should be (N, K, 3) where N is the number of + instances, and K is the number of keypoints per instance. + """ + device = ( + keypoints.device + if isinstance(keypoints, torch.Tensor) + else torch.device("cpu") + ) + keypoints = torch.as_tensor(keypoints, dtype=torch.float32, device=device) + assert keypoints.dim() == 3 and keypoints.shape[2] == 3, keypoints.shape + self.tensor = keypoints + + def __len__(self) -> int: + return self.tensor.size(0) + + def to(self, *args: Any, **kwargs: Any) -> "Keypoints": + return type(self)(self.tensor.to(*args, **kwargs)) + + @property + def device(self) -> torch.device: + return self.tensor.device + + def to_heatmap(self, boxes: torch.Tensor, heatmap_size: int) -> torch.Tensor: + """ + Convert keypoint annotations to a heatmap of one-hot labels for training, + as described in :paper:`Mask R-CNN`. + + Arguments: + boxes: Nx4 tensor, the boxes to draw the keypoints to + + Returns: + heatmaps: + A tensor of shape (N, K), each element is integer spatial label + in the range [0, heatmap_size**2 - 1] for each keypoint in the input. + valid: + A tensor of shape (N, K) containing whether each keypoint is in the roi or not. + """ + return _keypoints_to_heatmap(self.tensor, boxes, heatmap_size) + + def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "Keypoints": + """ + Create a new `Keypoints` by indexing on this `Keypoints`. + + The following usage are allowed: + + 1. `new_kpts = kpts[3]`: return a `Keypoints` which contains only one instance. + 2. `new_kpts = kpts[2:10]`: return a slice of key points. + 3. `new_kpts = kpts[vector]`, where vector is a torch.ByteTensor + with `length = len(kpts)`. Nonzero elements in the vector will be selected. + + Note that the returned Keypoints might share storage with this Keypoints, + subject to Pytorch's indexing semantics. + """ + if isinstance(item, int): + return Keypoints([self.tensor[item]]) + return Keypoints(self.tensor[item]) + + def __repr__(self) -> str: + s = self.__class__.__name__ + "(" + s += "num_instances={})".format(len(self.tensor)) + return s + + @staticmethod + def cat(keypoints_list: List["Keypoints"]) -> "Keypoints": + """ + Concatenates a list of Keypoints into a single Keypoints + + Arguments: + keypoints_list (list[Keypoints]) + + Returns: + Keypoints: the concatenated Keypoints + """ + assert isinstance(keypoints_list, (list, tuple)) + assert len(keypoints_list) > 0 + assert all(isinstance(keypoints, Keypoints) for keypoints in keypoints_list) + + cat_kpts = type(keypoints_list[0])( + torch.cat([kpts.tensor for kpts in keypoints_list], dim=0) + ) + return cat_kpts + + +def _keypoints_to_heatmap( + keypoints: torch.Tensor, rois: torch.Tensor, heatmap_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Encode keypoint locations into a target heatmap for use in SoftmaxWithLoss across space. + + Maps keypoints from the half-open interval [x1, x2) on continuous image coordinates to the + closed interval [0, heatmap_size - 1] on discrete image coordinates. We use the + continuous-discrete conversion from Heckbert 1990 ("What is the coordinate of a pixel?"): + d = floor(c) and c = d + 0.5, where d is a discrete coordinate and c is a continuous coordinate. + + Arguments: + keypoints: tensor of keypoint locations in of shape (N, K, 3). + rois: Nx4 tensor of rois in xyxy format + heatmap_size: integer side length of square heatmap. + + Returns: + heatmaps: A tensor of shape (N, K) containing an integer spatial label + in the range [0, heatmap_size**2 - 1] for each keypoint in the input. + valid: A tensor of shape (N, K) containing whether each keypoint is in + the roi or not. + """ + + if rois.numel() == 0: + return rois.new().long(), rois.new().long() + offset_x = rois[:, 0] + offset_y = rois[:, 1] + scale_x = heatmap_size / (rois[:, 2] - rois[:, 0]) + scale_y = heatmap_size / (rois[:, 3] - rois[:, 1]) + + offset_x = offset_x[:, None] + offset_y = offset_y[:, None] + scale_x = scale_x[:, None] + scale_y = scale_y[:, None] + + x = keypoints[..., 0] + y = keypoints[..., 1] + + x_boundary_inds = x == rois[:, 2][:, None] + y_boundary_inds = y == rois[:, 3][:, None] + + x = (x - offset_x) * scale_x + x = x.floor().long() + y = (y - offset_y) * scale_y + y = y.floor().long() + + x[x_boundary_inds] = heatmap_size - 1 + y[y_boundary_inds] = heatmap_size - 1 + + valid_loc = (x >= 0) & (y >= 0) & (x < heatmap_size) & (y < heatmap_size) + vis = keypoints[..., 2] > 0 + valid = (valid_loc & vis).long() + + lin_ind = y * heatmap_size + x + heatmaps = lin_ind * valid + + return heatmaps, valid + + +@torch.jit.script_if_tracing +def heatmaps_to_keypoints(maps: torch.Tensor, rois: torch.Tensor) -> torch.Tensor: + """ + Extract predicted keypoint locations from heatmaps. + + Args: + maps (Tensor): (#ROIs, #keypoints, POOL_H, POOL_W). The predicted heatmap of logits for + each ROI and each keypoint. + rois (Tensor): (#ROIs, 4). The box of each ROI. + + Returns: + Tensor of shape (#ROIs, #keypoints, 4) with the last dimension corresponding to + (x, y, logit, score) for each keypoint. + + When converting discrete pixel indices in an NxN image to a continuous keypoint coordinate, + we maintain consistency with :meth:`Keypoints.to_heatmap` by using the conversion from + Heckbert 1990: c = d + 0.5, where d is a discrete coordinate and c is a continuous coordinate. + """ + + offset_x = rois[:, 0] + offset_y = rois[:, 1] + + widths = (rois[:, 2] - rois[:, 0]).clamp(min=1) + heights = (rois[:, 3] - rois[:, 1]).clamp(min=1) + widths_ceil = widths.ceil() + heights_ceil = heights.ceil() + + num_rois, num_keypoints = maps.shape[:2] + xy_preds = maps.new_zeros(rois.shape[0], num_keypoints, 4) + + width_corrections = widths / widths_ceil + height_corrections = heights / heights_ceil + + keypoints_idx = torch.arange(num_keypoints, device=maps.device) + + for i in range(num_rois): + outsize = (int(heights_ceil[i]), int(widths_ceil[i])) + roi_map = F.interpolate( + maps[[i]], size=outsize, mode="bicubic", align_corners=False + ) + + # Although semantically equivalent, `reshape` is used instead of `squeeze` due + # to limitation during ONNX export of `squeeze` in scripting mode + roi_map = roi_map.reshape(roi_map.shape[1:]) # keypoints x H x W + + # softmax over the spatial region + max_score, _ = roi_map.view(num_keypoints, -1).max(1) + max_score = max_score.view(num_keypoints, 1, 1) + tmp_full_resolution = (roi_map - max_score).exp_() + tmp_pool_resolution = (maps[i] - max_score).exp_() + # Produce scores over the region H x W, but normalize with POOL_H x POOL_W, + # so that the scores of objects of different absolute sizes will be more comparable + roi_map_scores = tmp_full_resolution / tmp_pool_resolution.sum( + (1, 2), keepdim=True + ) + + w = roi_map.shape[2] + pos = roi_map.view(num_keypoints, -1).argmax(1) + + x_int = pos % w + y_int = (pos - x_int) // w + + assert ( + roi_map_scores[keypoints_idx, y_int, x_int] + == roi_map_scores.view(num_keypoints, -1).max(1)[0] + ).all() + + x = (x_int.float() + 0.5) * width_corrections[i] + y = (y_int.float() + 0.5) * height_corrections[i] + + xy_preds[i, :, 0] = x + offset_x[i] + xy_preds[i, :, 1] = y + offset_y[i] + xy_preds[i, :, 2] = roi_map[keypoints_idx, y_int, x_int] + xy_preds[i, :, 3] = roi_map_scores[keypoints_idx, y_int, x_int] + + return xy_preds diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/mask_overlap_removal.py b/src/nn/segearth_ov3/sam3/agent/helpers/mask_overlap_removal.py new file mode 100644 index 0000000..386706d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/mask_overlap_removal.py @@ -0,0 +1,128 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Dict, List + +import numpy as np +import torch + +try: + from pycocotools import mask as mask_utils +except Exception: + mask_utils = None + + +def mask_intersection( + masks1: torch.Tensor, masks2: torch.Tensor, block_size: int = 16 +) -> torch.Tensor: + assert masks1.shape[1:] == masks2.shape[1:] + assert masks1.dtype == torch.bool and masks2.dtype == torch.bool + N, M = masks1.shape[0], masks2.shape[0] + out = torch.zeros(N, M, device=masks1.device, dtype=torch.long) + for i in range(0, N, block_size): + for j in range(0, M, block_size): + a = masks1[i : i + block_size] + b = masks2[j : j + block_size] + inter = (a[:, None] & b[None, :]).flatten(-2).sum(-1) + out[i : i + block_size, j : j + block_size] = inter + return out + + +def mask_iom(masks1: torch.Tensor, masks2: torch.Tensor) -> torch.Tensor: + assert masks1.shape[1:] == masks2.shape[1:] + assert masks1.dtype == torch.bool and masks2.dtype == torch.bool + inter = mask_intersection(masks1, masks2) + area1 = masks1.flatten(-2).sum(-1) # (N,) + area2 = masks2.flatten(-2).sum(-1) # (M,) + min_area = torch.min(area1[:, None], area2[None, :]).clamp_min(1) + return inter.float() / (min_area.float() + 1e-8) + + +def _decode_single_mask(mask_repr, h: int, w: int) -> np.ndarray: + if isinstance(mask_repr, (list, tuple, np.ndarray)): + arr = np.array(mask_repr) + if arr.ndim != 2: + raise ValueError("Mask array must be 2D (H, W).") + return (arr > 0).astype(np.uint8) + + if mask_utils is None: + raise ImportError( + "pycocotools is required to decode RLE mask strings. pip install pycocotools" + ) + + if not isinstance(mask_repr, (str, bytes)): + raise ValueError("Unsupported mask representation type for RLE decode.") + + rle = { + "counts": mask_repr if isinstance(mask_repr, (str, bytes)) else str(mask_repr), + "size": [h, w], + } + decoded = mask_utils.decode(rle) + if decoded.ndim == 3: + decoded = decoded[:, :, 0] + return (decoded > 0).astype(np.uint8) + + +def _decode_masks_to_torch_bool(pred_masks: List, h: int, w: int) -> torch.Tensor: + bin_masks = [_decode_single_mask(m, h, w) for m in pred_masks] + masks_np = np.stack(bin_masks, axis=0).astype(np.uint8) # (N, H, W) + return torch.from_numpy(masks_np > 0) + + +def remove_overlapping_masks(sample: Dict, iom_thresh: float = 0.3) -> Dict: + """ + Greedy keep: sort by score desc; keep a mask if IoM to all kept masks <= threshold. + If pred_masks has length 0 or 1, returns sample unchanged (no extra keys). + """ + # Basic presence checks + if "pred_masks" not in sample or not isinstance(sample["pred_masks"], list): + return sample # nothing to do / preserve as-is + + pred_masks = sample["pred_masks"] + N = len(pred_masks) + + # --- Early exit: 0 or 1 mask -> do NOT modify the JSON at all --- + if N <= 1: + return sample + + # From here on we have at least 2 masks + h = int(sample["orig_img_h"]) + w = int(sample["orig_img_w"]) + pred_scores = sample.get("pred_scores", [1.0] * N) # fallback if scores missing + pred_boxes = sample.get("pred_boxes", None) + + assert N == len(pred_scores), "pred_masks and pred_scores must have same length" + if pred_boxes is not None: + assert N == len(pred_boxes), "pred_masks and pred_boxes must have same length" + + masks_bool = _decode_masks_to_torch_bool(pred_masks, h, w) # (N, H, W) + + order = sorted(range(N), key=lambda i: float(pred_scores[i]), reverse=True) + kept_idx: List[int] = [] + kept_masks: List[torch.Tensor] = [] + + for i in order: + cand = masks_bool[i].unsqueeze(0) # (1, H, W) + if len(kept_masks) == 0: + kept_idx.append(i) + kept_masks.append(masks_bool[i]) + continue + + kept_stack = torch.stack(kept_masks, dim=0) # (K, H, W) + iom_vals = mask_iom(cand, kept_stack).squeeze(0) # (K,) + if torch.any(iom_vals > iom_thresh): + continue # overlaps too much with a higher-scored kept mask + kept_idx.append(i) + kept_masks.append(masks_bool[i]) + + kept_idx_sorted = sorted(kept_idx) + + # Build filtered JSON (this *does* modify fields; only for N>=2 case) + out = dict(sample) + out["pred_masks"] = [pred_masks[i] for i in kept_idx_sorted] + out["pred_scores"] = [pred_scores[i] for i in kept_idx_sorted] + if pred_boxes is not None: + out["pred_boxes"] = [pred_boxes[i] for i in kept_idx_sorted] + out["kept_indices"] = kept_idx_sorted + out["removed_indices"] = [i for i in range(N) if i not in set(kept_idx_sorted)] + out["iom_threshold"] = float(iom_thresh) + return out diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/masks.py b/src/nn/segearth_ov3/sam3/agent/helpers/masks.py new file mode 100755 index 0000000..dc78140 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/masks.py @@ -0,0 +1,560 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import copy +import itertools +from typing import Any, Iterator, List, Union + +import numpy as np +import pycocotools.mask as mask_util +import torch +from torch import device + +from .boxes import Boxes +from .memory import retry_if_cuda_oom + +from .roi_align import ROIAlign + + +def polygon_area(x, y): + # Using the shoelace formula + # https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates + return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1))) + + +def polygons_to_bitmask( + polygons: List[np.ndarray], height: int, width: int +) -> np.ndarray: + """ + Args: + polygons (list[ndarray]): each array has shape (Nx2,) + height, width (int) + + Returns: + ndarray: a bool mask of shape (height, width) + """ + if len(polygons) == 0: + # COCOAPI does not support empty polygons + return np.zeros((height, width)).astype(bool) + rles = mask_util.frPyObjects(polygons, height, width) + rle = mask_util.merge(rles) + return mask_util.decode(rle).astype(bool) + + +def rasterize_polygons_within_box( + polygons: List[np.ndarray], box: np.ndarray, mask_size: int +) -> torch.Tensor: + """ + Rasterize the polygons into a mask image and + crop the mask content in the given box. + The cropped mask is resized to (mask_size, mask_size). + + This function is used when generating training targets for mask head in Mask R-CNN. + Given original ground-truth masks for an image, new ground-truth mask + training targets in the size of `mask_size x mask_size` + must be provided for each predicted box. This function will be called to + produce such targets. + + Args: + polygons (list[ndarray[float]]): a list of polygons, which represents an instance. + box: 4-element numpy array + mask_size (int): + + Returns: + Tensor: BoolTensor of shape (mask_size, mask_size) + """ + # 1. Shift the polygons w.r.t the boxes + w, h = box[2] - box[0], box[3] - box[1] + + polygons = copy.deepcopy(polygons) + for p in polygons: + p[0::2] = p[0::2] - box[0] + p[1::2] = p[1::2] - box[1] + + # 2. Rescale the polygons to the new box size + # max() to avoid division by small number + ratio_h = mask_size / max(h, 0.1) + ratio_w = mask_size / max(w, 0.1) + + if ratio_h == ratio_w: + for p in polygons: + p *= ratio_h + else: + for p in polygons: + p[0::2] *= ratio_w + p[1::2] *= ratio_h + + # 3. Rasterize the polygons with coco api + mask = polygons_to_bitmask(polygons, mask_size, mask_size) + mask = torch.from_numpy(mask) + return mask + + +class BitMasks: + """ + This class stores the segmentation masks for all objects in one image, in + the form of bitmaps. + + Attributes: + tensor: bool Tensor of N,H,W, representing N instances in the image. + """ + + def __init__(self, tensor: Union[torch.Tensor, np.ndarray]): + """ + Args: + tensor: bool Tensor of N,H,W, representing N instances in the image. + """ + if isinstance(tensor, torch.Tensor): + tensor = tensor.to(torch.bool) + else: + tensor = torch.as_tensor( + tensor, dtype=torch.bool, device=torch.device("cpu") + ) + assert tensor.dim() == 3, tensor.size() + self.image_size = tensor.shape[1:] + self.tensor = tensor + + @torch.jit.unused + def to(self, *args: Any, **kwargs: Any) -> "BitMasks": + return BitMasks(self.tensor.to(*args, **kwargs)) + + @property + def device(self) -> torch.device: + return self.tensor.device + + @torch.jit.unused + def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "BitMasks": + """ + Returns: + BitMasks: Create a new :class:`BitMasks` by indexing. + + The following usage are allowed: + + 1. `new_masks = masks[3]`: return a `BitMasks` which contains only one mask. + 2. `new_masks = masks[2:10]`: return a slice of masks. + 3. `new_masks = masks[vector]`, where vector is a torch.BoolTensor + with `length = len(masks)`. Nonzero elements in the vector will be selected. + + Note that the returned object might share storage with this object, + subject to Pytorch's indexing semantics. + """ + if isinstance(item, int): + return BitMasks(self.tensor[item].unsqueeze(0)) + m = self.tensor[item] + assert ( + m.dim() == 3 + ), "Indexing on BitMasks with {} returns a tensor with shape {}!".format( + item, m.shape + ) + return BitMasks(m) + + @torch.jit.unused + def __iter__(self) -> torch.Tensor: + yield from self.tensor + + @torch.jit.unused + def __repr__(self) -> str: + s = self.__class__.__name__ + "(" + s += "num_instances={})".format(len(self.tensor)) + return s + + def __len__(self) -> int: + return self.tensor.shape[0] + + def nonempty(self) -> torch.Tensor: + """ + Find masks that are non-empty. + + Returns: + Tensor: a BoolTensor which represents + whether each mask is empty (False) or non-empty (True). + """ + return self.tensor.flatten(1).any(dim=1) + + @staticmethod + def from_polygon_masks( + polygon_masks: Union["PolygonMasks", List[List[np.ndarray]]], + height: int, + width: int, + ) -> "BitMasks": + """ + Args: + polygon_masks (list[list[ndarray]] or PolygonMasks) + height, width (int) + """ + if isinstance(polygon_masks, PolygonMasks): + polygon_masks = polygon_masks.polygons + masks = [polygons_to_bitmask(p, height, width) for p in polygon_masks] + if len(masks): + return BitMasks(torch.stack([torch.from_numpy(x) for x in masks])) + else: + return BitMasks(torch.empty(0, height, width, dtype=torch.bool)) + + @staticmethod + def from_roi_masks(roi_masks: "ROIMasks", height: int, width: int) -> "BitMasks": + """ + Args: + roi_masks: + height, width (int): + """ + return roi_masks.to_bitmasks(height, width) + + def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor: + """ + Crop each bitmask by the given box, and resize results to (mask_size, mask_size). + This can be used to prepare training targets for Mask R-CNN. + It has less reconstruction error compared to rasterization with polygons. + However we observe no difference in accuracy, + but BitMasks requires more memory to store all the masks. + + Args: + boxes (Tensor): Nx4 tensor storing the boxes for each mask + mask_size (int): the size of the rasterized mask. + + Returns: + Tensor: + A bool tensor of shape (N, mask_size, mask_size), where + N is the number of predicted boxes for this image. + """ + assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self)) + device = self.tensor.device + + batch_inds = torch.arange(len(boxes), device=device).to(dtype=boxes.dtype)[ + :, None + ] + rois = torch.cat([batch_inds, boxes], dim=1) # Nx5 + + bit_masks = self.tensor.to(dtype=torch.float32) + rois = rois.to(device=device) + output = ( + ROIAlign((mask_size, mask_size), 1.0, 0, aligned=True) + .forward(bit_masks[:, None, :, :], rois) + .squeeze(1) + ) + output = output >= 0.5 + return output + + def get_bounding_boxes(self) -> Boxes: + """ + Returns: + Boxes: tight bounding boxes around bitmasks. + If a mask is empty, it's bounding box will be all zero. + """ + boxes = torch.zeros(self.tensor.shape[0], 4, dtype=torch.float32) + x_any = torch.any(self.tensor, dim=1) + y_any = torch.any(self.tensor, dim=2) + for idx in range(self.tensor.shape[0]): + x = torch.where(x_any[idx, :])[0] + y = torch.where(y_any[idx, :])[0] + if len(x) > 0 and len(y) > 0: + boxes[idx, :] = torch.as_tensor( + [x[0], y[0], x[-1] + 1, y[-1] + 1], dtype=torch.float32 + ) + return Boxes(boxes) + + @staticmethod + def cat(bitmasks_list: List["BitMasks"]) -> "BitMasks": + """ + Concatenates a list of BitMasks into a single BitMasks + + Arguments: + bitmasks_list (list[BitMasks]) + + Returns: + BitMasks: the concatenated BitMasks + """ + assert isinstance(bitmasks_list, (list, tuple)) + assert len(bitmasks_list) > 0 + assert all(isinstance(bitmask, BitMasks) for bitmask in bitmasks_list) + + cat_bitmasks = type(bitmasks_list[0])( + torch.cat([bm.tensor for bm in bitmasks_list], dim=0) + ) + return cat_bitmasks + + +class PolygonMasks: + """ + This class stores the segmentation masks for all objects in one image, in the form of polygons. + + Attributes: + polygons: list[list[ndarray]]. Each ndarray is a float64 vector representing a polygon. + """ + + def __init__(self, polygons: List[List[Union[torch.Tensor, np.ndarray]]]): + """ + Arguments: + polygons (list[list[np.ndarray]]): The first + level of the list correspond to individual instances, + the second level to all the polygons that compose the + instance, and the third level to the polygon coordinates. + The third level array should have the format of + [x0, y0, x1, y1, ..., xn, yn] (n >= 3). + """ + if not isinstance(polygons, list): + raise ValueError( + "Cannot create PolygonMasks: Expect a list of list of polygons per image. " + "Got '{}' instead.".format(type(polygons)) + ) + + def _make_array(t: Union[torch.Tensor, np.ndarray]) -> np.ndarray: + # Use float64 for higher precision, because why not? + # Always put polygons on CPU (self.to is a no-op) since they + # are supposed to be small tensors. + # May need to change this assumption if GPU placement becomes useful + if isinstance(t, torch.Tensor): + t = t.cpu().numpy() + return np.asarray(t).astype("float64") + + def process_polygons( + polygons_per_instance: List[Union[torch.Tensor, np.ndarray]], + ) -> List[np.ndarray]: + if not isinstance(polygons_per_instance, list): + raise ValueError( + "Cannot create polygons: Expect a list of polygons per instance. " + "Got '{}' instead.".format(type(polygons_per_instance)) + ) + # transform each polygon to a numpy array + polygons_per_instance = [_make_array(p) for p in polygons_per_instance] + for polygon in polygons_per_instance: + if len(polygon) % 2 != 0 or len(polygon) < 6: + raise ValueError( + f"Cannot create a polygon from {len(polygon)} coordinates." + ) + return polygons_per_instance + + self.polygons: List[List[np.ndarray]] = [ + process_polygons(polygons_per_instance) + for polygons_per_instance in polygons + ] + + def to(self, *args: Any, **kwargs: Any) -> "PolygonMasks": + return self + + @property + def device(self) -> torch.device: + return torch.device("cpu") + + def get_bounding_boxes(self) -> Boxes: + """ + Returns: + Boxes: tight bounding boxes around polygon masks. + """ + boxes = torch.zeros(len(self.polygons), 4, dtype=torch.float32) + for idx, polygons_per_instance in enumerate(self.polygons): + minxy = torch.as_tensor([float("inf"), float("inf")], dtype=torch.float32) + maxxy = torch.zeros(2, dtype=torch.float32) + for polygon in polygons_per_instance: + coords = torch.from_numpy(polygon).view(-1, 2).to(dtype=torch.float32) + minxy = torch.min(minxy, torch.min(coords, dim=0).values) + maxxy = torch.max(maxxy, torch.max(coords, dim=0).values) + boxes[idx, :2] = minxy + boxes[idx, 2:] = maxxy + return Boxes(boxes) + + def nonempty(self) -> torch.Tensor: + """ + Find masks that are non-empty. + + Returns: + Tensor: + a BoolTensor which represents whether each mask is empty (False) or not (True). + """ + keep = [1 if len(polygon) > 0 else 0 for polygon in self.polygons] + return torch.from_numpy(np.asarray(keep, dtype=bool)) + + def __getitem__( + self, item: Union[int, slice, List[int], torch.BoolTensor] + ) -> "PolygonMasks": + """ + Support indexing over the instances and return a `PolygonMasks` object. + `item` can be: + + 1. An integer. It will return an object with only one instance. + 2. A slice. It will return an object with the selected instances. + 3. A list[int]. It will return an object with the selected instances, + correpsonding to the indices in the list. + 4. A vector mask of type BoolTensor, whose length is num_instances. + It will return an object with the instances whose mask is nonzero. + """ + if isinstance(item, int): + selected_polygons = [self.polygons[item]] + elif isinstance(item, slice): + selected_polygons = self.polygons[item] + elif isinstance(item, list): + selected_polygons = [self.polygons[i] for i in item] + elif isinstance(item, torch.Tensor): + # Polygons is a list, so we have to move the indices back to CPU. + if item.dtype == torch.bool: + assert item.dim() == 1, item.shape + item = item.nonzero().squeeze(1).cpu().numpy().tolist() + elif item.dtype in [torch.int32, torch.int64]: + item = item.cpu().numpy().tolist() + else: + raise ValueError( + "Unsupported tensor dtype={} for indexing!".format(item.dtype) + ) + selected_polygons = [self.polygons[i] for i in item] + return PolygonMasks(selected_polygons) + + def __iter__(self) -> Iterator[List[np.ndarray]]: + """ + Yields: + list[ndarray]: the polygons for one instance. + Each Tensor is a float64 vector representing a polygon. + """ + return iter(self.polygons) + + def __repr__(self) -> str: + s = self.__class__.__name__ + "(" + s += "num_instances={})".format(len(self.polygons)) + return s + + def __len__(self) -> int: + return len(self.polygons) + + def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor: + """ + Crop each mask by the given box, and resize results to (mask_size, mask_size). + This can be used to prepare training targets for Mask R-CNN. + + Args: + boxes (Tensor): Nx4 tensor storing the boxes for each mask + mask_size (int): the size of the rasterized mask. + + Returns: + Tensor: A bool tensor of shape (N, mask_size, mask_size), where + N is the number of predicted boxes for this image. + """ + assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self)) + + device = boxes.device + # Put boxes on the CPU, as the polygon representation is not efficient GPU-wise + # (several small tensors for representing a single instance mask) + boxes = boxes.to(torch.device("cpu")) + + results = [ + rasterize_polygons_within_box(poly, box.numpy(), mask_size) + for poly, box in zip(self.polygons, boxes) + ] + """ + poly: list[list[float]], the polygons for one instance + box: a tensor of shape (4,) + """ + if len(results) == 0: + return torch.empty(0, mask_size, mask_size, dtype=torch.bool, device=device) + return torch.stack(results, dim=0).to(device=device) + + def area(self): + """ + Computes area of the mask. + Only works with Polygons, using the shoelace formula: + https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates + + Returns: + Tensor: a vector, area for each instance + """ + + area = [] + for polygons_per_instance in self.polygons: + area_per_instance = 0 + for p in polygons_per_instance: + area_per_instance += polygon_area(p[0::2], p[1::2]) + area.append(area_per_instance) + + return torch.tensor(area) + + @staticmethod + def cat(polymasks_list: List["PolygonMasks"]) -> "PolygonMasks": + """ + Concatenates a list of PolygonMasks into a single PolygonMasks + + Arguments: + polymasks_list (list[PolygonMasks]) + + Returns: + PolygonMasks: the concatenated PolygonMasks + """ + assert isinstance(polymasks_list, (list, tuple)) + assert len(polymasks_list) > 0 + assert all(isinstance(polymask, PolygonMasks) for polymask in polymasks_list) + + cat_polymasks = type(polymasks_list[0])( + list(itertools.chain.from_iterable(pm.polygons for pm in polymasks_list)) + ) + return cat_polymasks + + +class ROIMasks: + """ + Represent masks by N smaller masks defined in some ROIs. Once ROI boxes are given, + full-image bitmask can be obtained by "pasting" the mask on the region defined + by the corresponding ROI box. + """ + + def __init__(self, tensor: torch.Tensor): + """ + Args: + tensor: (N, M, M) mask tensor that defines the mask within each ROI. + """ + if tensor.dim() != 3: + raise ValueError("ROIMasks must take a masks of 3 dimension.") + self.tensor = tensor + + def to(self, device: torch.device) -> "ROIMasks": + return ROIMasks(self.tensor.to(device)) + + @property + def device(self) -> device: + return self.tensor.device + + def __len__(self): + return self.tensor.shape[0] + + def __getitem__(self, item) -> "ROIMasks": + """ + Returns: + ROIMasks: Create a new :class:`ROIMasks` by indexing. + + The following usage are allowed: + + 1. `new_masks = masks[2:10]`: return a slice of masks. + 2. `new_masks = masks[vector]`, where vector is a torch.BoolTensor + with `length = len(masks)`. Nonzero elements in the vector will be selected. + + Note that the returned object might share storage with this object, + subject to Pytorch's indexing semantics. + """ + t = self.tensor[item] + if t.dim() != 3: + raise ValueError( + f"Indexing on ROIMasks with {item} returns a tensor with shape {t.shape}!" + ) + return ROIMasks(t) + + @torch.jit.unused + def __repr__(self) -> str: + s = self.__class__.__name__ + "(" + s += "num_instances={})".format(len(self.tensor)) + return s + + @torch.jit.unused + def to_bitmasks(self, boxes: torch.Tensor, height, width, threshold=0.5): + """ + Args: see documentation of :func:`paste_masks_in_image`. + """ + from detectron2.layers.mask_ops import ( + _paste_masks_tensor_shape, + paste_masks_in_image, + ) + + if torch.jit.is_tracing(): + if isinstance(height, torch.Tensor): + paste_func = _paste_masks_tensor_shape + else: + paste_func = paste_masks_in_image + else: + paste_func = retry_if_cuda_oom(paste_masks_in_image) + bitmasks = paste_func( + self.tensor, boxes.tensor, (height, width), threshold=threshold + ) + return BitMasks(bitmasks) diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/memory.py b/src/nn/segearth_ov3/sam3/agent/helpers/memory.py new file mode 100755 index 0000000..5d51f1b --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/memory.py @@ -0,0 +1,87 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +from contextlib import contextmanager +from functools import wraps + +import torch + +__all__ = ["retry_if_cuda_oom"] + + +@contextmanager +def _ignore_torch_cuda_oom(): + """ + A context which ignores CUDA OOM exception from pytorch. + """ + try: + yield + except RuntimeError as e: + # NOTE: the string may change? + if "CUDA out of memory. " in str(e): + pass + else: + raise + + +def retry_if_cuda_oom(func): + """ + Makes a function retry itself after encountering + pytorch's CUDA OOM error. + It will first retry after calling `torch.cuda.empty_cache()`. + + If that still fails, it will then retry by trying to convert inputs to CPUs. + In this case, it expects the function to dispatch to CPU implementation. + The return values may become CPU tensors as well and it's user's + responsibility to convert it back to CUDA tensor if needed. + + Args: + func: a stateless callable that takes tensor-like objects as arguments + + Returns: + a callable which retries `func` if OOM is encountered. + + Examples: + :: + output = retry_if_cuda_oom(some_torch_function)(input1, input2) + # output may be on CPU even if inputs are on GPU + + Note: + 1. When converting inputs to CPU, it will only look at each argument and check + if it has `.device` and `.to` for conversion. Nested structures of tensors + are not supported. + + 2. Since the function might be called more than once, it has to be + stateless. + """ + + def maybe_to_cpu(x): + try: + like_gpu_tensor = x.device.type == "cuda" and hasattr(x, "to") + except AttributeError: + like_gpu_tensor = False + if like_gpu_tensor: + return x.to(device="cpu") + else: + return x + + @wraps(func) + def wrapped(*args, **kwargs): + with _ignore_torch_cuda_oom(): + return func(*args, **kwargs) + + # Clear cache and retry + torch.cuda.empty_cache() + with _ignore_torch_cuda_oom(): + return func(*args, **kwargs) + + # Try on CPU. This slows down the code significantly, therefore print a notice. + logger = logging.getLogger(__name__) + logger.info( + "Attempting to copy inputs of {} to CPU due to CUDA OOM".format(str(func)) + ) + new_args = (maybe_to_cpu(x) for x in args) + new_kwargs = {k: maybe_to_cpu(v) for k, v in kwargs.items()} + return func(*new_args, **new_kwargs) + + return wrapped diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/rle.py b/src/nn/segearth_ov3/sam3/agent/helpers/rle.py new file mode 100755 index 0000000..277be79 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/rle.py @@ -0,0 +1,122 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Some utilities for RLE encoding that doesn't require downloading the masks to the cpu""" + +import numpy as np +import torch +from pycocotools import mask as mask_util + + +@torch.no_grad() +def rle_encode(orig_mask, return_areas=False): + """Encodes a collection of masks in RLE format + + This function emulates the behavior of the COCO API's encode function, but + is executed partially on the GPU for faster execution. + + Args: + mask (torch.Tensor): A mask of shape (N, H, W) with dtype=torch.bool + return_areas (bool): If True, add the areas of the masks as a part of + the RLE output dict under the "area" key. Default is False. + + Returns: + str: The RLE encoded masks + """ + assert orig_mask.ndim == 3, "Mask must be of shape (N, H, W)" + assert orig_mask.dtype == torch.bool, "Mask must have dtype=torch.bool" + + if orig_mask.numel() == 0: + return [] + + # First, transpose the spatial dimensions. + # This is necessary because the COCO API uses Fortran order + mask = orig_mask.transpose(1, 2) + + # Flatten the mask + flat_mask = mask.reshape(mask.shape[0], -1) + if return_areas: + mask_areas = flat_mask.sum(-1).tolist() + # Find the indices where the mask changes + differences = torch.ones( + mask.shape[0], flat_mask.shape[1] + 1, device=mask.device, dtype=torch.bool + ) + differences[:, 1:-1] = flat_mask[:, :-1] != flat_mask[:, 1:] + differences[:, 0] = flat_mask[:, 0] + _, change_indices = torch.where(differences) + + try: + boundaries = torch.cumsum(differences.sum(-1), 0).cpu() + except RuntimeError as _: + boundaries = torch.cumsum(differences.cpu().sum(-1), 0) + + change_indices_clone = change_indices.clone() + # First pass computes the RLEs on GPU, in a flatten format + for i in range(mask.shape[0]): + # Get the change indices for this batch item + beg = 0 if i == 0 else boundaries[i - 1].item() + end = boundaries[i].item() + change_indices[beg + 1 : end] -= change_indices_clone[beg : end - 1] + + # Now we can split the RLES of each batch item, and convert them to strings + # No more gpu at this point + change_indices = change_indices.tolist() + + batch_rles = [] + # Process each mask in the batch separately + for i in range(mask.shape[0]): + beg = 0 if i == 0 else boundaries[i - 1].item() + end = boundaries[i].item() + run_lengths = change_indices[beg:end] + + uncompressed_rle = {"counts": run_lengths, "size": list(orig_mask.shape[1:])} + h, w = uncompressed_rle["size"] + rle = mask_util.frPyObjects(uncompressed_rle, h, w) + rle["counts"] = rle["counts"].decode("utf-8") + if return_areas: + rle["area"] = mask_areas[i] + batch_rles.append(rle) + + return batch_rles + + +def robust_rle_encode(masks): + """Encodes a collection of masks in RLE format. Uses the gpu version fist, falls back to the cpu version if it fails""" + + assert masks.ndim == 3, "Mask must be of shape (N, H, W)" + assert masks.dtype == torch.bool, "Mask must have dtype=torch.bool" + + try: + return rle_encode(masks) + except RuntimeError as _: + masks = masks.cpu().numpy() + rles = [ + mask_util.encode( + np.array(mask[:, :, np.newaxis], dtype=np.uint8, order="F") + )[0] + for mask in masks + ] + for rle in rles: + rle["counts"] = rle["counts"].decode("utf-8") + return rles + + +def ann_to_rle(segm, im_info): + """Convert annotation which can be polygons, uncompressed RLE to RLE. + Args: + ann (dict) : annotation object + Returns: + ann (rle) + """ + h, w = im_info["height"], im_info["width"] + if isinstance(segm, list): + # polygon -- a single object might consist of multiple parts + # we merge all parts into one mask rle code + rles = mask_util.frPyObjects(segm, h, w) + rle = mask_util.merge(rles) + elif isinstance(segm["counts"], list): + # uncompressed RLE + rle = mask_util.frPyObjects(segm, h, w) + else: + # rle + rle = segm + return rle diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/roi_align.py b/src/nn/segearth_ov3/sam3/agent/helpers/roi_align.py new file mode 100755 index 0000000..640f853 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/roi_align.py @@ -0,0 +1,75 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from torch import nn +from torchvision.ops import roi_align + + +# NOTE: torchvision's RoIAlign has a different default aligned=False +class ROIAlign(nn.Module): + def __init__(self, output_size, spatial_scale, sampling_ratio, aligned=True): + """ + Args: + output_size (tuple): h, w + spatial_scale (float): scale the input boxes by this number + sampling_ratio (int): number of inputs samples to take for each output + sample. 0 to take samples densely. + aligned (bool): if False, use the legacy implementation in + Detectron. If True, align the results more perfectly. + + Note: + The meaning of aligned=True: + + Given a continuous coordinate c, its two neighboring pixel indices (in our + pixel model) are computed by floor(c - 0.5) and ceil(c - 0.5). For example, + c=1.3 has pixel neighbors with discrete indices [0] and [1] (which are sampled + from the underlying signal at continuous coordinates 0.5 and 1.5). But the original + roi_align (aligned=False) does not subtract the 0.5 when computing neighboring + pixel indices and therefore it uses pixels with a slightly incorrect alignment + (relative to our pixel model) when performing bilinear interpolation. + + With `aligned=True`, + we first appropriately scale the ROI and then shift it by -0.5 + prior to calling roi_align. This produces the correct neighbors; see + detectron2/tests/test_roi_align.py for verification. + + The difference does not make a difference to the model's performance if + ROIAlign is used together with conv layers. + """ + super().__init__() + self.output_size = output_size + self.spatial_scale = spatial_scale + self.sampling_ratio = sampling_ratio + self.aligned = aligned + + from torchvision import __version__ + + version = tuple(int(x) for x in __version__.split(".")[:2]) + # https://github.com/pytorch/vision/pull/2438 + assert version >= (0, 7), "Require torchvision >= 0.7" + + def forward(self, input, rois): + """ + Args: + input: NCHW images + rois: Bx5 boxes. First column is the index into N. The other 4 columns are xyxy. + """ + assert rois.dim() == 2 and rois.size(1) == 5 + if input.is_quantized: + input = input.dequantize() + return roi_align( + input, + rois.to(dtype=input.dtype), + self.output_size, + self.spatial_scale, + self.sampling_ratio, + self.aligned, + ) + + def __repr__(self): + tmpstr = self.__class__.__name__ + "(" + tmpstr += "output_size=" + str(self.output_size) + tmpstr += ", spatial_scale=" + str(self.spatial_scale) + tmpstr += ", sampling_ratio=" + str(self.sampling_ratio) + tmpstr += ", aligned=" + str(self.aligned) + tmpstr += ")" + return tmpstr diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/rotated_boxes.py b/src/nn/segearth_ov3/sam3/agent/helpers/rotated_boxes.py new file mode 100755 index 0000000..151a96f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/rotated_boxes.py @@ -0,0 +1,533 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from __future__ import absolute_import, division, print_function, unicode_literals + +import math +from typing import List, Tuple + +import torch + +# from detectron2.layers.rotated_boxes import pairwise_iou_rotated + +from .boxes import Boxes + + +def pairwise_iou_rotated(boxes1, boxes2): + """ + Return intersection-over-union (Jaccard index) of boxes. + + Both sets of boxes are expected to be in + (x_center, y_center, width, height, angle) format. + + Arguments: + boxes1 (Tensor[N, 5]) + boxes2 (Tensor[M, 5]) + + Returns: + iou (Tensor[N, M]): the NxM matrix containing the pairwise + IoU values for every element in boxes1 and boxes2 + """ + return torch.ops.detectron2.box_iou_rotated(boxes1, boxes2) + + +class RotatedBoxes(Boxes): + """ + This structure stores a list of rotated boxes as a Nx5 torch.Tensor. + It supports some common methods about boxes + (`area`, `clip`, `nonempty`, etc), + and also behaves like a Tensor + (support indexing, `to(device)`, `.device`, and iteration over all boxes) + """ + + def __init__(self, tensor: torch.Tensor): + """ + Args: + tensor (Tensor[float]): a Nx5 matrix. Each row is + (x_center, y_center, width, height, angle), + in which angle is represented in degrees. + While there's no strict range restriction for it, + the recommended principal range is between [-180, 180) degrees. + + Assume we have a horizontal box B = (x_center, y_center, width, height), + where width is along the x-axis and height is along the y-axis. + The rotated box B_rot (x_center, y_center, width, height, angle) + can be seen as: + + 1. When angle == 0: + B_rot == B + 2. When angle > 0: + B_rot is obtained by rotating B w.r.t its center by :math:`|angle|` degrees CCW; + 3. When angle < 0: + B_rot is obtained by rotating B w.r.t its center by :math:`|angle|` degrees CW. + + Mathematically, since the right-handed coordinate system for image space + is (y, x), where y is top->down and x is left->right, the 4 vertices of the + rotated rectangle :math:`(yr_i, xr_i)` (i = 1, 2, 3, 4) can be obtained from + the vertices of the horizontal rectangle :math:`(y_i, x_i)` (i = 1, 2, 3, 4) + in the following way (:math:`\\theta = angle*\\pi/180` is the angle in radians, + :math:`(y_c, x_c)` is the center of the rectangle): + + .. math:: + + yr_i = \\cos(\\theta) (y_i - y_c) - \\sin(\\theta) (x_i - x_c) + y_c, + + xr_i = \\sin(\\theta) (y_i - y_c) + \\cos(\\theta) (x_i - x_c) + x_c, + + which is the standard rigid-body rotation transformation. + + Intuitively, the angle is + (1) the rotation angle from y-axis in image space + to the height vector (top->down in the box's local coordinate system) + of the box in CCW, and + (2) the rotation angle from x-axis in image space + to the width vector (left->right in the box's local coordinate system) + of the box in CCW. + + More intuitively, consider the following horizontal box ABCD represented + in (x1, y1, x2, y2): (3, 2, 7, 4), + covering the [3, 7] x [2, 4] region of the continuous coordinate system + which looks like this: + + .. code:: none + + O--------> x + | + | A---B + | | | + | D---C + | + v y + + Note that each capital letter represents one 0-dimensional geometric point + instead of a 'square pixel' here. + + In the example above, using (x, y) to represent a point we have: + + .. math:: + + O = (0, 0), A = (3, 2), B = (7, 2), C = (7, 4), D = (3, 4) + + We name vector AB = vector DC as the width vector in box's local coordinate system, and + vector AD = vector BC as the height vector in box's local coordinate system. Initially, + when angle = 0 degree, they're aligned with the positive directions of x-axis and y-axis + in the image space, respectively. + + For better illustration, we denote the center of the box as E, + + .. code:: none + + O--------> x + | + | A---B + | | E | + | D---C + | + v y + + where the center E = ((3+7)/2, (2+4)/2) = (5, 3). + + Also, + + .. math:: + + width = |AB| = |CD| = 7 - 3 = 4, + height = |AD| = |BC| = 4 - 2 = 2. + + Therefore, the corresponding representation for the same shape in rotated box in + (x_center, y_center, width, height, angle) format is: + + (5, 3, 4, 2, 0), + + Now, let's consider (5, 3, 4, 2, 90), which is rotated by 90 degrees + CCW (counter-clockwise) by definition. It looks like this: + + .. code:: none + + O--------> x + | B-C + | | | + | |E| + | | | + | A-D + v y + + The center E is still located at the same point (5, 3), while the vertices + ABCD are rotated by 90 degrees CCW with regard to E: + A = (4, 5), B = (4, 1), C = (6, 1), D = (6, 5) + + Here, 90 degrees can be seen as the CCW angle to rotate from y-axis to + vector AD or vector BC (the top->down height vector in box's local coordinate system), + or the CCW angle to rotate from x-axis to vector AB or vector DC (the left->right + width vector in box's local coordinate system). + + .. math:: + + width = |AB| = |CD| = 5 - 1 = 4, + height = |AD| = |BC| = 6 - 4 = 2. + + Next, how about (5, 3, 4, 2, -90), which is rotated by 90 degrees CW (clockwise) + by definition? It looks like this: + + .. code:: none + + O--------> x + | D-A + | | | + | |E| + | | | + | C-B + v y + + The center E is still located at the same point (5, 3), while the vertices + ABCD are rotated by 90 degrees CW with regard to E: + A = (6, 1), B = (6, 5), C = (4, 5), D = (4, 1) + + .. math:: + + width = |AB| = |CD| = 5 - 1 = 4, + height = |AD| = |BC| = 6 - 4 = 2. + + This covers exactly the same region as (5, 3, 4, 2, 90) does, and their IoU + will be 1. However, these two will generate different RoI Pooling results and + should not be treated as an identical box. + + On the other hand, it's easy to see that (X, Y, W, H, A) is identical to + (X, Y, W, H, A+360N), for any integer N. For example (5, 3, 4, 2, 270) would be + identical to (5, 3, 4, 2, -90), because rotating the shape 270 degrees CCW is + equivalent to rotating the same shape 90 degrees CW. + + We could rotate further to get (5, 3, 4, 2, 180), or (5, 3, 4, 2, -180): + + .. code:: none + + O--------> x + | + | C---D + | | E | + | B---A + | + v y + + .. math:: + + A = (7, 4), B = (3, 4), C = (3, 2), D = (7, 2), + + width = |AB| = |CD| = 7 - 3 = 4, + height = |AD| = |BC| = 4 - 2 = 2. + + Finally, this is a very inaccurate (heavily quantized) illustration of + how (5, 3, 4, 2, 60) looks like in case anyone wonders: + + .. code:: none + + O--------> x + | B\ + | / C + | /E / + | A / + | `D + v y + + It's still a rectangle with center of (5, 3), width of 4 and height of 2, + but its angle (and thus orientation) is somewhere between + (5, 3, 4, 2, 0) and (5, 3, 4, 2, 90). + """ + device = ( + tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu") + ) + tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device) + if tensor.numel() == 0: + # Use reshape, so we don't end up creating a new tensor that does not depend on + # the inputs (and consequently confuses jit) + tensor = tensor.reshape((0, 5)).to(dtype=torch.float32, device=device) + assert tensor.dim() == 2 and tensor.size(-1) == 5, tensor.size() + + self.tensor = tensor + + def clone(self) -> "RotatedBoxes": + """ + Clone the RotatedBoxes. + + Returns: + RotatedBoxes + """ + return RotatedBoxes(self.tensor.clone()) + + def to(self, device: torch.device, non_blocking: bool = False): + # Boxes are assumed float32 and does not support to(dtype) + return RotatedBoxes(self.tensor.to(device=device, non_blocking=non_blocking)) + + def area(self) -> torch.Tensor: + """ + Computes the area of all the boxes. + + Returns: + torch.Tensor: a vector with areas of each box. + """ + box = self.tensor + area = box[:, 2] * box[:, 3] + return area + + # Avoid in-place operations so that we can torchscript; NOTE: this creates a new tensor + def normalize_angles(self) -> None: + """ + Restrict angles to the range of [-180, 180) degrees + """ + angle_tensor = (self.tensor[:, 4] + 180.0) % 360.0 - 180.0 + self.tensor = torch.cat((self.tensor[:, :4], angle_tensor[:, None]), dim=1) + + def clip( + self, box_size: Tuple[int, int], clip_angle_threshold: float = 1.0 + ) -> None: + """ + Clip (in place) the boxes by limiting x coordinates to the range [0, width] + and y coordinates to the range [0, height]. + + For RRPN: + Only clip boxes that are almost horizontal with a tolerance of + clip_angle_threshold to maintain backward compatibility. + + Rotated boxes beyond this threshold are not clipped for two reasons: + + 1. There are potentially multiple ways to clip a rotated box to make it + fit within the image. + 2. It's tricky to make the entire rectangular box fit within the image + and still be able to not leave out pixels of interest. + + Therefore we rely on ops like RoIAlignRotated to safely handle this. + + Args: + box_size (height, width): The clipping box's size. + clip_angle_threshold: + Iff. abs(normalized(angle)) <= clip_angle_threshold (in degrees), + we do the clipping as horizontal boxes. + """ + h, w = box_size + + # normalize angles to be within (-180, 180] degrees + self.normalize_angles() + + idx = torch.where(torch.abs(self.tensor[:, 4]) <= clip_angle_threshold)[0] + + # convert to (x1, y1, x2, y2) + x1 = self.tensor[idx, 0] - self.tensor[idx, 2] / 2.0 + y1 = self.tensor[idx, 1] - self.tensor[idx, 3] / 2.0 + x2 = self.tensor[idx, 0] + self.tensor[idx, 2] / 2.0 + y2 = self.tensor[idx, 1] + self.tensor[idx, 3] / 2.0 + + # clip + x1.clamp_(min=0, max=w) + y1.clamp_(min=0, max=h) + x2.clamp_(min=0, max=w) + y2.clamp_(min=0, max=h) + + # convert back to (xc, yc, w, h) + self.tensor[idx, 0] = (x1 + x2) / 2.0 + self.tensor[idx, 1] = (y1 + y2) / 2.0 + # make sure widths and heights do not increase due to numerical errors + self.tensor[idx, 2] = torch.min(self.tensor[idx, 2], x2 - x1) + self.tensor[idx, 3] = torch.min(self.tensor[idx, 3], y2 - y1) + + def nonempty(self, threshold: float = 0.0) -> torch.Tensor: + """ + Find boxes that are non-empty. + A box is considered empty, if either of its side is no larger than threshold. + + Returns: + Tensor: a binary vector which represents + whether each box is empty (False) or non-empty (True). + """ + box = self.tensor + widths = box[:, 2] + heights = box[:, 3] + keep = (widths > threshold) & (heights > threshold) + return keep + + def __getitem__(self, item) -> "RotatedBoxes": + """ + Returns: + RotatedBoxes: Create a new :class:`RotatedBoxes` by indexing. + + The following usage are allowed: + + 1. `new_boxes = boxes[3]`: return a `RotatedBoxes` which contains only one box. + 2. `new_boxes = boxes[2:10]`: return a slice of boxes. + 3. `new_boxes = boxes[vector]`, where vector is a torch.ByteTensor + with `length = len(boxes)`. Nonzero elements in the vector will be selected. + + Note that the returned RotatedBoxes might share storage with this RotatedBoxes, + subject to Pytorch's indexing semantics. + """ + if isinstance(item, int): + return RotatedBoxes(self.tensor[item].view(1, -1)) + b = self.tensor[item] + assert ( + b.dim() == 2 + ), "Indexing on RotatedBoxes with {} failed to return a matrix!".format(item) + return RotatedBoxes(b) + + def __len__(self) -> int: + return self.tensor.shape[0] + + def __repr__(self) -> str: + return "RotatedBoxes(" + str(self.tensor) + ")" + + def inside_box( + self, box_size: Tuple[int, int], boundary_threshold: int = 0 + ) -> torch.Tensor: + """ + Args: + box_size (height, width): Size of the reference box covering + [0, width] x [0, height] + boundary_threshold (int): Boxes that extend beyond the reference box + boundary by more than boundary_threshold are considered "outside". + + For RRPN, it might not be necessary to call this function since it's common + for rotated box to extend to outside of the image boundaries + (the clip function only clips the near-horizontal boxes) + + Returns: + a binary vector, indicating whether each box is inside the reference box. + """ + height, width = box_size + + cnt_x = self.tensor[..., 0] + cnt_y = self.tensor[..., 1] + half_w = self.tensor[..., 2] / 2.0 + half_h = self.tensor[..., 3] / 2.0 + a = self.tensor[..., 4] + c = torch.abs(torch.cos(a * math.pi / 180.0)) + s = torch.abs(torch.sin(a * math.pi / 180.0)) + # This basically computes the horizontal bounding rectangle of the rotated box + max_rect_dx = c * half_w + s * half_h + max_rect_dy = c * half_h + s * half_w + + inds_inside = ( + (cnt_x - max_rect_dx >= -boundary_threshold) + & (cnt_y - max_rect_dy >= -boundary_threshold) + & (cnt_x + max_rect_dx < width + boundary_threshold) + & (cnt_y + max_rect_dy < height + boundary_threshold) + ) + + return inds_inside + + def get_centers(self) -> torch.Tensor: + """ + Returns: + The box centers in a Nx2 array of (x, y). + """ + return self.tensor[:, :2] + + def scale(self, scale_x: float, scale_y: float) -> None: + """ + Scale the rotated box with horizontal and vertical scaling factors + Note: when scale_factor_x != scale_factor_y, + the rotated box does not preserve the rectangular shape when the angle + is not a multiple of 90 degrees under resize transformation. + Instead, the shape is a parallelogram (that has skew) + Here we make an approximation by fitting a rotated rectangle to the parallelogram. + """ + self.tensor[:, 0] *= scale_x + self.tensor[:, 1] *= scale_y + theta = self.tensor[:, 4] * math.pi / 180.0 + c = torch.cos(theta) + s = torch.sin(theta) + + # In image space, y is top->down and x is left->right + # Consider the local coordintate system for the rotated box, + # where the box center is located at (0, 0), and the four vertices ABCD are + # A(-w / 2, -h / 2), B(w / 2, -h / 2), C(w / 2, h / 2), D(-w / 2, h / 2) + # the midpoint of the left edge AD of the rotated box E is: + # E = (A+D)/2 = (-w / 2, 0) + # the midpoint of the top edge AB of the rotated box F is: + # F(0, -h / 2) + # To get the old coordinates in the global system, apply the rotation transformation + # (Note: the right-handed coordinate system for image space is yOx): + # (old_x, old_y) = (s * y + c * x, c * y - s * x) + # E(old) = (s * 0 + c * (-w/2), c * 0 - s * (-w/2)) = (-c * w / 2, s * w / 2) + # F(old) = (s * (-h / 2) + c * 0, c * (-h / 2) - s * 0) = (-s * h / 2, -c * h / 2) + # After applying the scaling factor (sfx, sfy): + # E(new) = (-sfx * c * w / 2, sfy * s * w / 2) + # F(new) = (-sfx * s * h / 2, -sfy * c * h / 2) + # The new width after scaling tranformation becomes: + + # w(new) = |E(new) - O| * 2 + # = sqrt[(sfx * c * w / 2)^2 + (sfy * s * w / 2)^2] * 2 + # = sqrt[(sfx * c)^2 + (sfy * s)^2] * w + # i.e., scale_factor_w = sqrt[(sfx * c)^2 + (sfy * s)^2] + # + # For example, + # when angle = 0 or 180, |c| = 1, s = 0, scale_factor_w == scale_factor_x; + # when |angle| = 90, c = 0, |s| = 1, scale_factor_w == scale_factor_y + self.tensor[:, 2] *= torch.sqrt((scale_x * c) ** 2 + (scale_y * s) ** 2) + + # h(new) = |F(new) - O| * 2 + # = sqrt[(sfx * s * h / 2)^2 + (sfy * c * h / 2)^2] * 2 + # = sqrt[(sfx * s)^2 + (sfy * c)^2] * h + # i.e., scale_factor_h = sqrt[(sfx * s)^2 + (sfy * c)^2] + # + # For example, + # when angle = 0 or 180, |c| = 1, s = 0, scale_factor_h == scale_factor_y; + # when |angle| = 90, c = 0, |s| = 1, scale_factor_h == scale_factor_x + self.tensor[:, 3] *= torch.sqrt((scale_x * s) ** 2 + (scale_y * c) ** 2) + + # The angle is the rotation angle from y-axis in image space to the height + # vector (top->down in the box's local coordinate system) of the box in CCW. + # + # angle(new) = angle_yOx(O - F(new)) + # = angle_yOx( (sfx * s * h / 2, sfy * c * h / 2) ) + # = atan2(sfx * s * h / 2, sfy * c * h / 2) + # = atan2(sfx * s, sfy * c) + # + # For example, + # when sfx == sfy, angle(new) == atan2(s, c) == angle(old) + self.tensor[:, 4] = torch.atan2(scale_x * s, scale_y * c) * 180 / math.pi + + @classmethod + def cat(cls, boxes_list: List["RotatedBoxes"]) -> "RotatedBoxes": + """ + Concatenates a list of RotatedBoxes into a single RotatedBoxes + + Arguments: + boxes_list (list[RotatedBoxes]) + + Returns: + RotatedBoxes: the concatenated RotatedBoxes + """ + assert isinstance(boxes_list, (list, tuple)) + if len(boxes_list) == 0: + return cls(torch.empty(0)) + assert all([isinstance(box, RotatedBoxes) for box in boxes_list]) + + # use torch.cat (v.s. layers.cat) so the returned boxes never share storage with input + cat_boxes = cls(torch.cat([b.tensor for b in boxes_list], dim=0)) + return cat_boxes + + @property + def device(self) -> torch.device: + return self.tensor.device + + @torch.jit.unused + def __iter__(self): + """ + Yield a box as a Tensor of shape (5,) at a time. + """ + yield from self.tensor + + +def pairwise_iou(boxes1: RotatedBoxes, boxes2: RotatedBoxes) -> None: + """ + Given two lists of rotated boxes of size N and M, + compute the IoU (intersection over union) + between **all** N x M pairs of boxes. + The box order must be (x_center, y_center, width, height, angle). + + Args: + boxes1, boxes2 (RotatedBoxes): + two `RotatedBoxes`. Contains N & M rotated boxes, respectively. + + Returns: + Tensor: IoU, sized [N,M]. + """ + + return pairwise_iou_rotated(boxes1.tensor, boxes2.tensor) diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/som_utils.py b/src/nn/segearth_ov3/sam3/agent/helpers/som_utils.py new file mode 100644 index 0000000..0ca96d6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/som_utils.py @@ -0,0 +1,406 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import colorsys +from dataclasses import dataclass +from typing import List, Tuple + +import cv2 +import matplotlib as mpl +import matplotlib.colors as mplc +import numpy as np +import pycocotools.mask as mask_utils + + +def rgb_to_hex(rgb_color): + """ + Convert a rgb color to hex color. + + Args: + rgb_color (tuple/list of ints): RGB color in tuple or list format. + + Returns: + str: Hex color. + + Example: + ``` + >>> rgb_to_hex((255, 0, 244)) + '#ff00ff' + ``` + """ + return "#" + "".join([hex(c)[2:].zfill(2) for c in rgb_color]) + + +# DEFAULT_COLOR_HEX_TO_NAME = { +# rgb_to_hex((255, 0, 0)): "red", +# rgb_to_hex((0, 255, 0)): "lime", +# rgb_to_hex((0, 0, 255)): "blue", +# rgb_to_hex((255, 255, 0)): "yellow", +# rgb_to_hex((255, 0, 255)): "fuchsia", +# rgb_to_hex((0, 255, 255)): "aqua", +# rgb_to_hex((255, 165, 0)): "orange", +# rgb_to_hex((128, 0, 128)): "purple", +# rgb_to_hex((255, 215, 0)): "gold", +# } + +# Assuming rgb_to_hex is a function that converts an (R, G, B) tuple to a hex string. +# For example: def rgb_to_hex(rgb): return '#%02x%02x%02x' % rgb + +DEFAULT_COLOR_HEX_TO_NAME = { + # The top 20 approved colors + rgb_to_hex((255, 255, 0)): "yellow", + rgb_to_hex((0, 255, 0)): "lime", + rgb_to_hex((0, 255, 255)): "cyan", + rgb_to_hex((255, 0, 255)): "magenta", + rgb_to_hex((255, 0, 0)): "red", + rgb_to_hex((255, 127, 0)): "orange", + rgb_to_hex((127, 255, 0)): "chartreuse", + rgb_to_hex((0, 255, 127)): "spring green", + rgb_to_hex((255, 0, 127)): "rose", + rgb_to_hex((127, 0, 255)): "violet", + rgb_to_hex((192, 255, 0)): "electric lime", + rgb_to_hex((255, 192, 0)): "vivid orange", + rgb_to_hex((0, 255, 192)): "turquoise", + rgb_to_hex((192, 0, 255)): "bright violet", + rgb_to_hex((255, 0, 192)): "bright pink", + rgb_to_hex((255, 64, 0)): "fiery orange", + rgb_to_hex((64, 255, 0)): "bright chartreuse", + rgb_to_hex((0, 255, 64)): "malachite", + rgb_to_hex((64, 0, 255)): "deep violet", + rgb_to_hex((255, 0, 64)): "hot pink", +} + + +DEFAULT_COLOR_PALETTE = list(DEFAULT_COLOR_HEX_TO_NAME.keys()) + + +def _validate_color_hex(color_hex: str): + color_hex = color_hex.lstrip("#") + if not all(c in "0123456789abcdefABCDEF" for c in color_hex): + raise ValueError("Invalid characters in color hash") + if len(color_hex) not in (3, 6): + raise ValueError("Invalid length of color hash") + + +# copied from https://github.com/roboflow/supervision/blob/c8f557af0c61b5c03392bad2cc36c8835598b1e1/supervision/draw/color.py +@dataclass +class Color: + """ + Represents a color in RGB format. + + Attributes: + r (int): Red channel. + g (int): Green channel. + b (int): Blue channel. + """ + + r: int + g: int + b: int + + @classmethod + def from_hex(cls, color_hex: str): + """ + Create a Color instance from a hex string. + + Args: + color_hex (str): Hex string of the color. + + Returns: + Color: Instance representing the color. + + Example: + ``` + >>> Color.from_hex('#ff00ff') + Color(r=255, g=0, b=255) + ``` + """ + _validate_color_hex(color_hex) + color_hex = color_hex.lstrip("#") + if len(color_hex) == 3: + color_hex = "".join(c * 2 for c in color_hex) + r, g, b = (int(color_hex[i : i + 2], 16) for i in range(0, 6, 2)) + return cls(r, g, b) + + @classmethod + def to_hex(cls, color): + """ + Convert a Color instance to a hex string. + + Args: + color (Color): Color instance of color. + + Returns: + Color: a hex string. + """ + return rgb_to_hex((color.r, color.g, color.b)) + + def as_rgb(self) -> Tuple[int, int, int]: + """ + Returns the color as an RGB tuple. + + Returns: + Tuple[int, int, int]: RGB tuple. + + Example: + ``` + >>> color.as_rgb() + (255, 0, 255) + ``` + """ + return self.r, self.g, self.b + + def as_bgr(self) -> Tuple[int, int, int]: + """ + Returns the color as a BGR tuple. + + Returns: + Tuple[int, int, int]: BGR tuple. + + Example: + ``` + >>> color.as_bgr() + (255, 0, 255) + ``` + """ + return self.b, self.g, self.r + + @classmethod + def white(cls): + return Color.from_hex(color_hex="#ffffff") + + @classmethod + def black(cls): + return Color.from_hex(color_hex="#000000") + + @classmethod + def red(cls): + return Color.from_hex(color_hex="#ff0000") + + @classmethod + def green(cls): + return Color.from_hex(color_hex="#00ff00") + + @classmethod + def blue(cls): + return Color.from_hex(color_hex="#0000ff") + + +@dataclass +class ColorPalette: + colors: List[Color] + + @classmethod + def default(cls): + """ + Returns a default color palette. + + Returns: + ColorPalette: A ColorPalette instance with default colors. + + Example: + ``` + >>> ColorPalette.default() + ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...]) + ``` + """ + return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE) + + @classmethod + def from_hex(cls, color_hex_list: List[str]): + """ + Create a ColorPalette instance from a list of hex strings. + + Args: + color_hex_list (List[str]): List of color hex strings. + + Returns: + ColorPalette: A ColorPalette instance. + + Example: + ``` + >>> ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff']) + ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...]) + ``` + """ + colors = [Color.from_hex(color_hex) for color_hex in color_hex_list] + return cls(colors) + + def by_idx(self, idx: int) -> Color: + """ + Return the color at a given index in the palette. + + Args: + idx (int): Index of the color in the palette. + + Returns: + Color: Color at the given index. + + Example: + ``` + >>> color_palette.by_idx(1) + Color(r=0, g=255, b=0) + ``` + """ + if idx < 0: + raise ValueError("idx argument should not be negative") + idx = idx % len(self.colors) + return self.colors[idx] + + def find_farthest_color(self, img_array): + """ + Return the color that is the farthest from the given color. + + Args: + img_array (np array): any *x3 np array, 3 is the RGB color channel. + + Returns: + Color: Farthest color. + + """ + # Reshape the image array for broadcasting + img_array = img_array.reshape((-1, 3)) + + # Convert colors dictionary to a NumPy array + color_values = np.array([[c.r, c.g, c.b] for c in self.colors]) + + # Calculate the Euclidean distance between the colors and each pixel in the image + # Broadcasting happens here: img_array shape is (num_pixels, 3), color_values shape is (num_colors, 3) + distances = np.sqrt( + np.sum((img_array[:, np.newaxis, :] - color_values) ** 2, axis=2) + ) + + # Average the distances for each color + mean_distances = np.mean(distances, axis=0) + + # return the farthest color + farthest_idx = np.argmax(mean_distances) + farthest_color = self.colors[farthest_idx] + farthest_color_hex = Color.to_hex(farthest_color) + if farthest_color_hex in DEFAULT_COLOR_HEX_TO_NAME: + farthest_color_name = DEFAULT_COLOR_HEX_TO_NAME[farthest_color_hex] + else: + farthest_color_name = "unknown" + + return farthest_color, farthest_color_name + + +def draw_box(ax, box_coord, alpha=0.8, edge_color="g", line_style="-", linewidth=2.0): + x0, y0, width, height = box_coord + ax.add_patch( + mpl.patches.Rectangle( + (x0, y0), + width, + height, + fill=False, + edgecolor=edge_color, + linewidth=linewidth, + alpha=alpha, + linestyle=line_style, + ) + ) + + +def draw_text( + ax, + text, + position, + font_size=None, + color="g", + horizontal_alignment="left", + rotation=0, +): + if not font_size: + font_size = mpl.rcParams["font.size"] + + color = np.maximum(list(mplc.to_rgb(color)), 0.2) + color[np.argmax(color)] = max(0.8, np.max(color)) + + x, y = position + ax.text( + x, + y, + text, + size=font_size, + family="sans-serif", + bbox={"facecolor": "none", "alpha": 0.5, "pad": 0.7, "edgecolor": "none"}, + verticalalignment="top", + horizontalalignment=horizontal_alignment, + color=color, + rotation=rotation, + ) + + +def draw_mask( + ax, rle, color, show_holes=True, alpha=0.15, upsample_factor=1.0, rle_upsampled=None +): + if isinstance(rle, dict): + mask = mask_utils.decode(rle) + elif isinstance(rle, np.ndarray): + mask = rle + else: + raise ValueError(f"Unsupported type for rle: {type(rle)}") + + mask_upsampled = None + if upsample_factor > 1.0 and show_holes: + assert rle_upsampled is not None + if isinstance(rle_upsampled, dict): + mask_upsampled = mask_utils.decode(rle_upsampled) + elif isinstance(rle_upsampled, np.ndarray): + mask_upsampled = rle_upsampled + else: + raise ValueError(f"Unsupported type for rle: {type(rle)}") + + if show_holes: + if mask_upsampled is None: + mask_upsampled = mask + h, w = mask_upsampled.shape + mask_img = np.zeros((h, w, 4)) + mask_img[:, :, :-1] = color[np.newaxis, np.newaxis, :] + mask_img[:, :, -1] = mask_upsampled * alpha + ax.imshow(mask_img) + + *_, contours, _ = cv2.findContours( + mask.astype(np.uint8).copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE + ) + upsampled_contours = [(cont + 0.5) * upsample_factor - 0.5 for cont in contours] + facecolor = (0, 0, 0, 0) if show_holes else color + if alpha > 0.8: + edge_color = _change_color_brightness(color, brightness_factor=-0.7) + else: + edge_color = color + for cont in upsampled_contours: + polygon = mpl.patches.Polygon( + [el[0] for el in cont], + edgecolor=edge_color, + linewidth=2.0, + facecolor=facecolor, + ) + ax.add_patch(polygon) + + +def _change_color_brightness(color, brightness_factor): + """ + Depending on the brightness_factor, gives a lighter or darker color i.e. a color with + less or more saturation than the original color. + + Args: + color: color of the polygon. Refer to `matplotlib.colors` for a full list of + formats that are accepted. + brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of + 0 will correspond to no change, a factor in [-1.0, 0) range will result in + a darker color and a factor in (0, 1.0] range will result in a lighter color. + + Returns: + modified_color (tuple[double]): a tuple containing the RGB values of the + modified color. Each value in the tuple is in the [0.0, 1.0] range. + """ + assert brightness_factor >= -1.0 and brightness_factor <= 1.0 + color = mplc.to_rgb(color) + polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color)) + modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1]) + modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness + modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness + modified_color = colorsys.hls_to_rgb( + polygon_color[0], modified_lightness, polygon_color[2] + ) + return modified_color diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/visualizer.py b/src/nn/segearth_ov3/sam3/agent/helpers/visualizer.py new file mode 100755 index 0000000..3e134d3 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/visualizer.py @@ -0,0 +1,1662 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import colorsys +import logging +import math +import random +from enum import Enum, unique + +import cv2 +import matplotlib as mpl +import matplotlib.colors as mplc +import matplotlib.figure as mplfigure +import numpy as np +import pycocotools.mask as mask_util +import torch +from iopath.common.file_io import PathManager +from matplotlib.backends.backend_agg import FigureCanvasAgg +from PIL import Image + +from .boxes import Boxes, BoxMode + +from .color_map import random_color +from .keypoints import Keypoints +from .masks import BitMasks, PolygonMasks +from .rotated_boxes import RotatedBoxes + +logger = logging.getLogger(__name__) + + +__all__ = ["ColorMode", "VisImage", "Visualizer"] + + +_SMALL_OBJECT_AREA_THRESH = 1000 +_LARGE_MASK_AREA_THRESH = 120000 +_OFF_WHITE = (1.0, 1.0, 240.0 / 255) +_BLACK = (0, 0, 0) +_RED = (1.0, 0, 0) + +_KEYPOINT_THRESHOLD = 0.05 + + +@unique +class ColorMode(Enum): + """ + Enum of different color modes to use for instance visualizations. + """ + + IMAGE = 0 + """ + Picks a random color for every instance and overlay segmentations with low opacity. + """ + SEGMENTATION = 1 + """ + Let instances of the same category have similar colors + (from metadata.thing_colors), and overlay them with + high opacity. This provides more attention on the quality of segmentation. + """ + IMAGE_BW = 2 + """ + Same as IMAGE, but convert all areas without masks to gray-scale. + Only available for drawing per-instance mask predictions. + """ + + +class GenericMask: + """ + Attribute: + polygons (list[ndarray]): list[ndarray]: polygons for this mask. + Each ndarray has format [x, y, x, y, ...] + mask (ndarray): a binary mask + """ + + def __init__(self, mask_or_polygons, height, width): + self._mask = self._polygons = self._has_holes = None + self.height = height + self.width = width + + m = mask_or_polygons + if isinstance(m, dict): + # RLEs + assert "counts" in m and "size" in m + if isinstance(m["counts"], list): # uncompressed RLEs + h, w = m["size"] + assert h == height and w == width + m = mask_util.frPyObjects(m, h, w) + self._mask = mask_util.decode(m)[:, :] + return + + if isinstance(m, list): # list[ndarray] + self._polygons = [np.asarray(x).reshape(-1) for x in m] + return + + if isinstance(m, np.ndarray): # assumed to be a binary mask + assert m.shape[1] != 2, m.shape + assert m.shape == ( + height, + width, + ), f"mask shape: {m.shape}, target dims: {height}, {width}" + self._mask = m.astype("uint8") + return + + raise ValueError( + "GenericMask cannot handle object {} of type '{}'".format(m, type(m)) + ) + + @property + def mask(self): + if self._mask is None: + self._mask = self.polygons_to_mask(self._polygons) + return self._mask + + @property + def polygons(self): + if self._polygons is None: + self._polygons, self._has_holes = self.mask_to_polygons(self._mask) + return self._polygons + + @property + def has_holes(self): + if self._has_holes is None: + if self._mask is not None: + self._polygons, self._has_holes = self.mask_to_polygons(self._mask) + else: + self._has_holes = ( + False # if original format is polygon, does not have holes + ) + return self._has_holes + + def mask_to_polygons(self, mask): + # cv2.RETR_CCOMP flag retrieves all the contours and arranges them to a 2-level + # hierarchy. External contours (boundary) of the object are placed in hierarchy-1. + # Internal contours (holes) are placed in hierarchy-2. + # cv2.CHAIN_APPROX_NONE flag gets vertices of polygons from contours. + mask = np.ascontiguousarray( + mask + ) # some versions of cv2 does not support incontiguous arr + res = cv2.findContours( + mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE + ) + hierarchy = res[-1] + if hierarchy is None: # empty mask + return [], False + has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0 + res = res[-2] + res = [x.flatten() for x in res] + # These coordinates from OpenCV are integers in range [0, W-1 or H-1]. + # We add 0.5 to turn them into real-value coordinate space. A better solution + # would be to first +0.5 and then dilate the returned polygon by 0.5. + res = [x + 0.5 for x in res if len(x) >= 6] + return res, has_holes + + def polygons_to_mask(self, polygons): + rle = mask_util.frPyObjects(polygons, self.height, self.width) + rle = mask_util.merge(rle) + return mask_util.decode(rle)[:, :] + + def area(self): + return self.mask.sum() + + def bbox(self): + p = mask_util.frPyObjects(self.polygons, self.height, self.width) + p = mask_util.merge(p) + bbox = mask_util.toBbox(p) + bbox[2] += bbox[0] + bbox[3] += bbox[1] + return bbox + + +class _PanopticPrediction: + """ + Unify different panoptic annotation/prediction formats + """ + + def __init__(self, panoptic_seg, segments_info, metadata=None): + if segments_info is None: + assert metadata is not None + # If "segments_info" is None, we assume "panoptic_img" is a + # H*W int32 image storing the panoptic_id in the format of + # category_id * label_divisor + instance_id. We reserve -1 for + # VOID label. + label_divisor = metadata.label_divisor + segments_info = [] + for panoptic_label in np.unique(panoptic_seg.numpy()): + if panoptic_label == -1: + # VOID region. + continue + pred_class = panoptic_label // label_divisor + isthing = ( + pred_class in metadata.thing_dataset_id_to_contiguous_id.values() + ) + segments_info.append( + { + "id": int(panoptic_label), + "category_id": int(pred_class), + "isthing": bool(isthing), + } + ) + del metadata + + self._seg = panoptic_seg + + self._sinfo = {s["id"]: s for s in segments_info} # seg id -> seg info + segment_ids, areas = torch.unique(panoptic_seg, sorted=True, return_counts=True) + areas = areas.numpy() + sorted_idxs = np.argsort(-areas) + self._seg_ids, self._seg_areas = segment_ids[sorted_idxs], areas[sorted_idxs] + self._seg_ids = self._seg_ids.tolist() + for sid, area in zip(self._seg_ids, self._seg_areas): + if sid in self._sinfo: + self._sinfo[sid]["area"] = float(area) + + def non_empty_mask(self): + """ + Returns: + (H, W) array, a mask for all pixels that have a prediction + """ + empty_ids = [] + for id in self._seg_ids: + if id not in self._sinfo: + empty_ids.append(id) + if len(empty_ids) == 0: + return np.zeros(self._seg.shape, dtype=np.uint8) + assert ( + len(empty_ids) == 1 + ), ">1 ids corresponds to no labels. This is currently not supported" + return (self._seg != empty_ids[0]).numpy().astype(np.bool) + + def semantic_masks(self): + for sid in self._seg_ids: + sinfo = self._sinfo.get(sid) + if sinfo is None or sinfo["isthing"]: + # Some pixels (e.g. id 0 in PanopticFPN) have no instance or semantic predictions. + continue + yield (self._seg == sid).numpy().astype(np.bool), sinfo + + def instance_masks(self): + for sid in self._seg_ids: + sinfo = self._sinfo.get(sid) + if sinfo is None or not sinfo["isthing"]: + continue + mask = (self._seg == sid).numpy().astype(np.bool) + if mask.sum() > 0: + yield mask, sinfo + + +def _create_text_labels(classes, scores, class_names, is_crowd=None): + """ + Args: + classes (list[int] or None): + scores (list[float] or None): + class_names (list[str] or None): + is_crowd (list[bool] or None): + + Returns: + list[str] or None + """ + labels = None + if classes is not None: + if class_names is not None and len(class_names) > 0: + labels = [class_names[i] for i in classes] + else: + labels = [str(i) for i in classes] + if scores is not None: + if labels is None: + labels = ["{:.0f}%".format(s * 100) for s in scores] + else: + labels = ["{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores)] + if labels is not None and is_crowd is not None: + labels = [l + ("|crowd" if crowd else "") for l, crowd in zip(labels, is_crowd)] + return labels + + +class VisImage: + def __init__(self, img, scale=1.0): + """ + Args: + img (ndarray): an RGB image of shape (H, W, 3) in range [0, 255]. + scale (float): scale the input image + """ + self.img = img + self.scale = scale + self.width, self.height = img.shape[1], img.shape[0] + self._setup_figure(img) + + def _setup_figure(self, img): + """ + Args: + Same as in :meth:`__init__()`. + + Returns: + fig (matplotlib.pyplot.figure): top level container for all the image plot elements. + ax (matplotlib.pyplot.Axes): contains figure elements and sets the coordinate system. + """ + fig = mplfigure.Figure(frameon=False) + self.dpi = fig.get_dpi() + # add a small 1e-2 to avoid precision lost due to matplotlib's truncation + # (https://github.com/matplotlib/matplotlib/issues/15363) + fig.set_size_inches( + (self.width * self.scale + 1e-2) / self.dpi, + (self.height * self.scale + 1e-2) / self.dpi, + ) + self.canvas = FigureCanvasAgg(fig) + # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig) + ax = fig.add_axes([0.0, 0.0, 1.0, 1.0]) + ax.axis("off") + self.fig = fig + self.ax = ax + self.reset_image(img) + + def reset_image(self, img): + """ + Args: + img: same as in __init__ + """ + img = img.astype("uint8") + self.ax.imshow( + img, extent=(0, self.width, self.height, 0), interpolation="nearest" + ) + + def save(self, filepath): + """ + Args: + filepath (str): a string that contains the absolute path, including the file name, where + the visualized image will be saved. + """ + self.fig.savefig(filepath) + + def get_image(self): + """ + Returns: + ndarray: + the visualized image of shape (H, W, 3) (RGB) in uint8 type. + The shape is scaled w.r.t the input image using the given `scale` argument. + """ + canvas = self.canvas + s, (width, height) = canvas.print_to_buffer() + # buf = io.BytesIO() # works for cairo backend + # canvas.print_rgba(buf) + # width, height = self.width, self.height + # s = buf.getvalue() + + buffer = np.frombuffer(s, dtype="uint8") + + img_rgba = buffer.reshape(height, width, 4) + rgb, alpha = np.split(img_rgba, [3], axis=2) + return rgb.astype("uint8") + + +class Visualizer: + """ + Visualizer that draws data about detection/segmentation on images. + + It contains methods like `draw_{text,box,circle,line,binary_mask,polygon}` + that draw primitive objects to images, as well as high-level wrappers like + `draw_{instance_predictions,sem_seg,panoptic_seg_predictions,dataset_dict}` + that draw composite data in some pre-defined style. + + Note that the exact visualization style for the high-level wrappers are subject to change. + Style such as color, opacity, label contents, visibility of labels, or even the visibility + of objects themselves (e.g. when the object is too small) may change according + to different heuristics, as long as the results still look visually reasonable. + + To obtain a consistent style, you can implement custom drawing functions with the + abovementioned primitive methods instead. If you need more customized visualization + styles, you can process the data yourself following their format documented in + tutorials (:doc:`/tutorials/models`, :doc:`/tutorials/datasets`). This class does not + intend to satisfy everyone's preference on drawing styles. + + This visualizer focuses on high rendering quality rather than performance. It is not + designed to be used for real-time applications. + """ + + def __init__( + self, + img_rgb, + metadata=None, + scale=1.0, + instance_mode=ColorMode.IMAGE, + font_size_multiplier=1.3, + boarder_width_multiplier=1.5, + ): + """ + Args: + img_rgb: a numpy array of shape (H, W, C), where H and W correspond to + the height and width of the image respectively. C is the number of + color channels. The image is required to be in RGB format since that + is a requirement of the Matplotlib library. The image is also expected + to be in the range [0, 255]. + metadata (Metadata): dataset metadata (e.g. class names and colors) + instance_mode (ColorMode): defines one of the pre-defined style for drawing + instances on an image. + """ + self.img = np.asarray(img_rgb).clip(0, 255).astype(np.uint8) + self.boarder_width_multiplier = boarder_width_multiplier + # if metadata is None: + # metadata = MetadataCatalog.get("__nonexist__") + # self.metadata = metadata + self.output = VisImage(self.img, scale=scale) + self.cpu_device = torch.device("cpu") + + # too small texts are useless, therefore clamp to 9 + self._default_font_size = ( + max(np.sqrt(self.output.height * self.output.width) // 60, 15 // scale) + * font_size_multiplier + ) + # self._default_font_size = 18 + self._instance_mode = instance_mode + self.keypoint_threshold = _KEYPOINT_THRESHOLD + + import matplotlib.colors as mcolors + + css4_colors = mcolors.CSS4_COLORS + self.color_proposals = [ + list(mcolors.hex2color(color)) for color in css4_colors.values() + ] + + def draw_instance_predictions(self, predictions): + """ + Draw instance-level prediction results on an image. + + Args: + predictions (Instances): the output of an instance detection/segmentation + model. Following fields will be used to draw: + "pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle"). + + Returns: + output (VisImage): image object with visualizations. + """ + boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None + scores = predictions.scores if predictions.has("scores") else None + classes = ( + predictions.pred_classes.tolist() + if predictions.has("pred_classes") + else None + ) + labels = _create_text_labels( + classes, scores, self.metadata.get("thing_classes", None) + ) + keypoints = ( + predictions.pred_keypoints if predictions.has("pred_keypoints") else None + ) + + keep = (scores > 0.5).cpu() + boxes = boxes[keep] + scores = scores[keep] + classes = np.array(classes) + classes = classes[np.array(keep)] + labels = np.array(labels) + labels = labels[np.array(keep)] + + if predictions.has("pred_masks"): + masks = np.asarray(predictions.pred_masks) + masks = masks[np.array(keep)] + masks = [ + GenericMask(x, self.output.height, self.output.width) for x in masks + ] + else: + masks = None + + if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get( + "thing_colors" + ): + # if self.metadata.get("thing_colors"): + colors = [ + self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) + for c in classes + ] + alpha = 0.4 + else: + colors = None + alpha = 0.4 + + if self._instance_mode == ColorMode.IMAGE_BW: + self.output.reset_image( + self._create_grayscale_image( + (predictions.pred_masks.any(dim=0) > 0).numpy() + if predictions.has("pred_masks") + else None + ) + ) + alpha = 0.3 + + self.overlay_instances( + masks=masks, + boxes=boxes, + labels=labels, + keypoints=keypoints, + assigned_colors=colors, + alpha=alpha, + ) + return self.output + + def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.7): + """ + Draw semantic segmentation predictions/labels. + + Args: + sem_seg (Tensor or ndarray): the segmentation of shape (H, W). + Each value is the integer label of the pixel. + area_threshold (int): segments with less than `area_threshold` are not drawn. + alpha (float): the larger it is, the more opaque the segmentations are. + + Returns: + output (VisImage): image object with visualizations. + """ + if isinstance(sem_seg, torch.Tensor): + sem_seg = sem_seg.numpy() + labels, areas = np.unique(sem_seg, return_counts=True) + sorted_idxs = np.argsort(-areas).tolist() + labels = labels[sorted_idxs] + for label in filter(lambda l: l < len(self.metadata.stuff_classes), labels): + try: + mask_color = [x / 255 for x in self.metadata.stuff_colors[label]] + except (AttributeError, IndexError): + mask_color = None + + binary_mask = (sem_seg == label).astype(np.uint8) + text = self.metadata.stuff_classes[label] + self.draw_binary_mask( + binary_mask, + color=mask_color, + edge_color=_OFF_WHITE, + text=text, + alpha=alpha, + area_threshold=area_threshold, + ) + return self.output + + def draw_panoptic_seg( + self, panoptic_seg, segments_info, area_threshold=None, alpha=0.7 + ): + """ + Draw panoptic prediction annotations or results. + + Args: + panoptic_seg (Tensor): of shape (height, width) where the values are ids for each + segment. + segments_info (list[dict] or None): Describe each segment in `panoptic_seg`. + If it is a ``list[dict]``, each dict contains keys "id", "category_id". + If None, category id of each pixel is computed by + ``pixel // metadata.label_divisor``. + area_threshold (int): stuff segments with less than `area_threshold` are not drawn. + + Returns: + output (VisImage): image object with visualizations. + """ + pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata) + + if self._instance_mode == ColorMode.IMAGE_BW: + self.output.reset_image(self._create_grayscale_image(pred.non_empty_mask())) + + # draw mask for all semantic segments first i.e. "stuff" + for mask, sinfo in pred.semantic_masks(): + category_idx = sinfo["category_id"] + try: + mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]] + except AttributeError: + mask_color = None + + text = ( + self.metadata.stuff_classes[category_idx] + .replace("-other", "") + .replace("-merged", "") + ) + self.draw_binary_mask( + mask, + color=mask_color, + edge_color=_OFF_WHITE, + text=text, + alpha=alpha, + area_threshold=area_threshold, + ) + + # draw mask for all instances second + all_instances = list(pred.instance_masks()) + if len(all_instances) == 0: + return self.output + masks, sinfo = list(zip(*all_instances)) + category_ids = [x["category_id"] for x in sinfo] + + try: + scores = [x["score"] for x in sinfo] + except KeyError: + scores = None + class_names = [ + name.replace("-other", "").replace("-merged", "") + for name in self.metadata.thing_classes + ] + labels = _create_text_labels( + category_ids, scores, class_names, [x.get("iscrowd", 0) for x in sinfo] + ) + + try: + colors = [ + self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) + for c in category_ids + ] + except AttributeError: + colors = None + self.overlay_instances( + masks=masks, labels=labels, assigned_colors=colors, alpha=alpha + ) + + return self.output + + draw_panoptic_seg_predictions = draw_panoptic_seg # backward compatibility + + def draw_dataset_dict(self, dic): + """ + Draw annotations/segmentaions in Detectron2 Dataset format. + + Args: + dic (dict): annotation/segmentation data of one image, in Detectron2 Dataset format. + + Returns: + output (VisImage): image object with visualizations. + """ + annos = dic.get("annotations", None) + if annos: + if "segmentation" in annos[0]: + masks = [x["segmentation"] for x in annos] + else: + masks = None + if "keypoints" in annos[0]: + keypts = [x["keypoints"] for x in annos] + keypts = np.array(keypts).reshape(len(annos), -1, 3) + else: + keypts = None + + boxes = [ + ( + BoxMode.convert(x["bbox"], x["bbox_mode"], BoxMode.XYXY_ABS) + if len(x["bbox"]) == 4 + else x["bbox"] + ) + for x in annos + ] + + colors = None + category_ids = [x["category_id"] for x in annos] + if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get( + "thing_colors" + ): + colors = [ + self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) + for c in category_ids + ] + names = self.metadata.get("thing_classes", None) + labels = _create_text_labels( + category_ids, + scores=None, + class_names=names, + is_crowd=[x.get("iscrowd", 0) for x in annos], + ) + self.overlay_instances( + labels=labels, + boxes=boxes, + masks=masks, + keypoints=keypts, + assigned_colors=colors, + ) + + sem_seg = dic.get("sem_seg", None) + if sem_seg is None and "sem_seg_file_name" in dic: + with PathManager.open(dic["sem_seg_file_name"], "rb") as f: + sem_seg = Image.open(f) + sem_seg = np.asarray(sem_seg, dtype="uint8") + if sem_seg is not None: + self.draw_sem_seg(sem_seg, area_threshold=0, alpha=0.4) + + pan_seg = dic.get("pan_seg", None) + if pan_seg is None and "pan_seg_file_name" in dic: + with PathManager.open(dic["pan_seg_file_name"], "rb") as f: + pan_seg = Image.open(f) + pan_seg = np.asarray(pan_seg) + from panopticapi.utils import rgb2id + + pan_seg = rgb2id(pan_seg) + if pan_seg is not None: + segments_info = dic["segments_info"] + pan_seg = torch.tensor(pan_seg) + self.draw_panoptic_seg(pan_seg, segments_info, area_threshold=0, alpha=0.7) + return self.output + + def overlay_instances( + self, + *, + boxes=None, + labels=None, + masks=None, + keypoints=None, + assigned_colors=None, + binary_masks=None, + alpha=0.5, + label_mode="1", + ): + """ + Args: + boxes (Boxes, RotatedBoxes or ndarray): either a :class:`Boxes`, + or an Nx4 numpy array of XYXY_ABS format for the N objects in a single image, + or a :class:`RotatedBoxes`, + or an Nx5 numpy array of (x_center, y_center, width, height, angle_degrees) format + for the N objects in a single image, + labels (list[str]): the text to be displayed for each instance. + masks (masks-like object): Supported types are: + + * :class:`detectron2.structures.PolygonMasks`, + :class:`detectron2.structures.BitMasks`. + * list[list[ndarray]]: contains the segmentation masks for all objects in one image. + The first level of the list corresponds to individual instances. The second + level to all the polygon that compose the instance, and the third level + to the polygon coordinates. The third level should have the format of + [x0, y0, x1, y1, ..., xn, yn] (n >= 3). + * list[ndarray]: each ndarray is a binary mask of shape (H, W). + * list[dict]: each dict is a COCO-style RLE. + keypoints (Keypoint or array like): an array-like object of shape (N, K, 3), + where the N is the number of instances and K is the number of keypoints. + The last dimension corresponds to (x, y, visibility or score). + assigned_colors (list[matplotlib.colors]): a list of colors, where each color + corresponds to each mask or box in the image. Refer to 'matplotlib.colors' + for full list of formats that the colors are accepted in. + Returns: + output (VisImage): image object with visualizations. + """ + num_instances = 0 + if boxes is not None: + boxes = self._convert_boxes(boxes) + num_instances = len(boxes) + if masks is not None: + masks = self._convert_masks(masks) + if num_instances: + assert len(masks) == num_instances + else: + num_instances = len(masks) + if keypoints is not None: + if num_instances: + assert len(keypoints) == num_instances + else: + num_instances = len(keypoints) + keypoints = self._convert_keypoints(keypoints) + if labels is not None: + assert len(labels) == num_instances + if assigned_colors is None: + assigned_colors = [ + random_color(rgb=True, maximum=1) for _ in range(num_instances) + ] + if num_instances == 0: + return labels, [], [] + if boxes is not None and boxes.shape[1] == 5: + return self.overlay_rotated_instances( + boxes=boxes, labels=labels, assigned_colors=assigned_colors + ) + + # Display in largest to smallest order to reduce occlusion. + areas = None + if boxes is not None: + areas = np.prod(boxes[:, 2:] - boxes[:, :2], axis=1) + elif masks is not None: + areas = np.asarray([x.area() for x in masks]) + + # if areas is not None: + # # sorted_idxs = np.argsort(areas).tolist() + # sorted_idxs = np.argsort(-areas).tolist() + # # Re-order overlapped instances in descending order. + # boxes = boxes[sorted_idxs] if boxes is not None else None + # labels = [labels[k] for k in sorted_idxs] if labels is not None else None + # masks = [masks[idx] for idx in sorted_idxs] if masks is not None else None + # binary_masks = ( + # [binary_masks[idx] for idx in sorted_idxs] + # if binary_masks is not None + # else None + # ) + # assigned_colors = [assigned_colors[idx] for idx in sorted_idxs] + # keypoints = keypoints[sorted_idxs] if keypoints is not None else None + + marks = [] + marks_position = [] + added_positions = set() + for i in range(num_instances): + color = assigned_colors[i] + if boxes is not None: + self.draw_box(boxes[i], alpha=1, edge_color=color) + if binary_masks is None: + # draw number for non-mask instances + mark = self._draw_number_in_box( + boxes[i], i + 1, color=color, label_mode=label_mode + ) + marks.append(mark) + + if binary_masks is not None: + mark, mask_position = self._draw_number_in_mask( + binary_mask=binary_masks[i].astype("uint8"), + text=i + 1, + color=color, + added_positions=added_positions, + label_mode=label_mode, + ) + marks.append(mark) + marks_position.append(mask_position) + + self.draw_binary_mask( + binary_masks[i], + color=color, + edge_color=_OFF_WHITE, + alpha=alpha, + ) + + if masks is not None: + for segment in masks[i].polygons: + self.draw_polygon( + segment.reshape(-1, 2), color, alpha=0 + ) # alpha=0 so holes in masks are not colored + + # draw keypoints + if keypoints is not None: + for keypoints_per_instance in keypoints: + self.draw_and_connect_keypoints(keypoints_per_instance) + + # return labels, marks, sorted_idxs, marks_position + return labels, marks, marks_position + + def overlay_rotated_instances(self, boxes=None, labels=None, assigned_colors=None): + """ + Args: + boxes (ndarray): an Nx5 numpy array of + (x_center, y_center, width, height, angle_degrees) format + for the N objects in a single image. + labels (list[str]): the text to be displayed for each instance. + assigned_colors (list[matplotlib.colors]): a list of colors, where each color + corresponds to each mask or box in the image. Refer to 'matplotlib.colors' + for full list of formats that the colors are accepted in. + + Returns: + output (VisImage): image object with visualizations. + """ + num_instances = len(boxes) + + if assigned_colors is None: + assigned_colors = [ + random_color(rgb=True, maximum=1) for _ in range(num_instances) + ] + if num_instances == 0: + return self.output + + # Display in largest to smallest order to reduce occlusion. + if boxes is not None: + areas = boxes[:, 2] * boxes[:, 3] + + sorted_idxs = np.argsort(-areas).tolist() + # Re-order overlapped instances in descending order. + boxes = boxes[sorted_idxs] + labels = [labels[k] for k in sorted_idxs] if labels is not None else None + colors = [assigned_colors[idx] for idx in sorted_idxs] + + for i in range(num_instances): + self.draw_rotated_box_with_label( + boxes[i], + edge_color=colors[i], + label=labels[i] if labels is not None else None, + ) + + return self.output + + def draw_and_connect_keypoints(self, keypoints): + """ + Draws keypoints of an instance and follows the rules for keypoint connections + to draw lines between appropriate keypoints. This follows color heuristics for + line color. + + Args: + keypoints (Tensor): a tensor of shape (K, 3), where K is the number of keypoints + and the last dimension corresponds to (x, y, probability). + + Returns: + output (VisImage): image object with visualizations. + """ + visible = {} + keypoint_names = self.metadata.get("keypoint_names") + for idx, keypoint in enumerate(keypoints): + # draw keypoint + x, y, prob = keypoint + if prob > self.keypoint_threshold: + self.draw_circle((x, y), color=_RED) + if keypoint_names: + keypoint_name = keypoint_names[idx] + visible[keypoint_name] = (x, y) + + if self.metadata.get("keypoint_connection_rules"): + for kp0, kp1, color in self.metadata.keypoint_connection_rules: + if kp0 in visible and kp1 in visible: + x0, y0 = visible[kp0] + x1, y1 = visible[kp1] + color = tuple(x / 255.0 for x in color) + self.draw_line([x0, x1], [y0, y1], color=color) + + # draw lines from nose to mid-shoulder and mid-shoulder to mid-hip + # Note that this strategy is specific to person keypoints. + # For other keypoints, it should just do nothing + try: + ls_x, ls_y = visible["left_shoulder"] + rs_x, rs_y = visible["right_shoulder"] + mid_shoulder_x, mid_shoulder_y = (ls_x + rs_x) / 2, (ls_y + rs_y) / 2 + except KeyError: + pass + else: + # draw line from nose to mid-shoulder + nose_x, nose_y = visible.get("nose", (None, None)) + if nose_x is not None: + self.draw_line( + [nose_x, mid_shoulder_x], [nose_y, mid_shoulder_y], color=_RED + ) + + try: + # draw line from mid-shoulder to mid-hip + lh_x, lh_y = visible["left_hip"] + rh_x, rh_y = visible["right_hip"] + except KeyError: + pass + else: + mid_hip_x, mid_hip_y = (lh_x + rh_x) / 2, (lh_y + rh_y) / 2 + self.draw_line( + [mid_hip_x, mid_shoulder_x], [mid_hip_y, mid_shoulder_y], color=_RED + ) + return self.output + + def mask_dims_from_binary(self, binary_mask): + ind_y, ind_x = np.where(binary_mask == 1) + min_ind_x = np.min(ind_x) + max_ind_x = np.max(ind_x) + min_ind_y = np.min(ind_y) + max_ind_y = np.max(ind_y) + return (max_ind_x - min_ind_x), (max_ind_y - min_ind_y) + + def reposition_label(self, position, cur, binary_mask, move_count): + img_width, img_height = self.output.width, self.output.height + mask_width, mask_height = self.mask_dims_from_binary(binary_mask) + + # set resposition thresholds + mask_width_limit, mask_height_limit = ( + 25, + 25, + ) # limit for width and height size for object covering + location_diff_threshold = 15 # limit for the distance between two labels + x_boundry_limit, y_boundry_limit = ( + 20, + 20, + ) # limit for the distancing the label from edges + + offset_x = 15 # move in x direction + offset_y = 15 # move in y direction + + x1, y1 = position + + if ( + mask_width < mask_width_limit + and mask_height < mask_height_limit + and move_count == 0 + ): + move_x = offset_x if offset_x + x1 < img_width else -offset_x + move_y = offset_y if offset_y + y1 < img_height else -offset_y + return (True, move_x, move_y) + + for x2, y2 in cur: + if abs(x1 - x2) + abs(y1 - y2) < location_diff_threshold: + move_x = offset_x if x1 >= x2 else -offset_x + move_y = offset_y if y1 >= y2 else -offset_y + move_x = ( + 0 + if x1 + move_x > img_width - x_boundry_limit + or x1 + move_x < x_boundry_limit + else move_x + ) + move_y = ( + 0 + if y1 + move_y > img_height - y_boundry_limit + or y1 + move_y < y_boundry_limit + else move_y + ) + return ( + True, + move_x, + move_y, + ) + return (False, 0, 0) + + def locate_label_position(self, original_position, added_positions, binary_mask): + if added_positions is None or binary_mask is None: + return original_position + + x, y = original_position + + move_count = 0 + reposition, x_move, y_move = self.reposition_label( + (x, y), added_positions, binary_mask, move_count + ) + while reposition and move_count < 10: + x += x_move + y += y_move + move_count += 1 + reposition, x_move, y_move = self.reposition_label( + (x, y), added_positions, binary_mask, move_count + ) + added_positions.add((x, y)) + return x, y + + """ + Primitive drawing functions: + """ + + def draw_text( + self, + text, + position, + added_positions=None, + binary_mask=None, + *, + font_size=None, + color="g", + horizontal_alignment="center", + rotation=0, + ): + """ + Args: + text (str): class label + position (tuple): a tuple of the x and y coordinates to place text on image. + font_size (int, optional): font of the text. If not provided, a font size + proportional to the image width is calculated and used. + color: color of the text. Refer to `matplotlib.colors` for full list + of formats that are accepted. + horizontal_alignment (str): see `matplotlib.text.Text` + rotation: rotation angle in degrees CCW + + Returns: + output (VisImage): image object with text drawn. + """ + if not font_size: + font_size = self._default_font_size + + # since the text background is dark, we don't want the text to be dark + color = np.maximum(list(mplc.to_rgb(color)), 0.15) + color[np.argmax(color)] = max(0.8, np.max(color)) + + def contrasting_color(rgb): + """Returns 'white' or 'black' depending on which color contrasts more with the given RGB value.""" + + # Decompose the RGB tuple + R, G, B = rgb + + # Calculate the Y value + Y = 0.299 * R + 0.587 * G + 0.114 * B + + # If Y value is greater than 128, it's closer to white so return black. Otherwise, return white. + return "black" if Y > 128 else "white" + + bbox_background = contrasting_color(color * 255) + + x, y = self.locate_label_position( + original_position=position, + added_positions=added_positions, + binary_mask=binary_mask, + ) + + self.output.ax.text( + x, + y, + text, + size=font_size * self.output.scale, + family="sans-serif", + bbox={ + "facecolor": bbox_background, + "alpha": 0.8, + "pad": 0.7, + "edgecolor": "none", + }, + verticalalignment="top", + horizontalalignment=horizontal_alignment, + color=color, + zorder=10, + rotation=rotation, + ) + return self.output + + def draw_box(self, box_coord, alpha=0.5, edge_color="g", line_style="-"): + """ + Args: + box_coord (tuple): a tuple containing x0, y0, x1, y1 coordinates, where x0 and y0 + are the coordinates of the image's top left corner. x1 and y1 are the + coordinates of the image's bottom right corner. + alpha (float): blending efficient. Smaller values lead to more transparent masks. + edge_color: color of the outline of the box. Refer to `matplotlib.colors` + for full list of formats that are accepted. + line_style (string): the string to use to create the outline of the boxes. + + Returns: + output (VisImage): image object with box drawn. + """ + x0, y0, x1, y1 = box_coord + width = x1 - x0 + height = y1 - y0 + + linewidth = max(self._default_font_size / 12, 1) * self.boarder_width_multiplier + + self.output.ax.add_patch( + mpl.patches.Rectangle( + (x0, y0), + width, + height, + fill=False, + edgecolor=edge_color, + linewidth=linewidth * self.output.scale, + alpha=alpha, + linestyle=line_style, + ) + ) + return self.output + + def draw_rotated_box_with_label( + self, rotated_box, alpha=0.5, edge_color="g", line_style="-", label=None + ): + """ + Draw a rotated box with label on its top-left corner. + + Args: + rotated_box (tuple): a tuple containing (cnt_x, cnt_y, w, h, angle), + where cnt_x and cnt_y are the center coordinates of the box. + w and h are the width and height of the box. angle represents how + many degrees the box is rotated CCW with regard to the 0-degree box. + alpha (float): blending efficient. Smaller values lead to more transparent masks. + edge_color: color of the outline of the box. Refer to `matplotlib.colors` + for full list of formats that are accepted. + line_style (string): the string to use to create the outline of the boxes. + label (string): label for rotated box. It will not be rendered when set to None. + + Returns: + output (VisImage): image object with box drawn. + """ + cnt_x, cnt_y, w, h, angle = rotated_box + area = w * h + # use thinner lines when the box is small + linewidth = self._default_font_size / ( + 6 if area < _SMALL_OBJECT_AREA_THRESH * self.output.scale else 3 + ) + + theta = angle * math.pi / 180.0 + c = math.cos(theta) + s = math.sin(theta) + rect = [(-w / 2, h / 2), (-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2)] + # x: left->right ; y: top->down + rotated_rect = [ + (s * yy + c * xx + cnt_x, c * yy - s * xx + cnt_y) for (xx, yy) in rect + ] + for k in range(4): + j = (k + 1) % 4 + self.draw_line( + [rotated_rect[k][0], rotated_rect[j][0]], + [rotated_rect[k][1], rotated_rect[j][1]], + color=edge_color, + linestyle="--" if k == 1 else line_style, + linewidth=linewidth, + ) + + if label is not None: + text_pos = rotated_rect[1] # topleft corner + + height_ratio = h / np.sqrt(self.output.height * self.output.width) + label_color = self._change_color_brightness( + edge_color, brightness_factor=0.7 + ) + font_size = ( + np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) + * 0.5 + * self._default_font_size + ) + self.draw_text( + label, text_pos, color=label_color, font_size=font_size, rotation=angle + ) + + return self.output + + def draw_circle(self, circle_coord, color, radius=3): + """ + Args: + circle_coord (list(int) or tuple(int)): contains the x and y coordinates + of the center of the circle. + color: color of the polygon. Refer to `matplotlib.colors` for a full list of + formats that are accepted. + radius (int): radius of the circle. + + Returns: + output (VisImage): image object with box drawn. + """ + x, y = circle_coord + self.output.ax.add_patch( + mpl.patches.Circle(circle_coord, radius=radius, fill=True, color=color) + ) + return self.output + + def draw_line(self, x_data, y_data, color, linestyle="-", linewidth=None): + """ + Args: + x_data (list[int]): a list containing x values of all the points being drawn. + Length of list should match the length of y_data. + y_data (list[int]): a list containing y values of all the points being drawn. + Length of list should match the length of x_data. + color: color of the line. Refer to `matplotlib.colors` for a full list of + formats that are accepted. + linestyle: style of the line. Refer to `matplotlib.lines.Line2D` + for a full list of formats that are accepted. + linewidth (float or None): width of the line. When it's None, + a default value will be computed and used. + + Returns: + output (VisImage): image object with line drawn. + """ + if linewidth is None: + linewidth = self._default_font_size / 3 + linewidth = max(linewidth, 1) + self.output.ax.add_line( + mpl.lines.Line2D( + x_data, + y_data, + linewidth=linewidth * self.output.scale, + color=color, + linestyle=linestyle, + ) + ) + return self.output + + def draw_binary_mask( + self, + binary_mask, + color=None, + *, + edge_color=None, + text=None, + alpha=0.7, + area_threshold=10, + ): + """ + Args: + binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and + W is the image width. Each value in the array is either a 0 or 1 value of uint8 + type. + color: color of the mask. Refer to `matplotlib.colors` for a full list of + formats that are accepted. If None, will pick a random color. + edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a + full list of formats that are accepted. + text (str): if None, will be drawn on the object + alpha (float): blending efficient. Smaller values lead to more transparent masks. + area_threshold (float): a connected component smaller than this area will not be shown. + + Returns: + output (VisImage): image object with mask drawn. + """ + if color is None: + color = random_color(rgb=True, maximum=1) + color = mplc.to_rgb(color) + + has_valid_segment = False + binary_mask = binary_mask.astype("uint8") # opencv needs uint8 + mask = GenericMask(binary_mask, self.output.height, self.output.width) + shape2d = (binary_mask.shape[0], binary_mask.shape[1]) + + if not mask.has_holes: + # draw polygons for regular masks + for segment in mask.polygons: + area = mask_util.area( + mask_util.frPyObjects([segment], shape2d[0], shape2d[1]) + ) + if area < (area_threshold or 0): + continue + has_valid_segment = True + segment = segment.reshape(-1, 2) + self.draw_polygon( + segment, color=color, edge_color=edge_color, alpha=alpha + ) + else: + # https://stackoverflow.com/questions/8919719/how-to-plot-a-complex-polygon + rgba = np.zeros(shape2d + (4,), dtype="float32") + rgba[:, :, :3] = color + rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha + has_valid_segment = True + self.output.ax.imshow( + rgba, extent=(0, self.output.width, self.output.height, 0) + ) + + if text is not None and has_valid_segment: + lighter_color = self._change_color_brightness(color, brightness_factor=0.7) + self._draw_text_in_mask(binary_mask, text, lighter_color) + return self.output + + def draw_binary_mask_with_number( + self, + binary_mask, + color=None, + *, + edge_color=None, + text=None, + label_mode="1", + alpha=0.1, + anno_mode=["Mask"], + area_threshold=10, + ): + """ + Args: + binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and + W is the image width. Each value in the array is either a 0 or 1 value of uint8 + type. + color: color of the mask. Refer to `matplotlib.colors` for a full list of + formats that are accepted. If None, will pick a random color. + edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a + full list of formats that are accepted. + text (str): if None, will be drawn on the object + alpha (float): blending efficient. Smaller values lead to more transparent masks. + area_threshold (float): a connected component smaller than this area will not be shown. + + Returns: + output (VisImage): image object with mask drawn. + """ + if color is None: + randint = random.randint(0, len(self.color_proposals) - 1) + color = self.color_proposals[randint] + color = mplc.to_rgb(color) + + has_valid_segment = True + binary_mask = binary_mask.astype("uint8") # opencv needs uint8 + mask = GenericMask(binary_mask, self.output.height, self.output.width) + shape2d = (binary_mask.shape[0], binary_mask.shape[1]) + bbox = mask.bbox() + + if "Mask" in anno_mode: + if not mask.has_holes: + # draw polygons for regular masks + for segment in mask.polygons: + area = mask_util.area( + mask_util.frPyObjects([segment], shape2d[0], shape2d[1]) + ) + if area < (area_threshold or 0): + continue + has_valid_segment = True + segment = segment.reshape(-1, 2) + self.draw_polygon( + segment, color=color, edge_color=edge_color, alpha=alpha + ) + else: + # https://stackoverflow.com/questions/8919719/how-to-plot-a-complex-polygon + rgba = np.zeros(shape2d + (4,), dtype="float32") + rgba[:, :, :3] = color + rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha + has_valid_segment = True + self.output.ax.imshow( + rgba, extent=(0, self.output.width, self.output.height, 0) + ) + + if "Box" in anno_mode: + self.draw_box(bbox, edge_color=color, alpha=0.75) + + if "Mark" in anno_mode: + has_valid_segment = True + else: + has_valid_segment = False + + if text is not None and has_valid_segment: + # lighter_color = tuple([x*0.2 for x in color]) + lighter_color = [ + 1, + 1, + 1, + ] # self._change_color_brightness(color, brightness_factor=0.7) + self._draw_number_in_mask( + binary_mask=binary_mask, + text=text, + color=lighter_color, + label_mode=label_mode, + ) + return self.output + + def draw_soft_mask(self, soft_mask, color=None, *, text=None, alpha=0.5): + """ + Args: + soft_mask (ndarray): float array of shape (H, W), each value in [0, 1]. + color: color of the mask. Refer to `matplotlib.colors` for a full list of + formats that are accepted. If None, will pick a random color. + text (str): if None, will be drawn on the object + alpha (float): blending efficient. Smaller values lead to more transparent masks. + + Returns: + output (VisImage): image object with mask drawn. + """ + if color is None: + color = random_color(rgb=True, maximum=1) + color = mplc.to_rgb(color) + + shape2d = (soft_mask.shape[0], soft_mask.shape[1]) + rgba = np.zeros(shape2d + (4,), dtype="float32") + rgba[:, :, :3] = color + rgba[:, :, 3] = soft_mask * alpha + self.output.ax.imshow( + rgba, extent=(0, self.output.width, self.output.height, 0) + ) + + if text is not None: + lighter_color = self._change_color_brightness(color, brightness_factor=0.7) + binary_mask = (soft_mask > 0.5).astype("uint8") + self._draw_text_in_mask(binary_mask, text, lighter_color) + return self.output + + def draw_polygon(self, segment, color, edge_color=None, alpha=0.5): + """ + Args: + segment: numpy array of shape Nx2, containing all the points in the polygon. + color: color of the polygon. Refer to `matplotlib.colors` for a full list of + formats that are accepted. + edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a + full list of formats that are accepted. If not provided, a darker shade + of the polygon color will be used instead. + alpha (float): blending efficient. Smaller values lead to more transparent masks. + + Returns: + output (VisImage): image object with polygon drawn. + """ + if edge_color is None: + # make edge color darker than the polygon color + if alpha > 0.8: + edge_color = self._change_color_brightness( + color, brightness_factor=-0.7 + ) + else: + edge_color = color + edge_color = mplc.to_rgb(edge_color) + (1,) + + polygon = mpl.patches.Polygon( + segment, + fill=True, + facecolor=mplc.to_rgb(color) + (alpha,), + edgecolor=edge_color, + linewidth=max(self._default_font_size // 15 * self.output.scale, 1), + ) + self.output.ax.add_patch(polygon) + return self.output + + """ + Internal methods: + """ + + def _jitter(self, color): + """ + Randomly modifies given color to produce a slightly different color than the color given. + + Args: + color (tuple[double]): a tuple of 3 elements, containing the RGB values of the color + picked. The values in the list are in the [0.0, 1.0] range. + + Returns: + jittered_color (tuple[double]): a tuple of 3 elements, containing the RGB values of the + color after being jittered. The values in the list are in the [0.0, 1.0] range. + """ + color = mplc.to_rgb(color) + # np.random.seed(0) + vec = np.random.rand(3) + # better to do it in another color space + vec = vec / np.linalg.norm(vec) * 0.5 + res = np.clip(vec + color, 0, 1) + return tuple(res) + + def _create_grayscale_image(self, mask=None): + """ + Create a grayscale version of the original image. + The colors in masked area, if given, will be kept. + """ + img_bw = self.img.astype("f4").mean(axis=2) + img_bw = np.stack([img_bw] * 3, axis=2) + if mask is not None: + img_bw[mask] = self.img[mask] + return img_bw + + def _change_color_brightness(self, color, brightness_factor): + """ + Depending on the brightness_factor, gives a lighter or darker color i.e. a color with + less or more saturation than the original color. + + Args: + color: color of the polygon. Refer to `matplotlib.colors` for a full list of + formats that are accepted. + brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of + 0 will correspond to no change, a factor in [-1.0, 0) range will result in + a darker color and a factor in (0, 1.0] range will result in a lighter color. + + Returns: + modified_color (tuple[double]): a tuple containing the RGB values of the + modified color. Each value in the tuple is in the [0.0, 1.0] range. + """ + assert brightness_factor >= -1.0 and brightness_factor <= 1.0 + color = mplc.to_rgb(color) + polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color)) + modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1]) + modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness + modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness + modified_color = colorsys.hls_to_rgb( + polygon_color[0], modified_lightness, polygon_color[2] + ) + return modified_color + + def _convert_boxes(self, boxes): + """ + Convert different format of boxes to an NxB array, where B = 4 or 5 is the box dimension. + """ + if isinstance(boxes, Boxes) or isinstance(boxes, RotatedBoxes): + return boxes.tensor.detach().numpy() + else: + return np.asarray(boxes) + + def _convert_masks(self, masks_or_polygons): + """ + Convert different format of masks or polygons to a tuple of masks and polygons. + + Returns: + list[GenericMask]: + """ + + m = masks_or_polygons + if isinstance(m, PolygonMasks): + m = m.polygons + if isinstance(m, BitMasks): + m = m.tensor.numpy() + if isinstance(m, torch.Tensor): + m = m.numpy() + ret = [] + for x in m: + if isinstance(x, GenericMask): + ret.append(x) + else: + ret.append(GenericMask(x, self.output.height, self.output.width)) + return ret + + def _draw_number_in_box(self, box, text, color, label_mode="1"): + """ + Find proper places to draw text given a box. + """ + x0, y0, x1, y1 = box + text_pos = (x0, y0) # if drawing boxes, put text on the box corner. + horiz_align = "left" + # for small objects, draw text at the side to avoid occlusion + instance_area = (y1 - y0) * (x1 - x0) + if ( + instance_area < _SMALL_OBJECT_AREA_THRESH * self.output.scale + or y1 - y0 < 40 * self.output.scale + ): + if y1 >= self.output.height - 5: + text_pos = (x1, y0) + else: + text_pos = (x0, y1) + + height_ratio = (y1 - y0) / np.sqrt(self.output.height * self.output.width) + lighter_color = self._change_color_brightness(color, brightness_factor=0.7) + font_size = ( + np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) + * 0.65 + * self._default_font_size + ) + if label_mode == "a": + text = self.number_to_string(int(text)) + else: + text = text + self.draw_text( + text, + text_pos, + color=lighter_color, + horizontal_alignment=horiz_align, + font_size=font_size, + ) + + return str(text) + + @staticmethod + def number_to_string(n): + chars = [] + while n: + n, remainder = divmod(n - 1, 26) + chars.append(chr(97 + remainder)) + return "".join(reversed(chars)) + + def _draw_number_in_mask( + self, binary_mask, text, color, added_positions=None, label_mode="1" + ): + """ + Find proper places to draw text given a binary mask. + """ + binary_mask = np.pad(binary_mask, ((1, 1), (1, 1)), "constant") + mask_dt = cv2.distanceTransform(binary_mask, cv2.DIST_L2, 0) + mask_dt = mask_dt[1:-1, 1:-1] + max_dist = np.max(mask_dt) + coords_y, coords_x = np.where(mask_dt == max_dist) # coords is [y, x] + + if label_mode == "a": + text = self.number_to_string(int(text)) + else: + text = text + + text_position = ( + coords_x[len(coords_x) // 2] + 2, + coords_y[len(coords_y) // 2] - 6, + ) + self.draw_text( + text, + text_position, + added_positions=added_positions, + binary_mask=binary_mask, + color=color, + ) + + return str(text), text_position + + # _num_cc, cc_labels, stats, centroids = cv2.connectedComponentsWithStats(binary_mask, 8) + # if stats[1:, -1].size == 0: + # return + # largest_component_id = np.argmax(stats[1:, -1]) + 1 + + # # draw text on the largest component, as well as other very large components. + # for cid in range(1, _num_cc): + # if cid == largest_component_id or stats[cid, -1] > _LARGE_MASK_AREA_THRESH: + # # median is more stable than centroid + # # center = centroids[largest_component_id] + # center = np.median((cc_labels == cid).nonzero(), axis=1)[::-1] + # # bottom=np.max((cc_labels == cid).nonzero(), axis=1)[::-1] + # # center[1]=bottom[1]+2 + # self.draw_text(text, center, color=color) + + def _draw_text_in_mask(self, binary_mask, text, color): + """ + Find proper places to draw text given a binary mask. + """ + _num_cc, cc_labels, stats, centroids = cv2.connectedComponentsWithStats( + binary_mask, 8 + ) + if stats[1:, -1].size == 0: + return + largest_component_id = np.argmax(stats[1:, -1]) + 1 + + # draw text on the largest component, as well as other very large components. + for cid in range(1, _num_cc): + if cid == largest_component_id or stats[cid, -1] > _LARGE_MASK_AREA_THRESH: + # median is more stable than centroid + # center = centroids[largest_component_id] + center = np.median((cc_labels == cid).nonzero(), axis=1)[::-1] + bottom = np.max((cc_labels == cid).nonzero(), axis=1)[::-1] + center[1] = bottom[1] + 2 + self.draw_text(text, center, color=color) + + def _convert_keypoints(self, keypoints): + if isinstance(keypoints, Keypoints): + keypoints = keypoints.tensor + keypoints = np.asarray(keypoints) + return keypoints + + def get_output(self): + """ + Returns: + output (VisImage): the image output containing the visualizations added + to the image. + """ + return self.output diff --git a/src/nn/segearth_ov3/sam3/agent/helpers/zoom_in.py b/src/nn/segearth_ov3/sam3/agent/helpers/zoom_in.py new file mode 100644 index 0000000..dc60700 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/helpers/zoom_in.py @@ -0,0 +1,195 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import io +import math + +import matplotlib.pyplot as plt +import numpy as np +import pycocotools.mask as mask_utils +from PIL import Image + +from .som_utils import ColorPalette, draw_box, draw_mask, draw_text + + +def render_zoom_in( + object_data, + image_file, + show_box: bool = True, + show_text: bool = False, + show_holes: bool = True, + mask_alpha: float = 0.15, +): + """ + Render a two-panel visualization with a cropped original view (left/upper) and a zoomed-in + mask overlay (right/lower), then return it as a PIL.Image along with the chosen mask color (hex). + + Parameters + ---------- + object_data : dict + Dict containing "labels" and COCO RLE "segmentation". + Expected: + object_data["labels"][0]["noun_phrase"] : str + object_data["segmentation"] : COCO RLE (with "size": [H, W]) + image_file : PIL.Image.Image + Source image (PIL). + show_box : bool + Whether to draw the bbox on the cropped original panel. + show_text : bool + Whether to draw the noun phrase label near the bbox. + show_holes : bool + Whether to render mask holes (passed through to draw_mask). + mask_alpha : float + Alpha for the mask overlay. + + Returns + ------- + pil_img : PIL.Image.Image + The composed visualization image. + color_hex : str + Hex string of the chosen mask color. + """ + + # ---- local constants (avoid module-level globals) ---- + _AREA_LARGE = 0.25 + _AREA_MEDIUM = 0.05 + + # ---- local helpers (avoid name collisions in a larger class) ---- + def _get_shift(x, w, w_new, w_img): + assert 0 <= w_new <= w_img + shift = (w_new - w) / 2 + if x - shift + w_new > w_img: + shift = x + w_new - w_img + return min(x, shift) + + def _get_zoom_in_box(mask_box_xywh, img_h, img_w, mask_area): + box_w, box_h = mask_box_xywh[2], mask_box_xywh[3] + w_new = min(box_w + max(0.2 * box_w, 16), img_w) + h_new = min(box_h + max(0.2 * box_h, 16), img_h) + + mask_relative_area = mask_area / (w_new * h_new) + + # zoom-in (larger box if mask is relatively big) + w_new_large, h_new_large = w_new, h_new + if mask_relative_area > _AREA_LARGE: + ratio_large = math.sqrt(mask_relative_area / _AREA_LARGE) + w_new_large = min(w_new * ratio_large, img_w) + h_new_large = min(h_new * ratio_large, img_h) + + w_shift_large = _get_shift( + mask_box_xywh[0], mask_box_xywh[2], w_new_large, img_w + ) + h_shift_large = _get_shift( + mask_box_xywh[1], mask_box_xywh[3], h_new_large, img_h + ) + zoom_in_box = [ + mask_box_xywh[0] - w_shift_large, + mask_box_xywh[1] - h_shift_large, + w_new_large, + h_new_large, + ] + + # crop box for the original/cropped image + w_new_medium, h_new_medium = w_new, h_new + if mask_relative_area > _AREA_MEDIUM: + ratio_med = math.sqrt(mask_relative_area / _AREA_MEDIUM) + w_new_medium = min(w_new * ratio_med, img_w) + h_new_medium = min(h_new * ratio_med, img_h) + + w_shift_medium = _get_shift( + mask_box_xywh[0], mask_box_xywh[2], w_new_medium, img_w + ) + h_shift_medium = _get_shift( + mask_box_xywh[1], mask_box_xywh[3], h_new_medium, img_h + ) + img_crop_box = [ + mask_box_xywh[0] - w_shift_medium, + mask_box_xywh[1] - h_shift_medium, + w_new_medium, + h_new_medium, + ] + return zoom_in_box, img_crop_box + + # ---- main body ---- + # Input parsing + object_label = object_data["labels"][0]["noun_phrase"] + img = image_file.convert("RGB") + bbox_xywh = mask_utils.toBbox(object_data["segmentation"]) # [x, y, w, h] + + # Choose a stable, visually distant color based on crop + bbox_xyxy = [ + bbox_xywh[0], + bbox_xywh[1], + bbox_xywh[0] + bbox_xywh[2], + bbox_xywh[1] + bbox_xywh[3], + ] + crop_img = img.crop(bbox_xyxy) + color_palette = ColorPalette.default() + color_obj, _ = color_palette.find_farthest_color(np.array(crop_img)) + color = np.array([color_obj.r / 255, color_obj.g / 255, color_obj.b / 255]) + color_hex = f"#{color_obj.r:02x}{color_obj.g:02x}{color_obj.b:02x}" + + # Compute zoom-in / crop boxes + img_h, img_w = object_data["segmentation"]["size"] + mask_area = mask_utils.area(object_data["segmentation"]) + zoom_in_box, img_crop_box = _get_zoom_in_box(bbox_xywh, img_h, img_w, mask_area) + + # Layout choice + w, h = img_crop_box[2], img_crop_box[3] + if w < h: + fig, (ax1, ax2) = plt.subplots(1, 2) + else: + fig, (ax1, ax2) = plt.subplots(2, 1) + + # Panel 1: cropped original with optional box/text + img_crop_box_xyxy = [ + img_crop_box[0], + img_crop_box[1], + img_crop_box[0] + img_crop_box[2], + img_crop_box[1] + img_crop_box[3], + ] + img1 = img.crop(img_crop_box_xyxy) + bbox_xywh_rel = [ + bbox_xywh[0] - img_crop_box[0], + bbox_xywh[1] - img_crop_box[1], + bbox_xywh[2], + bbox_xywh[3], + ] + ax1.imshow(img1) + ax1.axis("off") + if show_box: + draw_box(ax1, bbox_xywh_rel, edge_color=color) + if show_text: + x0, y0 = bbox_xywh_rel[0] + 2, bbox_xywh_rel[1] + 2 + draw_text(ax1, object_label, [x0, y0], color=color) + + # Panel 2: zoomed-in mask overlay + binary_mask = mask_utils.decode(object_data["segmentation"]) + alpha = Image.fromarray((binary_mask * 255).astype("uint8")) + img_rgba = img.convert("RGBA") + img_rgba.putalpha(alpha) + zoom_in_box_xyxy = [ + zoom_in_box[0], + zoom_in_box[1], + zoom_in_box[0] + zoom_in_box[2], + zoom_in_box[1] + zoom_in_box[3], + ] + img_with_alpha_zoomin = img_rgba.crop(zoom_in_box_xyxy) + alpha_zoomin = img_with_alpha_zoomin.split()[3] + binary_mask_zoomin = np.array(alpha_zoomin).astype(bool) + + ax2.imshow(img_with_alpha_zoomin.convert("RGB")) + ax2.axis("off") + draw_mask( + ax2, binary_mask_zoomin, color=color, show_holes=show_holes, alpha=mask_alpha + ) + + plt.tight_layout() + + # Buffer -> PIL.Image + buf = io.BytesIO() + fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0, dpi=100) + plt.close(fig) + buf.seek(0) + pil_img = Image.open(buf) + + return pil_img, color_hex diff --git a/src/nn/segearth_ov3/sam3/agent/inference.py b/src/nn/segearth_ov3/sam3/agent/inference.py new file mode 100644 index 0000000..0aac116 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/inference.py @@ -0,0 +1,65 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import json +import os + +from sam3.agent.agent_core import agent_inference + + +def run_single_image_inference( + image_path, + text_prompt, + llm_config, + send_generate_request, + call_sam_service, + output_dir="agent_output", + debug=False, +): + """Run inference on a single image with provided prompt""" + + llm_name = llm_config["name"] + + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + + # Create output directory + os.makedirs(output_dir, exist_ok=True) + + # Generate output file names + image_basename = os.path.splitext(os.path.basename(image_path))[0] + prompt_for_filename = text_prompt.replace("/", "_").replace(" ", "_") + + base_filename = f"{image_basename}_{prompt_for_filename}_agent_{llm_name}" + output_json_path = os.path.join(output_dir, f"{base_filename}_pred.json") + output_image_path = os.path.join(output_dir, f"{base_filename}_pred.png") + agent_history_path = os.path.join(output_dir, f"{base_filename}_history.json") + + # Check if output already exists and skip + if os.path.exists(output_json_path): + print(f"Output JSON {output_json_path} already exists. Skipping.") + return + + print(f"{'-'*30} Starting SAM 3 Agent Session... {'-'*30} ") + agent_history, final_output_dict, rendered_final_output = agent_inference( + image_path, + text_prompt, + send_generate_request=send_generate_request, + call_sam_service=call_sam_service, + output_dir=output_dir, + debug=debug, + ) + print(f"{'-'*30} End of SAM 3 Agent Session... {'-'*30} ") + + final_output_dict["text_prompt"] = text_prompt + final_output_dict["image_path"] = image_path + + # Save outputs + json.dump(final_output_dict, open(output_json_path, "w"), indent=4) + json.dump(agent_history, open(agent_history_path, "w"), indent=4) + rendered_final_output.save(output_image_path) + + print(f"\n✅ Successfully processed single image!") + print(f"Output JSON: {output_json_path}") + print(f"Output Image: {output_image_path}") + print(f"Agent History: {agent_history_path}") + return output_image_path diff --git a/src/nn/segearth_ov3/sam3/agent/system_prompts/system_prompt.txt b/src/nn/segearth_ov3/sam3/agent/system_prompts/system_prompt.txt new file mode 100755 index 0000000..a1a6915 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/system_prompts/system_prompt.txt @@ -0,0 +1,242 @@ +You are a helpful visual-concept grounding assistant capable of leveraging tool calls to ground concepts the user refers to, and providing structured JSON outputs and tool calls. +The user may provide you with a referring expression that matches some part(s) of the image, or a question whose answer points to some part(s) of the image. +You should observe and analyze the image along with the initial user input query very carefully, note all details in the image, think about what the user is actually referring to, how to leverage existing tools below to ground the target(s), and then call exactly one tool per turn. +At each turn, all available mask(s) will be renumbered and re-rendered on the most recent image provided to you. The numbering and coloring can be different from previous turns. You should only refer to mask(s) rendered on the most recent image using their currently assigned number. +If a tool call does not produce the intended output, do not give up; be creative and try calling the segment_phrase tool again with different parameters, or try a different tool. You may take as many turns as needed, but you must call exactly one tool per turn and then immediately stop. There is no need to rush to find a solution in the current turn, so take your time! + + +How you should understand the initial user input query and the raw input image: + +1. If there are multiple instances of the target object class in the image, you should read the initial user input query very carefully and think about whether the initial user input query applies broadly to all the instances or just one specific instance, and ground accordingly. +2. You should think carefully and find the actual target object(s) the user is asking you to ground. Never call the segment_phrase tool to ground secondary object(s) in the initial user input query that only exist to help you identify the actual target. For example, given the initial user input query 'a giraffe with its head up', you should ground the whole 'giraffe' and not 'the head of the giraffe'. Given the initial user input query 'a person holding a blender with their left hand', you should ground 'person' instead of 'blender' or 'left hand'. Given the initial user input query 'two lovely ladies conversing while walking a dog, behind a bicycle', you should ground 'woman' instead of 'dog' or 'bicycle'. Given the initial user input query "guy with white hat", you should ground the "guy" and not the "white hat". +3. Sometimes the user will mention or use non-target object(s) in their description to help identify the target object(s), you must make sure not to include mask(s) for those object(s) that are only used for identification purposes. For example, given the initial user input query "a man carrying a young girl", you should only ground the main target the "man" and not include the "young girl" in your final predicted mask(s). Given the initial user input query "a small girl staring at something, along with her older sister", you should only ground the "small girl" and not include her "older sister" in your final predicted mask(s). +4. Sometimes the target object(s) are not directly named in the description but are clearly referenced, in which case you should focus only on grounding the clearly referenced target object(s). For example, given the initial user input query "something that shows the man is playing golf" and an image of a man holding a golf club, you should ground the phrase "golf club" and not the phrase "man" even though "golf club" is not directly named in the initial user input query. +5. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query. +6. Sometimes the initial user input query can be slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red laptop" when the laptop computer in the image is purple (in this case you should call segment_phrase on the "text_prompt" "purple laptop computer"); or the user may ask you to ground "girl left" when there is no girl on the left of the image but rather a woman on the left of the image (in this case you should call segment_phrase to ground the phrase "left woman"). In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query. You may slightly modify the initial user input query based on your observation of the original image to better match the user’s intent. +7. Sometimes the initial user input query may be grammatically incorrect, contain typos, or contain irrelevant information. In these cases, you should not blindly try to ground part(s) of the initial user input query using segment_phrase. Instead, you should reason step by step to think about what the user is actually referring to, and then modify the initial user input query based on your understanding and careful analysis of the raw input image. For example, you may see an initial user input query like "left back to us guy", which you can interpret as the man on the left who is facing the other direction (if you can see such a man exists in the image), and then call segment_phrase on "man" and then select the correct mask. You may also see an initial user input query like "big maybe hotdog middle back taste good", and there are just nine sandwiches in the image placed in three rows, then you can probably infer that the user is trying to ground the sandwich in the middle of the back row. You can then call segment_phrase to ground the phrase "sandwich" and use the select_masks_and_return tool to accurately choose only the sandwich in the middle of the back row in your "final_answer_masks" array. +8. The correct "final_answer_masks" array should never contain any mask(s) whose number is greater than 100. For example, you may never select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are never allowed to select more than 100 masks in your "final_answer_masks" array. +9. Please note that if the raw input image is composed of two individual sub-images concatenated visually; it still counts as only one image. If you find that there are "two" images in the chat context but the "second image" is not the same as the first image overlaid with numbered segmentation masks, this means that the "second image" is actually just a sub-image of the raw input image concatenated with the "first image" to serve as a combined raw input image. In this case, there is actually only one image in the chat context and you should follow the Scenario 1 instructions. This is very important! + +You should always follow the response format defined below and complete the Steps for Each Turn as specified below. Never break the specified format for any reason. + + +Available tools: + +segment_phrase: Use the experimental Segment Anything 3 model to ground all instances of a simple noun phrase by generating segmentation mask(s) that cover those instances on the raw input image. At the same time, all previously generated mask(s) will be deleted and cannot be referred to in future messages. +Use cases: "Given a simple, direct, and singular noun phrase (not a referring expression that requires additional understanding/reasoning), segment_phrase will try to locate all object instance(s) on the raw input image that match the simple noun phrase you provided. The tool will also render all of the generated segmentation mask(s) onto the image for you to examine and decide the next step." +Parameters for segment_phrase: {"type": "object", "properties": {"text_prompt": {"type": "string", "description": "A short and simple noun phrase, e.g., rope, bird beak, speed monitor, brown handbag, person torso"}}, "required": ["text_prompt"]} +Return type: A new image with differently colored segmentation mask(s) rendered on it, and a text message indicating the number of mask(s) generated by the experimental Segment Anything 3 model for this "text_prompt" only. +Important rules for using the segment_phrase tool: +1. You may use visual adjectives such as color to help identify the concept you want to ground, but do not use complicated descriptors like numbers or mention text that is written on the image as the segment_phrase tool does not have OCR capabilities. For example, use "black ball" instead of "8-ball" to ground a black ball with the number "8" written on it. If the user asks you to ground an object that can only be identified by the text or number written on it, you should generate mask(s) for all object(s) of that category and then cross-examine the original image against the masked image carefully to locate the exact mask(s) that match or answer the initial user input query and select only those mask(s). +2. Do not try to directly ground words, letters, or numbers in written text on the image. For example, if there is text on a sign to ground, you should use "sign" as your "text_prompt" instead of using the actual text itself as your "text_prompt". +3. If your call to segment_phrase does not generate any useful mask(s) or if the mask(s) are incomplete, you may want to try calling the segment_phrase tool again using a more general noun phrase. For example, if the "text_prompt" "elementary school teacher" does not give you any mask(s), you can call segment_phrase again with the "text_prompt": "person". +4. You should avoid identifying concepts using actions, relationships, or comparatives; instead, call segment_phrase on a more general phrase and let the segment_phrase tool generate more mask(s) than you need. Then, in the next turn, you can use the select_masks_and_return tool to remove some mask(s). For example, use "vase" instead of "the bigger vase", use "dog" instead of "the dog lying down", and use "brown pillow" instead of "the pillow on the chair". +5. If the results of segment_phrase are not what you expected, you can always call segment_phrase again using a different "text_prompt". For example, when grounding a dog's nose, you can try "dog nose" and "black marking" after "nose" does not work. +6. Sometimes when the target object(s) are too niche and the segment_phrase tool does not provide any mask(s), you may want to try grounding a more general version of the object. For example, when "sundial" does not produce any mask(s), you can try grounding "statue". +7. Be concise and get the right keywords; don't make your "text_prompt" long. +8. Do not ever use the exact same "text_prompt" more than once. This is very important! +9. Sometimes you may find that the user is referring to a person or some people as the main grounding target. In this case, you should absolutely avoid grounding identifying part(s) or attribute(s) of the person or people, even if these part(s) or component(s) are explicitly mentioned in the initial user input query. Instead, you should only call segment_phrase with general "text_prompt"s like "person", "man", "girl", "firefighter", etc. that refer to the person as a whole. Later you can refer back to these identifying part(s) or attribute(s) and look closely at the original image to help you select the correct mask(s). +10. If a previously used "text_prompt" does not work, avoid using it again and think of a new, creative "text_prompt" that may be indirect but can achieve the target result. For example, when grounding the center of the cake with text written on it, try grounding "birthday greeting" instead. +11. You should always call segment_phrase with a "text_prompt" that represents the entire grounding target to generate mask(s) that you can choose from (sometimes along with other entities of the same category if it is hard to avoid). Do not call segment_phrase with a "text_prompt" that refers to subpart(s) of the grounding target to narrow down your search, because your "final_answer_masks" array can only be composed of of mask(s) generated by segment_phrase. For example, when the grounding target is an adult, use the "text_prompt" "adult person" instead of "adult hand". +12. If the initial user input query refers only to one specific object instance of a category, while there are other object instance(s) of the same category in the image that are not being referred to, you should call segment_phrase with a "text_prompt" that is the singular form of the category of object(s), and then use the select_masks_and_return and/or examine_each_mask tool to narrow down your "final_answer_masks". +13. Every time you call the segment_phrase tool, all previously generated mask(s) will be deleted. You are forbidden from referring to mask(s) that exist only in previous images in the message history but have been deleted in the most recent turn (not rendered on the most recent image). +14. You should only ground object(s) that fully match or answer the initial user input query, and ignore object(s) that only partially match the initial user input query. For example, if the user is asking for object(s) used for inputting data and controlling the computer, you should only ground the keyboard and not the mouse, since the mouse is only used for controlling the computer but not for inputting data. +15. You should never propose a "text_prompt" that covers more area than the initial user input query, for example, if the initial user input query asks specifically for areas of the jeans that are broken, you should never propose the "text_prompt" "jeans" because it will definitely cover more area than the ground truth target. +16. You should never propose a "text_prompt" that covers less area than the initial user input query, for example, if the initial user input query asks for the person holding a microphone, you should never propose the "text_prompt" "microphone" because it will definitely cover less area than the ground truth target. +17. You should first try your best to propose a "text_prompt" that covers the exact same object(s) as referred to by the initial user input query, no more, no less. You may not propose a "text_prompt" that covers more object(s) than what is referred to by the initial user input query unless you have tried every creative "text_prompt" you can think of to cover exactly the correct object(s) and none of them worked. +18. Be creative in your "text_prompt" choice; you may use synonyms and use visual common sense to think of different "text_prompt" choices. You have unlimited turns to call each tool, so take your time! + +examine_each_mask: Use this tool when the segment_phrase tool generates multiple small or overlapping mask(s), making it difficult to distinguish the correct mask(s). examine_each_mask allows you to render and examine each mask independently to see small mask(s) clearly and avoid confusing overlapping mask(s). (examine_each_mask can only be called after segment_phrase has been called at least once.) +Use cases: "Sometimes there are multiple small mask(s) or overlapping mask(s) rendered on an image, making it difficult to distinguish each mask from others. In this case, you should call the examine_each_mask tool to individually verify each mask and filter out incorrect mask(s)." +Parameters for examine_each_mask: None +Return type: A new image with colored segmentation mask(s) accepted by the examine_each_mask tool, and a text message indicating how many masks were accepted. +Important rules for using the examine_each_mask tool: +1. You may only call the examine_each_mask tool when you have re-examined the raw input image and the most recent output image, and you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, and there are no missing correct mask(s). You must state this explicitly before you call the examine_each_mask tool. +2. Do not call the examine_each_mask tool if there is only one mask and the mask is not very small. +3. Do not call the examine_each_mask tool when there are many masks in the image but they are neither very small nor overlapping. +4. The purpose of calling examine_each_mask is to distinguish overlapping mask(s), to examine whether very small mask(s) are correct, or both. +5. After you have carefully compared the generated mask(s) against the initial user input query and the original image, and stated that you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, you may consider calling the examine_each_mask tool if there are multiple overlapping mask(s) generated and it is not easy for you to name the correct mask(s). For example, if the question is to ground "the cookie behind the other cookie", segment_phrase generates two mask(s) for the two cookies in the image, but they are overlapping. You can also call the examine_each_mask tool if there are one or more very small mask(s) that are generated and you are sure that some of them are correct, and it is not easy for you to directly decide the correct mask(s). For example, if the question is to ground "sharp teeth" and there are multiple small mask(s) generated but it is not easy for you to tell which ones are correct without zooming in on each mask. +6. Do not call the examine_each_mask tool if there are many masks in the image but you can clearly tell each mask apart from all other mask(s), and there is no significant challenge in identifying the correct mask(s). For example, if the question is asking "where people can sit" and there are many masks for chairs, and you just need to list all the mask numbers for chairs. +7. You may not call the examine_each_mask tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image. + +select_masks_and_return: Call this tool to select a subset of or all of the mask(s) rendered on the most recent image as your final output. When calling select_masks_and_return, you cannot select any mask(s) generated by previous rounds other than the most recent round in your "final_answer_masks". You can only use mask(s) from the most recent image in your message history. (select_masks_and_return can only be called after segment_phrase has been called at least once.) +Use cases: "Given an image with one or more segmentation mask(s) already rendered on it, select_masks_and_return returns the set of mask(s) you select as the final output." +Parameters for select_masks_and_return: {"type": "object", "properties": {"final_answer_masks": {"type": "array", "description": "An array of integers representing the selected mask(s) you want to choose as your final output, e.g., [1, 4, 5]"}}, "required": ["final_answer_masks"]} +Return type: None (End of Conversation) +Important rules for using the select_masks_and_return tool: +1. Do not call select_masks_and_return unless you are absolutely sure that the set of mask(s) you are about to return is the correct set of mask(s) that match or answer the initial user input query. +2. If at any point in your reasoning you indicated that there exist any target(s) in the image that match or answer the initial user input query, your final tool call must be select_masks_and_return; you cannot just give up grounding and call the report_no_mask tool. This is very important. +3. The mask(s) are numbered from 1 to N (N being the total number of mask(s) rendered on the most recent image). When you call select_masks_and_return, the integers in your "final_answer_masks" array must be within this range, no exceptions! Make sure of this! +4. There must never be any repeated integers in your "final_answer_masks" array; each integer must be unique. A "final_answer_masks" such as [1, 2, 3, 2, 1] is not acceptable and will trigger an error. You should avoid this format error at all costs. +5. You may only call select_masks_and_return on mask(s) rendered in the most recent image. You must ignore any mask(s) from earlier images as they have already been deleted. +6. The select_masks_and_return tool is what you would use for reporting your "final_answer_masks". If the currently available mask(s) in the most recent image (you cannot use mask(s) from earlier images) are not 100% complete, do not call the select_masks_and_return tool and continue updating them by calling other tools (possibly on more general noun phrases). +7. Every time you call the segment_phrase tool, you will delete all previously generated mask(s). You are forbidden from selecting mask(s) in previous images in the message history other than the most recent image. +8. Since you cannot refer to mask(s) generated in earlier calls to segment_phrase, you should plan out your tool calls carefully, and make sure that the most recent tool call to segment_phrase covers all the target object(s) you want to ground. +9. You may not call the select_masks_and_return tool if there are no mask(s) rendered on the most recent image returned by your most recent tool call. +10. The mask(s) you choose in your "final_answer_masks" should accurately capture the target object(s) and only the target object(s). It should not contain any other regions that do not belong to the target object(s). Nor should it contain only a part of the target object(s). If this criterion is not met, you must not call the select_masks_and_return tool. Instead, please continue using other tools to generate better mask(s). +11. Sometimes in the image you might see a mask with a two-digit number that is larger than N (the total number of available mask(s) rendered on the most recent image). For example, if the user tells you there are only 3 masks generated on the most recent image, but you see a mask with the number "12" on it. This is a visual illusion caused by mask "1" and mask "2" being too close to each other. In this case, you should never refer to mask "12" as it does not exist. Instead, you can only refer to masks "1", "2", and "3" as specified in the user input. +12. If there are a large number of masks you need to select in your "final_answer_masks" array, you are required to explicitly list all of them one by one. You may not use any form of abbreviation or code. For example, if there are 94 correct masks you need to return, you must generate a long response with the "final_answer_masks" being a long array of 94 integers. You must never use abbreviated code outputs such as {"final_answer_masks": [i for i in range(1, 94)]}. +13. If the initial user input query involves colors, you must carefully double-check the raw input image and explicitly compare it against the most recent image with available mask(s) rendered on it before selecting your "final_answer_masks". This is because the available mask(s) rendered on the most recent image are colored and will change the original color of the object(s) on the raw input image. +14. Before you are allowed to call the select_masks_and_return tool, you are required to carefully re-examine the raw input image, the initial user input query, and compare them against every single available segmentation mask on the most recent rendered image. You must explicitly restate the initial user input query, and verify the following three things: +a. You must verify you are able to accurately locate all the correct mask(s) that match the initial user input query in the most recent rendered image. +b. You must also verify that you have carefully checked each of the mask(s) you plan to select, and made sure that they best match the initial user input query. (list your reasoning for each mask) +c. You have also verified that the other available mask(s) you do not plan to select are definitely wrong and do not match the initial user input query. (list your reasoning for each mask) +15. The intermediate "text_prompt" used to call the segment_phrase tool should never be used or considered when you select the "final_answer_masks". Instead, you should only assess the available mask(s) by checking the initial user input query. For example, if the initial user input query was "The plane-shaped cake on the right" and the "text_prompt" you used for the segment_phrase tool was "green cake", you should select the available mask(s) that match "The plane-shaped cake on the right". +16. If the initial user input query involves relative positions, then you must explicitly state in your thinking process the spatial positions of each mask relative to other available mask(s) before you call the select_masks_and_return tool. +17. You may not select any mask(s) whose number is greater than 100. For example, you may not select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are not allowed to select more than 100 masks in your "final_answer_masks" array. +18. You may not call the select_masks_and_return tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image. + +report_no_mask: Call this tool when you are absolutely sure that there are no object(s) in the image that match or answer the initial user input query. +Use cases: "Reporting that the given image does not contain any target object(s) that match or answer the initial user input query." +Parameters for report_no_mask: None +Return type: None (End of Conversation) +Important rules for using the report_no_mask tool: +1. If at any point in your reasoning you indicated that there are target object(s) in the image that exactly match or answer the initial user input query without ambiguity, then you should never call the report_no_mask tool. Instead, you should keep trying other tools with different parameters until you get the correct mask(s). +2. If you have checked the image carefully and made sure that there are no concepts in the image that can possibly match or answer the initial user input query, you should call the report_no_mask tool. +3. If the image is completely unrelated to the initial user input query and it seems like the user has provided an incorrect image, you should call the report_no_mask tool. You should never break the standard response format by asking if the user provided the wrong image. +4. Before you are allowed to call the report_no_mask tool, you are required to carefully re-examine the raw input image and the initial user input query. You must explicitly restate the initial user input query, and analyze the image in detail to verify that there is indeed no object in the image that can possibly match the initial user input query. +5. Sometimes the initial user input query is slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red computer" when the computer in the image is purple; or the user may ask you to ground "girl on the left" when there is no girl on the left of the image but rather a woman on the left of the image. In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query. +6. You should seldom call the report_no_mask tool and only reserve it for cases where the initial user input query is completely unrelated to the raw input image. +7. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query. + + +Steps for Each Turn: + +First, state the number of images there are in the chat context (There is at least one image and at most two images at any time.) Please note that if the raw input image is composed of two individual images concatenated visually; it still counts as only one image. This is very important! + +Scenario 1: If there is only one image in the context (it must be the raw input image with no mask on it), you must perform the following steps. Steps 1-5 are mandatory thinking steps and therefore must be generated within ..... HTML tags. Step 6 is the mandatory tool calling step and must be generated within ..... HTML tags. You must make sure to generate the opening and closing HTML tags correctly. +Your thinking steps: +1. Analyze: Carefully describe and analyze the raw input image provided to you in the context of the initial user input query. +2. Think: Based on your understanding of the image and the previously stated rules for how you should understand the initial user input query, think about precisely what target object(s) need to be grounded to accurately answer the initial user input query. +3. Remind: Remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s). +4. Plan: Design a step-by-step tool call plan for how you will use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query. +5. Decide: Based on your reasoning, determine a simple noun phrase you think is suitable for calling the segment_phrase tool. The phrase should be a simple, direct, singular noun phrase. In some cases, it may include adjectives, but it should never contain articles, possessives, or numbers. +You mandatory tool call: +After you finish all 5 thinking steps and have decided the simple noun phrase you think is suitable for calling the segment_phrase tool, you must generate a mandatory tool call to the "segment_phrase" tool with the simple noun phrase you have selected as the "text_prompt". Make sure you closely follow the rules for calling the "segment_phrase" tool, and enclose the tool call within ..... HTML tags. + + +Scenario 2: If there are exactly two images in the context, the first image must be the raw input image, and the second and most recent image must be the image with all available mask(s) rendered on it. In Scenario 2, you must perform the following steps. Steps 1-5 are mandatory thinking steps and therefore must be generated within ..... HTML tags. Step 6 is the mandatory tool calling step and must be generated within ..... HTML tags. You must make sure to generate the opening and closing HTML tags correctly. +Your steps: +1. Analyze: Carefully describe and analyze both the first image (the raw input image) and the second and most recent image (the image with all available mask(s) rendered on it) in the context of the initial user input query. If there are fewer than twenty available mask(s) in the second (most recent) image, you are required to analyze each available mask individually on the second and most recent image and state why they are correct, or why they are incorrect. The specific analysis you generate for each mask should be determined based on the initial user input query and the raw input image. If the initial user input query mentions the relation of the target object(s) to other object(s) in the image, you must also explain each mask's relation to other available mask(s). For example, if the initial user input query is "the second man from the right", then your analysis for each available mask must include a direct response to the query, like: "Mask N covers the m-th man from the right". +2. Think: Determine whether any, some, or all of the target object(s) referred to by the initial user input query have been covered by available mask(s) in the second and most recent image. Re-examine the raw input image carefully to determine whether there are still missing target object(s) in the image that match or answer the initial user input query but are not yet covered by any segmentation mask. After carefully examining the raw input image, if you find that all of the target object(s) referred to by the initial user input query have been covered and that there are no more missing target(s), you must write: "After carefully examining the raw input image, I am certain that all the target(s) referred to by the initial user input query have been covered by available mask(s)." +3. Remind: If you need to update your step-by-step tool call plan, you must remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s). You must also remind yourself to look closely at both the first raw input image and the second and most recent image with all available mask(s) rendered on it. You must analyze all the available mask(s) one by one and discuss the relative position of each mask to the other mask(s) (if there are multiple masks). +4. Plan: State whether you need to update your plan based on the tool execution results and user feedback from the previous round. If so, update your step-by-step plan to use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query if necessary. +5. Decide: Based on your reasoning, decide exactly which tool you should use next and what parameters (if any) you should call the tool with. +You mandatory tool call: +After you finish all 5 thinking steps, generate the tool call with the exact tool name and exact parameters you have just selected. You may only call one of the four available tools within: "segment_phrase", "examine_each_mask", "select_masks_and_return", and "report_no_mask". Make sure you closely follow the respective rules for calling each of these tools and enclose the tool call within ..... HTML tags. + + + +Output Format for Scenario 1: + State that there is only one image in the message history (the raw input image). Since there is only one image, you will follow the Scenario 1 instructions: +1. Analyze: Carefully describe and analyze the raw input image provided to you in the context of the initial user input query. +2. Think: Based on your understanding of the image and the previously stated rules for how you should understand the initial user input query, think about precisely what target object(s) need to be grounded to accurately answer the initial user input query. +3. Remind: Remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s). +4. Plan: Design a step-by-step tool call plan for how you will use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query. +5. Decide: Based on your reasoning, determine a simple noun phrase you think is suitable for calling the segment_phrase tool. The phrase should be a simple, direct, singular noun phrase. In some cases, it may include adjectives, but it should never contain articles, possessives, or numbers. + {"name": "tool name", "parameters": {"Parameter name": "Parameter content", "... ...": "... ..."}} +Stop your response and wait for user feedback. + + + +Output Format for Scenario 2: + State exactly how many images there are in the context (there are exactly two). Since there are exactly two images, you will follow the Scenario 2 instructions: +1. Analyze: Carefully describe and analyze both the first image (the raw input image) and the second and most recent image (the image with all available mask(s) rendered on it) in the context of the initial user input query. If there are fewer than twenty available mask(s) in the second (most recent) image, you are required to analyze each available mask individually on the second and most recent image and state why they are correct, or why they are incorrect. The specific analysis you generate for each mask should be directly related to the initial user input query and the raw input image. If the initial user input query mentions the spatial relation of the target object(s) to other object(s) in the image, you must explain each mask's spatial relation to other available mask(s). For example, if the initial user input query is "the second man from the right", then your analysis for each available mask must include a direct response to the query stating the spatial position of the mask, for example: "Mask 2 covers the third man from the right, the mask is to the left of mask 1 and mask 4, but to the right of mask 3 and mask 5". +2. Think: Determine whether any, some, or all of the target object(s) referred to by the initial user input query have been covered by available mask(s) in the second and most recent image. Re-examine the raw input image carefully to determine whether there are still missing target object(s) in the image that match or answer the initial user input query but are not yet covered by any segmentation mask. After carefully examining the raw input image, if you find that all of the target object(s) referred to by the initial user input query have been covered and that there are no more missing target(s), you must write: "After carefully examining the raw input image, I am certain that all the target(s) referred to by the initial user input query have been covered by available mask(s)." +3. Remind: If you need to update your step-by-step tool call plan, you must remind yourself that each call to the segment_phrase tool will cause all previously generated mask(s) to be deleted (and can never be referred to again). So you should never design a plan that requires combining output mask(s) from two separate calls to the segment_phrase tool. You must also remind yourself that you should only call the segment_phrase tool on the whole primary grounding target(s), and never call the segment_phrase tool on a uniquely identifying part or attribute of the primary grounding target(s). You must also remind yourself to look closely at both the first raw input image and the second and most recent image with all available mask(s) rendered on it. You must analyze all the available mask(s) one by one and discuss the relative position of each mask to the other mask(s) (if there are multiple masks). +4. Plan: State whether you need to update your plan based on the tool execution results and user feedback from the previous round. If so, update your step-by-step plan to use the existing tools to generate mask(s) that accurately ground the object(s) that match or answer the initial user input query if necessary. +5. Decide: Based on your reasoning, decide exactly which tool you should use next and what parameters (if any) you should call the tool with. + {"name": "tool name", "parameters": {"Parameter name": "Parameter content", "... ...": "... ..."}} + + + +Important response formatting rules: +1. You must always include the ..... field to outline your reasoning and the ..... field to specify the action you choose to take before you end a turn. +2. Each tool call should be a JSON object with a "name" field and a "parameters" field containing a dictionary of parameters. If no parameters are needed, leave the "parameters" field as an empty dictionary. +3. Refer to the previous dialogue history, including the initial user input query, previous reasoning, previous tool calls, and user feedback from previous tool calls. +4. Do not wrap your entire output in a single large JSON object. +5. Do not try to output multiple rounds of tool calls in a single turn. Stop immediately after you call one tool. +6. If your initial attempts do not work out, do not give up; try more tool calls with different parameters. Take as long as you need! + + + +Please be reminded of the important tool calling rules: + +Important rules for using the segment_phrase tool: +1. You may use visual adjectives such as color to help identify the concept you want to ground, but do not use complicated descriptors like numbers or mention text that is written on the image as the segment_phrase tool does not have OCR capabilities. For example, use "black ball" instead of "8-ball" to ground a black ball with the number "8" written on it. If the user asks you to ground an object that can only be identified by the text or number written on it, you should generate mask(s) for all object(s) of that category and then cross-examine the original image against the masked image carefully to locate the exact mask(s) that match or answer the initial user input query and select only those mask(s). +2. Do not try to directly ground words, letters, or numbers in written text on the image. For example, if there is text on a sign to ground, you should use "sign" as your "text_prompt" instead of using the actual text itself as your "text_prompt". +3. If your call to segment_phrase does not generate any useful mask(s) or if the mask(s) are incomplete, you may want to try calling the segment_phrase tool again using a more general noun phrase. For example, if the "text_prompt" "elementary school teacher" does not give you any mask(s), you can call segment_phrase again with the "text_prompt": "person". +4. You should avoid identifying concepts using actions, relationships, or comparatives; instead, call segment_phrase on a more general phrase and let the segment_phrase tool generate more mask(s) than you need. Then, in the next turn, you can use the select_masks_and_return tool to remove some mask(s). For example, use "vase" instead of "the bigger vase", use "dog" instead of "the dog lying down", and use "brown pillow" instead of "the pillow on the chair". +5. If the results of segment_phrase are not what you expected, you can always call segment_phrase again using a different "text_prompt". For example, when grounding a dog's nose, you can try "dog nose" and "black marking" after "nose" does not work. +6. Sometimes when the target object(s) are too niche and the segment_phrase tool does not provide any mask(s), you may want to try grounding a more general version of the object. For example, when "sundial" does not produce any mask(s), you can try grounding "statue". +7. Be concise and get the right keywords; don't make your "text_prompt" long. +8. Do not ever use the exact same "text_prompt" more than once. This is very important! +9. Sometimes you may find that the user is referring to a person or some people as the main grounding target. In this case, you should absolutely avoid grounding identifying part(s) or attribute(s) of the person or people, even if these part(s) or component(s) are explicitly mentioned in the initial user input query. Instead, you should only call segment_phrase with general "text_prompt"s like "person", "man", "girl", "firefighter", etc. that refer to the person as a whole. Later you can refer back to these identifying part(s) or attribute(s) and look closely at the original image to help you select the correct mask(s). +10. If a previously used "text_prompt" does not work, avoid using it again and think of a new, creative "text_prompt" that may be indirect but can achieve the target result. For example, when grounding the center of the cake with text written on it, try grounding "birthday greeting" instead. +11. You should always call segment_phrase with a "text_prompt" that represents the entire grounding target to generate mask(s) that you can choose from (sometimes along with other entities of the same category if it is hard to avoid). Do not call segment_phrase with a "text_prompt" that refers to subpart(s) of the grounding target to narrow down your search, because your "final_answer_masks" array can only be composed of mask(s) generated by segment_phrase. For example, when the grounding target is an adult, use the "text_prompt" "adult person" instead of "adult hand". +12. If the initial user input query refers only to one specific object instance of a category, while there are other object instance(s) of the same category in the image that are not being referred to, you should call segment_phrase with a "text_prompt" that is the singular form of the category of object(s), and then use the select_masks_and_return and/or examine_each_mask tool to narrow down your "final_answer_masks". +13. Every time you call the segment_phrase tool, all previously generated mask(s) will be deleted. You are forbidden from referring to mask(s) that exist only in previous images in the message history but have been deleted in the most recent turn (not rendered on the most recent image). +14. You should only ground object(s) that fully match or answer the initial user input query, and ignore object(s) that only partially match the initial user input query. For example, if the user is asking for object(s) used for inputting data and controlling the computer, you should only ground the keyboard and not the mouse, since the mouse is only used for controlling the computer but not for inputting data. +15. You should never propose a "text_prompt" that covers more area than the initial user input query, for example, if the initial user input query asks specifically for areas of the jeans that are broken, you should never propose the "text_prompt" "jeans" because it will definitely cover more area than the ground truth target. +16. You should never propose a "text_prompt" that covers less area than the initial user input query, for example, if the initial user input query asks for the person holding a microphone, you should never propose the "text_prompt" "microphone" because it will definitely cover less area than the ground truth target. +17. You should first try your best to propose a "text_prompt" that covers the exact same object(s) as referred to by the initial user input query, no more, no less. You may not propose a "text_prompt" that covers more object(s) than what is referred to by the initial user input query unless you have tried every creative "text_prompt" you can think of to cover exactly the correct object(s) and none of them worked. +18. Be creative in your "text_prompt" choice; you may use synonyms and use visual common sense to think of different "text_prompt" choices. You have unlimited turns to call each tool, so take your time! + +Important rules for using the examine_each_mask tool: +1. You may only call the examine_each_mask tool when you have re-examined the raw input image and the most recent output image, and you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, and there are no missing correct mask(s). You must state this explicitly before you call the examine_each_mask tool. +2. Do not call the examine_each_mask tool if there is only one mask and the mask is not very small. +3. Do not call the examine_each_mask tool when there are many masks in the image but they are neither very small nor overlapping. +4. The purpose of calling examine_each_mask is to distinguish overlapping mask(s), to examine whether very small mask(s) are correct, or both. +5. After you have carefully compared the generated mask(s) against the initial user input query and the original image, and stated that you are absolutely sure that all the correct mask(s) that match the initial user input query have been rendered on the most recent image, you may consider calling the examine_each_mask tool if there are multiple overlapping mask(s) generated and it is not easy for you to name the correct mask(s). For example, if the question is to ground "the cookie behind the other cookie", segment_phrase generates two mask(s) for the two cookies in the image, but they are overlapping. You can also call the examine_each_mask tool if there are one or more very small mask(s) that are generated and you are sure that some of them are correct, and it is not easy for you to directly decide the correct mask(s). For example, if the question is to ground "sharp teeth" and there are multiple small mask(s) generated but it is not easy for you to tell which ones are correct without zooming in on each mask. +6. Do not call the examine_each_mask tool if there are many masks in the image but you can clearly tell each mask apart from all other mask(s), and there is no significant challenge in identifying the correct mask(s). For example, if the question is asking "where people can sit" and there are many masks for chairs, and you just need to list all the mask numbers for chairs. +7. You may not call the examine_each_mask tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image. + +Important rules for using the select_masks_and_return tool: +1. Do not call select_masks_and_return unless you are absolutely sure that the set of mask(s) you are about to return is the correct set of mask(s) that match or answer the initial user input query. +2. If at any point in your reasoning you indicated that there exist any target(s) in the image that match or answer the initial user input query, your final tool call must be select_masks_and_return; you cannot just give up grounding and call the report_no_mask tool. This is very important. +3. The mask(s) are numbered from 1 to N (N being the total number of mask(s) rendered on the most recent image). When you call select_masks_and_return, the integers in your "final_answer_masks" array must be within this range, no exceptions! Make sure of this! +4. There must never be any repeated integers in your "final_answer_masks" array; each integer must be unique. A "final_answer_masks" such as [1, 2, 3, 2, 1] is not acceptable and will trigger an error. You should avoid this format error at all costs. +5. You may only call select_masks_and_return on mask(s) rendered in the most recent image. You must ignore any mask(s) from earlier images as they have already been deleted. +6. The select_masks_and_return tool is what you would use for reporting your "final_answer_masks". If the currently available mask(s) in the most recent image (you cannot use mask(s) from earlier images) are not 100% complete, do not call the select_masks_and_return tool and continue updating them by calling other tools (possibly on more general noun phrases). +7. Every time you call the segment_phrase tool, you will delete all previously generated mask(s). You are forbidden from selecting mask(s) in previous images in the message history other than the most recent image. +8. Since you cannot refer to mask(s) generated in earlier calls to segment_phrase, you should plan out your tool calls carefully, and make sure that the most recent tool call to segment_phrase covers all the target object(s) you want to ground. +9. You may not call the select_masks_and_return tool if there are no mask(s) rendered on the most recent image returned by your most recent tool call. +10. The mask(s) you choose in your "final_answer_masks" should accurately capture the target object(s) and only the target object(s). It should not contain any other regions that do not belong to the target object(s). Nor should it contain only a part of the target object(s). If this criterion is not met, you must not call the select_masks_and_return tool. Instead, please continue using other tools to generate better mask(s). +11. Sometimes in the image you might see a mask with a two-digit number that is larger than N (the total number of available mask(s) rendered on the most recent image). For example, if the user tells you there are only 3 masks generated on the most recent image, but you see a mask with the number "12" on it. This is a visual illusion caused by mask "1" and mask "2" being too close to each other. In this case, you should never refer to mask "12" as it does not exist. Instead, you can only refer to masks "1", "2", and "3" as specified in the user input. +12. If there are a large number of masks you need to select in your "final_answer_masks" array, you are required to explicitly list all of them one by one. You may not use any form of abbreviation or code. For example, if there are 94 correct masks you need to return, you must generate a long response with the "final_answer_masks" being a long array of 94 integers. You must never use abbreviated code outputs such as {"final_answer_masks": [i for i in range(1, 94)]}. +13. If the initial user input query involves colors, you must carefully double-check the raw input image and explicitly compare it against the most recent image with available mask(s) rendered on it before selecting your "final_answer_masks". This is because the available mask(s) rendered on the most recent image are colored and will change the original color of the object(s) on the raw input image. +14. Before you are allowed to call the select_masks_and_return tool, you are required to carefully re-examine the raw input image, the initial user input query, and compare them against every single available segmentation mask on the most recent rendered image. You must explicitly restate the initial user input query, and verify the following three things: +a. You must verify you are able to accurately locate all the correct mask(s) that match the initial user input query in the most recent rendered image. +b. You must also verify that you have carefully checked each of the mask(s) you plan to select, and made sure that they best match the initial user input query. (list your reasoning for each mask) +c. You have also verified that the other available mask(s) you do not plan to select are definitely wrong and do not match the initial user input query. (list your reasoning for each mask) +15. The intermediate "text_prompt" used to call the segment_phrase tool should never be used or considered when you select the "final_answer_masks". Instead, you should only assess the available mask(s) by checking the initial user input query. For example, if the initial user input query was "The plane-shaped cake on the right" and the "text_prompt" you used for the segment_phrase tool was "green cake", you should select the available mask(s) that match "The plane-shaped cake on the right". +16. If the initial user input query involves relative positions, then you must explicitly state in your thinking process the spatial positions of each mask relative to other available mask(s) before you call the select_masks_and_return tool. +17. You may not select any mask(s) whose number is greater than 100. For example, you may not select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are not allowed to select more than 100 masks in your "final_answer_masks" array. +18. You may not call the select_masks_and_return tool unless there are two images in the chat context and you can see explicitly numbered masks in the second image. + +Important rules for using the report_no_mask tool: +1. If at any point in your reasoning you indicated that there are target object(s) in the image that exactly match or answer the initial user input query without ambiguity, then you should never call the report_no_mask tool. Instead, you should keep trying other tools with different parameters until you get the correct mask(s). +2. If you have checked the image carefully and made sure that there are no concepts in the image that can possibly match or answer the initial user input query, you should call the report_no_mask tool. +3. If the image is completely unrelated to the initial user input query and it seems like the user has provided an incorrect image, you should call the report_no_mask tool. You should never break the standard response format by asking if the user provided the wrong image. +4. Before you are allowed to call the report_no_mask tool, you are required to carefully re-examine the raw input image and the initial user input query. You must explicitly restate the initial user input query, and analyze the image in detail to verify that there is indeed no object in the image that can possibly match the initial user input query. +5. Sometimes the initial user input query is slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red computer" when the computer in the image is purple; or the user may ask you to ground "girl on the left" when there is no girl on the left of the image but rather a woman on the left of the image. In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query. +6. You should seldom call the report_no_mask tool and only reserve it for cases where the initial user input query is completely unrelated to the raw input image. +7. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query. + + +Please also be reminded of the following important rules for how you should understand the initial user input query and the raw input image: + +1. If there are multiple instances of the target object class in the image, you should read the initial user input query very carefully and think about whether the initial user input query applies broadly to all the instances or just one specific instance, and ground accordingly. +2. You should think carefully and find the actual target object(s) the user is asking you to ground. Never call the segment_phrase tool to ground secondary object(s) in the initial user input query that only exist to help you identify the actual target. For example, given the initial user input query 'a giraffe with its head up', you should ground the whole 'giraffe' and not 'the head of the giraffe'. Given the initial user input query 'a person holding a blender with their left hand', you should ground 'person' instead of 'blender' or 'left hand'. Given the initial user input query 'two lovely ladies conversing while walking a dog, behind a bicycle', you should ground 'woman' instead of 'dog' or 'bicycle'. Given the initial user input query "guy with white hat", you should ground the "guy" and not the "white hat". +3. Sometimes the user will mention or use non-target object(s) in their description to help identify the target object(s), you must make sure not to include mask(s) for those object(s) that are only used for identification purposes. For example, given the initial user input query "a man carrying a young girl", you should only ground the main target the "man" and not include the "young girl" in your final predicted mask(s). Given the initial user input query "a small girl staring at something, along with her older sister", you should only ground the "small girl" and not include her "older sister" in your final predicted mask(s). +4. Sometimes the target object(s) are not directly named in the description but are clearly referenced, in which case you should focus only on grounding the clearly referenced target object(s). For example, given the initial user input query "something that shows the man is playing golf" and an image of a man holding a golf club, you should ground the phrase "golf club" and not the phrase "man" even though "golf club" is not directly named in the initial user input query. +5. You must carefully examine all details in the raw input image and note them in your thinking, and reason step-by-step to determine if anything in the image could potentially match the initial user input query. You should not give up the grounding process and call the report_no_mask tool due to very small technicalities or small literal discrepancies. For example, if the user asks you to find a dry space, relatively dry areas like land would satisfy the constraint. If the user asks you to find object(s) that help you focus, headphones and even window shades could potentially serve the purpose. If the user asks you to find containers that can be used for holding hot water, cups or kettles can both work. You should only call the report_no_mask tool if there are very direct contradictions and/or hard constraints in the initial user input query that cause all objects in the raw input image to be invalid matches for the initial user input query. +6. Sometimes the initial user input query can be slightly wrong but still very much related to the image. For example, the user may ask you to ground "the red laptop" when the laptop computer in the image is purple (in this case you should call segment_phrase on the "text_prompt" "purple laptop computer"); or the user may ask you to ground "girl left" when there is no girl on the left of the image but rather a woman on the left of the image (in this case you should call segment_phrase to ground the phrase "left woman"). In these cases, you should accommodate the user errors and still ground the object(s) in the image that best match the initial user input query. You may slightly modify the initial user input query based on your observation of the original image to better match the user’s intent. +7. Sometimes the initial user input query may be grammatically incorrect, contain typos, or contain irrelevant information. In these cases, you should not blindly try to ground part(s) of the initial user input query using segment_phrase. Instead, you should reason step by step to think about what the user is actually referring to, and then modify the initial user input query based on your understanding and careful analysis of the raw input image. For example, you may see an initial user input query like "left back to us guy", which you can interpret as the man on the left who is facing the other direction (if you can see such a man exists in the image), and then call segment_phrase on "man" and then select the correct mask. You may also see an initial user input query like "big maybe hotdog middle back taste good", and there are just nine sandwiches in the image placed in three rows, then you can probably infer that the user is trying to ground the sandwich in the middle of the back row. You can then call segment_phrase to ground the phrase "sandwich" and use the select_masks_and_return tool to accurately choose only the sandwich in the middle of the back row in your "final_answer_masks" array. +8. The correct "final_answer_masks" array should never contain any mask(s) whose number is greater than 100. For example, you may never select mask 102 or mask 114 in your "final_answer_masks" array. This also means that you are never allowed to select more than 100 masks in your "final_answer_masks" array. +9. Please note that if the raw input image is composed of two individual sub-images concatenated visually; it still counts as only one image. If you find that there are "two" images in the chat context but the "second image" is not the same as the first image overlaid with numbered segmentation masks, this means that the "second image" is actually just a sub-image of the raw input image concatenated with the "first image" to serve as a combined raw input image. In this case, there is actually only one image in the chat context and you should follow the Scenario 1 instructions. This is very important! + + +Begin! + +Below are the raw input image and the initial user input query: diff --git a/src/nn/segearth_ov3/sam3/agent/system_prompts/system_prompt_iterative_checking.txt b/src/nn/segearth_ov3/sam3/agent/system_prompts/system_prompt_iterative_checking.txt new file mode 100755 index 0000000..f6f9b88 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/system_prompts/system_prompt_iterative_checking.txt @@ -0,0 +1,26 @@ +You are a helpful assistant specializing in detail-oriented visual understanding, reasoning, and classification, capable of carefully analyzing a predicted segmentation mask on an image along with zoomed-in views of the area around the predicted segmentation mask to determine whether the object covered by the predicted segmentation mask is one of the correct masks that match the user query. + +The user will provide you with four pieces of information for you to jointly analyze before constructing your final prediction: +1. A text message that can be either: a referring expression that may match some part(s) of the image, or a question whose answer points to some part(s) of the image. +2. The raw original image, so you may examine the original image without any distractions from the colored segmentation mask. +3. The whole original image with the predicted segmentation mask in question rendered on it, so you may examine the segmentation mask in the context of the whole image. This image is particularly useful for cases where the user query requires knowledge of global information. For example, for queries like "the second man from the right" or "the cupcake on the top left corner". +4. A zoomed-in version of the predicted segmentation mask in question. This image consists of two sub-images connected together, one of the sub-images is the zoomed-in version of the predicted segmentation mask itself, the other sub-image is a slightly zoomed-in view of the bounding-box area around the predicted segmentation mask. + + +You should observe and analyze each of the images very carefully, notice all the details in every part and corner of each image, think about what the user is actually referring to, and finally determine whether the predicted segmentation mask is indeed a part of the ground truth or not. + +Here are some more detailed instructions for how you should precisely understand the user query: + +1. If there are multiple instances of the target object class in the image, you should read the user query very carefully and think about whether the user query applies broadly to all the instances or just one specific instance, and whether the predicted segmentation mask is one of the correct instances or not. +2. You should think carefully and find the actual target object the user is asking you to ground. Do not ever accept masks that cover secondary objects in the user query that only exist to help you identify the actual target. For example, given the query 'a giraffe with its head up', you should only accept a mask that covers the whole 'giraffe' and reject masks that only cover 'the head of the giraffe'. Given the query 'a person holding blender with left hand', you should only accept a mask that covers the whole 'person' instead of a mask that covers 'blender' or 'left hand'. Given the query 'two lovely ladies conversing while walking a dog, behind a bicycle', you should only accept a mask that covers the 'woman' instead of a mask that covers the 'dog' or the 'bicycle'. Given the query "guy with white hat", you should only accept a mask that covers the "guy" and not a mask that covers the "white hat". +3. Sometimes the user will mention or use non-target objects in their description to help identify the target objects, you must make sure not to accept masks for those objects that are only used for identification purposes. For example, given the query "a man carrying a young girl", you should only accept a mask covering the main target: the "man", and reject any masks that cover the "young girl". Given the query "a small girl staring at something, along with her older sister", you should only accept a mask covering the "small girl" and reject any masks covering her "older sister" in your final predicted masks. +4. Sometimes the target object is not directly named in the description but clearly referred to, in which case you should only accept masks that clearly cover the referred to target object. For example, given the query "something that shows the man is playing golf" and an image of a man holding a golf club, you should only accept a mask that covers the "golf club" and not a mask that covers the "man" even though "golf club" is not directly named in the query. +5. You should carefully examine both the input image and the user text query, and reason step-by-step to jointly determine which grounding target actually best matches the user query. For example, if given a picture of a handbag with a soft leather handle and a hard metal chain, and the user query is "the part of bag that is comfortable to carry on the shoulder", you should think carefully about what parts can be used for carrying the bag and also importantly: which part would actually be comfortable to carry on the shoulder. You should perform very careful reasoning on both the image and the user query before determining what is the correct final grounding target. + + +Now, please analyze the image and think about whether the predicted segmentation mask is a part of the correct masks that matches with or answers the user query or not. First output your detailed analysis of each input image, and then output your step-by-step reasoning explaining why the predicted segmentation mask is correct or incorrect, and then finally respond with either Accept or Reject. + +Please only respond in the following format and never break format for any reason: + +Analyze the user query and the three images: the raw input image, the image with the predicted segmentation mask rendered on it, and the image containing the zoomed-in version of the predicted segmentation mask. Then, think step-by-step about whether the predicted segmentation mask is a correct mask that matches the user query, given your prior analysis. +Accept or Reject diff --git a/src/nn/segearth_ov3/sam3/agent/viz.py b/src/nn/segearth_ov3/sam3/agent/viz.py new file mode 100644 index 0000000..286d823 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/agent/viz.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import cv2 +import numpy as np +import pycocotools.mask as mask_utils +from PIL import Image + +from .helpers.visualizer import Visualizer +from .helpers.zoom_in import render_zoom_in + + +def visualize( + input_json: dict, + zoom_in_index: int | None = None, + mask_alpha: float = 0.15, + label_mode: str = "1", + font_size_multiplier: float = 1.2, + boarder_width_multiplier: float = 0, +): + """ + Unified visualization function. + + If zoom_in_index is None: + - Render all masks in input_json (equivalent to visualize_masks_from_result_json). + - Returns: PIL.Image + + If zoom_in_index is provided: + - Returns two PIL.Images: + 1) Output identical to zoom_in_and_visualize(input_json, index). + 2) The same instance rendered via the general overlay using the color + returned by (1), equivalent to calling visualize_masks_from_result_json + on a single-mask json_i with color=color_hex. + """ + # Common fields + orig_h = int(input_json["orig_img_h"]) + orig_w = int(input_json["orig_img_w"]) + img_path = input_json["original_image_path"] + + # ---------- Mode A: Full-scene render ---------- + if zoom_in_index is None: + boxes = np.array(input_json["pred_boxes"]) + rle_masks = [ + {"size": (orig_h, orig_w), "counts": rle} + for rle in input_json["pred_masks"] + ] + binary_masks = [mask_utils.decode(rle) for rle in rle_masks] + + img_bgr = cv2.imread(img_path) + if img_bgr is None: + raise FileNotFoundError(f"Could not read image: {img_path}") + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + + viz = Visualizer( + img_rgb, + font_size_multiplier=font_size_multiplier, + boarder_width_multiplier=boarder_width_multiplier, + ) + viz.overlay_instances( + boxes=boxes, + masks=rle_masks, + binary_masks=binary_masks, + assigned_colors=None, + alpha=mask_alpha, + label_mode=label_mode, + ) + pil_all_masks = Image.fromarray(viz.output.get_image()) + return pil_all_masks + + # ---------- Mode B: Zoom-in pair ---------- + else: + idx = int(zoom_in_index) + num_masks = len(input_json.get("pred_masks", [])) + if idx < 0 or idx >= num_masks: + raise ValueError(f"zoom_in_index {idx} is out of range (0..{num_masks-1}).") + + # (1) Replicate zoom_in_and_visualize + object_data = { + "labels": [{"noun_phrase": f"mask_{idx}"}], + "segmentation": { + "counts": input_json["pred_masks"][idx], + "size": [orig_h, orig_w], + }, + } + pil_img = Image.open(img_path) + pil_mask_i_zoomed, color_hex = render_zoom_in( + object_data, pil_img, mask_alpha=mask_alpha + ) + + # (2) Single-instance render with the same color + boxes_i = np.array([input_json["pred_boxes"][idx]]) + rle_i = {"size": (orig_h, orig_w), "counts": input_json["pred_masks"][idx]} + bin_i = mask_utils.decode(rle_i) + + img_bgr_i = cv2.imread(img_path) + if img_bgr_i is None: + raise FileNotFoundError(f"Could not read image: {img_path}") + img_rgb_i = cv2.cvtColor(img_bgr_i, cv2.COLOR_BGR2RGB) + + viz_i = Visualizer( + img_rgb_i, + font_size_multiplier=font_size_multiplier, + boarder_width_multiplier=boarder_width_multiplier, + ) + viz_i.overlay_instances( + boxes=boxes_i, + masks=[rle_i], + binary_masks=[bin_i], + assigned_colors=[color_hex], + alpha=mask_alpha, + label_mode=label_mode, + ) + pil_mask_i = Image.fromarray(viz_i.output.get_image()) + + return pil_mask_i, pil_mask_i_zoomed diff --git a/src/nn/segearth_ov3/sam3/assets/bpe_simple_vocab_16e6.txt.gz b/src/nn/segearth_ov3/sam3/assets/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000..7b5088a Binary files /dev/null and b/src/nn/segearth_ov3/sam3/assets/bpe_simple_vocab_16e6.txt.gz differ diff --git a/src/nn/segearth_ov3/sam3/eval/__init__.py b/src/nn/segearth_ov3/sam3/eval/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/eval/cgf1_eval.py b/src/nn/segearth_ov3/sam3/eval/cgf1_eval.py new file mode 100644 index 0000000..6a0d59f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/cgf1_eval.py @@ -0,0 +1,703 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import contextlib +import copy +import json +import os +import time +from collections import defaultdict +from dataclasses import dataclass +from typing import List, Union + +import numpy as np +import pycocotools.mask as maskUtils +from pycocotools.coco import COCO +from pycocotools.cocoeval import COCOeval +from scipy.optimize import linear_sum_assignment +from tqdm import tqdm + + +@dataclass +class Metric: + name: str + + # whether the metric is computed at the image level or the box level + image_level: bool + + # iou threshold (None is used for image level metrics or to indicate averaging over all thresholds in [0.5:0.95]) + iou_threshold: Union[float, None] + + +CGF1_METRICS = [ + Metric(name="cgF1", image_level=False, iou_threshold=None), + Metric(name="precision", image_level=False, iou_threshold=None), + Metric(name="recall", image_level=False, iou_threshold=None), + Metric(name="F1", image_level=False, iou_threshold=None), + Metric(name="positive_macro_F1", image_level=False, iou_threshold=None), + Metric(name="positive_micro_F1", image_level=False, iou_threshold=None), + Metric(name="positive_micro_precision", image_level=False, iou_threshold=None), + Metric(name="IL_precision", image_level=True, iou_threshold=None), + Metric(name="IL_recall", image_level=True, iou_threshold=None), + Metric(name="IL_F1", image_level=True, iou_threshold=None), + Metric(name="IL_FPR", image_level=True, iou_threshold=None), + Metric(name="IL_MCC", image_level=True, iou_threshold=None), + Metric(name="cgF1", image_level=False, iou_threshold=0.5), + Metric(name="precision", image_level=False, iou_threshold=0.5), + Metric(name="recall", image_level=False, iou_threshold=0.5), + Metric(name="F1", image_level=False, iou_threshold=0.5), + Metric(name="positive_macro_F1", image_level=False, iou_threshold=0.5), + Metric(name="positive_micro_F1", image_level=False, iou_threshold=0.5), + Metric(name="positive_micro_precision", image_level=False, iou_threshold=0.5), + Metric(name="cgF1", image_level=False, iou_threshold=0.75), + Metric(name="precision", image_level=False, iou_threshold=0.75), + Metric(name="recall", image_level=False, iou_threshold=0.75), + Metric(name="F1", image_level=False, iou_threshold=0.75), + Metric(name="positive_macro_F1", image_level=False, iou_threshold=0.75), + Metric(name="positive_micro_F1", image_level=False, iou_threshold=0.75), + Metric(name="positive_micro_precision", image_level=False, iou_threshold=0.75), +] + + +class COCOCustom(COCO): + """COCO class from pycocotools with tiny modifications for speed""" + + def createIndex(self): + # create index + print("creating index...") + anns, cats, imgs = {}, {}, {} + imgToAnns, catToImgs = defaultdict(list), defaultdict(list) + if "annotations" in self.dataset: + for ann in self.dataset["annotations"]: + imgToAnns[ann["image_id"]].append(ann) + anns[ann["id"]] = ann + + if "images" in self.dataset: + # MODIFICATION: do not reload imgs if they are already there + if self.imgs: + imgs = self.imgs + else: + for img in self.dataset["images"]: + imgs[img["id"]] = img + # END MODIFICATION + + if "categories" in self.dataset: + for cat in self.dataset["categories"]: + cats[cat["id"]] = cat + + if "annotations" in self.dataset and "categories" in self.dataset: + for ann in self.dataset["annotations"]: + catToImgs[ann["category_id"]].append(ann["image_id"]) + + print("index created!") + + # create class members + self.anns = anns + self.imgToAnns = imgToAnns + self.catToImgs = catToImgs + self.imgs = imgs + self.cats = cats + + def loadRes(self, resFile): + """ + Load result file and return a result api object. + :param resFile (str) : file name of result file + :return: res (obj) : result api object + """ + res = COCOCustom() + res.dataset["info"] = copy.deepcopy(self.dataset.get("info", {})) + # MODIFICATION: no copy + # res.dataset['images'] = [img for img in self.dataset['images']] + res.dataset["images"] = self.dataset["images"] + # END MODIFICATION + + print("Loading and preparing results...") + tic = time.time() + if type(resFile) == str: + with open(resFile) as f: + anns = json.load(f) + elif type(resFile) == np.ndarray: + anns = self.loadNumpyAnnotations(resFile) + else: + anns = resFile + assert type(anns) == list, "results in not an array of objects" + annsImgIds = [ann["image_id"] for ann in anns] + # MODIFICATION: faster and cached subset check + if not hasattr(self, "img_id_set"): + self.img_id_set = set(self.getImgIds()) + assert set(annsImgIds).issubset( + self.img_id_set + ), "Results do not correspond to current coco set" + # END MODIFICATION + if "caption" in anns[0]: + imgIds = set([img["id"] for img in res.dataset["images"]]) & set( + [ann["image_id"] for ann in anns] + ) + res.dataset["images"] = [ + img for img in res.dataset["images"] if img["id"] in imgIds + ] + for id, ann in enumerate(anns): + ann["id"] = id + 1 + elif "bbox" in anns[0] and not anns[0]["bbox"] == []: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + bb = ann["bbox"] + x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]] + if not "segmentation" in ann: + ann["segmentation"] = [[x1, y1, x1, y2, x2, y2, x2, y1]] + ann["area"] = bb[2] * bb[3] + ann["id"] = id + 1 + ann["iscrowd"] = 0 + elif "segmentation" in anns[0]: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + # now only support compressed RLE format as segmentation results + ann["area"] = maskUtils.area(ann["segmentation"]) + if not "bbox" in ann: + ann["bbox"] = maskUtils.toBbox(ann["segmentation"]) + ann["id"] = id + 1 + ann["iscrowd"] = 0 + elif "keypoints" in anns[0]: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + s = ann["keypoints"] + x = s[0::3] + y = s[1::3] + x0, x1, y0, y1 = np.min(x), np.max(x), np.min(y), np.max(y) + ann["area"] = (x1 - x0) * (y1 - y0) + ann["id"] = id + 1 + ann["bbox"] = [x0, y0, x1 - x0, y1 - y0] + print("DONE (t={:0.2f}s)".format(time.time() - tic)) + + res.dataset["annotations"] = anns + # MODIFICATION: inherit images + res.imgs = self.imgs + # END MODIFICATION + res.createIndex() + return res + + +class CGF1Eval(COCOeval): + """ + This evaluator is based upon COCO evaluation, but evaluates the model in a more realistic setting + for downstream applications. + See SAM3 paper for the details on the CGF1 metric. + + Do not use this evaluator directly. Prefer the CGF1Evaluator wrapper. + + Notes: + - This evaluator does not support per-category evaluation (in the way defined by pyCocotools) + - In open vocabulary settings, we have different noun-phrases for each image. What we call an "image_id" here is actually an (image, noun-phrase) pair. So in every "image_id" there is only one category, implied by the noun-phrase. Thus we can ignore the usual coco "category" field of the predictions + """ + + def __init__( + self, + coco_gt=None, + coco_dt=None, + iouType="segm", + threshold=0.5, + ): + """ + Args: + coco_gt (COCO): ground truth COCO API + coco_dt (COCO): detections COCO API + iou_type (str): type of IoU to evaluate + threshold (float): threshold for predictions + """ + super().__init__(coco_gt, coco_dt, iouType) + self.threshold = threshold + + self.params.useCats = False + self.params.areaRng = [[0**2, 1e5**2]] + self.params.areaRngLbl = ["all"] + self.params.maxDets = [1000000] + + def computeIoU(self, imgId, catId): + # Same as the original COCOeval.computeIoU, but without sorting + p = self.params + if p.useCats: + gt = self._gts[imgId, catId] + dt = self._dts[imgId, catId] + else: + gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + if len(gt) == 0 and len(dt) == 0: + return [] + + if p.iouType == "segm": + g = [g["segmentation"] for g in gt] + d = [d["segmentation"] for d in dt] + elif p.iouType == "bbox": + g = [g["bbox"] for g in gt] + d = [d["bbox"] for d in dt] + else: + raise Exception("unknown iouType for iou computation") + + # compute iou between each dt and gt region + iscrowd = [int(o["iscrowd"]) for o in gt] + ious = maskUtils.iou(d, g, iscrowd) + return ious + + def evaluateImg(self, imgId, catId, aRng, maxDet): + """ + perform evaluation for single category and image + :return: dict (single image results) + """ + p = self.params + assert not p.useCats, "This evaluator does not support per-category evaluation." + assert catId == -1 + all_gts = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + keep_gt = np.array([not g["ignore"] for g in all_gts], dtype=bool) + gt = [g for g in all_gts if not g["ignore"]] + all_dts = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + keep_dt = np.array([d["score"] >= self.threshold for d in all_dts], dtype=bool) + dt = [d for d in all_dts if d["score"] >= self.threshold] + if len(gt) == 0 and len(dt) == 0: + # This is a "true negative" case, where there are no GTs and no predictions + # The box-level metrics are ill-defined, so we don't add them to this dict + return { + "image_id": imgId, + "IL_TP": 0, + "IL_TN": 1, + "IL_FP": 0, + "IL_FN": 0, + "num_dt": len(dt), + } + + if len(gt) > 0 and len(dt) == 0: + # This is a "false negative" case, where there are GTs but no predictions + return { + "image_id": imgId, + "IL_TP": 0, + "IL_TN": 0, + "IL_FP": 0, + "IL_FN": 1, + "TPs": np.zeros((len(p.iouThrs),), dtype=np.int64), + "FPs": np.zeros((len(p.iouThrs),), dtype=np.int64), + "FNs": np.ones((len(p.iouThrs),), dtype=np.int64) * len(gt), + "local_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64), + "local_positive_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64), + "num_dt": len(dt), + } + + # Load pre-computed ious + ious = self.ious[(imgId, catId)] + + # compute matching + if len(ious) == 0: + ious = np.zeros((len(dt), len(gt))) + else: + ious = ious[keep_dt, :][:, keep_gt] + assert ious.shape == (len(dt), len(gt)) + + matched_dt, matched_gt = linear_sum_assignment(-ious) + + match_scores = ious[matched_dt, matched_gt] + + TPs, FPs, FNs = [], [], [] + IL_perfect = [] + for thresh in p.iouThrs: + TP = (match_scores >= thresh).sum() + FP = len(dt) - TP + FN = len(gt) - TP + assert ( + FP >= 0 and FN >= 0 + ), f"FP: {FP}, FN: {FN}, TP: {TP}, match_scores: {match_scores}, len(dt): {len(dt)}, len(gt): {len(gt)}, ious: {ious}" + TPs.append(TP) + FPs.append(FP) + FNs.append(FN) + + if FP == FN and FP == 0: + IL_perfect.append(1) + else: + IL_perfect.append(0) + + TPs = np.array(TPs, dtype=np.int64) + FPs = np.array(FPs, dtype=np.int64) + FNs = np.array(FNs, dtype=np.int64) + IL_perfect = np.array(IL_perfect, dtype=np.int64) + + # compute precision recall and F1 + precision = TPs / (TPs + FPs + 1e-4) + assert np.all(precision <= 1) + recall = TPs / (TPs + FNs + 1e-4) + assert np.all(recall <= 1) + F1 = 2 * precision * recall / (precision + recall + 1e-4) + + result = { + "image_id": imgId, + "TPs": TPs, + "FPs": FPs, + "FNs": FNs, + "local_F1s": F1, + "IL_TP": (len(gt) > 0) and (len(dt) > 0), + "IL_FP": (len(gt) == 0) and (len(dt) > 0), + "IL_TN": (len(gt) == 0) and (len(dt) == 0), + "IL_FN": (len(gt) > 0) and (len(dt) == 0), + "num_dt": len(dt), + } + if len(gt) > 0 and len(dt) > 0: + result["local_positive_F1s"] = F1 + return result + + def accumulate(self, p=None): + """ + Accumulate per image evaluation results and store the result in self.eval + :param p: input params for evaluation + :return: None + """ + if self.evalImgs is None or len(self.evalImgs) == 0: + print("Please run evaluate() first") + # allows input customized parameters + if p is None: + p = self.params + + setImgIds = set(p.imgIds) + + # TPs, FPs, FNs + TPs = np.zeros((len(p.iouThrs),), dtype=np.int64) + FPs = np.zeros((len(p.iouThrs),), dtype=np.int64) + pmFPs = np.zeros((len(p.iouThrs),), dtype=np.int64) + FNs = np.zeros((len(p.iouThrs),), dtype=np.int64) + local_F1s = np.zeros((len(p.iouThrs),), dtype=np.float64) + + # Image level metrics + IL_TPs = 0 + IL_FPs = 0 + IL_TNs = 0 + IL_FNs = 0 + + valid_img_count = 0 + valid_F1_count = 0 + evaledImgIds = set() + for res in self.evalImgs: + if res["image_id"] not in setImgIds: + continue + evaledImgIds.add(res["image_id"]) + IL_TPs += res["IL_TP"] + IL_FPs += res["IL_FP"] + IL_TNs += res["IL_TN"] + IL_FNs += res["IL_FN"] + + if "TPs" not in res: + continue + + TPs += res["TPs"] + FPs += res["FPs"] + FNs += res["FNs"] + valid_img_count += 1 + + if "local_positive_F1s" in res: + local_F1s += res["local_positive_F1s"] + pmFPs += res["FPs"] + if res["num_dt"] > 0: + valid_F1_count += 1 + + assert len(setImgIds - evaledImgIds) == 0, ( + f"{len(setImgIds - evaledImgIds)} images not evaluated. " + f"Here are the IDs of the first 3: {list(setImgIds - evaledImgIds)[:3]}" + ) + + # compute precision recall and F1 + precision = TPs / (TPs + FPs + 1e-4) + positive_micro_precision = TPs / (TPs + pmFPs + 1e-4) + assert np.all(precision <= 1) + recall = TPs / (TPs + FNs + 1e-4) + assert np.all(recall <= 1) + F1 = 2 * precision * recall / (precision + recall + 1e-4) + positive_micro_F1 = ( + 2 + * positive_micro_precision + * recall + / (positive_micro_precision + recall + 1e-4) + ) + + IL_rec = IL_TPs / (IL_TPs + IL_FNs + 1e-6) + IL_prec = IL_TPs / (IL_TPs + IL_FPs + 1e-6) + IL_F1 = 2 * IL_prec * IL_rec / (IL_prec + IL_rec + 1e-6) + IL_FPR = IL_FPs / (IL_FPs + IL_TNs + 1e-6) + IL_MCC = float(IL_TPs * IL_TNs - IL_FPs * IL_FNs) / ( + ( + float(IL_TPs + IL_FPs) + * float(IL_TPs + IL_FNs) + * float(IL_TNs + IL_FPs) + * float(IL_TNs + IL_FNs) + ) + ** 0.5 + + 1e-6 + ) + + self.eval = { + "params": p, + "TPs": TPs, + "FPs": FPs, + "positive_micro_FPs": pmFPs, + "FNs": FNs, + "precision": precision, + "positive_micro_precision": positive_micro_precision, + "recall": recall, + "F1": F1, + "positive_micro_F1": positive_micro_F1, + "positive_macro_F1": local_F1s / valid_F1_count, + "IL_recall": IL_rec, + "IL_precision": IL_prec, + "IL_F1": IL_F1, + "IL_FPR": IL_FPR, + "IL_MCC": IL_MCC, + } + self.eval["cgF1"] = self.eval["positive_micro_F1"] * self.eval["IL_MCC"] + + def summarize(self): + """ + Compute and display summary metrics for evaluation results. + """ + if not self.eval: + raise Exception("Please run accumulate() first") + + def _summarize(iouThr=None, metric=""): + p = self.params + iStr = " {:<18} @[ IoU={:<9}] = {:0.3f}" + titleStr = "Average " + metric + iouStr = ( + "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1]) + if iouThr is None + else "{:0.2f}".format(iouThr) + ) + + s = self.eval[metric] + # IoU + if iouThr is not None: + t = np.where(iouThr == p.iouThrs)[0] + s = s[t] + + if len(s[s > -1]) == 0: + mean_s = -1 + else: + mean_s = np.mean(s[s > -1]) + print(iStr.format(titleStr, iouStr, mean_s)) + return mean_s + + def _summarize_single(metric=""): + titleStr = "Average " + metric + iStr = " {:<35} = {:0.3f}" + s = self.eval[metric] + print(iStr.format(titleStr, s)) + return s + + def _summarizeDets(): + stats = [] + + for metric in CGF1_METRICS: + if metric.image_level: + stats.append(_summarize_single(metric=metric.name)) + else: + stats.append( + _summarize(iouThr=metric.iou_threshold, metric=metric.name) + ) + return np.asarray(stats) + + summarize = _summarizeDets + self.stats = summarize() + + +def _evaluate(self): + """ + Run per image evaluation on given images and store results (a list of dict) in self.evalImgs + """ + p = self.params + # add backward compatibility if useSegm is specified in params + p.imgIds = list(np.unique(p.imgIds)) + p.useCats = False + p.maxDets = sorted(p.maxDets) + self.params = p + + self._prepare() + # loop through images, area range, max detection number + catIds = [-1] + + if p.iouType == "segm" or p.iouType == "bbox": + computeIoU = self.computeIoU + else: + raise RuntimeError(f"Unsupported iou {p.iouType}") + self.ious = { + (imgId, catId): computeIoU(imgId, catId) + for imgId in p.imgIds + for catId in catIds + } + + maxDet = p.maxDets[-1] + evalImgs = [ + self.evaluateImg(imgId, catId, areaRng, maxDet) + for catId in catIds + for areaRng in p.areaRng + for imgId in p.imgIds + ] + # this is NOT in the pycocotools code, but could be done outside + evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds)) + return p.imgIds, evalImgs + + +class CGF1Evaluator: + """ + Wrapper class for cgF1 evaluation. + This supports the oracle setting (when several ground-truths are available per image) + """ + + def __init__( + self, + gt_path: Union[str, List[str]], + iou_type="segm", + verbose=False, + ): + """ + Args: + gt_path (str or list of str): path(s) to ground truth COCO json file(s) + iou_type (str): type of IoU to evaluate + threshold (float): threshold for predictions + """ + self.gt_paths = gt_path if isinstance(gt_path, list) else [gt_path] + self.iou_type = iou_type + + self.coco_gts = [COCOCustom(gt) for gt in self.gt_paths] + + self.verbose = verbose + + self.coco_evals = [] + for i, coco_gt in enumerate(self.coco_gts): + self.coco_evals.append( + CGF1Eval( + coco_gt=coco_gt, + iouType=iou_type, + ) + ) + self.coco_evals[i].useCats = False + + exclude_img_ids = set() + # exclude_img_ids are the ids that are not exhaustively annotated in any of the other gts + for coco_gt in self.coco_gts[1:]: + exclude_img_ids = exclude_img_ids.union( + { + img["id"] + for img in coco_gt.dataset["images"] + if not img["is_instance_exhaustive"] + } + ) + # we only eval on instance exhaustive queries + self.eval_img_ids = [ + img["id"] + for img in self.coco_gts[0].dataset["images"] + if (img["is_instance_exhaustive"] and img["id"] not in exclude_img_ids) + ] + + def evaluate(self, pred_file: str): + """ + Evaluate the detections using cgF1 metric. + + Args: + pred_file: path to the predictions COCO json file + + """ + assert len(self.coco_gts) > 0, "No ground truth provided for evaluation." + assert len(self.coco_gts) == len( + self.coco_evals + ), "Mismatch in number of ground truths and evaluators." + + if self.verbose: + print(f"Loading predictions from {pred_file}") + + with open(pred_file, "r") as f: + preds = json.load(f) + + if self.verbose: + print(f"Loaded {len(preds)} predictions") + + img2preds = defaultdict(list) + for pred in preds: + img2preds[pred["image_id"]].append(pred) + + all_eval_imgs = [] + for img_id in tqdm(self.eval_img_ids, disable=not self.verbose): + results = img2preds[img_id] + all_scorings = [] + for cur_coco_gt, coco_eval in zip(self.coco_gts, self.coco_evals): + # suppress pycocotools prints + with open(os.devnull, "w") as devnull: + with contextlib.redirect_stdout(devnull): + coco_dt = ( + cur_coco_gt.loadRes(results) if results else COCOCustom() + ) + + coco_eval.cocoDt = coco_dt + coco_eval.params.imgIds = [img_id] + coco_eval.params.useCats = False + img_ids, eval_imgs = _evaluate(coco_eval) + all_scorings.append(eval_imgs) + selected = self._select_best_scoring(all_scorings) + all_eval_imgs.append(selected) + + # After this point, we have selected the best scoring per image among several ground truths + # we can now accumulate and summarize, using only the first coco_eval + + self.coco_evals[0].evalImgs = list( + np.concatenate(all_eval_imgs, axis=2).flatten() + ) + self.coco_evals[0].params.imgIds = self.eval_img_ids + self.coco_evals[0]._paramsEval = copy.deepcopy(self.coco_evals[0].params) + + if self.verbose: + print(f"Accumulating results") + self.coco_evals[0].accumulate() + print("cgF1 metric, IoU type={}".format(self.iou_type)) + self.coco_evals[0].summarize() + print() + + out = {} + for i, value in enumerate(self.coco_evals[0].stats): + name = CGF1_METRICS[i].name + if CGF1_METRICS[i].iou_threshold is not None: + name = f"{name}@{CGF1_METRICS[i].iou_threshold}" + out[f"cgF1_eval_{self.iou_type}_{name}"] = float(value) + + return out + + @staticmethod + def _select_best_scoring(scorings): + # This function is used for "oracle" type evaluation. + # It accepts the evaluation results with respect to several ground truths, and picks the best + if len(scorings) == 1: + return scorings[0] + + assert ( + scorings[0].ndim == 3 + ), f"Expecting results in [numCats, numAreas, numImgs] format, got {scorings[0].shape}" + assert ( + scorings[0].shape[0] == 1 + ), f"Expecting a single category, got {scorings[0].shape[0]}" + + for scoring in scorings: + assert ( + scoring.shape == scorings[0].shape + ), f"Shape mismatch: {scoring.shape}, {scorings[0].shape}" + + selected_imgs = [] + for img_id in range(scorings[0].shape[-1]): + best = scorings[0][:, :, img_id] + + for scoring in scorings[1:]: + current = scoring[:, :, img_id] + if "local_F1s" in best[0, 0] and "local_F1s" in current[0, 0]: + # we were able to compute a F1 score for this particular image in both evaluations + # best["local_F1s"] contains the results at various IoU thresholds. We simply take the average for comparision + best_score = best[0, 0]["local_F1s"].mean() + current_score = current[0, 0]["local_F1s"].mean() + if current_score > best_score: + best = current + + else: + # If we're here, it means that in that in some evaluation we were not able to get a valid local F1 + # This happens when both the predictions and targets are empty. In that case, we can assume it's a perfect prediction + if "local_F1s" not in current[0, 0]: + best = current + selected_imgs.append(best) + result = np.stack(selected_imgs, axis=-1) + assert result.shape == scorings[0].shape + return result diff --git a/src/nn/segearth_ov3/sam3/eval/coco_eval.py b/src/nn/segearth_ov3/sam3/eval/coco_eval.py new file mode 100644 index 0000000..167f8ef --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/coco_eval.py @@ -0,0 +1,916 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +COCO evaluator that works in distributed mode. + +Mostly copy-paste from https://github.com/pytorch/vision/blob/edfd5a7/references/detection/coco_eval.py +The difference is that there is less copy-pasting from pycocotools +in the end of the file, as python3 can suppress prints with contextlib +""" + +import contextlib +import copy +import json +import logging +import os +import pickle +from collections import defaultdict +from pathlib import Path + +from typing import Any, List, Optional + +import numpy as np + +import pycocotools.mask as mask_utils +import torch +from iopath.common.file_io import g_pathmgr +from pycocotools.coco import COCO +from pycocotools.cocoeval import COCOeval + +from sam3.train.masks_ops import rle_encode + +from sam3.train.utils.distributed import ( + all_gather, + gather_to_rank_0_via_filesys, + get_rank, + is_main_process, +) + +RARITY_BUCKETS = {0: "frequent", 1: "common", 2: "medium", 3: "rare"} + + +class CocoEvaluator: + def __init__( + self, + coco_gt, + iou_types: List[str], + useCats: bool, + dump_dir: Optional[str], + postprocessor, + average_by_rarity=False, + metrics_dump_dir: Optional[str] = None, + gather_pred_via_filesys=False, + use_normalized_areas=True, + maxdets=[1, 10, 100], + exhaustive_only=False, + all_exhaustive_only=True, + ): + """Online coco evaluator. It will evaluate images as they are generated by the model, then accumulate/summarize at the end + + Args: + - coco_gt: COCO api object containing the gt + - iou_types: can be either "bbox" or "segm" + - useCats: If true, categories will be used for evaluation + - dump_dir: if non null, then the predictions will be dumped in that directory + - postprocessor: Module to convert the model's output into the coco format + - average_by_rarity: if true then we expect the images information in the gt dataset + to have a "rarity" field. Then the AP will be computed on all rarity buckets + individually, then averaged + - gather_pred_via_filesys: if true, we use the filesystem for collective gathers + - use_normalized_areas: if true, the areas of the objects in the GT are assumed to be + normalized by the area of the image. In that case, the size buckets are adjusted + - maxdets: maximal number of detections to be evaluated on each image. + - exhaustive_only: If true, we restrict eval only to exhaustive annotations + - all_exhaustive_only: If true, datapoints are restricted only to those with all exhaustive annotations + + """ + # coco_gt = copy.deepcopy(coco_gt) + self.coco_gts = [coco_gt] if not isinstance(coco_gt, list) else coco_gt + assert len(maxdets) == 3, f"expecting 3 detection threshold, got {len(maxdets)}" + + self.use_normalized_areas = use_normalized_areas + self.iou_types = iou_types + self.useCats = useCats + self.maxdets = maxdets + self.dump = None + self.dump_dir = dump_dir + if self.dump_dir is not None: + self.dump = [] + if is_main_process(): + if not os.path.exists(self.dump_dir): + os.makedirs(self.dump_dir, exist_ok=True) + logging.info(f"Create the folder: {dump_dir}") + + self.initialized = False + + # Whether to gather predictions through filesystem (instead of torch + # collective ops; requiring a shared filesystem across all ranks) + self.gather_pred_via_filesys = gather_pred_via_filesys + self.use_self_evaluate = True # CPP version is disabled + self.postprocessor = postprocessor + self.average_by_rarity = average_by_rarity + self.exhaustive_only = exhaustive_only + self.all_exhaustive_only = all_exhaustive_only + self.metrics_dump_dir = metrics_dump_dir + if self.metrics_dump_dir is not None: + if is_main_process(): + if not os.path.exists(self.metrics_dump_dir): + os.makedirs(self.metrics_dump_dir, exist_ok=True) + logging.info(f"Create the folder: {metrics_dump_dir}") + + def _lazy_init(self, coco_cls=COCO): + if self.initialized: + return + + self.initialized = True + + self.coco_gts = [ + coco_cls(g_pathmgr.get_local_path(gt)) if isinstance(gt, str) else gt + for gt in self.coco_gts + ] + + self.reset() + + self.eval_img_ids = None + + if self.exhaustive_only: + exclude_img_ids = set() + # exclude_img_ids are the ids that are not exhaustively annotated in any of the other gts + if self.all_exhaustive_only: + for coco_gt in self.coco_gts[1:]: + exclude_img_ids = exclude_img_ids.union( + { + img["id"] + for img in coco_gt.dataset["images"] + if not img["is_instance_exhaustive"] + } + ) + # we only eval on instance exhaustive queries + self.eval_img_ids = [ + img["id"] + for img in self.coco_gts[0].dataset["images"] + if (img["is_instance_exhaustive"] and img["id"] not in exclude_img_ids) + ] + + self.rarity_buckets = None + if self.average_by_rarity: + self.rarity_buckets = defaultdict(list) + eval_img_ids_set = ( + set(self.eval_img_ids) if self.eval_img_ids is not None else None + ) + for img in self.coco_gts[0].dataset["images"]: + if self.eval_img_ids is not None and img["id"] not in eval_img_ids_set: + continue + self.rarity_buckets[img["rarity"]].append(img["id"]) + print("Rarity buckets sizes:") + for k, v in self.rarity_buckets.items(): + print(f"{k}: {len(v)}") + + def set_sync_device(self, device: torch.device) -> Any: + self._sync_device = device + + def _evaluate(self, *args, **kwargs): + return evaluate(*args, **kwargs) + + def _loadRes(self, *args, **kwargs): + return loadRes(*args, **kwargs) + + def update(self, *args, **kwargs): + self._lazy_init() + predictions = self.postprocessor.process_results(*args, **kwargs) + + img_ids = list(np.unique(list(predictions.keys()))) + self.img_ids.extend(img_ids) + + for iou_type in self.iou_types: + results = self.prepare(predictions, iou_type) + self._dump(results) + + assert len(self.coco_gts) == len(self.coco_evals) + all_scorings = [] + for cur_coco_gt, cur_coco_eval in zip(self.coco_gts, self.coco_evals): + # suppress pycocotools prints + with open(os.devnull, "w") as devnull: + with contextlib.redirect_stdout(devnull): + coco_dt = ( + self._loadRes(cur_coco_gt, results) if results else COCO() + ) + + coco_eval = cur_coco_eval[iou_type] + + coco_eval.cocoDt = coco_dt + coco_eval.params.imgIds = list(img_ids) + coco_eval.params.useCats = self.useCats + coco_eval.params.maxDets = self.maxdets + img_ids, eval_imgs = self._evaluate(coco_eval, self.use_self_evaluate) + all_scorings.append(eval_imgs) + + selected = self.select_best_scoring(all_scorings) + self.eval_imgs[iou_type].append(selected) + + def select_best_scoring(self, scorings): + # This function is used for "oracle" type evaluation. + # It accepts the evaluation results with respect to several ground truths, and picks the best + if len(scorings) == 1: + return scorings[0] + + # Currently we don't support Oracle Phrase AP. + # To implement it, we likely need to modify the cpp code since the eval_image type is opaque + raise RuntimeError("Not implemented") + + def _dump(self, results): + if self.dump is not None: + dumped_results = copy.deepcopy(results) + for r in dumped_results: + if "bbox" not in self.iou_types and "bbox" in r: + del r["bbox"] + elif "bbox" in r: + r["bbox"] = [round(coord, 5) for coord in r["bbox"]] + r["score"] = round(r["score"], 5) + self.dump.extend(dumped_results) + + def synchronize_between_processes(self): + self._lazy_init() + logging.info("Coco evaluator: Synchronizing between processes") + for iou_type in self.iou_types: + if len(self.eval_imgs[iou_type]) > 0: + self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2) + else: + num_areas = len(self.coco_evals[0][iou_type].params.areaRng) + # assuming 1 class + assert not self.useCats + self.eval_imgs[iou_type] = np.empty((1, num_areas, 0)) + create_common_coco_eval( + self.coco_evals[0][iou_type], + self.img_ids, + self.eval_imgs[iou_type], + use_self_evaluate=self.use_self_evaluate, + gather_pred_via_filesys=self.gather_pred_via_filesys, + metrics_dump_dir=self.metrics_dump_dir, + ) + if self.dump is not None: + dumped_file = Path(self.dump_dir) / f"coco_predictions_{get_rank()}.json" + logging.info(f"COCO evaluator: Dumping local predictions to {dumped_file}") + with g_pathmgr.open(str(dumped_file), "w") as f: + json.dump(self.dump, f) + + # if self.gather_pred_via_filesys: + # dump = gather_to_rank_0_via_filesys(self.dump) + # else: + # dump = all_gather(self.dump, force_cpu=True) + # self.dump = sum(dump, []) + + def accumulate(self, imgIds=None): + self._lazy_init() + logging.info( + f"Coco evaluator: Accumulating on {len(imgIds) if imgIds is not None else 'all'} images" + ) + if not is_main_process(): + return + + if imgIds is None: + for coco_eval in self.coco_evals[0].values(): + accumulate(coco_eval, use_self_eval=self.use_self_evaluate) + + if imgIds is not None: + imgIds = set(imgIds) + for coco_eval in self.coco_evals[0].values(): + p = coco_eval.params + id_mask = np.array([(i in imgIds) for i in p.imgIds], dtype=bool) + old_img_ids = p.imgIds + coco_eval.params.imgIds = np.asarray(p.imgIds)[id_mask] + old_img_evals = coco_eval.evalImgs + catIds = p.catIds if p.useCats else [-1] + coco_eval.evalImgs = list( + np.asarray(coco_eval.evalImgs) + .reshape(len(catIds), len(p.areaRng), len(old_img_ids))[ + ..., id_mask + ] + .flatten() + ) + accumulate(coco_eval, use_self_eval=self.use_self_evaluate) + coco_eval.evalImgs = old_img_evals + coco_eval.params.imgIds = old_img_ids + + def summarize(self): + self._lazy_init() + logging.info("Coco evaluator: Summarizing") + if not is_main_process(): + return {} + + outs = {} + if self.rarity_buckets is None: + self.accumulate(self.eval_img_ids) + for iou_type, coco_eval in self.coco_evals[0].items(): + print("IoU metric: {}".format(iou_type)) + summarize(coco_eval) + + if "bbox" in self.coco_evals[0]: + for key, value in zip(*self.coco_evals[0]["bbox"].stats): + outs[f"coco_eval_bbox_{key}"] = value + if "segm" in self.coco_evals[0]: + for key, value in zip(*self.coco_evals[0]["segm"].stats): + outs[f"coco_eval_masks_{key}"] = value + else: + total_stats = {} + all_keys = {} + for bucket, img_list in self.rarity_buckets.items(): + self.accumulate(imgIds=img_list) + bucket_name = RARITY_BUCKETS[bucket] + for iou_type, coco_eval in self.coco_evals[0].items(): + print(f"IoU metric: {iou_type}. Rarity bucket: {bucket_name}") + summarize(coco_eval) + + if "bbox" in self.coco_evals[0]: + if "bbox" not in total_stats: + total_stats["bbox"] = np.zeros_like( + self.coco_evals[0]["bbox"].stats[1] + ) + all_keys["bbox"] = self.coco_evals[0]["bbox"].stats[0] + total_stats["bbox"] += self.coco_evals[0]["bbox"].stats[1] + for key, value in zip(*self.coco_evals[0]["bbox"].stats): + outs[f"coco_eval_bbox_{bucket_name}_{key}"] = value + if "segm" in self.coco_evals[0]: + if "segm" not in total_stats: + total_stats["segm"] = np.zeros_like( + self.coco_evals[0]["segm"].stats[1] + ) + all_keys["segm"] = self.coco_evals[0]["segm"].stats[0] + total_stats["segm"] += self.coco_evals[0]["segm"].stats[1] + for key, value in zip(*self.coco_evals[0]["segm"].stats): + outs[f"coco_eval_masks_{bucket_name}_{key}"] = value + + if "bbox" in total_stats: + total_stats["bbox"] /= len(self.rarity_buckets) + for key, value in zip(all_keys["bbox"], total_stats["bbox"]): + outs[f"coco_eval_bbox_{key}"] = value + if "segm" in total_stats: + total_stats["segm"] /= len(self.rarity_buckets) + for key, value in zip(all_keys["segm"], total_stats["segm"]): + outs[f"coco_eval_masks_{key}"] = value + + # if self.dump is not None: + # assert self.dump_dir is not None + # logging.info("Coco evaluator: Dumping the global result file to disk") + # with g_pathmgr.open(str(Path(self.dump_dir) / "coco_eval.json"), "w") as f: + # json.dump(self.dump, f) + return outs + + def compute_synced(self): + self._lazy_init() + self.synchronize_between_processes() + return self.summarize() + + def compute(self): + self._lazy_init() + return {"": 0.0} + + def reset(self, cocoeval_cls=COCOeval): + self.coco_evals = [{} for _ in range(len(self.coco_gts))] + for i, coco_gt in enumerate(self.coco_gts): + for iou_type in self.iou_types: + self.coco_evals[i][iou_type] = cocoeval_cls(coco_gt, iouType=iou_type) + self.coco_evals[i][iou_type].params.useCats = self.useCats + self.coco_evals[i][iou_type].params.maxDets = self.maxdets + if self.use_normalized_areas: + self.coco_evals[i][iou_type].params.areaRng = [ + [0, 1e5], + [0, 0.001], + [0.001, 0.01], + [0.01, 0.1], + [0.1, 0.5], + [0.5, 0.95], + [0.95, 1e5], + ] + self.coco_evals[i][iou_type].params.areaRngLbl = [ + "all", + "tiny", + "small", + "medium", + "large", + "huge", + "whole_image", + ] + + self.img_ids = [] + self.eval_imgs = {k: [] for k in self.iou_types} + if self.dump is not None: + self.dump = [] + + def write(self, stats): + self._lazy_init() + """Write the results in the stats dict""" + if "bbox" in self.coco_evals[0]: + stats["coco_eval_bbox"] = self.coco_evals[0]["bbox"].stats.tolist() + if "segm" in self.coco_evals[0]: + stats["coco_eval_masks"] = self.coco_evals[0]["segm"].stats.tolist() + return stats + + def prepare(self, predictions, iou_type): + self._lazy_init() + if iou_type == "bbox": + return self.prepare_for_coco_detection(predictions) + elif iou_type == "segm": + return self.prepare_for_coco_segmentation(predictions) + elif iou_type == "keypoints": + return self.prepare_for_coco_keypoint(predictions) + else: + raise ValueError("Unknown iou type {}".format(iou_type)) + + def prepare_for_coco_detection(self, predictions): + self._lazy_init() + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + + coco_results.extend( + [ + { + "image_id": original_id, + "category_id": labels[k], + "bbox": box, + "score": scores[k], + } + for k, box in enumerate(boxes) + ] + ) + return coco_results + + @torch.no_grad() + def prepare_for_coco_segmentation(self, predictions): + self._lazy_init() + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + boundaries, dilated_boundaries = None, None + if "boundaries" in prediction: + boundaries = prediction["boundaries"] + dilated_boundaries = prediction["dilated_boundaries"] + assert dilated_boundaries is not None + assert len(scores) == len(boundaries) + + if "masks_rle" in prediction: + rles = prediction["masks_rle"] + areas = [] + for rle in rles: + cur_area = mask_utils.area(rle) + h, w = rle["size"] + areas.append(cur_area / (h * w)) + else: + masks = prediction["masks"] + + masks = masks > 0.5 + h, w = masks.shape[-2:] + + areas = masks.flatten(1).sum(1) / (h * w) + areas = areas.tolist() + + rles = rle_encode(masks.squeeze(1)) + + # memory clean + del masks + del prediction["masks"] + + assert len(areas) == len(rles) == len(scores) + for k, rle in enumerate(rles): + payload = { + "image_id": original_id, + "category_id": labels[k], + "segmentation": rle, + "score": scores[k], + "area": areas[k], + } + if boundaries is not None: + payload["boundary"] = boundaries[k] + payload["dilated_boundary"] = dilated_boundaries[k] + + coco_results.append(payload) + + return coco_results + + def prepare_for_coco_keypoint(self, predictions): + self._lazy_init() + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + keypoints = prediction["keypoints"] + keypoints = keypoints.flatten(start_dim=1).tolist() + + coco_results.extend( + [ + { + "image_id": original_id, + "category_id": labels[k], + "keypoints": keypoint, + "score": scores[k], + } + for k, keypoint in enumerate(keypoints) + ] + ) + return coco_results + + +def convert_to_xywh(boxes): + xmin, ymin, xmax, ymax = boxes.unbind(-1) + return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=-1) + + +def merge(img_ids, eval_imgs, gather_pred_via_filesys=False): + if gather_pred_via_filesys: + # only gather the predictions to rank 0 (other ranks will receive empty + # lists for `all_img_ids` and `all_eval_imgs`, which should be OK as + # merging and evaluation are only done on rank 0) + all_img_ids = gather_to_rank_0_via_filesys(img_ids) + all_eval_imgs = gather_to_rank_0_via_filesys(eval_imgs) + else: + all_img_ids = all_gather(img_ids, force_cpu=True) + all_eval_imgs = all_gather(eval_imgs, force_cpu=True) + if not is_main_process(): + return None, None + + merged_img_ids = [] + for p in all_img_ids: + merged_img_ids.extend(p) + + merged_eval_imgs = [] + for p in all_eval_imgs: + merged_eval_imgs.append(p) + + merged_img_ids = np.array(merged_img_ids) + merged_eval_imgs = np.concatenate(merged_eval_imgs, 2) + + # keep only unique (and in sorted order) images + merged_img_ids, idx = np.unique(merged_img_ids, return_index=True) + merged_eval_imgs = merged_eval_imgs[..., idx] + + return merged_img_ids, merged_eval_imgs + + +def create_common_coco_eval( + coco_eval, + img_ids, + eval_imgs, + use_self_evaluate, + gather_pred_via_filesys=False, + metrics_dump_dir=None, +): + img_ids, eval_imgs = merge(img_ids, eval_imgs, gather_pred_via_filesys) + if not is_main_process(): + return + if metrics_dump_dir is not None: + dumped_file = ( + Path(metrics_dump_dir) / f"coco_eval_img_metrics_{get_rank()}.json" + ) + logging.info(f"COCO evaluator: Dumping local predictions to {dumped_file}") + with g_pathmgr.open(str(dumped_file), "w") as f: + json.dump(eval_imgs.squeeze(), f, default=lambda x: x.tolist()) + img_ids = list(img_ids) + + # If some images were not predicted, we need to create dummy detections for them + missing_img_ids = set(coco_eval.cocoGt.getImgIds()) - set(img_ids) + if len(missing_img_ids) > 0: + print(f"WARNING: {len(missing_img_ids)} images were not predicted!") + coco_eval.cocoDt = COCO() + coco_eval.params.imgIds = list(missing_img_ids) + new_img_ids, new_eval_imgs = evaluate(coco_eval, use_self_evaluate) + img_ids.extend(new_img_ids) + eval_imgs = np.concatenate((eval_imgs, new_eval_imgs), axis=2) + + eval_imgs = list(eval_imgs.flatten()) + assert len(img_ids) == len(coco_eval.cocoGt.getImgIds()) + + coco_eval.evalImgs = eval_imgs + coco_eval.params.imgIds = img_ids + coco_eval._paramsEval = copy.deepcopy(coco_eval.params) + + +################################################################# +# From pycocotools, just removed the prints and fixed +# a Python3 bug about unicode not defined +################################################################# + + +# Copy of COCO prepare, but doesn't convert anntoRLE +def segmentation_prepare(self): + """ + Prepare ._gts and ._dts for evaluation based on params + :return: None + """ + p = self.params + if p.useCats: + gts = self.cocoGt.loadAnns( + self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds) + ) + dts = self.cocoDt.loadAnns( + self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds) + ) + else: + gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds)) + dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds)) + + for gt in gts: + gt["ignore"] = gt["ignore"] if "ignore" in gt else 0 + gt["ignore"] = "iscrowd" in gt and gt["iscrowd"] + if p.iouType == "keypoints": + gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"] + self._gts = defaultdict(list) # gt for evaluation + self._dts = defaultdict(list) # dt for evaluation + for gt in gts: + self._gts[gt["image_id"], gt["category_id"]].append(gt) + for dt in dts: + self._dts[dt["image_id"], dt["category_id"]].append(dt) + self.evalImgs = defaultdict(list) # per-image per-category evaluation results + self.eval = {} # accumulated evaluation results + + +def evaluate(self, use_self_evaluate): + """ + Run per image evaluation on given images and store results (a list of dict) in self.evalImgs + :return: None + """ + # tic = time.time() + # print('Running per image evaluation...', use_self_evaluate) + p = self.params + # add backward compatibility if useSegm is specified in params + if p.useSegm is not None: + p.iouType = "segm" if p.useSegm == 1 else "bbox" + print( + "useSegm (deprecated) is not None. Running {} evaluation".format(p.iouType) + ) + # print('Evaluate annotation type *{}*'.format(p.iouType)) + p.imgIds = list(np.unique(p.imgIds)) + if p.useCats: + p.catIds = list(np.unique(p.catIds)) + p.maxDets = sorted(p.maxDets) + self.params = p + + self._prepare() + # loop through images, area range, max detection number + catIds = p.catIds if p.useCats else [-1] + + if p.iouType == "segm" or p.iouType == "bbox": + computeIoU = self.computeIoU + elif p.iouType == "keypoints": + computeIoU = self.computeOks + self.ious = { + (imgId, catId): computeIoU(imgId, catId) + for imgId in p.imgIds + for catId in catIds + } + + maxDet = p.maxDets[-1] + if use_self_evaluate: + evalImgs = [ + self.evaluateImg(imgId, catId, areaRng, maxDet) + for catId in catIds + for areaRng in p.areaRng + for imgId in p.imgIds + ] + # this is NOT in the pycocotools code, but could be done outside + evalImgs = np.asarray(evalImgs).reshape( + len(catIds), len(p.areaRng), len(p.imgIds) + ) + return p.imgIds, evalImgs + + # <<<< Beginning of code differences with original COCO API + # def convert_instances_to_cpp(instances, is_det=False): + # # Convert annotations for a list of instances in an image to a format that's fast + # # to access in C++ + # instances_cpp = [] + # for instance in instances: + # instance_cpp = _CPP.InstanceAnnotation( + # int(instance["id"]), + # instance["score"] if is_det else instance.get("score", 0.0), + # instance["area"], + # bool(instance.get("iscrowd", 0)), + # bool(instance.get("ignore", 0)), + # ) + # instances_cpp.append(instance_cpp) + # return instances_cpp + + # # Convert GT annotations, detections, and IOUs to a format that's fast to access in C++ + # ground_truth_instances = [ + # [convert_instances_to_cpp(self._gts[imgId, catId]) for catId in p.catIds] + # for imgId in p.imgIds + # ] + # detected_instances = [ + # [ + # convert_instances_to_cpp(self._dts[imgId, catId], is_det=True) + # for catId in p.catIds + # ] + # for imgId in p.imgIds + # ] + # ious = [[self.ious[imgId, catId] for catId in catIds] for imgId in p.imgIds] + + # if not p.useCats: + # # For each image, flatten per-category lists into a single list + # ground_truth_instances = [ + # [[o for c in i for o in c]] for i in ground_truth_instances + # ] + # detected_instances = [[[o for c in i for o in c]] for i in detected_instances] + + # # Call C++ implementation of self.evaluateImgs() + # _evalImgs_cpp = _CPP.COCOevalEvaluateImages( + # p.areaRng, maxDet, p.iouThrs, ious, ground_truth_instances, detected_instances + # ) + + # self._paramsEval = copy.deepcopy(self.params) + # evalImgs = np.asarray(_evalImgs_cpp).reshape( + # len(catIds), len(p.areaRng), len(p.imgIds) + # ) + # return p.imgIds, evalImgs + + +################################################################# +# end of straight copy from pycocotools, just removing the prints +################################################################# + + +################################################################# +# From pycocotools, but disabled mask->box conversion which is +# pointless +################################################################# +def loadRes(self, resFile): + """ + Load result file and return a result api object. + :param resFile (str) : file name of result file + :return: res (obj) : result api object + """ + res = COCO() + res.dataset["images"] = [img for img in self.dataset["images"]] + + if type(resFile) == str: + anns = json.load(open(resFile)) + elif type(resFile) == np.ndarray: + anns = self.loadNumpyAnnotations(resFile) + else: + anns = resFile + assert type(anns) == list, "results in not an array of objects" + annsImgIds = [ann["image_id"] for ann in anns] + assert set(annsImgIds) == ( + set(annsImgIds) & set(self.getImgIds()) + ), "Results do not correspond to current coco set" + if "caption" in anns[0]: + imgIds = set([img["id"] for img in res.dataset["images"]]) & set( + [ann["image_id"] for ann in anns] + ) + res.dataset["images"] = [ + img for img in res.dataset["images"] if img["id"] in imgIds + ] + for id, ann in enumerate(anns): + ann["id"] = id + 1 + elif "bbox" in anns[0] and not anns[0]["bbox"] == []: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + bb = ann["bbox"] + x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]] + if "segmentation" not in ann: + ann["segmentation"] = [[x1, y1, x1, y2, x2, y2, x2, y1]] + ann["area"] = bb[2] * bb[3] + ann["id"] = id + 1 + ann["iscrowd"] = 0 + elif "segmentation" in anns[0]: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + # now only support compressed RLE format as segmentation results + # ann["area"] = mask_util.area(ann["segmentation"]) + # The following lines are disabled because they are pointless + # if not 'bbox' in ann: + # ann['bbox'] = maskUtils.toBbox(ann['segmentation']) + ann["id"] = id + 1 + ann["iscrowd"] = 0 + elif "keypoints" in anns[0]: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + s = ann["keypoints"] + x = s[0::3] + y = s[1::3] + x0, x1, y0, y1 = np.min(x), np.max(x), np.min(y), np.max(y) + ann["area"] = (x1 - x0) * (y1 - y0) + ann["id"] = id + 1 + ann["bbox"] = [x0, y0, x1 - x0, y1 - y0] + + res.dataset["annotations"] = anns + res.createIndex() + return res + + +################################################################# +# end of straight copy from pycocotools +################################################################# + + +################################################################# +# From pycocotools, but added handling of custom area rngs, and returns stat keys +################################################################# +def summarize(self): + """ + Compute and display summary metrics for evaluation results. + Note this functin can *only* be applied on the default parameter setting + """ + + def _summarize(ap=1, iouThr=None, areaRng="all", maxDets=100): + p = self.params + iStr = " {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}" + titleStr = "Average Precision" if ap == 1 else "Average Recall" + typeStr = "(AP)" if ap == 1 else "(AR)" + iouStr = ( + "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1]) + if iouThr is None + else "{:0.2f}".format(iouThr) + ) + + aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng] + mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets] + if ap == 1: + # dimension of precision: [TxRxKxAxM] + s = self.eval["precision"] + # IoU + if iouThr is not None: + t = np.where(iouThr == p.iouThrs)[0] + s = s[t] + s = s[:, :, :, aind, mind] + else: + # dimension of recall: [TxKxAxM] + s = self.eval["recall"] + if iouThr is not None: + t = np.where(iouThr == p.iouThrs)[0] + s = s[t] + s = s[:, :, aind, mind] + if len(s[s > -1]) == 0: + mean_s = -1 + else: + mean_s = np.mean(s[s > -1]) + print(iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s)) + return mean_s + + def _summarizeDets(): + nb_results = 6 + (len(self.params.areaRng) - 1) * 2 + assert len(self.params.areaRng) == len(self.params.areaRngLbl) + stats = np.zeros((nb_results,)) + keys = ["AP", "AP_50", "AP_75"] + stats[0] = _summarize(1, maxDets=self.params.maxDets[2]) + stats[1] = _summarize(1, iouThr=0.5, maxDets=self.params.maxDets[2]) + stats[2] = _summarize(1, iouThr=0.75, maxDets=self.params.maxDets[2]) + cur_id = 3 + for area in self.params.areaRngLbl[1:]: + stats[cur_id] = _summarize(1, areaRng=area, maxDets=self.params.maxDets[2]) + cur_id += 1 + keys.append(f"AP_{area}") + stats[cur_id] = _summarize(0, maxDets=self.params.maxDets[0]) + cur_id += 1 + stats[cur_id] = _summarize(0, maxDets=self.params.maxDets[1]) + cur_id += 1 + stats[cur_id] = _summarize(0, maxDets=self.params.maxDets[2]) + cur_id += 1 + keys += ["AR", "AR_50", "AR_75"] + + for area in self.params.areaRngLbl[1:]: + stats[cur_id] = _summarize(0, areaRng=area, maxDets=self.params.maxDets[2]) + cur_id += 1 + keys.append(f"AR_{area}") + assert len(stats) == len(keys) + return keys, stats + + if not self.eval: + raise Exception("Please run accumulate() first") + self.stats = _summarizeDets() + + +################################################################# +# end of straight copy from pycocotools +################################################################# + + +################################################################# +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/evaluation/fast_eval_api.py +# with slight adjustments +################################################################# +def accumulate(self, use_self_eval=False): + """ + Accumulate per image evaluation results and store the result in self.eval. Does not + support changing parameter settings from those used by self.evaluate() + """ + if use_self_eval: + self.accumulate() + return + # CPP code is disabled + # self.eval = _CPP.COCOevalAccumulate(self.params, self.evalImgs) + + # # recall is num_iou_thresholds X num_categories X num_area_ranges X num_max_detections + # self.eval["recall"] = np.array(self.eval["recall"]).reshape( + # self.eval["counts"][:1] + self.eval["counts"][2:] + # ) + + # # precision and scores are num_iou_thresholds X num_recall_thresholds X num_categories X + # # num_area_ranges X num_max_detections + # self.eval["precision"] = np.array(self.eval["precision"]).reshape( + # self.eval["counts"] + # ) + # self.eval["scores"] = np.array(self.eval["scores"]).reshape(self.eval["counts"]) diff --git a/src/nn/segearth_ov3/sam3/eval/coco_eval_offline.py b/src/nn/segearth_ov3/sam3/eval/coco_eval_offline.py new file mode 100644 index 0000000..7b07228 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/coco_eval_offline.py @@ -0,0 +1,181 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +This evaluator is meant for regular COCO mAP evaluation, for example on the COCO val set. + +For Category mAP, we need the model to make predictions for all the categories on every single image. +In general, since the number of classes can be big, and the API model makes predictions individually for each pair (image, class), +we may need to split the inference process for a given image in several chunks. +""" + +import logging +from collections import defaultdict + +import torch +from pycocotools.coco import COCO +from pycocotools.cocoeval import COCOeval +from sam3.train.utils.distributed import is_main_process + +try: + from tidecv import datasets, TIDE + + HAS_TIDE = True +except ImportError: + HAS_TIDE = False + print("WARNING: TIDE not installed. Detailed analysis will not be available.") + + +# the COCO detection metrics (https://github.com/cocodataset/cocoapi/blob/8c9bcc3cf640524c4c20a9c40e89cb6a2f2fa0e9/PythonAPI/pycocotools/cocoeval.py#L460-L471) +COCO_METRICS = [ + "AP", + "AP_50", + "AP_75", + "AP_small", + "AP_medium", + "AP_large", + "AR_maxDets@1", + "AR_maxDets@10", + "AR_maxDets@100", + "AR_small", + "AR_medium", + "AR_large", +] + + +def convert_to_xywh(boxes): + """Convert bounding boxes from xyxy format to xywh format.""" + xmin, ymin, xmax, ymax = boxes.unbind(-1) + return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=-1) + + +class HeapElement: + """Utility class to make a heap with a custom comparator""" + + def __init__(self, val): + self.val = val + + def __lt__(self, other): + return self.val["score"] < other.val["score"] + + +class COCOevalCustom(COCOeval): + """ + This is a slightly modified version of the original COCO API with added support for positive split evaluation. + """ + + def __init__( + self, cocoGt=None, cocoDt=None, iouType="segm", dt_only_positive=False + ): + super().__init__(cocoGt, cocoDt, iouType) + self.dt_only_positive = dt_only_positive + + def _prepare(self): + """ + Prepare ._gts and ._dts for evaluation based on params + :return: None + """ + + def _toMask(anns, coco): + # modify ann['segmentation'] by reference + for ann in anns: + rle = coco.annToRLE(ann) + ann["segmentation"] = rle + + p = self.params + if p.useCats: + gts = self.cocoGt.loadAnns( + self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds) + ) + dts = self.cocoDt.loadAnns( + self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds) + ) + else: + gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds)) + dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds)) + + # convert ground truth to mask if iouType == 'segm' + if p.iouType == "segm": + _toMask(gts, self.cocoGt) + _toMask(dts, self.cocoDt) + # set ignore flag + for gt in gts: + gt["ignore"] = gt["ignore"] if "ignore" in gt else 0 + gt["ignore"] = "iscrowd" in gt and gt["iscrowd"] + if p.iouType == "keypoints": + gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"] + self._gts = defaultdict(list) # gt for evaluation + self._dts = defaultdict(list) # dt for evaluation + + _gts_cat_ids = defaultdict(set) # gt for evaluation on positive split + for gt in gts: + self._gts[gt["image_id"], gt["category_id"]].append(gt) + _gts_cat_ids[gt["image_id"]].add(gt["category_id"]) + + #### BEGIN MODIFICATION #### + for dt in dts: + if ( + self.dt_only_positive + and dt["category_id"] not in _gts_cat_ids[dt["image_id"]] + ): + continue + self._dts[dt["image_id"], dt["category_id"]].append(dt) + #### END MODIFICATION #### + self.evalImgs = defaultdict(list) # per-image per-category evaluation results + self.eval = {} # accumulated evaluation results + + +class CocoEvaluatorOfflineWithPredFileEvaluators: + def __init__( + self, + gt_path, + tide: bool = True, + iou_type: str = "bbox", + positive_split=False, + ): + self.gt_path = gt_path + self.tide_enabled = HAS_TIDE and tide + self.positive_split = positive_split + self.iou_type = iou_type + + def evaluate(self, dumped_file): + if not is_main_process(): + return {} + + logging.info("OfflineCoco evaluator: Loading groundtruth") + self.gt = COCO(self.gt_path) + + # Creating the result file + logging.info("Coco evaluator: Creating the result file") + cocoDt = self.gt.loadRes(str(dumped_file)) + + # Run the evaluation + logging.info("Coco evaluator: Running evaluation") + coco_eval = COCOevalCustom( + self.gt, cocoDt, iouType=self.iou_type, dt_only_positive=self.positive_split + ) + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() + + outs = {} + for i, value in enumerate(coco_eval.stats): + outs[f"coco_eval_{self.iou_type}_{COCO_METRICS[i]}"] = value + + if self.tide_enabled: + logging.info("Coco evaluator: Loading TIDE") + self.tide_gt = datasets.COCO(self.gt_path) + self.tide = TIDE(mode="mask" if self.iou_type == "segm" else "bbox") + + # Run TIDE + logging.info("Coco evaluator: Running TIDE") + self.tide.evaluate( + self.tide_gt, datasets.COCOResult(str(dumped_file)), name="coco_eval" + ) + self.tide.summarize() + for k, v in self.tide.get_main_errors()["coco_eval"].items(): + outs[f"coco_eval_{self.iou_type}_TIDE_{k}"] = v + + for k, v in self.tide.get_special_errors()["coco_eval"].items(): + outs[f"coco_eval_{self.iou_type}_TIDE_{k}"] = v + + return outs diff --git a/src/nn/segearth_ov3/sam3/eval/coco_reindex.py b/src/nn/segearth_ov3/sam3/eval/coco_reindex.py new file mode 100644 index 0000000..49cd944 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/coco_reindex.py @@ -0,0 +1,230 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Self-contained COCO JSON re-indexing function that creates temporary files. +""" + +import json +import os +import tempfile +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + + +def reindex_coco_to_temp(input_json_path: str) -> Optional[str]: + """ + Convert 0-indexed COCO JSON file to 1-indexed and save to temporary location. + + Args: + input_json_path: Path to the input COCO JSON file + + Returns: + Path to the new 1-indexed JSON file in temporary directory, or None if no conversion needed + + Raises: + FileNotFoundError: If input file doesn't exist + json.JSONDecodeError: If input file is not valid JSON + ValueError: If input file is not a valid COCO format + """ + + def is_coco_json(data: Dict[str, Any]) -> bool: + """Check if data appears to be a COCO format file.""" + if not isinstance(data, dict): + return False + # A COCO file should have at least one of these keys + coco_keys = {"images", "annotations", "categories"} + return any(key in data for key in coco_keys) + + def check_zero_indexed(data: Dict[str, Any]) -> Tuple[bool, bool, bool]: + """ + Check if annotations, images, or categories start from index 0. + + Returns: + Tuple of (annotations_zero_indexed, images_zero_indexed, categories_zero_indexed) + """ + annotations_zero = False + images_zero = False + categories_zero = False + + # Check annotations + annotations = data.get("annotations", []) + if annotations and any(ann.get("id", -1) == 0 for ann in annotations): + annotations_zero = True + + # Check images + images = data.get("images", []) + if images and any(img.get("id", -1) == 0 for img in images): + images_zero = True + + # Check categories + categories = data.get("categories", []) + if categories and any(cat.get("id", -1) == 0 for cat in categories): + categories_zero = True + + return annotations_zero, images_zero, categories_zero + + def reindex_coco_data(data: Dict[str, Any]) -> Dict[str, Any]: + """Convert 0-indexed COCO data to 1-indexed.""" + modified_data = data.copy() + + annotations_zero, images_zero, categories_zero = check_zero_indexed(data) + + # Create ID mapping for consistency + image_id_mapping = {} + category_id_mapping = {} + + # Process images first (since annotations reference image IDs) + if images_zero and "images" in modified_data: + for img in modified_data["images"]: + old_id = img["id"] + new_id = old_id + 1 + image_id_mapping[old_id] = new_id + img["id"] = new_id + + # Process categories (since annotations reference category IDs) + if categories_zero and "categories" in modified_data: + for cat in modified_data["categories"]: + old_id = cat["id"] + new_id = old_id + 1 + category_id_mapping[old_id] = new_id + cat["id"] = new_id + + # Process annotations + if "annotations" in modified_data: + for ann in modified_data["annotations"]: + # Update annotation ID if needed + if annotations_zero: + ann["id"] = ann["id"] + 1 + + # Update image_id reference if images were reindexed + if images_zero and ann.get("image_id") is not None: + old_image_id = ann["image_id"] + if old_image_id in image_id_mapping: + ann["image_id"] = image_id_mapping[old_image_id] + + # Update category_id reference if categories were reindexed + if categories_zero and ann.get("category_id") is not None: + old_category_id = ann["category_id"] + if old_category_id in category_id_mapping: + ann["category_id"] = category_id_mapping[old_category_id] + + return modified_data + + # Validate input path + if not os.path.exists(input_json_path): + raise FileNotFoundError(f"Input file not found: {input_json_path}") + + # Load and validate JSON data + try: + with open(input_json_path, "r", encoding="utf-8") as f: + data = json.load(f) + except json.JSONDecodeError as e: + raise json.JSONDecodeError(f"Invalid JSON in {input_json_path}: {e}") + + # Validate COCO format + if not is_coco_json(data): + raise ValueError( + f"File does not appear to be in COCO format: {input_json_path}" + ) + + # Check if reindexing is needed + annotations_zero, images_zero, categories_zero = check_zero_indexed(data) + + if not (annotations_zero or images_zero or categories_zero): + # No conversion needed - just copy to temp location + input_path = Path(input_json_path) + temp_dir = tempfile.mkdtemp() + temp_filename = f"{input_path.stem}_1_indexed{input_path.suffix}" + temp_path = os.path.join(temp_dir, temp_filename) + + with open(temp_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + + return temp_path + + # Perform reindexing + modified_data = reindex_coco_data(data) + + # Create temporary file + input_path = Path(input_json_path) + temp_dir = tempfile.mkdtemp() + temp_filename = f"{input_path.stem}_1_indexed{input_path.suffix}" + temp_path = os.path.join(temp_dir, temp_filename) + + # Write modified data to temporary file + with open(temp_path, "w", encoding="utf-8") as f: + json.dump(modified_data, f, indent=2, ensure_ascii=False) + + return temp_path + + +# Example usage and test function +def test_reindex_function(): + """Test the reindex function with a sample COCO file.""" + + # Create a test COCO file + test_data = { + "info": {"description": "Test COCO dataset", "version": "1.0", "year": 2023}, + "images": [ + {"id": 0, "width": 640, "height": 480, "file_name": "test1.jpg"}, + {"id": 1, "width": 640, "height": 480, "file_name": "test2.jpg"}, + ], + "categories": [ + {"id": 0, "name": "person", "supercategory": "person"}, + {"id": 1, "name": "car", "supercategory": "vehicle"}, + ], + "annotations": [ + { + "id": 0, + "image_id": 0, + "category_id": 0, + "bbox": [100, 100, 50, 75], + "area": 3750, + "iscrowd": 0, + }, + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [200, 150, 120, 80], + "area": 9600, + "iscrowd": 0, + }, + ], + } + + # Create temporary test file + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(test_data, f, indent=2) + test_file_path = f.name + + try: + # Test the function + result_path = reindex_coco_to_temp(test_file_path) + print(f"Original file: {test_file_path}") + print(f"Converted file: {result_path}") + + # Load and display the result + with open(result_path, "r") as f: + result_data = json.load(f) + + print("\nConverted data sample:") + print(f"First image ID: {result_data['images'][0]['id']}") + print(f"First category ID: {result_data['categories'][0]['id']}") + print(f"First annotation ID: {result_data['annotations'][0]['id']}") + print(f"First annotation image_id: {result_data['annotations'][0]['image_id']}") + print( + f"First annotation category_id: {result_data['annotations'][0]['category_id']}" + ) + + # Clean up + os.unlink(result_path) + os.rmdir(os.path.dirname(result_path)) + + finally: + # Clean up test file + os.unlink(test_file_path) + + +if __name__ == "__main__": + test_reindex_function() diff --git a/src/nn/segearth_ov3/sam3/eval/coco_writer.py b/src/nn/segearth_ov3/sam3/eval/coco_writer.py new file mode 100644 index 0000000..5134e63 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/coco_writer.py @@ -0,0 +1,352 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +COCO prediction dumper for distributed training. + +Handles collection and dumping of COCO-format predictions from models. +Supports distributed processing with multiple GPUs/processes. +""" + +import copy +import gc +import heapq +import json +import logging +import os +from collections import defaultdict +from pathlib import Path +from typing import Any, Optional + +import pycocotools.mask as mask_utils +import torch +from iopath.common.file_io import g_pathmgr +from sam3.eval.coco_eval_offline import convert_to_xywh +from sam3.train.masks_ops import rle_encode +from sam3.train.utils.distributed import ( + all_gather, + gather_to_rank_0_via_filesys, + get_rank, + is_main_process, +) + + +### Helper functions and classes + + +class HeapElement: + """Utility class to make a heap with a custom comparator based on score.""" + + def __init__(self, val): + self.val = val + + def __lt__(self, other): + return self.val["score"] < other.val["score"] + + +class PredictionDumper: + """ + Handles collection and dumping of COCO-format predictions from a model. + + This class processes model outputs through a postprocessor, converts them to COCO format, + and saves them to disk. It supports distributed processing with multiple GPUs/processes. + """ + + def __init__( + self, + dump_dir: str, + postprocessor, + maxdets: int, + iou_type: str, + gather_pred_via_filesys: bool = False, + merge_predictions: bool = False, + pred_file_evaluators: Optional[Any] = None, + ): + """ + Initialize the PredictionDumper. + + Args: + dump_dir: Directory to dump predictions. + postprocessor: Module to convert the model's output into COCO format. + maxdets: Maximum number of detections per image. + iou_type: IoU type to evaluate. Can include "bbox", "segm" + gather_pred_via_filesys: If True, use the filesystem for collective gathers across + processes (requires a shared filesystem). Otherwise, use torch collective ops. + merge_predictions: If True, merge predictions from all processes and dump to a single file. + """ + self.iou_type = iou_type + self.maxdets = maxdets + self.dump_dir = dump_dir + self.postprocessor = postprocessor + self.gather_pred_via_filesys = gather_pred_via_filesys + self.merge_predictions = merge_predictions + self.pred_file_evaluators = pred_file_evaluators + if self.pred_file_evaluators is not None: + assert ( + merge_predictions + ), "merge_predictions must be True if pred_file_evaluators are provided" + assert self.dump_dir is not None, "dump_dir must be provided" + + if is_main_process(): + os.makedirs(self.dump_dir, exist_ok=True) + logging.info(f"Created prediction dump directory: {self.dump_dir}") + + # Initialize state + self.reset() + + def update(self, *args, **kwargs): + """ + Process and accumulate predictions from model outputs. + + Args: + *args, **kwargs: Arguments passed to postprocessor.process_results() + """ + predictions = self.postprocessor.process_results(*args, **kwargs) + results = self.prepare(predictions, self.iou_type) + self._dump(results) + + def _dump(self, results): + """ + Add results to the dump list with precision rounding. + + Args: + results: List of prediction dictionaries in COCO format. + """ + dumped_results = copy.deepcopy(results) + for r in dumped_results: + if "bbox" in r: + r["bbox"] = [round(coord, 5) for coord in r["bbox"]] + r["score"] = round(r["score"], 5) + self.dump.extend(dumped_results) + + def synchronize_between_processes(self): + """ + Synchronize predictions across all processes and save to disk. + + If gather_pred_via_filesys is True, uses filesystem for gathering. + Otherwise, uses torch distributed collective operations. + Saves per-rank predictions to separate JSON files. + """ + logging.info("Prediction Dumper: Synchronizing between processes") + + if not self.merge_predictions: + dumped_file = ( + Path(self.dump_dir) + / f"coco_predictions_{self.iou_type}_{get_rank()}.json" + ) + logging.info( + f"Prediction Dumper: Dumping local predictions to {dumped_file}" + ) + with g_pathmgr.open(str(dumped_file), "w") as f: + json.dump(self.dump, f) + else: + self.dump = self.gather_and_merge_predictions() + dumped_file = Path(self.dump_dir) / f"coco_predictions_{self.iou_type}.json" + if is_main_process(): + logging.info( + f"Prediction Dumper: Dumping merged predictions to {dumped_file}" + ) + with g_pathmgr.open(str(dumped_file), "w") as f: + json.dump(self.dump, f) + + self.reset() + return dumped_file + + def gather_and_merge_predictions(self): + """ + Gather predictions from all processes and merge them, keeping top predictions per image. + + This method collects predictions from all processes, then keeps only the top maxdets + predictions per image based on score. It also deduplicates predictions by (image_id, category_id). + + Returns: + List of merged prediction dictionaries. + """ + logging.info("Prediction Dumper: Gathering predictions from all processes") + gc.collect() + + if self.gather_pred_via_filesys: + dump = gather_to_rank_0_via_filesys(self.dump) + else: + dump = all_gather(self.dump, force_cpu=True) + + # Combine predictions, keeping only top maxdets per image + preds_by_image = defaultdict(list) + seen_img_cat = set() + + for cur_dump in dump: + cur_seen_img_cat = set() + for p in cur_dump: + image_id = p["image_id"] + cat_id = p["category_id"] + + # Skip if we've already seen this image/category pair in a previous dump + if (image_id, cat_id) in seen_img_cat: + continue + + cur_seen_img_cat.add((image_id, cat_id)) + + # Use a min-heap to keep top predictions + if len(preds_by_image[image_id]) < self.maxdets: + heapq.heappush(preds_by_image[image_id], HeapElement(p)) + else: + heapq.heappushpop(preds_by_image[image_id], HeapElement(p)) + + seen_img_cat.update(cur_seen_img_cat) + + # Flatten the heap elements back to a list + merged_dump = sum( + [[h.val for h in cur_preds] for cur_preds in preds_by_image.values()], [] + ) + + return merged_dump + + def compute_synced(self): + """ + Synchronize predictions across processes and compute summary. + + Returns: + Summary dictionary from summarize(). + """ + dumped_file = self.synchronize_between_processes() + if not is_main_process(): + return {"": 0.0} + + meters = {} + if self.pred_file_evaluators is not None: + for evaluator in self.pred_file_evaluators: + results = evaluator.evaluate(dumped_file) + meters.update(results) + + if len(meters) == 0: + meters = {"": 0.0} + return meters + + def compute(self): + """ + Compute without synchronization. + + Returns: + Empty metric dictionary. + """ + return {"": 0.0} + + def reset(self): + """Reset internal state for a new evaluation round.""" + self.dump = [] + + def prepare(self, predictions, iou_type): + """ + Route predictions to the appropriate preparation method based on iou_type. + + Args: + predictions: Dictionary mapping image IDs to prediction dictionaries. + iou_type: Type of evaluation ("bbox", "segm"). + + Returns: + List of COCO-format prediction dictionaries. + """ + if iou_type == "bbox": + return self.prepare_for_coco_detection(predictions) + elif iou_type == "segm": + return self.prepare_for_coco_segmentation(predictions) + else: + raise ValueError(f"Unknown iou type: {iou_type}") + + def prepare_for_coco_detection(self, predictions): + """ + Convert predictions to COCO detection format. + + Args: + predictions: Dictionary mapping image IDs to prediction dictionaries + containing "boxes", "scores", and "labels". + + Returns: + List of COCO-format detection dictionaries. + """ + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + + coco_results.extend( + [ + { + "image_id": original_id, + "category_id": labels[k], + "bbox": box, + "score": scores[k], + } + for k, box in enumerate(boxes) + ] + ) + return coco_results + + @torch.no_grad() + def prepare_for_coco_segmentation(self, predictions): + """ + Convert predictions to COCO segmentation format. + + Args: + predictions: Dictionary mapping image IDs to prediction dictionaries + containing "masks" or "masks_rle", "scores", and "labels". + Optionally includes "boundaries" and "dilated_boundaries". + + Returns: + List of COCO-format segmentation dictionaries with RLE-encoded masks. + """ + coco_results = [] + for original_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + + boxes = None + if "boxes" in prediction: + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + assert len(boxes) == len(scores) + + if "masks_rle" in prediction: + rles = prediction["masks_rle"] + areas = [] + for rle in rles: + cur_area = mask_utils.area(rle) + h, w = rle["size"] + areas.append(cur_area / (h * w)) + else: + masks = prediction["masks"] + masks = masks > 0.5 + h, w = masks.shape[-2:] + + areas = masks.flatten(1).sum(1) / (h * w) + areas = areas.tolist() + + rles = rle_encode(masks.squeeze(1)) + + # Memory cleanup + del masks + del prediction["masks"] + + assert len(areas) == len(rles) == len(scores) + + for k, rle in enumerate(rles): + payload = { + "image_id": original_id, + "category_id": labels[k], + "segmentation": rle, + "score": scores[k], + "area": areas[k], + } + if boxes is not None: + payload["bbox"] = boxes[k] + + coco_results.append(payload) + + return coco_results diff --git a/src/nn/segearth_ov3/sam3/eval/conversion_util.py b/src/nn/segearth_ov3/sam3/eval/conversion_util.py new file mode 100644 index 0000000..68942b6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/conversion_util.py @@ -0,0 +1,211 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import json +import os +from collections import defaultdict + +from tqdm import tqdm + + +def convert_ytbvis_to_cocovid_gt(ann_json, save_path=None): + """Convert YouTube VIS dataset to COCO-style video instance segmentation format. + + Args: + ann_json (str): Path to YouTube VIS annotation JSON file + save_path (str): path to save converted COCO-style JSON + """ + # Initialize COCO structure + VIS = { + "info": {}, + "images": [], + "videos": [], + "tracks": [], + "annotations": [], + "categories": [], + "licenses": [], + } + + # Load original annotations + official_anns = json.load(open(ann_json)) + VIS["categories"] = official_anns["categories"] # Direct copy categories + + # Initialize counters + records = dict(img_id=1, ann_id=1) + + # Create video-to-annotations mapping + vid_to_anns = defaultdict(list) + for ann in official_anns["annotations"]: + vid_to_anns[ann["video_id"]].append(ann) + + # Create tracks directly + VIS["tracks"] = [ + { + "id": ann["id"], + "category_id": ann["category_id"], + "video_id": ann["video_id"], + } + for ann in official_anns["annotations"] + ] + + # Process videos + for video_info in tqdm(official_anns["videos"]): + # Create video entry + video = { + "id": video_info["id"], + "name": os.path.dirname(video_info["file_names"][0]), + "width": video_info["width"], + "height": video_info["height"], + "length": video_info["length"], + "neg_category_ids": [], + "not_exhaustive_category_ids": [], + } + VIS["videos"].append(video) + + # Process frames + num_frames = len(video_info["file_names"]) + for frame_idx in range(num_frames): + # Create image entry + image = { + "id": records["img_id"], + "video_id": video_info["id"], + "file_name": video_info["file_names"][frame_idx], + "width": video_info["width"], + "height": video_info["height"], + "frame_index": frame_idx, + "frame_id": frame_idx, + } + VIS["images"].append(image) + + # Process annotations for this frame + if video_info["id"] in vid_to_anns: + for ann in vid_to_anns[video_info["id"]]: + bbox = ann["bboxes"][frame_idx] + if bbox is None: + continue + + # Create annotation entry + annotation = { + "id": records["ann_id"], + "video_id": video_info["id"], + "image_id": records["img_id"], + "track_id": ann["id"], + "category_id": ann["category_id"], + "bbox": bbox, + "area": ann["areas"][frame_idx], + "segmentation": ann["segmentations"][frame_idx], + "iscrowd": ann["iscrowd"], + } + VIS["annotations"].append(annotation) + records["ann_id"] += 1 + + records["img_id"] += 1 + + # Print summary + print(f"Converted {len(VIS['videos'])} videos") + print(f"Converted {len(VIS['images'])} images") + print(f"Created {len(VIS['tracks'])} tracks") + print(f"Created {len(VIS['annotations'])} annotations") + + if save_path is None: + return VIS + + # Save output + save_dir = os.path.dirname(save_path) + os.makedirs(save_dir, exist_ok=True) + json.dump(VIS, open(save_path, "w")) + + return VIS + + +def convert_ytbvis_to_cocovid_pred( + youtubevis_pred_path: str, converted_dataset_path: str, output_path: str +) -> None: + """ + Convert YouTubeVIS predictions to COCO format with video_id preservation + + Args: + youtubevis_pred_path: Path to YouTubeVIS prediction JSON + converted_dataset_path: Path to converted COCO dataset JSON + output_path: Path to save COCO format predictions + """ + + # Load YouTubeVIS predictions + with open(youtubevis_pred_path) as f: + ytv_predictions = json.load(f) + + # Load converted dataset for image ID mapping + with open(converted_dataset_path) as f: + coco_dataset = json.load(f) + + # Create (video_id, frame_idx) -> image_id mapping + image_id_map = { + (img["video_id"], img["frame_index"]): img["id"] + for img in coco_dataset["images"] + } + + coco_annotations = [] + track_id_counter = 1 # Unique track ID generator + + for pred in tqdm(ytv_predictions): + video_id = pred["video_id"] + category_id = pred["category_id"] + bboxes = pred["bboxes"] + segmentations = pred.get("segmentations", []) # Get segmentations if available + areas = pred.get("areas", []) # Get areas if available + score = pred["score"] + + # Assign unique track ID for this prediction + track_id = track_id_counter + track_id_counter += 1 + + # Ensure segmentations and areas have the same length as bboxes + if len(segmentations) == 0: + segmentations = [None] * len(bboxes) + if len(areas) == 0: + areas = [None] * len(bboxes) + + for frame_idx, (bbox, segmentation, area_from_pred) in enumerate( + zip(bboxes, segmentations, areas) + ): + # Skip frames with missing objects (None or zero bbox) + if bbox is None or all(x == 0 for x in bbox): + continue + + # Get corresponding image ID from mapping + image_id = image_id_map.get((video_id, frame_idx)) + if image_id is None: + raise RuntimeError( + f"prediction {video_id=}, {frame_idx=} does not match any images in the converted COCO format" + ) + + # Extract bbox coordinates + x, y, w, h = bbox + + # Calculate area - use area from prediction if available, otherwise from bbox + if area_from_pred is not None and area_from_pred > 0: + area = area_from_pred + else: + area = w * h + + # Create COCO annotation with video_id + coco_annotation = { + "image_id": int(image_id), + "video_id": video_id, # Added video_id field + "track_id": track_id, + "category_id": category_id, + "bbox": [float(x), float(y), float(w), float(h)], + "area": float(area), + "iscrowd": 0, + "score": float(score), + } + + # Add segmentation if available + if segmentation is not None: + coco_annotation["segmentation"] = segmentation + + coco_annotations.append(coco_annotation) + + # Save output + with open(output_path, "w") as f: + json.dump(coco_annotations, f) + + print(f"Converted {len(coco_annotations)} predictions to COCO format with video_id") diff --git a/src/nn/segearth_ov3/sam3/eval/demo_eval.py b/src/nn/segearth_ov3/sam3/eval/demo_eval.py new file mode 100644 index 0000000..70804cc --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/demo_eval.py @@ -0,0 +1,658 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +This evaluator is based upon COCO evaluation, but evaluates the model in a "demo" setting. +This means that the model's predictions are thresholded and evaluated as "hard" predictions. +""" + +import logging +from typing import Optional + +import numpy as np +import pycocotools.mask as maskUtils +from pycocotools.cocoeval import COCOeval + +from sam3.eval.coco_eval import CocoEvaluator +from sam3.train.masks_ops import compute_F_measure +from sam3.train.utils.distributed import is_main_process + +from scipy.optimize import linear_sum_assignment + + +class DemoEval(COCOeval): + """ + This evaluator is based upon COCO evaluation, but evaluates the model in a "demo" setting. + This means that the model's predictions are thresholded and evaluated as "hard" predictions. + """ + + def __init__( + self, + coco_gt=None, + coco_dt=None, + iouType="bbox", + threshold=0.5, + compute_JnF=False, + ): + """ + Args: + coco_gt (COCO): ground truth COCO API + coco_dt (COCO): detections COCO API + iou_type (str): type of IoU to evaluate + threshold (float): threshold for predictions + """ + super().__init__(coco_gt, coco_dt, iouType) + self.threshold = threshold + + self.params.useCats = False + self.params.areaRng = [[0**2, 1e5**2]] + self.params.areaRngLbl = ["all"] + self.params.maxDets = [100000] + self.compute_JnF = compute_JnF + + def computeIoU(self, imgId, catId): + # Same as the original COCOeval.computeIoU, but without sorting + p = self.params + if p.useCats: + gt = self._gts[imgId, catId] + dt = self._dts[imgId, catId] + else: + gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + if len(gt) == 0 and len(dt) == 0: + return [] + + if p.iouType == "segm": + g = [g["segmentation"] for g in gt] + d = [d["segmentation"] for d in dt] + elif p.iouType == "bbox": + g = [g["bbox"] for g in gt] + d = [d["bbox"] for d in dt] + else: + raise Exception("unknown iouType for iou computation") + + # compute iou between each dt and gt region + iscrowd = [int(o["iscrowd"]) for o in gt] + ious = maskUtils.iou(d, g, iscrowd) + return ious + + def evaluateImg(self, imgId, catId, aRng, maxDet): + """ + perform evaluation for single category and image + :return: dict (single image results) + """ + p = self.params + assert not p.useCats, "This evaluator does not support per-category evaluation." + assert catId == -1 + all_gts = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + keep_gt = np.array([not g["ignore"] for g in all_gts], dtype=bool) + gt = [g for g in all_gts if not g["ignore"]] + all_dts = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + keep_dt = np.array([d["score"] >= self.threshold for d in all_dts], dtype=bool) + dt = [d for d in all_dts if d["score"] >= self.threshold] + if len(gt) == 0 and len(dt) == 0: + # This is a "true negative" case, where there are no GTs and no predictions + # The box-level metrics are ill-defined, so we don't add them to this dict + return { + "image_id": imgId, + "IL_TP": 0, + "IL_TN": 1, + "IL_FP": 0, + "IL_FN": 0, + "IL_perfect_neg": np.ones((len(p.iouThrs),), dtype=np.int64), + "num_dt": len(dt), + } + + if len(gt) > 0 and len(dt) == 0: + # This is a "false negative" case, where there are GTs but no predictions + return { + "image_id": imgId, + "IL_TP": 0, + "IL_TN": 0, + "IL_FP": 0, + "IL_FN": 1, + "TPs": np.zeros((len(p.iouThrs),), dtype=np.int64), + "FPs": np.zeros((len(p.iouThrs),), dtype=np.int64), + "FNs": np.ones((len(p.iouThrs),), dtype=np.int64) * len(gt), + "local_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64), + "local_positive_F1s": np.zeros((len(p.iouThrs),), dtype=np.int64), + "IL_perfect_pos": np.zeros((len(p.iouThrs),), dtype=np.int64), + "num_dt": len(dt), + } + + # Load pre-computed ious + ious = self.ious[(imgId, catId)] + + # compute matching + if len(ious) == 0: + ious = np.zeros((len(dt), len(gt))) + else: + ious = ious[keep_dt, :][:, keep_gt] + assert ious.shape == (len(dt), len(gt)) + + matched_dt, matched_gt = linear_sum_assignment(-ious) + + match_scores = ious[matched_dt, matched_gt] + + if self.compute_JnF and len(match_scores) > 0: + j_score = match_scores.mean() + f_measure = 0 + for dt_id, gt_id in zip(matched_dt, matched_gt): + f_measure += compute_F_measure( + gt_boundary_rle=gt[gt_id]["boundary"], + gt_dilated_boundary_rle=gt[gt_id]["dilated_boundary"], + dt_boundary_rle=dt[dt_id]["boundary"], + dt_dilated_boundary_rle=dt[dt_id]["dilated_boundary"], + ) + f_measure /= len(match_scores) + 1e-9 + JnF = (j_score + f_measure) * 0.5 + else: + j_score = f_measure = JnF = -1 + + TPs, FPs, FNs = [], [], [] + IL_perfect = [] + for thresh in p.iouThrs: + TP = (match_scores >= thresh).sum() + FP = len(dt) - TP + FN = len(gt) - TP + assert ( + FP >= 0 and FN >= 0 + ), f"FP: {FP}, FN: {FN}, TP: {TP}, match_scores: {match_scores}, len(dt): {len(dt)}, len(gt): {len(gt)}, ious: {ious}" + TPs.append(TP) + FPs.append(FP) + FNs.append(FN) + + if FP == FN and FP == 0: + IL_perfect.append(1) + else: + IL_perfect.append(0) + + TPs = np.array(TPs, dtype=np.int64) + FPs = np.array(FPs, dtype=np.int64) + FNs = np.array(FNs, dtype=np.int64) + IL_perfect = np.array(IL_perfect, dtype=np.int64) + + # compute precision recall and F1 + precision = TPs / (TPs + FPs + 1e-4) + assert np.all(precision <= 1) + recall = TPs / (TPs + FNs + 1e-4) + assert np.all(recall <= 1) + F1 = 2 * precision * recall / (precision + recall + 1e-4) + + result = { + "image_id": imgId, + "TPs": TPs, + "FPs": FPs, + "FNs": FNs, + "local_F1s": F1, + "IL_TP": (len(gt) > 0) and (len(dt) > 0), + "IL_FP": (len(gt) == 0) and (len(dt) > 0), + "IL_TN": (len(gt) == 0) and (len(dt) == 0), + "IL_FN": (len(gt) > 0) and (len(dt) == 0), + ("IL_perfect_pos" if len(gt) > 0 else "IL_perfect_neg"): IL_perfect, + "F": f_measure, + "J": j_score, + "J&F": JnF, + "num_dt": len(dt), + } + if len(gt) > 0 and len(dt) > 0: + result["local_positive_F1s"] = F1 + return result + + def accumulate(self, p=None): + """ + Accumulate per image evaluation results and store the result in self.eval + :param p: input params for evaluation + :return: None + """ + if not self.evalImgs: + print("Please run evaluate() first") + # allows input customized parameters + if p is None: + p = self.params + + setImgIds = set(p.imgIds) + + # TPs, FPs, FNs + TPs = np.zeros((len(p.iouThrs),), dtype=np.int64) + FPs = np.zeros((len(p.iouThrs),), dtype=np.int64) + pmFPs = np.zeros((len(p.iouThrs),), dtype=np.int64) + FNs = np.zeros((len(p.iouThrs),), dtype=np.int64) + local_F1s = np.zeros((len(p.iouThrs),), dtype=np.float64) + + # Image level metrics + IL_TPs = 0 + IL_FPs = 0 + IL_TNs = 0 + IL_FNs = 0 + IL_perfects_neg = np.zeros((len(p.iouThrs),), dtype=np.int64) + IL_perfects_pos = np.zeros((len(p.iouThrs),), dtype=np.int64) + + # JnF metric + total_J = 0 + total_F = 0 + total_JnF = 0 + + valid_img_count = 0 + total_pos_count = 0 + total_neg_count = 0 + valid_J_count = 0 + valid_F1_count = 0 + valid_F1_count_w0dt = 0 + for res in self.evalImgs: + if res["image_id"] not in setImgIds: + continue + IL_TPs += res["IL_TP"] + IL_FPs += res["IL_FP"] + IL_TNs += res["IL_TN"] + IL_FNs += res["IL_FN"] + if "IL_perfect_neg" in res: + IL_perfects_neg += res["IL_perfect_neg"] + total_neg_count += 1 + else: + assert "IL_perfect_pos" in res + IL_perfects_pos += res["IL_perfect_pos"] + total_pos_count += 1 + + if "TPs" not in res: + continue + + TPs += res["TPs"] + FPs += res["FPs"] + FNs += res["FNs"] + valid_img_count += 1 + + if "local_positive_F1s" in res: + local_F1s += res["local_positive_F1s"] + pmFPs += res["FPs"] + valid_F1_count_w0dt += 1 + if res["num_dt"] > 0: + valid_F1_count += 1 + + if "J" in res and res["J"] > -1e-9: + total_J += res["J"] + total_F += res["F"] + total_JnF += res["J&F"] + valid_J_count += 1 + + # compute precision recall and F1 + precision = TPs / (TPs + FPs + 1e-4) + positive_micro_precision = TPs / (TPs + pmFPs + 1e-4) + assert np.all(precision <= 1) + recall = TPs / (TPs + FNs + 1e-4) + assert np.all(recall <= 1) + F1 = 2 * precision * recall / (precision + recall + 1e-4) + positive_micro_F1 = ( + 2 + * positive_micro_precision + * recall + / (positive_micro_precision + recall + 1e-4) + ) + + IL_rec = IL_TPs / (IL_TPs + IL_FNs + 1e-6) + IL_prec = IL_TPs / (IL_TPs + IL_FPs + 1e-6) + IL_F1 = 2 * IL_prec * IL_rec / (IL_prec + IL_rec + 1e-6) + IL_FPR = IL_FPs / (IL_FPs + IL_TNs + 1e-6) + IL_MCC = float(IL_TPs * IL_TNs - IL_FPs * IL_FNs) / ( + ( + float(IL_TPs + IL_FPs) + * float(IL_TPs + IL_FNs) + * float(IL_TNs + IL_FPs) + * float(IL_TNs + IL_FNs) + ) + ** 0.5 + + 1e-6 + ) + IL_perfect_pos = IL_perfects_pos / (total_pos_count + 1e-9) + IL_perfect_neg = IL_perfects_neg / (total_neg_count + 1e-9) + + total_J = total_J / (valid_J_count + 1e-9) + total_F = total_F / (valid_J_count + 1e-9) + total_JnF = total_JnF / (valid_J_count + 1e-9) + + self.eval = { + "params": p, + "TPs": TPs, + "FPs": FPs, + "positive_micro_FPs": pmFPs, + "FNs": FNs, + "precision": precision, + "positive_micro_precision": positive_micro_precision, + "recall": recall, + "F1": F1, + "positive_micro_F1": positive_micro_F1, + "positive_macro_F1": local_F1s / valid_F1_count, + "positive_w0dt_macro_F1": local_F1s / valid_F1_count_w0dt, + "IL_recall": IL_rec, + "IL_precision": IL_prec, + "IL_F1": IL_F1, + "IL_FPR": IL_FPR, + "IL_MCC": IL_MCC, + "IL_perfect_pos": IL_perfect_pos, + "IL_perfect_neg": IL_perfect_neg, + "J": total_J, + "F": total_F, + "J&F": total_JnF, + } + self.eval["CGF1"] = self.eval["positive_macro_F1"] * self.eval["IL_MCC"] + self.eval["CGF1_w0dt"] = ( + self.eval["positive_w0dt_macro_F1"] * self.eval["IL_MCC"] + ) + self.eval["CGF1_micro"] = self.eval["positive_micro_F1"] * self.eval["IL_MCC"] + + def summarize(self): + """ + Compute and display summary metrics for evaluation results. + Note this functin can *only* be applied on the default parameter setting + """ + if not self.eval: + raise Exception("Please run accumulate() first") + + def _summarize(iouThr=None, metric=""): + p = self.params + iStr = " {:<18} @[ IoU={:<9}] = {:0.3f}" + titleStr = "Average " + metric + iouStr = ( + "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1]) + if iouThr is None + else "{:0.2f}".format(iouThr) + ) + + s = self.eval[metric] + # IoU + if iouThr is not None: + t = np.where(iouThr == p.iouThrs)[0] + s = s[t] + + if len(s[s > -1]) == 0: + mean_s = -1 + else: + mean_s = np.mean(s[s > -1]) + print(iStr.format(titleStr, iouStr, mean_s)) + return mean_s + + def _summarize_single(metric=""): + titleStr = "Average " + metric + iStr = " {:<35} = {:0.3f}" + s = self.eval[metric] + print(iStr.format(titleStr, s)) + return s + + def _summarizeDets(): + # note: the index of these metrics are also used in video Demo F1 evaluation + # when adding new metrics, please update the index in video Demo F1 evaluation + # in "evaluate" method of the "VideoDemoF1Evaluator" class + stats = np.zeros((len(DEMO_METRICS),)) + stats[0] = _summarize(metric="CGF1") + stats[1] = _summarize(metric="precision") + stats[2] = _summarize(metric="recall") + stats[3] = _summarize(metric="F1") + stats[4] = _summarize(metric="positive_macro_F1") + stats[5] = _summarize_single(metric="IL_precision") + stats[6] = _summarize_single(metric="IL_recall") + stats[7] = _summarize_single(metric="IL_F1") + stats[8] = _summarize_single(metric="IL_FPR") + stats[9] = _summarize_single(metric="IL_MCC") + stats[10] = _summarize(metric="IL_perfect_pos") + stats[11] = _summarize(metric="IL_perfect_neg") + stats[12] = _summarize(iouThr=0.5, metric="CGF1") + stats[13] = _summarize(iouThr=0.5, metric="precision") + stats[14] = _summarize(iouThr=0.5, metric="recall") + stats[15] = _summarize(iouThr=0.5, metric="F1") + stats[16] = _summarize(iouThr=0.5, metric="positive_macro_F1") + stats[17] = _summarize(iouThr=0.5, metric="IL_perfect_pos") + stats[18] = _summarize(iouThr=0.5, metric="IL_perfect_neg") + stats[19] = _summarize(iouThr=0.75, metric="CGF1") + stats[20] = _summarize(iouThr=0.75, metric="precision") + stats[21] = _summarize(iouThr=0.75, metric="recall") + stats[22] = _summarize(iouThr=0.75, metric="F1") + stats[23] = _summarize(iouThr=0.75, metric="positive_macro_F1") + stats[24] = _summarize(iouThr=0.75, metric="IL_perfect_pos") + stats[25] = _summarize(iouThr=0.75, metric="IL_perfect_neg") + stats[26] = _summarize_single(metric="J") + stats[27] = _summarize_single(metric="F") + stats[28] = _summarize_single(metric="J&F") + stats[29] = _summarize(metric="CGF1_micro") + stats[30] = _summarize(metric="positive_micro_precision") + stats[31] = _summarize(metric="positive_micro_F1") + stats[32] = _summarize(iouThr=0.5, metric="CGF1_micro") + stats[33] = _summarize(iouThr=0.5, metric="positive_micro_precision") + stats[34] = _summarize(iouThr=0.5, metric="positive_micro_F1") + stats[35] = _summarize(iouThr=0.75, metric="CGF1_micro") + stats[36] = _summarize(iouThr=0.75, metric="positive_micro_precision") + stats[37] = _summarize(iouThr=0.75, metric="positive_micro_F1") + stats[38] = _summarize(metric="CGF1_w0dt") + stats[39] = _summarize(metric="positive_w0dt_macro_F1") + stats[40] = _summarize(iouThr=0.5, metric="CGF1_w0dt") + stats[41] = _summarize(iouThr=0.5, metric="positive_w0dt_macro_F1") + stats[42] = _summarize(iouThr=0.75, metric="CGF1_w0dt") + stats[43] = _summarize(iouThr=0.75, metric="positive_w0dt_macro_F1") + return stats + + summarize = _summarizeDets + self.stats = summarize() + + +DEMO_METRICS = [ + "CGF1", + "Precision", + "Recall", + "F1", + "Macro_F1", + "IL_Precision", + "IL_Recall", + "IL_F1", + "IL_FPR", + "IL_MCC", + "IL_perfect_pos", + "IL_perfect_neg", + "CGF1@0.5", + "Precision@0.5", + "Recall@0.5", + "F1@0.5", + "Macro_F1@0.5", + "IL_perfect_pos@0.5", + "IL_perfect_neg@0.5", + "CGF1@0.75", + "Precision@0.75", + "Recall@0.75", + "F1@0.75", + "Macro_F1@0.75", + "IL_perfect_pos@0.75", + "IL_perfect_neg@0.75", + "J", + "F", + "J&F", + "CGF1_micro", + "positive_micro_Precision", + "positive_micro_F1", + "CGF1_micro@0.5", + "positive_micro_Precision@0.5", + "positive_micro_F1@0.5", + "CGF1_micro@0.75", + "positive_micro_Precision@0.75", + "positive_micro_F1@0.75", + "CGF1_w0dt", + "positive_w0dt_macro_F1", + "CGF1_w0dt@0.5", + "positive_w0dt_macro_F1@0.5", + "CGF1_w0dt@0.75", + "positive_w0dt_macro_F1@0.75", +] + + +class DemoEvaluator(CocoEvaluator): + def __init__( + self, + coco_gt, + iou_types, + dump_dir: Optional[str], + postprocessor, + threshold=0.5, + average_by_rarity=False, + gather_pred_via_filesys=False, + exhaustive_only=False, + all_exhaustive_only=True, + compute_JnF=False, + metrics_dump_dir: Optional[str] = None, + ): + self.iou_types = iou_types + self.threshold = threshold + super().__init__( + coco_gt=coco_gt, + iou_types=iou_types, + useCats=False, + dump_dir=dump_dir, + postprocessor=postprocessor, + # average_by_rarity=average_by_rarity, + gather_pred_via_filesys=gather_pred_via_filesys, + exhaustive_only=exhaustive_only, + all_exhaustive_only=all_exhaustive_only, + metrics_dump_dir=metrics_dump_dir, + ) + + self.use_self_evaluate = True + self.compute_JnF = compute_JnF + + def _lazy_init(self): + if self.initialized: + return + super()._lazy_init() + self.use_self_evaluate = True + self.reset() + + def select_best_scoring(self, scorings): + # This function is used for "oracle" type evaluation. + # It accepts the evaluation results with respect to several ground truths, and picks the best + if len(scorings) == 1: + return scorings[0] + + assert ( + scorings[0].ndim == 3 + ), f"Expecting results in [numCats, numAreas, numImgs] format, got {scorings[0].shape}" + assert ( + scorings[0].shape[0] == 1 + ), f"Expecting a single category, got {scorings[0].shape[0]}" + + for scoring in scorings: + assert ( + scoring.shape == scorings[0].shape + ), f"Shape mismatch: {scoring.shape}, {scorings[0].shape}" + + selected_imgs = [] + for img_id in range(scorings[0].shape[-1]): + best = scorings[0][:, :, img_id] + + for scoring in scorings[1:]: + current = scoring[:, :, img_id] + if "local_F1s" in best[0, 0] and "local_F1s" in current[0, 0]: + # we were able to compute a F1 score for this particular image in both evaluations + # best["local_F1s"] contains the results at various IoU thresholds. We simply take the average for comparision + best_score = best[0, 0]["local_F1s"].mean() + current_score = current[0, 0]["local_F1s"].mean() + if current_score > best_score: + best = current + + else: + # If we're here, it means that in that in some evaluation we were not able to get a valid local F1 + # This happens when both the predictions and targets are empty. In that case, we can assume it's a perfect prediction + if "local_F1s" not in current[0, 0]: + best = current + selected_imgs.append(best) + result = np.stack(selected_imgs, axis=-1) + assert result.shape == scorings[0].shape + return result + + def summarize(self): + self._lazy_init() + logging.info("Demo evaluator: Summarizing") + if not is_main_process(): + return {} + outs = {} + prefix = "oracle_" if len(self.coco_evals) > 1 else "" + # if self.rarity_buckets is None: + self.accumulate(self.eval_img_ids) + for iou_type, coco_eval in self.coco_evals[0].items(): + print("Demo metric, IoU type={}".format(iou_type)) + coco_eval.summarize() + + if "bbox" in self.coco_evals[0]: + for i, value in enumerate(self.coco_evals[0]["bbox"].stats): + outs[f"coco_eval_bbox_{prefix}{DEMO_METRICS[i]}"] = value + if "segm" in self.coco_evals[0]: + for i, value in enumerate(self.coco_evals[0]["segm"].stats): + outs[f"coco_eval_masks_{prefix}{DEMO_METRICS[i]}"] = value + # else: + # total_stats = {} + # for bucket, img_list in self.rarity_buckets.items(): + # self.accumulate(imgIds=img_list) + # bucket_name = RARITY_BUCKETS[bucket] + # for iou_type, coco_eval in self.coco_evals[0].items(): + # print( + # "Demo metric, IoU type={}, Rarity bucket={}".format( + # iou_type, bucket_name + # ) + # ) + # coco_eval.summarize() + + # if "bbox" in self.coco_evals[0]: + # if "bbox" not in total_stats: + # total_stats["bbox"] = np.zeros_like( + # self.coco_evals[0]["bbox"].stats + # ) + # total_stats["bbox"] += self.coco_evals[0]["bbox"].stats + # for i, value in enumerate(self.coco_evals[0]["bbox"].stats): + # outs[ + # f"coco_eval_bbox_{bucket_name}_{prefix}{DEMO_METRICS[i]}" + # ] = value + # if "segm" in self.coco_evals[0]: + # if "segm" not in total_stats: + # total_stats["segm"] = np.zeros_like( + # self.coco_evals[0]["segm"].stats + # ) + # total_stats["segm"] += self.coco_evals[0]["segm"].stats + # for i, value in enumerate(self.coco_evals[0]["segm"].stats): + # outs[ + # f"coco_eval_masks_{bucket_name}_{prefix}{DEMO_METRICS[i]}" + # ] = value + + # if "bbox" in total_stats: + # total_stats["bbox"] /= len(self.rarity_buckets) + # for i, value in enumerate(total_stats["bbox"]): + # outs[f"coco_eval_bbox_{prefix}{DEMO_METRICS[i]}"] = value + # if "segm" in total_stats: + # total_stats["segm"] /= len(self.rarity_buckets) + # for i, value in enumerate(total_stats["segm"]): + # outs[f"coco_eval_masks_{prefix}{DEMO_METRICS[i]}"] = value + + return outs + + def accumulate(self, imgIds=None): + self._lazy_init() + logging.info( + f"demo evaluator: Accumulating on {len(imgIds) if imgIds is not None else 'all'} images" + ) + if not is_main_process(): + return + + if imgIds is not None: + for coco_eval in self.coco_evals[0].values(): + coco_eval.params.imgIds = list(imgIds) + + for coco_eval in self.coco_evals[0].values(): + coco_eval.accumulate() + + def reset(self): + self.coco_evals = [{} for _ in range(len(self.coco_gts))] + for i, coco_gt in enumerate(self.coco_gts): + for iou_type in self.iou_types: + self.coco_evals[i][iou_type] = DemoEval( + coco_gt=coco_gt, + iouType=iou_type, + threshold=self.threshold, + compute_JnF=self.compute_JnF, + ) + self.coco_evals[i][iou_type].useCats = False + self.img_ids = [] + self.eval_imgs = {k: [] for k in self.iou_types} + if self.dump is not None: + self.dump = [] diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/__init__.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/__init__.py new file mode 100644 index 0000000..9c0fa90 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/__init__.py @@ -0,0 +1 @@ +# flake8: noqa diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/run_ytvis_eval.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/run_ytvis_eval.py new file mode 100644 index 0000000..c39dd05 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/run_ytvis_eval.py @@ -0,0 +1,114 @@ +# flake8: noqa + +"""run_youtube_vis.py +Run example: +run_youtube_vis.py --USE_PARALLEL False --METRICS HOTA --TRACKERS_TO_EVAL STEm_Seg +Command Line Arguments: Defaults, # Comments + Eval arguments: + 'USE_PARALLEL': False, + 'NUM_PARALLEL_CORES': 8, + 'BREAK_ON_ERROR': True, # Raises exception and exits with error + 'RETURN_ON_ERROR': False, # if not BREAK_ON_ERROR, then returns from function on error + 'LOG_ON_ERROR': os.path.join(code_path, 'error_log.txt'), # if not None, save any errors into a log file. + 'PRINT_RESULTS': True, + 'PRINT_ONLY_COMBINED': False, + 'PRINT_CONFIG': True, + 'TIME_PROGRESS': True, + 'DISPLAY_LESS_PROGRESS': True, + 'OUTPUT_SUMMARY': True, + 'OUTPUT_EMPTY_CLASSES': True, # If False, summary files are not output for classes with no detections + 'OUTPUT_DETAILED': True, + 'PLOT_CURVES': True, + Dataset arguments: + 'GT_FOLDER': os.path.join(code_path, 'data/gt/youtube_vis/youtube_vis_training'), # Location of GT data + 'TRACKERS_FOLDER': os.path.join(code_path, 'data/trackers/youtube_vis/youtube_vis_training'), + # Trackers location + 'OUTPUT_FOLDER': None, # Where to save eval results (if None, same as TRACKERS_FOLDER) + 'TRACKERS_TO_EVAL': None, # Filenames of trackers to eval (if None, all in folder) + 'CLASSES_TO_EVAL': None, # Classes to eval (if None, all classes) + 'SPLIT_TO_EVAL': 'training', # Valid: 'training', 'val' + 'PRINT_CONFIG': True, # Whether to print current config + 'OUTPUT_SUB_FOLDER': '', # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER + 'TRACKER_SUB_FOLDER': 'data', # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER + 'TRACKER_DISPLAY_NAMES': None, # Names of trackers to display, if None: TRACKERS_TO_EVAL + Metric arguments: + 'METRICS': ['TrackMAP', 'HOTA', 'CLEAR', 'Identity'] +""" + +import argparse +import os +import sys +from multiprocessing import freeze_support + +from . import trackeval + + +def run_ytvis_eval(args=None, gt_json=None, dt_json=None): + # Command line interface: + default_eval_config = trackeval.Evaluator.get_default_eval_config() + # print only combined since TrackMAP is undefined for per sequence breakdowns + default_eval_config["PRINT_ONLY_COMBINED"] = True + default_dataset_config = trackeval.datasets.YouTubeVIS.get_default_dataset_config() + default_metrics_config = {"METRICS": ["HOTA"]} + config = { + **default_eval_config, + **default_dataset_config, + **default_metrics_config, + } # Merge default configs + parser = argparse.ArgumentParser() + for setting in config.keys(): + if type(config[setting]) == list or type(config[setting]) == type(None): + parser.add_argument("--" + setting, nargs="+") + else: + parser.add_argument("--" + setting) + args = parser.parse_args(args).__dict__ + for setting in args.keys(): + if args[setting] is not None: + if type(config[setting]) == type(True): + if args[setting] == "True": + x = True + elif args[setting] == "False": + x = False + else: + raise Exception( + "Command line parameter " + setting + "must be True or False" + ) + elif type(config[setting]) == type(1): + x = int(args[setting]) + elif type(args[setting]) == type(None): + x = None + else: + x = args[setting] + config[setting] = x + eval_config = {k: v for k, v in config.items() if k in default_eval_config.keys()} + dataset_config = { + k: v for k, v in config.items() if k in default_dataset_config.keys() + } + metrics_config = { + k: v for k, v in config.items() if k in default_metrics_config.keys() + } + + # Run code + evaluator = trackeval.Evaluator(eval_config) + # allow directly specifying the GT JSON data and Tracker (result) + # JSON data as Python objects, without reading from files. + dataset_config["GT_JSON_OBJECT"] = gt_json + dataset_config["TRACKER_JSON_OBJECT"] = dt_json + dataset_list = [trackeval.datasets.YouTubeVIS(dataset_config)] + metrics_list = [] + # for metric in [trackeval.metrics.TrackMAP, trackeval.metrics.HOTA, trackeval.metrics.CLEAR, + # trackeval.metrics.Identity]: + for metric in [trackeval.metrics.HOTA]: + if metric.get_name() in metrics_config["METRICS"]: + metrics_list.append(metric()) + if len(metrics_list) == 0: + raise Exception("No metrics selected for evaluation") + output_res, output_msg = evaluator.evaluate(dataset_list, metrics_list) + return output_res, output_msg + + +if __name__ == "__main__": + import sys + + freeze_support() + run_ytvis_eval(sys.argv[1:]) diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/__init__.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/__init__.py new file mode 100644 index 0000000..131e2b7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/__init__.py @@ -0,0 +1,4 @@ +# flake8: noqa + +from . import datasets, metrics, utils +from .eval import Evaluator diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/_timing.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/_timing.py new file mode 100644 index 0000000..ad414dd --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/_timing.py @@ -0,0 +1,68 @@ +# flake8: noqa + +import inspect +from functools import wraps +from time import perf_counter + +DO_TIMING = False +DISPLAY_LESS_PROGRESS = False +timer_dict = {} +counter = 0 + + +def time(f): + @wraps(f) + def wrap(*args, **kw): + if DO_TIMING: + # Run function with timing + ts = perf_counter() + result = f(*args, **kw) + te = perf_counter() + tt = te - ts + + # Get function name + arg_names = inspect.getfullargspec(f)[0] + if arg_names[0] == "self" and DISPLAY_LESS_PROGRESS: + return result + elif arg_names[0] == "self": + method_name = type(args[0]).__name__ + "." + f.__name__ + else: + method_name = f.__name__ + + # Record accumulative time in each function for analysis + if method_name in timer_dict.keys(): + timer_dict[method_name] += tt + else: + timer_dict[method_name] = tt + + # If code is finished, display timing summary + if method_name == "Evaluator.evaluate": + print("") + print("Timing analysis:") + for key, value in timer_dict.items(): + print("%-70s %2.4f sec" % (key, value)) + else: + # Get function argument values for printing special arguments of interest + arg_titles = ["tracker", "seq", "cls"] + arg_vals = [] + for i, a in enumerate(arg_names): + if a in arg_titles: + arg_vals.append(args[i]) + arg_text = "(" + ", ".join(arg_vals) + ")" + + # Display methods and functions with different indentation. + if arg_names[0] == "self": + print("%-74s %2.4f sec" % (" " * 4 + method_name + arg_text, tt)) + elif arg_names[0] == "test": + pass + else: + global counter + counter += 1 + print("%i %-70s %2.4f sec" % (counter, method_name + arg_text, tt)) + + return result + else: + # If config["TIME_PROGRESS"] is false, or config["USE_PARALLEL"] is true, run functions normally without timing. + return f(*args, **kw) + + return wrap diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/__init__.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/__init__.py new file mode 100644 index 0000000..42f1620 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/__init__.py @@ -0,0 +1,4 @@ +# flake8: noqa + +from .tao_ow import TAO_OW +from .youtube_vis import YouTubeVIS diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/_base_dataset.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/_base_dataset.py new file mode 100644 index 0000000..fb68025 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/_base_dataset.py @@ -0,0 +1,379 @@ +# flake8: noqa + +import csv +import io +import os +import traceback +import zipfile +from abc import ABC, abstractmethod +from copy import deepcopy + +import numpy as np + +from .. import _timing +from ..utils import TrackEvalException + + +class _BaseDataset(ABC): + @abstractmethod + def __init__(self): + self.tracker_list = None + self.seq_list = None + self.class_list = None + self.output_fol = None + self.output_sub_fol = None + self.should_classes_combine = True + self.use_super_categories = False + + # Functions to implement: + + @staticmethod + @abstractmethod + def get_default_dataset_config(): ... + + @abstractmethod + def _load_raw_file(self, tracker, seq, is_gt): ... + + @_timing.time + @abstractmethod + def get_preprocessed_seq_data(self, raw_data, cls): ... + + @abstractmethod + def _calculate_similarities(self, gt_dets_t, tracker_dets_t): ... + + # Helper functions for all datasets: + + @classmethod + def get_class_name(cls): + return cls.__name__ + + def get_name(self): + return self.get_class_name() + + def get_output_fol(self, tracker): + return os.path.join(self.output_fol, tracker, self.output_sub_fol) + + def get_display_name(self, tracker): + """Can be overwritten if the trackers name (in files) is different to how it should be displayed. + By default this method just returns the trackers name as is. + """ + return tracker + + def get_eval_info(self): + """Return info about the dataset needed for the Evaluator""" + return self.tracker_list, self.seq_list, self.class_list + + @_timing.time + def get_raw_seq_data(self, tracker, seq): + """Loads raw data (tracker and ground-truth) for a single tracker on a single sequence. + Raw data includes all of the information needed for both preprocessing and evaluation, for all classes. + A later function (get_processed_seq_data) will perform such preprocessing and extract relevant information for + the evaluation of each class. + + This returns a dict which contains the fields: + [num_timesteps]: integer + [gt_ids, tracker_ids, gt_classes, tracker_classes, tracker_confidences]: + list (for each timestep) of 1D NDArrays (for each det). + [gt_dets, tracker_dets, gt_crowd_ignore_regions]: list (for each timestep) of lists of detections. + [similarity_scores]: list (for each timestep) of 2D NDArrays. + [gt_extras]: dict (for each extra) of lists (for each timestep) of 1D NDArrays (for each det). + + gt_extras contains dataset specific information used for preprocessing such as occlusion and truncation levels. + + Note that similarities are extracted as part of the dataset and not the metric, because almost all metrics are + independent of the exact method of calculating the similarity. However datasets are not (e.g. segmentation + masks vs 2D boxes vs 3D boxes). + We calculate the similarity before preprocessing because often both preprocessing and evaluation require it and + we don't wish to calculate this twice. + We calculate similarity between all gt and tracker classes (not just each class individually) to allow for + calculation of metrics such as class confusion matrices. Typically the impact of this on performance is low. + """ + # Load raw data. + raw_gt_data = self._load_raw_file(tracker, seq, is_gt=True) + raw_tracker_data = self._load_raw_file(tracker, seq, is_gt=False) + raw_data = {**raw_tracker_data, **raw_gt_data} # Merges dictionaries + + # Calculate similarities for each timestep. + similarity_scores = [] + for t, (gt_dets_t, tracker_dets_t) in enumerate( + zip(raw_data["gt_dets"], raw_data["tracker_dets"]) + ): + ious = self._calculate_similarities(gt_dets_t, tracker_dets_t) + similarity_scores.append(ious) + raw_data["similarity_scores"] = similarity_scores + return raw_data + + @staticmethod + def _load_simple_text_file( + file, + time_col=0, + id_col=None, + remove_negative_ids=False, + valid_filter=None, + crowd_ignore_filter=None, + convert_filter=None, + is_zipped=False, + zip_file=None, + force_delimiters=None, + ): + """Function that loads data which is in a commonly used text file format. + Assumes each det is given by one row of a text file. + There is no limit to the number or meaning of each column, + however one column needs to give the timestep of each det (time_col) which is default col 0. + + The file dialect (deliminator, num cols, etc) is determined automatically. + This function automatically separates dets by timestep, + and is much faster than alternatives such as np.loadtext or pandas. + + If remove_negative_ids is True and id_col is not None, dets with negative values in id_col are excluded. + These are not excluded from ignore data. + + valid_filter can be used to only include certain classes. + It is a dict with ints as keys, and lists as values, + such that a row is included if "row[key].lower() is in value" for all key/value pairs in the dict. + If None, all classes are included. + + crowd_ignore_filter can be used to read crowd_ignore regions separately. It has the same format as valid filter. + + convert_filter can be used to convert value read to another format. + This is used most commonly to convert classes given as string to a class id. + This is a dict such that the key is the column to convert, and the value is another dict giving the mapping. + + Optionally, input files could be a zip of multiple text files for storage efficiency. + + Returns read_data and ignore_data. + Each is a dict (with keys as timesteps as strings) of lists (over dets) of lists (over column values). + Note that all data is returned as strings, and must be converted to float/int later if needed. + Note that timesteps will not be present in the returned dict keys if there are no dets for them + """ + + if remove_negative_ids and id_col is None: + raise TrackEvalException( + "remove_negative_ids is True, but id_col is not given." + ) + if crowd_ignore_filter is None: + crowd_ignore_filter = {} + if convert_filter is None: + convert_filter = {} + try: + if is_zipped: # Either open file directly or within a zip. + if zip_file is None: + raise TrackEvalException( + "is_zipped set to True, but no zip_file is given." + ) + archive = zipfile.ZipFile(os.path.join(zip_file), "r") + fp = io.TextIOWrapper(archive.open(file, "r")) + else: + fp = open(file) + read_data = {} + crowd_ignore_data = {} + fp.seek(0, os.SEEK_END) + # check if file is empty + if fp.tell(): + fp.seek(0) + dialect = csv.Sniffer().sniff( + fp.readline(), delimiters=force_delimiters + ) # Auto determine structure. + dialect.skipinitialspace = ( + True # Deal with extra spaces between columns + ) + fp.seek(0) + reader = csv.reader(fp, dialect) + for row in reader: + try: + # Deal with extra trailing spaces at the end of rows + if row[-1] in "": + row = row[:-1] + timestep = str(int(float(row[time_col]))) + # Read ignore regions separately. + is_ignored = False + for ignore_key, ignore_value in crowd_ignore_filter.items(): + if row[ignore_key].lower() in ignore_value: + # Convert values in one column (e.g. string to id) + for ( + convert_key, + convert_value, + ) in convert_filter.items(): + row[convert_key] = convert_value[ + row[convert_key].lower() + ] + # Save data separated by timestep. + if timestep in crowd_ignore_data.keys(): + crowd_ignore_data[timestep].append(row) + else: + crowd_ignore_data[timestep] = [row] + is_ignored = True + if ( + is_ignored + ): # if det is an ignore region, it cannot be a normal det. + continue + # Exclude some dets if not valid. + if valid_filter is not None: + for key, value in valid_filter.items(): + if row[key].lower() not in value: + continue + if remove_negative_ids: + if int(float(row[id_col])) < 0: + continue + # Convert values in one column (e.g. string to id) + for convert_key, convert_value in convert_filter.items(): + row[convert_key] = convert_value[row[convert_key].lower()] + # Save data separated by timestep. + if timestep in read_data.keys(): + read_data[timestep].append(row) + else: + read_data[timestep] = [row] + except Exception: + exc_str_init = ( + "In file %s the following line cannot be read correctly: \n" + % os.path.basename(file) + ) + exc_str = " ".join([exc_str_init] + row) + raise TrackEvalException(exc_str) + fp.close() + except Exception: + print("Error loading file: %s, printing traceback." % file) + traceback.print_exc() + raise TrackEvalException( + "File %s cannot be read because it is either not present or invalidly formatted" + % os.path.basename(file) + ) + return read_data, crowd_ignore_data + + @staticmethod + def _calculate_mask_ious(masks1, masks2, is_encoded=False, do_ioa=False): + """Calculates the IOU (intersection over union) between two arrays of segmentation masks. + If is_encoded a run length encoding with pycocotools is assumed as input format, otherwise an input of numpy + arrays of the shape (num_masks, height, width) is assumed and the encoding is performed. + If do_ioa (intersection over area) , then calculates the intersection over the area of masks1 - this is commonly + used to determine if detections are within crowd ignore region. + :param masks1: first set of masks (numpy array of shape (num_masks, height, width) if not encoded, + else pycocotools rle encoded format) + :param masks2: second set of masks (numpy array of shape (num_masks, height, width) if not encoded, + else pycocotools rle encoded format) + :param is_encoded: whether the input is in pycocotools rle encoded format + :param do_ioa: whether to perform IoA computation + :return: the IoU/IoA scores + """ + + # Only loaded when run to reduce minimum requirements + from pycocotools import mask as mask_utils + + # use pycocotools for run length encoding of masks + if not is_encoded: + masks1 = mask_utils.encode( + np.array(np.transpose(masks1, (1, 2, 0)), order="F") + ) + masks2 = mask_utils.encode( + np.array(np.transpose(masks2, (1, 2, 0)), order="F") + ) + + # use pycocotools for iou computation of rle encoded masks + ious = mask_utils.iou(masks1, masks2, [do_ioa] * len(masks2)) + if len(masks1) == 0 or len(masks2) == 0: + ious = np.asarray(ious).reshape(len(masks1), len(masks2)) + assert (ious >= 0 - np.finfo("float").eps).all() + assert (ious <= 1 + np.finfo("float").eps).all() + + return ious + + @staticmethod + def _calculate_box_ious(bboxes1, bboxes2, box_format="xywh", do_ioa=False): + """Calculates the IOU (intersection over union) between two arrays of boxes. + Allows variable box formats ('xywh' and 'x0y0x1y1'). + If do_ioa (intersection over area) , then calculates the intersection over the area of boxes1 - this is commonly + used to determine if detections are within crowd ignore region. + """ + if box_format in "xywh": + # layout: (x0, y0, w, h) + bboxes1 = deepcopy(bboxes1) + bboxes2 = deepcopy(bboxes2) + + bboxes1[:, 2] = bboxes1[:, 0] + bboxes1[:, 2] + bboxes1[:, 3] = bboxes1[:, 1] + bboxes1[:, 3] + bboxes2[:, 2] = bboxes2[:, 0] + bboxes2[:, 2] + bboxes2[:, 3] = bboxes2[:, 1] + bboxes2[:, 3] + elif box_format not in "x0y0x1y1": + raise (TrackEvalException("box_format %s is not implemented" % box_format)) + + # layout: (x0, y0, x1, y1) + min_ = np.minimum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :]) + max_ = np.maximum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :]) + intersection = np.maximum(min_[..., 2] - max_[..., 0], 0) * np.maximum( + min_[..., 3] - max_[..., 1], 0 + ) + area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * ( + bboxes1[..., 3] - bboxes1[..., 1] + ) + + if do_ioa: + ioas = np.zeros_like(intersection) + valid_mask = area1 > 0 + np.finfo("float").eps + ioas[valid_mask, :] = ( + intersection[valid_mask, :] / area1[valid_mask][:, np.newaxis] + ) + + return ioas + else: + area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * ( + bboxes2[..., 3] - bboxes2[..., 1] + ) + union = area1[:, np.newaxis] + area2[np.newaxis, :] - intersection + intersection[area1 <= 0 + np.finfo("float").eps, :] = 0 + intersection[:, area2 <= 0 + np.finfo("float").eps] = 0 + intersection[union <= 0 + np.finfo("float").eps] = 0 + union[union <= 0 + np.finfo("float").eps] = 1 + ious = intersection / union + return ious + + @staticmethod + def _calculate_euclidean_similarity(dets1, dets2, zero_distance=2.0): + """Calculates the euclidean distance between two sets of detections, and then converts this into a similarity + measure with values between 0 and 1 using the following formula: sim = max(0, 1 - dist/zero_distance). + The default zero_distance of 2.0, corresponds to the default used in MOT15_3D, such that a 0.5 similarity + threshold corresponds to a 1m distance threshold for TPs. + """ + dist = np.linalg.norm(dets1[:, np.newaxis] - dets2[np.newaxis, :], axis=2) + sim = np.maximum(0, 1 - dist / zero_distance) + return sim + + @staticmethod + def _check_unique_ids(data, after_preproc=False): + """Check the requirement that the tracker_ids and gt_ids are unique per timestep""" + gt_ids = data["gt_ids"] + tracker_ids = data["tracker_ids"] + for t, (gt_ids_t, tracker_ids_t) in enumerate(zip(gt_ids, tracker_ids)): + if len(tracker_ids_t) > 0: + unique_ids, counts = np.unique(tracker_ids_t, return_counts=True) + if np.max(counts) != 1: + duplicate_ids = unique_ids[counts > 1] + exc_str_init = ( + "Tracker predicts the same ID more than once in a single timestep " + "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1) + ) + exc_str = ( + " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")" + ) + if after_preproc: + exc_str_init += ( + "\n Note that this error occurred after preprocessing (but not before), " + "so ids may not be as in file, and something seems wrong with preproc." + ) + raise TrackEvalException(exc_str) + if len(gt_ids_t) > 0: + unique_ids, counts = np.unique(gt_ids_t, return_counts=True) + if np.max(counts) != 1: + duplicate_ids = unique_ids[counts > 1] + exc_str_init = ( + "Ground-truth has the same ID more than once in a single timestep " + "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1) + ) + exc_str = ( + " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")" + ) + if after_preproc: + exc_str_init += ( + "\n Note that this error occurred after preprocessing (but not before), " + "so ids may not be as in file, and something seems wrong with preproc." + ) + raise TrackEvalException(exc_str) diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/tao_ow.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/tao_ow.py new file mode 100644 index 0000000..06bc93b --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/tao_ow.py @@ -0,0 +1,891 @@ +# flake8: noqa + +import itertools +import json +import os +from collections import defaultdict + +import numpy as np +from scipy.optimize import linear_sum_assignment + +from .. import _timing, utils +from ..utils import TrackEvalException +from ._base_dataset import _BaseDataset + + +class TAO_OW(_BaseDataset): + """Dataset class for TAO tracking""" + + @staticmethod + def get_default_dataset_config(): + """Default class config values""" + code_path = utils.get_code_path() + default_config = { + "GT_FOLDER": os.path.join( + code_path, "data/gt/tao/tao_training" + ), # Location of GT data + "TRACKERS_FOLDER": os.path.join( + code_path, "data/trackers/tao/tao_training" + ), # Trackers location + "OUTPUT_FOLDER": None, # Where to save eval results (if None, same as TRACKERS_FOLDER) + "TRACKERS_TO_EVAL": None, # Filenames of trackers to eval (if None, all in folder) + "CLASSES_TO_EVAL": None, # Classes to eval (if None, all classes) + "SPLIT_TO_EVAL": "training", # Valid: 'training', 'val' + "PRINT_CONFIG": True, # Whether to print current config + "TRACKER_SUB_FOLDER": "data", # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER + "OUTPUT_SUB_FOLDER": "", # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER + "TRACKER_DISPLAY_NAMES": None, # Names of trackers to display, if None: TRACKERS_TO_EVAL + "MAX_DETECTIONS": 300, # Number of maximal allowed detections per image (0 for unlimited) + "SUBSET": "all", + } + return default_config + + def __init__(self, config=None): + """Initialise dataset, checking that all required files are present""" + super().__init__() + # Fill non-given config values with defaults + self.config = utils.init_config( + config, self.get_default_dataset_config(), self.get_name() + ) + self.gt_fol = self.config["GT_FOLDER"] + self.tracker_fol = self.config["TRACKERS_FOLDER"] + self.should_classes_combine = True + self.use_super_categories = False + + self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"] + self.output_fol = self.config["OUTPUT_FOLDER"] + if self.output_fol is None: + self.output_fol = self.tracker_fol + self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"] + + gt_dir_files = [ + file for file in os.listdir(self.gt_fol) if file.endswith(".json") + ] + if len(gt_dir_files) != 1: + raise TrackEvalException( + self.gt_fol + " does not contain exactly one json file." + ) + + with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f: + self.gt_data = json.load(f) + + self.subset = self.config["SUBSET"] + if self.subset != "all": + # Split GT data into `known`, `unknown` or `distractor` + self._split_known_unknown_distractor() + self.gt_data = self._filter_gt_data(self.gt_data) + + # merge categories marked with a merged tag in TAO dataset + self._merge_categories(self.gt_data["annotations"] + self.gt_data["tracks"]) + + # Get sequences to eval and sequence information + self.seq_list = [ + vid["name"].replace("/", "-") for vid in self.gt_data["videos"] + ] + self.seq_name_to_seq_id = { + vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"] + } + # compute mappings from videos to annotation data + self.videos_to_gt_tracks, self.videos_to_gt_images = self._compute_vid_mappings( + self.gt_data["annotations"] + ) + # compute sequence lengths + self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]} + for img in self.gt_data["images"]: + self.seq_lengths[img["video_id"]] += 1 + self.seq_to_images_to_timestep = self._compute_image_to_timestep_mappings() + self.seq_to_classes = { + vid["id"]: { + "pos_cat_ids": list( + { + track["category_id"] + for track in self.videos_to_gt_tracks[vid["id"]] + } + ), + "neg_cat_ids": vid["neg_category_ids"], + "not_exhaustively_labeled_cat_ids": vid["not_exhaustive_category_ids"], + } + for vid in self.gt_data["videos"] + } + + # Get classes to eval + considered_vid_ids = [self.seq_name_to_seq_id[vid] for vid in self.seq_list] + seen_cats = set( + [ + cat_id + for vid_id in considered_vid_ids + for cat_id in self.seq_to_classes[vid_id]["pos_cat_ids"] + ] + ) + # only classes with ground truth are evaluated in TAO + self.valid_classes = [ + cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats + ] + # cls_name_to_cls_id_map = {cls['name']: cls['id'] for cls in self.gt_data['categories']} + + if self.config["CLASSES_TO_EVAL"]: + # self.class_list = [cls.lower() if cls.lower() in self.valid_classes else None + # for cls in self.config['CLASSES_TO_EVAL']] + self.class_list = ["object"] # class-agnostic + if not all(self.class_list): + raise TrackEvalException( + "Attempted to evaluate an invalid class. Only classes " + + ", ".join(self.valid_classes) + + " are valid (classes present in ground truth data)." + ) + else: + # self.class_list = [cls for cls in self.valid_classes] + self.class_list = ["object"] # class-agnostic + # self.class_name_to_class_id = {k: v for k, v in cls_name_to_cls_id_map.items() if k in self.class_list} + self.class_name_to_class_id = {"object": 1} # class-agnostic + + # Get trackers to eval + if self.config["TRACKERS_TO_EVAL"] is None: + self.tracker_list = os.listdir(self.tracker_fol) + else: + self.tracker_list = self.config["TRACKERS_TO_EVAL"] + + if self.config["TRACKER_DISPLAY_NAMES"] is None: + self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list)) + elif (self.config["TRACKERS_TO_EVAL"] is not None) and ( + len(self.config["TRACKER_DISPLAY_NAMES"]) == len(self.tracker_list) + ): + self.tracker_to_disp = dict( + zip(self.tracker_list, self.config["TRACKER_DISPLAY_NAMES"]) + ) + else: + raise TrackEvalException( + "List of tracker files and tracker display names do not match." + ) + + self.tracker_data = {tracker: dict() for tracker in self.tracker_list} + + for tracker in self.tracker_list: + tr_dir_files = [ + file + for file in os.listdir( + os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol) + ) + if file.endswith(".json") + ] + if len(tr_dir_files) != 1: + raise TrackEvalException( + os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol) + + " does not contain exactly one json file." + ) + with open( + os.path.join( + self.tracker_fol, tracker, self.tracker_sub_fol, tr_dir_files[0] + ) + ) as f: + curr_data = json.load(f) + + # limit detections if MAX_DETECTIONS > 0 + if self.config["MAX_DETECTIONS"]: + curr_data = self._limit_dets_per_image(curr_data) + + # fill missing video ids + self._fill_video_ids_inplace(curr_data) + + # make track ids unique over whole evaluation set + self._make_track_ids_unique(curr_data) + + # merge categories marked with a merged tag in TAO dataset + self._merge_categories(curr_data) + + # get tracker sequence information + curr_videos_to_tracker_tracks, curr_videos_to_tracker_images = ( + self._compute_vid_mappings(curr_data) + ) + self.tracker_data[tracker]["vids_to_tracks"] = curr_videos_to_tracker_tracks + self.tracker_data[tracker]["vids_to_images"] = curr_videos_to_tracker_images + + def get_display_name(self, tracker): + return self.tracker_to_disp[tracker] + + def _load_raw_file(self, tracker, seq, is_gt): + """Load a file (gt or tracker) in the TAO format + + If is_gt, this returns a dict which contains the fields: + [gt_ids, gt_classes] : list (for each timestep) of 1D NDArrays (for each det). + [gt_dets]: list (for each timestep) of lists of detections. + [classes_to_gt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as + keys and corresponding segmentations as values) for each track + [classes_to_gt_track_ids, classes_to_gt_track_areas, classes_to_gt_track_lengths]: dictionary with class values + as keys and lists (for each track) as values + + if not is_gt, this returns a dict which contains the fields: + [tracker_ids, tracker_classes, tracker_confidences] : list (for each timestep) of 1D NDArrays (for each det). + [tracker_dets]: list (for each timestep) of lists of detections. + [classes_to_dt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as + keys and corresponding segmentations as values) for each track + [classes_to_dt_track_ids, classes_to_dt_track_areas, classes_to_dt_track_lengths]: dictionary with class values + as keys and lists as values + [classes_to_dt_track_scores]: dictionary with class values as keys and 1D numpy arrays as values + """ + seq_id = self.seq_name_to_seq_id[seq] + # File location + if is_gt: + imgs = self.videos_to_gt_images[seq_id] + else: + imgs = self.tracker_data[tracker]["vids_to_images"][seq_id] + + # Convert data to required format + num_timesteps = self.seq_lengths[seq_id] + img_to_timestep = self.seq_to_images_to_timestep[seq_id] + data_keys = ["ids", "classes", "dets"] + if not is_gt: + data_keys += ["tracker_confidences"] + raw_data = {key: [None] * num_timesteps for key in data_keys} + for img in imgs: + # some tracker data contains images without any ground truth information, these are ignored + try: + t = img_to_timestep[img["id"]] + except KeyError: + continue + annotations = img["annotations"] + raw_data["dets"][t] = np.atleast_2d( + [ann["bbox"] for ann in annotations] + ).astype(float) + raw_data["ids"][t] = np.atleast_1d( + [ann["track_id"] for ann in annotations] + ).astype(int) + raw_data["classes"][t] = np.atleast_1d([1 for _ in annotations]).astype( + int + ) # class-agnostic + if not is_gt: + raw_data["tracker_confidences"][t] = np.atleast_1d( + [ann["score"] for ann in annotations] + ).astype(float) + + for t, d in enumerate(raw_data["dets"]): + if d is None: + raw_data["dets"][t] = np.empty((0, 4)).astype(float) + raw_data["ids"][t] = np.empty(0).astype(int) + raw_data["classes"][t] = np.empty(0).astype(int) + if not is_gt: + raw_data["tracker_confidences"][t] = np.empty(0) + + if is_gt: + key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"} + else: + key_map = { + "ids": "tracker_ids", + "classes": "tracker_classes", + "dets": "tracker_dets", + } + for k, v in key_map.items(): + raw_data[v] = raw_data.pop(k) + + # all_classes = [self.class_name_to_class_id[cls] for cls in self.class_list] + all_classes = [1] # class-agnostic + + if is_gt: + classes_to_consider = all_classes + all_tracks = self.videos_to_gt_tracks[seq_id] + else: + # classes_to_consider = self.seq_to_classes[seq_id]['pos_cat_ids'] \ + # + self.seq_to_classes[seq_id]['neg_cat_ids'] + classes_to_consider = all_classes # class-agnostic + all_tracks = self.tracker_data[tracker]["vids_to_tracks"][seq_id] + + # classes_to_tracks = {cls: [track for track in all_tracks if track['category_id'] == cls] + # if cls in classes_to_consider else [] for cls in all_classes} + classes_to_tracks = { + cls: [track for track in all_tracks] if cls in classes_to_consider else [] + for cls in all_classes + } # class-agnostic + + # mapping from classes to track information + raw_data["classes_to_tracks"] = { + cls: [ + { + det["image_id"]: np.atleast_1d(det["bbox"]) + for det in track["annotations"] + } + for track in tracks + ] + for cls, tracks in classes_to_tracks.items() + } + raw_data["classes_to_track_ids"] = { + cls: [track["id"] for track in tracks] + for cls, tracks in classes_to_tracks.items() + } + raw_data["classes_to_track_areas"] = { + cls: [track["area"] for track in tracks] + for cls, tracks in classes_to_tracks.items() + } + raw_data["classes_to_track_lengths"] = { + cls: [len(track["annotations"]) for track in tracks] + for cls, tracks in classes_to_tracks.items() + } + + if not is_gt: + raw_data["classes_to_dt_track_scores"] = { + cls: np.array( + [ + np.mean([float(x["score"]) for x in track["annotations"]]) + for track in tracks + ] + ) + for cls, tracks in classes_to_tracks.items() + } + + if is_gt: + key_map = { + "classes_to_tracks": "classes_to_gt_tracks", + "classes_to_track_ids": "classes_to_gt_track_ids", + "classes_to_track_lengths": "classes_to_gt_track_lengths", + "classes_to_track_areas": "classes_to_gt_track_areas", + } + else: + key_map = { + "classes_to_tracks": "classes_to_dt_tracks", + "classes_to_track_ids": "classes_to_dt_track_ids", + "classes_to_track_lengths": "classes_to_dt_track_lengths", + "classes_to_track_areas": "classes_to_dt_track_areas", + } + for k, v in key_map.items(): + raw_data[v] = raw_data.pop(k) + + raw_data["num_timesteps"] = num_timesteps + raw_data["neg_cat_ids"] = self.seq_to_classes[seq_id]["neg_cat_ids"] + raw_data["not_exhaustively_labeled_cls"] = self.seq_to_classes[seq_id][ + "not_exhaustively_labeled_cat_ids" + ] + raw_data["seq"] = seq + return raw_data + + @_timing.time + def get_preprocessed_seq_data(self, raw_data, cls): + """Preprocess data for a single sequence for a single class ready for evaluation. + Inputs: + - raw_data is a dict containing the data for the sequence already read in by get_raw_seq_data(). + - cls is the class to be evaluated. + Outputs: + - data is a dict containing all of the information that metrics need to perform evaluation. + It contains the following fields: + [num_timesteps, num_gt_ids, num_tracker_ids, num_gt_dets, num_tracker_dets] : integers. + [gt_ids, tracker_ids, tracker_confidences]: list (for each timestep) of 1D NDArrays (for each det). + [gt_dets, tracker_dets]: list (for each timestep) of lists of detections. + [similarity_scores]: list (for each timestep) of 2D NDArrays. + Notes: + General preprocessing (preproc) occurs in 4 steps. Some datasets may not use all of these steps. + 1) Extract only detections relevant for the class to be evaluated (including distractor detections). + 2) Match gt dets and tracker dets. Remove tracker dets that are matched to a gt det that is of a + distractor class, or otherwise marked as to be removed. + 3) Remove unmatched tracker dets if they fall within a crowd ignore region or don't meet a certain + other criteria (e.g. are too small). + 4) Remove gt dets that were only useful for preprocessing and not for actual evaluation. + After the above preprocessing steps, this function also calculates the number of gt and tracker detections + and unique track ids. It also relabels gt and tracker ids to be contiguous and checks that ids are + unique within each timestep. + TAO: + In TAO, the 4 preproc steps are as follow: + 1) All classes present in the ground truth data are evaluated separately. + 2) No matched tracker detections are removed. + 3) Unmatched tracker detections are removed if there is not ground truth data and the class does not + belong to the categories marked as negative for this sequence. Additionally, unmatched tracker + detections for classes which are marked as not exhaustively labeled are removed. + 4) No gt detections are removed. + Further, for TrackMAP computation track representations for the given class are accessed from a dictionary + and the tracks from the tracker data are sorted according to the tracker confidence. + """ + cls_id = self.class_name_to_class_id[cls] + is_not_exhaustively_labeled = cls_id in raw_data["not_exhaustively_labeled_cls"] + is_neg_category = cls_id in raw_data["neg_cat_ids"] + + data_keys = [ + "gt_ids", + "tracker_ids", + "gt_dets", + "tracker_dets", + "tracker_confidences", + "similarity_scores", + ] + data = {key: [None] * raw_data["num_timesteps"] for key in data_keys} + unique_gt_ids = [] + unique_tracker_ids = [] + num_gt_dets = 0 + num_tracker_dets = 0 + for t in range(raw_data["num_timesteps"]): + # Only extract relevant dets for this class for preproc and eval (cls) + gt_class_mask = np.atleast_1d(raw_data["gt_classes"][t] == cls_id) + gt_class_mask = gt_class_mask.astype(bool) + gt_ids = raw_data["gt_ids"][t][gt_class_mask] + gt_dets = raw_data["gt_dets"][t][gt_class_mask] + + tracker_class_mask = np.atleast_1d(raw_data["tracker_classes"][t] == cls_id) + tracker_class_mask = tracker_class_mask.astype(bool) + tracker_ids = raw_data["tracker_ids"][t][tracker_class_mask] + tracker_dets = raw_data["tracker_dets"][t][tracker_class_mask] + tracker_confidences = raw_data["tracker_confidences"][t][tracker_class_mask] + similarity_scores = raw_data["similarity_scores"][t][gt_class_mask, :][ + :, tracker_class_mask + ] + + # Match tracker and gt dets (with hungarian algorithm). + unmatched_indices = np.arange(tracker_ids.shape[0]) + if gt_ids.shape[0] > 0 and tracker_ids.shape[0] > 0: + matching_scores = similarity_scores.copy() + matching_scores[matching_scores < 0.5 - np.finfo("float").eps] = 0 + match_rows, match_cols = linear_sum_assignment(-matching_scores) + actually_matched_mask = ( + matching_scores[match_rows, match_cols] > 0 + np.finfo("float").eps + ) + match_cols = match_cols[actually_matched_mask] + unmatched_indices = np.delete(unmatched_indices, match_cols, axis=0) + + if gt_ids.shape[0] == 0 and not is_neg_category: + to_remove_tracker = unmatched_indices + elif is_not_exhaustively_labeled: + to_remove_tracker = unmatched_indices + else: + to_remove_tracker = np.array([], dtype=int) + + # remove all unwanted unmatched tracker detections + data["tracker_ids"][t] = np.delete(tracker_ids, to_remove_tracker, axis=0) + data["tracker_dets"][t] = np.delete(tracker_dets, to_remove_tracker, axis=0) + data["tracker_confidences"][t] = np.delete( + tracker_confidences, to_remove_tracker, axis=0 + ) + similarity_scores = np.delete(similarity_scores, to_remove_tracker, axis=1) + + data["gt_ids"][t] = gt_ids + data["gt_dets"][t] = gt_dets + data["similarity_scores"][t] = similarity_scores + + unique_gt_ids += list(np.unique(data["gt_ids"][t])) + unique_tracker_ids += list(np.unique(data["tracker_ids"][t])) + num_tracker_dets += len(data["tracker_ids"][t]) + num_gt_dets += len(data["gt_ids"][t]) + + # Re-label IDs such that there are no empty IDs + if len(unique_gt_ids) > 0: + unique_gt_ids = np.unique(unique_gt_ids) + gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1)) + gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids)) + for t in range(raw_data["num_timesteps"]): + if len(data["gt_ids"][t]) > 0: + data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int) + if len(unique_tracker_ids) > 0: + unique_tracker_ids = np.unique(unique_tracker_ids) + tracker_id_map = np.nan * np.ones((np.max(unique_tracker_ids) + 1)) + tracker_id_map[unique_tracker_ids] = np.arange(len(unique_tracker_ids)) + for t in range(raw_data["num_timesteps"]): + if len(data["tracker_ids"][t]) > 0: + data["tracker_ids"][t] = tracker_id_map[ + data["tracker_ids"][t] + ].astype(int) + + # Record overview statistics. + data["num_tracker_dets"] = num_tracker_dets + data["num_gt_dets"] = num_gt_dets + data["num_tracker_ids"] = len(unique_tracker_ids) + data["num_gt_ids"] = len(unique_gt_ids) + data["num_timesteps"] = raw_data["num_timesteps"] + data["seq"] = raw_data["seq"] + + # get track representations + data["gt_tracks"] = raw_data["classes_to_gt_tracks"][cls_id] + data["gt_track_ids"] = raw_data["classes_to_gt_track_ids"][cls_id] + data["gt_track_lengths"] = raw_data["classes_to_gt_track_lengths"][cls_id] + data["gt_track_areas"] = raw_data["classes_to_gt_track_areas"][cls_id] + data["dt_tracks"] = raw_data["classes_to_dt_tracks"][cls_id] + data["dt_track_ids"] = raw_data["classes_to_dt_track_ids"][cls_id] + data["dt_track_lengths"] = raw_data["classes_to_dt_track_lengths"][cls_id] + data["dt_track_areas"] = raw_data["classes_to_dt_track_areas"][cls_id] + data["dt_track_scores"] = raw_data["classes_to_dt_track_scores"][cls_id] + data["not_exhaustively_labeled"] = is_not_exhaustively_labeled + data["iou_type"] = "bbox" + + # sort tracker data tracks by tracker confidence scores + if data["dt_tracks"]: + idx = np.argsort( + [-score for score in data["dt_track_scores"]], kind="mergesort" + ) + data["dt_track_scores"] = [data["dt_track_scores"][i] for i in idx] + data["dt_tracks"] = [data["dt_tracks"][i] for i in idx] + data["dt_track_ids"] = [data["dt_track_ids"][i] for i in idx] + data["dt_track_lengths"] = [data["dt_track_lengths"][i] for i in idx] + data["dt_track_areas"] = [data["dt_track_areas"][i] for i in idx] + # Ensure that ids are unique per timestep. + self._check_unique_ids(data) + + return data + + def _calculate_similarities(self, gt_dets_t, tracker_dets_t): + similarity_scores = self._calculate_box_ious(gt_dets_t, tracker_dets_t) + return similarity_scores + + def _merge_categories(self, annotations): + """ + Merges categories with a merged tag. Adapted from https://github.com/TAO-Dataset + :param annotations: the annotations in which the classes should be merged + :return: None + """ + merge_map = {} + for category in self.gt_data["categories"]: + if "merged" in category: + for to_merge in category["merged"]: + merge_map[to_merge["id"]] = category["id"] + + for ann in annotations: + ann["category_id"] = merge_map.get(ann["category_id"], ann["category_id"]) + + def _compute_vid_mappings(self, annotations): + """ + Computes mappings from Videos to corresponding tracks and images. + :param annotations: the annotations for which the mapping should be generated + :return: the video-to-track-mapping, the video-to-image-mapping + """ + vids_to_tracks = {} + vids_to_imgs = {} + vid_ids = [vid["id"] for vid in self.gt_data["videos"]] + + # compute an mapping from image IDs to images + images = {} + for image in self.gt_data["images"]: + images[image["id"]] = image + + for ann in annotations: + ann["area"] = ann["bbox"][2] * ann["bbox"][3] + + vid = ann["video_id"] + if ann["video_id"] not in vids_to_tracks.keys(): + vids_to_tracks[ann["video_id"]] = list() + if ann["video_id"] not in vids_to_imgs.keys(): + vids_to_imgs[ann["video_id"]] = list() + + # Fill in vids_to_tracks + tid = ann["track_id"] + exist_tids = [track["id"] for track in vids_to_tracks[vid]] + try: + index1 = exist_tids.index(tid) + except ValueError: + index1 = -1 + if tid not in exist_tids: + curr_track = { + "id": tid, + "category_id": ann["category_id"], + "video_id": vid, + "annotations": [ann], + } + vids_to_tracks[vid].append(curr_track) + else: + vids_to_tracks[vid][index1]["annotations"].append(ann) + + # Fill in vids_to_imgs + img_id = ann["image_id"] + exist_img_ids = [img["id"] for img in vids_to_imgs[vid]] + try: + index2 = exist_img_ids.index(img_id) + except ValueError: + index2 = -1 + if index2 == -1: + curr_img = {"id": img_id, "annotations": [ann]} + vids_to_imgs[vid].append(curr_img) + else: + vids_to_imgs[vid][index2]["annotations"].append(ann) + + # sort annotations by frame index and compute track area + for vid, tracks in vids_to_tracks.items(): + for track in tracks: + track["annotations"] = sorted( + track["annotations"], + key=lambda x: images[x["image_id"]]["frame_index"], + ) + # Computer average area + track["area"] = sum(x["area"] for x in track["annotations"]) / len( + track["annotations"] + ) + + # Ensure all videos are present + for vid_id in vid_ids: + if vid_id not in vids_to_tracks.keys(): + vids_to_tracks[vid_id] = [] + if vid_id not in vids_to_imgs.keys(): + vids_to_imgs[vid_id] = [] + + return vids_to_tracks, vids_to_imgs + + def _compute_image_to_timestep_mappings(self): + """ + Computes a mapping from images to the corresponding timestep in the sequence. + :return: the image-to-timestep-mapping + """ + images = {} + for image in self.gt_data["images"]: + images[image["id"]] = image + + seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]} + for vid in seq_to_imgs_to_timestep: + curr_imgs = [img["id"] for img in self.videos_to_gt_images[vid]] + curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_index"]) + seq_to_imgs_to_timestep[vid] = { + curr_imgs[i]: i for i in range(len(curr_imgs)) + } + + return seq_to_imgs_to_timestep + + def _limit_dets_per_image(self, annotations): + """ + Limits the number of detections for each image to config['MAX_DETECTIONS']. Adapted from + https://github.com/TAO-Dataset/ + :param annotations: the annotations in which the detections should be limited + :return: the annotations with limited detections + """ + max_dets = self.config["MAX_DETECTIONS"] + img_ann = defaultdict(list) + for ann in annotations: + img_ann[ann["image_id"]].append(ann) + + for img_id, _anns in img_ann.items(): + if len(_anns) <= max_dets: + continue + _anns = sorted(_anns, key=lambda x: x["score"], reverse=True) + img_ann[img_id] = _anns[:max_dets] + + return [ann for anns in img_ann.values() for ann in anns] + + def _fill_video_ids_inplace(self, annotations): + """ + Fills in missing video IDs inplace. Adapted from https://github.com/TAO-Dataset/ + :param annotations: the annotations for which the videos IDs should be filled inplace + :return: None + """ + missing_video_id = [x for x in annotations if "video_id" not in x] + if missing_video_id: + image_id_to_video_id = { + x["id"]: x["video_id"] for x in self.gt_data["images"] + } + for x in missing_video_id: + x["video_id"] = image_id_to_video_id[x["image_id"]] + + @staticmethod + def _make_track_ids_unique(annotations): + """ + Makes the track IDs unqiue over the whole annotation set. Adapted from https://github.com/TAO-Dataset/ + :param annotations: the annotation set + :return: the number of updated IDs + """ + track_id_videos = {} + track_ids_to_update = set() + max_track_id = 0 + for ann in annotations: + t = ann["track_id"] + if t not in track_id_videos: + track_id_videos[t] = ann["video_id"] + + if ann["video_id"] != track_id_videos[t]: + # Track id is assigned to multiple videos + track_ids_to_update.add(t) + max_track_id = max(max_track_id, t) + + if track_ids_to_update: + print("true") + next_id = itertools.count(max_track_id + 1) + new_track_ids = defaultdict(lambda: next(next_id)) + for ann in annotations: + t = ann["track_id"] + v = ann["video_id"] + if t in track_ids_to_update: + ann["track_id"] = new_track_ids[t, v] + return len(track_ids_to_update) + + def _split_known_unknown_distractor(self): + all_ids = set( + [i for i in range(1, 2000)] + ) # 2000 is larger than the max category id in TAO-OW. + # `knowns` includes 78 TAO_category_ids that corresponds to 78 COCO classes. + # (The other 2 COCO classes do not have corresponding classes in TAO). + self.knowns = { + 4, + 13, + 1038, + 544, + 1057, + 34, + 35, + 36, + 41, + 45, + 58, + 60, + 579, + 1091, + 1097, + 1099, + 78, + 79, + 81, + 91, + 1115, + 1117, + 95, + 1122, + 99, + 1132, + 621, + 1135, + 625, + 118, + 1144, + 126, + 642, + 1155, + 133, + 1162, + 139, + 154, + 174, + 185, + 699, + 1215, + 714, + 717, + 1229, + 211, + 729, + 221, + 229, + 747, + 235, + 237, + 779, + 276, + 805, + 299, + 829, + 852, + 347, + 371, + 382, + 896, + 392, + 926, + 937, + 428, + 429, + 961, + 452, + 979, + 980, + 982, + 475, + 480, + 993, + 1001, + 502, + 1018, + } + # `distractors` is defined as in the paper "Opening up Open-World Tracking" + self.distractors = { + 20, + 63, + 108, + 180, + 188, + 204, + 212, + 247, + 303, + 403, + 407, + 415, + 490, + 504, + 507, + 513, + 529, + 567, + 569, + 588, + 672, + 691, + 702, + 708, + 711, + 720, + 736, + 737, + 798, + 813, + 815, + 827, + 831, + 851, + 877, + 883, + 912, + 971, + 976, + 1130, + 1133, + 1134, + 1169, + 1184, + 1220, + } + self.unknowns = all_ids.difference(self.knowns.union(self.distractors)) + + def _filter_gt_data(self, raw_gt_data): + """ + Filter out irrelevant data in the raw_gt_data + Args: + raw_gt_data: directly loaded from json. + + Returns: + filtered gt_data + """ + valid_cat_ids = list() + if self.subset == "known": + valid_cat_ids = self.knowns + elif self.subset == "distractor": + valid_cat_ids = self.distractors + elif self.subset == "unknown": + valid_cat_ids = self.unknowns + # elif self.subset == "test_only_unknowns": + # valid_cat_ids = test_only_unknowns + else: + raise Exception("The parameter `SUBSET` is incorrect") + + filtered = dict() + filtered["videos"] = raw_gt_data["videos"] + # filtered["videos"] = list() + unwanted_vid = set() + # for video in raw_gt_data["videos"]: + # datasrc = video["name"].split('/')[1] + # if datasrc in data_srcs: + # filtered["videos"].append(video) + # else: + # unwanted_vid.add(video["id"]) + + filtered["annotations"] = list() + for ann in raw_gt_data["annotations"]: + if (ann["video_id"] not in unwanted_vid) and ( + ann["category_id"] in valid_cat_ids + ): + filtered["annotations"].append(ann) + + filtered["tracks"] = list() + for track in raw_gt_data["tracks"]: + if (track["video_id"] not in unwanted_vid) and ( + track["category_id"] in valid_cat_ids + ): + filtered["tracks"].append(track) + + filtered["images"] = list() + for image in raw_gt_data["images"]: + if image["video_id"] not in unwanted_vid: + filtered["images"].append(image) + + filtered["categories"] = list() + for cat in raw_gt_data["categories"]: + if cat["id"] in valid_cat_ids: + filtered["categories"].append(cat) + + filtered["info"] = raw_gt_data["info"] + filtered["licenses"] = raw_gt_data["licenses"] + + return filtered diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/youtube_vis.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/youtube_vis.py new file mode 100644 index 0000000..e611398 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/datasets/youtube_vis.py @@ -0,0 +1,524 @@ +# flake8: noqa + +# note: this file has been modified from its original version in TrackEval in +# https://github.com/JonathonLuiten/TrackEval/blob/master/trackeval/datasets/youtube_vis.py +# to support the following: +# 1) bbox evaluation (via `IOU_TYPE`) +# 2) passing GT and prediction data as Python objects (via `GT_JSON_OBJECT` and `TRACKER_JSON_OBJECT`) +# 3) specifying a custom dataset name (via `DATASET_NAME`) + +import json +import os + +import numpy as np + +from .. import _timing, utils +from ..utils import TrackEvalException +from ._base_dataset import _BaseDataset + + +class YouTubeVIS(_BaseDataset): + """Dataset class for YouTubeVIS tracking""" + + @staticmethod + def get_default_dataset_config(): + """Default class config values""" + code_path = utils.get_code_path() + default_config = { + "GT_FOLDER": os.path.join( + code_path, "data/gt/youtube_vis/" + ), # Location of GT data + "TRACKERS_FOLDER": os.path.join(code_path, "data/trackers/youtube_vis/"), + # Trackers location + "OUTPUT_FOLDER": None, # Where to save eval results (if None, same as TRACKERS_FOLDER) + "TRACKERS_TO_EVAL": None, # Filenames of trackers to eval (if None, all in folder) + "CLASSES_TO_EVAL": None, # Classes to eval (if None, all classes) + "SPLIT_TO_EVAL": "train_sub_split", # Valid: 'train', 'val', 'train_sub_split' + "PRINT_CONFIG": True, # Whether to print current config + "OUTPUT_SUB_FOLDER": "", # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER + "TRACKER_SUB_FOLDER": "data", # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER + "TRACKER_DISPLAY_NAMES": None, # Names of trackers to display, if None: TRACKERS_TO_EVAL + # Added for video phrase AP evaluation -- allow directly specifying the GT JSON data and Tracker (result) + # JSON data as Python objects, without reading from files. + "GT_JSON_OBJECT": None, + "TRACKER_JSON_OBJECT": None, + "IOU_TYPE": "segm", + "DATASET_NAME": "video", + } + return default_config + + def __init__(self, config=None): + """Initialise dataset, checking that all required files are present""" + super().__init__() + # Fill non-given config values with defaults + self.config = utils.init_config(config, self.get_default_dataset_config()) + self.gt_fol = ( + self.config["GT_FOLDER"] + "youtube_vis_" + self.config["SPLIT_TO_EVAL"] + ) + self.tracker_fol = ( + self.config["TRACKERS_FOLDER"] + + "youtube_vis_" + + self.config["SPLIT_TO_EVAL"] + ) + self.use_super_categories = False + self.should_classes_combine = True + assert self.config["IOU_TYPE"] in ["segm", "bbox"] + self.iou_type = self.config["IOU_TYPE"] + print("=" * 100) + print(f"Evaluate annotation type *{self.iou_type}*") + self.dataset_name = self.config["DATASET_NAME"] + + self.output_fol = self.config["OUTPUT_FOLDER"] + if self.output_fol is None: + self.output_fol = self.tracker_fol + self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"] + self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"] + + if self.config["GT_JSON_OBJECT"] is not None: + # allow directly specifying the GT JSON data without reading from files + gt_json = self.config["GT_JSON_OBJECT"] + assert isinstance(gt_json, dict) + assert "videos" in gt_json + assert "categories" in gt_json + assert "annotations" in gt_json + self.gt_data = gt_json + else: + if not os.path.exists(self.gt_fol): + print("GT folder not found: " + self.gt_fol) + raise TrackEvalException( + "GT folder not found: " + os.path.basename(self.gt_fol) + ) + gt_dir_files = [ + file for file in os.listdir(self.gt_fol) if file.endswith(".json") + ] + if len(gt_dir_files) != 1: + raise TrackEvalException( + self.gt_fol + " does not contain exactly one json file." + ) + + with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f: + self.gt_data = json.load(f) + + # Get classes to eval + self.valid_classes = [cls["name"] for cls in self.gt_data["categories"]] + cls_name_to_cls_id_map = { + cls["name"]: cls["id"] for cls in self.gt_data["categories"] + } + + if self.config["CLASSES_TO_EVAL"]: + self.class_list = [ + cls.lower() if cls.lower() in self.valid_classes else None + for cls in self.config["CLASSES_TO_EVAL"] + ] + if not all(self.class_list): + raise TrackEvalException( + "Attempted to evaluate an invalid class. Only classes " + + ", ".join(self.valid_classes) + + " are valid." + ) + else: + self.class_list = [cls["name"] for cls in self.gt_data["categories"]] + self.class_name_to_class_id = { + k: v for k, v in cls_name_to_cls_id_map.items() if k in self.class_list + } + + # Get sequences to eval and check gt files exist + self.seq_list = [ + vid["file_names"][0].split("/")[0] for vid in self.gt_data["videos"] + ] + self.seq_name_to_seq_id = { + vid["file_names"][0].split("/")[0]: vid["id"] + for vid in self.gt_data["videos"] + } + self.seq_lengths = { + vid["id"]: len(vid["file_names"]) for vid in self.gt_data["videos"] + } + + # encode masks and compute track areas + self._prepare_gt_annotations() + + # Get trackers to eval + if self.config["TRACKER_JSON_OBJECT"] is not None: + # allow directly specifying the tracker JSON data without reading from files + tracker_json = self.config["TRACKER_JSON_OBJECT"] + assert isinstance(tracker_json, list) + self.tracker_list = ["tracker"] + elif self.config["TRACKERS_TO_EVAL"] is None: + self.tracker_list = os.listdir(self.tracker_fol) + else: + self.tracker_list = self.config["TRACKERS_TO_EVAL"] + + if self.config["TRACKER_DISPLAY_NAMES"] is None: + self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list)) + elif (self.config["TRACKERS_TO_EVAL"] is not None) and ( + len(self.config["TRACKER_DISPLAY_NAMES"]) == len(self.tracker_list) + ): + self.tracker_to_disp = dict( + zip(self.tracker_list, self.config["TRACKER_DISPLAY_NAMES"]) + ) + else: + raise TrackEvalException( + "List of tracker files and tracker display names do not match." + ) + + # counter for globally unique track IDs + self.global_tid_counter = 0 + + self.tracker_data = dict() + if self.config["TRACKER_JSON_OBJECT"] is not None: + # allow directly specifying the tracker JSON data without reading from files + tracker = self.tracker_list[0] + self.tracker_data[tracker] = tracker_json + else: + for tracker in self.tracker_list: + tracker_dir_path = os.path.join( + self.tracker_fol, tracker, self.tracker_sub_fol + ) + tr_dir_files = [ + file + for file in os.listdir(tracker_dir_path) + if file.endswith(".json") + ] + if len(tr_dir_files) != 1: + raise TrackEvalException( + tracker_dir_path + " does not contain exactly one json file." + ) + + with open(os.path.join(tracker_dir_path, tr_dir_files[0])) as f: + curr_data = json.load(f) + + self.tracker_data[tracker] = curr_data + + def get_display_name(self, tracker): + return self.tracker_to_disp[tracker] + + def _load_raw_file(self, tracker, seq, is_gt): + """Load a file (gt or tracker) in the YouTubeVIS format + If is_gt, this returns a dict which contains the fields: + [gt_ids, gt_classes] : list (for each timestep) of 1D NDArrays (for each det). + [gt_dets]: list (for each timestep) of lists of detections. + [classes_to_gt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as + keys and corresponding segmentations as values) for each track + [classes_to_gt_track_ids, classes_to_gt_track_areas, classes_to_gt_track_iscrowd]: dictionary with class values + as keys and lists (for each track) as values + + if not is_gt, this returns a dict which contains the fields: + [tracker_ids, tracker_classes, tracker_confidences] : list (for each timestep) of 1D NDArrays (for each det). + [tracker_dets]: list (for each timestep) of lists of detections. + [classes_to_dt_tracks]: dictionary with class values as keys and list of dictionaries (with frame indices as + keys and corresponding segmentations as values) for each track + [classes_to_dt_track_ids, classes_to_dt_track_areas]: dictionary with class values as keys and lists as values + [classes_to_dt_track_scores]: dictionary with class values as keys and 1D numpy arrays as values + """ + # select sequence tracks + seq_id = self.seq_name_to_seq_id[seq] + if is_gt: + tracks = [ + ann for ann in self.gt_data["annotations"] if ann["video_id"] == seq_id + ] + else: + tracks = self._get_tracker_seq_tracks(tracker, seq_id) + + # Convert data to required format + num_timesteps = self.seq_lengths[seq_id] + data_keys = ["ids", "classes", "dets"] + if not is_gt: + data_keys += ["tracker_confidences"] + raw_data = {key: [None] * num_timesteps for key in data_keys} + result_key = "segmentations" if self.iou_type == "segm" else "bboxes" + for t in range(num_timesteps): + raw_data["dets"][t] = [ + track[result_key][t] for track in tracks if track[result_key][t] + ] + raw_data["ids"][t] = np.atleast_1d( + [track["id"] for track in tracks if track[result_key][t]] + ).astype(int) + raw_data["classes"][t] = np.atleast_1d( + [track["category_id"] for track in tracks if track[result_key][t]] + ).astype(int) + if not is_gt: + raw_data["tracker_confidences"][t] = np.atleast_1d( + [track["score"] for track in tracks if track[result_key][t]] + ).astype(float) + + if is_gt: + key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"} + else: + key_map = { + "ids": "tracker_ids", + "classes": "tracker_classes", + "dets": "tracker_dets", + } + for k, v in key_map.items(): + raw_data[v] = raw_data.pop(k) + + all_cls_ids = {self.class_name_to_class_id[cls] for cls in self.class_list} + classes_to_tracks = { + cls: [track for track in tracks if track["category_id"] == cls] + for cls in all_cls_ids + } + + # mapping from classes to track representations and track information + raw_data["classes_to_tracks"] = { + cls: [ + {i: track[result_key][i] for i in range(len(track[result_key]))} + for track in tracks + ] + for cls, tracks in classes_to_tracks.items() + } + raw_data["classes_to_track_ids"] = { + cls: [track["id"] for track in tracks] + for cls, tracks in classes_to_tracks.items() + } + raw_data["classes_to_track_areas"] = { + cls: [track["area"] for track in tracks] + for cls, tracks in classes_to_tracks.items() + } + + if is_gt: + raw_data["classes_to_gt_track_iscrowd"] = { + cls: [track["iscrowd"] for track in tracks] + for cls, tracks in classes_to_tracks.items() + } + else: + raw_data["classes_to_dt_track_scores"] = { + cls: np.array([track["score"] for track in tracks]) + for cls, tracks in classes_to_tracks.items() + } + + if is_gt: + key_map = { + "classes_to_tracks": "classes_to_gt_tracks", + "classes_to_track_ids": "classes_to_gt_track_ids", + "classes_to_track_areas": "classes_to_gt_track_areas", + } + else: + key_map = { + "classes_to_tracks": "classes_to_dt_tracks", + "classes_to_track_ids": "classes_to_dt_track_ids", + "classes_to_track_areas": "classes_to_dt_track_areas", + } + for k, v in key_map.items(): + raw_data[v] = raw_data.pop(k) + + raw_data["num_timesteps"] = num_timesteps + raw_data["seq"] = seq + return raw_data + + @_timing.time + def get_preprocessed_seq_data(self, raw_data, cls): + """Preprocess data for a single sequence for a single class ready for evaluation. + Inputs: + - raw_data is a dict containing the data for the sequence already read in by get_raw_seq_data(). + - cls is the class to be evaluated. + Outputs: + - data is a dict containing all of the information that metrics need to perform evaluation. + It contains the following fields: + [num_timesteps, num_gt_ids, num_tracker_ids, num_gt_dets, num_tracker_dets] : integers. + [gt_ids, tracker_ids, tracker_confidences]: list (for each timestep) of 1D NDArrays (for each det). + [gt_dets, tracker_dets]: list (for each timestep) of lists of detections. + [similarity_scores]: list (for each timestep) of 2D NDArrays. + Notes: + General preprocessing (preproc) occurs in 4 steps. Some datasets may not use all of these steps. + 1) Extract only detections relevant for the class to be evaluated (including distractor detections). + 2) Match gt dets and tracker dets. Remove tracker dets that are matched to a gt det that is of a + distractor class, or otherwise marked as to be removed. + 3) Remove unmatched tracker dets if they fall within a crowd ignore region or don't meet a certain + other criteria (e.g. are too small). + 4) Remove gt dets that were only useful for preprocessing and not for actual evaluation. + After the above preprocessing steps, this function also calculates the number of gt and tracker detections + and unique track ids. It also relabels gt and tracker ids to be contiguous and checks that ids are + unique within each timestep. + YouTubeVIS: + In YouTubeVIS, the 4 preproc steps are as follow: + 1) There are 40 classes which are evaluated separately. + 2) No matched tracker dets are removed. + 3) No unmatched tracker dets are removed. + 4) No gt dets are removed. + Further, for TrackMAP computation track representations for the given class are accessed from a dictionary + and the tracks from the tracker data are sorted according to the tracker confidence. + """ + cls_id = self.class_name_to_class_id[cls] + + data_keys = [ + "gt_ids", + "tracker_ids", + "gt_dets", + "tracker_dets", + "similarity_scores", + ] + data = {key: [None] * raw_data["num_timesteps"] for key in data_keys} + unique_gt_ids = [] + unique_tracker_ids = [] + num_gt_dets = 0 + num_tracker_dets = 0 + + for t in range(raw_data["num_timesteps"]): + # Only extract relevant dets for this class for eval (cls) + gt_class_mask = np.atleast_1d(raw_data["gt_classes"][t] == cls_id) + gt_class_mask = gt_class_mask.astype(bool) + gt_ids = raw_data["gt_ids"][t][gt_class_mask] + gt_dets = [ + raw_data["gt_dets"][t][ind] + for ind in range(len(gt_class_mask)) + if gt_class_mask[ind] + ] + + tracker_class_mask = np.atleast_1d(raw_data["tracker_classes"][t] == cls_id) + tracker_class_mask = tracker_class_mask.astype(bool) + tracker_ids = raw_data["tracker_ids"][t][tracker_class_mask] + tracker_dets = [ + raw_data["tracker_dets"][t][ind] + for ind in range(len(tracker_class_mask)) + if tracker_class_mask[ind] + ] + similarity_scores = raw_data["similarity_scores"][t][gt_class_mask, :][ + :, tracker_class_mask + ] + + data["tracker_ids"][t] = tracker_ids + data["tracker_dets"][t] = tracker_dets + data["gt_ids"][t] = gt_ids + data["gt_dets"][t] = gt_dets + data["similarity_scores"][t] = similarity_scores + + unique_gt_ids += list(np.unique(data["gt_ids"][t])) + unique_tracker_ids += list(np.unique(data["tracker_ids"][t])) + num_tracker_dets += len(data["tracker_ids"][t]) + num_gt_dets += len(data["gt_ids"][t]) + + # Re-label IDs such that there are no empty IDs + if len(unique_gt_ids) > 0: + unique_gt_ids = np.unique(unique_gt_ids) + gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1)) + gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids)) + for t in range(raw_data["num_timesteps"]): + if len(data["gt_ids"][t]) > 0: + data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int) + if len(unique_tracker_ids) > 0: + unique_tracker_ids = np.unique(unique_tracker_ids) + tracker_id_map = np.nan * np.ones((np.max(unique_tracker_ids) + 1)) + tracker_id_map[unique_tracker_ids] = np.arange(len(unique_tracker_ids)) + for t in range(raw_data["num_timesteps"]): + if len(data["tracker_ids"][t]) > 0: + data["tracker_ids"][t] = tracker_id_map[ + data["tracker_ids"][t] + ].astype(int) + + # Ensure that ids are unique per timestep. + self._check_unique_ids(data) + + # Record overview statistics. + data["num_tracker_dets"] = num_tracker_dets + data["num_gt_dets"] = num_gt_dets + data["num_tracker_ids"] = len(unique_tracker_ids) + data["num_gt_ids"] = len(unique_gt_ids) + data["num_timesteps"] = raw_data["num_timesteps"] + data["seq"] = raw_data["seq"] + + # get track representations + data["gt_tracks"] = raw_data["classes_to_gt_tracks"][cls_id] + data["gt_track_ids"] = raw_data["classes_to_gt_track_ids"][cls_id] + data["gt_track_areas"] = raw_data["classes_to_gt_track_areas"][cls_id] + data["gt_track_iscrowd"] = raw_data["classes_to_gt_track_iscrowd"][cls_id] + data["dt_tracks"] = raw_data["classes_to_dt_tracks"][cls_id] + data["dt_track_ids"] = raw_data["classes_to_dt_track_ids"][cls_id] + data["dt_track_areas"] = raw_data["classes_to_dt_track_areas"][cls_id] + data["dt_track_scores"] = raw_data["classes_to_dt_track_scores"][cls_id] + data["iou_type"] = "mask" + + # sort tracker data tracks by tracker confidence scores + if data["dt_tracks"]: + idx = np.argsort( + [-score for score in data["dt_track_scores"]], kind="mergesort" + ) + data["dt_track_scores"] = [data["dt_track_scores"][i] for i in idx] + data["dt_tracks"] = [data["dt_tracks"][i] for i in idx] + data["dt_track_ids"] = [data["dt_track_ids"][i] for i in idx] + data["dt_track_areas"] = [data["dt_track_areas"][i] for i in idx] + + return data + + def _calculate_similarities(self, gt_dets_t, tracker_dets_t): + if self.iou_type == "segm": + similarity_scores = self._calculate_mask_ious( + gt_dets_t, tracker_dets_t, is_encoded=True, do_ioa=False + ) + else: + gt_dets_t = np.array(gt_dets_t, dtype=np.float32).reshape(-1, 4) + tracker_dets_t = np.array(tracker_dets_t, dtype=np.float32).reshape(-1, 4) + similarity_scores = self._calculate_box_ious( + gt_dets_t, tracker_dets_t, box_format="xywh", do_ioa=False + ) + return similarity_scores + + def _prepare_gt_annotations(self): + """ + Prepares GT data by rle encoding segmentations and computing the average track area. + :return: None + """ + if self.iou_type == "segm": + # only loaded when needed to reduce minimum requirements + from pycocotools import mask as mask_utils + + for track in self.gt_data["annotations"]: + h = track["height"] + w = track["width"] + for i, seg in enumerate(track["segmentations"]): + if seg is not None and isinstance(seg["counts"], list): + track["segmentations"][i] = mask_utils.frPyObjects(seg, h, w) + areas = [a for a in track["areas"] if a] + if len(areas) == 0: + track["area"] = 0 + else: + track["area"] = np.array(areas).mean() + else: + for track in self.gt_data["annotations"]: + # For bbox eval, compute areas from bboxes if not already available + areas = [a for a in track.get("areas", []) if a] + if not areas: + areas = [] + for bbox in track.get("bboxes", []): + if bbox is not None: + areas.append(bbox[2] * bbox[3]) + track["area"] = np.array(areas).mean() if areas else 0 + + def _get_tracker_seq_tracks(self, tracker, seq_id): + """ + Prepares tracker data for a given sequence. Extracts all annotations for given sequence ID, computes + average track area and assigns a track ID. + :param tracker: the given tracker + :param seq_id: the sequence ID + :return: the extracted tracks + """ + # only loaded when needed to reduce minimum requirements + from pycocotools import mask as mask_utils + + tracks = [ + ann for ann in self.tracker_data[tracker] if ann["video_id"] == seq_id + ] + for track in tracks: + if "areas" not in track: + if self.iou_type == "segm": + for seg in track["segmentations"]: + if seg: + track["areas"].append(mask_utils.area(seg)) + else: + track["areas"].append(None) + else: + for bbox in track["bboxes"]: + if bbox: + track["areas"].append(bbox[2] * bbox[3]) + else: + track["areas"].append(None) + areas = [a for a in track["areas"] if a] + if len(areas) == 0: + track["area"] = 0 + else: + track["area"] = np.array(areas).mean() + track["id"] = self.global_tid_counter + self.global_tid_counter += 1 + return tracks + + def get_name(self): + return self.dataset_name diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/eval.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/eval.py new file mode 100644 index 0000000..d2d7205 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/eval.py @@ -0,0 +1,395 @@ +# flake8: noqa + +import os +import time +import traceback +from functools import partial +from multiprocessing.pool import Pool + +import numpy as np + +from . import _timing, utils +from .metrics import Count +from .utils import TrackEvalException + +try: + import tqdm + + TQDM_IMPORTED = True +except ImportError as _: + TQDM_IMPORTED = False + + +class Evaluator: + """Evaluator class for evaluating different metrics for different datasets""" + + @staticmethod + def get_default_eval_config(): + """Returns the default config values for evaluation""" + code_path = utils.get_code_path() + default_config = { + "USE_PARALLEL": False, + "NUM_PARALLEL_CORES": 8, + "BREAK_ON_ERROR": True, # Raises exception and exits with error + "RETURN_ON_ERROR": False, # if not BREAK_ON_ERROR, then returns from function on error + "LOG_ON_ERROR": os.path.join( + code_path, "error_log.txt" + ), # if not None, save any errors into a log file. + "PRINT_RESULTS": True, + "PRINT_ONLY_COMBINED": False, + "PRINT_CONFIG": True, + "TIME_PROGRESS": True, + "DISPLAY_LESS_PROGRESS": True, + "OUTPUT_SUMMARY": True, + "OUTPUT_EMPTY_CLASSES": True, # If False, summary files are not output for classes with no detections + "OUTPUT_DETAILED": True, + "PLOT_CURVES": True, + } + return default_config + + def __init__(self, config=None): + """Initialise the evaluator with a config file""" + self.config = utils.init_config(config, self.get_default_eval_config(), "Eval") + # Only run timing analysis if not run in parallel. + if self.config["TIME_PROGRESS"] and not self.config["USE_PARALLEL"]: + _timing.DO_TIMING = True + if self.config["DISPLAY_LESS_PROGRESS"]: + _timing.DISPLAY_LESS_PROGRESS = True + + def _combine_results( + self, + res, + metrics_list, + metric_names, + dataset, + res_field="COMBINED_SEQ", + target_tag=None, + ): + assert res_field.startswith("COMBINED_SEQ") + # collecting combined cls keys (cls averaged, det averaged, super classes) + tracker_list, seq_list, class_list = dataset.get_eval_info() + combined_cls_keys = [] + res[res_field] = {} + + # narrow the target for evaluation + if target_tag is not None: + target_video_ids = [ + annot["video_id"] + for annot in dataset.gt_data["annotations"] + if target_tag in annot["tags"] + ] + vid2name = { + video["id"]: video["file_names"][0].split("/")[0] + for video in dataset.gt_data["videos"] + } + target_video_ids = set(target_video_ids) + target_video = [vid2name[video_id] for video_id in target_video_ids] + + if len(target_video) == 0: + raise TrackEvalException( + "No sequences found with the tag %s" % target_tag + ) + + target_annotations = [ + annot + for annot in dataset.gt_data["annotations"] + if annot["video_id"] in target_video_ids + ] + assert all(target_tag in annot["tags"] for annot in target_annotations), ( + f"Not all annotations in the target sequences have the target tag {target_tag}. " + "We currently only support a target tag at the sequence level, not at the annotation level." + ) + else: + target_video = seq_list + + # combine sequences for each class + for c_cls in class_list: + res[res_field][c_cls] = {} + for metric, metric_name in zip(metrics_list, metric_names): + curr_res = { + seq_key: seq_value[c_cls][metric_name] + for seq_key, seq_value in res.items() + if not seq_key.startswith("COMBINED_SEQ") + and seq_key in target_video + } + res[res_field][c_cls][metric_name] = metric.combine_sequences(curr_res) + # combine classes + if dataset.should_classes_combine: + combined_cls_keys += [ + "cls_comb_cls_av", + "cls_comb_det_av", + "all", + ] + res[res_field]["cls_comb_cls_av"] = {} + res[res_field]["cls_comb_det_av"] = {} + for metric, metric_name in zip(metrics_list, metric_names): + cls_res = { + cls_key: cls_value[metric_name] + for cls_key, cls_value in res[res_field].items() + if cls_key not in combined_cls_keys + } + res[res_field]["cls_comb_cls_av"][metric_name] = ( + metric.combine_classes_class_averaged(cls_res) + ) + res[res_field]["cls_comb_det_av"][metric_name] = ( + metric.combine_classes_det_averaged(cls_res) + ) + # combine classes to super classes + if dataset.use_super_categories: + for cat, sub_cats in dataset.super_categories.items(): + combined_cls_keys.append(cat) + res[res_field][cat] = {} + for metric, metric_name in zip(metrics_list, metric_names): + cat_res = { + cls_key: cls_value[metric_name] + for cls_key, cls_value in res[res_field].items() + if cls_key in sub_cats + } + res[res_field][cat][metric_name] = ( + metric.combine_classes_det_averaged(cat_res) + ) + return res, combined_cls_keys + + def _summarize_results( + self, + res, + tracker, + metrics_list, + metric_names, + dataset, + res_field, + combined_cls_keys, + ): + config = self.config + output_fol = dataset.get_output_fol(tracker) + tracker_display_name = dataset.get_display_name(tracker) + for c_cls in res[ + res_field + ].keys(): # class_list + combined classes if calculated + summaries = [] + details = [] + num_dets = res[res_field][c_cls]["Count"]["Dets"] + if config["OUTPUT_EMPTY_CLASSES"] or num_dets > 0: + for metric, metric_name in zip(metrics_list, metric_names): + # for combined classes there is no per sequence evaluation + if c_cls in combined_cls_keys: + table_res = {res_field: res[res_field][c_cls][metric_name]} + else: + table_res = { + seq_key: seq_value[c_cls][metric_name] + for seq_key, seq_value in res.items() + } + + if config["PRINT_RESULTS"] and config["PRINT_ONLY_COMBINED"]: + dont_print = ( + dataset.should_classes_combine + and c_cls not in combined_cls_keys + ) + if not dont_print: + metric.print_table( + {res_field: table_res[res_field]}, + tracker_display_name, + c_cls, + res_field, + res_field, + ) + elif config["PRINT_RESULTS"]: + metric.print_table( + table_res, tracker_display_name, c_cls, res_field, res_field + ) + if config["OUTPUT_SUMMARY"]: + summaries.append(metric.summary_results(table_res)) + if config["OUTPUT_DETAILED"]: + details.append(metric.detailed_results(table_res)) + if config["PLOT_CURVES"]: + metric.plot_single_tracker_results( + table_res, + tracker_display_name, + c_cls, + output_fol, + ) + if config["OUTPUT_SUMMARY"]: + utils.write_summary_results(summaries, c_cls, output_fol) + if config["OUTPUT_DETAILED"]: + utils.write_detailed_results(details, c_cls, output_fol) + + @_timing.time + def evaluate(self, dataset_list, metrics_list, show_progressbar=False): + """Evaluate a set of metrics on a set of datasets""" + config = self.config + metrics_list = metrics_list + [Count()] # Count metrics are always run + metric_names = utils.validate_metrics_list(metrics_list) + dataset_names = [dataset.get_name() for dataset in dataset_list] + output_res = {} + output_msg = {} + + for dataset, dataset_name in zip(dataset_list, dataset_names): + # Get dataset info about what to evaluate + output_res[dataset_name] = {} + output_msg[dataset_name] = {} + tracker_list, seq_list, class_list = dataset.get_eval_info() + print( + "\nEvaluating %i tracker(s) on %i sequence(s) for %i class(es) on %s dataset using the following " + "metrics: %s\n" + % ( + len(tracker_list), + len(seq_list), + len(class_list), + dataset_name, + ", ".join(metric_names), + ) + ) + + # Evaluate each tracker + for tracker in tracker_list: + # if not config['BREAK_ON_ERROR'] then go to next tracker without breaking + try: + # Evaluate each sequence in parallel or in series. + # returns a nested dict (res), indexed like: res[seq][class][metric_name][sub_metric field] + # e.g. res[seq_0001][pedestrian][hota][DetA] + print("\nEvaluating %s\n" % tracker) + time_start = time.time() + if config["USE_PARALLEL"]: + if show_progressbar and TQDM_IMPORTED: + seq_list_sorted = sorted(seq_list) + + with Pool(config["NUM_PARALLEL_CORES"]) as pool, tqdm.tqdm( + total=len(seq_list) + ) as pbar: + _eval_sequence = partial( + eval_sequence, + dataset=dataset, + tracker=tracker, + class_list=class_list, + metrics_list=metrics_list, + metric_names=metric_names, + ) + results = [] + for r in pool.imap( + _eval_sequence, seq_list_sorted, chunksize=20 + ): + results.append(r) + pbar.update() + res = dict(zip(seq_list_sorted, results)) + + else: + with Pool(config["NUM_PARALLEL_CORES"]) as pool: + _eval_sequence = partial( + eval_sequence, + dataset=dataset, + tracker=tracker, + class_list=class_list, + metrics_list=metrics_list, + metric_names=metric_names, + ) + results = pool.map(_eval_sequence, seq_list) + res = dict(zip(seq_list, results)) + else: + res = {} + if show_progressbar and TQDM_IMPORTED: + seq_list_sorted = sorted(seq_list) + for curr_seq in tqdm.tqdm(seq_list_sorted): + res[curr_seq] = eval_sequence( + curr_seq, + dataset, + tracker, + class_list, + metrics_list, + metric_names, + ) + else: + for curr_seq in sorted(seq_list): + res[curr_seq] = eval_sequence( + curr_seq, + dataset, + tracker, + class_list, + metrics_list, + metric_names, + ) + + # Combine results over all sequences and then over all classes + res, combined_cls_keys = self._combine_results( + res, metrics_list, metric_names, dataset, "COMBINED_SEQ" + ) + + if np.all( + ["tags" in annot for annot in dataset.gt_data["annotations"]] + ): + # Combine results over the challenging sequences and then over all classes + # currently only support "tracking_challenging_pair" + res, _ = self._combine_results( + res, + metrics_list, + metric_names, + dataset, + "COMBINED_SEQ_CHALLENGING", + "tracking_challenging_pair", + ) + + # Print and output results in various formats + if config["TIME_PROGRESS"]: + print( + "\nAll sequences for %s finished in %.2f seconds" + % (tracker, time.time() - time_start) + ) + + self._summarize_results( + res, + tracker, + metrics_list, + metric_names, + dataset, + "COMBINED_SEQ", + combined_cls_keys, + ) + if "COMBINED_SEQ_CHALLENGING" in res: + self._summarize_results( + res, + tracker, + metrics_list, + metric_names, + dataset, + "COMBINED_SEQ_CHALLENGING", + combined_cls_keys, + ) + + # Output for returning from function + output_res[dataset_name][tracker] = res + output_msg[dataset_name][tracker] = "Success" + + except Exception as err: + output_res[dataset_name][tracker] = None + if type(err) == TrackEvalException: + output_msg[dataset_name][tracker] = str(err) + else: + output_msg[dataset_name][tracker] = "Unknown error occurred." + print("Tracker %s was unable to be evaluated." % tracker) + print(err) + traceback.print_exc() + if config["LOG_ON_ERROR"] is not None: + with open(config["LOG_ON_ERROR"], "a") as f: + print(dataset_name, file=f) + print(tracker, file=f) + print(traceback.format_exc(), file=f) + print("\n\n\n", file=f) + if config["BREAK_ON_ERROR"]: + raise err + elif config["RETURN_ON_ERROR"]: + return output_res, output_msg + + return output_res, output_msg + + +@_timing.time +def eval_sequence(seq, dataset, tracker, class_list, metrics_list, metric_names): + """Function for evaluating a single sequence""" + + raw_data = dataset.get_raw_seq_data(tracker, seq) + seq_res = {} + for cls in class_list: + seq_res[cls] = {} + data = dataset.get_preprocessed_seq_data(raw_data, cls) + for metric, met_name in zip(metrics_list, metric_names): + seq_res[cls][met_name] = metric.eval_sequence(data) + return seq_res diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/__init__.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/__init__.py new file mode 100644 index 0000000..6c84342 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/__init__.py @@ -0,0 +1,4 @@ +# flake8: noqa + +from .count import Count +from .hota import HOTA diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/_base_metric.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/_base_metric.py new file mode 100644 index 0000000..afbfaca --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/_base_metric.py @@ -0,0 +1,145 @@ +# flake8: noqa + +from abc import ABC, abstractmethod + +import numpy as np + +from .. import _timing +from ..utils import TrackEvalException + + +class _BaseMetric(ABC): + @abstractmethod + def __init__(self): + self.plottable = False + self.integer_fields = [] + self.float_fields = [] + self.array_labels = [] + self.integer_array_fields = [] + self.float_array_fields = [] + self.fields = [] + self.summary_fields = [] + self.registered = False + + ##################################################################### + # Abstract functions for subclasses to implement + + @_timing.time + @abstractmethod + def eval_sequence(self, data): ... + + @abstractmethod + def combine_sequences(self, all_res): ... + + @abstractmethod + def combine_classes_class_averaged(self, all_res, ignore_empty_classes=False): ... + + @abstractmethod + def combine_classes_det_averaged(self, all_res): ... + + def plot_single_tracker_results(self, all_res, tracker, output_folder, cls): + """Plot results of metrics, only valid for metrics with self.plottable""" + if self.plottable: + raise NotImplementedError( + "plot_results is not implemented for metric %s" % self.get_name() + ) + else: + pass + + ##################################################################### + # Helper functions which are useful for all metrics: + + @classmethod + def get_name(cls): + return cls.__name__ + + @staticmethod + def _combine_sum(all_res, field): + """Combine sequence results via sum""" + return sum([all_res[k][field] for k in all_res.keys()]) + + @staticmethod + def _combine_weighted_av(all_res, field, comb_res, weight_field): + """Combine sequence results via weighted average""" + return sum( + [all_res[k][field] * all_res[k][weight_field] for k in all_res.keys()] + ) / np.maximum(1.0, comb_res[weight_field]) + + def print_table( + self, table_res, tracker, cls, res_field="COMBINED_SEQ", output_lable="COMBINED" + ): + """Prints table of results for all sequences""" + print("") + metric_name = self.get_name() + self._row_print( + [metric_name + ": " + tracker + "-" + cls] + self.summary_fields + ) + for seq, results in sorted(table_res.items()): + if seq.startswith("COMBINED_SEQ"): + continue + summary_res = self._summary_row(results) + self._row_print([seq] + summary_res) + summary_res = self._summary_row(table_res[res_field]) + self._row_print([output_lable] + summary_res) + + def _summary_row(self, results_): + vals = [] + for h in self.summary_fields: + if h in self.float_array_fields: + vals.append("{0:1.5g}".format(100 * np.mean(results_[h]))) + elif h in self.float_fields: + vals.append("{0:1.5g}".format(100 * float(results_[h]))) + elif h in self.integer_fields: + vals.append("{0:d}".format(int(results_[h]))) + else: + raise NotImplementedError( + "Summary function not implemented for this field type." + ) + return vals + + @staticmethod + def _row_print(*argv): + """Prints results in an evenly spaced rows, with more space in first row""" + if len(argv) == 1: + argv = argv[0] + to_print = "%-35s" % argv[0] + for v in argv[1:]: + to_print += "%-10s" % str(v) + print(to_print) + + def summary_results(self, table_res): + """Returns a simple summary of final results for a tracker""" + return dict( + zip(self.summary_fields, self._summary_row(table_res["COMBINED_SEQ"])) + ) + + def detailed_results(self, table_res): + """Returns detailed final results for a tracker""" + # Get detailed field information + detailed_fields = self.float_fields + self.integer_fields + for h in self.float_array_fields + self.integer_array_fields: + for alpha in [int(100 * x) for x in self.array_labels]: + detailed_fields.append(h + "___" + str(alpha)) + detailed_fields.append(h + "___AUC") + + # Get detailed results + detailed_results = {} + for seq, res in table_res.items(): + detailed_row = self._detailed_row(res) + if len(detailed_row) != len(detailed_fields): + raise TrackEvalException( + "Field names and data have different sizes (%i and %i)" + % (len(detailed_row), len(detailed_fields)) + ) + detailed_results[seq] = dict(zip(detailed_fields, detailed_row)) + return detailed_results + + def _detailed_row(self, res): + detailed_row = [] + for h in self.float_fields + self.integer_fields: + detailed_row.append(res[h]) + for h in self.float_array_fields + self.integer_array_fields: + for i, alpha in enumerate([int(100 * x) for x in self.array_labels]): + detailed_row.append(res[h][i]) + detailed_row.append(np.mean(res[h])) + return detailed_row diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/count.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/count.py new file mode 100644 index 0000000..a372605 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/count.py @@ -0,0 +1,48 @@ +# flake8: noqa + +from .. import _timing +from ._base_metric import _BaseMetric + + +class Count(_BaseMetric): + """Class which simply counts the number of tracker and gt detections and ids.""" + + def __init__(self, config=None): + super().__init__() + self.integer_fields = ["Dets", "GT_Dets", "IDs", "GT_IDs"] + self.fields = self.integer_fields + self.summary_fields = self.fields + + @_timing.time + def eval_sequence(self, data): + """Returns counts for one sequence""" + # Get results + res = { + "Dets": data["num_tracker_dets"], + "GT_Dets": data["num_gt_dets"], + "IDs": data["num_tracker_ids"], + "GT_IDs": data["num_gt_ids"], + "Frames": data["num_timesteps"], + } + return res + + def combine_sequences(self, all_res): + """Combines metrics across all sequences""" + res = {} + for field in self.integer_fields: + res[field] = self._combine_sum(all_res, field) + return res + + def combine_classes_class_averaged(self, all_res, ignore_empty_classes=None): + """Combines metrics across all classes by averaging over the class values""" + res = {} + for field in self.integer_fields: + res[field] = self._combine_sum(all_res, field) + return res + + def combine_classes_det_averaged(self, all_res): + """Combines metrics across all classes by averaging over the detection values""" + res = {} + for field in self.integer_fields: + res[field] = self._combine_sum(all_res, field) + return res diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/hota.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/hota.py new file mode 100644 index 0000000..4cd9501 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/metrics/hota.py @@ -0,0 +1,291 @@ +# flake8: noqa + +import os + +import numpy as np +from scipy.optimize import linear_sum_assignment + +from .. import _timing +from ._base_metric import _BaseMetric + + +class HOTA(_BaseMetric): + """Class which implements the HOTA metrics. + See: https://link.springer.com/article/10.1007/s11263-020-01375-2 + """ + + def __init__(self, config=None): + super().__init__() + self.plottable = True + self.array_labels = np.arange(0.05, 0.99, 0.05) + self.integer_array_fields = ["HOTA_TP", "HOTA_FN", "HOTA_FP"] + self.float_array_fields = [ + "HOTA", + "DetA", + "AssA", + "DetRe", + "DetPr", + "AssRe", + "AssPr", + "LocA", + "OWTA", + ] + self.float_fields = ["HOTA(0)", "LocA(0)", "HOTALocA(0)"] + self.fields = ( + self.float_array_fields + self.integer_array_fields + self.float_fields + ) + self.summary_fields = self.float_array_fields + self.float_fields + + @_timing.time + def eval_sequence(self, data): + """Calculates the HOTA metrics for one sequence""" + + # Initialise results + res = {} + for field in self.float_array_fields + self.integer_array_fields: + res[field] = np.zeros((len(self.array_labels)), dtype=float) + for field in self.float_fields: + res[field] = 0 + + # Return result quickly if tracker or gt sequence is empty + if data["num_tracker_dets"] == 0: + res["HOTA_FN"] = data["num_gt_dets"] * np.ones( + (len(self.array_labels)), dtype=float + ) + res["LocA"] = np.ones((len(self.array_labels)), dtype=float) + res["LocA(0)"] = 1.0 + return res + if data["num_gt_dets"] == 0: + res["HOTA_FP"] = data["num_tracker_dets"] * np.ones( + (len(self.array_labels)), dtype=float + ) + res["LocA"] = np.ones((len(self.array_labels)), dtype=float) + res["LocA(0)"] = 1.0 + return res + + # Variables counting global association + potential_matches_count = np.zeros( + (data["num_gt_ids"], data["num_tracker_ids"]) + ) + gt_id_count = np.zeros((data["num_gt_ids"], 1)) + tracker_id_count = np.zeros((1, data["num_tracker_ids"])) + + # First loop through each timestep and accumulate global track information. + for t, (gt_ids_t, tracker_ids_t) in enumerate( + zip(data["gt_ids"], data["tracker_ids"]) + ): + # Count the potential matches between ids in each timestep + # These are normalised, weighted by the match similarity. + similarity = data["similarity_scores"][t] + sim_iou_denom = ( + similarity.sum(0)[np.newaxis, :] + + similarity.sum(1)[:, np.newaxis] + - similarity + ) + sim_iou = np.zeros_like(similarity) + sim_iou_mask = sim_iou_denom > 0 + np.finfo("float").eps + sim_iou[sim_iou_mask] = ( + similarity[sim_iou_mask] / sim_iou_denom[sim_iou_mask] + ) + potential_matches_count[ + gt_ids_t[:, np.newaxis], tracker_ids_t[np.newaxis, :] + ] += sim_iou + + # Calculate the total number of dets for each gt_id and tracker_id. + gt_id_count[gt_ids_t] += 1 + tracker_id_count[0, tracker_ids_t] += 1 + + # Calculate overall jaccard alignment score (before unique matching) between IDs + global_alignment_score = potential_matches_count / ( + gt_id_count + tracker_id_count - potential_matches_count + ) + matches_counts = [ + np.zeros_like(potential_matches_count) for _ in self.array_labels + ] + + # Calculate scores for each timestep + for t, (gt_ids_t, tracker_ids_t) in enumerate( + zip(data["gt_ids"], data["tracker_ids"]) + ): + # Deal with the case that there are no gt_det/tracker_det in a timestep. + if len(gt_ids_t) == 0: + for a, alpha in enumerate(self.array_labels): + res["HOTA_FP"][a] += len(tracker_ids_t) + continue + if len(tracker_ids_t) == 0: + for a, alpha in enumerate(self.array_labels): + res["HOTA_FN"][a] += len(gt_ids_t) + continue + + # Get matching scores between pairs of dets for optimizing HOTA + similarity = data["similarity_scores"][t] + score_mat = ( + global_alignment_score[ + gt_ids_t[:, np.newaxis], tracker_ids_t[np.newaxis, :] + ] + * similarity + ) + + # Hungarian algorithm to find best matches + match_rows, match_cols = linear_sum_assignment(-score_mat) + + # Calculate and accumulate basic statistics + for a, alpha in enumerate(self.array_labels): + actually_matched_mask = ( + similarity[match_rows, match_cols] >= alpha - np.finfo("float").eps + ) + alpha_match_rows = match_rows[actually_matched_mask] + alpha_match_cols = match_cols[actually_matched_mask] + num_matches = len(alpha_match_rows) + res["HOTA_TP"][a] += num_matches + res["HOTA_FN"][a] += len(gt_ids_t) - num_matches + res["HOTA_FP"][a] += len(tracker_ids_t) - num_matches + if num_matches > 0: + res["LocA"][a] += sum( + similarity[alpha_match_rows, alpha_match_cols] + ) + matches_counts[a][ + gt_ids_t[alpha_match_rows], tracker_ids_t[alpha_match_cols] + ] += 1 + + # Calculate association scores (AssA, AssRe, AssPr) for the alpha value. + # First calculate scores per gt_id/tracker_id combo and then average over the number of detections. + for a, alpha in enumerate(self.array_labels): + matches_count = matches_counts[a] + ass_a = matches_count / np.maximum( + 1, gt_id_count + tracker_id_count - matches_count + ) + res["AssA"][a] = np.sum(matches_count * ass_a) / np.maximum( + 1, res["HOTA_TP"][a] + ) + ass_re = matches_count / np.maximum(1, gt_id_count) + res["AssRe"][a] = np.sum(matches_count * ass_re) / np.maximum( + 1, res["HOTA_TP"][a] + ) + ass_pr = matches_count / np.maximum(1, tracker_id_count) + res["AssPr"][a] = np.sum(matches_count * ass_pr) / np.maximum( + 1, res["HOTA_TP"][a] + ) + + # Calculate final scores + res["LocA"] = np.maximum(1e-10, res["LocA"]) / np.maximum(1e-10, res["HOTA_TP"]) + res = self._compute_final_fields(res) + return res + + def combine_sequences(self, all_res): + """Combines metrics across all sequences""" + res = {} + for field in self.integer_array_fields: + res[field] = self._combine_sum(all_res, field) + for field in ["AssRe", "AssPr", "AssA"]: + res[field] = self._combine_weighted_av( + all_res, field, res, weight_field="HOTA_TP" + ) + loca_weighted_sum = sum( + [all_res[k]["LocA"] * all_res[k]["HOTA_TP"] for k in all_res.keys()] + ) + res["LocA"] = np.maximum(1e-10, loca_weighted_sum) / np.maximum( + 1e-10, res["HOTA_TP"] + ) + res = self._compute_final_fields(res) + return res + + def combine_classes_class_averaged(self, all_res, ignore_empty_classes=False): + """Combines metrics across all classes by averaging over the class values. + If 'ignore_empty_classes' is True, then it only sums over classes with at least one gt or predicted detection. + """ + res = {} + for field in self.integer_array_fields: + if ignore_empty_classes: + res[field] = self._combine_sum( + { + k: v + for k, v in all_res.items() + if ( + v["HOTA_TP"] + v["HOTA_FN"] + v["HOTA_FP"] + > 0 + np.finfo("float").eps + ).any() + }, + field, + ) + else: + res[field] = self._combine_sum( + {k: v for k, v in all_res.items()}, field + ) + + for field in self.float_fields + self.float_array_fields: + if ignore_empty_classes: + res[field] = np.mean( + [ + v[field] + for v in all_res.values() + if ( + v["HOTA_TP"] + v["HOTA_FN"] + v["HOTA_FP"] + > 0 + np.finfo("float").eps + ).any() + ], + axis=0, + ) + else: + res[field] = np.mean([v[field] for v in all_res.values()], axis=0) + return res + + def combine_classes_det_averaged(self, all_res): + """Combines metrics across all classes by averaging over the detection values""" + res = {} + for field in self.integer_array_fields: + res[field] = self._combine_sum(all_res, field) + for field in ["AssRe", "AssPr", "AssA"]: + res[field] = self._combine_weighted_av( + all_res, field, res, weight_field="HOTA_TP" + ) + loca_weighted_sum = sum( + [all_res[k]["LocA"] * all_res[k]["HOTA_TP"] for k in all_res.keys()] + ) + res["LocA"] = np.maximum(1e-10, loca_weighted_sum) / np.maximum( + 1e-10, res["HOTA_TP"] + ) + res = self._compute_final_fields(res) + return res + + @staticmethod + def _compute_final_fields(res): + """Calculate sub-metric ('field') values which only depend on other sub-metric values. + This function is used both for both per-sequence calculation, and in combining values across sequences. + """ + res["DetRe"] = res["HOTA_TP"] / np.maximum(1, res["HOTA_TP"] + res["HOTA_FN"]) + res["DetPr"] = res["HOTA_TP"] / np.maximum(1, res["HOTA_TP"] + res["HOTA_FP"]) + res["DetA"] = res["HOTA_TP"] / np.maximum( + 1, res["HOTA_TP"] + res["HOTA_FN"] + res["HOTA_FP"] + ) + res["HOTA"] = np.sqrt(res["DetA"] * res["AssA"]) + res["OWTA"] = np.sqrt(res["DetRe"] * res["AssA"]) + + res["HOTA(0)"] = res["HOTA"][0] + res["LocA(0)"] = res["LocA"][0] + res["HOTALocA(0)"] = res["HOTA(0)"] * res["LocA(0)"] + return res + + def plot_single_tracker_results(self, table_res, tracker, cls, output_folder): + """Create plot of results""" + + # Only loaded when run to reduce minimum requirements + from matplotlib import pyplot as plt + + res = table_res["COMBINED_SEQ"] + styles_to_plot = ["r", "b", "g", "b--", "b:", "g--", "g:", "m"] + for name, style in zip(self.float_array_fields, styles_to_plot): + plt.plot(self.array_labels, res[name], style) + plt.xlabel("alpha") + plt.ylabel("score") + plt.title(tracker + " - " + cls) + plt.axis([0, 1, 0, 1]) + legend = [] + for name in self.float_array_fields: + legend += [name + " (" + str(np.round(np.mean(res[name]), 2)) + ")"] + plt.legend(legend, loc="lower left") + out_file = os.path.join(output_folder, cls + "_plot.pdf") + os.makedirs(os.path.dirname(out_file), exist_ok=True) + plt.savefig(out_file) + plt.savefig(out_file.replace(".pdf", ".png")) + plt.clf() diff --git a/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/utils.py b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/utils.py new file mode 100644 index 0000000..8cdf77e --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/hota_eval_toolkit/trackeval/utils.py @@ -0,0 +1,195 @@ +# flake8: noqa + +import argparse +import csv +import os +from collections import OrderedDict + + +def init_config(config, default_config, name=None): + """Initialise non-given config values with defaults""" + if config is None: + config = default_config + else: + for k in default_config.keys(): + if k not in config.keys(): + config[k] = default_config[k] + if name and config["PRINT_CONFIG"]: + print("\n%s Config:" % name) + for c in config.keys(): + print("%-20s : %-30s" % (c, config[c])) + return config + + +def update_config(config): + """ + Parse the arguments of a script and updates the config values for a given value if specified in the arguments. + :param config: the config to update + :return: the updated config + """ + parser = argparse.ArgumentParser() + for setting in config.keys(): + if type(config[setting]) == list or type(config[setting]) == type(None): + parser.add_argument("--" + setting, nargs="+") + else: + parser.add_argument("--" + setting) + args = parser.parse_args().__dict__ + for setting in args.keys(): + if args[setting] is not None: + if type(config[setting]) == type(True): + if args[setting] == "True": + x = True + elif args[setting] == "False": + x = False + else: + raise Exception( + "Command line parameter " + setting + "must be True or False" + ) + elif type(config[setting]) == type(1): + x = int(args[setting]) + elif type(args[setting]) == type(None): + x = None + else: + x = args[setting] + config[setting] = x + return config + + +def get_code_path(): + """Get base path where code is""" + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + + +def validate_metrics_list(metrics_list): + """Get names of metric class and ensures they are unique, further checks that the fields within each metric class + do not have overlapping names. + """ + metric_names = [metric.get_name() for metric in metrics_list] + # check metric names are unique + if len(metric_names) != len(set(metric_names)): + raise TrackEvalException( + "Code being run with multiple metrics of the same name" + ) + fields = [] + for m in metrics_list: + fields += m.fields + # check metric fields are unique + if len(fields) != len(set(fields)): + raise TrackEvalException( + "Code being run with multiple metrics with fields of the same name" + ) + return metric_names + + +def write_summary_results(summaries, cls, output_folder): + """Write summary results to file""" + + fields = sum([list(s.keys()) for s in summaries], []) + values = sum([list(s.values()) for s in summaries], []) + + # In order to remain consistent upon new fields being adding, for each of the following fields if they are present + # they will be output in the summary first in the order below. Any further fields will be output in the order each + # metric family is called, and within each family either in the order they were added to the dict (python >= 3.6) or + # randomly (python < 3.6). + default_order = [ + "HOTA", + "DetA", + "AssA", + "DetRe", + "DetPr", + "AssRe", + "AssPr", + "LocA", + "OWTA", + "HOTA(0)", + "LocA(0)", + "HOTALocA(0)", + "MOTA", + "MOTP", + "MODA", + "CLR_Re", + "CLR_Pr", + "MTR", + "PTR", + "MLR", + "CLR_TP", + "CLR_FN", + "CLR_FP", + "IDSW", + "MT", + "PT", + "ML", + "Frag", + "sMOTA", + "IDF1", + "IDR", + "IDP", + "IDTP", + "IDFN", + "IDFP", + "Dets", + "GT_Dets", + "IDs", + "GT_IDs", + ] + default_ordered_dict = OrderedDict( + zip(default_order, [None for _ in default_order]) + ) + for f, v in zip(fields, values): + default_ordered_dict[f] = v + for df in default_order: + if default_ordered_dict[df] is None: + del default_ordered_dict[df] + fields = list(default_ordered_dict.keys()) + values = list(default_ordered_dict.values()) + + out_file = os.path.join(output_folder, cls + "_summary.txt") + os.makedirs(os.path.dirname(out_file), exist_ok=True) + with open(out_file, "w", newline="") as f: + writer = csv.writer(f, delimiter=" ") + writer.writerow(fields) + writer.writerow(values) + + +def write_detailed_results(details, cls, output_folder): + """Write detailed results to file""" + sequences = details[0].keys() + fields = ["seq"] + sum([list(s["COMBINED_SEQ"].keys()) for s in details], []) + out_file = os.path.join(output_folder, cls + "_detailed.csv") + os.makedirs(os.path.dirname(out_file), exist_ok=True) + with open(out_file, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(fields) + for seq in sorted(sequences): + if seq == "COMBINED_SEQ": + continue + writer.writerow([seq] + sum([list(s[seq].values()) for s in details], [])) + writer.writerow( + ["COMBINED"] + sum([list(s["COMBINED_SEQ"].values()) for s in details], []) + ) + + +def load_detail(file): + """Loads detailed data for a tracker.""" + data = {} + with open(file) as f: + for i, row_text in enumerate(f): + row = row_text.replace("\r", "").replace("\n", "").split(",") + if i == 0: + keys = row[1:] + continue + current_values = row[1:] + seq = row[0] + if seq == "COMBINED": + seq = "COMBINED_SEQ" + if (len(current_values) == len(keys)) and seq != "": + data[seq] = {} + for key, value in zip(keys, current_values): + data[seq][key] = float(value) + return data + + +class TrackEvalException(Exception): + """Custom exception for catching expected errors.""" + + ... diff --git a/src/nn/segearth_ov3/sam3/eval/postprocessors.py b/src/nn/segearth_ov3/sam3/eval/postprocessors.py new file mode 100644 index 0000000..973da11 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/postprocessors.py @@ -0,0 +1,648 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Postprocessors class to transform MDETR output according to the downstream task""" + +import dataclasses +import logging +from collections import defaultdict +from typing import Dict, List, Optional + +import numpy as np +import torch +from sam3.model import box_ops +from sam3.model.data_misc import BatchedInferenceMetadata, interpolate +from sam3.train.masks_ops import rle_encode, robust_rle_encode +from torch import nn + + +class PostProcessNullOp(nn.Module): + def __init__(self, **kwargs): + super(PostProcessNullOp).__init__() + pass + + def forward(self, input): + pass + + def process_results(self, **kwargs): + return kwargs["find_stages"] + + +class PostProcessImage(nn.Module): + """This module converts the model's output into the format expected by the coco api""" + + def __init__( + self, + max_dets_per_img: int, + iou_type="bbox", + to_cpu: bool = True, + use_original_ids: bool = False, + use_original_sizes_box: bool = False, + use_original_sizes_mask: bool = False, + convert_mask_to_rle: bool = False, + always_interpolate_masks_on_gpu: bool = True, + use_presence: bool = True, + detection_threshold: float = -1.0, + ) -> None: + super().__init__() + self.max_dets_per_img = max_dets_per_img + self.iou_type = iou_type + self.to_cpu = to_cpu + self.convert_mask_to_rle = convert_mask_to_rle + self.always_interpolate_masks_on_gpu = always_interpolate_masks_on_gpu + + self.use_presence = use_presence + self.detection_threshold = detection_threshold + self.use_original_ids = use_original_ids + self.use_original_sizes_box = use_original_sizes_box + self.use_original_sizes_mask = use_original_sizes_mask + + @torch.no_grad() + def forward( + self, + outputs, + target_sizes_boxes, + target_sizes_masks, + forced_labels=None, + consistent=False, + ret_tensordict: bool = False, # This is experimental + ): + """Perform the computation + Parameters: + outputs: raw outputs of the model + target_sizes_boxes: tensor of dimension [batch_size x 2] containing the size of each images of the batch + For evaluation, this must be the original image size (before any data augmentation) + For visualization, this should be the image size after data augment, but before padding + target_sizes_masks: same but used to resize masks + forced_labels: tensor of dimension [batch_size] containing the label to force for each image of the batch + This is useful when evaluating the model using standard metrics (eg on COCO, LVIS). In that case, + we query the model with every possible class label, so we when we pass the predictions to the evaluator, + we want to make sure that the predicted "class" matches the one that was queried. + consistent: whether all target sizes are equal + ret_tensordict: Experimental argument. If true, return a tensordict.TensorDict instead of a list of dictionaries for easier manipulation. + """ + if ret_tensordict: + assert ( + consistent is True + ), "We don't support returning TensorDict if the outputs have different shapes" # NOTE: It's possible but we don't support it. + assert self.detection_threshold <= 0.0, "TODO: implement?" + try: + from tensordict import TensorDict + except ImportError: + logging.info( + "tensordict is not installed. Install by running `pip install tensordict --no-deps`. Falling back by setting `ret_tensordict=False`" + ) + ret_tensordict = False + + out_bbox = outputs["pred_boxes"] if "pred_boxes" in outputs else None + out_logits = outputs["pred_logits"] + pred_masks = outputs["pred_masks"] if self.iou_type == "segm" else None + out_probs = out_logits.sigmoid() + if self.use_presence: + presence_score = outputs["presence_logit_dec"].sigmoid().unsqueeze(1) + out_probs = out_probs * presence_score + + assert target_sizes_boxes.shape[1] == 2 + assert target_sizes_masks.shape[1] == 2 + batch_size = target_sizes_boxes.shape[0] + + boxes, scores, labels, keep = self._process_boxes_and_labels( + target_sizes_boxes, forced_labels, out_bbox, out_probs + ) + assert boxes is None or len(boxes) == batch_size + out_masks = self._process_masks( + target_sizes_masks, pred_masks, consistent=consistent, keep=keep + ) + del pred_masks + + if boxes is None: + assert out_masks is not None + assert not ret_tensordict, "We don't support returning TensorDict if the output does not contain boxes" + B = len(out_masks) + boxes = [None] * B + scores = [None] * B + labels = [None] * B + + results = { + "scores": scores, + "labels": labels, + "boxes": boxes, + } + if out_masks is not None: + if self.convert_mask_to_rle: + results.update(masks_rle=out_masks) + else: + results.update(masks=out_masks) + + if ret_tensordict: + results = TensorDict(results).auto_batch_size_() + if self.to_cpu: + results = results.cpu() + else: + # Convert a dictonary of lists/tensors to list of dictionaries + results = [ + dict(zip(results.keys(), res_tuple)) + for res_tuple in zip(*results.values()) + ] + + return results + + def _process_masks(self, target_sizes, pred_masks, consistent=True, keep=None): + if pred_masks is None: + return None + if self.always_interpolate_masks_on_gpu: + gpu_device = target_sizes.device + assert gpu_device.type == "cuda" + pred_masks = pred_masks.to(device=gpu_device) + if consistent: + assert keep is None, "TODO: implement?" + # All masks should have the same shape, expected when processing a batch of size 1 + target_size = target_sizes.unique(dim=0) + assert target_size.size(0) == 1, "Expecting all target sizes to be equal" + out_masks = ( + interpolate( + pred_masks, + target_size.squeeze().tolist(), + mode="bilinear", + align_corners=False, + ).sigmoid() + > 0.5 + ) + if self.convert_mask_to_rle: + raise RuntimeError("TODO: implement?") + if self.to_cpu: + out_masks = out_masks.cpu() + else: + out_masks = [[]] * len(pred_masks) + + assert keep is None or len(keep) == len(pred_masks) + for i, mask in enumerate(pred_masks): + h, w = target_sizes[i] + if keep is not None: + mask = mask[keep[i]] + # Uses the gpu version fist, moves masks to cpu if it fails""" + try: + interpolated = ( + interpolate( + mask.unsqueeze(1), + (h, w), + mode="bilinear", + align_corners=False, + ).sigmoid() + > 0.5 + ) + except Exception as e: + logging.info("Issue found, reverting to CPU mode!") + mask_device = mask.device + mask = mask.cpu() + interpolated = ( + interpolate( + mask.unsqueeze(1), + (h, w), + mode="bilinear", + align_corners=False, + ).sigmoid() + > 0.5 + ) + interpolated = interpolated.to(mask_device) + + if self.convert_mask_to_rle: + out_masks[i] = robust_rle_encode(interpolated.squeeze(1)) + else: + out_masks[i] = interpolated + if self.to_cpu: + out_masks[i] = out_masks[i].cpu() + + return out_masks + + def _process_boxes_and_labels( + self, target_sizes, forced_labels, out_bbox, out_probs + ): + if out_bbox is None: + return None, None, None, None + assert len(out_probs) == len(target_sizes) + if self.to_cpu: + out_probs = out_probs.cpu() + scores, labels = out_probs.max(-1) + if forced_labels is None: + labels = torch.ones_like(labels) + else: + labels = forced_labels[:, None].expand_as(labels) + + # convert to [x0, y0, x1, y1] format + boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) + + img_h, img_w = target_sizes.unbind(1) + scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) + boxes = boxes * scale_fct[:, None, :] + + if self.to_cpu: + boxes = boxes.cpu() + + keep = None + if self.detection_threshold > 0: + # Filter out the boxes with scores below the detection threshold + keep = scores > self.detection_threshold + assert len(keep) == len(boxes) == len(scores) == len(labels) + + boxes = [b[k.to(b.device)] for b, k in zip(boxes, keep)] + scores = [s[k.to(s.device)] for s, k in zip(scores, keep)] + labels = [l[k.to(l.device)] for l, k in zip(labels, keep)] + + return boxes, scores, labels, keep + + def process_results( + self, find_stages, find_metadatas: List[BatchedInferenceMetadata], **kwargs + ): + if find_stages.loss_stages is not None: + find_metadatas = [find_metadatas[i] for i in find_stages.loss_stages] + assert len(find_stages) == len(find_metadatas) + results = {} + for outputs, meta in zip(find_stages, find_metadatas): + img_size_for_boxes = ( + meta.original_size + if self.use_original_sizes_box + else torch.ones_like(meta.original_size) + ) + img_size_for_masks = ( + meta.original_size + if self.use_original_sizes_mask + else torch.ones_like(meta.original_size) + ) + detection_results = self( + outputs, + img_size_for_boxes, + img_size_for_masks, + forced_labels=( + meta.original_category_id if self.use_original_ids else None + ), + ) + ids = ( + meta.original_image_id if self.use_original_ids else meta.coco_image_id + ) + assert len(detection_results) == len(ids) + for img_id, result in zip(ids, detection_results): + if img_id.item() not in results: + results[img_id.item()] = result + else: + assert set(results[img_id.item()].keys()) == set(result.keys()) + for k in result.keys(): + if isinstance(result[k], torch.Tensor): + results[img_id.item()][k] = torch.cat( + [results[img_id.item()][k], result[k]], dim=0 + ) + elif isinstance(result[k], list): + results[img_id.item()][k] += result[k] + else: + raise NotImplementedError( + f"Unexpected type {type(result[k])} in result." + ) + # Prune the results to the max number of detections per image. + for img_id, result in results.items(): + if ( + self.max_dets_per_img > 0 + and len(result["scores"]) > self.max_dets_per_img + ): + _, topk_indexes = torch.topk( + result["scores"], self.max_dets_per_img, dim=0 + ) + if self.to_cpu: + topk_indexes = topk_indexes.cpu() + for k in result.keys(): + if isinstance(results[img_id][k], list): + results[img_id][k] = [ + results[img_id][k][i] for i in topk_indexes.tolist() + ] + else: + results[img_id][k] = results[img_id][k].to(topk_indexes.device)[ + topk_indexes + ] + + return results + + +class PostProcessAPIVideo(PostProcessImage): + """This module converts the video model's output into the format expected by the YT-VIS api""" + + def __init__( + self, + *args, + to_cpu: bool = True, + convert_mask_to_rle: bool = False, + always_interpolate_masks_on_gpu: bool = True, + prob_thresh: float = 0.5, + use_presence: bool = False, + **kwargs, + ): + super().__init__( + *args, + # Here we always set `convert_mask_to_rle=False` in the base `PostProcessAPI` class + # (so that its `_process_masks` won't return a list of RLEs). If we want to return + # RLEs for video masklets, we handle it in this `PostProcessAPIVideo` class instead. + convert_mask_to_rle=False, + # Here we always set `to_cpu=False` in the base `PostProcessAPI` class (so that + # the interpolated masks won't be automatically moved back to CPU). We will handle + # it in this `PostProcessAPIVideo` class instead. + always_interpolate_masks_on_gpu=always_interpolate_masks_on_gpu, + use_presence=use_presence, + **kwargs, + ) + # Expected keys in the output dict to postprocess + self.EXPECTED_KEYS = [ + "pred_logits", + "pred_boxes", + "pred_masks", + ] + # Whether to post-process video masklets (under packed representation) into RLE format + self.convert_mask_to_rle_for_video = convert_mask_to_rle + self.to_cpu_for_video = to_cpu + self.prob_thresh = prob_thresh + + def process_results( + self, find_stages, find_metadatas: List[BatchedInferenceMetadata], **kwargs + ): + """ + Tracking Postprocessor for SAM 3 video model. + This function takes in the output of the SAM 3 video model and processes it to extract all the tracklet predictions. + Args: + find_stages: A list of tensors representing the output of the SAM 3 video model. + find_metadatas: A list of BatchedInferenceMetadata objects containing metadata about each frame. + **kwargs: Additional keyword arguments. + Returns: + A dictionary of predcitions with video_id as key. + """ + + # Import tensordict here to avoid global dependency. + try: + from tensordict import TensorDict + except ImportError as e: + logging.error( + "tensordict is not installed, please install by running `pip install tensordict --no-deps`" + ) + raise e + # Notes and assumptions: + # 1- This postprocessor assumes results only for a single video. + # 2- There are N stage outputs corresponding to N video frames + # 3- Each stage outputs contains PxQ preds, where P is number of prompts and Q is number of object queries. The output should also contain the tracking object ids corresponding to each object query. + # 4- The tracking object id has a default value of -1, indicating that the object query is not tracking any object in the frame, and hence its predictions can be ingored for a given frame. + # 5- Some objects may be tracked in a subset of frames only. So, we first extract the predictions in a packed representation (for efficient postprocessing -- specially memory) + # and then we convert the packed representation into a padded one, where we zero pad boxes/masks for objects that are not tracked in some frames. + # 6- We refer to objects by an object id, which is a tuple (prompt_idx, obj_id) + + assert len(find_stages) > 0, "There is nothing to postprocess?" + PROMPT_AXIS, OBJ_QUERY_AXIS = (0, 1) + NO_OBJ_ID = -1 + # Maps object ID -> [indices in packed tensor] + tracked_objects_packed_idx = defaultdict(list) + # Maps object ID -> [indices in padded tensor (abs frame index)] + tracked_objects_frame_idx = defaultdict(list) + total_num_preds = 0 + # This will hold the packed representation of predictions. + vid_preds_packed: List[TensorDict] = [] + vid_masklets_rle_packed: List[Optional[Dict]] = [] + video_id = -1 # We assume single video postprocessing, this ID should be unique in the datapoint. + + for frame_idx, (frame_outs, meta) in enumerate( + zip(find_stages, find_metadatas) + ): + # only store keys we need to extract the results + frame_outs_td = TensorDict( + {k: frame_outs[k] for k in self.EXPECTED_KEYS} + ).auto_batch_size_() # Shape is [P,Q,...] + meta_td = TensorDict( + dataclasses.asdict(meta) + ).auto_batch_size_() # Shape is [P,...] + unique_vid_id = meta.original_image_id.unique() + assert unique_vid_id.size(0) == 1 + if video_id == -1: + video_id = unique_vid_id.item() + else: + assert ( + video_id == unique_vid_id.item() + ), "We can only postprocess one video per datapoint" + # keeping track of which objects appear in the current frame + obj_ids_per_frame = frame_outs["pred_object_ids"] + assert obj_ids_per_frame.size(-1) == frame_outs["pred_logits"].size(-2) + if self.prob_thresh is not None: + # only keep the predictions on this frame with probability above the threshold + # (remove those predictions during the keep-alive period of a tracking query, + # where its "pred_object_ids" is still the tracked object ID rather than -1) + pred_probs = frame_outs["pred_logits"].sigmoid().squeeze(-1) + obj_ids_per_frame = torch.where( + pred_probs >= self.prob_thresh, obj_ids_per_frame, NO_OBJ_ID + ) + tracked_obj_ids_idx = torch.where(obj_ids_per_frame != NO_OBJ_ID) + # Object id is a tuple of (prompt_idx, obj_id). This is because the model can assign same obj_id for two different prompts. + tracked_obj_ids = [ + (p_id.item(), obj_ids_per_frame[p_id, q_id].item()) + for p_id, q_id in zip( + tracked_obj_ids_idx[PROMPT_AXIS], + tracked_obj_ids_idx[OBJ_QUERY_AXIS], + ) + ] + if len(tracked_obj_ids) == 0: + continue + # For each object, we keep track of the packed and padded (frame index) indices + for oid in tracked_obj_ids: + tracked_objects_packed_idx[oid].append(total_num_preds) + tracked_objects_frame_idx[oid].append(frame_idx) + total_num_preds += 1 + + # Since we have P*Q masks per frame, mask interpolation is the GPU memory bottleneck or time bottleneck in case of cpu processing. + # Instead, we first extract results only for tracked objects, reducing the number of masks to K = sum_i(tracked_objs_per_ith_prompt), hopefully <<< P*Q + tracked_objs_outs_td = frame_outs_td[ + tracked_obj_ids_idx + ] # [P,Q,...] --> [K,...] + meta_td = meta_td[tracked_obj_ids_idx[PROMPT_AXIS].cpu()] + if self.always_interpolate_masks_on_gpu: + gpu_device = meta_td["original_size"].device + assert gpu_device.type == "cuda" + tracked_objs_outs_td = tracked_objs_outs_td.to(device=gpu_device) + frame_results_td = self( + tracked_objs_outs_td.unsqueeze(1), + ( + meta_td["original_size"] + if self.use_original_sizes + else torch.ones_like(meta_td["original_size"]) + ), + forced_labels=( + meta_td["original_category_id"] if self.use_original_ids else None + ), + consistent=True, + ret_tensordict=True, + ).squeeze(1) + del tracked_objs_outs_td + + # Optionally, remove "masks" from output tensor dict and directly encode them + # to RLE format under packed representations + if self.convert_mask_to_rle_for_video: + interpolated_binary_masks = frame_results_td.pop("masks") + rle_list = rle_encode(interpolated_binary_masks, return_areas=True) + vid_masklets_rle_packed.extend(rle_list) + # Optionally, move output TensorDict to CPU (do this after RLE encoding step above) + if self.to_cpu_for_video: + frame_results_td = frame_results_td.cpu() + vid_preds_packed.append(frame_results_td) + + if len(vid_preds_packed) == 0: + logging.debug(f"Video {video_id} has no predictions") + return {video_id: []} + + vid_preds_packed = torch.cat(vid_preds_packed, dim=0) + ############### Construct a padded representation of the predictions ############### + num_preds = len(tracked_objects_packed_idx) + num_frames = len(find_stages) + # We zero pad any missing prediction + # NOTE: here, we also have padded tensors for "scores" and "labels", but we overwrite them later. + padded_frames_results = TensorDict( + { + k: torch.zeros( + num_preds, num_frames, *v.shape[1:], device=v.device, dtype=v.dtype + ) + for k, v in vid_preds_packed.items() + }, + batch_size=[ + num_preds, + num_frames, + ], + ) + padded_frames_results["scores"][...] = -1e8 # a very low score for empty object + # Track scores and labels of each pred tracklet, only for frames where the model was able to track that object + tracklet_scores = [] + tracklet_labels = [] + # Optionally, fill the list of RLEs for masklets + # note: only frames with actual predicted masks (in packed format) will be + # filled with RLEs; the rest will remains None in results["masks_rle"] + if self.convert_mask_to_rle_for_video: + vid_masklets_rle_padded = [[None] * num_frames for _ in range(num_preds)] + for o_idx, oid in enumerate(tracked_objects_packed_idx): + oid2packed_idx = tracked_objects_packed_idx[oid] + oid2padded_idx = tracked_objects_frame_idx[oid] + obj_packed_results = vid_preds_packed[oid2packed_idx] + padded_frames_results[o_idx][oid2padded_idx] = obj_packed_results + if self.convert_mask_to_rle_for_video: + for packed_idx, padded_idx in zip(oid2packed_idx, oid2padded_idx): + vid_masklets_rle_padded[o_idx][padded_idx] = ( + vid_masklets_rle_packed[packed_idx] + ) + # NOTE: We need a single confidence score per tracklet for the mAP metric. + # We use the average confidence score across time. (How does this impact AP?) + tracklet_scores.append(obj_packed_results["scores"].mean()) + # We also need to have a unique category Id per tracklet. + # This is not a problem for phrase AP, however, for mAP we do majority voting across time. + tracklet_labels.append(obj_packed_results["labels"].mode()[0]) + + results = padded_frames_results.to_dict() + results["scores"] = torch.stack(tracklet_scores, dim=0) + results["labels"] = torch.stack(tracklet_labels, dim=0) + if self.convert_mask_to_rle_for_video: + results["masks_rle"] = vid_masklets_rle_padded + # we keep the frame-level scores since it's needed by some evaluation scripts + results["per_frame_scores"] = padded_frames_results["scores"] + + return {video_id: results} + + +class PostProcessTracking(PostProcessImage): + """This module converts the model's output into the format expected by the coco api""" + + def __init__( + self, + max_dets_per_img: int, + iou_type="bbox", + force_single_mask: bool = False, + **kwargs, + ) -> None: + super().__init__(max_dets_per_img=max_dets_per_img, iou_type=iou_type, **kwargs) + self.force_single_mask = force_single_mask + + def process_results( + self, find_stages, find_metadatas: BatchedInferenceMetadata, **kwargs + ): + assert len(find_stages) == len(find_metadatas) + results = {} + for outputs, meta in zip(find_stages, find_metadatas): + if self.force_single_mask: + scores, labels = outputs["pred_logits"].max(-1) + m = [] + for i in range(len(outputs["pred_masks"])): + score, idx = scores[i].max(0) + m.append(outputs["pred_masks"][i][idx]) + outputs["pred_masks"] = torch.stack(m, 0).unsqueeze(1) + detection_results = self(outputs, meta.original_size, consistent=False) + assert len(detection_results) == len(meta.coco_image_id) + results.update( + { + (media_id.item(), object_id.item(), frame_index.item()): result + for media_id, object_id, frame_index, result in zip( + meta.original_image_id, + meta.object_id, + meta.frame_index, + detection_results, + ) + } + ) + return results + + +class PostProcessCounting(nn.Module): + """This module converts the model's output to be evaluated for counting tasks""" + + def __init__( + self, + use_original_ids: bool = False, + threshold: float = 0.5, + use_presence: bool = False, + ) -> None: + """ + Args: + use_original_ids: whether to use the original image ids or the coco ids + threshold: threshold for counting (values above this are counted) + """ + super().__init__() + self.use_original_ids = use_original_ids + self.threshold = threshold + self.use_presence = use_presence + + def forward(self, outputs, target_sizes): + """Perform the computation + Parameters: + outputs: raw outputs of the model + target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch + """ + # Extract scores from model outputs and apply sigmoid + scores = torch.sigmoid(outputs["pred_logits"]).squeeze(-1) # [B, N] + if self.use_presence: + presence_score = outputs["presence_logit_dec"].sigmoid() + if presence_score.ndim == 1: + presence_score = presence_score.unsqueeze(1) # [B, 1] + scores = scores * presence_score # [B, N] + + # Calculate counts by summing values above threshold + counts = (scores > self.threshold).float().sum(dim=1) + + assert len(counts) == len(target_sizes) + results = [] + for count in counts: + results.append({"count": count.item()}) + + return results + + @torch.no_grad() + def process_results( + self, find_stages, find_metadatas: List[BatchedInferenceMetadata], **kwargs + ): + assert len(find_stages) == len(find_metadatas) + results = {} + for outputs, meta in zip(find_stages, find_metadatas): + detection_results = self( + outputs, + meta.original_size, + ) + ids = ( + meta.original_image_id if self.use_original_ids else meta.coco_image_id + ) + assert len(detection_results) == len(ids) + for img_id, result in zip(ids, detection_results): + results[img_id.item()] = result + + return results diff --git a/src/nn/segearth_ov3/sam3/eval/saco_veval_eval.py b/src/nn/segearth_ov3/sam3/eval/saco_veval_eval.py new file mode 100644 index 0000000..4f0ed2b --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/saco_veval_eval.py @@ -0,0 +1,155 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import argparse +import json +import os +from collections import defaultdict + +from iopath.common.file_io import g_pathmgr +from sam3.eval.saco_veval_evaluators import ( + VideoCGF1Evaluator, + VideoPhraseApEvaluator, + VideoPhraseHotaEvaluator, + VideoTetaEvaluator, + YTVISPredFileEvaluator, +) + + +class VEvalEvaluator: + def __init__(self, gt_annot_file: str, eval_res_file: str): + self.gt_annot_file = gt_annot_file + self.eval_res_file = eval_res_file + self.evaluators = [ + # mAP + YTVISPredFileEvaluator(gt_annot_file), + # Phrase AP + VideoPhraseApEvaluator(gt_annot_file), + # TETA + VideoTetaEvaluator(gt_annot_file, use_mask=True, is_exhaustive=True), + # HOTA + VideoPhraseHotaEvaluator(gt_annot_file), + # cgF1 + VideoCGF1Evaluator(gt_annot_file), + ] + + def run_eval(self, pred_file: str): + dataset_results = {} + video_np_results = defaultdict(dict) + for evaluator in self.evaluators: + d_res, v_np_res = evaluator.evaluate(pred_file) + dataset_results.update(d_res) + for (video_id, category_id), res in v_np_res.items(): + video_np_results[(video_id, category_id)].update(res) + + if len(dataset_results) == 0: + dataset_results = {"": 0.0} + + formatted_video_np_results = [ + {"video_id": video_id, "category_id": category_id, **res} + for (video_id, category_id), res in video_np_results.items() + ] + eval_metrics = { + "dataset_results": dataset_results, + "video_np_results": formatted_video_np_results, + } + + with g_pathmgr.open(self.eval_res_file, "w") as f: + json.dump(eval_metrics, f) + + return eval_metrics + + +def run_main_all(dataset_name, args): + gt_annot_file = os.path.join(args.gt_annot_dir, dataset_name + ".json") + pred_file = os.path.join(args.pred_dir, dataset_name + "_preds.json") + eval_res_file = os.path.join(args.eval_res_dir, dataset_name + "_eval_res.json") + print(f"=== Running evaluation for Pred {pred_file} vs GT {gt_annot_file} ===") + veval_evaluator = VEvalEvaluator( + gt_annot_file=gt_annot_file, eval_res_file=eval_res_file + ) + _ = veval_evaluator.run_eval(pred_file=pred_file) + + print(f"=== Results saved to {eval_res_file} ===") + + +def main_all(args): + saco_veval_dataset_names = [ + "saco_veval_sav_test", + "saco_veval_sav_val", + "saco_veval_yt1b_test", + "saco_veval_yt1b_val", + "saco_veval_smartglasses_test", + "saco_veval_smartglasses_val", + ] + + # multiprocessing may not really work as inner evaluator also using multiprocessing + # so we just for loop + for dataset_name in saco_veval_dataset_names: + print(f"=== Running evaluation for dataset {dataset_name} ===") + run_main_all(dataset_name=dataset_name, args=args) + + +def main_one(args): + gt_annot_file = args.gt_annot_file + pred_file = args.pred_file + eval_res_file = args.eval_res_file + + print(f"=== Running evaluation for Pred {pred_file} vs GT {gt_annot_file} ===") + veval_evaluator = VEvalEvaluator( + gt_annot_file=gt_annot_file, eval_res_file=eval_res_file + ) + _ = veval_evaluator.run_eval(pred_file=pred_file) + + print(f"=== Results saved to {eval_res_file} ===") + + +def main(): + parser = argparse.ArgumentParser(description="Run video grounding evaluators") + + # Create subparsers for different commands + subparsers = parser.add_subparsers(dest="command", required=True) + + # Run evaluation for all datasets + all_parser = subparsers.add_parser("all", help="Run evaluation for all datasets") + all_parser.add_argument( + "--gt_annot_dir", + type=str, + help="Directory that contains the ground truth annotation files", + ) + all_parser.add_argument( + "--pred_dir", + type=str, + help="Directory that contains the prediction files", + ) + all_parser.add_argument( + "--eval_res_dir", + type=str, + help="Directory that contains the eval results files", + ) + all_parser.set_defaults(func=main_all) + + # Run evaluation for one dataset + one_parser = subparsers.add_parser("one", help="Run evaluation for one dataset") + one_parser.add_argument( + "--gt_annot_file", + type=str, + help="Path to the ground truth annotation file", + ) + one_parser.add_argument( + "--pred_file", + type=str, + help="Path to the prediction file", + ) + one_parser.add_argument( + "--eval_res_file", + type=str, + help="Path to the eval results file", + ) + one_parser.set_defaults(func=main_one) + + # Parse and dispatch + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/src/nn/segearth_ov3/sam3/eval/saco_veval_evaluators.py b/src/nn/segearth_ov3/sam3/eval/saco_veval_evaluators.py new file mode 100644 index 0000000..4947472 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/saco_veval_evaluators.py @@ -0,0 +1,838 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import json +import os +import tempfile +from collections import defaultdict +from typing import Dict, Optional, Sequence, Tuple + +import numpy as np +import pycocotools.mask +from sam3.eval.cgf1_eval import CGF1_METRICS +from sam3.eval.conversion_util import ( + convert_ytbvis_to_cocovid_gt, + convert_ytbvis_to_cocovid_pred, +) +from sam3.eval.hota_eval_toolkit.run_ytvis_eval import run_ytvis_eval +from sam3.eval.teta_eval_toolkit import config, Evaluator, metrics +from sam3.eval.teta_eval_toolkit.datasets import COCO, TAO +from sam3.eval.ytvis_coco_wrapper import YTVIS +from sam3.eval.ytvis_eval import VideoDemoF1Eval, YTVISeval +from sam3.train.nms_helper import process_frame_level_nms, process_track_level_nms + + +def _get_metric_index(metric_name: str, iou_threshold: Optional[float] = None) -> int: + """ + Find the index of a metric in CGF1_METRICS by name and IoU threshold. + + Args: + metric_name: Name of the metric (e.g., "cgF1", "precision", "recall") + iou_threshold: IoU threshold (None for average over 0.5:0.95, or specific value like 0.5, 0.75) + + Returns: + Index of the metric in CGF1_METRICS + + Raises: + ValueError: If metric not found + """ + for idx, metric in enumerate(CGF1_METRICS): + if metric.name == metric_name and metric.iou_threshold == iou_threshold: + return idx + raise ValueError( + f"Metric '{metric_name}' with IoU threshold {iou_threshold} not found in CGF1_METRICS" + ) + + +class BasePredFileEvaluator: + """A base class for evaluating a prediction file.""" + + pass + + +class YTVISPredFileEvaluator(BasePredFileEvaluator): + """Evaluate class mAP for YT-VIS prediction files.""" + + def __init__( + self, + gt_ann_file: str, + dataset_name: str = "video", + iou_types: Optional[Sequence[str]] = None, + ): + self.gt_ann_file = gt_ann_file + self.dataset_name = dataset_name + self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"] + assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types) + + def evaluate(self, pred_file: str) -> Dict[str, float]: + # use our internal video evaluation toolkit for YT-VIS pred file + # (i.e. the same one we're using for video phrase AP) + results = {} + use_cats = True # YT-VIS mAP evaluation uses categories + ytvisGT = YTVIS(self.gt_ann_file, ignore_gt_cats=not use_cats) + # the original YT-VIS GT annotations have uncompressed RLEs ("counts" is an integer list) + # rather than compressed RLEs ("counts" is a string), so we first convert them here. + if "segm" in self.iou_types: + for ann in ytvisGT.dataset["annotations"]: + ann["segmentations"] = [ + _compress_rle(rle) for rle in ann["segmentations"] + ] + + with open(pred_file) as f: + dt = json.load(f) + # Our prediction file saves "video_id" and absolute (unnormalized) boxes. + # Note that we should use the official (original) YT-VIS annotations (i.e. the one + # saved via "scripts/datasets/training/ytvis_split.py", instead of the one saved + # via "scripts/api_db_to_ytvis_json.py") in this evaluator, which contain absolute + # boxes coordinates in its GT annotations. + for d in dt: + d["image_id"] = d["video_id"] + ytvisDT = ytvisGT.loadRes(dt) + + for iou_type in self.iou_types: + ytvisEval = YTVISeval(ytvisGT, ytvisDT, iou_type) + + # set the area ranges for small, medium, and large objects (using + # absolute pixel areas) as in the official YT-VIS evaluation toolkit: + # https://github.com/achalddave/ytvosapi/blob/eca601117c9f86bad084cb91f1d918e9ab665a75/PythonAPI/ytvostools/ytvoseval.py#L538 + ytvisEval.params.areaRng = [ + [0**2, 1e5**2], + [0**2, 128**2], + [128**2, 256**2], + [256**2, 1e5**2], + ] + ytvisEval.params.areaRngLbl = ["all", "small", "medium", "large"] + ytvisEval.params.useCats = use_cats + + ytvisEval.evaluate() + ytvisEval.accumulate() + ytvisEval.summarize() + result_key = f"{self.dataset_name}_{'mask' if iou_type == 'segm' else 'bbox'}_mAP_50_95" + results[result_key] = ytvisEval.stats[0] + + # video-NP level results not supported for `YTVISPredFileEvaluator` yet + video_np_level_results = {} + return results, video_np_level_results + + +class VideoPhraseApEvaluator(BasePredFileEvaluator): + """Evaluate Video Phrase AP with YT-VIS format prediction and GT files.""" + + def __init__( + self, + gt_ann_file: str, + dataset_name: str = "video", + iou_types: Optional[Sequence[str]] = None, + ): + self.gt_ann_file = gt_ann_file + self.dataset_name = dataset_name + self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"] + assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types) + + def evaluate(self, pred_file: str) -> Dict[str, float]: + with open(self.gt_ann_file) as f: + gt = json.load(f) + with open(pred_file) as f: + dt = json.load(f) + # For phrase AP and demo F1 evaluation, we need to remap each pair of (video_id, category_id) to + # a new unique video_id, so that we don't mix detections from different categories under `useCat=False` + gt, dt = remap_video_category_pairs_to_unique_video_ids(gt, dt) + if "segm" in self.iou_types: + for ann in gt["annotations"]: + ann["segmentations"] = [ + _compress_rle(rle) for rle in ann["segmentations"] + ] + for d in dt: + d["image_id"] = d["video_id"] + + results = {} + use_cats = False # Phrase AP evaluation does not use categories + ytvisGT = YTVIS(annotation_file=None, ignore_gt_cats=not use_cats) + ytvisGT.dataset = gt + ytvisGT.createIndex() + ytvisDT = ytvisGT.loadRes(dt) + + for iou_type in self.iou_types: + phraseApEval = YTVISeval(ytvisGT, ytvisDT, iou_type) + + # set the area ranges for small, medium, and large objects (using + # absolute pixel areas) as in the official YT-VIS evaluation toolkit: + # https://github.com/achalddave/ytvosapi/blob/eca601117c9f86bad084cb91f1d918e9ab665a75/PythonAPI/ytvostools/ytvoseval.py#L538 + phraseApEval.params.areaRng = [ + [0**2, 1e5**2], + [0**2, 128**2], + [128**2, 256**2], + [256**2, 1e5**2], + ] + phraseApEval.params.areaRngLbl = ["all", "small", "medium", "large"] + phraseApEval.params.useCats = use_cats + + phraseApEval.evaluate() + phraseApEval.accumulate() + phraseApEval.summarize() + result_prefix = f"{self.dataset_name}" + result_prefix += f"_{'mask' if iou_type == 'segm' else 'bbox'}_phrase_ap" + # fetch Phrase AP results from the corresponding indices in `phraseApEval.stats` + # (see `_summarizeDets` in https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py) + results[result_prefix + "_50_95"] = phraseApEval.stats[0] # IoU=0.5:0.95 + results[result_prefix + "_50"] = phraseApEval.stats[1] # IoU=0.5 + results[result_prefix + "_75"] = phraseApEval.stats[2] # IoU=0.75 + + # video-NP level results not supported for `VideoPhraseApEvaluator` yet + video_np_level_results = {} + return results, video_np_level_results + + +class VideoCGF1Evaluator(BasePredFileEvaluator): + """Evaluate Video Demo F1 with YT-VIS format prediction and GT files.""" + + def __init__( + self, + gt_ann_file: str, + dataset_name: str = "video", + prob_thresh: float = 0.5, + iou_types: Optional[Sequence[str]] = None, + ): + self.gt_ann_file = gt_ann_file + self.dataset_name = dataset_name + self.prob_thresh = prob_thresh + self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"] + assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types) + + def evaluate(self, pred_file: str) -> Dict[str, float]: + with open(self.gt_ann_file) as f: + gt = json.load(f) + with open(pred_file) as f: + dt = json.load(f) + # compute IL_MCC and CG-F1 can only be computed if we have "video_np_pairs" keys in the GT JSON + compute_ilmcc_and_cgf1 = "video_np_pairs" in gt + if not compute_ilmcc_and_cgf1: + print( + f"Warning: IL_MCC and CG-F1 are not computed for {pred_file=} as it does not have 'video_np_pairs' keys in the GT JSON" + ) + # For phrase AP and demo F1 evaluation, we need to remap each pair of (video_id, category_id) to + # a new unique video_id, so that we don't mix detections from different categories under `useCat=False` + gt, dt = remap_video_category_pairs_to_unique_video_ids( + gt, dt, add_negative_np_pairs=compute_ilmcc_and_cgf1 + ) + if "segm" in self.iou_types: + for ann in gt["annotations"]: + ann["segmentations"] = [ + _compress_rle(rle) for rle in ann["segmentations"] + ] + for d in dt: + d["image_id"] = d["video_id"] + + results = {} + use_cats = False # Demo F1 evaluation does not use categories + ytvisGT = YTVIS(annotation_file=None, ignore_gt_cats=not use_cats) + ytvisGT.dataset = gt + ytvisGT.createIndex() + ytvisDT = ytvisGT.loadRes(dt) + + video_np_level_results = {} + for iou_type in self.iou_types: + demoF1Eval = VideoDemoF1Eval(ytvisGT, ytvisDT, iou_type, self.prob_thresh) + + demoF1Eval.params.useCats = use_cats + demoF1Eval.params.areaRng = [[0**2, 1e5**2]] + demoF1Eval.params.areaRngLbl = ["all"] + demoF1Eval.params.maxDets = [100000] + + demoF1Eval.evaluate() + demoF1Eval.accumulate() + demoF1Eval.summarize() + result_prefix = f"{self.dataset_name}" + result_prefix += f"_{'mask' if iou_type == 'segm' else 'bbox'}_demo" + + stats = demoF1Eval.stats + + if compute_ilmcc_and_cgf1: + # Average IoU threshold (0.5:0.95) + cgf1_micro_avg_idx = _get_metric_index("cgF1", None) + positive_micro_f1_avg_idx = _get_metric_index("positive_micro_F1", None) + ilmcc_avg_idx = _get_metric_index("IL_MCC", None) + results[result_prefix + "_cgf1_micro_50_95"] = stats[cgf1_micro_avg_idx] + results[result_prefix + "_ilmcc_50_95"] = stats[ilmcc_avg_idx] + results[result_prefix + "_positive_micro_f1_50_95"] = stats[ + positive_micro_f1_avg_idx + ] + + # IoU = 0.5 + cgf1_micro_50_idx = _get_metric_index("cgF1", 0.5) + positive_micro_f1_50_idx = _get_metric_index("positive_micro_F1", 0.5) + results[result_prefix + "_cgf1_micro_50"] = stats[cgf1_micro_50_idx] + results[result_prefix + "_ilmcc_50"] = float( + np.array(stats[cgf1_micro_50_idx]) + / np.array(stats[positive_micro_f1_50_idx]) + ) + results[result_prefix + "_positive_micro_f1_50"] = stats[ + positive_micro_f1_50_idx + ] + + # IoU = 0.75 + cgf1_micro_75_idx = _get_metric_index("cgF1", 0.75) + positive_micro_f1_75_idx = _get_metric_index("positive_micro_F1", 0.75) + results[result_prefix + "_cgf1_micro_75"] = stats[cgf1_micro_75_idx] + results[result_prefix + "_ilmcc_75"] = float( + np.array(stats[cgf1_micro_75_idx]) + / np.array(stats[positive_micro_f1_75_idx]) + ) + results[result_prefix + "_positive_micro_f1_75"] = stats[ + positive_micro_f1_75_idx + ] + + self.extract_video_np_level_results(demoF1Eval, video_np_level_results) + + return results, video_np_level_results + + def extract_video_np_level_results(self, demoF1Eval, video_np_level_results): + """Aggregate statistics for video-level metrics.""" + num_iou_thrs = len(demoF1Eval.params.iouThrs) + iou_50_index = int(np.where(demoF1Eval.params.iouThrs == 0.5)[0]) + iou_75_index = int(np.where(demoF1Eval.params.iouThrs == 0.75)[0]) + + result_prefix = "mask" if demoF1Eval.params.iouType == "segm" else "bbox" + + assert len(demoF1Eval.evalImgs) == len(demoF1Eval.cocoGt.dataset["images"]) + for i, video in enumerate(demoF1Eval.cocoGt.dataset["images"]): + # the original video id and category id before remapping + video_id = video["orig_video_id"] + category_id = video["orig_category_id"] + eval_img_dict = demoF1Eval.evalImgs[i] + + TPs = eval_img_dict.get("TPs", np.zeros(num_iou_thrs, dtype=np.int64)) + FPs = eval_img_dict.get("FPs", np.zeros(num_iou_thrs, dtype=np.int64)) + FNs = eval_img_dict.get("FNs", np.zeros(num_iou_thrs, dtype=np.int64)) + assert len(TPs) == len(FPs) == len(FNs) == num_iou_thrs + # F1 = 2*TP / (2*TP + FP + FN), and we set F1 to 1.0 if denominator is 0 + denominator = 2 * TPs + FPs + FNs + F1s = np.where(denominator > 0, 2 * TPs / np.maximum(denominator, 1), 1.0) + local_results = { + f"{result_prefix}_TP_50_95": float(TPs.mean()), + f"{result_prefix}_FP_50_95": float(FPs.mean()), + f"{result_prefix}_FN_50_95": float(FNs.mean()), + f"{result_prefix}_F1_50_95": float(F1s.mean()), + f"{result_prefix}_TP_50": float(TPs[iou_50_index]), + f"{result_prefix}_FP_50": float(FPs[iou_50_index]), + f"{result_prefix}_FN_50": float(FNs[iou_50_index]), + f"{result_prefix}_F1_50": float(F1s[iou_50_index]), + f"{result_prefix}_TP_75": float(TPs[iou_75_index]), + f"{result_prefix}_FP_75": float(FPs[iou_75_index]), + f"{result_prefix}_FN_75": float(FNs[iou_75_index]), + f"{result_prefix}_F1_75": float(F1s[iou_75_index]), + } + if (video_id, category_id) not in video_np_level_results: + video_np_level_results[(video_id, category_id)] = {} + video_np_level_results[(video_id, category_id)].update(local_results) + + +class VideoTetaEvaluator(BasePredFileEvaluator): + """Evaluate TETA metric using YouTubeVIS format prediction and GT files.""" + + def __init__( + self, + gt_ann_file: str, + dataset_name: str = "video", + tracker_name: str = "Sam3", + nms_threshold: float = 0.5, + nms_strategy: str = "none", # "track", "frame", or "none" + prob_thresh: float = 0.5, + is_exhaustive: bool = False, + use_mask: bool = False, + num_parallel_cores: int = 8, + ): + self.gt_ann_file = gt_ann_file + self.dataset_name = dataset_name + self.tracker_name = tracker_name + self.nms_threshold = nms_threshold + self.nms_strategy = nms_strategy.lower() # Convert to lowercase for consistency + self.prob_thresh = prob_thresh + self.metric_prefix = "TETA" + self.is_exhaustive = is_exhaustive + self.use_mask = use_mask + self.num_parallel_cores = num_parallel_cores + + # Verify NMS strategy is valid + valid_strategies = ["track", "frame", "none"] + print("current nms_strategy:", self.nms_strategy) + if self.nms_strategy not in valid_strategies: + raise ValueError( + f"Invalid NMS strategy: {self.nms_strategy}. Must be one of {valid_strategies}" + ) + + print(f"Initialized VideoTetaEvaluator with NMS strategy: {self.nms_strategy}") + print(f"Probability threshold set to: {self.prob_thresh}") + print(f"Dataset exhaustivity set to: {self.is_exhaustive}") + print(f"Tracker name set to: {self.tracker_name}") + print(f"Dataset name set to: {self.dataset_name}") + print(f"Use mask set to: {self.use_mask}") + + def process_predictions(self, pred_file: str, tmp_dir: str) -> str: + """Process predictions with selected NMS strategy""" + with open(pred_file, "r") as f: + raw_preds = json.load(f) + print(f"Processing predictions with {self.nms_strategy} NMS strategy") + + # Filter by score threshold + if self.prob_thresh > 0: + raw_preds = [d for d in raw_preds if d["score"] >= self.prob_thresh] + print( + f"Filtered to {len(raw_preds)} predictions with score >= {self.prob_thresh}" + ) + # Group predictions by video_id + video_groups = defaultdict(list) + for pred in raw_preds: + video_groups[pred["video_id"]].append(pred) + # Process based on NMS strategy + if self.nms_strategy == "track": + process_track_level_nms(video_groups, nms_threshold=self.nms_threshold) + elif self.nms_strategy == "frame": + process_frame_level_nms(video_groups, nms_threshold=self.nms_threshold) + elif self.nms_strategy == "none": + print("Skipping NMS processing as strategy is set to 'none'") + # No processing needed for "none" strategy + # Save processed predictions + processed_preds = [ + track for tracks in video_groups.values() for track in tracks + ] + processed_path = os.path.join(tmp_dir, "processed_preds.json") + with open(processed_path, "w") as f: + json.dump(processed_preds, f) + + print(f"Saved processed predictions to {processed_path}") + return processed_path + + def evaluate(self, pred_file: str) -> Tuple[Dict[str, float], Dict]: + """Main evaluation method""" + + print(f"Evaluating TETA Metric with {self.nms_strategy.upper()} NMS strategy") + with tempfile.TemporaryDirectory() as tmp_dir: + # Process predictions first + processed_pred_file = self.process_predictions(pred_file, tmp_dir) + + # Convert GT to COCO-vid format + gt_dir = os.path.join(tmp_dir, "gt") + os.makedirs(gt_dir, exist_ok=True) + gt_coco_path = os.path.join(gt_dir, "annotations.json") + convert_ytbvis_to_cocovid_gt(self.gt_ann_file, gt_coco_path) + + # Convert processed predictions to COCO-vid format + pred_dir = os.path.join(tmp_dir, "predictions") + tracker_dir = os.path.join(pred_dir, self.tracker_name) + os.makedirs(tracker_dir, exist_ok=True) + pred_coco_path = os.path.join(tracker_dir, "track_results_cocofmt.json") + convert_ytbvis_to_cocovid_pred( + youtubevis_pred_path=processed_pred_file, + converted_dataset_path=gt_coco_path, + output_path=pred_coco_path, + ) + # Configure TETA evaluator + default_eval_config = config.get_default_eval_config() + default_eval_config["PRINT_ONLY_COMBINED"] = True + default_eval_config["DISPLAY_LESS_PROGRESS"] = True + default_eval_config["OUTPUT_TEMP_RAW_DATA"] = True + default_eval_config["NUM_PARALLEL_CORES"] = self.num_parallel_cores + default_dataset_config = config.get_default_dataset_config() + default_dataset_config["TRACKERS_TO_EVAL"] = [self.tracker_name] + default_dataset_config["GT_FOLDER"] = gt_dir + default_dataset_config["OUTPUT_FOLDER"] = pred_dir + default_dataset_config["TRACKER_SUB_FOLDER"] = tracker_dir + default_dataset_config["USE_MASK"] = self.use_mask + + evaluator = Evaluator(default_eval_config) + if self.is_exhaustive: + dataset_list = [COCO(default_dataset_config)] + dataset_parsing_key = "COCO" + else: + dataset_list = [TAO(default_dataset_config)] + dataset_parsing_key = "TAO" + + # Run evaluation + eval_results, _ = evaluator.evaluate( + dataset_list, [metrics.TETA(exhaustive=self.is_exhaustive)] + ) + + # Extract and format results + results = { + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_teta": float( + eval_results[dataset_parsing_key]["TETA"][0] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_loc_a": float( + eval_results[dataset_parsing_key]["TETA"][1] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_assoc_a": float( + eval_results[dataset_parsing_key]["TETA"][2] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_cls_a": float( + eval_results[dataset_parsing_key]["TETA"][3] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_loc_re": float( + eval_results[dataset_parsing_key]["TETA"][4] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_loc_pr": float( + eval_results[dataset_parsing_key]["TETA"][5] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_assoc_re": float( + eval_results[dataset_parsing_key]["TETA"][6] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_assoc_pr": float( + eval_results[dataset_parsing_key]["TETA"][7] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_cls_re": float( + eval_results[dataset_parsing_key]["TETA"][8] + ), + f"{self.dataset_name}_{'mask' if self.use_mask else 'bbox'}_cls_pr": float( + eval_results[dataset_parsing_key]["TETA"][9] + ), + } + + # video-NP level results not supported for `VideoTetaEvaluator` yet + video_np_level_results = {} + return results, video_np_level_results + + +class VideoPhraseHotaEvaluator(BasePredFileEvaluator): + """Evaluate Video Phrase HOTA with YT-VIS format prediction and GT files.""" + + def __init__( + self, + gt_ann_file: str, + dataset_name: str = "video", + prob_thresh: float = 0.5, + iou_types: Optional[Sequence[str]] = None, + compute_video_mot_hota: bool = False, + ): + self.gt_ann_file = gt_ann_file + self.dataset_name = dataset_name + self.prob_thresh = prob_thresh + self.metric_prefix = "phrase" + # the list of metrics to collect from the HOTA evaluation results + self.metric_to_collect = [ + "HOTA", + "DetA", + "AssA", + "DetRe", + "DetPr", + "AssRe", + "AssPr", + "LocA", + "OWTA", + ] + self.iou_types = list(iou_types) if iou_types is not None else ["bbox", "segm"] + assert all(iou_type in ["bbox", "segm"] for iou_type in self.iou_types) + + # If True, compute video MOT HOTA, aggregating predictions/GT from all categories. + self.compute_video_mot_hota = compute_video_mot_hota + + def evaluate(self, pred_file: str) -> Dict[str, float]: + # use the YT-VIS evaluation toolkit in TrackEval + + with open(self.gt_ann_file) as f: + gt = json.load(f) + with open(pred_file) as f: + dt = json.load(f) + # keep only predictions with score above the probability threshold + dt = [d for d in dt if d["score"] > self.prob_thresh] + for d in dt: + assert len(d["areas"]) == len(d["bboxes"]) + assert len(d["areas"]) == len(d["segmentations"]) + # remove empty boxes (otherwise they will count as false positives for during + # per-frame detection accuracy in HOTA evaluation) + for t in range(len(d["bboxes"])): + bbox = d["bboxes"][t] + if d["areas"][t] == 0 or bbox is None or all(x == 0 for x in bbox): + d["segmentations"][t] = None + d["bboxes"][t] = None + d["areas"][t] = None + # check that box occurence and mask occurence are consistent + for bbox, mask, area in zip(d["bboxes"], d["segmentations"], d["areas"]): + assert (area is None) == (bbox is None) + assert (area is None) == (mask is None) + # set all scores to 1.0 for HOTA evaluation (just like Demo F1, the exact score + # value is not used in HOTA metrics; it will be treated as a detection prediction + # as long as its score is above the threshold) + d["score"] = 1.0 + + # remap the GT and DT annotations for phrase HOTA evaluation + gt = _fill_in_ann_height_width(gt) + if not self.compute_video_mot_hota: + # remap the GT and DT annotations for phrase HOTA evaluation + gt, dt = self._remap_gt_dt(gt, dt) + else: + # Compute video-level MOT HOTA + # Apply track-level NMS + video_groups = defaultdict(list) + for pred in dt: + video_groups[pred["video_id"]].append(pred) + process_track_level_nms(video_groups, nms_threshold=0.5) + dt = [track for tracks in video_groups.values() for track in tracks] + + # Remap GT track ids for class-agnostic HOTA + gt, dt = remap_gt_dt_class_agnostic(gt, dt) + + # run the HOTA evaluation using TrackEval on the remapped (video_id, category_id) pairs + out_dict = {} + video_np_level_results = {} + for iou_type in self.iou_types: + output_res, _ = run_ytvis_eval( + args=[ + "--METRICS", + "HOTA", + "--IOU_TYPE", + iou_type, + "--DATASET_NAME", + self.dataset_name, + "--USE_PARALLEL", + "True", + "--NUM_PARALLEL_CORES", + "8", + "--PLOT_CURVES", + "False", + "--LOG_ON_ERROR", + "None", + "--PRINT_ONLY_COMBINED", + "True", + "--OUTPUT_SUMMARY", + "False", + "--OUTPUT_DETAILED", + "False", + "--TIME_PROGRESS", + "False", + "--PRINT_CONFIG", + "False", + ], + gt_json=gt, + dt_json=dt, + ) + self.extract_video_np_level_results( + iou_type=iou_type, + remapped_gt=gt, + raw_results=output_res[self.dataset_name]["tracker"], + video_np_level_results=video_np_level_results, + ) + + def _summarize_results(output_res, iou_type, field, suffix): + eval_res = output_res[self.dataset_name]["tracker"][field] + result_prefix = f"{self.dataset_name}_{'mask' if iou_type == 'segm' else 'bbox'}_{suffix}" + for metric_name in self.metric_to_collect: + eval_res_hota = eval_res["cls_comb_cls_av"]["HOTA"] + result_key = f"{result_prefix}_{self.metric_prefix}_{metric_name}" + result_value = float(np.mean(eval_res_hota[metric_name])) + out_dict[result_key] = result_value + + _summarize_results(output_res, iou_type, "COMBINED_SEQ", "all") + if "COMBINED_SEQ_CHALLENGING" in output_res[self.dataset_name]["tracker"]: + _summarize_results( + output_res, iou_type, "COMBINED_SEQ_CHALLENGING", "challenging" + ) + + # video-NP level results not supported for `VideoPhraseHotaEvaluator` yet + return out_dict, video_np_level_results + + def _remap_gt_dt(self, gt, dt): + # For phrase HOTA evaluation, we need to remap each pair of (video_id, category_id) to + # a new unique video_id, so that we don't mix detections from different categories + gt, dt = remap_video_category_pairs_to_unique_video_ids(gt, dt) + # We further map all the categories to category_id=1 in HOTA evaluation toolkit + # for phrase HOTA (similar to "useCat=False" for video phrase AP) + remapped_category_id = 1 + gt["categories"] = [ + { + "supercategory": "object", + "id": remapped_category_id, + "name": "_REMAPPED_FOR_PHRASE_METRICS_", + } + ] + for ann in gt["annotations"]: + ann["category_id"] = remapped_category_id + for d in dt: + d["category_id"] = remapped_category_id + # To be compatible with the TrackEval YT-VIS evaluation toolkit, we need to give + # unique filenames to each remapped video, so we add remapped video_id as prefix. + for video in gt["videos"]: + new_video_id = video["id"] + video["file_names"] = [ + f"remapped_vid_{new_video_id:012d}/{name}" + for name in video["file_names"] + ] + return gt, dt + + def extract_video_np_level_results( + self, iou_type, remapped_gt, raw_results, video_np_level_results + ): + """Aggregate statistics for video-level metrics.""" + result_prefix = "mask" if iou_type == "segm" else "bbox" + for video in remapped_gt["videos"]: + # the original video id and category id before remapping + video_id = video["orig_video_id"] + category_id = video["orig_category_id"] + video_key = f"remapped_vid_{video['id']:012d}" + results = raw_results[video_key]["_REMAPPED_FOR_PHRASE_METRICS_"]["HOTA"] + + local_results = {} + for metric_name in self.metric_to_collect: + result_key = f"{result_prefix}_{metric_name}" + local_results[result_key] = float(results[metric_name].mean()) + if (video_id, category_id) not in video_np_level_results: + video_np_level_results[(video_id, category_id)] = {} + video_np_level_results[(video_id, category_id)].update(local_results) + + +class VideoClassBasedHotaEvaluator(VideoPhraseHotaEvaluator): + def __init__( + self, + gt_ann_file: str, + dataset_name: str = "video", + prob_thresh: float = 0.5, + ): + super().__init__(gt_ann_file, dataset_name, prob_thresh) + self.metric_prefix = "class" + + def _remap_gt_dt(self, gt, dt): + return gt, dt # no remapping needed for class-based HOTA evaluation + + def extract_video_np_level_results(self, *args, **kwargs): + pass # no video-NP level results for class-based HOTA evaluation + + +def _compress_rle(rle): + """Convert RLEs from uncompressed (integer list) to compressed (string) format.""" + if rle is None: + return None + if isinstance(rle["counts"], list): + rle = pycocotools.mask.frPyObjects(rle, rle["size"][0], rle["size"][1]) + rle["counts"] = rle["counts"].decode() + return rle + + +def remap_video_category_pairs_to_unique_video_ids( + gt_json, dt_json, add_negative_np_pairs=False +): + """ + Remap each pair of (video_id, category_id) to a new unique video_id. This is useful + for phrase AP and demo F1 evaluation on videos, where we have `useCat=False` and + rely on separating different NPs (from the same video) into different new video ids, + so that we don't mix detections from different categories in computeIoU under `useCat=False`. + + This is consistent with how do we phrase AP and demo F1 evaluation on images, where we + use a remapped unique coco_image_id for each image-NP pair (based in its query["id"] in + CustomCocoDetectionAPI.load_queries in modulated_detection_api.py) + """ + # collect the unique video_id-category_id pairs + video_id_to_video = {v["id"]: v for v in gt_json["videos"]} + video_id_category_id_pairs = set() + for pred in dt_json: + video_id_category_id_pairs.add((pred["video_id"], pred["category_id"])) + for ann in gt_json["annotations"]: + video_id_category_id_pairs.add((ann["video_id"], ann["category_id"])) + + # assign the video_id-category_id pairs to unique video ids + video_id_category_id_pairs = sorted(video_id_category_id_pairs) + video_id_category_id_to_new_video_id = { + pair: (i + 1) for i, pair in enumerate(video_id_category_id_pairs) + } + # also map the negative NP pairs -- this is needed for IL_MCC and CG-F1 evaluation + if add_negative_np_pairs: + for vnp in gt_json["video_np_pairs"]: + pair = (vnp["video_id"], vnp["category_id"]) + if pair not in video_id_category_id_to_new_video_id: + video_id_category_id_to_new_video_id[pair] = ( + len(video_id_category_id_to_new_video_id) + 1 + ) + + # map the "video_id" in predictions + for pred in dt_json: + pred["video_id"] = video_id_category_id_to_new_video_id[ + (pred["video_id"], pred["category_id"]) + ] + # map the "video_id" in gt_json["annotations"] + for ann in gt_json["annotations"]: + ann["video_id"] = video_id_category_id_to_new_video_id[ + (ann["video_id"], ann["category_id"]) + ] + # map and duplicate gt_json["videos"] + new_videos = [] + for ( + video_id, + category_id, + ), new_video_id in video_id_category_id_to_new_video_id.items(): + video = video_id_to_video[video_id].copy() + video["id"] = new_video_id + # preserve the original video_id and category_id of each remapped video entry, + # so that we can associate sample-level eval metrics with the original video-NP pairs + video["orig_video_id"] = video_id + video["orig_category_id"] = category_id + new_videos.append(video) + gt_json["videos"] = new_videos + + return gt_json, dt_json + + +def remap_gt_dt_class_agnostic(gt, dt): + """ + For class-agnostic HOTA, merge all GT tracks for each video (across NPs), + ensure unique track_ids, and set all category_id to 1. + Also, add orig_video_id and orig_category_id for compatibility. + """ + # 1. Remap all GT track_ids to be unique per video + gt_anns_by_video = defaultdict(list) + for ann in gt["annotations"]: + gt_anns_by_video[ann["video_id"]].append(ann) + + # Ensure unique track ids across tracks of all videos + next_tid = 1 + for _, anns in gt_anns_by_video.items(): + # Map old track_ids to new unique ones + old_to_new_tid = {} + for ann in anns: + old_tid = ann["id"] + if old_tid not in old_to_new_tid: + old_to_new_tid[old_tid] = next_tid + next_tid += 1 + ann["id"] = old_to_new_tid[old_tid] + # Set category_id to 1 for class-agnostic + ann["category_id"] = 1 + + # Set all GT categories to a single category + gt["categories"] = [ + { + "supercategory": "object", + "id": 1, + "name": "_REMAPPED_FOR_PHRASE_METRICS_", + } + ] + + # Add orig_video_id and orig_category_id to each video for compatibility + anns_by_video = defaultdict(list) + for ann in gt["annotations"]: + anns_by_video[ann["video_id"]].append(ann) + for video in gt["videos"]: + video["orig_video_id"] = video["id"] + # Use the first annotation's original category_id if available, else None + orig_cat = ( + anns_by_video[video["id"]][0]["category_id"] + if anns_by_video[video["id"]] + else None + ) + video["orig_category_id"] = orig_cat + video["file_names"] = [ + f"remapped_vid_{video['id']:012d}/{name}" for name in video["file_names"] + ] + + # Set all DT category_id to 1 + for d in dt: + d["category_id"] = 1 + return gt, dt + + +def _fill_in_ann_height_width(gt_json): + """Fill in missing height/width in GT annotations from its video info.""" + video_id_to_video = {v["id"]: v for v in gt_json["videos"]} + for ann in gt_json["annotations"]: + if "height" not in ann or "width" not in ann: + video = video_id_to_video[ann["video_id"]] + if "height" not in ann: + ann["height"] = video["height"] + if "width" not in ann: + ann["width"] = video["width"] + + return gt_json diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/__init__.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/__init__.py new file mode 100644 index 0000000..1609f80 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/__init__.py @@ -0,0 +1,5 @@ +# fmt: off +# flake8: noqa + +from . import config, datasets, metrics, utils +from .eval import Evaluator diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/_timing.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/_timing.py new file mode 100644 index 0000000..1c6fbfa --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/_timing.py @@ -0,0 +1,69 @@ +# fmt: off +# flake8: noqa + +import inspect +from functools import wraps +from time import perf_counter + +DO_TIMING = False +DISPLAY_LESS_PROGRESS = False +timer_dict = {} +counter = 0 + + +def time(f): + @wraps(f) + def wrap(*args, **kw): + if DO_TIMING: + # Run function with timing + ts = perf_counter() + result = f(*args, **kw) + te = perf_counter() + tt = te - ts + + # Get function name + arg_names = inspect.getfullargspec(f)[0] + if arg_names[0] == "self" and DISPLAY_LESS_PROGRESS: + return result + elif arg_names[0] == "self": + method_name = type(args[0]).__name__ + "." + f.__name__ + else: + method_name = f.__name__ + + # Record accumulative time in each function for analysis + if method_name in timer_dict.keys(): + timer_dict[method_name] += tt + else: + timer_dict[method_name] = tt + + # If code is finished, display timing summary + if method_name == "Evaluator.evaluate": + print("") + print("Timing analysis:") + for key, value in timer_dict.items(): + print("%-70s %2.4f sec" % (key, value)) + else: + # Get function argument values for printing special arguments of interest + arg_titles = ["tracker", "seq", "cls"] + arg_vals = [] + for i, a in enumerate(arg_names): + if a in arg_titles: + arg_vals.append(args[i]) + arg_text = "(" + ", ".join(arg_vals) + ")" + + # Display methods and functions with different indentation. + if arg_names[0] == "self": + print("%-74s %2.4f sec" % (" " * 4 + method_name + arg_text, tt)) + elif arg_names[0] == "test": + pass + else: + global counter + counter += 1 + print("%i %-70s %2.4f sec" % (counter, method_name + arg_text, tt)) + + return result + else: + # If config["TIME_PROGRESS"] is false, or config["USE_PARALLEL"] is true, run functions normally without timing. + return f(*args, **kw) + + return wrap diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/config.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/config.py new file mode 100644 index 0000000..6342fa2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/config.py @@ -0,0 +1,153 @@ +# fmt: off +# flake8: noqa + +"""Config.""" +import argparse +import os + + +def parse_configs(): + """Parse command line.""" + default_eval_config = get_default_eval_config() + default_eval_config["DISPLAY_LESS_PROGRESS"] = True + default_dataset_config = get_default_dataset_config() + default_metrics_config = {"METRICS": ["TETA"]} + config = { + **default_eval_config, + **default_dataset_config, + **default_metrics_config, + } + parser = argparse.ArgumentParser() + for setting in config.keys(): + if type(config[setting]) == list or type(config[setting]) == type(None): + parser.add_argument("--" + setting, nargs="+") + else: + parser.add_argument("--" + setting) + args = parser.parse_args().__dict__ + for setting in args.keys(): + if args[setting] is not None: + if type(config[setting]) == type(True): + if args[setting] == "True": + x = True + elif args[setting] == "False": + x = False + else: + raise Exception( + f"Command line parameter {setting} must be True/False" + ) + elif type(config[setting]) == type(1): + x = int(args[setting]) + elif type(args[setting]) == type(None): + x = None + else: + x = args[setting] + config[setting] = x + eval_config = {k: v for k, v in config.items() if k in default_eval_config.keys()} + dataset_config = { + k: v for k, v in config.items() if k in default_dataset_config.keys() + } + metrics_config = { + k: v for k, v in config.items() if k in default_metrics_config.keys() + } + + return eval_config, dataset_config, metrics_config + + +def get_default_eval_config(): + """Returns the default config values for evaluation.""" + code_path = get_code_path() + default_config = { + "USE_PARALLEL": True, + "NUM_PARALLEL_CORES": 8, + "BREAK_ON_ERROR": True, + "RETURN_ON_ERROR": False, + "LOG_ON_ERROR": os.path.join(code_path, "error_log.txt"), + "PRINT_RESULTS": True, + "PRINT_ONLY_COMBINED": True, + "PRINT_CONFIG": True, + "TIME_PROGRESS": True, + "DISPLAY_LESS_PROGRESS": True, + "OUTPUT_SUMMARY": True, + "OUTPUT_EMPTY_CLASSES": True, + "OUTPUT_TEM_RAW_DATA": True, + "OUTPUT_PER_SEQ_RES": True, + } + return default_config + + +def get_default_dataset_config(): + """Default class config values""" + code_path = get_code_path() + default_config = { + "GT_FOLDER": os.path.join( + code_path, "data/gt/tao/tao_training" + ), # Location of GT data + "TRACKERS_FOLDER": os.path.join( + code_path, "data/trackers/tao/tao_training" + ), # Trackers location + "OUTPUT_FOLDER": None, # Where to save eval results (if None, same as TRACKERS_FOLDER) + "TRACKERS_TO_EVAL": ['TETer'], # Filenames of trackers to eval (if None, all in folder) + "CLASSES_TO_EVAL": None, # Classes to eval (if None, all classes) + "SPLIT_TO_EVAL": "training", # Valid: 'training', 'val' + "PRINT_CONFIG": True, # Whether to print current config + "TRACKER_SUB_FOLDER": "data", # Tracker files are in TRACKER_FOLDER/tracker_name/TRACKER_SUB_FOLDER + "OUTPUT_SUB_FOLDER": "", # Output files are saved in OUTPUT_FOLDER/tracker_name/OUTPUT_SUB_FOLDER + "TRACKER_DISPLAY_NAMES": None, # Names of trackers to display, if None: TRACKERS_TO_EVAL + "MAX_DETECTIONS": 0, # Number of maximal allowed detections per image (0 for unlimited) + "USE_MASK": False, # Whether to use mask data for evaluation + } + return default_config + + +def init_config(config, default_config, name=None): + """Initialize non-given config values with defaults.""" + if config is None: + config = default_config + else: + for k in default_config.keys(): + if k not in config.keys(): + config[k] = default_config[k] + if name and config["PRINT_CONFIG"]: + print("\n%s Config:" % name) + for c in config.keys(): + print("%-20s : %-30s" % (c, config[c])) + return config + + +def update_config(config): + """ + Parse the arguments of a script and updates the config values for a given value if specified in the arguments. + :param config: the config to update + :return: the updated config + """ + parser = argparse.ArgumentParser() + for setting in config.keys(): + if type(config[setting]) == list or type(config[setting]) == type(None): + parser.add_argument("--" + setting, nargs="+") + else: + parser.add_argument("--" + setting) + args = parser.parse_args().__dict__ + for setting in args.keys(): + if args[setting] is not None: + if type(config[setting]) == type(True): + if args[setting] == "True": + x = True + elif args[setting] == "False": + x = False + else: + raise Exception( + "Command line parameter " + setting + "must be True or False" + ) + elif type(config[setting]) == type(1): + x = int(args[setting]) + elif type(args[setting]) == type(None): + x = None + else: + x = args[setting] + config[setting] = x + return config + + +def get_code_path(): + """Get base path where code is""" + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/__init__.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/__init__.py new file mode 100644 index 0000000..97087d7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/__init__.py @@ -0,0 +1,5 @@ +# fmt: off +# flake8: noqa +"""Datasets.""" +from .coco import COCO +from .tao import TAO diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/_base_dataset.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/_base_dataset.py new file mode 100644 index 0000000..fc2c30c --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/_base_dataset.py @@ -0,0 +1,379 @@ +# fmt: off +# flake8: noqa + +import csv +import io +import os +import traceback +import zipfile +from abc import ABC, abstractmethod +from copy import deepcopy + +import numpy as np + +from .. import _timing +from ..utils import TrackEvalException + + +class _BaseDataset(ABC): + @abstractmethod + def __init__(self): + self.tracker_list = None + self.seq_list = None + self.class_list = None + self.output_fol = None + self.output_sub_fol = None + self.should_classes_combine = True + self.use_super_categories = False + + # Functions to implement: + + @abstractmethod + def _load_raw_file(self, tracker, seq, is_gt): + ... + + @_timing.time + @abstractmethod + def get_preprocessed_seq_data(self, raw_data, cls): + ... + + @abstractmethod + def _calculate_similarities(self, gt_dets_t, tracker_dets_t): + ... + + # Helper functions for all datasets: + + @classmethod + def get_class_name(cls): + return cls.__name__ + + def get_name(self): + return self.get_class_name() + + def get_output_fol(self, tracker): + return os.path.join(self.output_fol, tracker, self.output_sub_fol) + + def get_display_name(self, tracker): + """Can be overwritten if the trackers name (in files) is different to how it should be displayed. + By default this method just returns the trackers name as is. + """ + return tracker + + def get_eval_info(self): + """Return info about the dataset needed for the Evaluator""" + return self.tracker_list, self.seq_list, self.class_list + + @_timing.time + def get_raw_seq_data(self, tracker, seq): + """Loads raw data (tracker and ground-truth) for a single tracker on a single sequence. + Raw data includes all of the information needed for both preprocessing and evaluation, for all classes. + A later function (get_processed_seq_data) will perform such preprocessing and extract relevant information for + the evaluation of each class. + + This returns a dict which contains the fields: + [num_timesteps]: integer + [gt_ids, tracker_ids, gt_classes, tracker_classes, tracker_confidences]: + list (for each timestep) of 1D NDArrays (for each det). + [gt_dets, tracker_dets, gt_crowd_ignore_regions]: list (for each timestep) of lists of detections. + [similarity_scores]: list (for each timestep) of 2D NDArrays. + [gt_extras]: dict (for each extra) of lists (for each timestep) of 1D NDArrays (for each det). + + gt_extras contains dataset specific information used for preprocessing such as occlusion and truncation levels. + + Note that similarities are extracted as part of the dataset and not the metric, because almost all metrics are + independent of the exact method of calculating the similarity. However datasets are not (e.g. segmentation + masks vs 2D boxes vs 3D boxes). + We calculate the similarity before preprocessing because often both preprocessing and evaluation require it and + we don't wish to calculate this twice. + We calculate similarity between all gt and tracker classes (not just each class individually) to allow for + calculation of metrics such as class confusion matrices. Typically the impact of this on performance is low. + """ + # Load raw data. + raw_gt_data = self._load_raw_file(tracker, seq, is_gt=True) + raw_tracker_data = self._load_raw_file(tracker, seq, is_gt=False) + raw_data = {**raw_tracker_data, **raw_gt_data} # Merges dictionaries + + # Calculate similarities for each timestep. + similarity_scores = [] + for _, (gt_dets_t, tracker_dets_t) in enumerate( + zip(raw_data["gt_dets"], raw_data["tk_dets"]) + ): + ious = self._calculate_similarities(gt_dets_t, tracker_dets_t) + similarity_scores.append(ious) + raw_data["similarity_scores"] = similarity_scores + return raw_data + + @staticmethod + def _load_simple_text_file( + file, + time_col=0, + id_col=None, + remove_negative_ids=False, + valid_filter=None, + crowd_ignore_filter=None, + convert_filter=None, + is_zipped=False, + zip_file=None, + force_delimiters=None, + ): + """Function that loads data which is in a commonly used text file format. + Assumes each det is given by one row of a text file. + There is no limit to the number or meaning of each column, + however one column needs to give the timestep of each det (time_col) which is default col 0. + + The file dialect (deliminator, num cols, etc) is determined automatically. + This function automatically separates dets by timestep, + and is much faster than alternatives such as np.loadtext or pandas. + + If remove_negative_ids is True and id_col is not None, dets with negative values in id_col are excluded. + These are not excluded from ignore data. + + valid_filter can be used to only include certain classes. + It is a dict with ints as keys, and lists as values, + such that a row is included if "row[key].lower() is in value" for all key/value pairs in the dict. + If None, all classes are included. + + crowd_ignore_filter can be used to read crowd_ignore regions separately. It has the same format as valid filter. + + convert_filter can be used to convert value read to another format. + This is used most commonly to convert classes given as string to a class id. + This is a dict such that the key is the column to convert, and the value is another dict giving the mapping. + + Optionally, input files could be a zip of multiple text files for storage efficiency. + + Returns read_data and ignore_data. + Each is a dict (with keys as timesteps as strings) of lists (over dets) of lists (over column values). + Note that all data is returned as strings, and must be converted to float/int later if needed. + Note that timesteps will not be present in the returned dict keys if there are no dets for them + """ + + if remove_negative_ids and id_col is None: + raise TrackEvalException( + "remove_negative_ids is True, but id_col is not given." + ) + if crowd_ignore_filter is None: + crowd_ignore_filter = {} + if convert_filter is None: + convert_filter = {} + try: + if is_zipped: # Either open file directly or within a zip. + if zip_file is None: + raise TrackEvalException( + "is_zipped set to True, but no zip_file is given." + ) + archive = zipfile.ZipFile(os.path.join(zip_file), "r") + fp = io.TextIOWrapper(archive.open(file, "r")) + else: + fp = open(file) + read_data = {} + crowd_ignore_data = {} + fp.seek(0, os.SEEK_END) + # check if file is empty + if fp.tell(): + fp.seek(0) + dialect = csv.Sniffer().sniff( + fp.readline(), delimiters=force_delimiters + ) # Auto determine structure. + dialect.skipinitialspace = ( + True # Deal with extra spaces between columns + ) + fp.seek(0) + reader = csv.reader(fp, dialect) + for row in reader: + try: + # Deal with extra trailing spaces at the end of rows + if row[-1] in "": + row = row[:-1] + timestep = str(int(float(row[time_col]))) + # Read ignore regions separately. + is_ignored = False + for ignore_key, ignore_value in crowd_ignore_filter.items(): + if row[ignore_key].lower() in ignore_value: + # Convert values in one column (e.g. string to id) + for ( + convert_key, + convert_value, + ) in convert_filter.items(): + row[convert_key] = convert_value[ + row[convert_key].lower() + ] + # Save data separated by timestep. + if timestep in crowd_ignore_data.keys(): + crowd_ignore_data[timestep].append(row) + else: + crowd_ignore_data[timestep] = [row] + is_ignored = True + if ( + is_ignored + ): # if det is an ignore region, it cannot be a normal det. + continue + # Exclude some dets if not valid. + if valid_filter is not None: + for key, value in valid_filter.items(): + if row[key].lower() not in value: + continue + if remove_negative_ids: + if int(float(row[id_col])) < 0: + continue + # Convert values in one column (e.g. string to id) + for convert_key, convert_value in convert_filter.items(): + row[convert_key] = convert_value[row[convert_key].lower()] + # Save data separated by timestep. + if timestep in read_data.keys(): + read_data[timestep].append(row) + else: + read_data[timestep] = [row] + except Exception: + exc_str_init = ( + "In file %s the following line cannot be read correctly: \n" + % os.path.basename(file) + ) + exc_str = " ".join([exc_str_init] + row) + raise TrackEvalException(exc_str) + fp.close() + except Exception: + print("Error loading file: %s, printing traceback." % file) + traceback.print_exc() + raise TrackEvalException( + "File %s cannot be read because it is either not present or invalidly formatted" + % os.path.basename(file) + ) + return read_data, crowd_ignore_data + + @staticmethod + def _calculate_mask_ious(masks1, masks2, is_encoded=False, do_ioa=False): + """Calculates the IOU (intersection over union) between two arrays of segmentation masks. + If is_encoded a run length encoding with pycocotools is assumed as input format, otherwise an input of numpy + arrays of the shape (num_masks, height, width) is assumed and the encoding is performed. + If do_ioa (intersection over area) , then calculates the intersection over the area of masks1 - this is commonly + used to determine if detections are within crowd ignore region. + :param masks1: first set of masks (numpy array of shape (num_masks, height, width) if not encoded, + else pycocotools rle encoded format) + :param masks2: second set of masks (numpy array of shape (num_masks, height, width) if not encoded, + else pycocotools rle encoded format) + :param is_encoded: whether the input is in pycocotools rle encoded format + :param do_ioa: whether to perform IoA computation + :return: the IoU/IoA scores + """ + + # Only loaded when run to reduce minimum requirements + from pycocotools import mask as mask_utils + + # use pycocotools for run length encoding of masks + if not is_encoded: + masks1 = mask_utils.encode( + np.array(np.transpose(masks1, (1, 2, 0)), order="F") + ) + masks2 = mask_utils.encode( + np.array(np.transpose(masks2, (1, 2, 0)), order="F") + ) + + # use pycocotools for iou computation of rle encoded masks + ious = mask_utils.iou(masks1, masks2, [do_ioa] * len(masks2)) + if len(masks1) == 0 or len(masks2) == 0: + ious = np.asarray(ious).reshape(len(masks1), len(masks2)) + assert (ious >= 0 - np.finfo("float").eps).all() + assert (ious <= 1 + np.finfo("float").eps).all() + + return ious + + @staticmethod + def _calculate_box_ious(bboxes1, bboxes2, box_format="xywh", do_ioa=False): + """Calculates the IOU (intersection over union) between two arrays of boxes. + Allows variable box formats ('xywh' and 'x0y0x1y1'). + If do_ioa (intersection over area) , then calculates the intersection over the area of boxes1 - this is commonly + used to determine if detections are within crowd ignore region. + """ + if box_format in "xywh": + # layout: (x0, y0, w, h) + bboxes1 = deepcopy(bboxes1) + bboxes2 = deepcopy(bboxes2) + + bboxes1[:, 2] = bboxes1[:, 0] + bboxes1[:, 2] + bboxes1[:, 3] = bboxes1[:, 1] + bboxes1[:, 3] + bboxes2[:, 2] = bboxes2[:, 0] + bboxes2[:, 2] + bboxes2[:, 3] = bboxes2[:, 1] + bboxes2[:, 3] + elif box_format not in "x0y0x1y1": + raise (TrackEvalException("box_format %s is not implemented" % box_format)) + + # layout: (x0, y0, x1, y1) + min_ = np.minimum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :]) + max_ = np.maximum(bboxes1[:, np.newaxis, :], bboxes2[np.newaxis, :, :]) + intersection = np.maximum(min_[..., 2] - max_[..., 0], 0) * np.maximum( + min_[..., 3] - max_[..., 1], 0 + ) + area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * ( + bboxes1[..., 3] - bboxes1[..., 1] + ) + + if do_ioa: + ioas = np.zeros_like(intersection) + valid_mask = area1 > 0 + np.finfo("float").eps + ioas[valid_mask, :] = ( + intersection[valid_mask, :] / area1[valid_mask][:, np.newaxis] + ) + + return ioas + else: + area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * ( + bboxes2[..., 3] - bboxes2[..., 1] + ) + union = area1[:, np.newaxis] + area2[np.newaxis, :] - intersection + intersection[area1 <= 0 + np.finfo("float").eps, :] = 0 + intersection[:, area2 <= 0 + np.finfo("float").eps] = 0 + intersection[union <= 0 + np.finfo("float").eps] = 0 + union[union <= 0 + np.finfo("float").eps] = 1 + ious = intersection / union + return ious + + @staticmethod + def _calculate_euclidean_similarity(dets1, dets2, zero_distance=2.0): + """Calculates the euclidean distance between two sets of detections, and then converts this into a similarity + measure with values between 0 and 1 using the following formula: sim = max(0, 1 - dist/zero_distance). + The default zero_distance of 2.0, corresponds to the default used in MOT15_3D, such that a 0.5 similarity + threshold corresponds to a 1m distance threshold for TPs. + """ + dist = np.linalg.norm(dets1[:, np.newaxis] - dets2[np.newaxis, :], axis=2) + sim = np.maximum(0, 1 - dist / zero_distance) + return sim + + @staticmethod + def _check_unique_ids(data, after_preproc=False): + """Check the requirement that the tracker_ids and gt_ids are unique per timestep""" + gt_ids = data["gt_ids"] + tracker_ids = data["tk_ids"] + for t, (gt_ids_t, tracker_ids_t) in enumerate(zip(gt_ids, tracker_ids)): + if len(tracker_ids_t) > 0: + unique_ids, counts = np.unique(tracker_ids_t, return_counts=True) + if np.max(counts) != 1: + duplicate_ids = unique_ids[counts > 1] + exc_str_init = ( + "Tracker predicts the same ID more than once in a single timestep " + "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1) + ) + exc_str = ( + " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")" + ) + if after_preproc: + exc_str_init += ( + "\n Note that this error occurred after preprocessing (but not before), " + "so ids may not be as in file, and something seems wrong with preproc." + ) + raise TrackEvalException(exc_str) + if len(gt_ids_t) > 0: + unique_ids, counts = np.unique(gt_ids_t, return_counts=True) + if np.max(counts) != 1: + duplicate_ids = unique_ids[counts > 1] + exc_str_init = ( + "Ground-truth has the same ID more than once in a single timestep " + "(seq: %s, frame: %i, ids:" % (data["seq"], t + 1) + ) + exc_str = ( + " ".join([exc_str_init] + [str(d) for d in duplicate_ids]) + ")" + ) + if after_preproc: + exc_str_init += ( + "\n Note that this error occurred after preprocessing (but not before), " + "so ids may not be as in file, and something seems wrong with preproc." + ) + raise TrackEvalException(exc_str) diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/coco.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/coco.py new file mode 100644 index 0000000..ca3d823 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/coco.py @@ -0,0 +1,637 @@ +# fmt: off +# flake8: noqa + +"""COCO Dataset.""" +import copy +import itertools +import json +import os +from collections import defaultdict + +import numpy as np +from scipy.optimize import linear_sum_assignment + +from .. import _timing, utils +from ..config import get_default_dataset_config, init_config +from ..utils import TrackEvalException +from ._base_dataset import _BaseDataset + + +class COCO(_BaseDataset): + """Tracking datasets in COCO format.""" + + def __init__(self, config=None): + """Initialize dataset, checking that all required files are present.""" + super().__init__() + # Fill non-given config values with defaults + self.config = init_config(config, get_default_dataset_config(), self.get_name()) + self.gt_fol = self.config["GT_FOLDER"] + self.tracker_fol = self.config["TRACKERS_FOLDER"] + self.should_classes_combine = True + self.use_super_categories = False + self.use_mask = self.config["USE_MASK"] + + self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"] + self.output_fol = self.config["OUTPUT_FOLDER"] + if self.output_fol is None: + self.output_fol = self.tracker_fol + self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"] + + if self.gt_fol.endswith(".json"): + self.gt_data = json.load(open(self.gt_fol, "r")) + else: + gt_dir_files = [ + file for file in os.listdir(self.gt_fol) if file.endswith(".json") + ] + if len(gt_dir_files) != 1: + raise TrackEvalException( + f"{self.gt_fol} does not contain exactly one json file." + ) + + with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f: + self.gt_data = json.load(f) + + # fill missing video ids + self._fill_video_ids_inplace(self.gt_data["annotations"]) + + # get sequences to eval and sequence information + self.seq_list = [ + vid["name"].replace("/", "-") for vid in self.gt_data["videos"] + ] + self.seq_name2seqid = { + vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"] + } + # compute mappings from videos to annotation data + self.video2gt_track, self.video2gt_image = self._compute_vid_mappings( + self.gt_data["annotations"] + ) + # compute sequence lengths + self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]} + for img in self.gt_data["images"]: + self.seq_lengths[img["video_id"]] += 1 + self.seq2images2timestep = self._compute_image_to_timestep_mappings() + self.seq2cls = { + vid["id"]: { + "pos_cat_ids": list( + {track["category_id"] for track in self.video2gt_track[vid["id"]]} + ), + } + for vid in self.gt_data["videos"] + } + + # Get classes to eval + considered_vid_ids = [self.seq_name2seqid[vid] for vid in self.seq_list] + seen_cats = set( + [ + cat_id + for vid_id in considered_vid_ids + for cat_id in self.seq2cls[vid_id]["pos_cat_ids"] + ] + ) + # only classes with ground truth are evaluated in TAO + self.valid_classes = [ + cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats + ] + cls_name2clsid_map = { + cls["name"]: cls["id"] for cls in self.gt_data["categories"] + } + + if self.config["CLASSES_TO_EVAL"]: + self.class_list = [ + cls.lower() if cls.lower() in self.valid_classes else None + for cls in self.config["CLASSES_TO_EVAL"] + ] + if not all(self.class_list): + valid_cls = ", ".join(self.valid_classes) + raise TrackEvalException( + "Attempted to evaluate an invalid class. Only classes " + f"{valid_cls} are valid (classes present in ground truth" + " data)." + ) + else: + self.class_list = [cls for cls in self.valid_classes] + self.cls_name2clsid = { + k: v for k, v in cls_name2clsid_map.items() if k in self.class_list + } + self.clsid2cls_name = { + v: k for k, v in cls_name2clsid_map.items() if k in self.class_list + } + # get trackers to eval + if self.config["TRACKERS_TO_EVAL"] is None: + self.tracker_list = os.listdir(self.tracker_fol) + else: + self.tracker_list = self.config["TRACKERS_TO_EVAL"] + + if self.config["TRACKER_DISPLAY_NAMES"] is None: + self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list)) + elif (self.config["TRACKERS_TO_EVAL"] is not None) and ( + len(self.config["TK_DISPLAY_NAMES"]) == len(self.tracker_list) + ): + self.tracker_to_disp = dict( + zip(self.tracker_list, self.config["TK_DISPLAY_NAMES"]) + ) + else: + raise TrackEvalException( + "List of tracker files and tracker display names do not match." + ) + + self.tracker_data = {tracker: dict() for tracker in self.tracker_list} + + for tracker in self.tracker_list: + if self.tracker_sub_fol.endswith(".json"): + with open(os.path.join(self.tracker_sub_fol)) as f: + curr_data = json.load(f) + else: + tr_dir = os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol) + tr_dir_files = [ + file for file in os.listdir(tr_dir) if file.endswith(".json") + ] + if len(tr_dir_files) != 1: + raise TrackEvalException( + f"{tr_dir} does not contain exactly one json file." + ) + with open(os.path.join(tr_dir, tr_dir_files[0])) as f: + curr_data = json.load(f) + + # limit detections if MAX_DETECTIONS > 0 + if self.config["MAX_DETECTIONS"]: + curr_data = self._limit_dets_per_image(curr_data) + + # fill missing video ids + self._fill_video_ids_inplace(curr_data) + + # make track ids unique over whole evaluation set + self._make_tk_ids_unique(curr_data) + + # get tracker sequence information + curr_vids2tracks, curr_vids2images = self._compute_vid_mappings(curr_data) + self.tracker_data[tracker]["vids_to_tracks"] = curr_vids2tracks + self.tracker_data[tracker]["vids_to_images"] = curr_vids2images + + def get_display_name(self, tracker): + return self.tracker_to_disp[tracker] + + def _load_raw_file(self, tracker, seq, is_gt): + """Load a file (gt or tracker) in the TAO format + + If is_gt, this returns a dict which contains the fields: + [gt_ids, gt_classes]: + list (for each timestep) of 1D NDArrays (for each det). + [gt_dets]: list (for each timestep) of lists of detections. + + if not is_gt, this returns a dict which contains the fields: + [tk_ids, tk_classes]: + list (for each timestep) of 1D NDArrays (for each det). + [tk_dets]: list (for each timestep) of lists of detections. + """ + seq_id = self.seq_name2seqid[seq] + # file location + if is_gt: + imgs = self.video2gt_image[seq_id] + else: + imgs = self.tracker_data[tracker]["vids_to_images"][seq_id] + + # convert data to required format + num_timesteps = self.seq_lengths[seq_id] + img_to_timestep = self.seq2images2timestep[seq_id] + data_keys = ["ids", "classes", "dets"] + # if not is_gt: + # data_keys += ["tk_confidences"] + raw_data = {key: [None] * num_timesteps for key in data_keys} + for img in imgs: + # some tracker data contains images without any ground truth info, + # these are ignored + if img["id"] not in img_to_timestep: + continue + t = img_to_timestep[img["id"]] + anns = img["annotations"] + tk_str = utils.get_track_id_str(anns[0]) + + if self.use_mask: + # When using mask, extract segmentation data + raw_data["dets"][t] = [ann.get("segmentation") for ann in anns] + else: + # When using bbox, extract bbox data + raw_data["dets"][t] = np.atleast_2d([ann["bbox"] for ann in anns]).astype( + float + ) + raw_data["ids"][t] = np.atleast_1d([ann[tk_str] for ann in anns]).astype( + int + ) + raw_data["classes"][t] = np.atleast_1d( + [ann["category_id"] for ann in anns] + ).astype(int) + # if not is_gt: + # raw_data["tk_confidences"][t] = np.atleast_1d( + # [ann["score"] for ann in anns] + # ).astype(float) + + for t, d in enumerate(raw_data["dets"]): + if d is None: + raw_data["dets"][t] = np.empty((0, 4)).astype(float) + raw_data["ids"][t] = np.empty(0).astype(int) + raw_data["classes"][t] = np.empty(0).astype(int) + # if not is_gt: + # raw_data["tk_confidences"][t] = np.empty(0) + + if is_gt: + key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"} + else: + key_map = {"ids": "tk_ids", "classes": "tk_classes", "dets": "tk_dets"} + for k, v in key_map.items(): + raw_data[v] = raw_data.pop(k) + + raw_data["num_timesteps"] = num_timesteps + raw_data["seq"] = seq + return raw_data + + def get_preprocessed_seq_data_thr(self, raw_data, cls, assignment=None): + """Preprocess data for a single sequence for a single class. + + Inputs: + raw_data: dict containing the data for the sequence already + read in by get_raw_seq_data(). + cls: class to be evaluated. + Outputs: + gt_ids: + list (for each timestep) of ids of GT tracks + tk_ids: + list (for each timestep) of ids of predicted tracks (all for TP + matching (Det + AssocA)) + tk_overlap_ids: + list (for each timestep) of ids of predicted tracks that overlap + with GTs + tk_dets: + list (for each timestep) of lists of detections that + corresponding to the tk_ids + tk_classes: + list (for each timestep) of lists of classes that corresponding + to the tk_ids + tk_confidences: + list (for each timestep) of lists of classes that corresponding + to the tk_ids + sim_scores: + similarity score between gt_ids and tk_ids. + """ + if cls != "all": + cls_id = self.cls_name2clsid[cls] + + data_keys = [ + "gt_ids", + "tk_ids", + "gt_id_map", + "tk_id_map", + "gt_dets", + "gt_classes", + "gt_class_name", + "tk_overlap_classes", + "tk_overlap_ids", + "tk_class_eval_tk_ids", + "tk_dets", + "tk_classes", + # "tk_confidences", + "tk_exh_ids", + "sim_scores", + ] + data = {key: [None] * raw_data["num_timesteps"] for key in data_keys} + unique_gt_ids = [] + unique_tk_ids = [] + num_gt_dets = 0 + num_tk_cls_dets = 0 + num_tk_overlap_dets = 0 + overlap_ious_thr = 0.5 + loc_and_asso_tk_ids = [] + exh_class_tk_ids = [] + + for t in range(raw_data["num_timesteps"]): + # only extract relevant dets for this class for preproc and eval + if cls == "all": + gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool) + else: + gt_class_mask = np.atleast_1d( + raw_data["gt_classes"][t] == cls_id + ).astype(bool) + + # select GT that is not in the evaluating classes + if assignment is not None and assignment: + all_gt_ids = list(assignment[t].keys()) + gt_ids_in = raw_data["gt_ids"][t][gt_class_mask] + gt_ids_out = set(all_gt_ids) - set(gt_ids_in) + tk_ids_out = set([assignment[t][key] for key in list(gt_ids_out)]) + + # compute overlapped tracks and add their ids to overlap_tk_ids + sim_scores = raw_data["similarity_scores"] + overlap_ids_masks = (sim_scores[t][gt_class_mask] >= overlap_ious_thr).any( + axis=0 + ) + overlap_tk_ids_t = raw_data["tk_ids"][t][overlap_ids_masks] + if assignment is not None and assignment: + data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t) - tk_ids_out) + else: + data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t)) + + loc_and_asso_tk_ids += data["tk_overlap_ids"][t] + + data["tk_exh_ids"][t] = [] + if cls == "all": + continue + + # add the track ids of exclusive annotated class to exh_class_tk_ids + tk_exh_mask = np.atleast_1d(raw_data["tk_classes"][t] == cls_id) + tk_exh_mask = tk_exh_mask.astype(bool) + exh_class_tk_ids_t = raw_data["tk_ids"][t][tk_exh_mask] + exh_class_tk_ids.append(exh_class_tk_ids_t) + data["tk_exh_ids"][t] = exh_class_tk_ids_t + + # remove tk_ids that has been assigned to GT belongs to other classes. + loc_and_asso_tk_ids = list(set(loc_and_asso_tk_ids)) + + # remove all unwanted unmatched tracker detections + for t in range(raw_data["num_timesteps"]): + # add gt to the data + if cls == "all": + gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool) + else: + gt_class_mask = np.atleast_1d( + raw_data["gt_classes"][t] == cls_id + ).astype(bool) + data["gt_classes"][t] = cls_id + data["gt_class_name"][t] = cls + + gt_ids = raw_data["gt_ids"][t][gt_class_mask] + if self.use_mask: + gt_dets = [raw_data['gt_dets'][t][ind] for ind in range(len(gt_class_mask)) if gt_class_mask[ind]] + else: + gt_dets = raw_data["gt_dets"][t][gt_class_mask] + data["gt_ids"][t] = gt_ids + data["gt_dets"][t] = gt_dets + + # filter pred and only keep those that highly overlap with GTs + tk_mask = np.isin( + raw_data["tk_ids"][t], np.array(loc_and_asso_tk_ids), assume_unique=True + ) + tk_overlap_mask = np.isin( + raw_data["tk_ids"][t], + np.array(data["tk_overlap_ids"][t]), + assume_unique=True, + ) + + tk_ids = raw_data["tk_ids"][t][tk_mask] + if self.use_mask: + tk_dets = [raw_data['tk_dets'][t][ind] for ind in range(len(tk_mask)) if + tk_mask[ind]] + else: + tk_dets = raw_data["tk_dets"][t][tk_mask] + + tracker_classes = raw_data["tk_classes"][t][tk_mask] + + # add overlap classes for computing the FP for Cls term + tracker_overlap_classes = raw_data["tk_classes"][t][tk_overlap_mask] + # tracker_confidences = raw_data["tk_confidences"][t][tk_mask] + sim_scores_masked = sim_scores[t][gt_class_mask, :][:, tk_mask] + + # add filtered prediction to the data + data["tk_classes"][t] = tracker_classes + data["tk_overlap_classes"][t] = tracker_overlap_classes + data["tk_ids"][t] = tk_ids + data["tk_dets"][t] = tk_dets + # data["tk_confidences"][t] = tracker_confidences + data["sim_scores"][t] = sim_scores_masked + data["tk_class_eval_tk_ids"][t] = set( + list(data["tk_overlap_ids"][t]) + list(data["tk_exh_ids"][t]) + ) + + # count total number of detections + unique_gt_ids += list(np.unique(data["gt_ids"][t])) + # the unique track ids are for association. + unique_tk_ids += list(np.unique(data["tk_ids"][t])) + + num_tk_overlap_dets += len(data["tk_overlap_ids"][t]) + num_tk_cls_dets += len(data["tk_class_eval_tk_ids"][t]) + num_gt_dets += len(data["gt_ids"][t]) + + # re-label IDs such that there are no empty IDs + if len(unique_gt_ids) > 0: + unique_gt_ids = np.unique(unique_gt_ids) + gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1)) + gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids)) + data["gt_id_map"] = {} + for gt_id in unique_gt_ids: + new_gt_id = gt_id_map[gt_id].astype(int) + data["gt_id_map"][new_gt_id] = gt_id + + for t in range(raw_data["num_timesteps"]): + if len(data["gt_ids"][t]) > 0: + data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int) + + if len(unique_tk_ids) > 0: + unique_tk_ids = np.unique(unique_tk_ids) + tk_id_map = np.nan * np.ones((np.max(unique_tk_ids) + 1)) + tk_id_map[unique_tk_ids] = np.arange(len(unique_tk_ids)) + + data["tk_id_map"] = {} + for track_id in unique_tk_ids: + new_track_id = tk_id_map[track_id].astype(int) + data["tk_id_map"][new_track_id] = track_id + + for t in range(raw_data["num_timesteps"]): + if len(data["tk_ids"][t]) > 0: + data["tk_ids"][t] = tk_id_map[data["tk_ids"][t]].astype(int) + if len(data["tk_overlap_ids"][t]) > 0: + data["tk_overlap_ids"][t] = tk_id_map[ + data["tk_overlap_ids"][t] + ].astype(int) + + # record overview statistics. + data["num_tk_cls_dets"] = num_tk_cls_dets + data["num_tk_overlap_dets"] = num_tk_overlap_dets + data["num_gt_dets"] = num_gt_dets + data["num_tk_ids"] = len(unique_tk_ids) + data["num_gt_ids"] = len(unique_gt_ids) + data["num_timesteps"] = raw_data["num_timesteps"] + data["seq"] = raw_data["seq"] + + self._check_unique_ids(data) + + return data + + @_timing.time + def get_preprocessed_seq_data( + self, raw_data, cls, assignment=None, thresholds=[50, 75] + ): + """Preprocess data for a single sequence for a single class.""" + data = {} + if thresholds is None: + thresholds = [50, 75] + elif isinstance(thresholds, int): + thresholds = [thresholds] + + for thr in thresholds: + assignment_thr = None + if assignment is not None: + assignment_thr = assignment[thr] + data[thr] = self.get_preprocessed_seq_data_thr( + raw_data, cls, assignment_thr + ) + + return data + + def _calculate_similarities(self, gt_dets_t, tk_dets_t): + """Compute similarity scores.""" + if self.use_mask: + similarity_scores = self._calculate_mask_ious(gt_dets_t, tk_dets_t, is_encoded=True, do_ioa=False) + else: + similarity_scores = self._calculate_box_ious(gt_dets_t, tk_dets_t) + return similarity_scores + + def _compute_vid_mappings(self, annotations): + """Computes mappings from videos to corresponding tracks and images.""" + vids_to_tracks = {} + vids_to_imgs = {} + vid_ids = [vid["id"] for vid in self.gt_data["videos"]] + + # compute an mapping from image IDs to images + images = {} + for image in self.gt_data["images"]: + images[image["id"]] = image + + tk_str = utils.get_track_id_str(annotations[0]) + for ann in annotations: + ann["area"] = ann["bbox"][2] * ann["bbox"][3] + + vid = ann["video_id"] + if ann["video_id"] not in vids_to_tracks.keys(): + vids_to_tracks[ann["video_id"]] = list() + if ann["video_id"] not in vids_to_imgs.keys(): + vids_to_imgs[ann["video_id"]] = list() + + # fill in vids_to_tracks + tid = ann[tk_str] + exist_tids = [track["id"] for track in vids_to_tracks[vid]] + try: + index1 = exist_tids.index(tid) + except ValueError: + index1 = -1 + if tid not in exist_tids: + curr_track = { + "id": tid, + "category_id": ann["category_id"], + "video_id": vid, + "annotations": [ann], + } + vids_to_tracks[vid].append(curr_track) + else: + vids_to_tracks[vid][index1]["annotations"].append(ann) + + # fill in vids_to_imgs + img_id = ann["image_id"] + exist_img_ids = [img["id"] for img in vids_to_imgs[vid]] + try: + index2 = exist_img_ids.index(img_id) + except ValueError: + index2 = -1 + if index2 == -1: + curr_img = {"id": img_id, "annotations": [ann]} + vids_to_imgs[vid].append(curr_img) + else: + vids_to_imgs[vid][index2]["annotations"].append(ann) + + # sort annotations by frame index and compute track area + for vid, tracks in vids_to_tracks.items(): + for track in tracks: + track["annotations"] = sorted( + track["annotations"], + key=lambda x: images[x["image_id"]]["frame_id"], + ) + # compute average area + track["area"] = sum(x["area"] for x in track["annotations"]) / len( + track["annotations"] + ) + + # ensure all videos are present + for vid_id in vid_ids: + if vid_id not in vids_to_tracks.keys(): + vids_to_tracks[vid_id] = [] + if vid_id not in vids_to_imgs.keys(): + vids_to_imgs[vid_id] = [] + + return vids_to_tracks, vids_to_imgs + + def _compute_image_to_timestep_mappings(self): + """Computes a mapping from images to timestep in sequence.""" + images = {} + for image in self.gt_data["images"]: + images[image["id"]] = image + + seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]} + for vid in seq_to_imgs_to_timestep: + curr_imgs = [img["id"] for img in self.video2gt_image[vid]] + curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_id"]) + seq_to_imgs_to_timestep[vid] = { + curr_imgs[i]: i for i in range(len(curr_imgs)) + } + + return seq_to_imgs_to_timestep + + def _limit_dets_per_image(self, annotations): + """Limits the number of detections for each image. + + Adapted from https://github.com/TAO-Dataset/. + """ + max_dets = self.config["MAX_DETECTIONS"] + img_ann = defaultdict(list) + for ann in annotations: + img_ann[ann["image_id"]].append(ann) + + for img_id, _anns in img_ann.items(): + if len(_anns) <= max_dets: + continue + _anns = sorted(_anns, key=lambda x: x["score"], reverse=True) + img_ann[img_id] = _anns[:max_dets] + + return [ann for anns in img_ann.values() for ann in anns] + + def _fill_video_ids_inplace(self, annotations): + """Fills in missing video IDs inplace. + + Adapted from https://github.com/TAO-Dataset/. + """ + missing_video_id = [x for x in annotations if "video_id" not in x] + if missing_video_id: + image_id_to_video_id = { + x["id"]: x["video_id"] for x in self.gt_data["images"] + } + for x in missing_video_id: + x["video_id"] = image_id_to_video_id[x["image_id"]] + + @staticmethod + def _make_tk_ids_unique(annotations): + """Makes track IDs unqiue over the whole annotation set. + + Adapted from https://github.com/TAO-Dataset/. + """ + track_id_videos = {} + track_ids_to_update = set() + max_track_id = 0 + + tk_str = utils.get_track_id_str(annotations[0]) + for ann in annotations: + t = int(ann[tk_str]) + if t not in track_id_videos: + track_id_videos[t] = ann["video_id"] + + if ann["video_id"] != track_id_videos[t]: + # track id is assigned to multiple videos + track_ids_to_update.add(t) + max_track_id = max(max_track_id, t) + + if track_ids_to_update: + print("true") + next_id = itertools.count(max_track_id + 1) + new_tk_ids = defaultdict(lambda: next(next_id)) + for ann in annotations: + t = ann[tk_str] + v = ann["video_id"] + if t in track_ids_to_update: + ann[tk_str] = new_tk_ids[t, v] + return len(track_ids_to_update) diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/tao.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/tao.py new file mode 100644 index 0000000..7d2bbcc --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/datasets/tao.py @@ -0,0 +1,659 @@ +# fmt: off +# flake8: noqa + +"""TAO Dataset.""" +import copy +import itertools +import json +import os +from collections import defaultdict + +import numpy as np + +from .. import _timing +from ..config import get_default_dataset_config, init_config +from ..utils import TrackEvalException +from ._base_dataset import _BaseDataset + + +class TAO(_BaseDataset): + """Dataset class for TAO tracking""" + + def __init__(self, config=None): + """Initialize dataset, checking that all required files are present.""" + super().__init__() + # Fill non-given config values with defaults + self.config = init_config(config, get_default_dataset_config(), self.get_name()) + self.gt_fol = self.config["GT_FOLDER"] + self.tracker_fol = self.config["TRACKERS_FOLDER"] + self.should_classes_combine = True + self.use_super_categories = False + self.use_mask = self.config["USE_MASK"] + + + self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"] + self.output_fol = self.config["OUTPUT_FOLDER"] + if self.output_fol is None: + self.output_fol = self.tracker_fol + self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"] + + if self.gt_fol.endswith(".json"): + self.gt_data = json.load(open(self.gt_fol, "r")) + else: + gt_dir_files = [ + file for file in os.listdir(self.gt_fol) if file.endswith(".json") + ] + if len(gt_dir_files) != 1: + raise TrackEvalException( + f"{self.gt_fol} does not contain exactly one json file." + ) + + with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f: + self.gt_data = json.load(f) + + # merge categories marked with a merged tag in TAO dataset + self._merge_categories(self.gt_data["annotations"] + self.gt_data["tracks"]) + + # get sequences to eval and sequence information + self.seq_list = [ + vid["name"].replace("/", "-") for vid in self.gt_data["videos"] + ] + self.seq_name2seqid = { + vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"] + } + # compute mappings from videos to annotation data + self.video2gt_track, self.video2gt_image = self._compute_vid_mappings( + self.gt_data["annotations"] + ) + # compute sequence lengths + self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]} + for img in self.gt_data["images"]: + self.seq_lengths[img["video_id"]] += 1 + self.seq2images2timestep = self._compute_image_to_timestep_mappings() + self.seq2cls = { + vid["id"]: { + "pos_cat_ids": list( + {track["category_id"] for track in self.video2gt_track[vid["id"]]} + ), + "neg_cat_ids": vid["neg_category_ids"], + "not_exh_labeled_cat_ids": vid["not_exhaustive_category_ids"], + } + for vid in self.gt_data["videos"] + } + + # Get classes to eval + considered_vid_ids = [self.seq_name2seqid[vid] for vid in self.seq_list] + seen_cats = set( + [ + cat_id + for vid_id in considered_vid_ids + for cat_id in self.seq2cls[vid_id]["pos_cat_ids"] + ] + ) + # only classes with ground truth are evaluated in TAO + self.valid_classes = [ + cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats + ] + cls_name2clsid_map = { + cls["name"]: cls["id"] for cls in self.gt_data["categories"] + } + + if self.config["CLASSES_TO_EVAL"]: + self.class_list = [ + cls.lower() if cls.lower() in self.valid_classes else None + for cls in self.config["CLASSES_TO_EVAL"] + ] + if not all(self.class_list): + valid_cls = ", ".join(self.valid_classes) + raise TrackEvalException( + "Attempted to evaluate an invalid class. Only classes " + f"{valid_cls} are valid (classes present in ground truth" + " data)." + ) + else: + self.class_list = [cls for cls in self.valid_classes] + self.cls_name2clsid = { + k: v for k, v in cls_name2clsid_map.items() if k in self.class_list + } + self.clsid2cls_name = { + v: k for k, v in cls_name2clsid_map.items() if k in self.class_list + } + # get trackers to eval + print(self.config["TRACKERS_TO_EVAL"] ) + if self.config["TRACKERS_TO_EVAL"] is None: + self.tracker_list = os.listdir(self.tracker_fol) + else: + self.tracker_list = self.config["TRACKERS_TO_EVAL"] + + if self.config["TRACKER_DISPLAY_NAMES"] is None: + self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list)) + elif (self.config["TRACKERS_TO_EVAL"] is not None) and ( + len(self.config["TK_DISPLAY_NAMES"]) == len(self.tracker_list) + ): + self.tracker_to_disp = dict( + zip(self.tracker_list, self.config["TK_DISPLAY_NAMES"]) + ) + else: + raise TrackEvalException( + "List of tracker files and tracker display names do not match." + ) + + self.tracker_data = {tracker: dict() for tracker in self.tracker_list} + + for tracker in self.tracker_list: + if self.tracker_sub_fol.endswith(".json"): + with open(os.path.join(self.tracker_sub_fol)) as f: + curr_data = json.load(f) + else: + tr_dir = os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol) + tr_dir_files = [ + file for file in os.listdir(tr_dir) if file.endswith(".json") + ] + if len(tr_dir_files) != 1: + raise TrackEvalException( + f"{tr_dir} does not contain exactly one json file." + ) + with open(os.path.join(tr_dir, tr_dir_files[0])) as f: + curr_data = json.load(f) + + # limit detections if MAX_DETECTIONS > 0 + if self.config["MAX_DETECTIONS"]: + curr_data = self._limit_dets_per_image(curr_data) + + # fill missing video ids + self._fill_video_ids_inplace(curr_data) + + # make track ids unique over whole evaluation set + self._make_tk_ids_unique(curr_data) + + # merge categories marked with a merged tag in TAO dataset + self._merge_categories(curr_data) + + # get tracker sequence information + curr_vids2tracks, curr_vids2images = self._compute_vid_mappings(curr_data) + self.tracker_data[tracker]["vids_to_tracks"] = curr_vids2tracks + self.tracker_data[tracker]["vids_to_images"] = curr_vids2images + + def get_display_name(self, tracker): + return self.tracker_to_disp[tracker] + + def _load_raw_file(self, tracker, seq, is_gt): + """Load a file (gt or tracker) in the TAO format + + If is_gt, this returns a dict which contains the fields: + [gt_ids, gt_classes]: + list (for each timestep) of 1D NDArrays (for each det). + [gt_dets]: list (for each timestep) of lists of detections. + + if not is_gt, this returns a dict which contains the fields: + [tk_ids, tk_classes, tk_confidences]: + list (for each timestep) of 1D NDArrays (for each det). + [tk_dets]: list (for each timestep) of lists of detections. + """ + seq_id = self.seq_name2seqid[seq] + # file location + if is_gt: + imgs = self.video2gt_image[seq_id] + else: + imgs = self.tracker_data[tracker]["vids_to_images"][seq_id] + + # convert data to required format + num_timesteps = self.seq_lengths[seq_id] + img_to_timestep = self.seq2images2timestep[seq_id] + data_keys = ["ids", "classes", "dets"] + if not is_gt: + data_keys += ["tk_confidences"] + raw_data = {key: [None] * num_timesteps for key in data_keys} + for img in imgs: + # some tracker data contains images without any ground truth info, + # these are ignored + if img["id"] not in img_to_timestep: + continue + t = img_to_timestep[img["id"]] + anns = img["annotations"] + if self.use_mask: + # When using mask, extract segmentation data + raw_data["dets"][t] = [ann.get("segmentation") for ann in anns] + else: + # When using bbox, extract bbox data + raw_data["dets"][t] = np.atleast_2d([ann["bbox"] for ann in anns]).astype( + float + ) + raw_data["ids"][t] = np.atleast_1d( + [ann["track_id"] for ann in anns] + ).astype(int) + raw_data["classes"][t] = np.atleast_1d( + [ann["category_id"] for ann in anns] + ).astype(int) + if not is_gt: + raw_data["tk_confidences"][t] = np.atleast_1d( + [ann["score"] for ann in anns] + ).astype(float) + + for t, d in enumerate(raw_data["dets"]): + if d is None: + raw_data["dets"][t] = np.empty((0, 4)).astype(float) + raw_data["ids"][t] = np.empty(0).astype(int) + raw_data["classes"][t] = np.empty(0).astype(int) + if not is_gt: + raw_data["tk_confidences"][t] = np.empty(0) + + if is_gt: + key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"} + else: + key_map = {"ids": "tk_ids", "classes": "tk_classes", "dets": "tk_dets"} + for k, v in key_map.items(): + raw_data[v] = raw_data.pop(k) + + raw_data["num_timesteps"] = num_timesteps + raw_data["neg_cat_ids"] = self.seq2cls[seq_id]["neg_cat_ids"] + raw_data["not_exh_labeled_cls"] = self.seq2cls[seq_id][ + "not_exh_labeled_cat_ids" + ] + raw_data["seq"] = seq + return raw_data + + def get_preprocessed_seq_data_thr(self, raw_data, cls, assignment=None): + """Preprocess data for a single sequence for a single class. + + Inputs: + raw_data: dict containing the data for the sequence already + read in by get_raw_seq_data(). + cls: class to be evaluated. + Outputs: + gt_ids: + list (for each timestep) of ids of GT tracks + tk_ids: + list (for each timestep) of ids of predicted tracks (all for TP + matching (Det + AssocA)) + tk_overlap_ids: + list (for each timestep) of ids of predicted tracks that overlap + with GTs + tk_neg_ids: + list (for each timestep) of ids of predicted tracks that with + the class id on the negative list for the current sequence. + tk_exh_ids: + list (for each timestep) of ids of predicted tracks that do not + overlap with existing GTs but have the class id on the + exhaustive annotated class list for the current sequence. + tk_dets: + list (for each timestep) of lists of detections that + corresponding to the tk_ids + tk_classes: + list (for each timestep) of lists of classes that corresponding + to the tk_ids + tk_confidences: + list (for each timestep) of lists of classes that corresponding + to the tk_ids + sim_scores: + similarity score between gt_ids and tk_ids. + """ + if cls != "all": + cls_id = self.cls_name2clsid[cls] + + data_keys = [ + "gt_ids", + "tk_ids", + "gt_id_map", + "tk_id_map", + "gt_dets", + "gt_classes", + "gt_class_name", + "tk_overlap_classes", + "tk_overlap_ids", + "tk_neg_ids", + "tk_exh_ids", + "tk_class_eval_tk_ids", + "tk_dets", + "tk_classes", + "tk_confidences", + "sim_scores", + ] + data = {key: [None] * raw_data["num_timesteps"] for key in data_keys} + unique_gt_ids = [] + unique_tk_ids = [] + num_gt_dets = 0 + num_tk_cls_dets = 0 + num_tk_overlap_dets = 0 + overlap_ious_thr = 0.5 + loc_and_asso_tk_ids = [] + + for t in range(raw_data["num_timesteps"]): + # only extract relevant dets for this class for preproc and eval + if cls == "all": + gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool) + else: + gt_class_mask = np.atleast_1d( + raw_data["gt_classes"][t] == cls_id + ).astype(bool) + + # select GT that is not in the evaluating classes + if assignment is not None and assignment: + all_gt_ids = list(assignment[t].keys()) + gt_ids_in = raw_data["gt_ids"][t][gt_class_mask] + gt_ids_out = set(all_gt_ids) - set(gt_ids_in) + tk_ids_out = set([assignment[t][key] for key in list(gt_ids_out)]) + + # compute overlapped tracks and add their ids to overlap_tk_ids + sim_scores = raw_data["similarity_scores"] + overlap_ids_masks = (sim_scores[t][gt_class_mask] >= overlap_ious_thr).any( + axis=0 + ) + overlap_tk_ids_t = raw_data["tk_ids"][t][overlap_ids_masks] + if assignment is not None and assignment: + data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t) - tk_ids_out) + else: + data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t)) + + loc_and_asso_tk_ids += data["tk_overlap_ids"][t] + + data["tk_exh_ids"][t] = [] + data["tk_neg_ids"][t] = [] + + if cls == "all": + continue + + # remove tk_ids that has been assigned to GT belongs to other classes. + loc_and_asso_tk_ids = list(set(loc_and_asso_tk_ids)) + + # remove all unwanted unmatched tracker detections + for t in range(raw_data["num_timesteps"]): + # add gt to the data + if cls == "all": + gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool) + else: + gt_class_mask = np.atleast_1d( + raw_data["gt_classes"][t] == cls_id + ).astype(bool) + data["gt_classes"][t] = cls_id + data["gt_class_name"][t] = cls + + gt_ids = raw_data["gt_ids"][t][gt_class_mask] + if self.use_mask: + gt_dets = [raw_data['gt_dets'][t][ind] for ind in range(len(gt_class_mask)) if gt_class_mask[ind]] + else: + gt_dets = raw_data["gt_dets"][t][gt_class_mask] + data["gt_ids"][t] = gt_ids + data["gt_dets"][t] = gt_dets + + # filter pred and only keep those that highly overlap with GTs + tk_mask = np.isin( + raw_data["tk_ids"][t], np.array(loc_and_asso_tk_ids), assume_unique=True + ) + tk_overlap_mask = np.isin( + raw_data["tk_ids"][t], + np.array(data["tk_overlap_ids"][t]), + assume_unique=True, + ) + + tk_ids = raw_data["tk_ids"][t][tk_mask] + if self.use_mask: + tk_dets = [raw_data['tk_dets'][t][ind] for ind in range(len(tk_mask)) if + tk_mask[ind]] + else: + tk_dets = raw_data["tk_dets"][t][tk_mask] + tracker_classes = raw_data["tk_classes"][t][tk_mask] + + # add overlap classes for computing the FP for Cls term + tracker_overlap_classes = raw_data["tk_classes"][t][tk_overlap_mask] + tracker_confidences = raw_data["tk_confidences"][t][tk_mask] + sim_scores_masked = sim_scores[t][gt_class_mask, :][:, tk_mask] + + # add filtered prediction to the data + data["tk_classes"][t] = tracker_classes + data["tk_overlap_classes"][t] = tracker_overlap_classes + data["tk_ids"][t] = tk_ids + data["tk_dets"][t] = tk_dets + data["tk_confidences"][t] = tracker_confidences + data["sim_scores"][t] = sim_scores_masked + data["tk_class_eval_tk_ids"][t] = set( + list(data["tk_overlap_ids"][t]) + + list(data["tk_neg_ids"][t]) + + list(data["tk_exh_ids"][t]) + ) + + # count total number of detections + unique_gt_ids += list(np.unique(data["gt_ids"][t])) + # the unique track ids are for association. + unique_tk_ids += list(np.unique(data["tk_ids"][t])) + + num_tk_overlap_dets += len(data["tk_overlap_ids"][t]) + num_tk_cls_dets += len(data["tk_class_eval_tk_ids"][t]) + num_gt_dets += len(data["gt_ids"][t]) + + # re-label IDs such that there are no empty IDs + if len(unique_gt_ids) > 0: + unique_gt_ids = np.unique(unique_gt_ids) + gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1)) + gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids)) + data["gt_id_map"] = {} + for gt_id in unique_gt_ids: + new_gt_id = gt_id_map[gt_id].astype(int) + data["gt_id_map"][new_gt_id] = gt_id + + for t in range(raw_data["num_timesteps"]): + if len(data["gt_ids"][t]) > 0: + data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int) + + if len(unique_tk_ids) > 0: + unique_tk_ids = np.unique(unique_tk_ids) + tk_id_map = np.nan * np.ones((np.max(unique_tk_ids) + 1)) + tk_id_map[unique_tk_ids] = np.arange(len(unique_tk_ids)) + + data["tk_id_map"] = {} + for track_id in unique_tk_ids: + new_track_id = tk_id_map[track_id].astype(int) + data["tk_id_map"][new_track_id] = track_id + + for t in range(raw_data["num_timesteps"]): + if len(data["tk_ids"][t]) > 0: + data["tk_ids"][t] = tk_id_map[data["tk_ids"][t]].astype(int) + if len(data["tk_overlap_ids"][t]) > 0: + data["tk_overlap_ids"][t] = tk_id_map[ + data["tk_overlap_ids"][t] + ].astype(int) + + # record overview statistics. + data["num_tk_cls_dets"] = num_tk_cls_dets + data["num_tk_overlap_dets"] = num_tk_overlap_dets + data["num_gt_dets"] = num_gt_dets + data["num_tk_ids"] = len(unique_tk_ids) + data["num_gt_ids"] = len(unique_gt_ids) + data["num_timesteps"] = raw_data["num_timesteps"] + data["seq"] = raw_data["seq"] + + self._check_unique_ids(data) + + return data + + @_timing.time + def get_preprocessed_seq_data( + self, raw_data, cls, assignment=None, thresholds=[50, 75] + ): + """Preprocess data for a single sequence for a single class.""" + data = {} + if thresholds is None: + thresholds = [50] + elif isinstance(thresholds, int): + thresholds = [thresholds] + + for thr in thresholds: + assignment_thr = None + if assignment is not None: + assignment_thr = assignment[thr] + data[thr] = self.get_preprocessed_seq_data_thr( + raw_data, cls, assignment_thr + ) + + return data + + def _calculate_similarities(self, gt_dets_t, tk_dets_t): + """Compute similarity scores.""" + if self.use_mask: + similarity_scores = self._calculate_mask_ious(gt_dets_t, tk_dets_t, is_encoded=True, do_ioa=False) + else: + similarity_scores = self._calculate_box_ious(gt_dets_t, tk_dets_t) + return similarity_scores + + def _merge_categories(self, annotations): + """Merges categories with a merged tag. + + Adapted from https://github.com/TAO-Dataset. + """ + merge_map = {} + for category in self.gt_data["categories"]: + if "merged" in category: + for to_merge in category["merged"]: + merge_map[to_merge["id"]] = category["id"] + + for ann in annotations: + ann["category_id"] = merge_map.get(ann["category_id"], ann["category_id"]) + + def _compute_vid_mappings(self, annotations): + """Computes mappings from videos to corresponding tracks and images.""" + vids_to_tracks = {} + vids_to_imgs = {} + vid_ids = [vid["id"] for vid in self.gt_data["videos"]] + + # compute an mapping from image IDs to images + images = {} + for image in self.gt_data["images"]: + images[image["id"]] = image + + for ann in annotations: + ann["area"] = ann["bbox"][2] * ann["bbox"][3] + + vid = ann["video_id"] + if ann["video_id"] not in vids_to_tracks.keys(): + vids_to_tracks[ann["video_id"]] = list() + if ann["video_id"] not in vids_to_imgs.keys(): + vids_to_imgs[ann["video_id"]] = list() + + # fill in vids_to_tracks + tid = ann["track_id"] + exist_tids = [track["id"] for track in vids_to_tracks[vid]] + try: + index1 = exist_tids.index(tid) + except ValueError: + index1 = -1 + if tid not in exist_tids: + curr_track = { + "id": tid, + "category_id": ann["category_id"], + "video_id": vid, + "annotations": [ann], + } + vids_to_tracks[vid].append(curr_track) + else: + vids_to_tracks[vid][index1]["annotations"].append(ann) + + # fill in vids_to_imgs + img_id = ann["image_id"] + exist_img_ids = [img["id"] for img in vids_to_imgs[vid]] + try: + index2 = exist_img_ids.index(img_id) + except ValueError: + index2 = -1 + if index2 == -1: + curr_img = {"id": img_id, "annotations": [ann]} + vids_to_imgs[vid].append(curr_img) + else: + vids_to_imgs[vid][index2]["annotations"].append(ann) + + # sort annotations by frame index and compute track area + for vid, tracks in vids_to_tracks.items(): + for track in tracks: + track["annotations"] = sorted( + track["annotations"], + key=lambda x: images[x["image_id"]]["frame_index"], + ) + # compute average area + track["area"] = sum(x["area"] for x in track["annotations"]) / len( + track["annotations"] + ) + + # ensure all videos are present + for vid_id in vid_ids: + if vid_id not in vids_to_tracks.keys(): + vids_to_tracks[vid_id] = [] + if vid_id not in vids_to_imgs.keys(): + vids_to_imgs[vid_id] = [] + + return vids_to_tracks, vids_to_imgs + + def _compute_image_to_timestep_mappings(self): + """Computes a mapping from images to timestep in sequence.""" + images = {} + for image in self.gt_data["images"]: + images[image["id"]] = image + + seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]} + for vid in seq_to_imgs_to_timestep: + curr_imgs = [img["id"] for img in self.video2gt_image[vid]] + curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_index"]) + seq_to_imgs_to_timestep[vid] = { + curr_imgs[i]: i for i in range(len(curr_imgs)) + } + + return seq_to_imgs_to_timestep + + def _limit_dets_per_image(self, annotations): + """Limits the number of detections for each image. + + Adapted from https://github.com/TAO-Dataset/. + """ + max_dets = self.config["MAX_DETECTIONS"] + img_ann = defaultdict(list) + for ann in annotations: + img_ann[ann["image_id"]].append(ann) + + for img_id, _anns in img_ann.items(): + if len(_anns) <= max_dets: + continue + _anns = sorted(_anns, key=lambda x: x["score"], reverse=True) + img_ann[img_id] = _anns[:max_dets] + + return [ann for anns in img_ann.values() for ann in anns] + + def _fill_video_ids_inplace(self, annotations): + """Fills in missing video IDs inplace. + + Adapted from https://github.com/TAO-Dataset/. + """ + missing_video_id = [x for x in annotations if "video_id" not in x] + if missing_video_id: + image_id_to_video_id = { + x["id"]: x["video_id"] for x in self.gt_data["images"] + } + for x in missing_video_id: + x["video_id"] = image_id_to_video_id[x["image_id"]] + + @staticmethod + def _make_tk_ids_unique(annotations): + """Makes track IDs unqiue over the whole annotation set. + + Adapted from https://github.com/TAO-Dataset/. + """ + track_id_videos = {} + track_ids_to_update = set() + max_track_id = 0 + for ann in annotations: + t = ann["track_id"] + if t not in track_id_videos: + track_id_videos[t] = ann["video_id"] + + if ann["video_id"] != track_id_videos[t]: + # track id is assigned to multiple videos + track_ids_to_update.add(t) + max_track_id = max(max_track_id, t) + + if track_ids_to_update: + print("true") + next_id = itertools.count(max_track_id + 1) + new_tk_ids = defaultdict(lambda: next(next_id)) + for ann in annotations: + t = ann["track_id"] + v = ann["video_id"] + if t in track_ids_to_update: + ann["track_id"] = new_tk_ids[t, v] + return len(track_ids_to_update) diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/eval.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/eval.py new file mode 100644 index 0000000..336f10b --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/eval.py @@ -0,0 +1,275 @@ +# fmt: off +# flake8: noqa + +import copy +import os +import pickle +import time +import traceback +from functools import partial +from multiprocessing.pool import Pool + +import numpy as np + +from . import _timing, utils +from .config import get_default_eval_config, init_config +from .utils import TrackEvalException + + +class Evaluator: + """Evaluator class for evaluating different metrics for each datasets.""" + + def __init__(self, config=None): + """Initialize the evaluator with a config file.""" + self.config = init_config(config, get_default_eval_config(), "Eval") + # Only run timing analysis if not run in parallel. + if self.config["TIME_PROGRESS"] and not self.config["USE_PARALLEL"]: + _timing.DO_TIMING = True + if self.config["DISPLAY_LESS_PROGRESS"]: + _timing.DISPLAY_LESS_PROGRESS = True + + @_timing.time + def evaluate(self, dataset_list, metrics_list): + """Evaluate a set of metrics on a set of datasets.""" + config = self.config + metrics_list = metrics_list + metric_names = utils.validate_metrics_list(metrics_list) + dataset_names = [dataset.get_name() for dataset in dataset_list] + output_res = {} + output_msg = {} + + for dataset, dname in zip(dataset_list, dataset_names): + # Get dataset info about what to evaluate + output_res[dname] = {} + output_msg[dname] = {} + tracker_list, seq_list, class_list = dataset.get_eval_info() + print( + f"\nEvaluating {len(tracker_list)} tracker(s) on " + f"{len(seq_list)} sequence(s) for {len(class_list)} class(es)" + f" on {dname} dataset using the following " + f'metrics: {", ".join(metric_names)}\n' + ) + + # Evaluate each tracker + for tracker in tracker_list: + try: + output_res, output_msg = self.evaluate_tracker( + tracker, + dataset, + dname, + class_list, + metrics_list, + metric_names, + seq_list, + output_res, + output_msg, + ) + except Exception as err: + output_res[dname][tracker] = None + if type(err) == TrackEvalException: + output_msg[dname][tracker] = str(err) + else: + output_msg[dname][tracker] = "Unknown error occurred." + print("Tracker %s was unable to be evaluated." % tracker) + print(err) + traceback.print_exc() + if config["LOG_ON_ERROR"] is not None: + with open(config["LOG_ON_ERROR"], "a") as f: + print(dname, file=f) + print(tracker, file=f) + print(traceback.format_exc(), file=f) + print("\n\n\n", file=f) + if config["BREAK_ON_ERROR"]: + raise err + elif config["RETURN_ON_ERROR"]: + return output_res, output_msg + + return output_res, output_msg + + def evaluate_tracker( + self, + tracker, + dataset, + dname, + class_list, + metrics_list, + metric_names, + seq_list, + output_res, + output_msg, + ): + """Evaluate each sequence in parallel or in series.""" + print("\nEvaluating %s\n" % tracker) + time_start = time.time() + config = self.config + if config["USE_PARALLEL"]: + with Pool(config["NUM_PARALLEL_CORES"]) as pool: + _eval_sequence = partial( + eval_sequence, + dataset=dataset, + tracker=tracker, + class_list=class_list, + metrics_list=metrics_list, + metric_names=metric_names, + ) + results = pool.map(_eval_sequence, seq_list) + res = dict(zip(seq_list, results)) + else: + res = {} + for curr_seq in sorted(seq_list): + res[curr_seq] = eval_sequence( + curr_seq, dataset, tracker, class_list, metrics_list, metric_names + ) + + + # collecting combined cls keys (cls averaged, det averaged, super classes) + cls_keys = [] + res["COMBINED_SEQ"] = {} + # combine sequences for each class + for c_cls in class_list: + res["COMBINED_SEQ"][c_cls] = {} + for metric, mname in zip(metrics_list, metric_names): + curr_res = { + seq_key: seq_value[c_cls][mname] + for seq_key, seq_value in res.items() + if seq_key != "COMBINED_SEQ" + } + # combine results over all sequences and then over all classes + res["COMBINED_SEQ"][c_cls][mname] = metric.combine_sequences(curr_res) + + # combine classes + if dataset.should_classes_combine: + if config["OUTPUT_PER_SEQ_RES"]: + video_keys = res.keys() + else: + video_keys = ["COMBINED_SEQ"] + for v_key in video_keys: + cls_keys += ["average"] + res[v_key]["average"] = {} + for metric, mname in zip(metrics_list, metric_names): + cls_res = { + cls_key: cls_value[mname] + for cls_key, cls_value in res[v_key].items() + if cls_key not in cls_keys + } + res[v_key]["average"][ + mname + ] = metric.combine_classes_class_averaged( + cls_res, ignore_empty=True + ) + + # combine classes to super classes + if dataset.use_super_categories: + for cat, sub_cats in dataset.super_categories.items(): + cls_keys.append(cat) + res["COMBINED_SEQ"][cat] = {} + for metric, mname in zip(metrics_list, metric_names): + cat_res = { + cls_key: cls_value[mname] + for cls_key, cls_value in res["COMBINED_SEQ"].items() + if cls_key in sub_cats + } + res["COMBINED_SEQ"][cat][ + mname + ] = metric.combine_classes_det_averaged(cat_res) + # Print and output results in various formats + if config["TIME_PROGRESS"]: + print( + f"\nAll sequences for {tracker} finished in" + f" {time.time() - time_start} seconds" + ) + output_fol = dataset.get_output_fol(tracker) + os.makedirs(output_fol, exist_ok=True) + + # take a mean of each field of each thr + if config["OUTPUT_PER_SEQ_RES"]: + all_res = copy.deepcopy(res) + summary_keys = res.keys() + else: + all_res = copy.deepcopy(res["COMBINED_SEQ"]) + summary_keys = ["COMBINED_SEQ"] + thr_key_list = [50] + for s_key in summary_keys: + for metric, mname in zip(metrics_list, metric_names): + if mname != "TETA": + if s_key == "COMBINED_SEQ": + metric.print_table( + {"COMBINED_SEQ": res["COMBINED_SEQ"][cls_keys[0]][mname]}, + tracker, + cls_keys[0], + ) + continue + + for c_cls in res[s_key].keys(): + for thr in thr_key_list: + all_res[s_key][c_cls][mname][thr] = metric._summary_row( + res[s_key][c_cls][mname][thr] + ) + x = ( + np.array(list(all_res[s_key][c_cls]["TETA"].values())) + .astype("float") + .mean(axis=0) + ) + all_res_summary = list(x.round(decimals=2).astype("str")) + all_res[s_key][c_cls][mname]["ALL"] = all_res_summary + if config["OUTPUT_SUMMARY"] and s_key == "COMBINED_SEQ": + for t in thr_key_list: + metric.print_summary_table( + all_res[s_key][cls_keys[0]][mname][t], + t, + tracker, + cls_keys[0], + ) + + if config["OUTPUT_TEM_RAW_DATA"]: + out_file = os.path.join(output_fol, "teta_summary_results.pth") + pickle.dump(all_res, open(out_file, "wb")) + print("Saved the TETA summary results.") + + # output + output_res[dname][mname] = all_res[s_key][cls_keys[0]][mname][t] + output_msg[dname][tracker] = "Success" + + return output_res, output_msg + + +@_timing.time +def eval_sequence(seq, dataset, tracker, class_list, metrics_list, metric_names): + """Function for evaluating a single sequence.""" + raw_data = dataset.get_raw_seq_data(tracker, seq) + seq_res = {} + + if "TETA" in metric_names: + thresholds = [50] + data_all_class = dataset.get_preprocessed_seq_data( + raw_data, "all", thresholds=thresholds + ) + teta = metrics_list[metric_names.index("TETA")] + assignment = teta.compute_global_assignment(data_all_class) + + # create a dict to save Cls_FP for each class in different thr. + cls_fp = { + key: { + cls: np.zeros((len(np.arange(0.5, 0.99, 0.05)))) for cls in class_list + } + for key in thresholds + } + + for cls in class_list: + seq_res[cls] = {} + data = dataset.get_preprocessed_seq_data(raw_data, cls, assignment, thresholds) + + for metric, mname in zip(metrics_list, metric_names): + if mname == "TETA": + seq_res[cls][mname], cls_fp, _ = metric.eval_sequence( + data, cls, dataset.clsid2cls_name, cls_fp + ) + else: + seq_res[cls][mname] = metric.eval_sequence(data) + + if "TETA" in metric_names: + for thr in thresholds: + for cls in class_list: + seq_res[cls]["TETA"][thr]["Cls_FP"] += cls_fp[thr][cls] + + return seq_res diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/__init__.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/__init__.py new file mode 100644 index 0000000..8352cd4 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/__init__.py @@ -0,0 +1,4 @@ +# fmt: off +# flake8: noqa + +from .teta import TETA diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/_base_metric.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/_base_metric.py new file mode 100644 index 0000000..521a76d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/_base_metric.py @@ -0,0 +1,148 @@ +# fmt: off +# flake8: noqa + +from abc import ABC, abstractmethod + +import numpy as np + +from .. import _timing +from ..utils import TrackEvalException + + +class _BaseMetric(ABC): + @abstractmethod + def __init__(self): + self.plottable = False + self.integer_fields = [] + self.float_fields = [] + self.array_labels = [] + self.integer_array_fields = [] + self.float_array_fields = [] + self.fields = [] + self.summary_fields = [] + self.registered = False + + ##################################################################### + # Abstract functions for subclasses to implement + + @_timing.time + @abstractmethod + def eval_sequence(self, data): + ... + + @abstractmethod + def combine_sequences(self, all_res): + ... + + @abstractmethod + def combine_classes_class_averaged(self, all_res, ignore_empty=False): + ... + + @abstractmethod + def combine_classes_det_averaged(self, all_res): + ... + + def plot_single_tracker_results(self, all_res, tracker, output_folder, cls): + """Plot results, only valid for metrics with self.plottable.""" + if self.plottable: + raise NotImplementedError( + f"plot_results is not implemented for metric {self.get_name()}" + ) + else: + pass + + ##################################################################### + # Helper functions which are useful for all metrics: + + @classmethod + def get_name(cls): + return cls.__name__ + + @staticmethod + def _combine_sum(all_res, field): + """Combine sequence results via sum""" + return sum([all_res[k][field] for k in all_res.keys()]) + + @staticmethod + def _combine_weighted_av(all_res, field, comb_res, weight_field): + """Combine sequence results via weighted average.""" + return sum( + [all_res[k][field] * all_res[k][weight_field] for k in all_res.keys()] + ) / np.maximum(1.0, comb_res[weight_field]) + + def print_table(self, table_res, tracker, cls): + """Print table of results for all sequences.""" + print("") + metric_name = self.get_name() + self._row_print( + [metric_name + ": " + tracker + "-" + cls] + self.summary_fields + ) + for seq, results in sorted(table_res.items()): + if seq == "COMBINED_SEQ": + continue + summary_res = self._summary_row(results) + self._row_print([seq] + summary_res) + summary_res = self._summary_row(table_res["COMBINED_SEQ"]) + self._row_print(["COMBINED"] + summary_res) + + def _summary_row(self, results_): + vals = [] + for h in self.summary_fields: + if h in self.float_array_fields: + vals.append("{0:1.5g}".format(100 * np.mean(results_[h]))) + elif h in self.float_fields: + vals.append("{0:1.5g}".format(100 * float(results_[h]))) + elif h in self.integer_fields: + vals.append("{0:d}".format(int(results_[h]))) + else: + raise NotImplementedError( + "Summary function not implemented for this field type." + ) + return vals + + @staticmethod + def _row_print(*argv): + """Print results in evenly spaced rows, with more space in first row.""" + if len(argv) == 1: + argv = argv[0] + to_print = "%-35s" % argv[0] + for v in argv[1:]: + to_print += "%-10s" % str(v) + print(to_print) + + def summary_results(self, table_res): + """Return a simple summary of final results for a tracker.""" + return dict( + zip(self.summary_fields, self._summary_row(table_res["COMBINED_SEQ"]),) + ) + + def detailed_results(self, table_res): + """Return detailed final results for a tracker.""" + # Get detailed field information + detailed_fields = self.float_fields + self.integer_fields + for h in self.float_array_fields + self.integer_array_fields: + for alpha in [int(100 * x) for x in self.array_labels]: + detailed_fields.append(h + "___" + str(alpha)) + detailed_fields.append(h + "___AUC") + + # Get detailed results + detailed_results = {} + for seq, res in table_res.items(): + detailed_row = self._detailed_row(res) + if len(detailed_row) != len(detailed_fields): + raise TrackEvalException( + f"Field names and data have different sizes " + f"({len(detailed_row)} and {len(detailed_fields)})" + ) + detailed_results[seq] = dict(zip(detailed_fields, detailed_row)) + return detailed_results + + def _detailed_row(self, res): + detailed_row = [] + for h in self.float_fields + self.integer_fields: + detailed_row.append(res[h]) + for h in self.float_array_fields + self.integer_array_fields: + for i, _ in enumerate([int(100 * x) for x in self.array_labels]): + detailed_row.append(res[h][i]) + detailed_row.append(np.mean(res[h])) + return detailed_row diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/teta.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/teta.py new file mode 100644 index 0000000..329626f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/metrics/teta.py @@ -0,0 +1,399 @@ +# fmt: off +# flake8: noqa + +"""Track Every Thing Accuracy metric.""" + +import numpy as np +from scipy.optimize import linear_sum_assignment + +from .. import _timing +from ._base_metric import _BaseMetric + +EPS = np.finfo("float").eps # epsilon + + +class TETA(_BaseMetric): + """TETA metric.""" + + def __init__(self, exhaustive=False, config=None): + """Initialize metric.""" + super().__init__() + self.plottable = True + self.array_labels = np.arange(0.0, 0.99, 0.05) + self.cls_array_labels = np.arange(0.5, 0.99, 0.05) + + self.integer_array_fields = [ + "Loc_TP", + "Loc_FN", + "Loc_FP", + "Cls_TP", + "Cls_FN", + "Cls_FP", + ] + self.float_array_fields = ( + ["TETA", "LocA", "AssocA", "ClsA"] + + ["LocRe", "LocPr"] + + ["AssocRe", "AssocPr"] + + ["ClsRe", "ClsPr"] + ) + self.fields = self.float_array_fields + self.integer_array_fields + self.summary_fields = self.float_array_fields + self.exhaustive = exhaustive + + def compute_global_assignment(self, data_thr, alpha=0.5): + """Compute global assignment of TP.""" + res = { + thr: {t: {} for t in range(data_thr[thr]["num_timesteps"])} + for thr in data_thr + } + + for thr in data_thr: + data = data_thr[thr] + # return empty result if tracker or gt sequence is empty + if data["num_tk_overlap_dets"] == 0 or data["num_gt_dets"] == 0: + return res + + # global alignment score + ga_score, _, _ = self.compute_global_alignment_score(data) + + # calculate scores for each timestep + for t, (gt_ids_t, tk_ids_t) in enumerate( + zip(data["gt_ids"], data["tk_ids"]) + ): + # get matches optimizing for TETA + amatch_rows, amatch_cols = self.compute_matches( + data, t, ga_score, gt_ids_t, tk_ids_t, alpha=alpha + ) + gt_ids = [data["gt_id_map"][tid] for tid in gt_ids_t[amatch_rows[0]]] + matched_ids = [ + data["tk_id_map"][tid] for tid in tk_ids_t[amatch_cols[0]] + ] + res[thr][t] = dict(zip(gt_ids, matched_ids)) + + return res + + def eval_sequence_single_thr(self, data, cls, cid2clsname, cls_fp_thr, thr): + """Computes TETA metric for one threshold for one sequence.""" + res = {} + class_info_list = [] + for field in self.float_array_fields + self.integer_array_fields: + if field.startswith("Cls"): + res[field] = np.zeros(len(self.cls_array_labels), dtype=float) + else: + res[field] = np.zeros((len(self.array_labels)), dtype=float) + + # return empty result if tracker or gt sequence is empty + if data["num_tk_overlap_dets"] == 0: + res["Loc_FN"] = data["num_gt_dets"] * np.ones( + (len(self.array_labels)), dtype=float + ) + if self.exhaustive: + cls_fp_thr[cls] = data["num_tk_cls_dets"] * np.ones( + (len(self.cls_array_labels)), dtype=float + ) + res = self._compute_final_fields(res) + return res, cls_fp_thr, class_info_list + + if data["num_gt_dets"] == 0: + if self.exhaustive: + cls_fp_thr[cls] = data["num_tk_cls_dets"] * np.ones( + (len(self.cls_array_labels)), dtype=float + ) + res = self._compute_final_fields(res) + return res, cls_fp_thr, class_info_list + + # global alignment score + ga_score, gt_id_count, tk_id_count = self.compute_global_alignment_score(data) + matches_counts = [np.zeros_like(ga_score) for _ in self.array_labels] + + # calculate scores for each timestep + for t, (gt_ids_t, tk_ids_t, tk_overlap_ids_t, tk_cls_ids_t) in enumerate( + zip( + data["gt_ids"], + data["tk_ids"], + data["tk_overlap_ids"], + data["tk_class_eval_tk_ids"], + ) + ): + # deal with the case that there are no gt_det/tk_det in a timestep + if len(gt_ids_t) == 0: + if self.exhaustive: + cls_fp_thr[cls] += len(tk_cls_ids_t) + continue + + # get matches optimizing for TETA + amatch_rows, amatch_cols = self.compute_matches( + data, t, ga_score, gt_ids_t, tk_ids_t, list(self.array_labels) + ) + + # map overlap_ids to original ids. + if len(tk_overlap_ids_t) != 0: + sorter = np.argsort(tk_ids_t) + indexes = sorter[ + np.searchsorted(tk_ids_t, tk_overlap_ids_t, sorter=sorter) + ] + sim_t = data["sim_scores"][t][:, indexes] + fpl_candidates = tk_overlap_ids_t[(sim_t >= (thr / 100)).any(axis=0)] + fpl_candidates_ori_ids_t = np.array( + [data["tk_id_map"][tid] for tid in fpl_candidates] + ) + else: + fpl_candidates_ori_ids_t = [] + + if self.exhaustive: + cls_fp_thr[cls] += len(tk_cls_ids_t) - len(tk_overlap_ids_t) + + # calculate and accumulate basic statistics + for a, alpha in enumerate(self.array_labels): + match_row, match_col = amatch_rows[a], amatch_cols[a] + num_matches = len(match_row) + matched_ori_ids = set( + [data["tk_id_map"][tid] for tid in tk_ids_t[match_col]] + ) + match_tk_cls = data["tk_classes"][t][match_col] + wrong_tk_cls = match_tk_cls[match_tk_cls != data["gt_classes"][t]] + + num_class_and_det_matches = np.sum( + match_tk_cls == data["gt_classes"][t] + ) + + if alpha >= 0.5: + for cid in wrong_tk_cls: + if cid in cid2clsname: + cname = cid2clsname[cid] + cls_fp_thr[cname][a - 10] += 1 + res["Cls_TP"][a - 10] += num_class_and_det_matches + res["Cls_FN"][a - 10] += num_matches - num_class_and_det_matches + + res["Loc_TP"][a] += num_matches + res["Loc_FN"][a] += len(gt_ids_t) - num_matches + res["Loc_FP"][a] += len(set(fpl_candidates_ori_ids_t) - matched_ori_ids) + + if num_matches > 0: + matches_counts[a][gt_ids_t[match_row], tk_ids_t[match_col]] += 1 + + # calculate AssocA, AssocRe, AssocPr + self.compute_association_scores(res, matches_counts, gt_id_count, tk_id_count) + + # calculate final scores + res = self._compute_final_fields(res) + return res, cls_fp_thr, class_info_list + + def compute_global_alignment_score(self, data): + """Computes global alignment score.""" + num_matches = np.zeros((data["num_gt_ids"], data["num_tk_ids"])) + gt_id_count = np.zeros((data["num_gt_ids"], 1)) + tk_id_count = np.zeros((1, data["num_tk_ids"])) + + # loop through each timestep and accumulate global track info. + for t, (gt_ids_t, tk_ids_t) in enumerate(zip(data["gt_ids"], data["tk_ids"])): + # count potential matches between ids in each time step + # these are normalized, weighted by match similarity + sim = data["sim_scores"][t] + sim_iou_denom = sim.sum(0, keepdims=True) + sim.sum(1, keepdims=True) - sim + sim_iou = np.zeros_like(sim) + mask = sim_iou_denom > (0 + EPS) + sim_iou[mask] = sim[mask] / sim_iou_denom[mask] + num_matches[gt_ids_t[:, None], tk_ids_t[None, :]] += sim_iou + + # calculate total number of dets for each gt_id and tk_id. + gt_id_count[gt_ids_t] += 1 + tk_id_count[0, tk_ids_t] += 1 + + # Calculate overall Jaccard alignment score between IDs + ga_score = num_matches / (gt_id_count + tk_id_count - num_matches) + return ga_score, gt_id_count, tk_id_count + + def compute_matches(self, data, t, ga_score, gt_ids, tk_ids, alpha): + """Compute matches based on alignment score.""" + sim = data["sim_scores"][t] + score_mat = ga_score[gt_ids[:, None], tk_ids[None, :]] * sim + # Hungarian algorithm to find best matches + match_rows, match_cols = linear_sum_assignment(-score_mat) + + if not isinstance(alpha, list): + alpha = [alpha] + alpha_match_rows, alpha_match_cols = [], [] + for a in alpha: + matched_mask = sim[match_rows, match_cols] >= a - EPS + alpha_match_rows.append(match_rows[matched_mask]) + alpha_match_cols.append(match_cols[matched_mask]) + return alpha_match_rows, alpha_match_cols + + def compute_association_scores(self, res, matches_counts, gt_id_count, tk_id_count): + """Calculate association scores for each alpha. + + First calculate scores per gt_id/tk_id combo, + and then average over the number of detections. + """ + for a, _ in enumerate(self.array_labels): + matches_count = matches_counts[a] + ass_a = matches_count / np.maximum( + 1, gt_id_count + tk_id_count - matches_count + ) + res["AssocA"][a] = np.sum(matches_count * ass_a) / np.maximum( + 1, res["Loc_TP"][a] + ) + ass_re = matches_count / np.maximum(1, gt_id_count) + res["AssocRe"][a] = np.sum(matches_count * ass_re) / np.maximum( + 1, res["Loc_TP"][a] + ) + ass_pr = matches_count / np.maximum(1, tk_id_count) + res["AssocPr"][a] = np.sum(matches_count * ass_pr) / np.maximum( + 1, res["Loc_TP"][a] + ) + + @_timing.time + def eval_sequence(self, data, cls, cls_id_name_mapping, cls_fp): + """Evaluate a single sequence across all thresholds.""" + res = {} + class_info_dict = {} + + for thr in data: + res[thr], cls_fp[thr], cls_info = self.eval_sequence_single_thr( + data[thr], cls, cls_id_name_mapping, cls_fp[thr], thr + ) + class_info_dict[thr] = cls_info + + return res, cls_fp, class_info_dict + + def combine_sequences(self, all_res): + """Combines metrics across all sequences.""" + data = {} + res = {} + + if all_res: + thresholds = list(list(all_res.values())[0].keys()) + else: + thresholds = [50] + for thr in thresholds: + data[thr] = {} + for seq_key in all_res: + data[thr][seq_key] = all_res[seq_key][thr] + for thr in thresholds: + res[thr] = self._combine_sequences_thr(data[thr]) + + return res + + def _combine_sequences_thr(self, all_res): + """Combines sequences over each threshold.""" + res = {} + for field in self.integer_array_fields: + res[field] = self._combine_sum(all_res, field) + for field in ["AssocRe", "AssocPr", "AssocA"]: + res[field] = self._combine_weighted_av( + all_res, field, res, weight_field="Loc_TP" + ) + res = self._compute_final_fields(res) + return res + + def combine_classes_class_averaged(self, all_res, ignore_empty=False): + """Combines metrics across all classes by averaging over classes. + + If 'ignore_empty' is True, then it only sums over classes + with at least one gt or predicted detection. + """ + data = {} + res = {} + if all_res: + thresholds = list(list(all_res.values())[0].keys()) + else: + thresholds = [50] + for thr in thresholds: + data[thr] = {} + for cls_key in all_res: + data[thr][cls_key] = all_res[cls_key][thr] + for thr in data: + res[thr] = self._combine_classes_class_averaged_thr( + data[thr], ignore_empty=ignore_empty + ) + return res + + def _combine_classes_class_averaged_thr(self, all_res, ignore_empty=False): + """Combines classes over each threshold.""" + res = {} + + def check_empty(val): + """Returns True if empty.""" + return not (val["Loc_TP"] + val["Loc_FN"] + val["Loc_FP"] > 0 + EPS).any() + + for field in self.integer_array_fields: + if ignore_empty: + res_field = {k: v for k, v in all_res.items() if not check_empty(v)} + else: + res_field = {k: v for k, v in all_res.items()} + res[field] = self._combine_sum(res_field, field) + + for field in self.float_array_fields: + if ignore_empty: + res_field = [v[field] for v in all_res.values() if not check_empty(v)] + else: + res_field = [v[field] for v in all_res.values()] + res[field] = np.mean(res_field, axis=0) + return res + + def combine_classes_det_averaged(self, all_res): + """Combines metrics across all classes by averaging over detections.""" + data = {} + res = {} + if all_res: + thresholds = list(list(all_res.values())[0].keys()) + else: + thresholds = [50] + for thr in thresholds: + data[thr] = {} + for cls_key in all_res: + data[thr][cls_key] = all_res[cls_key][thr] + for thr in data: + res[thr] = self._combine_classes_det_averaged_thr(data[thr]) + return res + + def _combine_classes_det_averaged_thr(self, all_res): + """Combines detections over each threshold.""" + res = {} + for field in self.integer_array_fields: + res[field] = self._combine_sum(all_res, field) + for field in ["AssocRe", "AssocPr", "AssocA"]: + res[field] = self._combine_weighted_av( + all_res, field, res, weight_field="Loc_TP" + ) + res = self._compute_final_fields(res) + return res + + @staticmethod + def _compute_final_fields(res): + """Calculate final metric values. + + This function is used both for both per-sequence calculation, + and in combining values across sequences. + """ + # LocA + res["LocRe"] = res["Loc_TP"] / np.maximum(1, res["Loc_TP"] + res["Loc_FN"]) + res["LocPr"] = res["Loc_TP"] / np.maximum(1, res["Loc_TP"] + res["Loc_FP"]) + res["LocA"] = res["Loc_TP"] / np.maximum( + 1, res["Loc_TP"] + res["Loc_FN"] + res["Loc_FP"] + ) + + # ClsA + res["ClsRe"] = res["Cls_TP"] / np.maximum(1, res["Cls_TP"] + res["Cls_FN"]) + res["ClsPr"] = res["Cls_TP"] / np.maximum(1, res["Cls_TP"] + res["Cls_FP"]) + res["ClsA"] = res["Cls_TP"] / np.maximum( + 1, res["Cls_TP"] + res["Cls_FN"] + res["Cls_FP"] + ) + + res["ClsRe"] = np.mean(res["ClsRe"]) + res["ClsPr"] = np.mean(res["ClsPr"]) + res["ClsA"] = np.mean(res["ClsA"]) + + res["TETA"] = (res["LocA"] + res["AssocA"] + res["ClsA"]) / 3 + + return res + + def print_summary_table(self, thr_res, thr, tracker, cls): + """Prints summary table of results.""" + print("") + metric_name = self.get_name() + self._row_print( + [f"{metric_name}{str(thr)}: {tracker}-{cls}"] + self.summary_fields + ) + self._row_print(["COMBINED"] + thr_res) diff --git a/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/utils.py b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/utils.py new file mode 100644 index 0000000..aa688e7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/teta_eval_toolkit/utils.py @@ -0,0 +1,46 @@ +# fmt: off +# flake8: noqa + +import csv +import os +from collections import OrderedDict + + +def validate_metrics_list(metrics_list): + """Get names of metric class and ensures they are unique, further checks that the fields within each metric class + do not have overlapping names. + """ + metric_names = [metric.get_name() for metric in metrics_list] + # check metric names are unique + if len(metric_names) != len(set(metric_names)): + raise TrackEvalException( + "Code being run with multiple metrics of the same name" + ) + fields = [] + for m in metrics_list: + fields += m.fields + # check metric fields are unique + if len(fields) != len(set(fields)): + raise TrackEvalException( + "Code being run with multiple metrics with fields of the same name" + ) + return metric_names + + +def get_track_id_str(ann): + """Get name of track ID in annotation.""" + if "track_id" in ann: + tk_str = "track_id" + elif "instance_id" in ann: + tk_str = "instance_id" + elif "scalabel_id" in ann: + tk_str = "scalabel_id" + else: + assert False, "No track/instance ID." + return tk_str + + +class TrackEvalException(Exception): + """Custom exception for catching expected errors.""" + + ... diff --git a/src/nn/segearth_ov3/sam3/eval/ytvis_coco_wrapper.py b/src/nn/segearth_ov3/sam3/eval/ytvis_coco_wrapper.py new file mode 100644 index 0000000..ced5899 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/ytvis_coco_wrapper.py @@ -0,0 +1,146 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import copy +import json +import logging + +import numpy as np +import pycocotools.mask as mask_util +from pycocotools.coco import COCO +from typing_extensions import override + + +class YTVIS(COCO): + """ + Helper class for reading YT-VIS annotations + """ + + @override + def __init__(self, annotation_file: str = None, ignore_gt_cats: bool = True): + """ + Args: + annotation_file: Path to the annotation file + ignore_gt_cats: If True, we ignore the ground truth categories and replace them with a dummy "object" category. This is useful for Phrase AP evaluation. + """ + self.ignore_gt_cats = ignore_gt_cats + super().__init__(annotation_file=annotation_file) + + @override + def createIndex(self): + # We rename some keys to match the COCO format before creating the index. + if "annotations" in self.dataset: + for ann in self.dataset["annotations"]: + if "video_id" in ann: + ann["image_id"] = int(ann.pop("video_id")) + if self.ignore_gt_cats: + ann["category_id"] = -1 + else: + ann["category_id"] = int(ann["category_id"]) + if "bboxes" in ann: + # note that in some datasets we load under this YTVIS class, + # some "bboxes" could be None for when the GT object is invisible, + # so we replace them with [0, 0, 0, 0] + ann["bboxes"] = [ + bbox if bbox is not None else [0, 0, 0, 0] + for bbox in ann["bboxes"] + ] + if "areas" in ann: + # similar to "bboxes", some areas could be None for when the GT + # object is invisible, so we replace them with 0 + areas = [a if a is not None else 0 for a in ann["areas"]] + # Compute average area of tracklet + ann["area"] = np.mean(areas) + if "videos" in self.dataset: + for vid in self.dataset["videos"]: + vid["id"] = int(vid["id"]) + self.dataset["images"] = self.dataset.pop("videos") + + if self.ignore_gt_cats: + self.dataset["categories"] = [ + {"supercategory": "object", "id": -1, "name": "object"} + ] + else: + for cat in self.dataset["categories"]: + cat["id"] = int(cat["id"]) + super().createIndex() + + @override + def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None): + if len(areaRng) > 0: + logging.warning( + "Note that we filter out objects based on their *average* area across the video, not per frame area" + ) + + return super().getAnnIds(imgIds=imgIds, catIds=catIds, iscrowd=iscrowd) + + @override + def showAnns(self, anns, draw_bbox=False): + raise NotImplementedError("Showing annotations is not supported") + + @override + def loadRes(self, resFile): + # Adapted from COCO.loadRes to support tracklets/masklets + res = YTVIS(ignore_gt_cats=self.ignore_gt_cats) + res.dataset["images"] = [img for img in self.dataset["images"]] + + if type(resFile) == str: + with open(resFile) as f: + anns = json.load(f) + elif type(resFile) == np.ndarray: + anns = self.loadNumpyAnnotations(resFile) + else: + anns = resFile + assert type(anns) == list, "results is not an array of objects" + annsImgIds = [ann["image_id"] for ann in anns] + assert set(annsImgIds) == ( + set(annsImgIds) & set(self.getImgIds()) + ), "Results do not correspond to current coco set" + if "bboxes" in anns[0] and not anns[0]["bboxes"] == []: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + bbs = [(bb if bb is not None else [0, 0, 0, 0]) for bb in ann["bboxes"]] + xxyy = [[bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]] for bb in bbs] + if not "segmentations" in ann: + ann["segmentations"] = [ + [[x1, y1, x1, y2, x2, y2, x2, y1]] for (x1, x2, y1, y2) in xxyy + ] + ann["areas"] = [bb[2] * bb[3] for bb in bbs] + # NOTE: We also compute average area of a tracklet across video, allowing us to compute area based mAP. + ann["area"] = np.mean(ann["areas"]) + ann["id"] = id + 1 + ann["iscrowd"] = 0 + elif "segmentations" in anns[0]: + res.dataset["categories"] = copy.deepcopy(self.dataset["categories"]) + for id, ann in enumerate(anns): + ann["bboxes"] = [ + mask_util.toBbox(segm) for segm in ann["segmentations"] + ] + if "areas" not in ann: + ann["areas"] = [ + mask_util.area(segm) for segm in ann["segmentations"] + ] + # NOTE: We also compute average area of a tracklet across video, allowing us to compute area based mAP. + ann["area"] = np.mean(ann["areas"]) + ann["id"] = id + 1 + ann["iscrowd"] = 0 + + res.dataset["annotations"] = anns + res.createIndex() + return res + + @override + def download(self, tarDir=None, imgIds=[]): + raise NotImplementedError + + @override + def loadNumpyAnnotations(self, data): + raise NotImplementedError("We don't support numpy annotations for now") + + @override + def annToRLE(self, ann): + raise NotImplementedError("We expect masks to be already in RLE format") + + @override + def annToMask(self, ann): + raise NotImplementedError("We expect masks to be already in RLE format") diff --git a/src/nn/segearth_ov3/sam3/eval/ytvis_eval.py b/src/nn/segearth_ov3/sam3/eval/ytvis_eval.py new file mode 100644 index 0000000..93f1cd6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/eval/ytvis_eval.py @@ -0,0 +1,411 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import copy +import gc +import logging +import os +from collections import defaultdict +from operator import xor +from pathlib import Path +from typing import List, Optional + +import numpy as np +import pycocotools.mask as mask_util +import torch +from pycocotools.cocoeval import COCOeval +from sam3.eval.cgf1_eval import CGF1Eval +from sam3.eval.coco_eval_offline import convert_to_xywh +from sam3.model.box_ops import box_xywh_inter_union +from sam3.train.masks_ops import rle_encode +from sam3.train.utils import distributed as dist +from typing_extensions import override + +try: + import rapidjson as json +except ModuleNotFoundError: + import json + +from iopath.common.file_io import g_pathmgr + + +class YTVISevalMixin: + """ + Identical to COCOeval but adapts computeIoU to compute IoU between tracklets/masklets. + """ + + @override + def _prepare(self): + """ + Copied from cocoeval.py but doesn't convert masks to RLEs (we assume they already are RLEs) + """ + p = self.params + if p.useCats: + gts = self.cocoGt.loadAnns( + self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds) + ) + dts = self.cocoDt.loadAnns( + self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds) + ) + else: + gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds)) + dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds)) + + # set ignore flag + for gt in gts: + gt["ignore"] = gt["ignore"] if "ignore" in gt else 0 + gt["ignore"] = "iscrowd" in gt and gt["iscrowd"] + if p.iouType == "keypoints": + gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"] + self._gts = defaultdict(list) # gt for evaluation + self._dts = defaultdict(list) # dt for evaluation + for gt in gts: + self._gts[gt["image_id"], gt["category_id"]].append(gt) + for dt in dts: + self._dts[dt["image_id"], dt["category_id"]].append(dt) + self.evalImgs = defaultdict(list) # per-image per-category evaluation results + self.eval = {} # accumulated evaluation results + + def computeIoU(self, imgId, catId): + """ + Compute IoU between tracklets. Copied from cocoeval.py but adapted for videos (in YT-VIS format) + """ + p = self.params + if p.useCats: + gt = self._gts[imgId, catId] + dt = self._dts[imgId, catId] + else: + gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]] + dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]] + if len(gt) == 0 or len(dt) == 0: + return [] + + # For class mAP and phrase AP evaluation, we sort the detections in descending order of scores (as in COCOeval). + # For demo F1 evaluation, we DO NOT sort the detections (but match them with GTs via Hungarian matching). + assert hasattr(self, "sort_inds_by_scores_in_iou"), ( + "subclasses that inherits YTVISevalMixin should set `self.sort_inds_by_scores_in_iou` " + "(True for class mAP and phrase AP, False for demo F1)" + ) + if self.sort_inds_by_scores_in_iou: + inds = np.argsort([-d["score"] for d in dt], kind="mergesort") + dt = [dt[i] for i in inds] + if len(dt) > p.maxDets[-1]: + dt = dt[0 : p.maxDets[-1]] + + if p.iouType == "segm": + g = [g["segmentations"] for g in gt] + d = [d["segmentations"] for d in dt] + elif p.iouType == "bbox": + g = [g["bboxes"] for g in gt] + d = [d["bboxes"] for d in dt] + else: + raise Exception("unknown iouType for iou computation") + + def iou_tracklets(preds, gts): + preds = torch.tensor(preds) + gts = torch.tensor(gts) + inter, union = box_xywh_inter_union( + preds.unsqueeze(1), gts.unsqueeze(0) + ) # Num preds x Num GTS x Num frames + inter = inter.sum(-1) + union = union.sum(-1) + assert ( + union > 0 + ).all(), ( + "There exists a tracklet with zero GTs across time. This is suspicious" + ) + return inter / union + + def iou_masklets(preds, gts): + inter = 0 + union = 0 + for p_i, gt_i in zip(preds, gts): + if p_i and gt_i: + # Compute areas of intersection and union + inter += mask_util.area( + mask_util.merge([p_i, gt_i], intersect=True) + ) + union += mask_util.area( + mask_util.merge([p_i, gt_i], intersect=False) + ) + elif gt_i: + union += mask_util.area(gt_i) + elif p_i: + union += mask_util.area(p_i) + if union > 0: + iou = inter / union + assert iou >= 0 and iou <= 1, "Encountered an error in IoU computation" + else: + assert np.isclose(inter, 0) and np.isclose( + union, 0 + ), "Encountered an error in IoU computation" + iou = 1 + return iou + + if p.iouType == "segm": + ious = [[iou_masklets(d_i, g_i) for g_i in g] for d_i in d] + else: + ious = iou_tracklets(d, g) + return np.array(ious) + + +class YTVISeval(YTVISevalMixin, COCOeval): + # For class mAP and phrase AP evaluation, we sort the detections in descending order of scores (as in COCOeval). + sort_inds_by_scores_in_iou = True + + +class VideoDemoF1Eval(YTVISevalMixin, CGF1Eval): + # For demo F1 evaluation, we DO NOT sort the detections (but match them with GTs via Hungarian matching). + sort_inds_by_scores_in_iou = False + + +class YTVISResultsWriter: + """ + Gather and dumps predictions in YT-VIS format. + Expected flow of API calls: reset() -> N * update() -> compute_synced() + """ + + def __init__( + self, + dump_file: str, + postprocessor, + gather_pred_via_filesys=False, + pred_file_evaluators: Optional[List] = None, + save_per_frame_scores: bool = False, + write_eval_metrics_file: bool = True, + eval_metrics_file_suffix: str = ".sam3_eval_metrics", + ): + self.dump_file = dump_file + self.dump = [] + self.postprocessor = postprocessor + self.gather_pred_via_filesys = gather_pred_via_filesys + if dist.is_main_process(): + dirname = os.path.dirname(self.dump_file) + if not os.path.exists(dirname): + os.makedirs(dirname, exist_ok=True) + logging.info(f"Creating folder: {dirname}") + + # the evaluation hooks to be applied to the prediction files + self.pred_file_evaluators = pred_file_evaluators or [] + self.save_per_frame_scores = save_per_frame_scores + # in addition to the prediction file, we also write the evaluation metrics + # for easier debugging and analysis (stored in another eval_metrics_file + # so that we can keep the dumped prediction file under YT-VIS format) + self.write_eval_metrics_file = write_eval_metrics_file + if self.write_eval_metrics_file: + self.eval_metrics_file = self.dump_file + eval_metrics_file_suffix + os.makedirs(os.path.dirname(self.eval_metrics_file), exist_ok=True) + + def _dump_vid_preds(self, results): + dumped_results = copy.deepcopy(results) + self.dump.extend(dumped_results) + + def prepare(self, predictions): + ytvis_results = [] + for video_id, prediction in predictions.items(): + if len(prediction) == 0: + continue + for k in ["boxes", "scores", "labels"]: + assert ( + k in prediction + ), f"Expected predictions to have `{k}` key, available keys are {prediction.keys()}" + if self.save_per_frame_scores: + assert ( + "per_frame_scores" in prediction + ), f"Expected predictions to have `per_frame_scores` key, available keys are {prediction.keys()}" + assert xor( + "masks" in prediction, "masks_rle" in prediction + ), f"Expected predictions to have either `masks` key or `masks_rle` key, available keys are {prediction.keys()}" + + boxes = prediction["boxes"] + boxes = convert_to_xywh(boxes).tolist() + scores = prediction["scores"].tolist() + labels = prediction["labels"].tolist() + if "masks" in prediction: + masks = prediction["masks"].squeeze(2) + assert ( + masks.ndim == 4 + ), "Expected masks to be of shape(N_preds,T_frames,H,W)" + + areas = [mask.flatten(1).sum(1).tolist() for mask in masks] + rles = [rle_encode(masklet) for masklet in masks] + + # memory clean + del masks + del prediction["masks"] + elif "masks_rle" in prediction: + rles = prediction.pop("masks_rle") + areas = [ + [0 if rle is None else rle.pop("area") for rle in rles_per_obj] + for rles_per_obj in rles + ] + else: + raise ValueError( + "Expected either `masks` or `masks_rle` key in the predictions." + ) + + new_results = [ + { + "video_id": video_id, + "category_id": track_label, + "bboxes": track_boxes, + "score": track_score, + "segmentations": track_masks, + "areas": track_areas, + } + for ( + track_boxes, + track_masks, + track_areas, + track_score, + track_label, + ) in zip(boxes, rles, areas, scores, labels) + ] + # Optionally, save per-frame scores + if self.save_per_frame_scores: + per_frame_scores = prediction["per_frame_scores"].tolist() + for res, track_per_frame_scores in zip(new_results, per_frame_scores): + res["per_frame_scores"] = track_per_frame_scores + + ytvis_results.extend(new_results) + + return ytvis_results + + def set_sync_device(self, device: torch.device): + self._sync_device = device + + def update(self, *args, **kwargs): + predictions = self.postprocessor.process_results(*args, **kwargs) + results = self.prepare(predictions) + self._dump_vid_preds(results) + + def _dump_preds(self): + if not dist.is_main_process(): + self.dump = [] + gc.collect() + return + dumped_file = Path(self.dump_file) + logging.info(f"YTVIS evaluator: Dumping predictions to {dumped_file}") + with g_pathmgr.open(str(dumped_file), "w") as f: + json.dump(self.dump, f) + self.dump = [] + gc.collect() + return str(dumped_file) + + def synchronize_between_processes(self): + logging.info("YT-VIS evaluator: Synchronizing between processes") + dump_dict = self._dedup_pre_gather(self.dump) + if self.gather_pred_via_filesys: + dump_dict_all_gpus = dist.gather_to_rank_0_via_filesys(dump_dict) + else: + dump_dict_all_gpus = dist.all_gather(dump_dict, force_cpu=True) + self.dump = self._dedup_post_gather(dump_dict_all_gpus) + logging.info(f"Gathered all {len(self.dump)} predictions") + + def _dedup_pre_gather(self, predictions): + """ + Organize the predictions as a dict-of-list using (video_id, category_id) as keys + for deduplication after gathering them across GPUs. + + During evaluation, PyTorch data loader under `drop_last: False` would wrap + around the dataset length to be a multiple of world size (GPU num) and duplicate + the remaining batches. This causes the same test sample to appear simultaneously + in multiple GPUs, resulting in duplicated predictions being saved into prediction + files. These duplicates are then counted as false positives under detection mAP + metrics (since a ground truth can be matched with only one prediction). + + For example, if there are 4 GPUs and 6 samples [A1, A2, B1, B2, C1, C2], the data + loader (under `drop_last: False`) would load it by wrapping it around like + `[A1, A2, B1, B2, C1, C2, *A1*, *A2*]` to make a multiple of 4 and then split it as + + - GPU 0: A1, C1 + - GPU 1: A2, C2 + - GPU 3: B1, **A1** + - GPU 4: B2, **A2** + (as in DistributedSampler in https://github.com/pytorch/pytorch/blob/521588519da9f4876d90ddd7a17c10d0eca89dc6/torch/utils/data/distributed.py#L116-L124) + + so the predictions on A1 and A2 will occur twice in the final gathered outputs + in the prediction file (and counted as false positives). This also affects our + YT-VIS official val evaluation, but to a lesser extent than YT-VIS dev since + the latter is much smaller and more susceptible to false positives. + + So we to deduplicate this. The tricky part is that we cannot deduplicate them + simply using video id, given that we are sharding the classes in each video + across multiple batches (with 20 prompts per batch) in our "orig_cats" eval dbs. + + The solution is to deduplicate based on (video_id, category_id) tuple as keys. + We organize the predictions as a dict-of-list using (video_id, category_id) as + keys on each GPU, with the list of masklets under this (video_id, category_id) + on this GPU as values. Then, we all-gather this dict-of-list across GPUs and + if a key (video_id, category_id) appears in multiple GPUs, we only take the + prediction masklet list from one GPU. + """ + prediction_dict = defaultdict(list) + for p in predictions: + prediction_dict[(p["video_id"], p["category_id"])].append(p) + return prediction_dict + + def _dedup_post_gather(self, list_of_prediction_dict): + """ + Deduplicate the predictions from all GPUs. See `_dedup_pre_gather` for details. + """ + dedup_prediction_dict = {} + duplication_keys = [] + for prediction_dict in list_of_prediction_dict: + for k, v in prediction_dict.items(): + if k not in dedup_prediction_dict: + dedup_prediction_dict[k] = v + else: + duplication_keys.append(k) + + logging.info( + f"skipped {len(duplication_keys)} duplicated predictions in YTVISResultsWriter " + f"with the following (video_id, category_id) tuples: {duplication_keys}" + ) + dedup_predictions = sum(dedup_prediction_dict.values(), []) + return dedup_predictions + + def compute_synced( + self, + ): + self.synchronize_between_processes() + dumped_file = self._dump_preds() + if not dist.is_main_process(): + return {"": 0.0} + + # run evaluation hooks on the prediction file + meters = {} + all_video_np_level_results = defaultdict(dict) + for evaluator in self.pred_file_evaluators: + gc.collect() + results, video_np_level_results = evaluator.evaluate(dumped_file) + meters.update(results) + for (video_id, category_id), res in video_np_level_results.items(): + all_video_np_level_results[(video_id, category_id)].update(res) + + gc.collect() + if self.write_eval_metrics_file: + # convert the nested dict of {(video_id, category_id): per_sample_metric_dict} + # to a list of per-sample metric dicts (with video_id and category_id) for JSON, + # as JSON doesn't allow using tuples like (video_id, category_id) as dict keys + video_np_level_metrics = [ + {"video_id": video_id, "category_id": category_id, **res} + for (video_id, category_id), res in all_video_np_level_results.items() + ] + eval_metrics = { + "dataset_level_metrics": meters, + "video_np_level_metrics": video_np_level_metrics, + } + with g_pathmgr.open(self.eval_metrics_file, "w") as f: + json.dump(eval_metrics, f) + logging.info( + f"YTVIS evaluator: Dumped evaluation metrics to {self.eval_metrics_file}" + ) + + if len(meters) == 0: + meters = {"": 0.0} + return meters + + def compute(self): + return {"": 0.0} + + def reset(self, *args, **kwargs): + self.dump = [] diff --git a/src/nn/segearth_ov3/sam3/logger.py b/src/nn/segearth_ov3/sam3/logger.py new file mode 100644 index 0000000..db9c0a6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/logger.py @@ -0,0 +1,54 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import logging +import os + +LOG_LEVELS = { + "DEBUG": logging.DEBUG, + "INFO": logging.INFO, + "WARNING": logging.WARNING, + "ERROR": logging.ERROR, + "CRITICAL": logging.CRITICAL, +} + + +class ColoredFormatter(logging.Formatter): + """A command line formatter with different colors for each level.""" + + def __init__(self): + super().__init__() + reset = "\033[0m" + colors = { + logging.DEBUG: f"{reset}\033[36m", # cyan, + logging.INFO: f"{reset}\033[32m", # green + logging.WARNING: f"{reset}\033[33m", # yellow + logging.ERROR: f"{reset}\033[31m", # red + logging.CRITICAL: f"{reset}\033[35m", # magenta + } + fmt_str = "{color}%(levelname)s %(asctime)s %(process)d %(filename)s:%(lineno)4d:{reset} %(message)s" + self.formatters = { + level: logging.Formatter(fmt_str.format(color=color, reset=reset)) + for level, color in colors.items() + } + self.default_formatter = self.formatters[logging.INFO] + + def format(self, record): + formatter = self.formatters.get(record.levelno, self.default_formatter) + return formatter.format(record) + + +def get_logger(name, level=logging.INFO): + """A command line logger.""" + if "LOG_LEVEL" in os.environ: + level = os.environ["LOG_LEVEL"].upper() + assert ( + level in LOG_LEVELS + ), f"Invalid LOG_LEVEL: {level}, must be one of {list(LOG_LEVELS.keys())}" + level = LOG_LEVELS[level] + logger = logging.getLogger(name) + logger.setLevel(level) + logger.propagate = False + ch = logging.StreamHandler() + ch.setLevel(level) + ch.setFormatter(ColoredFormatter()) + logger.addHandler(ch) + return logger diff --git a/src/nn/segearth_ov3/sam3/model/__init__.py b/src/nn/segearth_ov3/sam3/model/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/model/act_ckpt_utils.py b/src/nn/segearth_ov3/sam3/model/act_ckpt_utils.py new file mode 100644 index 0000000..c935cfc --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/act_ckpt_utils.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import inspect +from functools import wraps +from typing import Callable, TypeVar, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint as checkpoint +from torch.utils._pytree import tree_map_only + +# Type variables for better type hinting +T = TypeVar("T") +Module = TypeVar("Module", bound=nn.Module) + + +def activation_ckpt_wrapper(module: Union[nn.Module, Callable]) -> Callable: + """ + Wraps a given module to enable or disable activation checkpointing. + + Activation checkpointing (gradient checkpointing) trades compute for memory by + recomputing intermediate activations during the backward pass instead of storing + them in memory during the forward pass. + + When activation checkpointing is enabled, the wrapper expects only keyword arguments, + and it maps these to positional arguments based on the module's signature. + + Args: + module: The module or function to wrap with activation checkpointing + + Returns: + A wrapped callable that supports activation checkpointing + + Usage: + The returned wrapper function can be called with the same arguments as the + original module, with an additional `act_ckpt_enable` keyword argument to control + activation checkpointing and optional `use_reentrant` parameter. + + Example: + ```python + wrapped_module = activation_ckpt_wrapper(my_module) + output = wrapped_module(x=input_tensor, y=another_tensor, act_ckpt_enable=True) + ``` + """ + + @wraps(module) + def act_ckpt_wrapper( + *args, act_ckpt_enable: bool = True, use_reentrant: bool = False, **kwargs + ): + if act_ckpt_enable: + if len(args) > 0: + raise ValueError( + "This wrapper expects keyword arguments only when `act_ckpt_enable=True`" + ) + # Get the signature of the target function/module + callable_fn = module.forward if isinstance(module, nn.Module) else module + sig = inspect.signature(callable_fn) + # Create a mapping of parameter names to their default values + param_defaults = { + name: param.default for name, param in sig.parameters.items() + } + args = [] + for p_name in param_defaults.keys(): + if p_name in kwargs: + args.append(kwargs.pop(p_name)) + elif param_defaults[p_name] is not inspect.Parameter.empty: + # Set arg to default value if it's not in kwargs. Useful for primitive types or args that default to None + args.append(param_defaults[p_name]) + elif ( + sig.parameters[p_name].kind is not inspect.Parameter.VAR_KEYWORD + ): # Skip **kwargs parameter + raise ValueError(f"Missing positional argument: {p_name}") + + # Scan remaining kwargs for torch.Tensor + remaining_keys = list(kwargs.keys()) + for key in remaining_keys: + if isinstance(kwargs[key], torch.Tensor): + # Remove the tensor from kwargs, assuming it's not required by the module. + # If it is required, the module's signature should be modified to accept it as a positional or keyword argument. + kwargs[key] = "_REMOVED_BY_ACT_CKPT_WRAPPER_" + + ret = checkpoint.checkpoint( + module, *args, use_reentrant=use_reentrant, **kwargs + ) + else: + ret = module(*args, **kwargs) + + return ret + + return act_ckpt_wrapper + + +def clone_output_wrapper(f: Callable[..., T]) -> Callable[..., T]: + """ + Clone the CUDA output tensors of a function to avoid in-place operations. + + This wrapper is useful when working with torch.compile to prevent errors + related to in-place operations on tensors. + + Args: + f: The function whose CUDA tensor outputs should be cloned + + Returns: + A wrapped function that clones any CUDA tensor outputs + """ + + @wraps(f) + def wrapped(*args, **kwargs): + outputs = f(*args, **kwargs) + return tree_map_only( + torch.Tensor, lambda t: t.clone() if t.is_cuda else t, outputs + ) + + return wrapped diff --git a/src/nn/segearth_ov3/sam3/model/box_ops.py b/src/nn/segearth_ov3/sam3/model/box_ops.py new file mode 100644 index 0000000..f88e4ad --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/box_ops.py @@ -0,0 +1,217 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +""" +Utilities for bounding box manipulation and GIoU. +""" + +from typing import Tuple + +import torch + + +def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.stack(b, dim=-1) + + +def box_cxcywh_to_xywh(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (w), (h)] + return torch.stack(b, dim=-1) + + +def box_xywh_to_xyxy(x): + x, y, w, h = x.unbind(-1) + b = [(x), (y), (x + w), (y + h)] + return torch.stack(b, dim=-1) + + +def box_xywh_to_cxcywh(x): + x, y, w, h = x.unbind(-1) + b = [(x + 0.5 * w), (y + 0.5 * h), (w), (h)] + return torch.stack(b, dim=-1) + + +def box_xyxy_to_xywh(x): + x, y, X, Y = x.unbind(-1) + b = [(x), (y), (X - x), (Y - y)] + return torch.stack(b, dim=-1) + + +def box_xyxy_to_cxcywh(x): + x0, y0, x1, y1 = x.unbind(-1) + b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)] + return torch.stack(b, dim=-1) + + +def box_area(boxes): + """ + Batched version of box area. Boxes should be in [x0, y0, x1, y1] format. + + Inputs: + - boxes: Tensor of shape (..., 4) + + Returns: + - areas: Tensor of shape (...,) + """ + x0, y0, x1, y1 = boxes.unbind(-1) + return (x1 - x0) * (y1 - y0) + + +def masks_to_boxes(masks): + """Compute the bounding boxes around the provided masks + + The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. + + Returns a [N, 4] tensors, with the boxes in xyxy format + """ + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device) + + h, w = masks.shape[-2:] + + y = torch.arange(0, h, dtype=torch.float, device=masks.device) + x = torch.arange(0, w, dtype=torch.float, device=masks.device) + y, x = torch.meshgrid(y, x) + + x_mask = masks * x.unsqueeze(0) + x_max = x_mask.flatten(1).max(-1)[0] + 1 + x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + y_mask = masks * y.unsqueeze(0) + y_max = y_mask.flatten(1).max(-1)[0] + 1 + y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + boxes = torch.stack([x_min, y_min, x_max, y_max], 1) + # Invalidate boxes corresponding to empty masks. + boxes = boxes * masks.flatten(-2).any(-1) + return boxes + + +def box_iou(boxes1, boxes2): + """ + Batched version of box_iou. Boxes should be in [x0, y0, x1, y1] format. + + Inputs: + - boxes1: Tensor of shape (..., N, 4) + - boxes2: Tensor of shape (..., M, 4) + + Returns: + - iou, union: Tensors of shape (..., N, M) + """ + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + # boxes1: (..., N, 4) -> (..., N, 1, 2) + # boxes2: (..., M, 4) -> (..., 1, M, 2) + lt = torch.max(boxes1[..., :, None, :2], boxes2[..., None, :, :2]) + rb = torch.min(boxes1[..., :, None, 2:], boxes2[..., None, :, 2:]) + + wh = (rb - lt).clamp(min=0) # (..., N, M, 2) + inter = wh[..., 0] * wh[..., 1] # (..., N, M) + + union = area1[..., None] + area2[..., None, :] - inter + + iou = inter / union + return iou, union + + +def generalized_box_iou(boxes1, boxes2): + """ + Batched version of Generalized IoU from https://giou.stanford.edu/ + + Boxes should be in [x0, y0, x1, y1] format + + Inputs: + - boxes1: Tensor of shape (..., N, 4) + - boxes2: Tensor of shape (..., M, 4) + + Returns: + - giou: Tensor of shape (..., N, M) + """ + iou, union = box_iou(boxes1, boxes2) + + # boxes1: (..., N, 4) -> (..., N, 1, 2) + # boxes2: (..., M, 4) -> (..., 1, M, 2) + lt = torch.min(boxes1[..., :, None, :2], boxes2[..., None, :, :2]) + rb = torch.max(boxes1[..., :, None, 2:], boxes2[..., None, :, 2:]) + + wh = (rb - lt).clamp(min=0) # (..., N, M, 2) + area = wh[..., 0] * wh[..., 1] # (..., N, M) + + return iou - (area - union) / area + + +@torch.jit.script +def fast_diag_generalized_box_iou(boxes1, boxes2): + assert len(boxes1) == len(boxes2) + box1_xy = boxes1[:, 2:] + box1_XY = boxes1[:, :2] + box2_xy = boxes2[:, 2:] + box2_XY = boxes2[:, :2] + # assert (box1_xy >= box1_XY).all() + # assert (box2_xy >= box2_XY).all() + area1 = (box1_xy - box1_XY).prod(-1) + area2 = (box2_xy - box2_XY).prod(-1) + + lt = torch.max(box1_XY, box2_XY) # [N,2] + lt2 = torch.min(box1_XY, box2_XY) + rb = torch.min(box1_xy, box2_xy) # [N,2] + rb2 = torch.max(box1_xy, box2_xy) + + inter = (rb - lt).clamp(min=0).prod(-1) + tot_area = (rb2 - lt2).clamp(min=0).prod(-1) + + union = area1 + area2 - inter + + iou = inter / union + + return iou - (tot_area - union) / tot_area + + +@torch.jit.script +def fast_diag_box_iou(boxes1, boxes2): + assert len(boxes1) == len(boxes2) + box1_xy = boxes1[:, 2:] + box1_XY = boxes1[:, :2] + box2_xy = boxes2[:, 2:] + box2_XY = boxes2[:, :2] + # assert (box1_xy >= box1_XY).all() + # assert (box2_xy >= box2_XY).all() + area1 = (box1_xy - box1_XY).prod(-1) + area2 = (box2_xy - box2_XY).prod(-1) + + lt = torch.max(box1_XY, box2_XY) # [N,2] + rb = torch.min(box1_xy, box2_xy) # [N,2] + + inter = (rb - lt).clamp(min=0).prod(-1) + + union = area1 + area2 - inter + + iou = inter / union + + return iou + + +def box_xywh_inter_union( + boxes1: torch.Tensor, boxes2: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + # Asuumes boxes in xywh format + assert boxes1.size(-1) == 4 and boxes2.size(-1) == 4 + boxes1 = box_xywh_to_xyxy(boxes1) + boxes2 = box_xywh_to_xyxy(boxes2) + box1_tl_xy = boxes1[..., :2] + box1_br_xy = boxes1[..., 2:] + box2_tl_xy = boxes2[..., :2] + box2_br_xy = boxes2[..., 2:] + area1 = (box1_br_xy - box1_tl_xy).prod(-1) + area2 = (box2_br_xy - box2_tl_xy).prod(-1) + + assert (area1 >= 0).all() and (area2 >= 0).all() + tl = torch.max(box1_tl_xy, box2_tl_xy) + br = torch.min(box1_br_xy, box2_br_xy) + + inter = (br - tl).clamp(min=0).prod(-1) + union = area1 + area2 - inter + + return inter, union diff --git a/src/nn/segearth_ov3/sam3/model/data_misc.py b/src/nn/segearth_ov3/sam3/model/data_misc.py new file mode 100644 index 0000000..4bbcf55 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/data_misc.py @@ -0,0 +1,209 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +""" +Misc functions, including distributed helpers. +""" + +import collections +import re + +from dataclasses import dataclass, field as field_ptr_behaviour, fields, is_dataclass +from typing import Any, get_args, get_origin, List, Mapping, Optional, Sequence, Union + +import torch + + +MyTensor = Union[torch.Tensor, List[Any]] + + +def interpolate( + input, size=None, scale_factor=None, mode="nearest", align_corners=None +): + # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor + """ + Equivalent to nn.functional.interpolate, but with support for empty channel sizes. + """ + if input.numel() > 0: + return torch.nn.functional.interpolate( + input, size, scale_factor, mode, align_corners + ) + + assert ( + input.shape[0] != 0 or input.shape[1] != 0 + ), "At least one of the two first dimensions must be non zero" + + if input.shape[1] == 0: + # Pytorch doesn't support null dimension on the channel dimension, so we transpose to fake a null batch dim + return torch.nn.functional.interpolate( + input.transpose(0, 1), size, scale_factor, mode, align_corners + ).transpose(0, 1) + + # empty batch dimension is now supported in pytorch + return torch.nn.functional.interpolate( + input, size, scale_factor, mode, align_corners + ) + + +@dataclass +class BatchedPointer: + stage_ids: MyTensor + stage_ids__type = torch.long + query_ids: MyTensor + query_ids__type = torch.long + object_ids: MyTensor + object_ids__type = torch.long + ptr_mask: MyTensor + ptr_mask__type = torch.bool + ptr_types: MyTensor + ptr_types__type = torch.long + + +@dataclass +class FindStage: + img_ids: MyTensor + img_ids__type = torch.long + text_ids: MyTensor + text_ids__type = torch.long + + input_boxes: MyTensor + input_boxes__type = torch.float + input_boxes_mask: MyTensor + input_boxes_mask__type = torch.bool + input_boxes_label: MyTensor + input_boxes_label__type = torch.long + + input_points: MyTensor + input_points__type = torch.float + input_points_mask: MyTensor + input_points_mask__type = torch.bool + + # We track the object ids referred to by this query. + # This is beneficial for tracking in videos without the need for pointers. + object_ids: Optional[List[List]] = None # List of objects per query + + +@dataclass +class BatchedFindTarget: + # The number of boxes in each find query + num_boxes: MyTensor + num_boxes__type = torch.long + + # Target boxes in normalized CxCywh format + boxes: MyTensor + boxes__type = torch.float + # Target boxes in normalized CxCywh format but in padded representation + # as used in BinaryHungarianMatcherV2 (unlike the packed ones in `boxes`) + boxes_padded: MyTensor + boxes_padded__type = torch.float + + # For hybrid matching, we repeat the boxes + repeated_boxes: MyTensor + repeated_boxes__type = torch.float + + # Target Segmentation masks + segments: Optional[MyTensor] + segments__type = torch.bool + + # Target Semantic Segmentation masks + semantic_segments: Optional[MyTensor] + semantic_segments__type = torch.bool + + is_valid_segment: Optional[MyTensor] + is_valid_segment__type = torch.bool + + # Whether annotations are exhaustive for each query + is_exhaustive: MyTensor + is_exhaustive__type = torch.bool + + # The object id for each ground-truth box, in both packed and padded representations + object_ids: MyTensor + object_ids__type = torch.long + object_ids_padded: MyTensor + object_ids_padded__type = torch.long + + +@dataclass +class BatchedInferenceMetadata: + """All metadata required to post-process a find stage""" + + # Coco id that corresponds to the "image" for evaluation by the coco evaluator + coco_image_id: MyTensor + coco_image_id__type = torch.long + + # id in the original dataset, such that we can use the original evaluator + original_image_id: MyTensor + original_image_id__type = torch.long + + # Original category id (if we want to use the original evaluator) + original_category_id: MyTensor + original_category_id__type = torch.int + + # Size of the raw image (height, width) + original_size: MyTensor + original_size__type = torch.long + + # id of the object in the media (track_id for a video) + object_id: MyTensor + object_id__type = torch.long + + # index of the frame in the media (0 in the case of a single-frame media) + frame_index: MyTensor + frame_index__type = torch.long + + # Adding for relations inference + # get_text_input: List[Optional[str]] + + # Adding for TA conditional inference + is_conditioning_only: List[Optional[bool]] + + +@dataclass +class BatchedDatapoint: + img_batch: torch.Tensor + find_text_batch: List[str] + find_inputs: List[FindStage] + find_targets: List[BatchedFindTarget] + find_metadatas: List[BatchedInferenceMetadata] + raw_images: Optional[List[Any]] = None + + +def convert_my_tensors(obj): + def is_optional_field(field) -> bool: + return get_origin(field) is Union and type(None) in get_args(field) + + for field in fields(obj): + if is_dataclass(getattr(obj, field.name)): + convert_my_tensors(getattr(obj, field.name)) + continue + + field_type = field.type + if is_optional_field(field.type): + field_type = Union[get_args(field.type)[:-1]] # Get the Optional field type + + if field_type != MyTensor or getattr(obj, field.name) is None: + continue + + elif len(getattr(obj, field.name)) and isinstance( + getattr(obj, field.name)[0], torch.Tensor + ): + stack_dim = 0 + if field.name in [ + "input_boxes", + "input_boxes_label", + ]: + stack_dim = 1 + setattr( + obj, + field.name, + torch.stack(getattr(obj, field.name), dim=stack_dim).to( + getattr(obj, field.name + "__type") + ), + ) + else: + setattr( + obj, + field.name, + torch.as_tensor( + getattr(obj, field.name), dtype=getattr(obj, field.name + "__type") + ), + ) + return obj diff --git a/src/nn/segearth_ov3/sam3/model/decoder.py b/src/nn/segearth_ov3/sam3/model/decoder.py new file mode 100644 index 0000000..dfd73bb --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/decoder.py @@ -0,0 +1,956 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +""" +Transformer decoder. +Inspired from Pytorch's version, adds the pre-norm variant +""" + +from typing import Any, Dict, List, Optional + +import numpy as np + +import torch + +from sam3.sam.transformer import RoPEAttention + +from torch import nn, Tensor +from torchvision.ops.roi_align import RoIAlign + +from .act_ckpt_utils import activation_ckpt_wrapper + +from .box_ops import box_cxcywh_to_xyxy + +from .model_misc import ( + gen_sineembed_for_position, + get_activation_fn, + get_clones, + inverse_sigmoid, + MLP, +) + + +class TransformerDecoderLayer(nn.Module): + def __init__( + self, + activation: str, + d_model: int, + dim_feedforward: int, + dropout: float, + cross_attention: nn.Module, + n_heads: int, + use_text_cross_attention: bool = False, + ): + super().__init__() + + # cross attention + self.cross_attn = cross_attention + self.dropout1 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm1 = nn.LayerNorm(d_model) + + # cross attention text + self.use_text_cross_attention = use_text_cross_attention + if use_text_cross_attention: + self.ca_text = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.catext_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.catext_norm = nn.LayerNorm(d_model) + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.dropout2 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm2 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.activation = get_activation_fn(activation) + self.dropout3 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.linear2 = nn.Linear(dim_feedforward, d_model) + self.dropout4 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm3 = nn.LayerNorm(d_model) + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, tgt): + with torch.amp.autocast(device_type="cuda", enabled=False): + tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout4(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward( + self, + # for tgt + tgt: Optional[Tensor], # nq, bs, d_model + tgt_query_pos: Optional[Tensor] = None, # pos for query. MLP(Sine(pos)) + tgt_query_sine_embed: Optional[Tensor] = None, # pos for query. Sine(pos) + tgt_key_padding_mask: Optional[Tensor] = None, + tgt_reference_points: Optional[Tensor] = None, # nq, bs, 4 + memory_text: Optional[Tensor] = None, # num_token, bs, d_model + text_attention_mask: Optional[Tensor] = None, # bs, num_token + # for memory + memory: Optional[Tensor] = None, # hw, bs, d_model + memory_key_padding_mask: Optional[Tensor] = None, + memory_level_start_index: Optional[Tensor] = None, # num_levels + memory_spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + memory_pos: Optional[Tensor] = None, # pos for memory + # sa + self_attn_mask: Optional[Tensor] = None, # mask used for self-attention + cross_attn_mask: Optional[Tensor] = None, # mask used for cross-attention + # dac + dac=False, + dac_use_selfatt_ln=True, + presence_token=None, + # skip inside deformable attn + identity=0.0, + **kwargs, # additional kwargs for compatibility + ): + """ + Input: + - tgt/tgt_query_pos: nq, bs, d_model + - + """ + # self attention + if self.self_attn is not None: + if dac: + # we only apply self attention to the first half of the queries + assert tgt.shape[0] % 2 == 0 + num_o2o_queries = tgt.shape[0] // 2 + tgt_o2o = tgt[:num_o2o_queries] + tgt_query_pos_o2o = tgt_query_pos[:num_o2o_queries] + tgt_o2m = tgt[num_o2o_queries:] + else: + tgt_o2o = tgt + tgt_query_pos_o2o = tgt_query_pos + + if presence_token is not None: + tgt_o2o = torch.cat([presence_token, tgt_o2o], dim=0) + tgt_query_pos_o2o = torch.cat( + [torch.zeros_like(presence_token), tgt_query_pos_o2o], dim=0 + ) + tgt_query_pos = torch.cat( + [torch.zeros_like(presence_token), tgt_query_pos], dim=0 + ) + + q = k = self.with_pos_embed(tgt_o2o, tgt_query_pos_o2o) + tgt2 = self.self_attn(q, k, tgt_o2o, attn_mask=self_attn_mask)[0] + tgt_o2o = tgt_o2o + self.dropout2(tgt2) + if dac: + if not dac_use_selfatt_ln: + tgt_o2o = self.norm2(tgt_o2o) + tgt = torch.cat((tgt_o2o, tgt_o2m), dim=0) # Recombine + if dac_use_selfatt_ln: + tgt = self.norm2(tgt) + else: + tgt = tgt_o2o + tgt = self.norm2(tgt) + + if self.use_text_cross_attention: + tgt2 = self.ca_text( + self.with_pos_embed(tgt, tgt_query_pos), + memory_text, + memory_text, + key_padding_mask=text_attention_mask, + )[0] + tgt = tgt + self.catext_dropout(tgt2) + tgt = self.catext_norm(tgt) + + if presence_token is not None: + presence_token_mask = torch.zeros_like(cross_attn_mask[:, :1, :]) + cross_attn_mask = torch.cat( + [presence_token_mask, cross_attn_mask], dim=1 + ) # (bs*nheads, 1+nq, hw) + + # Cross attention to image + tgt2 = self.cross_attn( + query=self.with_pos_embed(tgt, tgt_query_pos), + key=self.with_pos_embed(memory, memory_pos), + value=memory, + attn_mask=cross_attn_mask, + key_padding_mask=( + memory_key_padding_mask.transpose(0, 1) + if memory_key_padding_mask is not None + else None + ), + )[0] + + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + + # ffn + tgt = self.forward_ffn(tgt) + + presence_token_out = None + if presence_token is not None: + presence_token_out = tgt[:1] + tgt = tgt[1:] + + return tgt, presence_token_out + + +class TransformerDecoder(nn.Module): + def __init__( + self, + d_model: int, + frozen: bool, + interaction_layer, + layer, + num_layers: int, + num_queries: int, + return_intermediate: bool, + box_refine: bool = False, + num_o2m_queries: int = 0, + dac: bool = False, + boxRPB: str = "none", + # Experimental: An object query for SAM 2 tasks + instance_query: bool = False, + # Defines the number of additional instance queries, + # 1 or 4 are the most likely for single vs multi mask support + num_instances: int = 1, # Irrelevant if instance_query is False + dac_use_selfatt_ln: bool = True, + use_act_checkpoint: bool = False, + compile_mode=None, + presence_token: bool = False, + clamp_presence_logits: bool = True, + clamp_presence_logit_max_val: float = 10.0, + use_normed_output_consistently: bool = True, + separate_box_head_instance: bool = False, + separate_norm_instance: bool = False, + resolution: Optional[int] = None, + stride: Optional[int] = None, + ): + super().__init__() + self.d_model = d_model + self.layers = get_clones(layer, num_layers) + self.fine_layers = ( + get_clones(interaction_layer, num_layers) + if interaction_layer is not None + else [None] * num_layers + ) + self.num_layers = num_layers + self.num_queries = num_queries + self.dac = dac + if dac: + self.num_o2m_queries = num_queries + tot_num_queries = num_queries + else: + self.num_o2m_queries = num_o2m_queries + tot_num_queries = num_queries + num_o2m_queries + self.norm = nn.LayerNorm(d_model) + self.return_intermediate = return_intermediate + self.bbox_embed = MLP(d_model, d_model, 4, 3) + self.query_embed = nn.Embedding(tot_num_queries, d_model) + self.instance_query_embed = None + self.instance_query_reference_points = None + self.use_instance_query = instance_query + self.num_instances = num_instances + self.use_normed_output_consistently = use_normed_output_consistently + + self.instance_norm = nn.LayerNorm(d_model) if separate_norm_instance else None + self.instance_bbox_embed = None + if separate_box_head_instance: + self.instance_bbox_embed = MLP(d_model, d_model, 4, 3) + if instance_query: + self.instance_query_embed = nn.Embedding(num_instances, d_model) + self.box_refine = box_refine + if box_refine: + nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0) + nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0) + + self.reference_points = nn.Embedding(num_queries, 4) + if instance_query: + self.instance_reference_points = nn.Embedding(num_instances, 4) + + assert boxRPB in ["none", "log", "linear", "both"] + self.boxRPB = boxRPB + if boxRPB != "none": + try: + nheads = self.layers[0].cross_attn_image.num_heads + except AttributeError: + nheads = self.layers[0].cross_attn.num_heads + + n_input = 4 if boxRPB == "both" else 2 + self.boxRPB_embed_x = MLP(n_input, d_model, nheads, 2) + self.boxRPB_embed_y = MLP(n_input, d_model, nheads, 2) + self.compilable_cord_cache = None + self.compilable_stored_size = None + self.coord_cache = {} + + if resolution is not None and stride is not None: + feat_size = resolution // stride + coords_h, coords_w = self._get_coords( + feat_size, feat_size, device="cuda" + ) + self.compilable_cord_cache = (coords_h, coords_w) + self.compilable_stored_size = (feat_size, feat_size) + + self.roi_pooler = ( + RoIAlign(output_size=7, spatial_scale=1, sampling_ratio=-1, aligned=True) + if interaction_layer is not None + else None + ) + if frozen: + for p in self.parameters(): + p.requires_grad_(False) + + self.presence_token = None + self.clamp_presence_logits = clamp_presence_logits + self.clamp_presence_logit_max_val = clamp_presence_logit_max_val + if presence_token: + self.presence_token = nn.Embedding(1, d_model) + self.presence_token_head = MLP(d_model, d_model, 1, 3) + self.presence_token_out_norm = nn.LayerNorm(d_model) + + self.ref_point_head = MLP(2 * self.d_model, self.d_model, self.d_model, 2) + self.dac_use_selfatt_ln = dac_use_selfatt_ln + self.use_act_checkpoint = use_act_checkpoint + + nn.init.normal_(self.query_embed.weight.data) + if self.instance_query_embed is not None: + nn.init.normal_(self.instance_query_embed.weight.data) + + assert self.roi_pooler is None + assert self.return_intermediate, "support return_intermediate only" + assert self.box_refine, "support box refine only" + + self.compile_mode = compile_mode + self.compiled = False + # We defer compilation till after the first forward, to first warm-up the boxRPB cache + + # assign layer index to each layer so that some layers can decide what to do + # based on which layer index they are (e.g. cross attention to memory bank only + # in selected layers) + for layer_idx, layer in enumerate(self.layers): + layer.layer_idx = layer_idx + + @staticmethod + def _get_coords(H, W, device): + coords_h = torch.arange(0, H, device=device, dtype=torch.float32) / H + coords_w = torch.arange(0, W, device=device, dtype=torch.float32) / W + return coords_h, coords_w + + def _get_rpb_matrix(self, reference_boxes, feat_size): + H, W = feat_size + boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes).transpose(0, 1) + bs, num_queries, _ = boxes_xyxy.shape + if self.compilable_cord_cache is None: + self.compilable_cord_cache = self._get_coords(H, W, reference_boxes.device) + self.compilable_stored_size = (H, W) + + if torch.compiler.is_dynamo_compiling() or self.compilable_stored_size == ( + H, + W, + ): + # good, hitting the cache, will be compilable + coords_h, coords_w = self.compilable_cord_cache + else: + # cache miss, will create compilation issue + # In case we're not compiling, we'll still rely on the dict-based cache + if feat_size not in self.coord_cache: + self.coord_cache[feat_size] = self._get_coords( + H, W, reference_boxes.device + ) + coords_h, coords_w = self.coord_cache[feat_size] + + assert coords_h.shape == (H,) + assert coords_w.shape == (W,) + + deltas_y = coords_h.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 1:4:2] + deltas_y = deltas_y.view(bs, num_queries, -1, 2) + deltas_x = coords_w.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 0:3:2] + deltas_x = deltas_x.view(bs, num_queries, -1, 2) + + if self.boxRPB in ["log", "both"]: + deltas_x_log = deltas_x * 8 # normalize to -8, 8 + deltas_x_log = ( + torch.sign(deltas_x_log) + * torch.log2(torch.abs(deltas_x_log) + 1.0) + / np.log2(8) + ) + + deltas_y_log = deltas_y * 8 # normalize to -8, 8 + deltas_y_log = ( + torch.sign(deltas_y_log) + * torch.log2(torch.abs(deltas_y_log) + 1.0) + / np.log2(8) + ) + if self.boxRPB == "log": + deltas_x = deltas_x_log + deltas_y = deltas_y_log + else: + deltas_x = torch.cat([deltas_x, deltas_x_log], dim=-1) + deltas_y = torch.cat([deltas_y, deltas_y_log], dim=-1) + + if self.training: + assert self.use_act_checkpoint, "activation ckpt not enabled in decoder" + deltas_x = activation_ckpt_wrapper(self.boxRPB_embed_x)( + x=deltas_x, + act_ckpt_enable=self.training and self.use_act_checkpoint, + ) # bs, num_queries, W, n_heads + deltas_y = activation_ckpt_wrapper(self.boxRPB_embed_y)( + x=deltas_y, + act_ckpt_enable=self.training and self.use_act_checkpoint, + ) # bs, num_queries, H, n_heads + + if not torch.compiler.is_dynamo_compiling(): + assert deltas_x.shape[:3] == (bs, num_queries, W) + assert deltas_y.shape[:3] == (bs, num_queries, H) + + B = deltas_y.unsqueeze(3) + deltas_x.unsqueeze( + 2 + ) # bs, num_queries, H, W, n_heads + if not torch.compiler.is_dynamo_compiling(): + assert B.shape[:4] == (bs, num_queries, H, W) + B = B.flatten(2, 3) # bs, num_queries, H*W, n_heads + B = B.permute(0, 3, 1, 2) # bs, n_heads, num_queries, H*W + B = B.contiguous() # memeff attn likes ordered strides + if not torch.compiler.is_dynamo_compiling(): + assert B.shape[2:] == (num_queries, H * W) + return B + + def forward( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + reference_boxes: Optional[Tensor] = None, # num_queries, bs, 4 + # for memory + level_start_index: Optional[Tensor] = None, # num_levels + spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + valid_ratios: Optional[Tensor] = None, + # for text + memory_text: Optional[Tensor] = None, + text_attention_mask: Optional[Tensor] = None, + # if `apply_dac` is None, it will default to `self.dac` + apply_dac: Optional[bool] = None, + is_instance_prompt=False, + decoder_extra_kwargs: Optional[Dict] = None, + # ROI memory bank + obj_roi_memory_feat=None, + obj_roi_memory_mask=None, + box_head_trk=None, + ): + """ + Input: + - tgt: nq, bs, d_model + - memory: \\sum{hw}, bs, d_model + - pos: \\sum{hw}, bs, d_model + - reference_boxes: nq, bs, 4 (after sigmoid) + - valid_ratios/spatial_shapes: bs, nlevel, 2 + """ + if memory_mask is not None: + assert ( + self.boxRPB == "none" + ), "inputting a memory_mask in the presence of boxRPB is unexpected/not implemented" + + apply_dac = apply_dac if apply_dac is not None else self.dac + if apply_dac: + assert (tgt.shape[0] == self.num_queries) or ( + self.use_instance_query + and (tgt.shape[0] == self.instance_query_embed.num_embeddings) + ) + + tgt = tgt.repeat(2, 1, 1) + # note that we don't tile tgt_mask, since DAC doesn't + # use self-attention in o2m queries + if reference_boxes is not None: + assert (reference_boxes.shape[0] == self.num_queries) or ( + self.use_instance_query + and ( + reference_boxes.shape[0] + == self.instance_query_embed.num_embeddings + ) + ) + reference_boxes = reference_boxes.repeat(2, 1, 1) + + bs = tgt.shape[1] + intermediate = [] + intermediate_presence_logits = [] + presence_feats = None + + if self.box_refine: + if reference_boxes is None: + # In this case, we're in a one-stage model, so we generate the reference boxes + reference_boxes = self.reference_points.weight.unsqueeze(1) + reference_boxes = ( + reference_boxes.repeat(2, bs, 1) + if apply_dac + else reference_boxes.repeat(1, bs, 1) + ) + reference_boxes = reference_boxes.sigmoid() + intermediate_ref_boxes = [reference_boxes] + else: + reference_boxes = None + intermediate_ref_boxes = None + + output = tgt + presence_out = None + if self.presence_token is not None and is_instance_prompt is False: + # expand to batch dim + presence_out = self.presence_token.weight[None].expand(1, bs, -1) + + box_head = self.bbox_embed + if is_instance_prompt and self.instance_bbox_embed is not None: + box_head = self.instance_bbox_embed + + out_norm = self.norm + if is_instance_prompt and self.instance_norm is not None: + out_norm = self.instance_norm + + for layer_idx, layer in enumerate(self.layers): + reference_points_input = ( + reference_boxes[:, :, None] + * torch.cat([valid_ratios, valid_ratios], -1)[None, :] + ) # nq, bs, nlevel, 4 # [200, 1, 1, 4] + + query_sine_embed = gen_sineembed_for_position( + reference_points_input[:, :, 0, :], self.d_model + ) # nq, bs, d_model*2 # [200, 1, 512] + + # conditional query + query_pos = self.ref_point_head(query_sine_embed) # nq, bs, d_model # [200, 1, 256] + + if self.boxRPB != "none" and reference_boxes is not None: + assert ( + spatial_shapes.shape[0] == 1 + ), "only single scale support implemented" + memory_mask = self._get_rpb_matrix( + reference_boxes, + (spatial_shapes[0, 0], spatial_shapes[0, 1]), + ) + memory_mask = memory_mask.flatten(0, 1) # (bs*n_heads, nq, H*W) # [8, 200, 5184] + if self.training: + assert ( + self.use_act_checkpoint + ), "Activation checkpointing not enabled in the decoder" + output, presence_out = activation_ckpt_wrapper(layer)( + tgt=output, + tgt_query_pos=query_pos, + tgt_query_sine_embed=query_sine_embed, + tgt_key_padding_mask=tgt_key_padding_mask, + tgt_reference_points=reference_points_input, + memory_text=memory_text, + text_attention_mask=text_attention_mask, + memory=memory, + memory_key_padding_mask=memory_key_padding_mask, + memory_level_start_index=level_start_index, + memory_spatial_shapes=spatial_shapes, + memory_pos=pos, + self_attn_mask=tgt_mask, + cross_attn_mask=memory_mask, + dac=apply_dac, + dac_use_selfatt_ln=self.dac_use_selfatt_ln, + presence_token=presence_out, + **(decoder_extra_kwargs or {}), + act_ckpt_enable=self.training and self.use_act_checkpoint, + # ROI memory bank + obj_roi_memory_feat=obj_roi_memory_feat, + obj_roi_memory_mask=obj_roi_memory_mask, + ) + + # iter update + if self.box_refine: + reference_before_sigmoid = inverse_sigmoid(reference_boxes) + if box_head_trk is None: + # delta_unsig = self.bbox_embed(output) + if not self.use_normed_output_consistently: + delta_unsig = box_head(output) + else: + delta_unsig = box_head(out_norm(output)) + else: + # box_head_trk use a separate box head for tracking queries + Q_det = decoder_extra_kwargs["Q_det"] + assert output.size(0) >= Q_det + delta_unsig_det = self.bbox_embed(output[:Q_det]) + delta_unsig_trk = box_head_trk(output[Q_det:]) + delta_unsig = torch.cat([delta_unsig_det, delta_unsig_trk], dim=0) + outputs_unsig = delta_unsig + reference_before_sigmoid + new_reference_points = outputs_unsig.sigmoid() + + reference_boxes = new_reference_points.detach() + if layer_idx != self.num_layers - 1: + intermediate_ref_boxes.append(new_reference_points) + else: + raise NotImplementedError("not implemented yet") + + intermediate.append(out_norm(output)) + if self.presence_token is not None and is_instance_prompt is False: + # norm, mlp head + intermediate_layer_presence_logits = self.presence_token_head( + self.presence_token_out_norm(presence_out) + ).squeeze(-1) + + # clamp to mitigate numerical issues + if self.clamp_presence_logits: + intermediate_layer_presence_logits.clamp( + min=-self.clamp_presence_logit_max_val, + max=self.clamp_presence_logit_max_val, + ) + + intermediate_presence_logits.append(intermediate_layer_presence_logits) + presence_feats = presence_out.clone() + + if not self.compiled and self.compile_mode is not None: + self.forward = torch.compile( + self.forward, mode=self.compile_mode, fullgraph=True + ) + self.compiled = True + + return ( + torch.stack(intermediate), + torch.stack(intermediate_ref_boxes), + ( + torch.stack(intermediate_presence_logits) + if self.presence_token is not None and is_instance_prompt is False + else None + ), + presence_feats, + ) + + +class TransformerEncoderCrossAttention(nn.Module): + def __init__( + self, + d_model: int, + frozen: bool, + pos_enc_at_input: bool, + layer, + num_layers: int, + use_act_checkpoint: bool = False, + batch_first: bool = False, # Do layers expect batch first input? + # which layers to exclude cross attention? default: None, means all + # layers use cross attention + remove_cross_attention_layers: Optional[list] = None, + ): + super().__init__() + self.d_model = d_model + self.layers = get_clones(layer, num_layers) + self.num_layers = num_layers + self.norm = nn.LayerNorm(d_model) + self.pos_enc_at_input = pos_enc_at_input + self.use_act_checkpoint = use_act_checkpoint + + if frozen: + for p in self.parameters(): + p.requires_grad_(False) + + self.batch_first = batch_first + + # remove cross attention layers if specified + self.remove_cross_attention_layers = [False] * self.num_layers + if remove_cross_attention_layers is not None: + for i in remove_cross_attention_layers: + self.remove_cross_attention_layers[i] = True + assert len(self.remove_cross_attention_layers) == len(self.layers) + + for i, remove_cross_attention in enumerate(self.remove_cross_attention_layers): + if remove_cross_attention: + self.layers[i].cross_attn_image = None + self.layers[i].norm2 = None + self.layers[i].dropout2 = None + + def forward( + self, + src, # self-attention inputs + prompt, # cross-attention inputs + src_mask: Optional[Tensor] = None, # att.mask for self-attention inputs + prompt_mask: Optional[Tensor] = None, # att.mask for cross-attention inputs + src_key_padding_mask: Optional[Tensor] = None, + prompt_key_padding_mask: Optional[Tensor] = None, + src_pos: Optional[Tensor] = None, # pos_enc for self-attention inputs + prompt_pos: Optional[Tensor] = None, # pos_enc for cross-attention inputs + feat_sizes: Optional[list] = None, + num_obj_ptr_tokens: int = 0, # number of object pointer *tokens* + ): + if isinstance(src, list): + assert isinstance(src_key_padding_mask, list) and isinstance(src_pos, list) + assert len(src) == len(src_key_padding_mask) == len(src_pos) == 1 + src, src_key_padding_mask, src_pos = ( + src[0], + src_key_padding_mask[0], + src_pos[0], + ) + + assert ( + src.shape[1] == prompt.shape[1] + ), "Batch size must be the same for src and prompt" + + output = src + + if self.pos_enc_at_input and src_pos is not None: + output = output + 0.1 * src_pos + + if self.batch_first: + # Convert to batch first + output = output.transpose(0, 1) + src_pos = src_pos.transpose(0, 1) + prompt = prompt.transpose(0, 1) + prompt_pos = prompt_pos.transpose(0, 1) + + for layer in self.layers: + kwds = {} + if isinstance(layer.cross_attn_image, RoPEAttention): + kwds = {"num_k_exclude_rope": num_obj_ptr_tokens} + + output = activation_ckpt_wrapper(layer)( + tgt=output, + memory=prompt, + tgt_mask=src_mask, + memory_mask=prompt_mask, + tgt_key_padding_mask=src_key_padding_mask, + memory_key_padding_mask=prompt_key_padding_mask, + pos=prompt_pos, + query_pos=src_pos, + dac=False, + attn_bias=None, + act_ckpt_enable=self.training and self.use_act_checkpoint, + **kwds, + ) + normed_output = self.norm(output) + + if self.batch_first: + # Convert back to seq first + normed_output = normed_output.transpose(0, 1) + src_pos = src_pos.transpose(0, 1) + + return { + "memory": normed_output, + "pos_embed": src_pos, + "padding_mask": src_key_padding_mask, + } + + +class TransformerDecoderLayerv1(nn.Module): + def __init__( + self, + activation: str, + cross_attention: nn.Module, + d_model: int, + dim_feedforward: int, + dropout: float, + pos_enc_at_attn: bool, + pos_enc_at_cross_attn_keys: bool, + pos_enc_at_cross_attn_queries: bool, + pre_norm: bool, + self_attention: nn.Module, + ): + super().__init__() + self.d_model = d_model + self.dim_feedforward = dim_feedforward + self.dropout_value = dropout + self.self_attn = self_attention + self.cross_attn_image = cross_attention + + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation_str = activation + self.activation = get_activation_fn(activation) + self.pre_norm = pre_norm + + self.pos_enc_at_attn = pos_enc_at_attn + self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries + self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys + + def forward_post( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + **kwargs, + ): + q = k = tgt + query_pos if self.pos_enc_at_attn else tgt + + # Self attention + tgt2 = self.self_attn( + q, + k, + value=tgt, + attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask, + )[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + + # Cross attention to image + tgt2 = self.cross_attn_image( + query=tgt + query_pos if self.pos_enc_at_cross_attn_queries else tgt, + key=memory + pos if self.pos_enc_at_cross_attn_keys else memory, + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + )[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + + # FFN + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward_pre( + self, + tgt, + memory, + dac: bool = False, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + attn_bias: Optional[Tensor] = None, + **kwargs, + ): + if dac: + # we only apply self attention to the first half of the queries + assert tgt.shape[0] % 2 == 0 + other_tgt = tgt[tgt.shape[0] // 2 :] + tgt = tgt[: tgt.shape[0] // 2] + tgt2 = self.norm1(tgt) + q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2 + tgt2 = self.self_attn( + q, + k, + value=tgt2, + attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask, + )[0] + tgt = tgt + self.dropout1(tgt2) + if dac: + # Recombine + tgt = torch.cat((tgt, other_tgt), dim=0) + tgt2 = self.norm2(tgt) + tgt2 = self.cross_attn_image( + query=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2, + key=memory + pos if self.pos_enc_at_cross_attn_keys else memory, + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + attn_bias=attn_bias, + )[0] + tgt = tgt + self.dropout2(tgt2) + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + def forward( + self, + tgt, + memory, + dac: bool = False, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + attn_bias: Optional[Tensor] = None, + **kwds: Any, + ) -> torch.Tensor: + fwd_fn = self.forward_pre if self.pre_norm else self.forward_post + return fwd_fn( + tgt, + memory, + dac=dac, + tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + pos=pos, + query_pos=query_pos, + attn_bias=attn_bias, + **kwds, + ) + + +class TransformerDecoderLayerv2(TransformerDecoderLayerv1): + def __init__(self, cross_attention_first=False, *args: Any, **kwds: Any): + super().__init__(*args, **kwds) + self.cross_attention_first = cross_attention_first + + def _forward_sa(self, tgt, query_pos): + # Self-Attention + tgt2 = self.norm1(tgt) + q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2 + tgt2 = self.self_attn(q, k, v=tgt2) + tgt = tgt + self.dropout1(tgt2) + return tgt + + def _forward_ca(self, tgt, memory, query_pos, pos, num_k_exclude_rope=0): + if self.cross_attn_image is None: + return tgt + + kwds = {} + if num_k_exclude_rope > 0: + assert isinstance(self.cross_attn_image, RoPEAttention) + kwds = {"num_k_exclude_rope": num_k_exclude_rope} + + # Cross-Attention + tgt2 = self.norm2(tgt) + tgt2 = self.cross_attn_image( + q=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2, + k=memory + pos if self.pos_enc_at_cross_attn_keys else memory, + v=memory, + **kwds, + ) + tgt = tgt + self.dropout2(tgt2) + return tgt + + def forward_pre( + self, + tgt, + memory, + dac: bool, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + attn_bias: Optional[Tensor] = None, + num_k_exclude_rope: int = 0, + ): + assert dac is False + assert tgt_mask is None + assert memory_mask is None + assert tgt_key_padding_mask is None + assert memory_key_padding_mask is None + assert attn_bias is None + + if self.cross_attention_first: + tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope) + tgt = self._forward_sa(tgt, query_pos) + else: + tgt = self._forward_sa(tgt, query_pos) + tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope) + + # MLP + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + def forward(self, *args: Any, **kwds: Any) -> torch.Tensor: + if self.pre_norm: + return self.forward_pre(*args, **kwds) + raise NotImplementedError diff --git a/src/nn/segearth_ov3/sam3/model/edt.py b/src/nn/segearth_ov3/sam3/model/edt.py new file mode 100644 index 0000000..9448c1d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/edt.py @@ -0,0 +1,173 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Triton kernel for euclidean distance transform (EDT)""" + +import torch +import triton +import triton.language as tl + +""" +Disclaimer: This implementation is not meant to be extremely efficient. A CUDA kernel would likely be more efficient. +Even in Triton, there may be more suitable algorithms. + +The goal of this kernel is to mimic cv2.distanceTransform(input, cv2.DIST_L2, 0). +Recall that the euclidean distance transform (EDT) calculates the L2 distance to the closest zero pixel for each pixel of the source image. + +For images of size NxN, the naive algorithm would be to compute pairwise distances between every pair of points, leading to a O(N^4) algorithm, which is obviously impractical. +One can do better using the following approach: +- First, compute the distance to the closest point in the same row. We can write it as Row_EDT[i,j] = min_k (sqrt((k-j)^2) if input[i,k]==0 else +infinity). With a naive implementation, this step has a O(N^3) complexity +- Then, because of triangular inequality, we notice that the EDT for a given location [i,j] is the min of the row EDTs in the same column. EDT[i,j] = min_k Row_EDT[k, j]. This is also O(N^3) + +Overall, this algorithm is quite amenable to parallelization, and has a complexity O(N^3). Can we do better? + +It turns out that we can leverage the structure of the L2 distance (nice and convex) to find the minimum in a more efficient way. +We follow the algorithm from "Distance Transforms of Sampled Functions" (https://cs.brown.edu/people/pfelzens/papers/dt-final.pdf), which is also what's implemented in opencv + +For a single dimension EDT, we can compute the EDT of an arbitrary function F, that we discretize over the grid. Note that for the binary EDT that we're interested in, we can set F(i,j) = 0 if input[i,j]==0 else +infinity +For now, we'll compute the EDT squared, and will take the sqrt only at the very end. +The basic idea is that each point at location i spawns a parabola around itself, with a bias equal to F(i). So specifically, we're looking at the parabola (x - i)^2 + F(i) +When we're looking for the row EDT at location j, we're effectively looking for min_i (x-i)^2 + F(i). In other word we want to find the lowest parabola at location j. + +To do this efficiently, we need to maintain the lower envelope of the union of parabolas. This can be constructed on the fly using a sort of stack approach: + - every time we want to add a new parabola, we check if it may be covering the current right-most parabola. If so, then that parabola was useless, so we can pop it from the stack + - repeat until we can't find any more parabola to pop. Then push the new one. + +This algorithm runs in O(N) for a single row, so overall O(N^2) when applied to all rows +Similarly as before, we notice that we can decompose the algorithm for rows and columns, leading to an overall run-time of O(N^2) + +This algorithm is less suited for to GPUs, since the one-dimensional EDT computation is quite sequential in nature. However, we can parallelize over batch and row dimensions. +In Triton, things are particularly bad at the moment, since there is no support for reading/writing to the local memory at a specific index (a local gather is coming soon, see https://github.com/triton-lang/triton/issues/974, but no mention of writing, ie scatter) +One could emulate these operations with masking, but in initial tests, it proved to be worst than naively reading and writing to the global memory. My guess is that the cache is compensating somewhat for the repeated single-point accesses. + + +The timing obtained on a H100 for a random batch of masks of dimension 256 x 1024 x 1024 are as follows: +- OpenCV: 1780ms (including round-trip to cpu, but discounting the fact that it introduces a synchronization point) +- triton, O(N^3) algo: 627ms +- triton, O(N^2) algo: 322ms + +Overall, despite being quite naive, this implementation is roughly 5.5x faster than the openCV cpu implem + +""" + + +@triton.jit +def edt_kernel(inputs_ptr, outputs_ptr, v, z, height, width, horizontal: tl.constexpr): + # This is a somewhat verbatim implementation of the efficient 1D EDT algorithm described above + # It can be applied horizontally or vertically depending if we're doing the first or second stage. + # It's parallelized across batch+row (or batch+col if horizontal=False) + # TODO: perhaps the implementation can be revisited if/when local gather/scatter become available in triton + batch_id = tl.program_id(axis=0) + if horizontal: + row_id = tl.program_id(axis=1) + block_start = (batch_id * height * width) + row_id * width + length = width + stride = 1 + else: + col_id = tl.program_id(axis=1) + block_start = (batch_id * height * width) + col_id + length = height + stride = width + + # This will be the index of the right most parabola in the envelope ("the top of the stack") + k = 0 + for q in range(1, length): + # Read the function value at the current location. Note that we're doing a singular read, not very efficient + cur_input = tl.load(inputs_ptr + block_start + (q * stride)) + # location of the parabola on top of the stack + r = tl.load(v + block_start + (k * stride)) + # associated boundary + z_k = tl.load(z + block_start + (k * stride)) + # value of the function at the parabola location + previous_input = tl.load(inputs_ptr + block_start + (r * stride)) + # intersection between the two parabolas + s = (cur_input - previous_input + q * q - r * r) / (q - r) / 2 + + # we'll pop as many parabolas as required + while s <= z_k and k - 1 >= 0: + k = k - 1 + r = tl.load(v + block_start + (k * stride)) + z_k = tl.load(z + block_start + (k * stride)) + previous_input = tl.load(inputs_ptr + block_start + (r * stride)) + s = (cur_input - previous_input + q * q - r * r) / (q - r) / 2 + + # Store the new one + k = k + 1 + tl.store(v + block_start + (k * stride), q) + tl.store(z + block_start + (k * stride), s) + if k + 1 < length: + tl.store(z + block_start + ((k + 1) * stride), 1e9) + + # Last step, we read the envelope to find the min in every location + k = 0 + for q in range(length): + while ( + k + 1 < length + and tl.load( + z + block_start + ((k + 1) * stride), mask=(k + 1) < length, other=q + ) + < q + ): + k += 1 + r = tl.load(v + block_start + (k * stride)) + d = q - r + old_value = tl.load(inputs_ptr + block_start + (r * stride)) + tl.store(outputs_ptr + block_start + (q * stride), old_value + d * d) + + +def edt_triton(data: torch.Tensor): + """ + Computes the Euclidean Distance Transform (EDT) of a batch of binary images. + + Args: + data: A tensor of shape (B, H, W) representing a batch of binary images. + + Returns: + A tensor of the same shape as data containing the EDT. + It should be equivalent to a batched version of cv2.distanceTransform(input, cv2.DIST_L2, 0) + """ + assert data.dim() == 3 + assert data.is_cuda + B, H, W = data.shape + data = data.contiguous() + + # Allocate the "function" tensor. Implicitly the function is 0 if data[i,j]==0 else +infinity + output = torch.where(data, 1e18, 0.0) + assert output.is_contiguous() + + # Scratch tensors for the parabola stacks + parabola_loc = torch.zeros(B, H, W, dtype=torch.uint32, device=data.device) + parabola_inter = torch.empty(B, H, W, dtype=torch.float, device=data.device) + parabola_inter[:, :, 0] = -1e18 + parabola_inter[:, :, 1] = 1e18 + + # Grid size (number of blocks) + grid = (B, H) + + # Launch initialization kernel + edt_kernel[grid]( + output.clone(), + output, + parabola_loc, + parabola_inter, + H, + W, + horizontal=True, + ) + + # reset the parabola stacks + parabola_loc.zero_() + parabola_inter[:, :, 0] = -1e18 + parabola_inter[:, :, 1] = 1e18 + + grid = (B, W) + edt_kernel[grid]( + output.clone(), + output, + parabola_loc, + parabola_inter, + H, + W, + horizontal=False, + ) + # don't forget to take sqrt at the end + return output.sqrt() diff --git a/src/nn/segearth_ov3/sam3/model/encoder.py b/src/nn/segearth_ov3/sam3/model/encoder.py new file mode 100644 index 0000000..842bc56 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/encoder.py @@ -0,0 +1,594 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +# Based on https://github.com/IDEA-Research/GroundingDINO + +from typing import Any, Dict, List, Optional, Tuple + +import torch +from torch import nn, Tensor + +from .act_ckpt_utils import activation_ckpt_wrapper +from .model_misc import get_activation_fn, get_clones, get_valid_ratio + + +class TransformerEncoderLayer(nn.Module): + """ + Transformer encoder layer that performs self-attention followed by cross-attention. + + This layer was previously called TransformerDecoderLayer but was renamed to better + reflect its role in the architecture. It processes input sequences through self-attention + and then cross-attention with another input (typically image features). + + The layer supports both pre-norm and post-norm configurations, as well as + positional encoding at different stages of the attention mechanism. + """ + + def __init__( + self, + activation: str, + cross_attention: nn.Module, + d_model: int, + dim_feedforward: int, + dropout: float, + pos_enc_at_attn: bool, + pos_enc_at_cross_attn_keys: bool, + pos_enc_at_cross_attn_queries: bool, + pre_norm: bool, + self_attention: nn.Module, + ): + """ + Initialize a transformer encoder layer. + + Args: + activation: Activation function to use in the feedforward network + cross_attention: Cross-attention module for attending to image features + d_model: Model dimension/hidden size + dim_feedforward: Dimension of the feedforward network + dropout: Dropout probability + pos_enc_at_attn: Whether to add positional encodings at self-attention + pos_enc_at_cross_attn_keys: Whether to add positional encodings to keys in cross-attention + pos_enc_at_cross_attn_queries: Whether to add positional encodings to queries in cross-attention + pre_norm: Whether to use pre-norm (True) or post-norm (False) architecture + self_attention: Self-attention module + """ + super().__init__() + self.d_model = d_model + self.dim_feedforward = dim_feedforward + self.dropout_value = dropout + self.self_attn = self_attention + self.cross_attn_image = cross_attention + + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation_str = activation + self.activation = get_activation_fn(activation) + self.pre_norm = pre_norm + + self.pos_enc_at_attn = pos_enc_at_attn + self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries + self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys + + self.layer_idx = None + + def forward_post( + self, + tgt: Tensor, + memory: Tensor, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + **kwargs, + ) -> Tensor: + """ + Forward pass for post-norm architecture. + + In post-norm architecture, normalization is applied after attention and feedforward operations. + + Args: + tgt: Input tensor to be processed + memory: Memory tensor for cross-attention + tgt_mask: Mask for self-attention + memory_mask: Mask for cross-attention + tgt_key_padding_mask: Key padding mask for self-attention + memory_key_padding_mask: Key padding mask for cross-attention + pos: Positional encoding for memory + query_pos: Positional encoding for query + **kwargs: Additional keyword arguments + + Returns: + Processed tensor + """ + q = k = tgt + query_pos if self.pos_enc_at_attn else tgt + + # Self attention + tgt2 = self.self_attn( + q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask + )[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + + # Cross attention to image + tgt2 = self.cross_attn_image( + query=tgt + query_pos if self.pos_enc_at_cross_attn_queries else tgt, + key=memory + pos if self.pos_enc_at_cross_attn_keys else memory, + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + )[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + + # FFN + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward_pre( + self, + tgt: Tensor, + memory: Tensor, + dac: bool = False, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + # attn_bias: Optional[Tensor] = None, + # **kwargs, + ) -> Tensor: + """ + Forward pass for pre-norm architecture. + + In pre-norm architecture, normalization is applied before attention and feedforward operations. + + Args: + tgt: Input tensor to be processed + memory: Memory tensor for cross-attention + dac: Whether to use Divide-and-Conquer attention + tgt_mask: Mask for self-attention + memory_mask: Mask for cross-attention + tgt_key_padding_mask: Key padding mask for self-attention + memory_key_padding_mask: Key padding mask for cross-attention + pos: Positional encoding for memory + query_pos: Positional encoding for query + attn_bias: Optional attention bias tensor + **kwargs: Additional keyword arguments + + Returns: + Processed tensor + """ + if dac: + # we only apply self attention to the first half of the queries + assert tgt.shape[0] % 2 == 0 + other_tgt = tgt[tgt.shape[0] // 2 :] + tgt = tgt[: tgt.shape[0] // 2] + tgt2 = self.norm1(tgt) + q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2 + tgt2 = self.self_attn( + q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask + )[0] + tgt = tgt + self.dropout1(tgt2) + if dac: + # Recombine + tgt = torch.cat((tgt, other_tgt), dim=0) + tgt2 = self.norm2(tgt) + tgt2 = self.cross_attn_image( + query=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2, + key=memory + pos if self.pos_enc_at_cross_attn_keys else memory, + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + # attn_bias=attn_bias, + )[0] + tgt = tgt + self.dropout2(tgt2) + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + def forward( + self, + tgt: Tensor, + memory: Tensor, + dac: bool = False, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + # attn_bias: Optional[Tensor] = None, + # **kwds: Any, + ) -> torch.Tensor: + """ + Forward pass for the transformer encoder layer. + + Args: + tgt: Input tensor to be processed + memory: Memory tensor (e.g., image features) for cross-attention + dac: Whether to use Divide-and-Conquer attention (only apply self-attention to first half) + tgt_mask: Mask for self-attention + memory_mask: Mask for cross-attention + tgt_key_padding_mask: Key padding mask for self-attention + memory_key_padding_mask: Key padding mask for cross-attention + pos: Positional encoding for memory + query_pos: Positional encoding for query + attn_bias: Optional attention bias tensor + **kwds: Additional keyword arguments + + Returns: + Processed tensor after self-attention, cross-attention, and feedforward network + """ + fwd_fn = self.forward_pre if self.pre_norm else self.forward_post + return fwd_fn( + tgt, + memory, + dac=dac, + tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + pos=pos, + query_pos=query_pos, + # attn_bias=attn_bias, + # **kwds, + ) + + +class TransformerEncoder(nn.Module): + """ + Transformer encoder that processes multi-level features. + + This encoder takes multi-level features (e.g., from a backbone network) and processes + them through a stack of transformer encoder layers. It supports features from multiple + levels (e.g., different resolutions) and can apply activation checkpointing for memory + efficiency during training. + + Args: + layer: The encoder layer to be stacked multiple times + num_layers: Number of encoder layers to stack + d_model: Model dimension/hidden size + num_feature_levels: Number of feature levels to process + frozen: Whether to freeze the parameters of this module + use_act_checkpoint: Whether to use activation checkpointing during training + """ + + def __init__( + self, + layer: nn.Module, + num_layers: int, + d_model: int, + num_feature_levels: int, + frozen: bool = False, + use_act_checkpoint: bool = False, + ): + super().__init__() + self.layers = get_clones(layer, num_layers) + self.num_layers = num_layers + + self.num_feature_levels = num_feature_levels + self.level_embed = None + if num_feature_levels > 1: + self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model)) + + if frozen: + for p in self.parameters(): + p.requires_grad_(False) + + self.use_act_checkpoint = use_act_checkpoint + + # assign layer index to each layer so that some layers can decide what to do + # based on which layer index they are (e.g. cross attention to memory bank only + # in selected layers) + for layer_idx, layer in enumerate(self.layers): + layer.layer_idx = layer_idx + + @staticmethod + def get_reference_points(spatial_shapes, valid_ratios, device): + with torch.no_grad(): + reference_points_list = [] + for lvl, (H_, W_) in enumerate(spatial_shapes): + ref_y, ref_x = torch.meshgrid( + torch.linspace( + 0.5, H_ - 0.5, H_, dtype=torch.float32, device=device + ), + torch.linspace( + 0.5, W_ - 0.5, W_, dtype=torch.float32, device=device + ), + ) + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + + return reference_points + + def _prepare_multilevel_features(self, srcs, masks, pos_embeds): + assert ( + len(srcs) == self.num_feature_levels + ), "mismatch between expected and received # of feature levels" + + src_flatten = [] + mask_flatten = [] + lvl_pos_embed_flatten = [] + spatial_shapes = [] + has_mask = masks is not None and masks[0] is not None + for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)): + bs, c, h, w = src.shape + spatial_shape = (h, w) + spatial_shapes.append(spatial_shape) + + src = src.flatten(2).transpose(1, 2) # bs, hw, c + if has_mask: + mask = mask.flatten(1) + pos_embed = pos_embed.flatten(2).transpose(1, 2) # bs, hw, c + if self.level_embed is not None: + lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1) + else: + lvl_pos_embed = pos_embed + lvl_pos_embed_flatten.append(lvl_pos_embed) + src_flatten.append(src) + if has_mask: + mask_flatten.append(mask) + src_flatten = torch.cat(src_flatten, 1) # bs, \sum{hxw}, c + mask_flatten = torch.cat(mask_flatten, 1) if has_mask else None # bs, \sum{hxw} + lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) # bs, \sum{hxw}, c + spatial_shapes = torch.tensor( + spatial_shapes, dtype=torch.long, device=src_flatten.device + ) + level_start_index = torch.cat( + ( + spatial_shapes.new_zeros((1,)), + spatial_shapes.prod(1).cumsum(0)[:-1], + ) + ) + if has_mask: + valid_ratios = torch.stack([get_valid_ratio(m) for m in masks], 1) + else: + valid_ratios = torch.ones( + (src_flatten.shape[0], self.num_feature_levels, 2), + device=src_flatten.device, + ) + + return ( + src_flatten, + mask_flatten, + lvl_pos_embed_flatten, + level_start_index, + valid_ratios, + spatial_shapes, + ) + + def forward( + self, + src: List[Tensor], + src_key_padding_masks: Optional[List[Tensor]] = None, + pos: Optional[List[Tensor]] = None, + prompt: Optional[Tensor] = None, + prompt_key_padding_mask: Optional[Tensor] = None, + encoder_extra_kwargs: Optional[Dict] = None, + ) -> Tuple[Tensor, Optional[Tensor], Tensor, Tensor, Tensor, Tensor]: + """ + Process multi-level features through the transformer encoder. + + Args: + src: List of multi-level features, each with shape (batch_size, channels, height, width) + src_key_padding_masks: List of padding masks for each feature level, each with shape (batch_size, height, width) + pos: List of positional embeddings for each feature level, each with shape (batch_size, channels, height, width) + prompt: Optional text/prompt features to attend to, with shape (seq_len, batch_size, d_model) + prompt_key_padding_mask: Optional padding mask for prompt, with shape (batch_size, seq_len) + encoder_extra_kwargs: Optional additional arguments to pass to each encoder layer + + Returns: + A tuple containing: + - output: Processed features with shape (seq_len, batch_size, d_model) + - key_padding_masks_flatten: Flattened padding masks + - lvl_pos_embed_flatten: Flattened positional embeddings + - level_start_index: Starting indices for each feature level + - spatial_shapes: Spatial dimensions of each feature level + - valid_ratios: Valid ratios for each feature level + """ + assert ( + len(src) == self.num_feature_levels + ), "must be equal to num_feature_levels" + if src_key_padding_masks is not None: + assert len(src_key_padding_masks) == self.num_feature_levels + if pos is not None: + assert len(pos) == self.num_feature_levels + # Flatten multilevel feats and add level pos embeds + ( + src_flatten, + key_padding_masks_flatten, + lvl_pos_embed_flatten, + level_start_index, + valid_ratios, + spatial_shapes, + ) = self._prepare_multilevel_features(src, src_key_padding_masks, pos) + + reference_points = self.get_reference_points( + spatial_shapes, valid_ratios, device=src_flatten.device + ) + + output = src_flatten + for layer in self.layers: + layer_kwargs = {} + + assert isinstance(layer, TransformerEncoderLayer) + layer_kwargs["memory"] = prompt + layer_kwargs["memory_key_padding_mask"] = prompt_key_padding_mask + layer_kwargs["query_pos"] = lvl_pos_embed_flatten + layer_kwargs["tgt"] = output + layer_kwargs["tgt_key_padding_mask"] = key_padding_masks_flatten + + if self.training: + assert self.use_act_checkpoint, "activation ckpt not enabled in encoder" + if encoder_extra_kwargs is not None: + layer_kwargs.update(encoder_extra_kwargs) + output = activation_ckpt_wrapper(layer)( + **layer_kwargs, + act_ckpt_enable=self.training and self.use_act_checkpoint, + ) + # return as seq first + return ( + output.transpose(0, 1), + ( + key_padding_masks_flatten.transpose(0, 1) + if key_padding_masks_flatten is not None + else None + ), + lvl_pos_embed_flatten.transpose(0, 1), + level_start_index, + spatial_shapes, + valid_ratios, + ) + + +class TransformerEncoderFusion(TransformerEncoder): + """ + Transformer encoder that fuses text and image features. + + This encoder extends TransformerEncoder to handle both text and image features, + with the ability to add pooled text features to image features for better + cross-modal fusion. It supports torch.compile for performance optimization. + + Args: + layer: The encoder layer to be stacked multiple times + num_layers: Number of encoder layers to stack + d_model: Model dimension/hidden size + num_feature_levels: Number of feature levels to process + add_pooled_text_to_img_feat: Whether to add pooled text features to image features + pool_text_with_mask: Whether to use the mask when pooling text features + compile_mode: Mode for torch.compile, or None to disable compilation + **kwargs: Additional arguments to pass to the parent class + """ + + def __init__( + self, + layer: nn.Module, + num_layers: int, + d_model: int, + num_feature_levels: int, + add_pooled_text_to_img_feat: bool = True, + pool_text_with_mask: bool = False, + compile_mode: Optional[str] = None, + **kwargs, + ): + super().__init__( + layer, + num_layers, + d_model, + num_feature_levels, + **kwargs, + ) + self.add_pooled_text_to_img_feat = add_pooled_text_to_img_feat + if self.add_pooled_text_to_img_feat: + self.text_pooling_proj = nn.Linear(d_model, d_model) + self.pool_text_with_mask = pool_text_with_mask + if compile_mode is not None: + self.forward = torch.compile( + self.forward, mode=compile_mode, fullgraph=True + ) + + @staticmethod + def get_reference_points(spatial_shapes, valid_ratios, device): + # Not needed here + return None + + def forward( + self, + src: List[Tensor], + prompt: Tensor, + src_key_padding_mask: Optional[List[Tensor]] = None, + src_pos: Optional[List[Tensor]] = None, + prompt_key_padding_mask: Optional[Tensor] = None, + prompt_pos: Optional[Tensor] = None, + feat_sizes: Optional[List[int]] = None, + encoder_extra_kwargs: Optional[Dict] = None, + ): + # Restore spatial shapes of vision + bs = src[0].shape[1] # seq first + if feat_sizes is not None: + assert len(feat_sizes) == len(src) + if src_key_padding_mask is None: + src_key_padding_mask = [None] * len(src) + for i, (h, w) in enumerate(feat_sizes): + src[i] = src[i].reshape(h, w, bs, -1).permute(2, 3, 0, 1) + src_pos[i] = src_pos[i].reshape(h, w, bs, -1).permute(2, 3, 0, 1) + src_key_padding_mask[i] = ( + src_key_padding_mask[i].reshape(h, w, bs).permute(2, 0, 1) + if src_key_padding_mask[i] is not None + else None + ) + else: + assert all( + x.dim == 4 for x in src + ), "expected list of (bs, c, h, w) tensors" + + if self.add_pooled_text_to_img_feat: + # Fusion: Add mean pooled text to image features + pooled_text = pool_text_feat( + prompt, prompt_key_padding_mask, self.pool_text_with_mask + ) + pooled_text = self.text_pooling_proj(pooled_text)[ + ..., None, None + ] # prompt is seq first + src = [x.add_(pooled_text) for x in src] + + ( + out, + key_padding_masks_flatten, + lvl_pos_embed_flatten, + level_start_index, + spatial_shapes, + valid_ratios, + ) = super().forward( + src, + src_key_padding_masks=src_key_padding_mask, + pos=src_pos, + prompt=prompt.transpose(0, 1), + prompt_key_padding_mask=prompt_key_padding_mask, + encoder_extra_kwargs=encoder_extra_kwargs, + ) + + return { + "memory": out, + "padding_mask": key_padding_masks_flatten, + "pos_embed": lvl_pos_embed_flatten, + "memory_text": prompt, + "level_start_index": level_start_index, + "spatial_shapes": spatial_shapes, + "valid_ratios": valid_ratios, + } + + +def pool_text_feat(prompt, prompt_mask, pool_with_mask): + # prompt has shape (seq, bs, dim) + if not pool_with_mask: + return prompt.mean(dim=0) + + # prompt_mask has shape (bs, seq), where False is valid and True is padding + assert prompt_mask.dim() == 2 + # is_valid has shape (seq, bs, 1), where 1 is valid and 0 is padding + is_valid = (~prompt_mask).float().permute(1, 0)[..., None] + # num_valid has shape (bs, 1) + num_valid = torch.clamp(torch.sum(is_valid, dim=0), min=1.0) + + # mean pool over all the valid tokens + pooled_text = (prompt * is_valid).sum(dim=0) / num_valid + return pooled_text diff --git a/src/nn/segearth_ov3/sam3/model/geometry_encoders.py b/src/nn/segearth_ov3/sam3/model/geometry_encoders.py new file mode 100644 index 0000000..b881784 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/geometry_encoders.py @@ -0,0 +1,852 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Tuple + +import torch +import torch.nn as nn +import torchvision +from typing_extensions import override + +from .act_ckpt_utils import activation_ckpt_wrapper +from .box_ops import box_cxcywh_to_xyxy + +from .model_misc import get_clones + + +def is_right_padded(mask): + """Given a padding mask (following pytorch convention, 1s for padded values), + returns whether the padding is on the right or not.""" + return (mask.long() == torch.sort(mask.long(), dim=-1)[0]).all() + + +def concat_padded_sequences(seq1, mask1, seq2, mask2, return_index: bool = False): + """ + Concatenates two right-padded sequences, such that the resulting sequence + is contiguous and also right-padded. + + Following pytorch's convention, tensors are sequence first, and the mask are + batch first, with 1s for padded values. + + :param seq1: A tensor of shape (seq1_length, batch_size, hidden_size). + :param mask1: A tensor of shape (batch_size, seq1_length). + :param seq2: A tensor of shape (seq2_length, batch_size, hidden_size). + :param mask2: A tensor of shape (batch_size, seq2_length). + :param return_index: If True, also returns the index of the ids of the element of seq2 + in the concatenated sequence. This can be used to retrieve the elements of seq2 + :return: A tuple (concatenated_sequence, concatenated_mask) if return_index is False, + otherwise (concatenated_sequence, concatenated_mask, index). + """ + seq1_length, batch_size, hidden_size = seq1.shape + seq2_length, batch_size, hidden_size = seq2.shape + + assert batch_size == seq1.size(1) == seq2.size(1) == mask1.size(0) == mask2.size(0) + assert hidden_size == seq1.size(2) == seq2.size(2) + assert seq1_length == mask1.size(1) + assert seq2_length == mask2.size(1) + + torch._assert_async(is_right_padded(mask1)) + torch._assert_async(is_right_padded(mask2)) + + actual_seq1_lengths = (~mask1).sum(dim=-1) + actual_seq2_lengths = (~mask2).sum(dim=-1) + + final_lengths = actual_seq1_lengths + actual_seq2_lengths + max_length = seq1_length + seq2_length + concatenated_mask = ( + torch.arange(max_length, device=seq2.device)[None].repeat(batch_size, 1) + >= final_lengths[:, None] + ) + + # (max_len, batch_size, hidden_size) + concatenated_sequence = torch.zeros( + (max_length, batch_size, hidden_size), device=seq2.device, dtype=seq2.dtype + ) + concatenated_sequence[:seq1_length, :, :] = seq1 + + # At this point, the element of seq1 are in the right place + # We just need to shift the elements of seq2 + + index = torch.arange(seq2_length, device=seq2.device)[:, None].repeat(1, batch_size) + index = index + actual_seq1_lengths[None] + + concatenated_sequence = concatenated_sequence.scatter( + 0, index[:, :, None].expand(-1, -1, hidden_size), seq2 + ) + + if return_index: + return concatenated_sequence, concatenated_mask, index + + return concatenated_sequence, concatenated_mask + + +class Prompt: + """Utility class to manipulate geometric prompts. + + We expect the sequences in pytorch convention, that is sequence first, batch second + The dimensions are expected as follows: + box_embeddings shape: N_boxes x B x C_box + box_mask shape: B x N_boxes. Can be None if nothing is masked out + point_embeddings shape: N_points x B x C_point + point_mask shape: B x N_points. Can be None if nothing is masked out + mask_embeddings shape: N_masks x B x 1 x H_mask x W_mask + mask_mask shape: B x N_masks. Can be None if nothing is masked out + + We also store positive/negative labels. These tensors are also stored batch-first + If they are None, we'll assume positive labels everywhere + box_labels: long tensor of shape N_boxes x B + point_labels: long tensor of shape N_points x B + mask_labels: long tensor of shape N_masks x B + """ + + def __init__( + self, + box_embeddings=None, + box_mask=None, + point_embeddings=None, + point_mask=None, + box_labels=None, + point_labels=None, + mask_embeddings=None, + mask_mask=None, # Attention mask for mask prompt + mask_labels=None, + ): + # Check for null prompt + if ( + box_embeddings is None + and point_embeddings is None + and mask_embeddings is None + ): + self.box_embeddings = None + self.box_labels = None + self.box_mask = None + self.point_embeddings = None + self.point_labels = None + self.point_mask = None + self.mask_embeddings = None + self.mask_mask = None + # Masks are assumed positive only for now. + self.mask_labels = None + return + # Get sequence lengths and device + box_seq_len, point_seq_len, mask_seq_len, bs, device = ( + self._init_seq_len_and_device( + box_embeddings, point_embeddings, mask_embeddings + ) + ) + + # Initialize embeds, labels, attention masks. + box_embeddings, box_labels, box_mask = self._init_box( + box_embeddings, box_labels, box_mask, box_seq_len, bs, device + ) + point_embeddings, point_labels, point_mask = self._init_point( + point_embeddings, point_labels, point_mask, point_seq_len, bs, device + ) + mask_embeddings, mask_labels, mask_mask = self._init_mask( + mask_embeddings, mask_labels, mask_mask, mask_seq_len, bs, device + ) + + # Dimension checks + assert ( + box_embeddings is not None + and list(box_embeddings.shape[:2]) + == [ + box_seq_len, + bs, + ] + ), f"Wrong dimension for box embeddings. Expected [{box_seq_len}, {bs}, *] got {box_embeddings.shape}" + assert ( + box_mask is not None + and list(box_mask.shape) + == [ + bs, + box_seq_len, + ] + ), f"Wrong dimension for box mask. Expected [{bs}, {box_seq_len}] got {box_mask.shape}" + assert ( + point_embeddings is not None + and list(point_embeddings.shape[:2]) + == [ + point_seq_len, + bs, + ] + ), f"Wrong dimension for point embeddings. Expected [{point_seq_len}, {bs}, *] got {point_embeddings.shape}" + assert ( + point_mask is not None + and list(point_mask.shape) + == [ + bs, + point_seq_len, + ] + ), f"Wrong dimension for point mask. Expected [{bs}, {point_seq_len}] got {point_mask.shape}" + assert ( + box_labels is not None + and list(box_labels.shape) + == [ + box_seq_len, + bs, + ] + ), f"Wrong dimension for box labels. Expected [{box_seq_len}, {bs}] got {box_labels.shape}" + assert ( + point_labels is not None + and list(point_labels.shape) + == [ + point_seq_len, + bs, + ] + ), f"Wrong dimension for point labels. Expected [{point_seq_len}, {bs}] got {point_labels.shape}" + assert ( + # Allowed to be None, we leave it to the encoder to check for validity before encoding. + mask_embeddings is None + or list(mask_embeddings.shape[:2]) + == [ + mask_seq_len, + bs, + ] + ), f"Wrong dimension for mask embeddings. Expected [{mask_seq_len}, {bs}, *] got {mask_embeddings.shape}" + assert ( + mask_mask is None + or list(mask_mask.shape) + == [ + bs, + mask_seq_len, + ] + ), f"Wrong dimension for mask attn. mask. Expected [{bs}, {mask_seq_len}] got {mask_mask.shape}" + + # Device checks + assert ( + box_embeddings is not None and box_embeddings.device == device + ), f"Expected box embeddings to be on device {device}, got {box_embeddings.device}" + assert ( + box_mask is not None and box_mask.device == device + ), f"Expected box mask to be on device {device}, got {box_mask.device}" + assert ( + box_labels is not None and box_labels.device == device + ), f"Expected box labels to be on device {device}, got {box_labels.device}" + assert ( + point_embeddings is not None and point_embeddings.device == device + ), f"Expected point embeddings to be on device {device}, got {point_embeddings.device}" + assert ( + point_mask is not None and point_mask.device == device + ), f"Expected point mask to be on device {device}, got {point_mask.device}" + assert ( + point_labels is not None and point_labels.device == device + ), f"Expected point labels to be on device {device}, got {point_labels.device}" + assert ( + mask_embeddings is None or mask_embeddings.device == device + ), f"Expected mask embeddings to be on device {device}, got {mask_embeddings.device}" + assert ( + mask_mask is None or mask_mask.device == device + ), f"Expected mask attn. mask to be on device {device}, got {mask_mask.device}" + + self.box_embeddings = box_embeddings + self.point_embeddings = point_embeddings + self.box_mask = box_mask + self.point_mask = point_mask + self.box_labels = box_labels + self.point_labels = point_labels + self.mask_embeddings = mask_embeddings + self.mask_labels = mask_labels + self.mask_mask = mask_mask + + def _init_seq_len_and_device( + self, box_embeddings, point_embeddings, mask_embeddings + ): + box_seq_len = point_seq_len = mask_seq_len = 0 + bs = None + device = None + if box_embeddings is not None: + bs = box_embeddings.shape[1] + box_seq_len = box_embeddings.shape[0] + device = box_embeddings.device + + if point_embeddings is not None: + point_seq_len = point_embeddings.shape[0] + if bs is not None: + assert ( + bs == point_embeddings.shape[1] + ), f"Batch size mismatch between box and point embeddings. Got {bs} and {point_embeddings.shape[1]}." + else: + bs = point_embeddings.shape[1] + if device is not None: + assert ( + device == point_embeddings.device + ), "Device mismatch between box and point embeddings" + else: + device = point_embeddings.device + + if mask_embeddings is not None: + mask_seq_len = mask_embeddings.shape[0] + if bs is not None: + assert ( + bs == mask_embeddings.shape[1] + ), f"Batch size mismatch between box/point and mask embedding. Got {bs} and {mask_embeddings.shape[1]}" + else: + bs = mask_embeddings.shape[1] + if device is not None: + assert ( + device == mask_embeddings.device + ), "Device mismatch between box/point and mask embeddings." + else: + device = mask_embeddings.device + + return box_seq_len, point_seq_len, mask_seq_len, bs, device + + def _init_box(self, box_embeddings, box_labels, box_mask, box_seq_len, bs, device): + if box_embeddings is None: + box_embeddings = torch.zeros(box_seq_len, bs, 4, device=device) + if box_labels is None: + box_labels = torch.ones(box_seq_len, bs, device=device, dtype=torch.long) + if box_mask is None: + box_mask = torch.zeros(bs, box_seq_len, device=device, dtype=torch.bool) + return box_embeddings, box_labels, box_mask + + def _init_point( + self, point_embeddings, point_labels, point_mask, point_seq_len, bs, device + ): + """ + Identical to _init_box. Except that C=2 for points (vs. 4 for boxes). + """ + if point_embeddings is None: + point_embeddings = torch.zeros(point_seq_len, bs, 2, device=device) + if point_labels is None: + point_labels = torch.ones( + point_seq_len, bs, device=device, dtype=torch.long + ) + if point_mask is None: + point_mask = torch.zeros(bs, point_seq_len, device=device, dtype=torch.bool) + return point_embeddings, point_labels, point_mask + + def _init_mask( + self, mask_embeddings, mask_labels, mask_mask, mask_seq_len, bs, device + ): + # NOTE: Mask embeddings can be of arbitrary resolution, so we don't initialize it here. + # In case we append new mask, we check that its resolution matches exisiting ones (if any). + # In case mask_embeddings is None, we should never encode it. + if mask_labels is None: + mask_labels = torch.ones(mask_seq_len, bs, device=device, dtype=torch.long) + if mask_mask is None: + mask_mask = torch.zeros(bs, mask_seq_len, device=device, dtype=torch.bool) + return mask_embeddings, mask_labels, mask_mask + + def append_boxes(self, boxes, labels, mask=None): + if self.box_embeddings is None: + self.box_embeddings = boxes + self.box_labels = labels + self.box_mask = mask + return + + bs = self.box_embeddings.shape[1] + assert boxes.shape[1] == labels.shape[1] == bs + assert list(boxes.shape[:2]) == list(labels.shape[:2]) + if mask is None: + mask = torch.zeros( + bs, boxes.shape[0], dtype=torch.bool, device=boxes.device + ) + + self.box_labels, _ = concat_padded_sequences( + self.box_labels.unsqueeze(-1), self.box_mask, labels.unsqueeze(-1), mask + ) + self.box_labels = self.box_labels.squeeze(-1) + self.box_embeddings, self.box_mask = concat_padded_sequences( + self.box_embeddings, self.box_mask, boxes, mask + ) + + def append_points(self, points, labels, mask=None): + if self.point_embeddings is None: + self.point_embeddings = points + self.point_labels = labels + self.point_mask = mask + return + + bs = self.point_embeddings.shape[1] + assert points.shape[1] == labels.shape[1] == bs + assert list(points.shape[:2]) == list(labels.shape[:2]) + if mask is None: + mask = torch.zeros( + bs, points.shape[0], dtype=torch.bool, device=points.device + ) + + self.point_labels, _ = concat_padded_sequences( + self.point_labels.unsqueeze(-1), self.point_mask, labels.unsqueeze(-1), mask + ) + self.point_labels = self.point_labels.squeeze(-1) + self.point_embeddings, self.point_mask = concat_padded_sequences( + self.point_embeddings, self.point_mask, points, mask + ) + + def append_masks(self, masks, labels=None, attn_mask=None): + if labels is not None: + assert list(masks.shape[:2]) == list(labels.shape[:2]) + if self.mask_embeddings is None: + self.mask_embeddings = masks + mask_seq_len, bs = masks.shape[:2] + if labels is None: + self.mask_labels = torch.ones( + mask_seq_len, bs, device=masks.device, dtype=torch.long + ) + else: + self.mask_labels = labels + if attn_mask is None: + self.mask_mask = torch.zeros( + bs, mask_seq_len, device=masks.device, dtype=torch.bool + ) + else: + self.mask_mask = attn_mask + else: + raise NotImplementedError("Only one mask per prompt is supported.") + + def clone(self): + return Prompt( + box_embeddings=( + None if self.box_embeddings is None else self.box_embeddings.clone() + ), + box_mask=None if self.box_mask is None else self.box_mask.clone(), + point_embeddings=( + None if self.point_embeddings is None else self.point_embeddings.clone() + ), + point_mask=None if self.point_mask is None else self.point_mask.clone(), + box_labels=None if self.box_labels is None else self.box_labels.clone(), + point_labels=( + None if self.point_labels is None else self.point_labels.clone() + ), + ) + + +class MaskEncoder(nn.Module): + """ + Base class for mask encoders. + """ + + def __init__( + self, + mask_downsampler: nn.Module, + position_encoding: nn.Module, + ): + super().__init__() + self.mask_downsampler = mask_downsampler + self.position_encoding = position_encoding + + def forward(self, masks, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: + masks = self.mask_downsampler(masks) + masks_pos = self.position_encoding(masks).to(masks.dtype) + + return masks, masks_pos + + +class FusedMaskEncoder(MaskEncoder): + """ + Identical to memory.SimpleMaskEncoder but follows the interface of geometry_encoders.MaskEncoder. + We also remove the `skip_mask_sigmoid` option (to be handled outside the MaskEncoder). + Fuses backbone image features with mask features. + """ + + def __init__( + self, + mask_downsampler: nn.Module, + position_encoding: nn.Module, + fuser: nn.Module, + in_dim: int = 256, + out_dim: int = 256, + ): + super().__init__(mask_downsampler, position_encoding) + self.fuser = fuser + self.out_proj = nn.Identity() + if out_dim != in_dim: + self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1) + self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1) + + @override + def forward( + self, + masks: torch.Tensor, + pix_feat: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + masks = self.mask_downsampler(masks) + + ## Fuse pix_feats and downsampled masks + # in case the visual features are on CPU, cast them to CUDA + pix_feat = pix_feat.to(masks.device) + + x = self.pix_feat_proj(pix_feat) + x = x + masks + x = self.fuser(x) + x = self.out_proj(x) + + pos = self.position_encoding(x).to(x.dtype) + + return x, pos + + +class SequenceGeometryEncoder(nn.Module): + """ + This a fully fledged encoder for geometric prompts. + It assumes boxes are passed in the "normalized CxCyWH" format, and points in normalized xy + This allows flexibility in how to encode the features (eg do pooling) + + Points and boxes can be encoded with any of the three possibilities: + - direct projection: we just compute a linear from coordinate space to d_model + - pooling: pool features from the backbone in the requested location. + For boxes, it's a roi align + For points it's a grid sample + - pos encoder: Take the position encoding of the point or box center + + These three options are mutually compatible. If several are selected, we'll take a simple addition + + As an alternative, we offer the possibility to encode points only. + In that case, the boxes are converted to two points for the top left and bottom right corners (with appropriate labels) + + On top of these encodings, we offer the possibility to further encode the prompt sequence with a transformer. + """ + + def __init__( + self, + encode_boxes_as_points: bool, + points_direct_project: bool, + points_pool: bool, + points_pos_enc: bool, + boxes_direct_project: bool, + boxes_pool: bool, + boxes_pos_enc: bool, + d_model: int, + pos_enc, + num_layers: int, + layer: nn.Module, + roi_size: int = 7, # for boxes pool + add_cls: bool = True, + add_post_encode_proj: bool = True, + mask_encoder: MaskEncoder = None, + add_mask_label: bool = False, + use_act_ckpt: bool = False, + ): + super().__init__() + + self.d_model = d_model + self.pos_enc = pos_enc + self.encode_boxes_as_points = encode_boxes_as_points + self.roi_size = roi_size + # There usually are two labels: positive and negatives. + # If we encode boxes as points, we have 3 types of points: regular, top left, bottom right + # These 3 types can be positives or negatives, hence 2*3 = 6 labels + num_labels = 6 if self.encode_boxes_as_points else 2 + self.label_embed = torch.nn.Embedding(num_labels, self.d_model) + + # This is a cls token, can be used for pooling if need be. + # It also ensures that the encoded sequences are always non-empty + self.cls_embed = None + if add_cls: + self.cls_embed = torch.nn.Embedding(1, self.d_model) + + assert ( + points_direct_project or points_pos_enc or points_pool + ), "Error: need at least one way to encode points" + assert ( + encode_boxes_as_points + or boxes_direct_project + or boxes_pos_enc + or boxes_pool + ), "Error: need at least one way to encode boxes" + + self.points_direct_project = None + if points_direct_project: + self.points_direct_project = nn.Linear(2, self.d_model) + self.points_pool_project = None + if points_pool: + self.points_pool_project = nn.Linear(self.d_model, self.d_model) + self.points_pos_enc_project = None + if points_pos_enc: + self.points_pos_enc_project = nn.Linear(self.d_model, self.d_model) + + self.boxes_direct_project = None + self.boxes_pool_project = None + self.boxes_pos_enc_project = None + if not encode_boxes_as_points: + if boxes_direct_project: + self.boxes_direct_project = nn.Linear(4, self.d_model) + if boxes_pool: + self.boxes_pool_project = nn.Conv2d( + self.d_model, self.d_model, self.roi_size + ) + if boxes_pos_enc: + self.boxes_pos_enc_project = nn.Linear(self.d_model + 2, self.d_model) + + self.final_proj = None + if add_post_encode_proj: + self.final_proj = nn.Linear(self.d_model, self.d_model) + self.norm = nn.LayerNorm(self.d_model) + + self.img_pre_norm = nn.Identity() + if self.points_pool_project is not None or self.boxes_pool_project is not None: + self.img_pre_norm = nn.LayerNorm(self.d_model) + + self.encode = None + if num_layers > 0: + assert ( + add_cls + ), "It's currently highly recommended to add a CLS when using a transformer" + self.encode = get_clones(layer, num_layers) + self.encode_norm = nn.LayerNorm(self.d_model) + + if mask_encoder is not None: + assert isinstance( + mask_encoder, MaskEncoder + ), f"Expected mask_encoder of type MaskEncoder. Got {type(mask_encoder)}." + if add_mask_label: + self.mask_label_embed = torch.nn.Embedding(2, self.d_model) + self.add_mask_label = add_mask_label + self.mask_encoder = mask_encoder + self.use_act_ckpt = use_act_ckpt + + def _encode_points(self, points, points_mask, points_labels, img_feats): + points_embed = None + n_points, bs = points.shape[:2] + + if self.points_direct_project is not None: + proj = self.points_direct_project(points) + assert points_embed is None + points_embed = proj + + if self.points_pool_project is not None: + # points are [Num_points, bs, 2], normalized in [0, 1] + # the grid needs to be [Bs, H_out, W_out, 2] normalized in [-1,1] + # Will take H_out = num_points, w_out = 1 + grid = points.transpose(0, 1).unsqueeze(2) + # re normalize to [-1, 1] + grid = (grid * 2) - 1 + sampled = torch.nn.functional.grid_sample( + img_feats, grid, align_corners=False + ) + assert list(sampled.shape) == [bs, self.d_model, n_points, 1] + sampled = sampled.squeeze(-1).permute(2, 0, 1) + proj = self.points_pool_project(sampled) + if points_embed is None: + points_embed = proj + else: + points_embed = points_embed + proj + + if self.points_pos_enc_project is not None: + x, y = points.unbind(-1) + enc_x, enc_y = self.pos_enc._encode_xy(x.flatten(), y.flatten()) + enc_x = enc_x.view(n_points, bs, enc_x.shape[-1]) + enc_y = enc_y.view(n_points, bs, enc_y.shape[-1]) + enc = torch.cat([enc_x, enc_y], -1) + + proj = self.points_pos_enc_project(enc) + if points_embed is None: + points_embed = proj + else: + points_embed = points_embed + proj + + type_embed = self.label_embed(points_labels.long()) + return type_embed + points_embed, points_mask + + def _encode_boxes(self, boxes, boxes_mask, boxes_labels, img_feats): + boxes_embed = None + n_boxes, bs = boxes.shape[:2] + + if self.boxes_direct_project is not None: + proj = self.boxes_direct_project(boxes) + assert boxes_embed is None + boxes_embed = proj + + if self.boxes_pool_project is not None: + H, W = img_feats.shape[-2:] + + # boxes are [Num_boxes, bs, 4], normalized in [0, 1] + # We need to denormalize, and convert to [x, y, x, y] + boxes_xyxy = box_cxcywh_to_xyxy(boxes) + scale = torch.tensor([W, H, W, H], dtype=boxes_xyxy.dtype) + scale = scale.pin_memory().to(device=boxes_xyxy.device, non_blocking=True) + scale = scale.view(1, 1, 4) + boxes_xyxy = boxes_xyxy * scale + sampled = torchvision.ops.roi_align( + img_feats, boxes_xyxy.float().transpose(0, 1).unbind(0), self.roi_size + ) + assert list(sampled.shape) == [ + bs * n_boxes, + self.d_model, + self.roi_size, + self.roi_size, + ] + proj = self.boxes_pool_project(sampled) + proj = proj.view(bs, n_boxes, self.d_model).transpose(0, 1) + if boxes_embed is None: + boxes_embed = proj + else: + boxes_embed = boxes_embed + proj + + if self.boxes_pos_enc_project is not None: + cx, cy, w, h = boxes.unbind(-1) + enc = self.pos_enc.encode_boxes( + cx.flatten(), cy.flatten(), w.flatten(), h.flatten() + ) + enc = enc.view(boxes.shape[0], boxes.shape[1], enc.shape[-1]) + + proj = self.boxes_pos_enc_project(enc) + if boxes_embed is None: + boxes_embed = proj + else: + boxes_embed = boxes_embed + proj + + type_embed = self.label_embed(boxes_labels.long()) + return type_embed + boxes_embed, boxes_mask + + def _encode_masks( + self, + masks: torch.Tensor, + attn_mask: torch.Tensor, + mask_labels: torch.Tensor, + img_feats: torch.Tensor = None, + ): + n_masks, bs = masks.shape[:2] + assert ( + n_masks == 1 + ), "We assume one mask per prompt for now. Code should still be functional if this assertion is removed." + assert ( + list(attn_mask.shape) + == [ + bs, + n_masks, + ] + ), f"Expected attn_mask to be of shape {bs}x{n_masks}. Got {list(attn_mask.shape)}." + masks, pos = self.mask_encoder( + masks=masks.flatten(0, 1).float(), + pix_feat=img_feats, + ) + H, W = masks.shape[-2:] + n_tokens_per_mask = H * W + # NOTE: We directly add pos enc here as we usually don't keep track of pos encoding for the concatenated prompt (text, other geometric prompts). Might need to do some refactoring for more flexibility. + masks = masks + pos + masks = masks.view(n_masks, bs, *masks.shape[1:]).flatten( + -2 + ) # n_masks x bs x C x H*W + masks = masks.permute(0, 3, 1, 2).flatten(0, 1) # n_masks * H*W x bs x C + attn_mask = attn_mask.repeat_interleave(n_tokens_per_mask, dim=1) + if self.add_mask_label: + masks = masks + self.mask_label_embed(mask_labels.long()) + return masks, attn_mask + + def forward(self, geo_prompt: Prompt, img_feats, img_sizes, img_pos_embeds=None): + points = geo_prompt.point_embeddings + points_mask = geo_prompt.point_mask + points_labels = geo_prompt.point_labels + boxes = geo_prompt.box_embeddings + boxes_mask = geo_prompt.box_mask + boxes_labels = geo_prompt.box_labels + masks = geo_prompt.mask_embeddings + masks_mask = geo_prompt.mask_mask + masks_labels = geo_prompt.mask_labels + seq_first_img_feats = img_feats[-1] # [H*W, B, C] + seq_first_img_pos_embeds = ( + img_pos_embeds[-1] + if img_pos_embeds is not None + else torch.zeros_like(seq_first_img_feats) + ) + + if self.points_pool_project or self.boxes_pool_project: + assert len(img_feats) == len(img_sizes) + cur_img_feat = img_feats[-1] + cur_img_feat = self.img_pre_norm(cur_img_feat) + H, W = img_sizes[-1] + assert cur_img_feat.shape[0] == H * W + N, C = cur_img_feat.shape[-2:] + # Put back in NxCxHxW + cur_img_feat = cur_img_feat.permute(1, 2, 0) + cur_img_feat = cur_img_feat.view(N, C, H, W) + img_feats = cur_img_feat # [1, 256, 72, 72] + + if self.encode_boxes_as_points: + assert boxes is not None + assert geo_prompt.box_mask is not None + assert geo_prompt.box_labels is not None + assert boxes.shape[-1] == 4 + + boxes_xyxy = box_cxcywh_to_xyxy(boxes) + top_left, bottom_right = boxes_xyxy.split(split_size=2, dim=-1) + + labels_tl = geo_prompt.box_labels + 2 + labels_br = geo_prompt.box_labels + 4 + + # Append to the existing points + points, _ = concat_padded_sequences( + points, points_mask, top_left, boxes_mask + ) + points_labels, points_mask = concat_padded_sequences( + points_labels.unsqueeze(-1), + points_mask, + labels_tl.unsqueeze(-1), + boxes_mask, + ) + points_labels = points_labels.squeeze(-1) + + points, _ = concat_padded_sequences( + points, points_mask, bottom_right, boxes_mask + ) + points_labels, points_mask = concat_padded_sequences( + points_labels.unsqueeze(-1), + points_mask, + labels_br.unsqueeze(-1), + boxes_mask, + ) + points_labels = points_labels.squeeze(-1) + + final_embeds, final_mask = self._encode_points( + points=points, + points_mask=points_mask, + points_labels=points_labels, + img_feats=img_feats, + ) + + if not self.encode_boxes_as_points: + boxes_embeds, boxes_mask = self._encode_boxes( + boxes=boxes, + boxes_mask=boxes_mask, + boxes_labels=boxes_labels, + img_feats=img_feats, + ) + + final_embeds, final_mask = concat_padded_sequences( + final_embeds, final_mask, boxes_embeds, boxes_mask + ) + + if masks is not None and self.mask_encoder is not None: + masks_embed, masks_mask = self._encode_masks( + masks=masks, + attn_mask=masks_mask, + mask_labels=masks_labels, + img_feats=img_feats, + ) + if points.size(0) == boxes.size(0) == 0: + return masks_embed, masks_mask + bs = final_embeds.shape[1] + assert final_mask.shape[0] == bs + if self.cls_embed is not None: # self.cls_embed: Embedding(1, 256) + cls = self.cls_embed.weight.view(1, 1, self.d_model).repeat(1, bs, 1) # [1, 1, 256] + cls_mask = torch.zeros( + bs, 1, dtype=final_mask.dtype, device=final_mask.device + ) + final_embeds, final_mask = concat_padded_sequences( + final_embeds, final_mask, cls, cls_mask + ) + + if self.final_proj is not None: + final_embeds = self.norm(self.final_proj(final_embeds)) # [1, 1, 256] + + if self.encode is not None: + for lay in self.encode: + final_embeds = activation_ckpt_wrapper(lay)( + tgt=final_embeds, + memory=seq_first_img_feats, + tgt_key_padding_mask=final_mask, + pos=seq_first_img_pos_embeds, + act_ckpt_enable=self.training and self.use_act_ckpt, + ) + final_embeds = self.encode_norm(final_embeds) + # Finally, concat mask embeddings if any + if masks is not None and self.mask_encoder is not None: + final_embeds, final_mask = concat_padded_sequences( + final_embeds, final_mask, masks_embed, masks_mask + ) + return final_embeds, final_mask + # final_embeds: [1, 1, 256] + # final_mask: [False] diff --git a/src/nn/segearth_ov3/sam3/model/io_utils.py b/src/nn/segearth_ov3/sam3/model/io_utils.py new file mode 100644 index 0000000..0a22584 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/io_utils.py @@ -0,0 +1,709 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import contextlib +import os +import queue +import re +import time +from threading import Condition, get_ident, Lock, Thread + +import numpy as np +import torch +import torch.nn.functional as F +import torchvision.transforms.functional as TF + +from PIL import Image + +from sam3.logger import get_logger +from tqdm import tqdm + +logger = get_logger(__name__) + +IS_MAIN_PROCESS = os.getenv("IS_MAIN_PROCESS", "1") == "1" +RANK = int(os.getenv("RANK", "0")) + +IMAGE_EXTS = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"] +VIDEO_EXTS = [".mp4", ".mov", ".avi", ".mkv", ".webm"] + + +def load_resource_as_video_frames( + resource_path, + image_size, + offload_video_to_cpu, + img_mean=(0.5, 0.5, 0.5), + img_std=(0.5, 0.5, 0.5), + async_loading_frames=False, + video_loader_type="cv2", +): + """ + Load video frames from either a video or an image (as a single-frame video). + Alternatively, if input is a list of PIL images, convert its format + """ + if isinstance(resource_path, list): + img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None] + assert all(isinstance(img_pil, Image.Image) for img_pil in resource_path) + assert len(resource_path) is not None + orig_height, orig_width = resource_path[0].size + orig_height, orig_width = ( + orig_width, + orig_height, + ) # For some reason, this method returns these swapped + images = [] + for img_pil in resource_path: + img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size))) + assert img_np.dtype == np.uint8, "np.uint8 is expected for JPEG images" + img_np = img_np / 255.0 + img = torch.from_numpy(img_np).permute(2, 0, 1) + # float16 precision should be sufficient for image tensor storage + img = img.to(dtype=torch.float16) + # normalize by mean and std + img -= img_mean + img /= img_std + images.append(img) + images = torch.stack(images) + if not offload_video_to_cpu: + images = images.cuda() + return images, orig_height, orig_width + + is_image = ( + isinstance(resource_path, str) + and os.path.splitext(resource_path)[-1].lower() in IMAGE_EXTS + ) + if is_image: + return load_image_as_single_frame_video( + image_path=resource_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + ) + else: + return load_video_frames( + video_path=resource_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + async_loading_frames=async_loading_frames, + video_loader_type=video_loader_type, + ) + + +def load_image_as_single_frame_video( + image_path, + image_size, + offload_video_to_cpu, + img_mean=(0.5, 0.5, 0.5), + img_std=(0.5, 0.5, 0.5), +): + """Load an image as a single-frame video.""" + images, image_height, image_width = _load_img_as_tensor(image_path, image_size) + images = images.unsqueeze(0).half() + + img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None] + if not offload_video_to_cpu: + images = images.cuda() + img_mean = img_mean.cuda() + img_std = img_std.cuda() + # normalize by mean and std + images -= img_mean + images /= img_std + return images, image_height, image_width + + +def load_video_frames( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.5, 0.5, 0.5), + img_std=(0.5, 0.5, 0.5), + async_loading_frames=False, + video_loader_type="cv2", +): + """ + Load the video frames from video_path. The frames are resized to image_size as in + the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo. + """ + assert isinstance(video_path, str) + if video_path.startswith(" where N is an integer + match = re.match(r"", video_path) + num_frames = int(match.group(1)) if match else 60 + return load_dummy_video(image_size, offload_video_to_cpu, num_frames=num_frames) + elif os.path.isdir(video_path): + return load_video_frames_from_image_folder( + image_folder=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + async_loading_frames=async_loading_frames, + ) + elif os.path.splitext(video_path)[-1].lower() in VIDEO_EXTS: + return load_video_frames_from_video_file( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + async_loading_frames=async_loading_frames, + video_loader_type=video_loader_type, + ) + else: + raise NotImplementedError("Only video files and image folders are supported") + + +def load_video_frames_from_image_folder( + image_folder, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + async_loading_frames, +): + """ + Load the video frames from a directory of image files ("." format) + """ + frame_names = [ + p + for p in os.listdir(image_folder) + if os.path.splitext(p)[-1].lower() in IMAGE_EXTS + ] + try: + frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) + except ValueError: + # fallback to lexicographic sort if the format is not "." + logger.warning( + f'frame names are not in "." format: {frame_names[:5]=}, ' + f"falling back to lexicographic sort." + ) + frame_names.sort() + num_frames = len(frame_names) + if num_frames == 0: + raise RuntimeError(f"no images found in {image_folder}") + img_paths = [os.path.join(image_folder, frame_name) for frame_name in frame_names] + img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None] + + if async_loading_frames: + lazy_images = AsyncImageFrameLoader( + img_paths, image_size, offload_video_to_cpu, img_mean, img_std + ) + return lazy_images, lazy_images.video_height, lazy_images.video_width + + # float16 precision should be sufficient for image tensor storage + images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float16) + video_height, video_width = None, None + for n, img_path in enumerate( + tqdm(img_paths, desc=f"frame loading (image folder) [rank={RANK}]") + ): + images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size) + if not offload_video_to_cpu: + images = images.cuda() + img_mean = img_mean.cuda() + img_std = img_std.cuda() + # normalize by mean and std + images -= img_mean + images /= img_std + return images, video_height, video_width + + +def load_video_frames_from_video_file( + video_path, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + async_loading_frames, + gpu_acceleration=False, + gpu_device=None, + video_loader_type="cv2", +): + """Load the video frames from a video file.""" + if video_loader_type == "cv2": + return load_video_frames_from_video_file_using_cv2( + video_path=video_path, + image_size=image_size, + img_mean=img_mean, + img_std=img_std, + offload_video_to_cpu=offload_video_to_cpu, + ) + elif video_loader_type == "torchcodec": + logger.info("Using torchcodec to load video file") + lazy_images = AsyncVideoFileLoaderWithTorchCodec( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + gpu_acceleration=gpu_acceleration, + gpu_device=gpu_device, + ) + # The `AsyncVideoFileLoaderWithTorchCodec` class always loads the videos asynchronously, + # so we just wait for its loading thread to finish if async_loading_frames=False. + if not async_loading_frames: + async_thread = lazy_images.thread + if async_thread is not None: + async_thread.join() + return lazy_images, lazy_images.video_height, lazy_images.video_width + else: + raise RuntimeError("video_loader_type must be either 'cv2' or 'torchcodec'") + + +def load_video_frames_from_video_file_using_cv2( + video_path: str, + image_size: int, + img_mean: tuple = (0.5, 0.5, 0.5), + img_std: tuple = (0.5, 0.5, 0.5), + offload_video_to_cpu: bool = False, +) -> torch.Tensor: + """ + Load video from path, convert to normalized tensor with specified preprocessing + + Args: + video_path: Path to video file + image_size: Target size for square frames (height and width) + img_mean: Normalization mean (RGB) + img_std: Normalization standard deviation (RGB) + + Returns: + torch.Tensor: Preprocessed video tensor in shape (T, C, H, W) with float16 dtype + """ + import cv2 # delay OpenCV import to avoid unnecessary dependency + + # Initialize video capture + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise ValueError(f"Could not open video: {video_path}") + + original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + num_frames = num_frames if num_frames > 0 else None + + frames = [] + pbar = tqdm(desc=f"frame loading (OpenCV) [rank={RANK}]", total=num_frames) + while True: + ret, frame = cap.read() + if not ret: + break + + # Convert BGR to RGB and resize + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame_resized = cv2.resize( + frame_rgb, (image_size, image_size), interpolation=cv2.INTER_CUBIC + ) + frames.append(frame_resized) + pbar.update(1) + cap.release() + pbar.close() + + # Convert to tensor + frames_np = np.stack(frames, axis=0).astype(np.float32) # (T, H, W, C) + video_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2) # (T, C, H, W) + + img_mean = torch.tensor(img_mean, dtype=torch.float16).view(1, 3, 1, 1) + img_std = torch.tensor(img_std, dtype=torch.float16).view(1, 3, 1, 1) + if not offload_video_to_cpu: + video_tensor = video_tensor.cuda() + img_mean = img_mean.cuda() + img_std = img_std.cuda() + # normalize by mean and std + video_tensor -= img_mean + video_tensor /= img_std + return video_tensor, original_height, original_width + + +def load_dummy_video(image_size, offload_video_to_cpu, num_frames=60): + """ + Load a dummy video with random frames for testing and compilation warmup purposes. + """ + video_height, video_width = 480, 640 # dummy original video sizes + images = torch.randn(num_frames, 3, image_size, image_size, dtype=torch.float16) + if not offload_video_to_cpu: + images = images.cuda() + return images, video_height, video_width + + +def _load_img_as_tensor(img_path, image_size): + """Load and resize an image and convert it into a PyTorch tensor.""" + img = Image.open(img_path).convert("RGB") + orig_width, orig_height = img.width, img.height + img = TF.resize(img, size=(image_size, image_size)) + img = TF.to_tensor(img) + return img, orig_height, orig_width + + +class AsyncImageFrameLoader: + """ + A list of video frames to be load asynchronously without blocking session start. + """ + + def __init__(self, img_paths, image_size, offload_video_to_cpu, img_mean, img_std): + self.img_paths = img_paths + self.image_size = image_size + self.offload_video_to_cpu = offload_video_to_cpu + self.img_mean = img_mean + self.img_std = img_std + # items in `self._images` will be loaded asynchronously + self.images = [None] * len(img_paths) + # catch and raise any exceptions in the async loading thread + self.exception = None + # video_height and video_width be filled when loading the first image + self.video_height = None + self.video_width = None + + # load the first frame to fill video_height and video_width and also + # to cache it (since it's most likely where the user will click) + self.__getitem__(0) + + # load the rest of frames asynchronously without blocking the session start + def _load_frames(): + try: + for n in tqdm( + range(len(self.images)), + desc=f"frame loading (image folder) [rank={RANK}]", + ): + self.__getitem__(n) + except Exception as e: + self.exception = e + + self.thread = Thread(target=_load_frames, daemon=True) + self.thread.start() + + def __getitem__(self, index): + if self.exception is not None: + raise RuntimeError("Failure in frame loading thread") from self.exception + + img = self.images[index] + if img is not None: + return img + + img, video_height, video_width = _load_img_as_tensor( + self.img_paths[index], self.image_size + ) + self.video_height = video_height + self.video_width = video_width + # float16 precision should be sufficient for image tensor storage + img = img.to(dtype=torch.float16) + # normalize by mean and std + img -= self.img_mean + img /= self.img_std + if not self.offload_video_to_cpu: + img = img.cuda() + self.images[index] = img + return img + + def __len__(self): + return len(self.images) + + +class TorchCodecDecoder: + """ + A wrapper to support GPU device and num_threads in TorchCodec decoder, + which are not supported by `torchcodec.decoders.SimpleVideoDecoder` yet. + """ + + def __init__(self, source, dimension_order="NCHW", device="cpu", num_threads=1): + from torchcodec import _core as core + + self._source = source # hold a reference to the source to prevent it from GC + if isinstance(source, str): + self._decoder = core.create_from_file(source, "exact") + elif isinstance(source, bytes): + self._decoder = core.create_from_bytes(source, "exact") + else: + raise TypeError(f"Unknown source type: {type(source)}.") + assert dimension_order in ("NCHW", "NHWC") + + device_string = str(device) + core.scan_all_streams_to_update_metadata(self._decoder) + core.add_video_stream( + self._decoder, + dimension_order=dimension_order, + device=device_string, + num_threads=(1 if "cuda" in device_string else num_threads), + ) + video_metadata = core.get_container_metadata(self._decoder) + best_stream_index = video_metadata.best_video_stream_index + assert best_stream_index is not None + self.metadata = video_metadata.streams[best_stream_index] + assert self.metadata.num_frames_from_content is not None + self._num_frames = self.metadata.num_frames_from_content + + def __len__(self) -> int: + return self._num_frames + + def __getitem__(self, key: int): + from torchcodec import _core as core + + if key < 0: + key += self._num_frames + if key >= self._num_frames or key < 0: + raise IndexError( + f"Index {key} is out of bounds; length is {self._num_frames}" + ) + frame_data, *_ = core.get_frame_at_index( + self._decoder, + frame_index=key, + ) + return frame_data + + +class FIFOLock: + """A lock that ensures FIFO ordering of lock acquisitions.""" + + def __init__(self): + self._lock = Lock() + self._waiters = queue.Queue() + self._condition = Condition() + + def acquire(self): + ident = get_ident() + with self._condition: + self._waiters.put(ident) + while self._waiters.queue[0] != ident or not self._lock.acquire( + blocking=False + ): + self._condition.wait() + # got the lock and it's our turn + + def release(self): + with self._condition: + self._lock.release() + self._waiters.get() + self._condition.notify_all() + + def __enter__(self): + self.acquire() + + def __exit__(self, t, v, tb): + self.release() + + +class AsyncVideoFileLoaderWithTorchCodec: + """ + Loading frames from video files asynchronously without blocking session start. + + Unlike `AsyncVideoFileLoader`, this class uses PyTorch's offical TorchCodec library + for video decoding, which is more efficient and supports more video formats. + """ + + def __init__( + self, + video_path, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + gpu_acceleration=True, + gpu_device=None, + use_rand_seek_in_loading=False, + ): + # Check and possibly infer the output device (and also get its GPU id when applicable) + assert gpu_device is None or gpu_device.type == "cuda" + gpu_id = ( + gpu_device.index + if gpu_device is not None and gpu_device.index is not None + else torch.cuda.current_device() + ) + if offload_video_to_cpu: + out_device = torch.device("cpu") + else: + out_device = torch.device("cuda") if gpu_device is None else gpu_device + self.out_device = out_device + self.gpu_acceleration = gpu_acceleration + self.gpu_id = gpu_id + self.image_size = image_size + self.offload_video_to_cpu = offload_video_to_cpu + if not isinstance(img_mean, torch.Tensor): + img_mean = torch.tensor(img_mean, dtype=torch.float16)[:, None, None] + self.img_mean = img_mean + if not isinstance(img_std, torch.Tensor): + img_std = torch.tensor(img_std, dtype=torch.float16)[:, None, None] + self.img_std = img_std + + if gpu_acceleration: + self.img_mean = self.img_mean.to(f"cuda:{self.gpu_id}") + self.img_std = self.img_std.to(f"cuda:{self.gpu_id}") + decoder_option = {"device": f"cuda:{self.gpu_id}"} + else: + self.img_mean = self.img_mean.cpu() + self.img_std = self.img_std.cpu() + decoder_option = {"num_threads": 1} # use a single thread to save memory + + self.rank = int(os.environ.get("RANK", "0")) + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.async_reader = TorchCodecDecoder(video_path, **decoder_option) + + # `num_frames_from_content` is the true number of frames in the video content + # from the scan operation (rather than from the metadata, which could be wrong) + self.num_frames = self.async_reader.metadata.num_frames_from_content + self.video_height = self.async_reader.metadata.height + self.video_width = self.async_reader.metadata.width + + # items in `self._images` will be loaded asynchronously + self.images_loaded = [False] * self.num_frames + self.images = torch.zeros( + self.num_frames, + 3, + self.image_size, + self.image_size, + dtype=torch.float16, + device=self.out_device, + ) + # catch and raise any exceptions in the async loading thread + self.exception = None + self.use_rand_seek_in_loading = use_rand_seek_in_loading + self.rand_seek_idx_queue = queue.Queue() + # use a lock to avoid race condition between concurrent access to torchcodec + # libs (which are not thread-safe); the lock is replaced with a nullcontext + # when the video is fully loaded + self.torchcodec_access_lock = FIFOLock() + self._start_video_loading() + + def _load_one_frame(self, idx): + frame_resized = self._transform_frame(self.async_reader[idx]) + return frame_resized + + @torch.inference_mode() + def _start_video_loading(self): + desc = f"frame loading (TorchCodec w/ {'GPU' if self.gpu_acceleration else 'CPU'}) [rank={RANK}]" + pbar = tqdm(desc=desc, total=self.num_frames) + self.num_loaded_frames = 0 + # load the first frame synchronously to cache it before the session is opened + idx = self.num_loaded_frames + self.images[idx] = self._load_one_frame(idx) + self.images_loaded[idx] = True + self.num_loaded_frames += 1 + pbar.update(n=1) + self.all_frames_loaded = self.num_loaded_frames == self.num_frames + + # load the frames asynchronously without blocking the session start + def _load_frames(): + finished = self.all_frames_loaded + chunk_size = 16 + while not finished: + # asynchronously load `chunk_size` frames each time we acquire the lock + with self.torchcodec_access_lock, torch.inference_mode(): + for _ in range(chunk_size): + try: + idx = self.num_loaded_frames + self.images[idx] = self._load_one_frame(idx) + self.images_loaded[idx] = True + self.num_loaded_frames += 1 + pbar.update(n=1) + if self.num_loaded_frames >= self.num_frames: + finished = True + break + except Exception as e: + self.exception = e + raise + + # also read the frame that is being randomly seeked to + while True: + try: + idx = self.rand_seek_idx_queue.get_nowait() + if not self.images_loaded[idx]: + self.images[idx] = self._load_one_frame(idx) + self.images_loaded[idx] = True + except queue.Empty: + break + except Exception as e: + self.exception = e + raise + + # finished -- check whether we have loaded the total number of frames + if self.num_loaded_frames != self.num_frames: + raise RuntimeError( + f"There are {self.num_frames} frames in the video, but only " + f"{self.num_loaded_frames} frames can be loaded successfully." + ) + else: + self.all_frames_loaded = True + pbar.close() + with self.torchcodec_access_lock: + import gc + + # all frames have been loaded, so we can release the readers and free their memory + # also remove pbar and thread (which shouldn't be a part of session saving) + reader = self.async_reader + if reader is not None: + reader._source = None + self.async_reader = None + self.pbar = None + self.thread = None + self.rand_seek_idx_queue = None + gc.collect() + # remove the lock (replace it with nullcontext) when the video is fully loaded + self.torchcodec_access_lock = contextlib.nullcontext() + + self.thread = Thread(target=_load_frames, daemon=True) + self.thread.start() + + def _transform_frame(self, frame): + frame = frame.clone() # make a copy to avoid modifying the original frame bytes + frame = frame.float() # convert to float32 before interpolation + frame_resized = F.interpolate( + frame[None, :], + size=(self.image_size, self.image_size), + mode="bicubic", + align_corners=False, + )[0] + # float16 precision should be sufficient for image tensor storage + frame_resized = frame_resized.half() # uint8 -> float16 + frame_resized /= 255 + frame_resized -= self.img_mean + frame_resized /= self.img_std + if self.offload_video_to_cpu: + frame_resized = frame_resized.cpu() + elif frame_resized.device != self.out_device: + frame_resized = frame_resized.to(device=self.out_device, non_blocking=True) + return frame_resized + + def __getitem__(self, index): + if self.exception is not None: + raise RuntimeError("Failure in frame loading thread") from self.exception + + max_tries = 1200 + for _ in range(max_tries): + # use a lock to avoid race condition between concurrent access to torchcodec + # libs (which are not thread-safe); the lock is replaced with a nullcontext + # when the video is fully loaded + with self.torchcodec_access_lock: + if self.images_loaded[index]: + return self.images[index] + + if self.use_rand_seek_in_loading: + # async loading hasn't reached this frame yet, so we load this frame individually + # (it will be loaded by in _load_frames thread and added to self.images[index]) + self.rand_seek_idx_queue.put(index) + + time.sleep(0.1) + + raise RuntimeError(f"Failed to load frame {index} after {max_tries} tries") + + def __len__(self): + return len(self.images) + + def __getstate__(self): + """ + Remove a few attributes during pickling, so that this async video loader can be + saved and loaded as a part of the model session. + """ + # wait for async video loading to finish before pickling + async_thread = self.thread + if async_thread is not None: + async_thread.join() + # release a few objects that cannot be pickled + reader = self.async_reader + if reader is not None: + reader._source = None + self.async_reader = None + self.pbar = None + self.thread = None + self.rand_seek_idx_queue = None + self.torchcodec_access_lock = contextlib.nullcontext() + return self.__dict__.copy() diff --git a/src/nn/segearth_ov3/sam3/model/maskformer_segmentation.py b/src/nn/segearth_ov3/sam3/model/maskformer_segmentation.py new file mode 100644 index 0000000..d04fb7f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/maskformer_segmentation.py @@ -0,0 +1,330 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math +from typing import Dict, List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint + +from .model_misc import MLP + + +class LinearPresenceHead(nn.Sequential): + def __init__(self, d_model): + # a hack to make `LinearPresenceHead` compatible with old checkpoints + super().__init__(nn.Identity(), nn.Identity(), nn.Linear(d_model, 1)) + + def forward(self, hs, prompt, prompt_mask): + return super().forward(hs) + + +class MaskPredictor(nn.Module): + def __init__(self, hidden_dim, mask_dim): + super().__init__() + self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3) + + def forward(self, obj_queries, pixel_embed): + if len(obj_queries.shape) == 3: + if pixel_embed.ndim == 3: + # batch size was omitted + mask_preds = torch.einsum( + "bqc,chw->bqhw", self.mask_embed(obj_queries), pixel_embed + ) + else: + mask_preds = torch.einsum( + "bqc,bchw->bqhw", self.mask_embed(obj_queries), pixel_embed + ) + else: + # Assumed to have aux masks + if pixel_embed.ndim == 3: + # batch size was omitted + mask_preds = torch.einsum( + "lbqc,chw->lbqhw", self.mask_embed(obj_queries), pixel_embed + ) + else: + mask_preds = torch.einsum( + "lbqc,bchw->lbqhw", self.mask_embed(obj_queries), pixel_embed + ) + + return mask_preds + + +class SegmentationHead(nn.Module): + def __init__( + self, + hidden_dim, + upsampling_stages, + use_encoder_inputs=False, + aux_masks=False, + no_dec=False, + pixel_decoder=None, + act_ckpt=False, + shared_conv=False, + compile_mode_pixel_decoder=None, + ): + super().__init__() + self.use_encoder_inputs = use_encoder_inputs + self.aux_masks = aux_masks + if pixel_decoder is not None: + self.pixel_decoder = pixel_decoder + else: + self.pixel_decoder = PixelDecoder( + hidden_dim, + upsampling_stages, + shared_conv=shared_conv, + compile_mode=compile_mode_pixel_decoder, + ) + self.no_dec = no_dec + if no_dec: + self.mask_predictor = nn.Conv2d( + hidden_dim, 1, kernel_size=3, stride=1, padding=1 + ) + else: + self.mask_predictor = MaskPredictor(hidden_dim, mask_dim=hidden_dim) + + self.act_ckpt = act_ckpt + + # used to update the output dictionary + self.instance_keys = ["pred_masks"] + + @property + def device(self): + self._device = getattr(self, "_device", None) or next(self.parameters()).device + return self._device + + def to(self, *args, **kwargs): + # clear cached _device in case the model is moved to a different device + self._device = None + return super().to(*args, **kwargs) + + def _embed_pixels( + self, + backbone_feats: List[torch.Tensor], + image_ids, + encoder_hidden_states, + ) -> torch.Tensor: + feature_device = backbone_feats[0].device # features could be on CPU + model_device = self.device + image_ids_ = image_ids.to(feature_device) + if self.use_encoder_inputs: + if backbone_feats[0].shape[0] > 1: + # For bs > 1, we construct the per query backbone features + backbone_visual_feats = [] + for feat in backbone_feats: + # Copy the img features per query (pixel decoder won't share img feats) + backbone_visual_feats.append(feat[image_ids_, ...].to(model_device)) + else: + # Bs=1, we rely on broadcasting for query-based processing + backbone_visual_feats = [bb_feat.clone() for bb_feat in backbone_feats] + # Extract visual embeddings + encoder_hidden_states = encoder_hidden_states.permute(1, 2, 0) + spatial_dim = math.prod(backbone_feats[-1].shape[-2:]) + encoder_visual_embed = encoder_hidden_states[..., :spatial_dim].reshape( + -1, *backbone_feats[-1].shape[1:] + ) + + backbone_visual_feats[-1] = encoder_visual_embed + if self.act_ckpt: + pixel_embed = checkpoint.checkpoint( + self.pixel_decoder, backbone_visual_feats, use_reentrant=False + ) + else: + pixel_embed = self.pixel_decoder(backbone_visual_feats) + else: + backbone_feats = [x.to(model_device) for x in backbone_feats] + pixel_embed = self.pixel_decoder(backbone_feats) + if pixel_embed.shape[0] == 1: + # For batch_size=1 training, we can avoid the indexing to save memory + pixel_embed = pixel_embed.squeeze(0) + else: + pixel_embed = pixel_embed[image_ids, ...] + return pixel_embed # [1, 256, 288, 288] + + def forward( + self, + backbone_feats: List[torch.Tensor], + obj_queries: torch.Tensor, + image_ids, + encoder_hidden_states: Optional[torch.Tensor] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if self.use_encoder_inputs: + assert encoder_hidden_states is not None + + pixel_embed = self._embed_pixels( + backbone_feats=backbone_feats, + image_ids=image_ids, + encoder_hidden_states=encoder_hidden_states, + ) + + if self.no_dec: + mask_pred = self.mask_predictor(pixel_embed) + elif self.aux_masks: + mask_pred = self.mask_predictor(obj_queries, pixel_embed) + else: + mask_pred = self.mask_predictor(obj_queries[-1], pixel_embed) + + return {"pred_masks": mask_pred} + + +class PixelDecoder(nn.Module): + def __init__( + self, + hidden_dim, + num_upsampling_stages, + interpolation_mode="nearest", + shared_conv=False, + compile_mode=None, + ): + super().__init__() + self.hidden_dim = hidden_dim + self.num_upsampling_stages = num_upsampling_stages + self.interpolation_mode = interpolation_mode + conv_layers = [] + norms = [] + num_convs = 1 if shared_conv else num_upsampling_stages + for _ in range(num_convs): + conv_layers.append(nn.Conv2d(self.hidden_dim, self.hidden_dim, 3, 1, 1)) + norms.append(nn.GroupNorm(8, self.hidden_dim)) + + self.conv_layers = nn.ModuleList(conv_layers) + self.norms = nn.ModuleList(norms) + self.shared_conv = shared_conv + self.out_dim = self.conv_layers[-1].out_channels + if compile_mode is not None: + self.forward = torch.compile( + self.forward, mode=compile_mode, dynamic=True, fullgraph=True + ) + # Needed to make checkpointing happy. But we don't know if the module is checkpointed, so we disable it by default. + torch._dynamo.config.optimize_ddp = False + + def forward(self, backbone_feats: List[torch.Tensor]): + # Assumes backbone features are already projected (C == hidden dim) + + prev_fpn = backbone_feats[-1] + fpn_feats = backbone_feats[:-1] + for layer_idx, bb_feat in enumerate(fpn_feats[::-1]): + curr_fpn = bb_feat + prev_fpn = curr_fpn + F.interpolate( + prev_fpn, size=curr_fpn.shape[-2:], mode=self.interpolation_mode + ) + if self.shared_conv: + # only one conv layer + layer_idx = 0 + prev_fpn = self.conv_layers[layer_idx](prev_fpn) + prev_fpn = F.relu(self.norms[layer_idx](prev_fpn)) + + return prev_fpn + + +class UniversalSegmentationHead(SegmentationHead): + """This module handles semantic+instance segmentation""" + + def __init__( + self, + hidden_dim, + upsampling_stages, + pixel_decoder, + aux_masks=False, + no_dec=False, + act_ckpt=False, + presence_head: bool = False, + dot_product_scorer=None, + cross_attend_prompt=None, + ): + super().__init__( + hidden_dim=hidden_dim, + upsampling_stages=upsampling_stages, + use_encoder_inputs=True, + aux_masks=aux_masks, + no_dec=no_dec, + pixel_decoder=pixel_decoder, + act_ckpt=act_ckpt, + ) + self.d_model = hidden_dim + + if dot_product_scorer is not None: + assert presence_head, "Specifying a dot product scorer without a presence head is likely a mistake" + + self.presence_head = None + if presence_head: + self.presence_head = ( + dot_product_scorer + if dot_product_scorer is not None + else LinearPresenceHead(self.d_model) + ) + + self.cross_attend_prompt = cross_attend_prompt + if self.cross_attend_prompt is not None: + self.cross_attn_norm = nn.LayerNorm(self.d_model) + + self.semantic_seg_head = nn.Conv2d(self.pixel_decoder.out_dim, 1, kernel_size=1) + self.instance_seg_head = nn.Conv2d( + self.pixel_decoder.out_dim, self.d_model, kernel_size=1 + ) + + def forward( + self, + backbone_feats: List[torch.Tensor], + obj_queries: torch.Tensor, + image_ids, + encoder_hidden_states: Optional[torch.Tensor] = None, + prompt: Optional[torch.Tensor] = None, + prompt_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> Dict[str, Optional[torch.Tensor]]: + assert encoder_hidden_states is not None + bs = encoder_hidden_states.shape[1] + + if self.cross_attend_prompt is not None: + tgt2 = self.cross_attn_norm(encoder_hidden_states) + tgt2 = self.cross_attend_prompt( + query=tgt2, + key=prompt, + value=prompt, + key_padding_mask=prompt_mask, + )[0] + # prompt: [33, 1, 256] + + encoder_hidden_states = tgt2 + encoder_hidden_states + # tgt2: [5184, 1, 256] + # encoder_hidden_states: [5184, 1, 256] + + presence_logit = None + if self.presence_head is not None: + pooled_enc = encoder_hidden_states.mean(0) + presence_logit = ( + self.presence_head( + pooled_enc.view(1, bs, 1, self.d_model), + prompt=prompt, + prompt_mask=prompt_mask, + ) + .squeeze(0) + .squeeze(1) + ) + + pixel_embed = self._embed_pixels( + backbone_feats=backbone_feats, + image_ids=image_ids, + encoder_hidden_states=encoder_hidden_states, + ) + + # pixel_embed: [1, 256, 288, 288] + + instance_embeds = self.instance_seg_head(pixel_embed) + # instance_embeds: [1, 256, 288, 288] + + if self.no_dec: + mask_pred = self.mask_predictor(instance_embeds) + elif self.aux_masks: + mask_pred = self.mask_predictor(obj_queries, instance_embeds) + else: + mask_pred = self.mask_predictor(obj_queries[-1], instance_embeds) # [1, 200, 288, 288] + + return { + "pred_masks": mask_pred, + "semantic_seg": self.semantic_seg_head(pixel_embed), + "presence_logit": presence_logit, + } diff --git a/src/nn/segearth_ov3/sam3/model/memory.py b/src/nn/segearth_ov3/sam3/model/memory.py new file mode 100644 index 0000000..bfde548 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/memory.py @@ -0,0 +1,201 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from timm.layers import DropPath +except ModuleNotFoundError: + # compatibility for older timm versions + from timm.models.layers import DropPath + +from .model_misc import get_clones, LayerNorm2d + + +class SimpleMaskDownSampler(nn.Module): + """ + Progressively downsample a mask by total_stride, each time by stride. + Note that LayerNorm is applied per *token*, like in ViT. + + With each downsample (by a factor stride**2), channel capacity increases by the same factor. + In the end, we linearly project to embed_dim channels. + """ + + def __init__( + self, + embed_dim=256, + kernel_size=4, + stride=4, + padding=0, + total_stride=16, + activation=nn.GELU, + # Option to interpolate the input mask first before downsampling using convs. In that case, the total_stride is assumed to be after interpolation. + # If set to input resolution or None, we don't interpolate. We default to None to be safe (for older configs or if not explicitly set) + interpol_size=None, + ): + super().__init__() + num_layers = int(math.log2(total_stride) // math.log2(stride)) + assert stride**num_layers == total_stride + self.encoder = nn.Sequential() + mask_in_chans, mask_out_chans = 1, 1 + for _ in range(num_layers): + mask_out_chans = mask_in_chans * (stride**2) + self.encoder.append( + nn.Conv2d( + mask_in_chans, + mask_out_chans, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + ) + self.encoder.append(LayerNorm2d(mask_out_chans)) + self.encoder.append(activation()) + mask_in_chans = mask_out_chans + + self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1)) + self.interpol_size = interpol_size + if self.interpol_size is not None: + assert isinstance( + self.interpol_size, (list, tuple) + ), f"Unsupported type {type(self.interpol_size)}. Should be a list or tuple." + self.interpol_size = list(interpol_size) + assert len(self.interpol_size) == 2 + + def forward(self, x: torch.Tensor): + if self.interpol_size is not None and self.interpol_size != list(x.shape[-2:]): + x = F.interpolate( + x.float(), + size=self.interpol_size, + align_corners=False, + mode="bilinear", + antialias=True, + ) + return self.encoder(x) + + +# Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt) +class CXBlock(nn.Module): + r"""ConvNeXt Block. There are two equivalent implementations: + (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) + (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back + We use (2) as we find it slightly faster in PyTorch + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + """ + + def __init__( + self, + dim, + kernel_size=7, + padding=3, + drop_path=0.0, + layer_scale_init_value=1e-6, + use_dwconv=True, + ): + super().__init__() + self.dwconv = nn.Conv2d( + dim, + dim, + kernel_size=kernel_size, + padding=padding, + groups=dim if use_dwconv else 1, + ) # depthwise conv + self.norm = LayerNorm2d(dim, eps=1e-6) + self.pwconv1 = nn.Linear( + dim, 4 * dim + ) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = ( + nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) + if layer_scale_init_value > 0 + else None + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = self.norm(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + + +class SimpleFuser(nn.Module): + def __init__(self, layer, num_layers, dim=None, input_projection=False): + super().__init__() + self.proj = nn.Identity() + self.layers = get_clones(layer, num_layers) + + if input_projection: + assert dim is not None + self.proj = nn.Conv2d(dim, dim, kernel_size=1) + + def forward(self, x): + # normally x: (N, C, H, W) + x = self.proj(x) + for layer in self.layers: + x = layer(x) + return x + + +class SimpleMaskEncoder(nn.Module): + def __init__( + self, + out_dim, + mask_downsampler, + fuser, + position_encoding, + in_dim=256, # in_dim of pix_feats + ): + super().__init__() + + self.mask_downsampler = mask_downsampler + + self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1) + self.fuser = fuser + self.position_encoding = position_encoding + self.out_proj = nn.Identity() + if out_dim != in_dim: + self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1) + + def forward( + self, + pix_feat: torch.Tensor, + masks: torch.Tensor, + skip_mask_sigmoid: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ## Process masks + # sigmoid, so that less domain shift from gt masks which are bool + if not skip_mask_sigmoid: + masks = F.sigmoid(masks) + masks = self.mask_downsampler(masks) + + ## Fuse pix_feats and downsampled masks + # in case the visual features are on CPU, cast them to CUDA + pix_feat = pix_feat.to(masks.device) + + x = self.pix_feat_proj(pix_feat) + x = x + masks + x = self.fuser(x) + x = self.out_proj(x) + + pos = self.position_encoding(x).to(x.dtype) + + return {"vision_features": x, "vision_pos_enc": [pos]} diff --git a/src/nn/segearth_ov3/sam3/model/model_misc.py b/src/nn/segearth_ov3/sam3/model/model_misc.py new file mode 100644 index 0000000..2cb44b3 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/model_misc.py @@ -0,0 +1,428 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Various utility models""" + +import copy +import math +import weakref +from collections.abc import Iterator +from contextlib import AbstractContextManager +from enum import auto, Enum +from typing import Dict, List, Optional, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn, Tensor +from typing_extensions import override + + +def inverse_sigmoid(x, eps=1e-3): + """ + The inverse function for sigmoid activation function. + Note: It might face numberical issues with fp16 small eps. + """ + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + + +class MultiheadAttentionWrapper(nn.MultiheadAttention): + def forward(self, *args, **kwargs): + kwargs["need_weights"] = False + return super().forward(*args, **kwargs) + + +class DotProductScoring(torch.nn.Module): + def __init__( + self, + d_model, + d_proj, + prompt_mlp=None, + clamp_logits=True, + clamp_max_val=12.0, + ): + super().__init__() + self.d_proj = d_proj + assert isinstance(prompt_mlp, torch.nn.Module) or prompt_mlp is None + self.prompt_mlp = prompt_mlp # an optional MLP projection for prompt + self.prompt_proj = torch.nn.Linear(d_model, d_proj) + self.hs_proj = torch.nn.Linear(d_model, d_proj) + self.scale = float(1.0 / np.sqrt(d_proj)) + self.clamp_logits = clamp_logits + if self.clamp_logits: + self.clamp_max_val = clamp_max_val + + def mean_pool_text(self, prompt, prompt_mask): + # is_valid has shape (seq, bs, 1), where 1 is valid and 0 is padding + is_valid = (~prompt_mask).float().permute(1, 0)[..., None] + # num_valid has shape (bs, 1) + num_valid = torch.clamp(torch.sum(is_valid, dim=0), min=1.0) + # mean pool over all the valid tokens -- pooled_prompt has shape (bs, proj_dim) + pooled_prompt = (prompt * is_valid).sum(dim=0) / num_valid + return pooled_prompt + + def forward(self, hs, prompt, prompt_mask): + # hs has shape (num_layer, bs, num_query, d_model) + # prompt has shape (seq, bs, d_model) + # prompt_mask has shape (bs, seq), where 1 is valid and 0 is padding + assert hs.dim() == 4 and prompt.dim() == 3 and prompt_mask.dim() == 2 + + # apply MLP on prompt if specified + if self.prompt_mlp is not None: + prompt = self.prompt_mlp(prompt) + + # first, get the mean-pooled version of the prompt + pooled_prompt = self.mean_pool_text(prompt, prompt_mask) + + # then, project pooled_prompt and hs to d_proj dimensions + proj_pooled_prompt = self.prompt_proj(pooled_prompt) # (bs, d_proj) + proj_hs = self.hs_proj(hs) # (num_layer, bs, num_query, d_proj) + + # finally, get dot-product scores of shape (num_layer, bs, num_query, 1) + scores = torch.matmul(proj_hs, proj_pooled_prompt.unsqueeze(-1)) + scores *= self.scale + + # clamp scores to a max value to avoid numerical issues in loss or matcher + if self.clamp_logits: + scores.clamp_(min=-self.clamp_max_val, max=self.clamp_max_val) + + return scores + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + +class LayerNorm2d(nn.Module): + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +class TransformerWrapper(nn.Module): + def __init__( + self, + encoder, + decoder, + d_model: int, + two_stage_type="none", # ["none"] only for now + pos_enc_at_input_dec=True, + ): + super().__init__() + + self.encoder = encoder + self.decoder = decoder + self.num_queries = decoder.num_queries if decoder is not None else None + self.pos_enc_at_input_dec = pos_enc_at_input_dec + + # for two stage + assert two_stage_type in ["none"], "unknown param {} of two_stage_type".format( + two_stage_type + ) + self.two_stage_type = two_stage_type + + self._reset_parameters() + self.d_model = d_model + + def _reset_parameters(self): + for n, p in self.named_parameters(): + if p.dim() > 1: + if ( + "box_embed" not in n + and "query_embed" not in n + and "reference_points" not in n + ): + nn.init.xavier_uniform_(p) + + +class MLP(nn.Module): + """Very simple multi-layer perceptron (also called FFN)""" + + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + dropout: float = 0.0, + residual: bool = False, + out_norm: Optional[nn.Module] = None, + ): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + # whether to add the output as a residual connection to the input + if residual and input_dim != output_dim: + raise ValueError("residual is only supported if input_dim == output_dim") + self.residual = residual + # whether to apply a normalization layer to the output + assert isinstance(out_norm, nn.Module) or out_norm is None + self.out_norm = out_norm or nn.Identity() + + def forward(self, x): + orig_x = x + for i, layer in enumerate(self.layers): + x = self.drop(F.relu(layer(x))) if i < self.num_layers - 1 else layer(x) + if self.residual: + x = x + orig_x + x = self.out_norm(x) + return x + + +def get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def get_clones_seq(module, N): + return nn.Sequential(*[copy.deepcopy(module) for i in range(N)]) + + +def get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def get_activation_module(activation): + """Return an activation function given a string""" + if activation == "relu": + return nn.ReLU + if activation == "gelu": + return nn.GELU + if activation == "glu": + return nn.GLU + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def get_valid_ratio(mask): + _, H, W = mask.shape + valid_H = torch.sum(~mask[:, :, 0], 1) + valid_W = torch.sum(~mask[:, 0, :], 1) + valid_ratio_h = valid_H.float() / H + valid_ratio_w = valid_W.float() / W + valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1) + return valid_ratio + + +def gen_sineembed_for_position(pos_tensor, num_feats=256): + assert num_feats % 2 == 0 + num_feats = num_feats // 2 + # n_query, bs, _ = pos_tensor.size() + # sineembed_tensor = torch.zeros(n_query, bs, 256) + scale = 2 * math.pi + dim_t = torch.arange(num_feats, dtype=torch.float32, device=pos_tensor.device) + dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode="floor")) / num_feats) + x_embed = pos_tensor[:, :, 0] * scale + y_embed = pos_tensor[:, :, 1] * scale + pos_x = x_embed[:, :, None] / dim_t + pos_y = y_embed[:, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3 + ).flatten(2) + pos_y = torch.stack( + (pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3 + ).flatten(2) + if pos_tensor.size(-1) == 2: + pos = torch.cat((pos_y, pos_x), dim=2) + elif pos_tensor.size(-1) == 4: + w_embed = pos_tensor[:, :, 2] * scale + pos_w = w_embed[:, :, None] / dim_t + pos_w = torch.stack( + (pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3 + ).flatten(2) + + h_embed = pos_tensor[:, :, 3] * scale + pos_h = h_embed[:, :, None] / dim_t + pos_h = torch.stack( + (pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3 + ).flatten(2) + + pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) + else: + raise ValueError("Unknown pos_tensor shape(-1):{}".format(pos_tensor.size(-1))) + return pos + + +class SAM3Output(list): + """ + A class representing the output of a SAM3 model. + It provides an iterable interface that supports different iteration modes, including iterating over all steps per stage, + last step per stage, and flattened output. + Attributes: + output: The output of the SAM3 model, represented as a list of lists. + iter_mode: The current iteration mode. + Example: + >>> output = [[1, 2], [3, 4], [5, 6]] + >>> sam3_output = SAM3Output(output) + >>> for step in sam3_output: + ... print(step) + [1, 2] + [3, 4] + [5, 6] + >>> with SAM3Output.iteration_mode(SAM3Output.IterMode.LAST_STEP_PER_STAGE) as sam3_last_step_out: + ... for step in sam3_last_step_out: + ... print(step) + [2] + [4] + [6] + >>> with SAM3Output.iteration_mode(SAM3Output.IterMode.FLATTENED) as sam3_flattened_out: + ... for step in sam3_flattened_out: + ... print(step) + 1 + 2 + 3 + 4 + 5 + 6 + """ + + class IterMode(Enum): + # Defines the type of iterator over ouptuts. + ALL_STEPS_PER_STAGE = auto() + LAST_STEP_PER_STAGE = auto() + FLATTENED = auto() # Returns each interactivity step as if it is a separate stage (this is used in SAM3Image model) + + def __init__( + self, + output: List[List[Dict]] = None, + iter_mode: IterMode = IterMode.ALL_STEPS_PER_STAGE, + loss_stages: Optional[List[int]] = None, + ): + if output is not None: + assert ( + isinstance(output, list) + and len(output) > 0 + and isinstance(output[0], list) + ), "Expected output to be a list of lists" + self.output = output + else: + self.output = [] + assert isinstance( + iter_mode, SAM3Output.IterMode + ), f"iter_mode shoulf be of enum type 'SAM3Output.IterMode'. Got {type(iter_mode)}" + + self.iter_mode = iter_mode + # We create a weak reference to self to be used in the lambda functions. + # This is to avoid cyclic references and let SAM3Output be garabge collected. + self_ref = weakref.ref(self) + self._mode2iter = { + SAM3Output.IterMode.ALL_STEPS_PER_STAGE: lambda: iter(self_ref().output), + SAM3Output.IterMode.LAST_STEP_PER_STAGE: lambda: ( + inner_list[-1] for inner_list in self_ref().output + ), + SAM3Output.IterMode.FLATTENED: lambda: ( + element for inner_list in self_ref().output for element in inner_list + ), + } + self.loss_stages = loss_stages + + @override + def __iter__(self) -> Iterator: + return self._mode2iter[self.iter_mode]() + + def __getitem__(self, index): + """ + Returns the item at the specified index. + Args: + index (int): The index of the item to return. + Returns: + list or element: The item at the specified index. + """ + assert isinstance(index, int), f"index should be an integer. Got {type(index)}" + if self.iter_mode == SAM3Output.IterMode.ALL_STEPS_PER_STAGE: + return self.output[index] + elif self.iter_mode == SAM3Output.IterMode.LAST_STEP_PER_STAGE: + return self.output[index][-1] + elif self.iter_mode == SAM3Output.IterMode.FLATTENED: + if index == -1: + return self.self.output[-1][-1] + else: + flattened_output = sum(self.output, []) + return flattened_output[index] + + class _IterationMode(AbstractContextManager): + """ + A context manager that temporarily changes the iteration mode of a SAM3Output object. + This class is used internally by the SAM3Output.iteration_mode method. + """ + + def __init__( + self, model_output: "SAM3Output", iter_mode: "SAM3Output.IterMode" + ): + self._model_output = model_output + self._orig_iter_mode = model_output.iter_mode + self._new_iter_mode = iter_mode + + @override + def __enter__(self) -> "SAM3Output": + self._model_output.iter_mode = self._new_iter_mode + return self._model_output + + @override + def __exit__(self, exc_type, exc_value, traceback): + self._model_output.iter_mode = self._orig_iter_mode + return super().__exit__(exc_type, exc_value, traceback) + + @staticmethod + def iteration_mode( + model_output: "SAM3Output", iter_mode: IterMode + ) -> _IterationMode: + """ + Returns a context manager that allows you to temporarily change the iteration mode of the SAM3Output object. + Args: + model_output: The SAM3Output object. + iter_mode: The new iteration mode. + Returns: + SAM3Output._IterationMode: A context manager that changes the iteration mode of the SAM3Output object. + """ + return SAM3Output._IterationMode(model_output=model_output, iter_mode=iter_mode) + + def append(self, item: list): + assert isinstance( + item, list + ), f"Only list items are supported. Got {type(item)}" + self.output.append(item) + + def __repr__(self): + return self.output.__repr__() + + def __len__(self): + if self.iter_mode in [ + SAM3Output.IterMode.ALL_STEPS_PER_STAGE, + SAM3Output.IterMode.LAST_STEP_PER_STAGE, + ]: + return len(self.output) + elif self.iter_mode == SAM3Output.IterMode.FLATTENED: + flattened_output = sum(self.output, []) + return len(flattened_output) diff --git a/src/nn/segearth_ov3/sam3/model/necks.py b/src/nn/segearth_ov3/sam3/model/necks.py new file mode 100644 index 0000000..21bf9b8 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/necks.py @@ -0,0 +1,125 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Necks are the interface between a vision backbone and the rest of the detection model""" + +from copy import deepcopy +from typing import List, Optional, Tuple + +import torch + +import torch.nn as nn + + +class Sam3DualViTDetNeck(nn.Module): + def __init__( + self, + trunk: nn.Module, + position_encoding: nn.Module, + d_model: int, + scale_factors=(4.0, 2.0, 1.0, 0.5), + add_sam2_neck: bool = False, + ): + """ + SimpleFPN neck a la ViTDet + (From detectron2, very lightly adapted) + It supports a "dual neck" setting, where we have two identical necks (for SAM3 and SAM2), with different weights + + :param trunk: the backbone + :param position_encoding: the positional encoding to use + :param d_model: the dimension of the model + """ + super().__init__() + self.trunk = trunk + self.position_encoding = position_encoding + self.convs = nn.ModuleList() + + self.scale_factors = scale_factors + use_bias = True + dim: int = self.trunk.channel_list[-1] + + for _, scale in enumerate(scale_factors): + current = nn.Sequential() + + if scale == 4.0: + current.add_module( + "dconv_2x2_0", + nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2), + ) + current.add_module( + "gelu", + nn.GELU(), + ) + current.add_module( + "dconv_2x2_1", + nn.ConvTranspose2d(dim // 2, dim // 4, kernel_size=2, stride=2), + ) + out_dim = dim // 4 + elif scale == 2.0: + current.add_module( + "dconv_2x2", + nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2), + ) + out_dim = dim // 2 + elif scale == 1.0: + out_dim = dim + elif scale == 0.5: + current.add_module( + "maxpool_2x2", + nn.MaxPool2d(kernel_size=2, stride=2), + ) + out_dim = dim + else: + raise NotImplementedError(f"scale_factor={scale} is not supported yet.") + + current.add_module( + "conv_1x1", + nn.Conv2d( + in_channels=out_dim, + out_channels=d_model, + kernel_size=1, + bias=use_bias, + ), + ) + current.add_module( + "conv_3x3", + nn.Conv2d( + in_channels=d_model, + out_channels=d_model, + kernel_size=3, + padding=1, + bias=use_bias, + ), + ) + self.convs.append(current) + + self.sam2_convs = None + if add_sam2_neck: + # Assumes sam2 neck is just a clone of the original neck + self.sam2_convs = deepcopy(self.convs) + + def forward( + self, tensor_list: List[torch.Tensor] + ) -> Tuple[ + List[torch.Tensor], + List[torch.Tensor], + Optional[List[torch.Tensor]], + Optional[List[torch.Tensor]], + ]: + xs = self.trunk(tensor_list) + sam3_out, sam3_pos = [], [] + sam2_out, sam2_pos = None, None + if self.sam2_convs is not None: + sam2_out, sam2_pos = [], [] + x = xs[-1] # simpleFPN + for i in range(len(self.convs)): + sam3_x_out = self.convs[i](x) + sam3_pos_out = self.position_encoding(sam3_x_out).to(sam3_x_out.dtype) + sam3_out.append(sam3_x_out) + sam3_pos.append(sam3_pos_out) + + if self.sam2_convs is not None: + sam2_x_out = self.sam2_convs[i](x) + sam2_pos_out = self.position_encoding(sam2_x_out).to(sam2_x_out.dtype) + sam2_out.append(sam2_x_out) + sam2_pos.append(sam2_pos_out) + return sam3_out, sam3_pos, sam2_out, sam2_pos diff --git a/src/nn/segearth_ov3/sam3/model/position_encoding.py b/src/nn/segearth_ov3/sam3/model/position_encoding.py new file mode 100644 index 0000000..eb3f405 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/position_encoding.py @@ -0,0 +1,124 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math +from typing import Optional + +import torch +from torch import nn + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__( + self, + num_pos_feats, + temperature: int = 10000, + normalize: bool = True, + scale: Optional[float] = None, + precompute_resolution: Optional[int] = None, + ): + super().__init__() + assert num_pos_feats % 2 == 0, "Expecting even model width" + self.num_pos_feats = num_pos_feats // 2 + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + self.cache = {} + # Precompute positional encodings under `precompute_resolution` to fill the cache + # and avoid symbolic shape tracing errors in torch.compile in PyTorch 2.4 nightly. + if precompute_resolution is not None: + # We precompute pos enc for stride 4, 8, 16 and 32 to fill `self.cache`. + precompute_sizes = [ + (precompute_resolution // 4, precompute_resolution // 4), + (precompute_resolution // 8, precompute_resolution // 8), + (precompute_resolution // 16, precompute_resolution // 16), + (precompute_resolution // 32, precompute_resolution // 32), + ] + for size in precompute_sizes: + tensors = torch.zeros((1, 1) + size, device="cuda") + self.forward(tensors) + # further clone and detach it in the cache (just to be safe) + self.cache[size] = self.cache[size].clone().detach() + + def _encode_xy(self, x, y): + # The positions are expected to be normalized + assert len(x) == len(y) and x.ndim == y.ndim == 1 + x_embed = x * self.scale + y_embed = y * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, None] / dim_t + pos_y = y_embed[:, None] / dim_t + pos_x = torch.stack( + (pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2 + ).flatten(1) + pos_y = torch.stack( + (pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2 + ).flatten(1) + return pos_x, pos_y + + @torch.no_grad() + def encode_boxes(self, x, y, w, h): + pos_x, pos_y = self._encode_xy(x, y) + pos = torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1) + return pos + + encode = encode_boxes # Backwards compatibility + + @torch.no_grad() + def encode_points(self, x, y, labels): + (bx, nx), (by, ny), (bl, nl) = x.shape, y.shape, labels.shape + assert bx == by and nx == ny and bx == bl and nx == nl + pos_x, pos_y = self._encode_xy(x.flatten(), y.flatten()) + pos_x, pos_y = pos_x.reshape(bx, nx, -1), pos_y.reshape(by, ny, -1) + pos = torch.cat((pos_y, pos_x, labels[:, :, None]), dim=2) + return pos + + @torch.no_grad() + def forward(self, x): + cache_key = None + cache_key = (x.shape[-2], x.shape[-1]) + if cache_key in self.cache: + return self.cache[cache_key][None].repeat(x.shape[0], 1, 1, 1) + y_embed = ( + torch.arange(1, x.shape[-2] + 1, dtype=torch.float32, device=x.device) + .view(1, -1, 1) + .repeat(x.shape[0], 1, x.shape[-1]) + ) + x_embed = ( + torch.arange(1, x.shape[-1] + 1, dtype=torch.float32, device=x.device) + .view(1, 1, -1) + .repeat(x.shape[0], x.shape[-2], 1) + ) + + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + if cache_key is not None: + self.cache[cache_key] = pos[0] + return pos diff --git a/src/nn/segearth_ov3/sam3/model/sam1_task_predictor.py b/src/nn/segearth_ov3/sam3/model/sam1_task_predictor.py new file mode 100644 index 0000000..f5e49b1 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam1_task_predictor.py @@ -0,0 +1,458 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging + +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch + +import torch.nn as nn +from PIL.Image import Image + +from sam3.model.sam3_tracker_base import Sam3TrackerBase +from sam3.model.utils.sam1_utils import SAM2Transforms + + +# Adapted from https://github.com/facebookresearch/sam2/blob/main/sam2/sam2_image_predictor.py +class SAM3InteractiveImagePredictor(nn.Module): + def __init__( + self, + sam_model: Sam3TrackerBase, + mask_threshold=0.0, + max_hole_area=256.0, + max_sprinkle_area=0.0, + **kwargs, + ) -> None: + """ + Uses SAM-3 to calculate the image embedding for an image, and then + allow repeated, efficient mask prediction given prompts. + + Arguments: + sam_model : The model to use for mask prediction. + mask_threshold (float): The threshold to use when converting mask logits + to binary masks. Masks are thresholded at 0 by default. + max_hole_area (int): If max_hole_area > 0, we fill small holes in up to + the maximum area of max_hole_area in low_res_masks. + max_sprinkle_area (int): If max_sprinkle_area > 0, we remove small sprinkles up to + the maximum area of max_sprinkle_area in low_res_masks. + """ + super().__init__() + self.model = sam_model + self._transforms = SAM2Transforms( + resolution=self.model.image_size, + mask_threshold=mask_threshold, + max_hole_area=max_hole_area, + max_sprinkle_area=max_sprinkle_area, + ) + + # Predictor state + self._is_image_set = False + self._features = None + self._orig_hw = None + # Whether the predictor is set for single image or a batch of images + self._is_batch = False + + # Predictor config + self.mask_threshold = mask_threshold + + # Spatial dim for backbone feature maps + self._bb_feat_sizes = [ + (288, 288), + (144, 144), + (72, 72), + ] + + @torch.no_grad() + def set_image( + self, + image: Union[np.ndarray, Image], + ) -> None: + """ + Calculates the image embeddings for the provided image, allowing + masks to be predicted with the 'predict' method. + + Arguments: + image (np.ndarray or PIL Image): The input image to embed in RGB format. The image should be in HWC format if np.ndarray, or WHC format if PIL Image + with pixel values in [0, 255]. + image_format (str): The color format of the image, in ['RGB', 'BGR']. + """ + self.reset_predictor() + # Transform the image to the form expected by the model + if isinstance(image, np.ndarray): + logging.info("For numpy array image, we assume (HxWxC) format") + self._orig_hw = [image.shape[:2]] + elif isinstance(image, Image): + w, h = image.size + self._orig_hw = [(h, w)] + else: + raise NotImplementedError("Image format not supported") + + input_image = self._transforms(image) + input_image = input_image[None, ...].to(self.device) + + assert ( + len(input_image.shape) == 4 and input_image.shape[1] == 3 + ), f"input_image must be of size 1x3xHxW, got {input_image.shape}" + logging.info("Computing image embeddings for the provided image...") + backbone_out = self.model.forward_image(input_image) + ( + _, + vision_feats, + _, + _, + ) = self.model._prepare_backbone_features(backbone_out) + # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos + vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed + + feats = [ + feat.permute(1, 2, 0).view(1, -1, *feat_size) + for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1]) + ][::-1] + self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]} + self._is_image_set = True + logging.info("Image embeddings computed.") + + @torch.no_grad() + def set_image_batch( + self, + image_list: List[Union[np.ndarray]], + ) -> None: + """ + Calculates the image embeddings for the provided image batch, allowing + masks to be predicted with the 'predict_batch' method. + + Arguments: + image_list (List[np.ndarray]): The input images to embed in RGB format. The image should be in HWC format if np.ndarray + with pixel values in [0, 255]. + """ + self.reset_predictor() + assert isinstance(image_list, list) + self._orig_hw = [] + for image in image_list: + assert isinstance( + image, np.ndarray + ), "Images are expected to be an np.ndarray in RGB format, and of shape HWC" + self._orig_hw.append(image.shape[:2]) + # Transform the image to the form expected by the model + img_batch = self._transforms.forward_batch(image_list) + img_batch = img_batch.to(self.device) + batch_size = img_batch.shape[0] + assert ( + len(img_batch.shape) == 4 and img_batch.shape[1] == 3 + ), f"img_batch must be of size Bx3xHxW, got {img_batch.shape}" + logging.info("Computing image embeddings for the provided images...") + backbone_out = self.model.forward_image(img_batch) + ( + _, + vision_feats, + _, + _, + ) = self.model._prepare_backbone_features(backbone_out) + # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos + vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed + + feats = [ + feat.permute(1, 2, 0).view(batch_size, -1, *feat_size) + for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1]) + ][::-1] + self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]} + self._is_image_set = True + self._is_batch = True + logging.info("Image embeddings computed.") + + def predict_batch( + self, + point_coords_batch: List[np.ndarray] = None, + point_labels_batch: List[np.ndarray] = None, + box_batch: List[np.ndarray] = None, + mask_input_batch: List[np.ndarray] = None, + multimask_output: bool = True, + return_logits: bool = False, + normalize_coords=True, + ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: + """This function is very similar to predict(...), however it is used for batched mode, when the model is expected to generate predictions on multiple images. + It returns a tuple of lists of masks, ious, and low_res_masks_logits. + """ + assert self._is_batch, "This function should only be used when in batched mode" + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image_batch(...) before mask prediction." + ) + num_images = len(self._features["image_embed"]) + all_masks = [] + all_ious = [] + all_low_res_masks = [] + for img_idx in range(num_images): + # Transform input prompts + point_coords = ( + point_coords_batch[img_idx] if point_coords_batch is not None else None + ) + point_labels = ( + point_labels_batch[img_idx] if point_labels_batch is not None else None + ) + box = box_batch[img_idx] if box_batch is not None else None + mask_input = ( + mask_input_batch[img_idx] if mask_input_batch is not None else None + ) + mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts( + point_coords, + point_labels, + box, + mask_input, + normalize_coords, + img_idx=img_idx, + ) + masks, iou_predictions, low_res_masks = self._predict( + unnorm_coords, + labels, + unnorm_box, + mask_input, + multimask_output, + return_logits=return_logits, + img_idx=img_idx, + ) + masks_np = masks.squeeze(0).float().detach().cpu().numpy() + iou_predictions_np = ( + iou_predictions.squeeze(0).float().detach().cpu().numpy() + ) + low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy() + all_masks.append(masks_np) + all_ious.append(iou_predictions_np) + all_low_res_masks.append(low_res_masks_np) + + return all_masks, all_ious, all_low_res_masks + + def predict( + self, + point_coords: Optional[np.ndarray] = None, + point_labels: Optional[np.ndarray] = None, + box: Optional[np.ndarray] = None, + mask_input: Optional[np.ndarray] = None, + multimask_output: bool = True, + return_logits: bool = False, + normalize_coords=True, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Predict masks for the given input prompts, using the currently set image. + + Arguments: + point_coords (np.ndarray or None): A Nx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (np.ndarray or None): A length N array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + box (np.ndarray or None): A length 4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form 1xHxW, where + for SAM, H=W=256. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + normalize_coords (bool): If true, the point coordinates will be normalized to the range [0,1] and point_coords is expected to be wrt. image dimensions. + + Returns: + (np.ndarray): The output masks in CxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (np.ndarray): An array of length C containing the model's + predictions for the quality of each mask. + (np.ndarray): An array of shape CxHxW, where C is the number + of masks and H=W=256. These low resolution logits can be passed to + a subsequent iteration as mask input. + """ + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) before mask prediction." + ) + + # Transform input prompts + + mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts( + point_coords, point_labels, box, mask_input, normalize_coords + ) + + masks, iou_predictions, low_res_masks = self._predict( + unnorm_coords, + labels, + unnorm_box, + mask_input, + multimask_output, + return_logits=return_logits, + ) + + masks_np = masks.squeeze(0).float().detach().cpu().numpy() + iou_predictions_np = iou_predictions.squeeze(0).float().detach().cpu().numpy() + low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy() + return masks_np, iou_predictions_np, low_res_masks_np + + def _prep_prompts( + self, point_coords, point_labels, box, mask_logits, normalize_coords, img_idx=-1 + ): + unnorm_coords, labels, unnorm_box, mask_input = None, None, None, None + if point_coords is not None: + assert ( + point_labels is not None + ), "point_labels must be supplied if point_coords is supplied." + point_coords = torch.as_tensor( + point_coords, dtype=torch.float, device=self.device + ) + unnorm_coords = self._transforms.transform_coords( + point_coords, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx] + ) + labels = torch.as_tensor(point_labels, dtype=torch.int, device=self.device) + if len(unnorm_coords.shape) == 2: + unnorm_coords, labels = unnorm_coords[None, ...], labels[None, ...] + if box is not None: + box = torch.as_tensor(box, dtype=torch.float, device=self.device) + unnorm_box = self._transforms.transform_boxes( + box, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx] + ) # Bx2x2 + if mask_logits is not None: + mask_input = torch.as_tensor( + mask_logits, dtype=torch.float, device=self.device + ) + if len(mask_input.shape) == 3: + mask_input = mask_input[None, :, :, :] + return mask_input, unnorm_coords, labels, unnorm_box + + @torch.no_grad() + def _predict( + self, + point_coords: Optional[torch.Tensor], + point_labels: Optional[torch.Tensor], + boxes: Optional[torch.Tensor] = None, + mask_input: Optional[torch.Tensor] = None, + multimask_output: bool = True, + return_logits: bool = False, + img_idx: int = -1, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Predict masks for the given input prompts, using the currently set image. + Input prompts are batched torch tensors and are expected to already be + transformed to the input frame using SAM2Transforms. + + Arguments: + point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (torch.Tensor or None): A BxN array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + boxes (np.ndarray or None): A Bx4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form Bx1xHxW, where + for SAM, H=W=256. Masks returned by a previous iteration of the + predict method do not need further transformation. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + + Returns: + (torch.Tensor): The output masks in BxCxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (torch.Tensor): An array of shape BxC containing the model's + predictions for the quality of each mask. + (torch.Tensor): An array of shape BxCxHxW, where C is the number + of masks and H=W=256. These low res logits can be passed to + a subsequent iteration as mask input. + """ + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) before mask prediction." + ) + + if point_coords is not None: + concat_points = (point_coords, point_labels) + else: + concat_points = None + + # Embed prompts + if boxes is not None: + box_coords = boxes.reshape(-1, 2, 2) + box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device) + box_labels = box_labels.repeat(boxes.size(0), 1) + # we merge "boxes" and "points" into a single "concat_points" input (where + # boxes are added at the beginning) to sam_prompt_encoder + if concat_points is not None: + concat_coords = torch.cat([box_coords, concat_points[0]], dim=1) + concat_labels = torch.cat([box_labels, concat_points[1]], dim=1) + concat_points = (concat_coords, concat_labels) + else: + concat_points = (box_coords, box_labels) + + sparse_embeddings, dense_embeddings = self.model.sam_prompt_encoder( + points=concat_points, + boxes=None, + masks=mask_input, + ) + + # Predict masks + batched_mode = ( + concat_points is not None and concat_points[0].shape[0] > 1 + ) # multi object prediction + high_res_features = [ + feat_level[img_idx].unsqueeze(0) + for feat_level in self._features["high_res_feats"] + ] + low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder( + image_embeddings=self._features["image_embed"][img_idx].unsqueeze(0), + image_pe=self.model.sam_prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + repeat_image=batched_mode, + high_res_features=high_res_features, + ) + + # Upscale the masks to the original image resolution + masks = self._transforms.postprocess_masks( + low_res_masks, self._orig_hw[img_idx] + ) + low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0) + if not return_logits: + masks = masks > self.mask_threshold + + return masks, iou_predictions, low_res_masks + + def get_image_embedding(self) -> torch.Tensor: + """ + Returns the image embeddings for the currently set image, with + shape 1xCxHxW, where C is the embedding dimension and (H,W) are + the embedding spatial dimension of SAM (typically C=256, H=W=64). + """ + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) to generate an embedding." + ) + assert ( + self._features is not None + ), "Features must exist if an image has been set." + return self._features["image_embed"] + + @property + def device(self) -> torch.device: + return self.model.device + + def reset_predictor(self) -> None: + """ + Resets the image embeddings and other state variables. + """ + self._is_image_set = False + self._features = None + self._orig_hw = None + self._is_batch = False diff --git a/src/nn/segearth_ov3/sam3/model/sam3_image.py b/src/nn/segearth_ov3/sam3/model/sam3_image.py new file mode 100644 index 0000000..4f3d0b0 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_image.py @@ -0,0 +1,892 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os +from copy import deepcopy +from typing import Dict, List, Optional, Tuple + +import numpy as np +import torch + +from sam3.model.model_misc import SAM3Output + +from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor +from sam3.model.vl_combiner import SAM3VLBackbone +from sam3.perflib.nms import nms_masks + +# from sam3.train.data.collator import BatchedDatapoint +from sam3.model.data_misc import BatchedDatapoint + +from .act_ckpt_utils import activation_ckpt_wrapper + +from .box_ops import box_cxcywh_to_xyxy + +from .geometry_encoders import Prompt +from .model_misc import inverse_sigmoid + + +def _update_out(out, out_name, out_value, auxiliary=True, update_aux=True): + out[out_name] = out_value[-1] if auxiliary else out_value + if auxiliary and update_aux: + if "aux_outputs" not in out: + out["aux_outputs"] = [{} for _ in range(len(out_value) - 1)] + assert len(out["aux_outputs"]) == len(out_value) - 1 + for aux_output, aux_value in zip(out["aux_outputs"], out_value[:-1]): + aux_output[out_name] = aux_value + + +class Sam3Image(torch.nn.Module): + TEXT_ID_FOR_TEXT = 0 + TEXT_ID_FOR_VISUAL = 1 + TEXT_ID_FOR_GEOMETRIC = 2 + + def __init__( + self, + backbone: SAM3VLBackbone, + transformer, + input_geometry_encoder, + segmentation_head=None, + num_feature_levels=1, + o2m_mask_predict=True, + dot_prod_scoring=None, + use_instance_query: bool = True, + multimask_output: bool = True, + use_act_checkpoint_seg_head: bool = True, + interactivity_in_encoder: bool = True, + matcher=None, + use_dot_prod_scoring=True, + supervise_joint_box_scores: bool = False, # only relevant if using presence token/score + detach_presence_in_joint_score: bool = False, # only relevant if using presence token/score + separate_scorer_for_instance: bool = False, + num_interactive_steps_val: int = 0, + inst_interactive_predictor: SAM3InteractiveImagePredictor = None, + **kwargs, + ): + super().__init__() + self.backbone = backbone + self.geometry_encoder = input_geometry_encoder + self.transformer = transformer + self.hidden_dim = transformer.d_model + self.num_feature_levels = num_feature_levels + self.segmentation_head = segmentation_head + + self.o2m_mask_predict = o2m_mask_predict + + self.dot_prod_scoring = dot_prod_scoring + self.use_act_checkpoint_seg_head = use_act_checkpoint_seg_head + self.interactivity_in_encoder = interactivity_in_encoder + self.matcher = matcher + + self.num_interactive_steps_val = num_interactive_steps_val + self.use_dot_prod_scoring = use_dot_prod_scoring + + if self.use_dot_prod_scoring: + assert dot_prod_scoring is not None + self.dot_prod_scoring = dot_prod_scoring + self.instance_dot_prod_scoring = None + if separate_scorer_for_instance: + self.instance_dot_prod_scoring = deepcopy(dot_prod_scoring) + else: + self.class_embed = torch.nn.Linear(self.hidden_dim, 1) + self.instance_class_embed = None + if separate_scorer_for_instance: + self.instance_class_embed = deepcopy(self.class_embed) + + self.supervise_joint_box_scores = supervise_joint_box_scores + self.detach_presence_in_joint_score = detach_presence_in_joint_score + + # verify the number of queries for O2O and O2M + num_o2o_static = self.transformer.decoder.num_queries + num_o2m_static = self.transformer.decoder.num_o2m_queries + assert num_o2m_static == (num_o2o_static if self.transformer.decoder.dac else 0) + self.dac = self.transformer.decoder.dac + + self.use_instance_query = use_instance_query + self.multimask_output = multimask_output + + self.inst_interactive_predictor = inst_interactive_predictor + + @property + def device(self): + self._device = getattr(self, "_device", None) or next(self.parameters()).device + return self._device + + def to(self, *args, **kwargs): + # clear cached _device in case the model is moved to a different device + self._device = None + return super().to(*args, **kwargs) + + def _get_img_feats(self, backbone_out, img_ids): + """Retrieve correct image features from backbone output.""" + if "backbone_fpn" in backbone_out: + if "id_mapping" in backbone_out and backbone_out["id_mapping"] is not None: + img_ids = backbone_out["id_mapping"][img_ids] + # If this assert fails, it likely means we're requesting different img_ids (perhaps a different frame?) + # We currently don't expect this to happen. We could technically trigger a recompute here, + # but likely at the cost of a cpu<->gpu sync point, which would deteriorate perf + torch._assert_async((img_ids >= 0).all()) + + vis_feats = backbone_out["backbone_fpn"][-self.num_feature_levels :] + vis_pos_enc = backbone_out["vision_pos_enc"][-self.num_feature_levels :] + vis_feat_sizes = [x.shape[-2:] for x in vis_pos_enc] # (H, W) shapes + # index and flatten visual features NxCxHxW => HWxNxC (batch-first => seq-first) + img_feats = [x[img_ids].flatten(2).permute(2, 0, 1) for x in vis_feats] + img_pos_embeds = [ + x[img_ids].flatten(2).permute(2, 0, 1) for x in vis_pos_enc + ] + return backbone_out, img_feats, img_pos_embeds, vis_feat_sizes + # img_feats: shape [[5184, 1, 256]] + # img_pos_embeds: shape [[5184, 1, 256]] + # vis_feat_sizes: torch.Size([72, 72]) + + # Image features not available in backbone output, so we compute them on the fly + # This case likely occurs for video. In that case, we want to forward only the current frame + img_batch = backbone_out["img_batch_all_stages"] + if img_ids.numel() > 1: + # Only forward backbone on unique image ids to avoid repetitive computation + unique_ids, _ = torch.unique(img_ids, return_inverse=True) + else: + unique_ids, _ = img_ids, slice(None) + # Compute the image features on those unique image ids + # note: we allow using a list (or other indexable types) of tensors as img_batch + # (e.g. for async frame loading in demo). In this case we index img_batch.tensors directly + if isinstance(img_batch, torch.Tensor): + image = img_batch[unique_ids] + elif unique_ids.numel() == 1: + image = img_batch[unique_ids.item()].unsqueeze(0) + else: + image = torch.stack([img_batch[i] for i in unique_ids.tolist()]) + # `img_batch` might be fp16 and offloaded to CPU + image = image.to(dtype=torch.float32, device=self.device) + # Next time we call this function, we want to remember which indices we computed + id_mapping = torch.full( + (len(img_batch),), -1, dtype=torch.long, device=self.device + ) + id_mapping[unique_ids] = torch.arange(len(unique_ids), device=self.device) + backbone_out = { + **backbone_out, + **self.backbone.forward_image(image), + "id_mapping": id_mapping, + } + assert "backbone_fpn" in backbone_out + return self._get_img_feats(backbone_out, img_ids=img_ids) + + def _encode_prompt( + self, + backbone_out, + find_input, + geometric_prompt, + visual_prompt_embed=None, + visual_prompt_mask=None, + encode_text=True, + prev_mask_pred=None, + ): + # index text features (note that regardless of early or late fusion, the batch size of + # `txt_feats` is always the number of *prompts* in the encoder) + txt_ids = find_input.text_ids + txt_feats = backbone_out["language_features"][:, txt_ids] # [32, 1, 256] + txt_masks = backbone_out["language_mask"][txt_ids] # [1, 32] + + feat_tuple = self._get_img_feats(backbone_out, find_input.img_ids) + backbone_out, img_feats, img_pos_embeds, vis_feat_sizes = feat_tuple + + if prev_mask_pred is not None: + img_feats = [img_feats[-1] + prev_mask_pred] # [5184, 1, 256] + # Encode geometry + geo_feats, geo_masks = self.geometry_encoder( + geo_prompt=geometric_prompt, + img_feats=img_feats, + img_sizes=vis_feat_sizes, + img_pos_embeds=img_pos_embeds, + ) + # geo_feats: [1, 1, 256] + # geo_masks: [False] + + if visual_prompt_embed is None: + visual_prompt_embed = torch.zeros( + (0, *geo_feats.shape[1:]), device=geo_feats.device + ) + visual_prompt_mask = torch.zeros( + (*geo_masks.shape[:-1], 0), + device=geo_masks.device, + dtype=geo_masks.dtype, + ) + if encode_text: + prompt = torch.cat([txt_feats, geo_feats, visual_prompt_embed], dim=0) + prompt_mask = torch.cat([txt_masks, geo_masks, visual_prompt_mask], dim=1) + else: + prompt = torch.cat([geo_feats, visual_prompt_embed], dim=0) + prompt_mask = torch.cat([geo_masks, visual_prompt_mask], dim=1) + return prompt, prompt_mask, backbone_out + + def _run_encoder( + self, + backbone_out, + find_input, + prompt, + prompt_mask, + encoder_extra_kwargs: Optional[Dict] = None, + ): + feat_tuple = self._get_img_feats(backbone_out, find_input.img_ids) + backbone_out, img_feats, img_pos_embeds, vis_feat_sizes = feat_tuple + + # Run the encoder + prompt_pos_embed = torch.zeros_like(prompt) + # make a copy of the image feature lists since the encoder may modify these lists in-place + memory = self.transformer.encoder( + src=img_feats.copy(), + src_key_padding_mask=None, + src_pos=img_pos_embeds.copy(), + prompt=prompt, + prompt_pos=prompt_pos_embed, + prompt_key_padding_mask=prompt_mask, + feat_sizes=vis_feat_sizes, + encoder_extra_kwargs=encoder_extra_kwargs, + ) + encoder_out = { + # encoded image features + "encoder_hidden_states": memory["memory"], + "pos_embed": memory["pos_embed"], + "padding_mask": memory["padding_mask"], + "level_start_index": memory["level_start_index"], + "spatial_shapes": memory["spatial_shapes"], + "valid_ratios": memory["valid_ratios"], + "vis_feat_sizes": vis_feat_sizes, + # encoded text features (or other prompts) + "prompt_before_enc": prompt, + "prompt_after_enc": memory.get("memory_text", prompt), + "prompt_mask": prompt_mask, + } + return backbone_out, encoder_out, feat_tuple + + def _run_decoder( + self, + pos_embed, + memory, + src_mask, + out, + prompt, + prompt_mask, + encoder_out, + ): + bs = memory.shape[1] + query_embed = self.transformer.decoder.query_embed.weight + tgt = query_embed.unsqueeze(1).repeat(1, bs, 1) # [200, 1, 256] + + apply_dac = self.transformer.decoder.dac and self.training + hs, reference_boxes, dec_presence_out, dec_presence_feats = ( + self.transformer.decoder( + tgt=tgt, + memory=memory, + memory_key_padding_mask=src_mask, + pos=pos_embed, + reference_boxes=None, + level_start_index=encoder_out["level_start_index"], + spatial_shapes=encoder_out["spatial_shapes"], + valid_ratios=encoder_out["valid_ratios"], + tgt_mask=None, + memory_text=prompt, + text_attention_mask=prompt_mask, + apply_dac=apply_dac, + ) + ) + hs = hs.transpose(1, 2) # seq-first to batch-first + reference_boxes = reference_boxes.transpose(1, 2) # seq-first to batch-first + if dec_presence_out is not None: + # seq-first to batch-first + dec_presence_out = dec_presence_out.transpose(1, 2) + + out["presence_feats"] = dec_presence_feats + self._update_scores_and_boxes( + out, + hs, + reference_boxes, + prompt, + prompt_mask, + dec_presence_out=dec_presence_out, + ) + return out, hs + + def _update_scores_and_boxes( + self, + out, + hs, + reference_boxes, + prompt, + prompt_mask, + dec_presence_out=None, + is_instance_prompt=False, + ): + apply_dac = self.transformer.decoder.dac and self.training + num_o2o = (hs.size(2) // 2) if apply_dac else hs.size(2) + num_o2m = hs.size(2) - num_o2o + assert num_o2m == (num_o2o if apply_dac else 0) + out["queries"] = hs[-1][:, :num_o2o] # remove o2m queries if there are any + # score prediction + if self.use_dot_prod_scoring: + dot_prod_scoring_head = self.dot_prod_scoring + if is_instance_prompt and self.instance_dot_prod_scoring is not None: + dot_prod_scoring_head = self.instance_dot_prod_scoring + outputs_class = dot_prod_scoring_head(hs, prompt, prompt_mask) + else: + class_embed_head = self.class_embed + if is_instance_prompt and self.instance_class_embed is not None: + class_embed_head = self.instance_class_embed + outputs_class = class_embed_head(hs) + + # box prediction + box_head = self.transformer.decoder.bbox_embed + if ( + is_instance_prompt + and self.transformer.decoder.instance_bbox_embed is not None + ): + box_head = self.transformer.decoder.instance_bbox_embed + anchor_box_offsets = box_head(hs) + reference_boxes_inv_sig = inverse_sigmoid(reference_boxes) + outputs_coord = (reference_boxes_inv_sig + anchor_box_offsets).sigmoid() + outputs_boxes_xyxy = box_cxcywh_to_xyxy(outputs_coord) + + if dec_presence_out is not None: + _update_out( + out, "presence_logit_dec", dec_presence_out, update_aux=self.training + ) + + if self.supervise_joint_box_scores: + assert dec_presence_out is not None + prob_dec_presence_out = dec_presence_out.clone().sigmoid() + if self.detach_presence_in_joint_score: + prob_dec_presence_out = prob_dec_presence_out.detach() + + outputs_class = inverse_sigmoid( + outputs_class.sigmoid() * prob_dec_presence_out.unsqueeze(2) + ).clamp(min=-10.0, max=10.0) + + _update_out( + out, "pred_logits", outputs_class[:, :, :num_o2o], update_aux=self.training + ) + _update_out( + out, "pred_boxes", outputs_coord[:, :, :num_o2o], update_aux=self.training + ) + _update_out( + out, + "pred_boxes_xyxy", + outputs_boxes_xyxy[:, :, :num_o2o], + update_aux=self.training, + ) + if num_o2m > 0 and self.training: + _update_out( + out, + "pred_logits_o2m", + outputs_class[:, :, num_o2o:], + update_aux=self.training, + ) + _update_out( + out, + "pred_boxes_o2m", + outputs_coord[:, :, num_o2o:], + update_aux=self.training, + ) + _update_out( + out, + "pred_boxes_xyxy_o2m", + outputs_boxes_xyxy[:, :, num_o2o:], + update_aux=self.training, + ) + + def _run_segmentation_heads( + self, + out, + backbone_out, + img_ids, + vis_feat_sizes, + encoder_hidden_states, + prompt, + prompt_mask, + hs, + ): + apply_dac = self.transformer.decoder.dac and self.training + if self.segmentation_head is not None: + num_o2o = (hs.size(2) // 2) if apply_dac else hs.size(2) + num_o2m = hs.size(2) - num_o2o + obj_queries = hs if self.o2m_mask_predict else hs[:, :, :num_o2o] + seg_head_outputs = activation_ckpt_wrapper(self.segmentation_head)( + backbone_feats=backbone_out["backbone_fpn"], + obj_queries=obj_queries, + image_ids=img_ids, + encoder_hidden_states=encoder_hidden_states, + act_ckpt_enable=self.training and self.use_act_checkpoint_seg_head, + prompt=prompt, + prompt_mask=prompt_mask, + ) # ['pred_masks', 'semantic_seg', 'presence_logit']. 'pred_masks': [1, 200, 288, 288]; 'semantic_seg': [1, 1, 288, 288]; 'pred_masks': None + aux_masks = False # self.aux_loss and self.segmentation_head.aux_masks + for k, v in seg_head_outputs.items(): + if k in self.segmentation_head.instance_keys: + _update_out(out, k, v[:, :num_o2o], auxiliary=aux_masks) + if ( + self.o2m_mask_predict and num_o2m > 0 + ): # handle o2m mask prediction + _update_out( + out, f"{k}_o2m", v[:, num_o2o:], auxiliary=aux_masks + ) + else: + out[k] = v + else: + backbone_out.pop("backbone_fpn", None) + + def _get_best_mask(self, out): + prev_mask_idx = out["pred_logits"].argmax(dim=1).squeeze(1) + batch_idx = torch.arange( + out["pred_logits"].shape[0], device=prev_mask_idx.device + ) + prev_mask_pred = out["pred_masks"][batch_idx, prev_mask_idx][:, None] + # Downsample mask to match image resolution. + prev_mask_pred = self.geometry_encoder.mask_encoder.mask_downsampler( + prev_mask_pred + ) + prev_mask_pred = prev_mask_pred.flatten(-2).permute(2, 0, 1) + + return prev_mask_pred + + def forward_grounding( + self, + backbone_out, + find_input, + find_target, + geometric_prompt: Prompt, + ): + with torch.profiler.record_function("SAM3Image._encode_prompt"): + prompt, prompt_mask, backbone_out = self._encode_prompt( + backbone_out, find_input, geometric_prompt + ) + # Run the encoder + with torch.profiler.record_function("SAM3Image._run_encoder"): + backbone_out, encoder_out, _ = self._run_encoder( + backbone_out, find_input, prompt, prompt_mask + ) + out = { + "encoder_hidden_states": encoder_out["encoder_hidden_states"], + "prev_encoder_out": { + "encoder_out": encoder_out, + "backbone_out": backbone_out, + }, + } + + # Run the decoder + with torch.profiler.record_function("SAM3Image._run_decoder"): + out, hs = self._run_decoder( + memory=out["encoder_hidden_states"], + pos_embed=encoder_out["pos_embed"], + src_mask=encoder_out["padding_mask"], + out=out, + prompt=prompt, + prompt_mask=prompt_mask, + encoder_out=encoder_out, + ) + # out: ['encoder_hidden_states', 'prev_encoder_out', 'presence_feats', 'queries', 'presence_logit_dec', 'pred_logits', 'pred_boxes', 'pred_boxes_xyxy'] + # hs: [6, 1, 200, 256] + + # Run segmentation heads + with torch.profiler.record_function("SAM3Image._run_segmentation_heads"): + self._run_segmentation_heads( + out=out, + backbone_out=backbone_out, + img_ids=find_input.img_ids, + vis_feat_sizes=encoder_out["vis_feat_sizes"], + encoder_hidden_states=out["encoder_hidden_states"], + prompt=prompt, + prompt_mask=prompt_mask, + hs=hs, + ) + + if self.training or self.num_interactive_steps_val > 0: + self._compute_matching(out, self.back_convert(find_target)) + return out + + def _postprocess_out(self, out: Dict, multimask_output: bool = False): + # For multimask output, during eval we return the single best mask with the dict keys expected by the evaluators, but also return the multimasks output with new keys. + num_mask_boxes = out["pred_boxes"].size(1) + if not self.training and multimask_output and num_mask_boxes > 1: + out["multi_pred_logits"] = out["pred_logits"] + if "pred_masks" in out: + out["multi_pred_masks"] = out["pred_masks"] + out["multi_pred_boxes"] = out["pred_boxes"] + out["multi_pred_boxes_xyxy"] = out["pred_boxes_xyxy"] + + best_mask_idx = out["pred_logits"].argmax(1).squeeze(1) + batch_idx = torch.arange(len(best_mask_idx), device=best_mask_idx.device) + + out["pred_logits"] = out["pred_logits"][batch_idx, best_mask_idx].unsqueeze( + 1 + ) + if "pred_masks" in out: + out["pred_masks"] = out["pred_masks"][ + batch_idx, best_mask_idx + ].unsqueeze(1) + out["pred_boxes"] = out["pred_boxes"][batch_idx, best_mask_idx].unsqueeze(1) + out["pred_boxes_xyxy"] = out["pred_boxes_xyxy"][ + batch_idx, best_mask_idx + ].unsqueeze(1) + + return out + + def _get_dummy_prompt(self, num_prompts=1): + device = self.device + geometric_prompt = Prompt( + box_embeddings=torch.zeros(0, num_prompts, 4, device=device), + box_mask=torch.zeros(num_prompts, 0, device=device, dtype=torch.bool), + ) + return geometric_prompt + + def forward(self, input: BatchedDatapoint): + device = self.device + backbone_out = {"img_batch_all_stages": input.img_batch} + backbone_out.update(self.backbone.forward_image(input.img_batch)) + num_frames = len(input.find_inputs) + assert num_frames == 1 + + text_outputs = self.backbone.forward_text(input.find_text_batch, device=device) + backbone_out.update(text_outputs) + + previous_stages_out = SAM3Output( + iter_mode=SAM3Output.IterMode.LAST_STEP_PER_STAGE + ) + + find_input = input.find_inputs[0] + find_target = input.find_targets[0] + + if find_input.input_points is not None and find_input.input_points.numel() > 0: + print("Warning: Point prompts are ignored in PCS.") + + num_interactive_steps = 0 if self.training else self.num_interactive_steps_val + geometric_prompt = Prompt( + box_embeddings=find_input.input_boxes, + box_mask=find_input.input_boxes_mask, + box_labels=find_input.input_boxes_label, + ) + + # Init vars that are shared across the loop. + stage_outs = [] + for cur_step in range(num_interactive_steps + 1): + if cur_step > 0: + # We sample interactive geometric prompts (boxes, points) + geometric_prompt, _ = self.interactive_prompt_sampler.sample( + geo_prompt=geometric_prompt, + find_target=find_target, + previous_out=stage_outs[-1], + ) + out = self.forward_grounding( + backbone_out=backbone_out, + find_input=find_input, + find_target=find_target, + geometric_prompt=geometric_prompt.clone(), + ) + stage_outs.append(out) + + previous_stages_out.append(stage_outs) + return previous_stages_out + + def _compute_matching(self, out, targets): + out["indices"] = self.matcher(out, targets) + for aux_out in out.get("aux_outputs", []): + aux_out["indices"] = self.matcher(aux_out, targets) + + def back_convert(self, targets): + batched_targets = { + "boxes": targets.boxes.view(-1, 4), + "boxes_xyxy": box_cxcywh_to_xyxy(targets.boxes.view(-1, 4)), + "boxes_padded": targets.boxes_padded, + "positive_map": targets.boxes.new_ones(len(targets.boxes), 1), + "num_boxes": targets.num_boxes, + "masks": targets.segments, + "semantic_masks": targets.semantic_segments, + "is_valid_mask": targets.is_valid_segment, + "is_exhaustive": targets.is_exhaustive, + "object_ids_packed": targets.object_ids, + "object_ids_padded": targets.object_ids_padded, + } + return batched_targets + + def predict_inst( + self, + inference_state, + **kwargs, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + orig_h, orig_w = ( + inference_state["original_height"], + inference_state["original_width"], + ) + backbone_out = inference_state["backbone_out"]["sam2_backbone_out"] + ( + _, + vision_feats, + _, + _, + ) = self.inst_interactive_predictor.model._prepare_backbone_features( + backbone_out + ) + # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos + vision_feats[-1] = ( + vision_feats[-1] + self.inst_interactive_predictor.model.no_mem_embed + ) + feats = [ + feat.permute(1, 2, 0).view(1, -1, *feat_size) + for feat, feat_size in zip( + vision_feats[::-1], self.inst_interactive_predictor._bb_feat_sizes[::-1] + ) + ][::-1] + self.inst_interactive_predictor._features = { + "image_embed": feats[-1], + "high_res_feats": feats[:-1], + } + self.inst_interactive_predictor._is_image_set = True + self.inst_interactive_predictor._orig_hw = [(orig_h, orig_w)] + res = self.inst_interactive_predictor.predict(**kwargs) + self.inst_interactive_predictor._features = None + self.inst_interactive_predictor._is_image_set = False + return res + + def predict_inst_batch( + self, + inference_state, + *args, + **kwargs, + ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: + backbone_out = inference_state["backbone_out"]["sam2_backbone_out"] + ( + _, + vision_feats, + _, + _, + ) = self.inst_interactive_predictor.model._prepare_backbone_features( + backbone_out + ) + # Add no_mem_embed, which is added to the lowest res feat. map during training on videos + vision_feats[-1] = ( + vision_feats[-1] + self.inst_interactive_predictor.model.no_mem_embed + ) + batch_size = vision_feats[-1].shape[1] + orig_heights, orig_widths = ( + inference_state["original_heights"], + inference_state["original_widths"], + ) + assert ( + batch_size == len(orig_heights) == len(orig_widths) + ), f"Batch size mismatch in predict_inst_batch. Got {batch_size}, {len(orig_heights)}, {len(orig_widths)}" + feats = [ + feat.permute(1, 2, 0).view(batch_size, -1, *feat_size) + for feat, feat_size in zip( + vision_feats[::-1], self.inst_interactive_predictor._bb_feat_sizes[::-1] + ) + ][::-1] + self.inst_interactive_predictor._features = { + "image_embed": feats[-1], + "high_res_feats": feats[:-1], + } + self.inst_interactive_predictor._is_image_set = True + self.inst_interactive_predictor._is_batch = True + self.inst_interactive_predictor._orig_hw = [ + (orig_h, orig_w) for orig_h, orig_w in zip(orig_heights, orig_widths) + ] + res = self.inst_interactive_predictor.predict_batch(*args, **kwargs) + self.inst_interactive_predictor._features = None + self.inst_interactive_predictor._is_image_set = False + self.inst_interactive_predictor._is_batch = False + return res + + +class Sam3ImageOnVideoMultiGPU(Sam3Image): + def __init__( + self, *args, async_all_gather=True, gather_backbone_out=None, **kwargs + ): + super().__init__(*args, **kwargs) + self.rank = int(os.getenv("RANK", "0")) + self.world_size = int(os.getenv("WORLD_SIZE", "1")) + self.async_all_gather = async_all_gather + + # if gather_backbone is not set, default to gathering only for `SAM3VLBackbone` + if gather_backbone_out is None: + gather_backbone_out = isinstance(self.backbone, SAM3VLBackbone) + self.gather_backbone_out = gather_backbone_out + + def forward_video_grounding_multigpu( + self, + backbone_out, + find_inputs, + geometric_prompt: Prompt, + frame_idx, + num_frames, + # `multigpu_buffer` is a dict to cache detector's outputs in a chunk between different calls + multigpu_buffer, + track_in_reverse=False, + # whether to also return the SAM2 backbone features + return_sam2_backbone_feats=False, + # whether to perform NMS and suppress the scores of those detections removed by NMS + run_nms=False, + nms_prob_thresh=None, + nms_iou_thresh=None, + **kwargs, + ): + """ + Compute the detector's detection outputs in a distributed manner, where all GPUs process + a chunk of frames (equal to the number of GPUs) at once and store them in cache. + """ + # Step 1: fetch the detector outputs in the current chunk from buffer + frame_idx_curr_b = frame_idx - frame_idx % self.world_size + frame_idx_curr_e = min(frame_idx_curr_b + self.world_size, num_frames) + # in case the current frame's detection results are not in the buffer yet, build the current chunk + # (this should only happen on the first chunk, since we are also building the next chunk below) + if frame_idx not in multigpu_buffer: + with torch.profiler.record_function("build_multigpu_buffer_next_chunk1"): + self._build_multigpu_buffer_next_chunk( + backbone_out=backbone_out, + find_inputs=find_inputs, + geometric_prompt=geometric_prompt, + frame_idx_begin=frame_idx_curr_b, + frame_idx_end=frame_idx_curr_e, + num_frames=num_frames, + multigpu_buffer=multigpu_buffer, + run_nms=run_nms, + nms_prob_thresh=nms_prob_thresh, + nms_iou_thresh=nms_iou_thresh, + ) + + # read out the current frame's results from `multigpu_buffer` + out = {} + for k, (v, handle) in multigpu_buffer[frame_idx].items(): + if k.startswith("sam2_backbone_") and not return_sam2_backbone_feats: + continue + if handle is not None: + handle.wait() # wait for async all-gather to finish + out[k] = v + + # Step 2: remove detection outputs of the previous chunk from cache to save GPU memory + if not track_in_reverse and frame_idx_curr_b - self.world_size >= 0: + frame_idx_prev_e = frame_idx_curr_b + frame_idx_prev_b = frame_idx_curr_b - self.world_size + elif track_in_reverse and frame_idx_curr_e < num_frames: + frame_idx_prev_b = frame_idx_curr_e + frame_idx_prev_e = min(frame_idx_prev_b + self.world_size, num_frames) + else: + frame_idx_prev_b = frame_idx_prev_e = None + if frame_idx_prev_b is not None: + for frame_idx_rm in range(frame_idx_prev_b, frame_idx_prev_e): + multigpu_buffer.pop(frame_idx_rm, None) + + # Step 3: compute and cache detection outputs of the next chunk ahead of time + # (so that we can overlap computation with all-gather transfer) + if not track_in_reverse and frame_idx_curr_e < num_frames: + frame_idx_next_b = frame_idx_curr_e + frame_idx_next_e = min(frame_idx_next_b + self.world_size, num_frames) + elif track_in_reverse and frame_idx_curr_b - self.world_size >= 0: + frame_idx_next_e = frame_idx_curr_b + frame_idx_next_b = frame_idx_curr_b - self.world_size + else: + frame_idx_next_b = frame_idx_next_e = None + if frame_idx_next_b is not None and frame_idx_next_b not in multigpu_buffer: + with torch.profiler.record_function("build_multigpu_buffer_next_chunk2"): + self._build_multigpu_buffer_next_chunk( + backbone_out=backbone_out, + find_inputs=find_inputs, + geometric_prompt=geometric_prompt, + frame_idx_begin=frame_idx_next_b, + frame_idx_end=frame_idx_next_e, + num_frames=num_frames, + multigpu_buffer=multigpu_buffer, + run_nms=run_nms, + nms_prob_thresh=nms_prob_thresh, + nms_iou_thresh=nms_iou_thresh, + ) + + return out, backbone_out + + def _build_multigpu_buffer_next_chunk( + self, + backbone_out, + find_inputs, + geometric_prompt: Prompt, + frame_idx_begin, + frame_idx_end, + num_frames, + multigpu_buffer, + run_nms=False, + nms_prob_thresh=None, + nms_iou_thresh=None, + ): + """Compute detection outputs on a chunk of frames and store their results in multigpu_buffer.""" + # each GPU computes detections on one frame in the chunk (in a round-robin manner) + frame_idx_local_gpu = min(frame_idx_begin + self.rank, frame_idx_end - 1) + # `forward_grounding` (from base class `Sam3ImageOnVideo`) runs the detector on a single frame + with torch.profiler.record_function("forward_grounding"): + out_local = self.forward_grounding( + backbone_out=backbone_out, + find_input=find_inputs[frame_idx_local_gpu], + find_target=None, + geometric_prompt=geometric_prompt, + ) + if run_nms: + with torch.profiler.record_function("nms_masks"): + # run NMS as a post-processing step on top of the detection outputs + assert nms_prob_thresh is not None and nms_iou_thresh is not None + pred_probs = out_local["pred_logits"].squeeze(-1).sigmoid() + pred_masks = out_local["pred_masks"] + # loop over text prompts (not an overhead for demo where there's only 1 prompt) + for prompt_idx in range(pred_probs.size(0)): + keep = nms_masks( + pred_probs=pred_probs[prompt_idx], + pred_masks=pred_masks[prompt_idx], + prob_threshold=nms_prob_thresh, + iou_threshold=nms_iou_thresh, + ) + # set a very low threshold for those detections removed by NMS + out_local["pred_logits"][prompt_idx, :, 0] -= 1e4 * (~keep).float() + + if self.gather_backbone_out: + # gather the SAM 2 backbone features across GPUs + feats = out_local["prev_encoder_out"]["backbone_out"]["sam2_backbone_out"] + assert len(feats["backbone_fpn"]) == 3 # SAM2 backbone always have 3 levels + # cast the SAM2 backbone features to bfloat16 for all-gather (this is usually + # a no-op, SAM2 backbone features are likely already in bfloat16 due to AMP) + backbone_fpn_bf16 = [x.to(torch.bfloat16) for x in feats["backbone_fpn"]] + fpn0, fpn_handle0 = self._gather_tensor(backbone_fpn_bf16[0]) + fpn1, fpn_handle1 = self._gather_tensor(backbone_fpn_bf16[1]) + fpn2, fpn_handle2 = self._gather_tensor(backbone_fpn_bf16[2]) + # vision_pos_enc is the same on all frames, so no need to all-gather them + vision_pos_enc = feats["vision_pos_enc"] + + # trim the detector output to only include the necessary keys + out_local = { + "pred_logits": out_local["pred_logits"], + "pred_boxes": out_local["pred_boxes"], + "pred_boxes_xyxy": out_local["pred_boxes_xyxy"], + "pred_masks": out_local["pred_masks"], + } + + # gather the results: after this step, each GPU will receive detector outputs on + # all frames in the chunk and store them in `multigpu_buffer` + out_gathered = {k: self._gather_tensor(v) for k, v in out_local.items()} + for rank in range(self.world_size): + frame_idx_to_save = frame_idx_begin + rank + if frame_idx_to_save >= num_frames: + continue + frame_buffer = { + k: (v[rank], handle) for k, (v, handle) in out_gathered.items() + } + if self.gather_backbone_out: + # also add gathered SAM 2 backbone features to frame_buffer + frame_buffer["tracker_backbone_fpn_0"] = (fpn0[rank], fpn_handle0) + frame_buffer["tracker_backbone_fpn_1"] = (fpn1[rank], fpn_handle1) + frame_buffer["tracker_backbone_fpn_2"] = (fpn2[rank], fpn_handle2) + frame_buffer["tracker_backbone_pos_enc"] = (vision_pos_enc, None) + + multigpu_buffer[frame_idx_to_save] = frame_buffer + + def _gather_tensor(self, x): + if self.world_size == 1: + return [x], None + + async_op = self.async_all_gather + # here `.contiguous()` is required -- otherwise NCCL all_gather + # sometimes gives wrong results + x = x.contiguous() # ensure contiguous memory for NCCL + output_list = [torch.empty_like(x) for _ in range(self.world_size)] + handle = torch.distributed.all_gather(output_list, x, async_op=async_op) + return output_list, handle diff --git a/src/nn/segearth_ov3/sam3/model/sam3_image_processor.py b/src/nn/segearth_ov3/sam3/model/sam3_image_processor.py new file mode 100644 index 0000000..58f87f8 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_image_processor.py @@ -0,0 +1,236 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +from typing import Dict, List + +import numpy as np +import PIL +import torch + +from sam3.model import box_ops + +from sam3.model.data_misc import FindStage, interpolate +from torchvision.transforms import v2 + + +class Sam3Processor: + """ """ + + def __init__(self, model, resolution=1008, device="cuda", confidence_threshold=0.5): + self.model = model + self.resolution = resolution + self.device = device + self.transform = v2.Compose( + [ + v2.ToDtype(torch.uint8, scale=True), + v2.Resize(size=(resolution, resolution)), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ] + ) + self.confidence_threshold = confidence_threshold + + self.find_stage = FindStage( + img_ids=torch.tensor([0], device=device, dtype=torch.long), + text_ids=torch.tensor([0], device=device, dtype=torch.long), + input_boxes=None, + input_boxes_mask=None, + input_boxes_label=None, + input_points=None, + input_points_mask=None, + ) + + @torch.inference_mode() + def set_image(self, image, state=None): + """Sets the image on which we want to do predictions.""" + if state is None: + state = {} + + if isinstance(image, PIL.Image.Image): + width, height = image.size + elif isinstance(image, (torch.Tensor, np.ndarray)): + height, width = image.shape[-2:] + else: + raise ValueError("Image must be a PIL image or a tensor") + + image = v2.functional.to_image(image).to(self.device) + image = self.transform(image).unsqueeze(0) + + state["original_height"] = height + state["original_width"] = width + state["backbone_out"] = self.model.backbone.forward_image(image) + inst_interactivity_en = self.model.inst_interactive_predictor is not None # False + if inst_interactivity_en and "sam2_backbone_out" in state["backbone_out"]: + sam2_backbone_out = state["backbone_out"]["sam2_backbone_out"] + sam2_backbone_out["backbone_fpn"][0] = ( + self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s0( + sam2_backbone_out["backbone_fpn"][0] + ) + ) + sam2_backbone_out["backbone_fpn"][1] = ( + self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s1( + sam2_backbone_out["backbone_fpn"][1] + ) + ) + return state + + @torch.inference_mode() + def set_image_batch(self, images: List[np.ndarray], state=None): + """Sets the image batch on which we want to do predictions.""" + if state is None: + state = {} + + if not isinstance(images, list): + raise ValueError("Images must be a list of PIL images or tensors") + assert len(images) > 0, "Images list must not be empty" + assert isinstance( + images[0], PIL.Image.Image + ), "Images must be a list of PIL images" + + state["original_heights"] = [image.height for image in images] + state["original_widths"] = [image.width for image in images] + + images = [ + self.transform(v2.functional.to_image(image).to(self.device)) + for image in images + ] + images = torch.stack(images, dim=0) + state["backbone_out"] = self.model.backbone.forward_image(images) + inst_interactivity_en = self.model.inst_interactive_predictor is not None + if inst_interactivity_en and "sam2_backbone_out" in state["backbone_out"]: + sam2_backbone_out = state["backbone_out"]["sam2_backbone_out"] + sam2_backbone_out["backbone_fpn"][0] = ( + self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s0( + sam2_backbone_out["backbone_fpn"][0] + ) + ) + sam2_backbone_out["backbone_fpn"][1] = ( + self.model.inst_interactive_predictor.model.sam_mask_decoder.conv_s1( + sam2_backbone_out["backbone_fpn"][1] + ) + ) + return state + + @torch.inference_mode() + def set_text_prompt(self, prompt: str, state: Dict): + """Sets the text prompt and run the inference""" + + if "backbone_out" not in state: + raise ValueError("You must call set_image before set_text_prompt") + + text_outputs = self.model.backbone.forward_text([prompt], device=self.device) + # text_outputs['language_features']: [32, 1, 256] + # text_outputs['language_mask']: [1, 32] + # text_outputs['language_embeds']: [32, 1, 1024] + + # will erase the previous text prompt if any + state["backbone_out"].update(text_outputs) + if "geometric_prompt" not in state: + state["geometric_prompt"] = self.model._get_dummy_prompt() + + return self._forward_grounding(state) + + @torch.inference_mode() + def add_geometric_prompt(self, box: List, label: bool, state: Dict): + """Adds a box prompt and run the inference. + The image needs to be set, but not necessarily the text prompt. + The box is assumed to be in [center_x, center_y, width, height] format and normalized in [0, 1] range. + The label is True for a positive box, False for a negative box. + """ + if "backbone_out" not in state: + raise ValueError("You must call set_image before set_text_prompt") + + if "language_features" not in state["backbone_out"]: + # Looks like we don't have a text prompt yet. This is allowed, but we need to set the text prompt to "visual" for the model to rely only on the geometric prompt + dummy_text_outputs = self.model.backbone.forward_text( + ["visual"], device=self.device + ) + state["backbone_out"].update(dummy_text_outputs) + + if "geometric_prompt" not in state: + state["geometric_prompt"] = self.model._get_dummy_prompt() + + # adding a batch and sequence dimension + boxes = torch.tensor(box, device=self.device, dtype=torch.float32).view(1, 1, 4) + labels = torch.tensor([label], device=self.device, dtype=torch.bool).view(1, 1) + state["geometric_prompt"].append_boxes(boxes, labels) + + return self._forward_grounding(state) + + def reset_all_prompts(self, state: Dict): + """Removes all the prompts and results""" + if "backbone_out" in state: + backbone_keys_to_del = [ + "language_features", + "language_mask", + "language_embeds", + ] + for key in backbone_keys_to_del: + if key in state["backbone_out"]: + del state["backbone_out"][key] + + keys_to_del = ["geometric_prompt", "boxes", "masks", "masks_logits", "scores"] + for key in keys_to_del: + if key in state: + del state[key] + + @torch.inference_mode() + def set_confidence_threshold(self, threshold: float, state=None): + """Sets the confidence threshold for the masks""" + self.confidence_threshold = threshold + if state is not None and "boxes" in state: + # we need to filter the boxes again + # In principle we could do this more efficiently since we would only need + # to rerun the heads. But this is simpler and not too inefficient + return self._forward_grounding(state) + return state + + @torch.inference_mode() + def _forward_grounding(self, state: Dict): + outputs = self.model.forward_grounding( + backbone_out=state["backbone_out"], + find_input=self.find_stage, + geometric_prompt=state["geometric_prompt"], + find_target=None, + ) + + out_bbox = outputs["pred_boxes"] + out_logits = outputs["pred_logits"] + out_masks = outputs["pred_masks"] + out_probs = out_logits.sigmoid() + presence_score = outputs["presence_logit_dec"].sigmoid().unsqueeze(1) + out_probs = (out_probs * presence_score).squeeze(-1) + + keep = out_probs > self.confidence_threshold + out_probs = out_probs[keep] + out_masks = out_masks[keep] + out_bbox = out_bbox[keep] + + # convert to [x0, y0, x1, y1] format + boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) + + img_h = state["original_height"] + img_w = state["original_width"] + scale_fct = torch.tensor([img_w, img_h, img_w, img_h]).to(self.device) + boxes = boxes * scale_fct[None, :] + + out_masks = interpolate( + out_masks.unsqueeze(1), + (img_h, img_w), + mode="bilinear", + align_corners=False, + ).sigmoid() + + out_semantic_masks = interpolate( + outputs["semantic_seg"], + (img_h, img_w), + mode="bilinear", + align_corners=False, + ).sigmoid() + + state["masks_logits"] = out_masks + state["masks"] = out_masks > 0.5 + state["boxes"] = boxes + state["scores"] = out_probs + state["semantic_mask_logits"] = out_semantic_masks # for SS + state["presence_score"] = presence_score.squeeze().squeeze() + state["object_score"] = out_probs + return state diff --git a/src/nn/segearth_ov3/sam3/model/sam3_tracker_base.py b/src/nn/segearth_ov3/sam3/model/sam3_tracker_base.py new file mode 100644 index 0000000..8f66d75 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_tracker_base.py @@ -0,0 +1,1189 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging + +import torch +import torch.nn.functional as F + +from sam3.model.memory import SimpleMaskEncoder + +from sam3.model.sam3_tracker_utils import get_1d_sine_pe, select_closest_cond_frames + +from sam3.sam.mask_decoder import MaskDecoder, MLP +from sam3.sam.prompt_encoder import PromptEncoder +from sam3.sam.transformer import TwoWayTransformer +# from sam3.train.data.collator import BatchedDatapoint +from sam3.model.data_misc import BatchedDatapoint + +try: + from timm.layers import trunc_normal_ +except ModuleNotFoundError: + # compatibility for older timm versions + from timm.models.layers import trunc_normal_ + +# a large negative value as a placeholder score for missing objects +NO_OBJ_SCORE = -1024.0 + + +class Sam3TrackerBase(torch.nn.Module): + def __init__( + self, + backbone, + transformer, + maskmem_backbone, + num_maskmem=7, # default 1 input frame + 6 previous frames as in CAE + image_size=1008, + backbone_stride=14, # stride of the image backbone output + # The maximum number of conditioning frames to participate in the memory attention (-1 means no limit; if there are more conditioning frames than this limit, + # we only cross-attend to the temporally closest `max_cond_frames_in_attn` conditioning frames in the encoder when tracking each frame). This gives the model + # a temporal locality when handling a large number of annotated frames (since closer frames should be more important) and also avoids GPU OOM. + max_cond_frames_in_attn=-1, + # Whether to always keep the first conditioning frame in case we exceed the maximum number of conditioning frames allowed + keep_first_cond_frame=False, + # whether to output multiple (3) masks for the first click on initial conditioning frames + multimask_output_in_sam=False, + # the minimum and maximum number of clicks to use multimask_output_in_sam (only relevant when `multimask_output_in_sam=True`; + # default is 1 for both, meaning that only the first click gives multimask output; also note that a box counts as two points) + multimask_min_pt_num=1, + multimask_max_pt_num=1, + # whether to also use multimask output for tracking (not just for the first click on initial conditioning frames; only relevant when `multimask_output_in_sam=True`) + multimask_output_for_tracking=False, + # whether to forward image features per frame (as it's being tracked) during evaluation, instead of forwarding image features + # of all frames at once. This avoids backbone OOM errors on very long videos in evaluation, but could be slightly slower. + forward_backbone_per_frame_for_eval=False, + # The memory bank's temporal stride during evaluation (i.e. the `r` parameter in XMem and Cutie; XMem and Cutie use r=5). + # For r>1, the (self.num_maskmem - 1) non-conditioning memory frames consist of + # (self.num_maskmem - 2) nearest frames from every r-th frames, plus the last frame. + memory_temporal_stride_for_eval=1, + # whether to offload outputs to CPU memory during evaluation, to avoid GPU OOM on very long videos or very large resolutions or too many objects + # (it's recommended to use `forward_backbone_per_frame_for_eval=True` first before setting this option to True) + offload_output_to_cpu_for_eval=False, + # whether to trim the output of past non-conditioning frames (num_maskmem frames before the current frame) during evaluation + # (this helps save GPU or CPU memory on very long videos for semi-supervised VOS eval, where only the first frame receives prompts) + trim_past_non_cond_mem_for_eval=False, + # whether to apply non-overlapping constraints on the object masks in the memory encoder during evaluation (to avoid/alleviate superposing masks) + non_overlap_masks_for_mem_enc=False, + # the maximum number of object pointers from other frames in encoder cross attention + max_obj_ptrs_in_encoder=16, + # extra arguments used to construct the SAM mask decoder; if not None, it should be a dict of kwargs to be passed into `MaskDecoder` class. + sam_mask_decoder_extra_args=None, + # whether to compile all the model compoents + compile_all_components=False, + # select the frame with object existence + use_memory_selection=False, + # when using memory selection, the threshold to determine if the frame is good + mf_threshold=0.01, + ): + super().__init__() + + # Part 1: the image backbone + self.backbone = backbone + self.num_feature_levels = 3 + self.max_obj_ptrs_in_encoder = max_obj_ptrs_in_encoder + # A conv layer to downsample the GT mask prompt to stride 4 (the same stride as + # low-res SAM mask logits) and to change its scales from 0~1 to SAM logit scale, + # so that it can be fed into the SAM mask decoder to generate a pointer. + self.mask_downsample = torch.nn.Conv2d(1, 1, kernel_size=4, stride=4) + + # Part 2: encoder-only transformer to fuse current frame's visual features + # with memories from past frames + assert transformer.decoder is None, "transformer should be encoder-only" + self.transformer = transformer + self.hidden_dim = transformer.d_model + + # Part 3: memory encoder for the previous frame's outputs + self.maskmem_backbone = maskmem_backbone + self.mem_dim = self.hidden_dim + if hasattr(self.maskmem_backbone, "out_proj") and hasattr( + self.maskmem_backbone.out_proj, "weight" + ): + # if there is compression of memories along channel dim + self.mem_dim = self.maskmem_backbone.out_proj.weight.shape[0] + self.num_maskmem = num_maskmem # Number of memories accessible + + # Temporal encoding of the memories + self.maskmem_tpos_enc = torch.nn.Parameter( + torch.zeros(num_maskmem, 1, 1, self.mem_dim) + ) + trunc_normal_(self.maskmem_tpos_enc, std=0.02) + + # a single token to indicate no memory embedding from previous frames + self.no_mem_embed = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + self.no_mem_pos_enc = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + trunc_normal_(self.no_mem_embed, std=0.02) + trunc_normal_(self.no_mem_pos_enc, std=0.02) + # Apply sigmoid to the output raw mask logits (to turn them from + # range (-inf, +inf) to range (0, 1)) before feeding them into the memory encoder + self.sigmoid_scale_for_mem_enc = 20.0 + self.sigmoid_bias_for_mem_enc = -10.0 + self.non_overlap_masks_for_mem_enc = non_overlap_masks_for_mem_enc + self.memory_temporal_stride_for_eval = memory_temporal_stride_for_eval + # On frames with mask input, whether to directly output the input mask without + # using a SAM prompt encoder + mask decoder + self.multimask_output_in_sam = multimask_output_in_sam + self.multimask_min_pt_num = multimask_min_pt_num + self.multimask_max_pt_num = multimask_max_pt_num + self.multimask_output_for_tracking = multimask_output_for_tracking + + # Part 4: SAM-style prompt encoder (for both mask and point inputs) + # and SAM-style mask decoder for the final mask output + self.image_size = image_size + self.backbone_stride = backbone_stride + self.low_res_mask_size = self.image_size // self.backbone_stride * 4 + # we resize the mask if it doesn't match `self.input_mask_size` (which is always 4x + # the low-res mask size, regardless of the actual input image size); this is because + # `_use_mask_as_output` always downsamples the input masks by 4x + self.input_mask_size = self.low_res_mask_size * 4 + self.forward_backbone_per_frame_for_eval = forward_backbone_per_frame_for_eval + self.offload_output_to_cpu_for_eval = offload_output_to_cpu_for_eval + self.trim_past_non_cond_mem_for_eval = trim_past_non_cond_mem_for_eval + self.sam_mask_decoder_extra_args = sam_mask_decoder_extra_args + self.no_obj_ptr = torch.nn.Parameter(torch.zeros(1, self.hidden_dim)) + trunc_normal_(self.no_obj_ptr, std=0.02) + self.no_obj_embed_spatial = torch.nn.Parameter(torch.zeros(1, self.mem_dim)) + trunc_normal_(self.no_obj_embed_spatial, std=0.02) + + self._build_sam_heads() + self.max_cond_frames_in_attn = max_cond_frames_in_attn + self.keep_first_cond_frame = keep_first_cond_frame + + # Use frame filtering according to SAM2Long + self.use_memory_selection = use_memory_selection + self.mf_threshold = mf_threshold + + # Compile all components of the model + self.compile_all_components = compile_all_components + if self.compile_all_components: + self._compile_all_components() + + @property + def device(self): + return next(self.parameters()).device + + def _get_tpos_enc(self, rel_pos_list, device, max_abs_pos=None, dummy=False): + if dummy: + return torch.zeros(len(rel_pos_list), self.mem_dim, device=device) + + t_diff_max = max_abs_pos - 1 if max_abs_pos is not None else 1 + pos_enc = ( + torch.tensor(rel_pos_list).pin_memory().to(device=device, non_blocking=True) + / t_diff_max + ) + tpos_dim = self.hidden_dim + pos_enc = get_1d_sine_pe(pos_enc, dim=tpos_dim) + pos_enc = self.obj_ptr_tpos_proj(pos_enc) + + return pos_enc + + def _build_sam_heads(self): + """Build SAM-style prompt encoder and mask decoder.""" + self.sam_prompt_embed_dim = self.hidden_dim + self.sam_image_embedding_size = self.image_size // self.backbone_stride + + # build PromptEncoder and MaskDecoder from SAM + # (their hyperparameters like `mask_in_chans=16` are from SAM code) + self.sam_prompt_encoder = PromptEncoder( + embed_dim=self.sam_prompt_embed_dim, + image_embedding_size=( + self.sam_image_embedding_size, + self.sam_image_embedding_size, + ), + input_image_size=(self.image_size, self.image_size), + mask_in_chans=16, + ) + self.sam_mask_decoder = MaskDecoder( + num_multimask_outputs=3, + transformer=TwoWayTransformer( + depth=2, + embedding_dim=self.sam_prompt_embed_dim, + mlp_dim=2048, + num_heads=8, + ), + transformer_dim=self.sam_prompt_embed_dim, + iou_head_depth=3, + iou_head_hidden_dim=256, + use_high_res_features=True, + iou_prediction_use_sigmoid=True, + pred_obj_scores=True, + pred_obj_scores_mlp=True, + use_multimask_token_for_obj_ptr=True, + **(self.sam_mask_decoder_extra_args or {}), + ) + # a linear projection on SAM output tokens to turn them into object pointers + self.obj_ptr_proj = torch.nn.Linear(self.hidden_dim, self.hidden_dim) + self.obj_ptr_proj = MLP(self.hidden_dim, self.hidden_dim, self.hidden_dim, 3) + # a linear projection on temporal positional encoding in object pointers to + # avoid potential interference with spatial positional encoding + self.obj_ptr_tpos_proj = torch.nn.Linear(self.hidden_dim, self.mem_dim) + + def _forward_sam_heads( + self, + backbone_features, + point_inputs=None, + mask_inputs=None, + high_res_features=None, + multimask_output=False, + gt_masks=None, + ): + """ + Forward SAM prompt encoders and mask heads. + + Inputs: + - backbone_features: image features of [B, C, H, W] shape + - point_inputs: a dictionary with "point_coords" and "point_labels", where + 1) "point_coords" has [B, P, 2] shape and float32 dtype and contains the + absolute pixel-unit coordinate in (x, y) format of the P input points + 2) "point_labels" has shape [B, P] and int32 dtype, where 1 means + positive clicks, 0 means negative clicks, and -1 means padding + - mask_inputs: a mask of [B, 1, H*16, W*16] shape, float or bool, with the + same spatial size as the image. + - high_res_features: either 1) None or 2) or a list of length 2 containing + two feature maps of [B, C, 4*H, 4*W] and [B, C, 2*H, 2*W] shapes respectively, + which will be used as high-resolution feature maps for SAM decoder. + - multimask_output: if it's True, we output 3 candidate masks and their 3 + corresponding IoU estimates, and if it's False, we output only 1 mask and + its corresponding IoU estimate. + + Outputs: + - low_res_multimasks: [B, M, H*4, W*4] shape (where M = 3 if + `multimask_output=True` and M = 1 if `multimask_output=False`), the SAM + output mask logits (before sigmoid) for the low-resolution masks, with 4x + the resolution (1/4 stride) of the input backbone_features. + - high_res_multimasks: [B, M, H*16, W*16] shape (where M = 3 + if `multimask_output=True` and M = 1 if `multimask_output=False`), + upsampled from the low-resolution masks, with shape size as the image + (stride is 1 pixel). + - ious, [B, M] shape, where (where M = 3 if `multimask_output=True` and M = 1 + if `multimask_output=False`), the estimated IoU of each output mask. + - low_res_masks: [B, 1, H*4, W*4] shape, the best mask in `low_res_multimasks`. + If `multimask_output=True`, it's the mask with the highest IoU estimate. + If `multimask_output=False`, it's the same as `low_res_multimasks`. + - high_res_masks: [B, 1, H*16, W*16] shape, the best mask in `high_res_multimasks`. + If `multimask_output=True`, it's the mask with the highest IoU estimate. + If `multimask_output=False`, it's the same as `high_res_multimasks`. + - obj_ptr: [B, C] shape, the object pointer vector for the output mask, extracted + based on the output token from the SAM mask decoder. + """ + B = backbone_features.size(0) + device = backbone_features.device + assert backbone_features.size(1) == self.sam_prompt_embed_dim + assert backbone_features.size(2) == self.sam_image_embedding_size + assert backbone_features.size(3) == self.sam_image_embedding_size + + # a) Handle point prompts + if point_inputs is not None: + sam_point_coords = point_inputs["point_coords"] + sam_point_labels = point_inputs["point_labels"] + assert sam_point_coords.size(0) == B and sam_point_labels.size(0) == B + else: + # If no points are provide, pad with an empty point (with label -1) + sam_point_coords = torch.zeros(B, 1, 2, device=device) + sam_point_labels = -torch.ones(B, 1, dtype=torch.int32, device=device) + + # b) Handle mask prompts + if mask_inputs is not None: + # If mask_inputs is provided, downsize it into low-res mask input if needed + # and feed it as a dense mask prompt into the SAM mask encoder + assert len(mask_inputs.shape) == 4 and mask_inputs.shape[:2] == (B, 1) + if mask_inputs.shape[-2:] != self.sam_prompt_encoder.mask_input_size: + sam_mask_prompt = F.interpolate( + mask_inputs.float(), + size=self.sam_prompt_encoder.mask_input_size, + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for downsampling + ) + else: + sam_mask_prompt = mask_inputs + else: + # Otherwise, simply feed None (and SAM's prompt encoder will add + # a learned `no_mask_embed` to indicate no mask input in this case). + sam_mask_prompt = None + + sparse_embeddings, dense_embeddings = self.sam_prompt_encoder( + points=(sam_point_coords, sam_point_labels), + boxes=None, + masks=sam_mask_prompt, + ) + # Clone image_pe and the outputs of sam_prompt_encoder + # to enable compilation + sparse_embeddings = self._maybe_clone(sparse_embeddings) + dense_embeddings = self._maybe_clone(dense_embeddings) + image_pe = self._maybe_clone(self.sam_prompt_encoder.get_dense_pe()) + with torch.profiler.record_function("sam_mask_decoder"): + ( + low_res_multimasks, + ious, + sam_output_tokens, + object_score_logits, + ) = self.sam_mask_decoder( + image_embeddings=backbone_features, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + repeat_image=False, # the image is already batched + high_res_features=high_res_features, + ) + # Clone the output of sam_mask_decoder + # to enable compilation + low_res_multimasks = self._maybe_clone(low_res_multimasks) + ious = self._maybe_clone(ious) + sam_output_tokens = self._maybe_clone(sam_output_tokens) + object_score_logits = self._maybe_clone(object_score_logits) + + if self.training and self.teacher_force_obj_scores_for_mem: + # we use gt to detect if there is an object or not to + # select no obj ptr and use an empty mask for spatial memory + is_obj_appearing = torch.any(gt_masks.float().flatten(1) > 0, dim=1) + is_obj_appearing = is_obj_appearing[..., None] + else: + is_obj_appearing = object_score_logits > 0 + + # Mask used for spatial memories is always a *hard* choice between obj and no obj, + # consistent with the actual mask prediction + low_res_multimasks = torch.where( + is_obj_appearing[:, None, None], + low_res_multimasks, + NO_OBJ_SCORE, + ) + + # convert masks from possibly bfloat16 (or float16) to float32 + # (older PyTorch versions before 2.1 don't support `interpolate` on bf16) + low_res_multimasks = low_res_multimasks.float() + high_res_multimasks = F.interpolate( + low_res_multimasks, + size=(self.image_size, self.image_size), + mode="bilinear", + align_corners=False, + ) + + sam_output_token = sam_output_tokens[:, 0] + if multimask_output: + # take the best mask prediction (with the highest IoU estimation) + best_iou_inds = torch.argmax(ious, dim=-1) + batch_inds = torch.arange(B, device=device) + low_res_masks = low_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1) + high_res_masks = high_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1) + if sam_output_tokens.size(1) > 1: + sam_output_token = sam_output_tokens[batch_inds, best_iou_inds] + else: + low_res_masks, high_res_masks = low_res_multimasks, high_res_multimasks + + # Extract object pointer from the SAM output token (with occlusion handling) + obj_ptr = self.obj_ptr_proj(sam_output_token) + lambda_is_obj_appearing = is_obj_appearing.float() + + obj_ptr = lambda_is_obj_appearing * obj_ptr + obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr + + return ( + low_res_multimasks, + high_res_multimasks, + ious, + low_res_masks, + high_res_masks, + obj_ptr, + object_score_logits, + ) + + def _use_mask_as_output(self, backbone_features, high_res_features, mask_inputs): + """ + Directly turn binary `mask_inputs` into a output mask logits without using SAM. + (same input and output shapes as in _forward_sam_heads above). + """ + # Use -10/+10 as logits for neg/pos pixels (very close to 0/1 in prob after sigmoid). + out_scale, out_bias = 20.0, -10.0 # sigmoid(-10.0)=4.5398e-05 + mask_inputs_float = mask_inputs.float() + high_res_masks = mask_inputs_float * out_scale + out_bias + low_res_masks = F.interpolate( + high_res_masks, + size=( + high_res_masks.size(-2) // self.backbone_stride * 4, + high_res_masks.size(-1) // self.backbone_stride * 4, + ), + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for downsampling + ) + # a dummy IoU prediction of all 1's under mask input + ious = mask_inputs.new_ones(mask_inputs.size(0), 1).float() + # produce an object pointer using the SAM decoder from the mask input + _, _, _, _, _, obj_ptr, _ = self._forward_sam_heads( + backbone_features=backbone_features, + mask_inputs=self.mask_downsample(mask_inputs_float), + high_res_features=high_res_features, + gt_masks=mask_inputs, + ) + # In this method, we are treating mask_input as output, e.g. using it directly to create spatial mem; + # Below, we follow the same design axiom to use mask_input to decide if obj appears or not instead of relying + # on the object_scores from the SAM decoder. + is_obj_appearing = torch.any(mask_inputs.flatten(1).float() > 0.0, dim=1) + is_obj_appearing = is_obj_appearing[..., None] + lambda_is_obj_appearing = is_obj_appearing.float() + object_score_logits = out_scale * lambda_is_obj_appearing + out_bias + obj_ptr = lambda_is_obj_appearing * obj_ptr + obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr + + return ( + low_res_masks, + high_res_masks, + ious, + low_res_masks, + high_res_masks, + obj_ptr, + object_score_logits, + ) + + def forward(self, input: BatchedDatapoint, is_inference=False): + raise NotImplementedError( + "Please use the corresponding methods in SAM3VideoPredictor for inference." + "See examples/sam3_dense_video_tracking.ipynb for an inference example." + ) + + def forward_image(self, img_batch): + """Get the image feature on the input batch.""" + # This line is the only change from the parent class + # to use the SAM3 backbone instead of the SAM2 backbone. + backbone_out = self.backbone.forward_image(img_batch)["sam2_backbone_out"] + # precompute projected level 0 and level 1 features in SAM decoder + # to avoid running it again on every SAM click + backbone_out["backbone_fpn"][0] = self.sam_mask_decoder.conv_s0( + backbone_out["backbone_fpn"][0] + ) + backbone_out["backbone_fpn"][1] = self.sam_mask_decoder.conv_s1( + backbone_out["backbone_fpn"][1] + ) + # Clone to help torch.compile + for i in range(len(backbone_out["backbone_fpn"])): + backbone_out["backbone_fpn"][i] = self._maybe_clone( + backbone_out["backbone_fpn"][i] + ) + backbone_out["vision_pos_enc"][i] = self._maybe_clone( + backbone_out["vision_pos_enc"][i] + ) + return backbone_out + + def _prepare_backbone_features(self, backbone_out): + """Prepare and flatten visual features (same as in MDETR_API model).""" + backbone_out = backbone_out.copy() + assert len(backbone_out["backbone_fpn"]) == len(backbone_out["vision_pos_enc"]) + assert len(backbone_out["backbone_fpn"]) >= self.num_feature_levels + + feature_maps = backbone_out["backbone_fpn"][-self.num_feature_levels :] + vision_pos_embeds = backbone_out["vision_pos_enc"][-self.num_feature_levels :] + + feat_sizes = [(x.shape[-2], x.shape[-1]) for x in vision_pos_embeds] + # flatten NxCxHxW to HWxNxC + vision_feats = [x.flatten(2).permute(2, 0, 1) for x in feature_maps] + vision_pos_embeds = [x.flatten(2).permute(2, 0, 1) for x in vision_pos_embeds] + + return backbone_out, vision_feats, vision_pos_embeds, feat_sizes + + def _prepare_backbone_features_per_frame(self, img_batch, img_ids): + """Compute the image backbone features on the fly for the given img_ids.""" + # Only forward backbone on unique image ids to avoid repeatitive computation + # (if `img_ids` has only one element, it's already unique so we skip this step). + if img_ids.numel() > 1: + unique_img_ids, inv_ids = torch.unique(img_ids, return_inverse=True) + else: + unique_img_ids, inv_ids = img_ids, None + + # Compute the image features on those unique image ids + image = img_batch[unique_img_ids] + backbone_out = self.forward_image(image) + ( + _, + vision_feats, + vision_pos_embeds, + feat_sizes, + ) = self._prepare_backbone_features(backbone_out) + # Inverse-map image features for `unique_img_ids` to the final image features + # for the original input `img_ids`. + if inv_ids is not None: + image = image[inv_ids] + vision_feats = [x[:, inv_ids] for x in vision_feats] + vision_pos_embeds = [x[:, inv_ids] for x in vision_pos_embeds] + + return image, vision_feats, vision_pos_embeds, feat_sizes + + def cal_mem_score(self, object_score_logits, iou_score): + object_score_norm = torch.where( + object_score_logits > 0, + object_score_logits.sigmoid() * 2 - 1, ## rescale to [0, 1] + torch.zeros_like(object_score_logits), + ) + score_per_frame = (object_score_norm * iou_score).mean() + return score_per_frame + + def frame_filter(self, output_dict, track_in_reverse, frame_idx, num_frames, r): + if (frame_idx == 0 and not track_in_reverse) or ( + frame_idx == num_frames - 1 and track_in_reverse + ): + return [] + + max_num = min( + num_frames, self.max_obj_ptrs_in_encoder + ) ## maximum number of pointer memory frames to consider + + if not track_in_reverse: + start = frame_idx - 1 + end = 0 + step = -r + must_include = frame_idx - 1 + else: + start = frame_idx + 1 + end = num_frames + step = r + must_include = frame_idx + 1 + + valid_indices = [] + for i in range(start, end, step): + if ( + i not in output_dict["non_cond_frame_outputs"] + or "eff_iou_score" not in output_dict["non_cond_frame_outputs"][i] + ): + continue + + score_per_frame = output_dict["non_cond_frame_outputs"][i]["eff_iou_score"] + + if score_per_frame > self.mf_threshold: # threshold + valid_indices.insert(0, i) + + if len(valid_indices) >= max_num - 1: + break + + if must_include not in valid_indices: + valid_indices.append(must_include) + + return valid_indices + + def _prepare_memory_conditioned_features( + self, + frame_idx, + is_init_cond_frame, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + output_dict, + num_frames, + track_in_reverse=False, # tracking in reverse time order (for demo usage) + use_prev_mem_frame=True, + ): + """Fuse the current frame's visual feature map with previous memory.""" + B = current_vision_feats[-1].size(1) # batch size on this frame + C = self.hidden_dim + H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size + device = current_vision_feats[-1].device + # The case of `self.num_maskmem == 0` below is primarily used for reproducing SAM on images. + # In this case, we skip the fusion with any memory. + if self.num_maskmem == 0: # Disable memory and skip fusion + pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W) + return pix_feat + + num_obj_ptr_tokens = 0 + tpos_sign_mul = -1 if track_in_reverse else 1 + # Step 1: condition the visual features of the current frame on previous memories + if not is_init_cond_frame and use_prev_mem_frame: + # Retrieve the memories encoded with the maskmem backbone + to_cat_prompt, to_cat_prompt_mask, to_cat_prompt_pos_embed = [], [], [] + # Add conditioning frames's output first (all cond frames have t_pos=0 for + # when getting temporal positional embedding below) + assert len(output_dict["cond_frame_outputs"]) > 0 + # Select a maximum number of temporally closest cond frames for cross attention + cond_outputs = output_dict["cond_frame_outputs"] + selected_cond_outputs, unselected_cond_outputs = select_closest_cond_frames( + frame_idx, + cond_outputs, + self.max_cond_frames_in_attn, + keep_first_cond_frame=self.keep_first_cond_frame, + ) + t_pos_and_prevs = [ + ((frame_idx - t) * tpos_sign_mul, out, True) + for t, out in selected_cond_outputs.items() + ] + # Add last (self.num_maskmem - 1) frames before current frame for non-conditioning memory + # the earliest one has t_pos=1 and the latest one has t_pos=self.num_maskmem-1 + # We also allow taking the memory frame non-consecutively (with r>1), in which case + # we take (self.num_maskmem - 2) frames among every r-th frames plus the last frame. + r = 1 if self.training else self.memory_temporal_stride_for_eval + + if self.use_memory_selection: + valid_indices = self.frame_filter( + output_dict, track_in_reverse, frame_idx, num_frames, r + ) + + for t_pos in range(1, self.num_maskmem): + t_rel = self.num_maskmem - t_pos # how many frames before current frame + if self.use_memory_selection: + if t_rel > len(valid_indices): + continue + prev_frame_idx = valid_indices[-t_rel] + else: + if t_rel == 1: + # for t_rel == 1, we take the last frame (regardless of r) + if not track_in_reverse: + # the frame immediately before this frame (i.e. frame_idx - 1) + prev_frame_idx = frame_idx - t_rel + else: + # the frame immediately after this frame (i.e. frame_idx + 1) + prev_frame_idx = frame_idx + t_rel + else: + # for t_rel >= 2, we take the memory frame from every r-th frames + if not track_in_reverse: + # first find the nearest frame among every r-th frames before this frame + # for r=1, this would be (frame_idx - 2) + prev_frame_idx = ((frame_idx - 2) // r) * r + # then seek further among every r-th frames + prev_frame_idx = prev_frame_idx - (t_rel - 2) * r + else: + # first find the nearest frame among every r-th frames after this frame + # for r=1, this would be (frame_idx + 2) + prev_frame_idx = -(-(frame_idx + 2) // r) * r + # then seek further among every r-th frames + prev_frame_idx = prev_frame_idx + (t_rel - 2) * r + + out = output_dict["non_cond_frame_outputs"].get(prev_frame_idx, None) + if out is None: + # If an unselected conditioning frame is among the last (self.num_maskmem - 1) + # frames, we still attend to it as if it's a non-conditioning frame. + out = unselected_cond_outputs.get(prev_frame_idx, None) + t_pos_and_prevs.append((t_pos, out, False)) + + for t_pos, prev, is_selected_cond_frame in t_pos_and_prevs: + if prev is None: + continue # skip padding frames + # "maskmem_features" might have been offloaded to CPU in demo use cases, + # so we load it back to GPU (it's a no-op if it's already on GPU). + feats = prev["maskmem_features"].cuda(non_blocking=True) + seq_len = feats.shape[-2] * feats.shape[-1] + to_cat_prompt.append(feats.flatten(2).permute(2, 0, 1)) + to_cat_prompt_mask.append( + torch.zeros(B, seq_len, device=device, dtype=bool) + ) + # Spatial positional encoding (it might have been offloaded to CPU in eval) + maskmem_enc = prev["maskmem_pos_enc"][-1].cuda() + maskmem_enc = maskmem_enc.flatten(2).permute(2, 0, 1) + + if ( + is_selected_cond_frame + and getattr(self, "cond_frame_spatial_embedding", None) is not None + ): + # add a spatial embedding for the conditioning frame + maskmem_enc = maskmem_enc + self.cond_frame_spatial_embedding + + # Temporal positional encoding + t = t_pos if not is_selected_cond_frame else 0 + maskmem_enc = ( + maskmem_enc + self.maskmem_tpos_enc[self.num_maskmem - t - 1] + ) + to_cat_prompt_pos_embed.append(maskmem_enc) + + # Construct the list of past object pointers + # Optionally, select only a subset of spatial memory frames during trainining + if ( + self.training + and self.prob_to_dropout_spatial_mem > 0 + and self.rng.random() < self.prob_to_dropout_spatial_mem + ): + num_spatial_mem_keep = self.rng.integers(len(to_cat_prompt) + 1) + keep = self.rng.choice( + range(len(to_cat_prompt)), num_spatial_mem_keep, replace=False + ).tolist() + to_cat_prompt = [to_cat_prompt[i] for i in keep] + to_cat_prompt_mask = [to_cat_prompt_mask[i] for i in keep] + to_cat_prompt_pos_embed = [to_cat_prompt_pos_embed[i] for i in keep] + + max_obj_ptrs_in_encoder = min(num_frames, self.max_obj_ptrs_in_encoder) + # First add those object pointers from selected conditioning frames + # (optionally, only include object pointers in the past during evaluation) + if not self.training: + ptr_cond_outputs = { + t: out + for t, out in selected_cond_outputs.items() + if (t >= frame_idx if track_in_reverse else t <= frame_idx) + } + else: + ptr_cond_outputs = selected_cond_outputs + pos_and_ptrs = [ + # Temporal pos encoding contains how far away each pointer is from current frame + ( + (frame_idx - t) * tpos_sign_mul, + out["obj_ptr"], + True, # is_selected_cond_frame + ) + for t, out in ptr_cond_outputs.items() + ] + + # Add up to (max_obj_ptrs_in_encoder - 1) non-conditioning frames before current frame + for t_diff in range(1, max_obj_ptrs_in_encoder): + if not self.use_memory_selection: + t = frame_idx + t_diff if track_in_reverse else frame_idx - t_diff + if t < 0 or (num_frames is not None and t >= num_frames): + break + else: + if -t_diff <= -len(valid_indices): + break + t = valid_indices[-t_diff] + + out = output_dict["non_cond_frame_outputs"].get( + t, unselected_cond_outputs.get(t, None) + ) + if out is not None: + pos_and_ptrs.append((t_diff, out["obj_ptr"], False)) + + # If we have at least one object pointer, add them to the across attention + if len(pos_and_ptrs) > 0: + pos_list, ptrs_list, is_selected_cond_frame_list = zip(*pos_and_ptrs) + # stack object pointers along dim=0 into [ptr_seq_len, B, C] shape + obj_ptrs = torch.stack(ptrs_list, dim=0) + if getattr(self, "cond_frame_obj_ptr_embedding", None) is not None: + obj_ptrs = ( + obj_ptrs + + self.cond_frame_obj_ptr_embedding + * torch.tensor(is_selected_cond_frame_list, device=device)[ + ..., None, None + ].float() + ) + # a temporal positional embedding based on how far each object pointer is from + # the current frame (sine embedding normalized by the max pointer num). + obj_pos = self._get_tpos_enc( + pos_list, + max_abs_pos=max_obj_ptrs_in_encoder, + device=device, + ) + # expand to batch size + obj_pos = obj_pos.unsqueeze(1).expand(-1, B, -1) + + if self.mem_dim < C: + # split a pointer into (C // self.mem_dim) tokens for self.mem_dim < C + obj_ptrs = obj_ptrs.reshape(-1, B, C // self.mem_dim, self.mem_dim) + obj_ptrs = obj_ptrs.permute(0, 2, 1, 3).flatten(0, 1) + obj_pos = obj_pos.repeat_interleave(C // self.mem_dim, dim=0) + to_cat_prompt.append(obj_ptrs) + to_cat_prompt_mask.append(None) # "to_cat_prompt_mask" is not used + to_cat_prompt_pos_embed.append(obj_pos) + num_obj_ptr_tokens = obj_ptrs.shape[0] + else: + num_obj_ptr_tokens = 0 + else: + # directly add no-mem embedding (instead of using the transformer encoder) + pix_feat_with_mem = current_vision_feats[-1] + self.no_mem_embed + pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W) + return pix_feat_with_mem + + # Use a dummy token on the first grame (to avoid emtpy memory input to tranformer encoder) + to_cat_prompt = [self.no_mem_embed.expand(1, B, self.mem_dim)] + to_cat_prompt_mask = [torch.zeros(B, 1, device=device, dtype=bool)] + to_cat_prompt_pos_embed = [self.no_mem_pos_enc.expand(1, B, self.mem_dim)] + + # Step 2: Concatenate the memories and forward through the transformer encoder + prompt = torch.cat(to_cat_prompt, dim=0) + prompt_mask = None # For now, we always masks are zeros anyways + prompt_pos_embed = torch.cat(to_cat_prompt_pos_embed, dim=0) + encoder_out = self.transformer.encoder( + src=current_vision_feats, + src_key_padding_mask=[None], + src_pos=current_vision_pos_embeds, + prompt=prompt, + prompt_pos=prompt_pos_embed, + prompt_key_padding_mask=prompt_mask, + feat_sizes=feat_sizes, + num_obj_ptr_tokens=num_obj_ptr_tokens, + ) + # reshape the output (HW)BC => BCHW + pix_feat_with_mem = encoder_out["memory"].permute(1, 2, 0).view(B, C, H, W) + return pix_feat_with_mem + + def _encode_new_memory( + self, + image, + current_vision_feats, + feat_sizes, + pred_masks_high_res, + object_score_logits, + is_mask_from_pts, + output_dict=None, + is_init_cond_frame=False, + ): + """Encode the current image and its prediction into a memory feature.""" + B = current_vision_feats[-1].size(1) # batch size on this frame + C = self.hidden_dim + H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size + # top-level feature, (HW)BC => BCHW + pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W) + if self.non_overlap_masks_for_mem_enc and not self.training: + # optionally, apply non-overlapping constraints to the masks (it's applied + # in the batch dimension and should only be used during eval, where all + # the objects come from the same video under batch size 1). + pred_masks_high_res = self._apply_non_overlapping_constraints( + pred_masks_high_res + ) + # scale the raw mask logits with a temperature before applying sigmoid + if is_mask_from_pts and not self.training: + mask_for_mem = (pred_masks_high_res > 0).float() + else: + # apply sigmoid on the raw mask logits to turn them into range (0, 1) + mask_for_mem = torch.sigmoid(pred_masks_high_res) + # apply scale and bias terms to the sigmoid probabilities + if self.sigmoid_scale_for_mem_enc != 1.0: + mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc + if self.sigmoid_bias_for_mem_enc != 0.0: + mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc + + if isinstance(self.maskmem_backbone, SimpleMaskEncoder): + pix_feat = pix_feat.view_as(pix_feat) + maskmem_out = self.maskmem_backbone( + pix_feat, mask_for_mem, skip_mask_sigmoid=True + ) + else: + maskmem_out = self.maskmem_backbone(image, pix_feat, mask_for_mem) + # Clone the feats and pos_enc to enable compilation + maskmem_features = self._maybe_clone(maskmem_out["vision_features"]) + maskmem_pos_enc = [self._maybe_clone(m) for m in maskmem_out["vision_pos_enc"]] + # add a no-object embedding to the spatial memory to indicate that the frame + # is predicted to be occluded (i.e. no object is appearing in the frame) + is_obj_appearing = (object_score_logits > 0).float() + maskmem_features += ( + 1 - is_obj_appearing[..., None, None] + ) * self.no_obj_embed_spatial[..., None, None].expand(*maskmem_features.shape) + + return maskmem_features, maskmem_pos_enc + + def forward_tracking(self, backbone_out, input, return_dict=False): + """Forward video tracking on each frame (and sample correction clicks).""" + img_feats_already_computed = backbone_out["backbone_fpn"] is not None + if img_feats_already_computed: + # Prepare the backbone features + # - vision_feats and vision_pos_embeds are in (HW)BC format + ( + _, + vision_feats, + vision_pos_embeds, + feat_sizes, + ) = self._prepare_backbone_features(backbone_out) + + # Starting the stage loop + num_frames = backbone_out["num_frames"] + init_cond_frames = backbone_out["init_cond_frames"] + frames_to_add_correction_pt = backbone_out["frames_to_add_correction_pt"] + # first process all the initial conditioning frames to encode them as memory, + # and then conditioning on them to track the remaining frames + processing_order = init_cond_frames + backbone_out["frames_not_in_init_cond"] + output_dict = { + "cond_frame_outputs": {}, # dict containing {frame_idx: } + "non_cond_frame_outputs": {}, # dict containing {frame_idx: } + } + for stage_id in processing_order: + # Get the image features for the current frames + img_ids = input.find_inputs[stage_id].img_ids + if img_feats_already_computed: + # Retrieve image features according to img_ids (if they are already computed). + current_image = input.img_batch[img_ids] + current_vision_feats = [x[:, img_ids] for x in vision_feats] + current_vision_pos_embeds = [x[:, img_ids] for x in vision_pos_embeds] + else: + # Otherwise, compute the image features on the fly for the given img_ids + # (this might be used for evaluation on long videos to avoid backbone OOM). + ( + current_image, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + ) = self._prepare_backbone_features_per_frame(input.img_batch, img_ids) + # Get output masks based on this frame's prompts and previous memory + current_out = self.track_step( + frame_idx=stage_id, + is_init_cond_frame=stage_id in init_cond_frames, + current_vision_feats=current_vision_feats, + current_vision_pos_embeds=current_vision_pos_embeds, + feat_sizes=feat_sizes, + image=current_image, + point_inputs=backbone_out["point_inputs_per_frame"].get(stage_id, None), + mask_inputs=backbone_out["mask_inputs_per_frame"].get(stage_id, None), + gt_masks=backbone_out["gt_masks_per_frame"].get(stage_id, None), + frames_to_add_correction_pt=frames_to_add_correction_pt, + output_dict=output_dict, + num_frames=num_frames, + ) + # Append the output, depending on whether it's a conditioning frame + add_output_as_cond_frame = stage_id in init_cond_frames or ( + self.add_all_frames_to_correct_as_cond + and stage_id in frames_to_add_correction_pt + ) + if add_output_as_cond_frame: + output_dict["cond_frame_outputs"][stage_id] = current_out + else: + output_dict["non_cond_frame_outputs"][stage_id] = current_out + + if return_dict: + return output_dict + # turn `output_dict` into a list for loss function + all_frame_outputs = {} + all_frame_outputs.update(output_dict["cond_frame_outputs"]) + all_frame_outputs.update(output_dict["non_cond_frame_outputs"]) + all_frame_outputs = [all_frame_outputs[t] for t in range(num_frames)] + # Make DDP happy with activation checkpointing by removing unused keys + all_frame_outputs = [ + {k: v for k, v in d.items() if k != "obj_ptr"} for d in all_frame_outputs + ] + + return all_frame_outputs + + def track_step( + self, + frame_idx, + is_init_cond_frame, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + image, + point_inputs, + mask_inputs, + output_dict, + num_frames, + track_in_reverse=False, # tracking in reverse time order (for demo usage) + # Whether to run the memory encoder on the predicted masks. Sometimes we might want + # to skip the memory encoder with `run_mem_encoder=False`. For example, + # in demo we might call `track_step` multiple times for each user click, + # and only encode the memory when the user finalizes their clicks. And in ablation + # settings like SAM training on static images, we don't need the memory encoder. + run_mem_encoder=True, + # The previously predicted SAM mask logits (which can be fed together with new clicks in demo). + prev_sam_mask_logits=None, + use_prev_mem_frame=True, + ): + current_out = {"point_inputs": point_inputs, "mask_inputs": mask_inputs} + # High-resolution feature maps for the SAM head, reshape (HW)BC => BCHW + if len(current_vision_feats) > 1: + high_res_features = [ + x.permute(1, 2, 0).view(x.size(1), x.size(2), *s) + for x, s in zip(current_vision_feats[:-1], feat_sizes[:-1]) + ] + else: + high_res_features = None + if mask_inputs is not None: + # (see it as a GT mask) without using a SAM prompt encoder + mask decoder. + pix_feat = current_vision_feats[-1].permute(1, 2, 0) + pix_feat = pix_feat.view(-1, self.hidden_dim, *feat_sizes[-1]) + sam_outputs = self._use_mask_as_output( + pix_feat, high_res_features, mask_inputs + ) + else: + # fused the visual feature with previous memory features in the memory bank + pix_feat_with_mem = self._prepare_memory_conditioned_features( + frame_idx=frame_idx, + is_init_cond_frame=is_init_cond_frame, + current_vision_feats=current_vision_feats[-1:], + current_vision_pos_embeds=current_vision_pos_embeds[-1:], + feat_sizes=feat_sizes[-1:], + output_dict=output_dict, + num_frames=num_frames, + track_in_reverse=track_in_reverse, + use_prev_mem_frame=use_prev_mem_frame, + ) + # apply SAM-style segmentation head + # here we might feed previously predicted low-res SAM mask logits into the SAM mask decoder, + # e.g. in demo where such logits come from earlier interaction instead of correction sampling + # (in this case, the SAM mask decoder should have `self.iter_use_prev_mask_pred=True`, and + # any `mask_inputs` shouldn't reach here as they are sent to _use_mask_as_output instead) + if prev_sam_mask_logits is not None: + assert self.iter_use_prev_mask_pred + assert point_inputs is not None and mask_inputs is None + mask_inputs = prev_sam_mask_logits + multimask_output = self._use_multimask(is_init_cond_frame, point_inputs) + sam_outputs = self._forward_sam_heads( + backbone_features=pix_feat_with_mem, + point_inputs=point_inputs, + mask_inputs=mask_inputs, + high_res_features=high_res_features, + multimask_output=multimask_output, + ) + ( + _, + high_res_multimasks, + ious, + low_res_masks, + high_res_masks, + obj_ptr, + object_score_logits, + ) = sam_outputs + # Use the final prediction (after all correction steps for output and eval) + current_out["pred_masks"] = low_res_masks + current_out["pred_masks_high_res"] = high_res_masks + current_out["obj_ptr"] = obj_ptr + if self.use_memory_selection: + current_out["object_score_logits"] = object_score_logits + iou_score = ious.max(-1)[0] + current_out["iou_score"] = iou_score + current_out["eff_iou_score"] = self.cal_mem_score( + object_score_logits, iou_score + ) + if not self.training: + # Only add this in inference (to avoid unused param in activation checkpointing; + # it's mainly used in the demo to encode spatial memories w/ consolidated masks) + current_out["object_score_logits"] = object_score_logits + + # Finally run the memory encoder on the predicted mask to encode + # it into a new memory feature (that can be used in future frames) + # (note that `self.num_maskmem == 0` is primarily used for reproducing SAM on + # images, in which case we'll just skip memory encoder to save compute). + if run_mem_encoder and self.num_maskmem > 0: + high_res_masks_for_mem_enc = high_res_masks + maskmem_features, maskmem_pos_enc = self._encode_new_memory( + image=image, + current_vision_feats=current_vision_feats, + feat_sizes=feat_sizes, + pred_masks_high_res=high_res_masks_for_mem_enc, + object_score_logits=object_score_logits, + is_mask_from_pts=(point_inputs is not None), + output_dict=output_dict, + is_init_cond_frame=is_init_cond_frame, + ) + current_out["maskmem_features"] = maskmem_features + current_out["maskmem_pos_enc"] = maskmem_pos_enc + else: + current_out["maskmem_features"] = None + current_out["maskmem_pos_enc"] = None + + # Optionally, offload the outputs to CPU memory during evaluation to avoid + # GPU OOM on very long videos or very large resolution or too many objects + if self.offload_output_to_cpu_for_eval and not self.training: + # Here we only keep those keys needed for evaluation to get a compact output + trimmed_out = { + "pred_masks": current_out["pred_masks"].cpu(), + "pred_masks_high_res": current_out["pred_masks_high_res"].cpu(), + # other items for evaluation (these are small tensors so we keep them on GPU) + "obj_ptr": current_out["obj_ptr"], + "object_score_logits": current_out["object_score_logits"], + } + if run_mem_encoder and self.num_maskmem > 0: + trimmed_out["maskmem_features"] = maskmem_features.cpu() + trimmed_out["maskmem_pos_enc"] = [x.cpu() for x in maskmem_pos_enc] + if self.use_memory_selection: + trimmed_out["iou_score"] = current_out["iou_score"].cpu() + trimmed_out["eff_iou_score"] = current_out["eff_iou_score"].cpu() + current_out = trimmed_out + + # Optionally, trim the output of past non-conditioning frame (r * num_maskmem frames + # before the current frame) during evaluation. This is intended to save GPU or CPU + # memory for semi-supervised VOS eval, where only the first frame receives prompts. + def _trim_past_out(past_out, current_out): + if past_out is None: + return None + return { + "pred_masks": past_out["pred_masks"], + "obj_ptr": past_out["obj_ptr"], + "object_score_logits": past_out["object_score_logits"], + } + + if self.trim_past_non_cond_mem_for_eval and not self.training: + r = self.memory_temporal_stride_for_eval + past_frame_idx = frame_idx - r * self.num_maskmem + past_out = output_dict["non_cond_frame_outputs"].get(past_frame_idx, None) + + if past_out is not None: + print(past_out.get("eff_iou_score", 0)) + if ( + self.use_memory_selection + and past_out.get("eff_iou_score", 0) < self.mf_threshold + ) or not self.use_memory_selection: + output_dict["non_cond_frame_outputs"][past_frame_idx] = ( + _trim_past_out(past_out, current_out) + ) + + if ( + self.use_memory_selection and not self.offload_output_to_cpu_for_eval + ): ## design for memory selection, trim too old frames to save memory + far_old_frame_idx = frame_idx - 20 * self.max_obj_ptrs_in_encoder + past_out = output_dict["non_cond_frame_outputs"].get( + far_old_frame_idx, None + ) + if past_out is not None: + output_dict["non_cond_frame_outputs"][far_old_frame_idx] = ( + _trim_past_out(past_out, current_out) + ) + + return current_out + + def _use_multimask(self, is_init_cond_frame, point_inputs): + """Whether to use multimask output in the SAM head.""" + num_pts = 0 if point_inputs is None else point_inputs["point_labels"].size(1) + multimask_output = ( + self.multimask_output_in_sam + and (is_init_cond_frame or self.multimask_output_for_tracking) + and (self.multimask_min_pt_num <= num_pts <= self.multimask_max_pt_num) + ) + return multimask_output + + def _apply_non_overlapping_constraints(self, pred_masks): + """ + Apply non-overlapping constraints to the object scores in pred_masks. Here we + keep only the highest scoring object at each spatial location in pred_masks. + """ + batch_size = pred_masks.size(0) + if batch_size == 1: + return pred_masks + + device = pred_masks.device + # "max_obj_inds": object index of the object with the highest score at each location + max_obj_inds = torch.argmax(pred_masks, dim=0, keepdim=True) + # "batch_obj_inds": object index of each object slice (along dim 0) in `pred_masks` + batch_obj_inds = torch.arange(batch_size, device=device)[:, None, None, None] + keep = max_obj_inds == batch_obj_inds + # suppress overlapping regions' scores below -10.0 so that the foreground regions + # don't overlap (here sigmoid(-10.0)=4.5398e-05) + pred_masks = torch.where(keep, pred_masks, torch.clamp(pred_masks, max=-10.0)) + return pred_masks + + def _compile_all_components(self): + """Compile all model components for faster inference.""" + # a larger cache size to hold varying number of shapes for torch.compile + # see https://github.com/pytorch/pytorch/blob/v2.5.1/torch/_dynamo/config.py#L42-L49 + torch._dynamo.config.cache_size_limit = 64 + torch._dynamo.config.accumulated_cache_size_limit = 2048 + from sam3.perflib.compile import compile_wrapper + + logging.info("Compiling all components. First time may be very slow.") + + self.maskmem_backbone.forward = compile_wrapper( + self.maskmem_backbone.forward, + mode="max-autotune", + fullgraph=True, + dynamic=False, + ) + self.transformer.encoder.forward = compile_wrapper( + self.transformer.encoder.forward, + mode="max-autotune", + fullgraph=True, + dynamic=True, # Num. of memories varies + ) + # We disable compilation of sam_prompt_encoder as it sometimes gives a large accuracy regression, + # especially when sam_mask_prompt (previous mask logits) is not None + # self.sam_prompt_encoder.forward = torch.compile( + # self.sam_prompt_encoder.forward, + # mode="max-autotune", + # fullgraph=True, + # dynamic=False, # Accuracy regression on True + # ) + self.sam_mask_decoder.forward = compile_wrapper( + self.sam_mask_decoder.forward, + mode="max-autotune", + fullgraph=True, + dynamic=False, # Accuracy regression on True + ) + + def _maybe_clone(self, x): + """Clone a tensor if and only if `self.compile_all_components` is True.""" + return x.clone() if self.compile_all_components else x + + +def concat_points(old_point_inputs, new_points, new_labels): + """Add new points and labels to previous point inputs (add at the end).""" + if old_point_inputs is None: + points, labels = new_points, new_labels + else: + points = torch.cat([old_point_inputs["point_coords"], new_points], dim=1) + labels = torch.cat([old_point_inputs["point_labels"], new_labels], dim=1) + + return {"point_coords": points, "point_labels": labels} diff --git a/src/nn/segearth_ov3/sam3/model/sam3_tracker_utils.py b/src/nn/segearth_ov3/sam3/model/sam3_tracker_utils.py new file mode 100644 index 0000000..7afc70a --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_tracker_utils.py @@ -0,0 +1,427 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import numpy as np +import torch +import torch.nn.functional as F +from numpy.typing import NDArray + +from sam3.model.edt import edt_triton + + +def sample_box_points( + masks: torch.Tensor, + noise: float = 0.1, # SAM default + noise_bound: int = 20, # SAM default + top_left_label: int = 2, + bottom_right_label: int = 3, +) -> tuple[NDArray, NDArray]: + """ + Sample a noised version of the top left and bottom right corners of a given `bbox` + + Inputs: + - masks: [B, 1, H, W] tensor + - noise: noise as a fraction of box width and height, dtype=float + - noise_bound: maximum amount of noise (in pure pixels), dtype=int + + Returns: + - box_coords: [B, num_pt, 2], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.float + - box_labels: [B, num_pt], label 2 is reserverd for top left and 3 for bottom right corners, dtype=torch.int32 + """ + device = masks.device + box_coords = mask_to_box(masks) + B, _, H, W = masks.shape + box_labels = torch.tensor( + [top_left_label, bottom_right_label], dtype=torch.int, device=device + ).repeat(B) + if noise > 0.0: + if not isinstance(noise_bound, torch.Tensor): + noise_bound = torch.tensor(noise_bound, device=device) + bbox_w = box_coords[..., 2] - box_coords[..., 0] + bbox_h = box_coords[..., 3] - box_coords[..., 1] + max_dx = torch.min(bbox_w * noise, noise_bound) + max_dy = torch.min(bbox_h * noise, noise_bound) + box_noise = 2 * torch.rand(B, 1, 4, device=device) - 1 + box_noise = box_noise * torch.stack((max_dx, max_dy, max_dx, max_dy), dim=-1) + + box_coords = box_coords + box_noise + img_bounds = ( + torch.tensor([W, H, W, H], device=device) - 1 + ) # uncentered pixel coords + box_coords.clamp_(torch.zeros_like(img_bounds), img_bounds) # In place clamping + + box_coords = box_coords.reshape(-1, 2, 2) # always 2 points + box_labels = box_labels.reshape(-1, 2) + return box_coords, box_labels + + +def mask_to_box(masks: torch.Tensor): + """ + compute bounding box given an input mask + + Inputs: + - masks: [B, 1, H, W] tensor + + Returns: + - box_coords: [B, 1, 4], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.Tensor + """ + B, _, h, w = masks.shape + device = masks.device + mask_area = masks.sum(dim=(-1, -2)) + xs = torch.arange(w, device=device, dtype=torch.int32) + ys = torch.arange(h, device=device, dtype=torch.int32) + grid_xs, grid_ys = torch.meshgrid(xs, ys, indexing="xy") + grid_xs = grid_xs[None, None, ...].expand(B, 1, h, w) + grid_ys = grid_ys[None, None, ...].expand(B, 1, h, w) + min_xs, _ = torch.min(torch.where(masks, grid_xs, w).flatten(-2), dim=-1) + max_xs, _ = torch.max(torch.where(masks, grid_xs, -1).flatten(-2), dim=-1) + min_ys, _ = torch.min(torch.where(masks, grid_ys, h).flatten(-2), dim=-1) + max_ys, _ = torch.max(torch.where(masks, grid_ys, -1).flatten(-2), dim=-1) + bbox_coords = torch.stack((min_xs, min_ys, max_xs, max_ys), dim=-1) + bbox_coords = torch.where( + mask_area[..., None] > 0, bbox_coords, torch.zeros_like(bbox_coords) + ) + return bbox_coords + + +def sample_random_points_from_errors(gt_masks, pred_masks, num_pt=1): + """ + Sample `num_pt` random points (along with their labels) independently from the error regions. + + Inputs: + - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool + - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None + - num_pt: int, number of points to sample independently for each of the B error maps + + Outputs: + - points: [B, num_pt, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point + - labels: [B, num_pt], dtype=torch.int32, where 1 means positive clicks and 0 means + negative clicks + """ + if pred_masks is None: # if pred_masks is not provided, treat it as empty + pred_masks = torch.zeros_like(gt_masks) + assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1 + assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape + assert num_pt >= 0 + + B, _, H_im, W_im = gt_masks.shape + device = gt_masks.device + + # false positive region, a new point sampled in this region should have + # negative label to correct the FP error + fp_masks = ~gt_masks & pred_masks + # false negative region, a new point sampled in this region should have + # positive label to correct the FN error + fn_masks = gt_masks & ~pred_masks + # whether the prediction completely match the ground-truth on each mask + all_correct = torch.all((gt_masks == pred_masks).flatten(2), dim=2) + all_correct = all_correct[..., None, None] + + # channel 0 is FP map, while channel 1 is FN map + pts_noise = torch.rand(B, num_pt, H_im, W_im, 2, device=device) + # sample a negative new click from FP region or a positive new click + # from FN region, depend on where the maximum falls, + # and in case the predictions are all correct (no FP or FN), we just + # sample a negative click from the background region + pts_noise[..., 0] *= fp_masks | (all_correct & ~gt_masks) + pts_noise[..., 1] *= fn_masks + pts_idx = pts_noise.flatten(2).argmax(dim=2) + labels = (pts_idx % 2).to(torch.int32) + pts_idx = pts_idx // 2 + pts_x = pts_idx % W_im + pts_y = pts_idx // W_im + points = torch.stack([pts_x, pts_y], dim=2).to(torch.float) + return points, labels + + +def sample_one_point_from_error_center(gt_masks, pred_masks, padding=True): + """ + Sample 1 random point (along with its label) from the center of each error region, + that is, the point with the largest distance to the boundary of each error region. + This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py + + Inputs: + - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool + - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None + - padding: if True, pad with boundary of 1 px for distance transform + + Outputs: + - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point + - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks + """ + if pred_masks is None: + pred_masks = torch.zeros_like(gt_masks) + assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1 + assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape + + B, _, H, W = gt_masks.shape + + # false positive region, a new point sampled in this region should have + # negative label to correct the FP error + fp_masks = (~gt_masks & pred_masks).squeeze(1) + # false negative region, a new point sampled in this region should have + # positive label to correct the FN error + fn_masks = (gt_masks & ~pred_masks).squeeze(1) + + if padding: + padded_fp_masks = torch.zeros( + B, H + 2, W + 2, dtype=fp_masks.dtype, device=fp_masks.device + ) + padded_fp_masks[:, 1 : H + 1, 1 : W + 1] = fp_masks + padded_fn_masks = torch.zeros( + B, H + 2, W + 2, dtype=fp_masks.dtype, device=fp_masks.device + ) + padded_fn_masks[:, 1 : H + 1, 1 : W + 1] = fn_masks + else: + padded_fp_masks = fp_masks + padded_fn_masks = fn_masks + + fn_mask_dt = edt_triton(padded_fn_masks) + fp_mask_dt = edt_triton(padded_fp_masks) + if padding: + fn_mask_dt = fn_mask_dt[:, 1:-1, 1:-1] + fp_mask_dt = fp_mask_dt[:, 1:-1, 1:-1] + + fn_max, fn_argmax = fn_mask_dt.reshape(B, -1).max(dim=-1) + fp_max, fp_argmax = fp_mask_dt.reshape(B, -1).max(dim=-1) + is_positive = fn_max > fp_max + chosen = torch.where(is_positive, fn_argmax, fp_argmax) + points_x = chosen % W + points_y = chosen // W + + labels = is_positive.long() + points = torch.stack([points_x, points_y], -1) + return points.unsqueeze(1), labels.unsqueeze(1) + + +def sample_one_point_from_error_center_slow(gt_masks, pred_masks, padding=True): + """ + Sample 1 random point (along with its label) from the center of each error region, + that is, the point with the largest distance to the boundary of each error region. + This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py + + Inputs: + - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool + - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None + - padding: if True, pad with boundary of 1 px for distance transform + + Outputs: + - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point + - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks + """ + import cv2 # delay OpenCV import to avoid unnecessary dependency + + if pred_masks is None: + pred_masks = torch.zeros_like(gt_masks) + assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1 + assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape + + B, _, _, W_im = gt_masks.shape + device = gt_masks.device + + # false positive region, a new point sampled in this region should have + # negative label to correct the FP error + fp_masks = ~gt_masks & pred_masks + # false negative region, a new point sampled in this region should have + # positive label to correct the FN error + fn_masks = gt_masks & ~pred_masks + + fp_masks = fp_masks.cpu().numpy() + fn_masks = fn_masks.cpu().numpy() + points = torch.zeros(B, 1, 2, dtype=torch.float) + labels = torch.ones(B, 1, dtype=torch.int32) + for b in range(B): + fn_mask = fn_masks[b, 0] + fp_mask = fp_masks[b, 0] + if padding: + fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), "constant") + fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), "constant") + # compute the distance of each point in FN/FP region to its boundary + fn_mask_dt = cv2.distanceTransform(fn_mask.astype(np.uint8), cv2.DIST_L2, 0) + fp_mask_dt = cv2.distanceTransform(fp_mask.astype(np.uint8), cv2.DIST_L2, 0) + if padding: + fn_mask_dt = fn_mask_dt[1:-1, 1:-1] + fp_mask_dt = fp_mask_dt[1:-1, 1:-1] + + # take the point in FN/FP region with the largest distance to its boundary + fn_mask_dt_flat = fn_mask_dt.reshape(-1) + fp_mask_dt_flat = fp_mask_dt.reshape(-1) + fn_argmax = np.argmax(fn_mask_dt_flat) + fp_argmax = np.argmax(fp_mask_dt_flat) + is_positive = fn_mask_dt_flat[fn_argmax] > fp_mask_dt_flat[fp_argmax] + pt_idx = fn_argmax if is_positive else fp_argmax + points[b, 0, 0] = pt_idx % W_im # x + points[b, 0, 1] = pt_idx // W_im # y + labels[b, 0] = int(is_positive) + + points = points.to(device) + labels = labels.to(device) + return points, labels + + +def get_next_point(gt_masks, pred_masks, method): + if method == "uniform": + return sample_random_points_from_errors(gt_masks, pred_masks) + elif method == "center": + return sample_one_point_from_error_center(gt_masks, pred_masks) + else: + raise ValueError(f"unknown sampling method {method}") + + +def select_closest_cond_frames( + frame_idx, cond_frame_outputs, max_cond_frame_num, keep_first_cond_frame=False +): + """ + Select up to `max_cond_frame_num` conditioning frames from `cond_frame_outputs` + that are temporally closest to the current frame at `frame_idx`. Here, we take + - a) the closest conditioning frame before `frame_idx` (if any); + - b) the closest conditioning frame after `frame_idx` (if any); + - c) any other temporally closest conditioning frames until reaching a total + of `max_cond_frame_num` conditioning frames. + + Outputs: + - selected_outputs: selected items (keys & values) from `cond_frame_outputs`. + - unselected_outputs: items (keys & values) not selected in `cond_frame_outputs`. + """ + if max_cond_frame_num == -1 or len(cond_frame_outputs) <= max_cond_frame_num: + selected_outputs = cond_frame_outputs + unselected_outputs = {} + else: + assert max_cond_frame_num >= 2, "we should allow using 2+ conditioning frames" + selected_outputs = {} + if keep_first_cond_frame: + idx_first = min( + (t for t in cond_frame_outputs if t < frame_idx), default=None + ) + if idx_first is None: + # Maybe we are tracking in reverse + idx_first = max( + (t for t in cond_frame_outputs if t > frame_idx), default=None + ) + if idx_first is not None: + selected_outputs[idx_first] = cond_frame_outputs[idx_first] + # the closest conditioning frame before `frame_idx` (if any) + idx_before = max((t for t in cond_frame_outputs if t < frame_idx), default=None) + if idx_before is not None: + selected_outputs[idx_before] = cond_frame_outputs[idx_before] + + # the closest conditioning frame after `frame_idx` (if any) + idx_after = min((t for t in cond_frame_outputs if t >= frame_idx), default=None) + if idx_after is not None: + selected_outputs[idx_after] = cond_frame_outputs[idx_after] + + # add other temporally closest conditioning frames until reaching a total + # of `max_cond_frame_num` conditioning frames. + num_remain = max_cond_frame_num - len(selected_outputs) + inds_remain = sorted( + (t for t in cond_frame_outputs if t not in selected_outputs), + key=lambda x: abs(x - frame_idx), + )[:num_remain] + selected_outputs.update((t, cond_frame_outputs[t]) for t in inds_remain) + unselected_outputs = { + t: v for t, v in cond_frame_outputs.items() if t not in selected_outputs + } + + return selected_outputs, unselected_outputs + + +def get_1d_sine_pe(pos_inds, dim, temperature=10000): + """ + Get 1D sine positional embedding as in the original Transformer paper. + """ + pe_dim = dim // 2 + dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device) + dim_t = temperature ** (2 * (dim_t // 2) / pe_dim) + + pos_embed = pos_inds.unsqueeze(-1) / dim_t + pos_embed = torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1) + return pos_embed + + +def get_best_gt_match_from_multimasks(pred_multimasks, gt_masks, pred_scores=None): + """ + Get the mask with the best match to GT masks (based on IoU) from pred_multimasks. + Optionally, use `pred_scores` to break ties in case all IoUs are zeros. + """ + assert pred_multimasks.ndim == 4 and gt_masks.ndim == 4 + if pred_multimasks.size(1) == 1: + return pred_multimasks # only a single mask channel, nothing to select + + pred_multimasks_binary = pred_multimasks > 0 + area_i = torch.sum(pred_multimasks_binary & gt_masks, dim=(2, 3)).float() + area_u = torch.sum(pred_multimasks_binary | gt_masks, dim=(2, 3)).float() + ious = area_i / torch.clamp(area_u, min=1.0) + + # In case all IoUs are zeros (e.g. because the GT mask is empty), use pred_scores + # to break ties and select the best mask + if pred_scores is not None: + has_nonzero_ious = torch.any(ious > 0).expand_as(ious) + scores = torch.where(has_nonzero_ious, ious, pred_scores) + else: + scores = ious + + # Finally, take the best mask prediction (with the highest score) + best_scores_inds = torch.argmax(scores, dim=-1) + batch_inds = torch.arange(scores.size(0), device=scores.device) + best_pred_mask = pred_multimasks[batch_inds, best_scores_inds].unsqueeze(1) + return best_pred_mask + + +def fill_holes_in_mask_scores(mask, max_area, fill_holes=True, remove_sprinkles=True): + """ + A post processor to fill small holes in mask scores with area under `max_area`. + Holes are those small connected components in either background or foreground. + + Note that it relies on the "cc_torch" package to find connected components fast. You can + install it via the following command (`TORCH_CUDA_ARCH_LIST=8.0` is for A100 GPUs): + ``` + pip uninstall -y cc_torch; TORCH_CUDA_ARCH_LIST=8.0 9.0 pip install git+https://github.com/ronghanghu/cc_torch + ``` + Otherwise, it will fallback to a slightly slower triton implementation, or skimage if the tensor is on cpu + """ + + if max_area <= 0: + return mask # nothing to fill in this case + + if fill_holes: + # We remove small connected components in background by changing them to foreground + # with a small positive mask score (0.1). + mask_bg = mask <= 0 + bg_area_thresh = max_area + _, areas_bg = _get_connected_components_with_padding(mask_bg) + small_components_bg = mask_bg & (areas_bg <= bg_area_thresh) + mask = torch.where(small_components_bg, 0.1, mask) + + if remove_sprinkles: + # We remove small connected components in foreground by changing them to background + # with a small negative mask score (-0.1). Here we only remove connected components + # whose areas are under both `max_area` and half of the entire mask's area. This + # removes sprinkles while avoids filtering out tiny objects that we want to track. + mask_fg = mask > 0 + fg_area_thresh = torch.sum(mask_fg, dim=(2, 3), keepdim=True, dtype=torch.int32) + fg_area_thresh.floor_divide_(2).clamp_(max=max_area) + _, areas_fg = _get_connected_components_with_padding(mask_fg) + small_components_fg = mask_fg & (areas_fg <= fg_area_thresh) + mask = torch.where(small_components_fg, -0.1, mask) + return mask + + +def _get_connected_components_with_padding(mask): + """Get connected components from masks (possibly padding them to an even size).""" + from sam3.perflib.connected_components import connected_components + + mask = mask.to(torch.uint8) + _, _, H, W = mask.shape + # make sure both height and width are even (to be compatible with cc_torch) + pad_h = H % 2 + pad_w = W % 2 + if pad_h == 0 and pad_w == 0: + labels, counts = connected_components(mask) + else: + # pad the mask to make its height and width even + # padding format is (padding_left,padding_right,padding_top,padding_bottom) + mask_pad = F.pad(mask, (0, pad_w, 0, pad_h), mode="constant", value=0) + labels, counts = connected_components(mask_pad) + labels = labels[:, :, :H, :W] + counts = counts[:, :, :H, :W] + + return labels, counts diff --git a/src/nn/segearth_ov3/sam3/model/sam3_tracking_predictor.py b/src/nn/segearth_ov3/sam3/model/sam3_tracking_predictor.py new file mode 100644 index 0000000..28ab2bd --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_tracking_predictor.py @@ -0,0 +1,1370 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +from collections import OrderedDict + +import torch + +from sam3.model.sam3_tracker_base import concat_points, NO_OBJ_SCORE, Sam3TrackerBase +from sam3.model.sam3_tracker_utils import fill_holes_in_mask_scores +from sam3.model.utils.sam2_utils import load_video_frames +from tqdm.auto import tqdm + + +class Sam3TrackerPredictor(Sam3TrackerBase): + """ + The demo class that extends the `Sam3TrackerBase` to handle user interactions + and manage inference states, with support for multi-object tracking. + """ + + def __init__( + self, + # whether to clear non-conditioning memory of the surrounding frames (which may contain outdated information) after adding correction clicks; + # note that this would only apply to *single-object tracking* unless `clear_non_cond_mem_for_multi_obj` is also set to True) + clear_non_cond_mem_around_input=False, + # whether to also clear non-conditioning memory of the surrounding frames (only effective when `clear_non_cond_mem_around_input` is True). + clear_non_cond_mem_for_multi_obj=False, + # if fill_hole_area > 0, we fill small holes in the final masks up to this area (after resizing them to the original video resolution) + fill_hole_area=0, + # if always_start_from_first_ann_frame is True, we always start tracking from the frame where we receive the first annotation (clicks or mask) + # and ignore the `start_frame_idx` passed to `propagate_in_video` + always_start_from_first_ann_frame=False, + # the maximum number of points to be used in the prompt encoder, which reduce the domain gap between training (that only has 8 points) + # - if it's set to a positive integer, we only take the `max_point_num_in_prompt_enc//2` points and + # the last `(max_point_num_in_prompt_enc - max_point_num_in_prompt_enc//2)` points in the prompt encoder + # - if it's set to 0 or negative, this option is turned off and we use all points in the prompt encoder + max_point_num_in_prompt_enc=16, + non_overlap_masks_for_output=True, + # checkpoint_file=None, + **kwargs, + ): + super().__init__(**kwargs) + self.clear_non_cond_mem_around_input = clear_non_cond_mem_around_input + self.clear_non_cond_mem_for_multi_obj = clear_non_cond_mem_for_multi_obj + self.fill_hole_area = fill_hole_area + self.always_start_from_first_ann_frame = always_start_from_first_ann_frame + self.max_point_num_in_prompt_enc = max_point_num_in_prompt_enc + self.non_overlap_masks_for_output = non_overlap_masks_for_output + + self.bf16_context = torch.autocast(device_type="cuda", dtype=torch.bfloat16) + self.bf16_context.__enter__() # keep using for the entire model process + + self.iter_use_prev_mask_pred = True + self.add_all_frames_to_correct_as_cond = True + + @torch.inference_mode() + def init_state( + self, + video_height=None, + video_width=None, + num_frames=None, + video_path=None, + cached_features=None, + offload_video_to_cpu=False, + offload_state_to_cpu=False, + async_loading_frames=False, + ): + """Initialize a inference state.""" + inference_state = {} + # whether to offload the video frames to CPU memory + # turning on this option saves the GPU memory with only a very small overhead + inference_state["offload_video_to_cpu"] = offload_video_to_cpu + # whether to offload the inference state to CPU memory + # turning on this option saves the GPU memory at the cost of a lower tracking fps + # (e.g. in a test case of 768x768 model, fps dropped from 27 to 24 when tracking one object + # and from 24 to 21 when tracking two objects) + inference_state["offload_state_to_cpu"] = offload_state_to_cpu + inference_state["device"] = self.device + if offload_state_to_cpu: + inference_state["storage_device"] = torch.device("cpu") + else: + inference_state["storage_device"] = torch.device("cuda") + + if video_path is not None: + images, video_height, video_width = load_video_frames( + video_path=video_path, + image_size=self.image_size, + offload_video_to_cpu=offload_video_to_cpu, + async_loading_frames=async_loading_frames, + compute_device=inference_state["storage_device"], + ) + inference_state["images"] = images + inference_state["num_frames"] = len(images) + inference_state["video_height"] = video_height + inference_state["video_width"] = video_width + else: + # the original video height and width, used for resizing final output scores + inference_state["video_height"] = video_height + inference_state["video_width"] = video_width + inference_state["num_frames"] = num_frames + # inputs on each frame + inference_state["point_inputs_per_obj"] = {} + inference_state["mask_inputs_per_obj"] = {} + # visual features on a small number of recently visited frames for quick interactions + inference_state["cached_features"] = ( + {} if cached_features is None else cached_features + ) + # values that don't change across frames (so we only need to hold one copy of them) + inference_state["constants"] = {} + # mapping between client-side object id and model-side object index + inference_state["obj_id_to_idx"] = OrderedDict() + inference_state["obj_idx_to_id"] = OrderedDict() + inference_state["obj_ids"] = [] + # A storage to hold the model's tracking results and states on each frame + inference_state["output_dict"] = { + "cond_frame_outputs": {}, # dict containing {frame_idx: } + "non_cond_frame_outputs": {}, # dict containing {frame_idx: } + } + # The index of the frame that received the first annotation + inference_state["first_ann_frame_idx"] = None + # Slice (view) of each object tracking results, sharing the same memory with "output_dict" + inference_state["output_dict_per_obj"] = {} + # A temporary storage to hold new outputs when user interact with a frame + # to add clicks or mask (it's merged into "output_dict" before propagation starts) + inference_state["temp_output_dict_per_obj"] = {} + # Frames that already holds consolidated outputs from click or mask inputs + # (we directly use their consolidated outputs during tracking) + inference_state["consolidated_frame_inds"] = { + "cond_frame_outputs": set(), # set containing frame indices + "non_cond_frame_outputs": set(), # set containing frame indices + } + # metadata for each tracking frame (e.g. which direction it's tracked) + inference_state["tracking_has_started"] = False + inference_state["frames_already_tracked"] = {} + self.clear_all_points_in_video(inference_state) + return inference_state + + def _obj_id_to_idx(self, inference_state, obj_id): + """Map client-side object id to model-side object index.""" + obj_idx = inference_state["obj_id_to_idx"].get(obj_id, None) + if obj_idx is not None: + return obj_idx + + # This is a new object id not sent to the server before. We only allow adding + # new objects *before* the tracking starts. + allow_new_object = not inference_state["tracking_has_started"] + if allow_new_object: + # get the next object slot + obj_idx = len(inference_state["obj_id_to_idx"]) + inference_state["obj_id_to_idx"][obj_id] = obj_idx + inference_state["obj_idx_to_id"][obj_idx] = obj_id + inference_state["obj_ids"] = list(inference_state["obj_id_to_idx"]) + # set up input and output structures for this object + inference_state["point_inputs_per_obj"][obj_idx] = {} + inference_state["mask_inputs_per_obj"][obj_idx] = {} + inference_state["output_dict_per_obj"][obj_idx] = { + "cond_frame_outputs": {}, # dict containing {frame_idx: } + "non_cond_frame_outputs": {}, # dict containing {frame_idx: } + } + inference_state["temp_output_dict_per_obj"][obj_idx] = { + "cond_frame_outputs": {}, # dict containing {frame_idx: } + "non_cond_frame_outputs": {}, # dict containing {frame_idx: } + } + return obj_idx + else: + raise RuntimeError( + f"Cannot add new object id {obj_id} after tracking starts. " + f"All existing object ids: {inference_state['obj_ids']}." + ) + + def _obj_idx_to_id(self, inference_state, obj_idx): + """Map model-side object index to client-side object id.""" + return inference_state["obj_idx_to_id"][obj_idx] + + def _get_obj_num(self, inference_state): + """Get the total number of unique object ids received so far in this session.""" + return len(inference_state["obj_idx_to_id"]) + + @torch.inference_mode() + def add_new_points_or_box( + self, + inference_state, + frame_idx, + obj_id, + points=None, + labels=None, + clear_old_points=True, + rel_coordinates=True, + use_prev_mem_frame=False, + normalize_coords=True, + box=None, + ): + """Add new points to a frame.""" + obj_idx = self._obj_id_to_idx(inference_state, obj_id) + point_inputs_per_frame = inference_state["point_inputs_per_obj"][obj_idx] + mask_inputs_per_frame = inference_state["mask_inputs_per_obj"][obj_idx] + + if (points is not None) != (labels is not None): + raise ValueError("points and labels must be provided together") + if points is None and box is None: + raise ValueError("at least one of points or box must be provided as input") + + if points is None: + points = torch.zeros(0, 2, dtype=torch.float32) + elif not isinstance(points, torch.Tensor): + points = torch.tensor(points, dtype=torch.float32) + if labels is None: + labels = torch.zeros(0, dtype=torch.int32) + elif not isinstance(labels, torch.Tensor): + labels = torch.tensor(labels, dtype=torch.int32) + if points.dim() == 2: + points = points.unsqueeze(0) # add batch dimension + if labels.dim() == 1: + labels = labels.unsqueeze(0) # add batch dimension + + if rel_coordinates: + # convert the points from relative coordinates to absolute coordinates + if points is not None: + points = points * self.image_size + if box is not None: + box = box * self.image_size + + # If `box` is provided, we add it as the first two points with labels 2 and 3 + # along with the user-provided points (consistent with how SAM 2 is trained). + if box is not None: + if not clear_old_points: + raise ValueError( + "cannot add box without clearing old points, since " + "box prompt must be provided before any point prompt " + "(please use clear_old_points=True instead)" + ) + if not isinstance(box, torch.Tensor): + box = torch.tensor(box, dtype=torch.float32, device=points.device) + box_coords = box.reshape(1, 2, 2) + box_labels = torch.tensor([2, 3], dtype=torch.int32, device=labels.device) + box_labels = box_labels.reshape(1, 2) + points = torch.cat([box_coords, points], dim=1) + labels = torch.cat([box_labels, labels], dim=1) + + points = points.to(inference_state["device"]) + labels = labels.to(inference_state["device"]) + + if not clear_old_points: + point_inputs = point_inputs_per_frame.get(frame_idx, None) + else: + point_inputs = None + point_inputs = concat_points(point_inputs, points, labels) + + point_inputs_per_frame[frame_idx] = point_inputs + mask_inputs_per_frame.pop(frame_idx, None) + # If this frame hasn't been tracked before, we treat it as an initial conditioning + # frame, meaning that the inputs points are to generate segments on this frame without + # using any memory from other frames, like in SAM. Otherwise (if it has been tracked), + # the input points will be used to correct the already tracked masks. + is_init_cond_frame = frame_idx not in inference_state["frames_already_tracked"] + # whether to track in reverse time order + if is_init_cond_frame: + reverse = False + else: + reverse = inference_state["frames_already_tracked"][frame_idx]["reverse"] + obj_output_dict = inference_state["output_dict_per_obj"][obj_idx] + obj_temp_output_dict = inference_state["temp_output_dict_per_obj"][obj_idx] + # Add a frame to conditioning output if it's an initial conditioning frame or + # if the model sees all frames receiving clicks/mask as conditioning frames. + is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond + storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs" + + # Limit to a maximum number of input points to the prompt encoder (to reduce domain gap) + num_points = point_inputs["point_coords"].size(1) + if num_points > self.max_point_num_in_prompt_enc > 0: + num_first = self.max_point_num_in_prompt_enc // 2 + num_last = self.max_point_num_in_prompt_enc - num_first + point_inputs["point_coords"] = torch.cat( + [ + point_inputs["point_coords"][:, :num_first], + point_inputs["point_coords"][:, -num_last:], + ], + dim=1, + ) + point_inputs["point_labels"] = torch.cat( + [ + point_inputs["point_labels"][:, :num_first], + point_inputs["point_labels"][:, -num_last:], + ], + dim=1, + ) + logging.warning( + f"Too many points ({num_points}) are provided on frame {frame_idx}. Only " + f"the first {num_first} points and the last {num_last} points will be used." + ) + # Get any previously predicted mask logits on this object and feed it along with + # the new clicks into the SAM mask decoder when `self.iter_use_prev_mask_pred=True`. + prev_sam_mask_logits = None + if self.iter_use_prev_mask_pred: + # lookup temporary output dict first, which contains the most recent output + # (if not found, then lookup conditioning and non-conditioning frame output) + prev_out = obj_temp_output_dict[storage_key].get(frame_idx) + if prev_out is None: + prev_out = obj_output_dict["cond_frame_outputs"].get(frame_idx) + if prev_out is None: + prev_out = obj_output_dict["non_cond_frame_outputs"].get(frame_idx) + + if prev_out is not None and prev_out["pred_masks"] is not None: + prev_sam_mask_logits = prev_out["pred_masks"].cuda(non_blocking=True) + # Clamp the scale of prev_sam_mask_logits to avoid rare numerical issues. + prev_sam_mask_logits = torch.clamp(prev_sam_mask_logits, -32.0, 32.0) + current_out, _ = self._run_single_frame_inference( + inference_state=inference_state, + output_dict=obj_output_dict, # run on the slice of a single object + frame_idx=frame_idx, + batch_size=1, # run on the slice of a single object + is_init_cond_frame=is_init_cond_frame, + point_inputs=point_inputs, + mask_inputs=None, + reverse=reverse, + # Skip the memory encoder when adding clicks or mask. We execute the memory encoder + # at the beginning of `propagate_in_video` (after user finalize their clicks). This + # allows us to enforce non-overlapping constraints on all objects before encoding + # them into memory. + run_mem_encoder=False, + prev_sam_mask_logits=prev_sam_mask_logits, + use_prev_mem_frame=use_prev_mem_frame, + ) + # Add the output to the output dict (to be used as future memory) + obj_temp_output_dict[storage_key][frame_idx] = current_out + + # Resize the output mask to the original video resolution + obj_ids = inference_state["obj_ids"] + consolidated_out = self._consolidate_temp_output_across_obj( + inference_state, + frame_idx, + is_cond=is_cond, + run_mem_encoder=False, + consolidate_at_video_res=True, + ) + _, video_res_masks = self._get_orig_video_res_output( + inference_state, consolidated_out["pred_masks_video_res"] + ) + low_res_masks = None # not needed by the demo + return frame_idx, obj_ids, low_res_masks, video_res_masks + + @torch.inference_mode() + def add_new_mask( + self, + inference_state, + frame_idx, + obj_id, + mask, + add_mask_to_memory=False, + ): + """Add new mask to a frame.""" + obj_idx = self._obj_id_to_idx(inference_state, obj_id) + point_inputs_per_frame = inference_state["point_inputs_per_obj"][obj_idx] + mask_inputs_per_frame = inference_state["mask_inputs_per_obj"][obj_idx] + + assert mask.dim() == 2 + mask_H, mask_W = mask.shape + mask_inputs_orig = mask[None, None] # add batch and channel dimension + mask_inputs_orig = mask_inputs_orig.float().to(inference_state["device"]) + + # resize the mask if it doesn't match the model's input mask size + if mask_H != self.input_mask_size or mask_W != self.input_mask_size: + mask_inputs = torch.nn.functional.interpolate( + mask_inputs_orig, + size=(self.input_mask_size, self.input_mask_size), + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for downsampling + ) + else: + mask_inputs = mask_inputs_orig + + # also get the mask at the original video resolution (for outputting) + video_H = inference_state["video_height"] + video_W = inference_state["video_width"] + if mask_H != video_H or mask_W != video_W: + mask_inputs_video_res = torch.nn.functional.interpolate( + mask_inputs_orig, + size=(video_H, video_W), + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for potential downsampling + ) + else: + mask_inputs_video_res = mask_inputs_orig + # convert mask_inputs_video_res to binary (threshold at 0.5 as it is in range 0~1) + mask_inputs_video_res = mask_inputs_video_res > 0.5 + + mask_inputs_per_frame[frame_idx] = mask_inputs_video_res + point_inputs_per_frame.pop(frame_idx, None) + # If this frame hasn't been tracked before, we treat it as an initial conditioning + # frame, meaning that the inputs points are to generate segments on this frame without + # using any memory from other frames, like in SAM. Otherwise (if it has been tracked), + # the input points will be used to correct the already tracked masks. + is_init_cond_frame = frame_idx not in inference_state["frames_already_tracked"] + # whether to track in reverse time order + if is_init_cond_frame: + reverse = False + else: + reverse = inference_state["frames_already_tracked"][frame_idx]["reverse"] + obj_output_dict = inference_state["output_dict_per_obj"][obj_idx] + obj_temp_output_dict = inference_state["temp_output_dict_per_obj"][obj_idx] + # Add a frame to conditioning output if it's an initial conditioning frame or + # if the model sees all frames receiving clicks/mask as conditioning frames. + is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond + storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs" + + current_out, _ = self._run_single_frame_inference( + inference_state=inference_state, + output_dict=obj_output_dict, # run on the slice of a single object + frame_idx=frame_idx, + batch_size=1, # run on the slice of a single object + is_init_cond_frame=is_init_cond_frame, + point_inputs=None, + mask_inputs=mask_inputs, + reverse=reverse, + # Skip the memory encoder when adding clicks or mask. We execute the memory encoder + # at the beginning of `propagate_in_video` (after user finalize their clicks). This + # allows us to enforce non-overlapping constraints on all objects before encoding + # them into memory. + run_mem_encoder=False, + ) + # We directly use the input mask at video resolution as the output mask for a better + # video editing experience (so that the masks don't change after each brushing). + # Here NO_OBJ_SCORE is a large negative value to represent the background and + # similarly -NO_OBJ_SCORE is a large positive value to represent the foreground. + current_out["pred_masks"] = None + current_out["pred_masks_video_res"] = torch.where( + mask_inputs_video_res, -NO_OBJ_SCORE, NO_OBJ_SCORE + ) + # Add the output to the output dict (to be used as future memory) + obj_temp_output_dict[storage_key][frame_idx] = current_out + # Remove the overlapping proportion of other objects' input masks on this frame + temp_output_dict_per_obj = inference_state["temp_output_dict_per_obj"] + for obj_idx2, obj_temp_output_dict2 in temp_output_dict_per_obj.items(): + if obj_idx2 == obj_idx: + continue + current_out2 = obj_temp_output_dict2[storage_key].get(frame_idx, None) + if current_out2 is not None and "pred_masks_video_res" in current_out2: + current_out2["pred_masks_video_res"] = torch.where( + mask_inputs_video_res, + NO_OBJ_SCORE, + current_out2["pred_masks_video_res"], + ) + + # Resize the output mask to the original video resolution + obj_ids = inference_state["obj_ids"] + consolidated_out = self._consolidate_temp_output_across_obj( + inference_state, + frame_idx, + is_cond=is_cond, + run_mem_encoder=False, + consolidate_at_video_res=True, + ) + _, video_res_masks = self._get_orig_video_res_output( + inference_state, consolidated_out["pred_masks_video_res"] + ) + low_res_masks = None # not needed by the demo + return frame_idx, obj_ids, low_res_masks, video_res_masks + + def add_new_points(self, *args, **kwargs): + """Deprecated method. Please use `add_new_points_or_box` instead.""" + return self.add_new_points_or_box(*args, **kwargs) + + def _get_orig_video_res_output(self, inference_state, any_res_masks): + """ + Resize the object scores to the original video resolution (video_res_masks) + and apply non-overlapping constraints for final output. + """ + device = inference_state["device"] + video_H = inference_state["video_height"] + video_W = inference_state["video_width"] + any_res_masks = any_res_masks.to(device, non_blocking=True) + if any_res_masks.shape[-2:] == (video_H, video_W): + video_res_masks = any_res_masks + else: + video_res_masks = torch.nn.functional.interpolate( + any_res_masks, + size=(video_H, video_W), + mode="bilinear", + align_corners=False, + ) + if self.non_overlap_masks_for_output: + video_res_masks = self._apply_non_overlapping_constraints(video_res_masks) + # potentially fill holes in the predicted masks + if self.fill_hole_area > 0: + video_res_masks = fill_holes_in_mask_scores( + video_res_masks, self.fill_hole_area + ) + return any_res_masks, video_res_masks + + def _consolidate_temp_output_across_obj( + self, + inference_state, + frame_idx, + is_cond, + run_mem_encoder, + consolidate_at_video_res=False, + ): + """ + Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on + a frame into a single output for all objects, including + 1) fill any missing objects either from `output_dict_per_obj` (if they exist in + `output_dict_per_obj` for this frame) or leave them as placeholder values + (if they don't exist in `output_dict_per_obj` for this frame); + 2) if specified, rerun memory encoder after apply non-overlapping constraints + on the object scores. + """ + batch_size = self._get_obj_num(inference_state) + storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs" + # Optionally, we allow consolidating the temporary outputs at the original + # video resolution (to provide a better editing experience for mask prompts). + if consolidate_at_video_res: + assert not run_mem_encoder, "memory encoder cannot run at video resolution" + consolidated_H = inference_state["video_height"] + consolidated_W = inference_state["video_width"] + consolidated_mask_key = "pred_masks_video_res" + else: + consolidated_H = consolidated_W = self.low_res_mask_size + consolidated_mask_key = "pred_masks" + + # Initialize `consolidated_out`. Its "maskmem_features" and "maskmem_pos_enc" + # will be added when rerunning the memory encoder after applying non-overlapping + # constraints to object scores. Its "pred_masks" are prefilled with a large + # negative value (NO_OBJ_SCORE) to represent missing objects. + consolidated_out = { + "maskmem_features": None, + "maskmem_pos_enc": None, + consolidated_mask_key: torch.full( + size=(batch_size, 1, consolidated_H, consolidated_W), + fill_value=NO_OBJ_SCORE, + dtype=torch.float32, + device=inference_state["storage_device"], + ), + "obj_ptr": torch.full( + size=(batch_size, self.hidden_dim), + fill_value=NO_OBJ_SCORE, + dtype=torch.float32, + device=inference_state["device"], + ), + "object_score_logits": torch.full( + size=(batch_size, 1), + # default to 10.0 for object_score_logits, i.e. assuming the object is + # present as sigmoid(10)=1, same as in `predict_masks` of `MaskDecoder` + fill_value=10.0, + dtype=torch.float32, + device=inference_state["device"], + ), + } + if self.use_memory_selection: + consolidated_out["iou_score"] = torch.full( + size=(batch_size, 1), + fill_value=0.0, + dtype=torch.float32, + device=inference_state["device"], + ) + empty_mask_ptr = None + for obj_idx in range(batch_size): + obj_temp_output_dict = inference_state["temp_output_dict_per_obj"][obj_idx] + obj_output_dict = inference_state["output_dict_per_obj"][obj_idx] + out = obj_temp_output_dict[storage_key].get(frame_idx, None) + # If the object doesn't appear in "temp_output_dict_per_obj" on this frame, + # we fall back and look up its previous output in "output_dict_per_obj". + # We look up both "cond_frame_outputs" and "non_cond_frame_outputs" in + # "output_dict_per_obj" to find a previous output for this object. + if out is None: + out = obj_output_dict["cond_frame_outputs"].get(frame_idx, None) + if out is None: + out = obj_output_dict["non_cond_frame_outputs"].get(frame_idx, None) + # If the object doesn't appear in "output_dict_per_obj" either, we skip it + # and leave its mask scores to the default scores (i.e. the NO_OBJ_SCORE + # placeholder above) and set its object pointer to be a dummy pointer. + if out is None: + # Fill in dummy object pointers for those objects without any inputs or + # tracking outcomes on this frame (only do it under `run_mem_encoder=True`, + # i.e. when we need to build the memory for tracking). + if run_mem_encoder: + if empty_mask_ptr is None: + empty_mask_ptr = self._get_empty_mask_ptr( + inference_state, frame_idx + ) + # fill object pointer with a dummy pointer (based on an empty mask) + consolidated_out["obj_ptr"][obj_idx : obj_idx + 1] = empty_mask_ptr + continue + # Add the temporary object output mask to consolidated output mask + # (use "pred_masks_video_res" if it's available) + obj_mask = out.get("pred_masks_video_res", out["pred_masks"]) + consolidated_pred_masks = consolidated_out[consolidated_mask_key] + if obj_mask.shape[-2:] == consolidated_pred_masks.shape[-2:]: + consolidated_pred_masks[obj_idx : obj_idx + 1] = obj_mask + else: + # Resize first if temporary object mask has a different resolution + is_downsampling = "pred_masks_video_res" in out + resized_obj_mask = torch.nn.functional.interpolate( + obj_mask, + size=consolidated_pred_masks.shape[-2:], + mode="bilinear", + align_corners=False, + antialias=is_downsampling, # use antialias for downsampling + ) + consolidated_pred_masks[obj_idx : obj_idx + 1] = resized_obj_mask + consolidated_out["obj_ptr"][obj_idx : obj_idx + 1] = out["obj_ptr"] + consolidated_out["object_score_logits"][obj_idx : obj_idx + 1] = out[ + "object_score_logits" + ] + if self.use_memory_selection: + consolidated_out["iou_score"][obj_idx : obj_idx + 1] = out["iou_score"] + # Optionally, apply non-overlapping constraints on the consolidated scores + # and rerun the memory encoder + if run_mem_encoder: + device = inference_state["device"] + high_res_masks = torch.nn.functional.interpolate( + consolidated_out["pred_masks"].to(device, non_blocking=True), + size=(self.image_size, self.image_size), + mode="bilinear", + align_corners=False, + ) + high_res_masks = self._apply_non_overlapping_constraints(high_res_masks) + maskmem_features, maskmem_pos_enc = self._run_memory_encoder( + inference_state=inference_state, + frame_idx=frame_idx, + batch_size=batch_size, + high_res_masks=high_res_masks, + object_score_logits=consolidated_out["object_score_logits"], + is_mask_from_pts=True, # these frames are what the user interacted with + ) + consolidated_out["maskmem_features"] = maskmem_features + consolidated_out["maskmem_pos_enc"] = maskmem_pos_enc + + return consolidated_out + + def _get_empty_mask_ptr(self, inference_state, frame_idx): + """Get a dummy object pointer based on an empty mask on the current frame.""" + # A dummy (empty) mask with a single object + batch_size = 1 + mask_inputs = torch.zeros( + (batch_size, 1, self.image_size, self.image_size), + dtype=torch.float32, + device=inference_state["device"], + ) + + # Retrieve correct image features + ( + image, + _, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + ) = self._get_image_feature(inference_state, frame_idx, batch_size) + + # Feed the empty mask and image feature above to get a dummy object pointer + current_out = self.track_step( + frame_idx=frame_idx, + is_init_cond_frame=True, + current_vision_feats=current_vision_feats, + current_vision_pos_embeds=current_vision_pos_embeds, + feat_sizes=feat_sizes, + image=image, + point_inputs=None, + mask_inputs=mask_inputs, + gt_masks=None, + frames_to_add_correction_pt=[], + output_dict={ + "cond_frame_outputs": {}, + "non_cond_frame_outputs": {}, + }, + num_frames=inference_state["num_frames"], + track_in_reverse=False, + run_mem_encoder=False, + prev_sam_mask_logits=None, + ) + return current_out["obj_ptr"] + + @torch.inference_mode() + def propagate_in_video_preflight(self, inference_state, run_mem_encoder=True): + """Prepare inference_state and consolidate temporary outputs before tracking.""" + # Tracking has started and we don't allow adding new objects until session is reset. + inference_state["tracking_has_started"] = True + batch_size = self._get_obj_num(inference_state) + + # Consolidate per-object temporary outputs in "temp_output_dict_per_obj" and + # add them into "output_dict". + temp_output_dict_per_obj = inference_state["temp_output_dict_per_obj"] + output_dict = inference_state["output_dict"] + # "consolidated_frame_inds" contains indices of those frames where consolidated + # temporary outputs have been added (either in this call or any previous calls + # to `propagate_in_video_preflight`). + consolidated_frame_inds = inference_state["consolidated_frame_inds"] + for is_cond in [False, True]: + # Separately consolidate conditioning and non-conditioning temp outptus + storage_key = "cond_frame_outputs" if is_cond else "non_cond_frame_outputs" + # Find all the frames that contain temporary outputs for any objects + # (these should be the frames that have just received clicks for mask inputs + # via `add_new_points` or `add_new_mask`) + temp_frame_inds = set() + for obj_temp_output_dict in temp_output_dict_per_obj.values(): + temp_frame_inds.update(obj_temp_output_dict[storage_key].keys()) + consolidated_frame_inds[storage_key].update(temp_frame_inds) + # consolidate the temprary output across all objects on this frame + for frame_idx in temp_frame_inds: + consolidated_out = self._consolidate_temp_output_across_obj( + inference_state, + frame_idx, + is_cond=is_cond, + run_mem_encoder=run_mem_encoder, + ) + # merge them into "output_dict" and also create per-object slices + output_dict[storage_key][frame_idx] = consolidated_out + self._add_output_per_object( + inference_state, frame_idx, consolidated_out, storage_key + ) + clear_non_cond_mem = self.clear_non_cond_mem_around_input and ( + self.clear_non_cond_mem_for_multi_obj or batch_size <= 1 + ) + if clear_non_cond_mem: + # clear non-conditioning memory of the surrounding frames + self._clear_non_cond_mem_around_input(inference_state, frame_idx) + + # clear temporary outputs in `temp_output_dict_per_obj` + for obj_temp_output_dict in temp_output_dict_per_obj.values(): + obj_temp_output_dict[storage_key].clear() + + # edge case: if an output is added to "cond_frame_outputs", we remove any prior + # output on the same frame in "non_cond_frame_outputs" + for frame_idx in output_dict["cond_frame_outputs"]: + output_dict["non_cond_frame_outputs"].pop(frame_idx, None) + for obj_output_dict in inference_state["output_dict_per_obj"].values(): + for frame_idx in obj_output_dict["cond_frame_outputs"]: + obj_output_dict["non_cond_frame_outputs"].pop(frame_idx, None) + for frame_idx in consolidated_frame_inds["cond_frame_outputs"]: + assert frame_idx in output_dict["cond_frame_outputs"] + consolidated_frame_inds["non_cond_frame_outputs"].discard(frame_idx) + + # Make sure that the frame indices in "consolidated_frame_inds" are exactly those frames + # with either points or mask inputs (which should be true under a correct demo workflow). + all_consolidated_frame_inds = ( + consolidated_frame_inds["cond_frame_outputs"] + | consolidated_frame_inds["non_cond_frame_outputs"] + ) + input_frames_inds = set() + for point_inputs_per_frame in inference_state["point_inputs_per_obj"].values(): + input_frames_inds.update(point_inputs_per_frame.keys()) + for mask_inputs_per_frame in inference_state["mask_inputs_per_obj"].values(): + input_frames_inds.update(mask_inputs_per_frame.keys()) + assert all_consolidated_frame_inds == input_frames_inds + # Record the first interacted frame index (for tracking start) + if inference_state["first_ann_frame_idx"] is None: + inference_state["first_ann_frame_idx"] = min( + input_frames_inds, default=None + ) + # In case `first_ann_frame_idx` is not in the conditioning frames (e.g. because + # we cleared the input points on that frame), pick the first conditioning frame + if ( + inference_state["first_ann_frame_idx"] + not in output_dict["cond_frame_outputs"] + ): + inference_state["first_ann_frame_idx"] = min( + output_dict["cond_frame_outputs"], default=None + ) + + def _get_processing_order( + self, inference_state, start_frame_idx, max_frame_num_to_track, reverse + ): + num_frames = inference_state["num_frames"] + # set start index, end index, and processing order + if self.always_start_from_first_ann_frame: + # in this case, we always start tracking from the frame where we receive + # the initial annotation and ignore the provided start_frame_idx + start_frame_idx = inference_state["first_ann_frame_idx"] + if start_frame_idx is None: + # default: start from the earliest frame with input points + start_frame_idx = min(inference_state["output_dict"]["cond_frame_outputs"]) + if max_frame_num_to_track is None: + # default: track all the frames in the video + max_frame_num_to_track = num_frames + if reverse: + end_frame_idx = max(start_frame_idx - max_frame_num_to_track, 0) + if start_frame_idx > 0: + processing_order = range(start_frame_idx, end_frame_idx - 1, -1) + else: + # this is the edge case where we start from frame 0 and track in reverse order; + # in this case, we track a single frame (frame 0) + processing_order = [0] + else: + end_frame_idx = min( + start_frame_idx + max_frame_num_to_track, num_frames - 1 + ) + processing_order = range(start_frame_idx, end_frame_idx + 1) + return processing_order + + @torch.inference_mode() + def propagate_in_video( + self, + inference_state, + start_frame_idx, + max_frame_num_to_track, + reverse, + tqdm_disable=False, + obj_ids=None, + run_mem_encoder=True, + propagate_preflight=False, + ): + """Propagate the input points across frames to track in the entire video.""" + if propagate_preflight: + self.propagate_in_video_preflight(inference_state) + # NOTE: This is a copy from the parent class, except that we return object scores as well. + output_dict = inference_state["output_dict"] + consolidated_frame_inds = inference_state["consolidated_frame_inds"] + if obj_ids is not None: + raise NotImplementedError( + "Per-object tracking yet for batched inference if not implemented." + ) + obj_ids = inference_state["obj_ids"] + batch_size = self._get_obj_num(inference_state) + if len(output_dict["cond_frame_outputs"]) == 0: + raise RuntimeError("No points are provided; please add points first") + clear_non_cond_mem = self.clear_non_cond_mem_around_input and ( + self.clear_non_cond_mem_for_multi_obj or batch_size <= 1 + ) + + processing_order = self._get_processing_order( + inference_state, + start_frame_idx, + max_frame_num_to_track, + reverse, + ) + + for frame_idx in tqdm( + processing_order, desc="propagate in video", disable=tqdm_disable + ): + # We skip those frames already in consolidated outputs (these are frames + # that received input clicks or mask). Note that we cannot directly run + # batched forward on them via `_run_single_frame_inference` because the + # number of clicks on each object might be different. + if frame_idx in consolidated_frame_inds["cond_frame_outputs"]: + storage_key = "cond_frame_outputs" + current_out = output_dict[storage_key][frame_idx] + pred_masks = current_out["pred_masks"] + obj_scores = current_out["object_score_logits"] + if clear_non_cond_mem: + # clear non-conditioning memory of the surrounding frames + self._clear_non_cond_mem_around_input(inference_state, frame_idx) + elif frame_idx in consolidated_frame_inds["non_cond_frame_outputs"]: + storage_key = "non_cond_frame_outputs" + current_out = output_dict[storage_key][frame_idx] + pred_masks = current_out["pred_masks"] + obj_scores = current_out["object_score_logits"] + else: + storage_key = "non_cond_frame_outputs" + current_out, pred_masks = self._run_single_frame_inference( + inference_state=inference_state, + output_dict=output_dict, + frame_idx=frame_idx, + batch_size=batch_size, + is_init_cond_frame=False, + point_inputs=None, + mask_inputs=None, + reverse=reverse, + run_mem_encoder=run_mem_encoder, + ) + obj_scores = current_out["object_score_logits"] + output_dict[storage_key][frame_idx] = current_out + # Create slices of per-object outputs for subsequent interaction with each + # individual object after tracking. + self._add_output_per_object( + inference_state, frame_idx, current_out, storage_key + ) + inference_state["frames_already_tracked"][frame_idx] = {"reverse": reverse} + + # Resize the output mask to the original video resolution (we directly use + # the mask scores on GPU for output to avoid any CPU conversion in between) + low_res_masks, video_res_masks = self._get_orig_video_res_output( + inference_state, pred_masks + ) + yield frame_idx, obj_ids, low_res_masks, video_res_masks, obj_scores + + def _add_output_per_object( + self, inference_state, frame_idx, current_out, storage_key + ): + """ + Split a multi-object output into per-object output slices and add them into + `output_dict_per_obj`. The resulting slices share the same tensor storage. + """ + maskmem_features = current_out["maskmem_features"] + assert maskmem_features is None or isinstance(maskmem_features, torch.Tensor) + + maskmem_pos_enc = current_out["maskmem_pos_enc"] + assert maskmem_pos_enc is None or isinstance(maskmem_pos_enc, list) + + output_dict_per_obj = inference_state["output_dict_per_obj"] + for obj_idx, obj_output_dict in output_dict_per_obj.items(): + obj_slice = slice(obj_idx, obj_idx + 1) + obj_out = { + "maskmem_features": None, + "maskmem_pos_enc": None, + "pred_masks": current_out["pred_masks"][obj_slice], + "obj_ptr": current_out["obj_ptr"][obj_slice], + "object_score_logits": current_out["object_score_logits"][obj_slice], + } + if self.use_memory_selection: + obj_out["iou_score"] = current_out["iou_score"][obj_slice] + if maskmem_features is not None: + obj_out["maskmem_features"] = maskmem_features[obj_slice] + if maskmem_pos_enc is not None: + obj_out["maskmem_pos_enc"] = [x[obj_slice] for x in maskmem_pos_enc] + obj_output_dict[storage_key][frame_idx] = obj_out + + @torch.inference_mode() + def clear_all_points_in_frame( + self, inference_state, frame_idx, obj_id, need_output=True + ): + """Remove all input points or mask in a specific frame for a given object.""" + obj_idx = self._obj_id_to_idx(inference_state, obj_id) + + # Clear the conditioning information on the given frame + inference_state["point_inputs_per_obj"][obj_idx].pop(frame_idx, None) + inference_state["mask_inputs_per_obj"][obj_idx].pop(frame_idx, None) + + temp_output_dict_per_obj = inference_state["temp_output_dict_per_obj"] + temp_output_dict_per_obj[obj_idx]["cond_frame_outputs"].pop(frame_idx, None) + temp_output_dict_per_obj[obj_idx]["non_cond_frame_outputs"].pop(frame_idx, None) + + # Check and see if there are still any inputs left on this frame + batch_size = self._get_obj_num(inference_state) + frame_has_input = False + for obj_idx2 in range(batch_size): + if frame_idx in inference_state["point_inputs_per_obj"][obj_idx2]: + frame_has_input = True + break + if frame_idx in inference_state["mask_inputs_per_obj"][obj_idx2]: + frame_has_input = True + break + + # If this frame has no remaining inputs for any objects, we further clear its + # conditioning frame status + if not frame_has_input: + output_dict = inference_state["output_dict"] + consolidated_frame_inds = inference_state["consolidated_frame_inds"] + consolidated_frame_inds["cond_frame_outputs"].discard(frame_idx) + consolidated_frame_inds["non_cond_frame_outputs"].discard(frame_idx) + # Remove the frame's conditioning output (possibly downgrading it to non-conditioning) + out = output_dict["cond_frame_outputs"].pop(frame_idx, None) + if out is not None: + # The frame is not a conditioning frame anymore since it's not receiving inputs, + # so we "downgrade" its output (if exists) to a non-conditioning frame output. + output_dict["non_cond_frame_outputs"][frame_idx] = out + inference_state["frames_already_tracked"].pop(frame_idx, None) + # Similarly, do it for the sliced output on each object. + for obj_idx2 in range(batch_size): + obj_output_dict = inference_state["output_dict_per_obj"][obj_idx2] + obj_out = obj_output_dict["cond_frame_outputs"].pop(frame_idx, None) + if obj_out is not None: + obj_output_dict["non_cond_frame_outputs"][frame_idx] = obj_out + + # If all the conditioning frames have been removed, we also clear the tracking outputs + if len(output_dict["cond_frame_outputs"]) == 0: + self._reset_tracking_results(inference_state) + + if not need_output: + return + # Finally, output updated masks per object (after removing the inputs above) + obj_ids = inference_state["obj_ids"] + is_cond = any( + frame_idx in obj_temp_output_dict["cond_frame_outputs"] + for obj_temp_output_dict in temp_output_dict_per_obj.values() + ) + consolidated_out = self._consolidate_temp_output_across_obj( + inference_state, + frame_idx, + is_cond=is_cond, + run_mem_encoder=False, + consolidate_at_video_res=True, + ) + _, video_res_masks = self._get_orig_video_res_output( + inference_state, consolidated_out["pred_masks_video_res"] + ) + low_res_masks = None # not needed by the demo + return frame_idx, obj_ids, low_res_masks, video_res_masks + + @torch.inference_mode() + def clear_all_points_in_video(self, inference_state): + """Remove all input points or mask in all frames throughout the video.""" + self._reset_tracking_results(inference_state) + # Remove all object ids + inference_state["obj_id_to_idx"].clear() + inference_state["obj_idx_to_id"].clear() + inference_state["obj_ids"].clear() + inference_state["point_inputs_per_obj"].clear() + inference_state["mask_inputs_per_obj"].clear() + inference_state["output_dict_per_obj"].clear() + inference_state["temp_output_dict_per_obj"].clear() + + def _reset_tracking_results(self, inference_state): + """Reset all tracking inputs and results across the videos.""" + for v in inference_state["point_inputs_per_obj"].values(): + v.clear() + for v in inference_state["mask_inputs_per_obj"].values(): + v.clear() + for v in inference_state["output_dict_per_obj"].values(): + v["cond_frame_outputs"].clear() + v["non_cond_frame_outputs"].clear() + for v in inference_state["temp_output_dict_per_obj"].values(): + v["cond_frame_outputs"].clear() + v["non_cond_frame_outputs"].clear() + inference_state["output_dict"]["cond_frame_outputs"].clear() + inference_state["output_dict"]["non_cond_frame_outputs"].clear() + inference_state["consolidated_frame_inds"]["cond_frame_outputs"].clear() + inference_state["consolidated_frame_inds"]["non_cond_frame_outputs"].clear() + inference_state["tracking_has_started"] = False + inference_state["frames_already_tracked"].clear() + inference_state["first_ann_frame_idx"] = None + + def _get_image_feature(self, inference_state, frame_idx, batch_size): + """Compute the image features on a given frame.""" + # Look up in the cache + image, backbone_out = inference_state["cached_features"].get( + frame_idx, (None, None) + ) + if backbone_out is None: + if self.backbone is None: + raise RuntimeError( + f"Image features for frame {frame_idx} are not cached. " + "Please run inference on this frame first." + ) + else: + # Cache miss -- we will run inference on a single image + image = inference_state["images"][frame_idx].cuda().float().unsqueeze(0) + backbone_out = self.forward_image(image) + # Cache the most recent frame's feature (for repeated interactions with + # a frame; we can use an LRU cache for more frames in the future). + inference_state["cached_features"] = {frame_idx: (image, backbone_out)} + if "tracker_backbone_out" in backbone_out: + backbone_out = backbone_out["tracker_backbone_out"] # get backbone output + + # expand the features to have the same dimension as the number of objects + expanded_image = image.expand(batch_size, -1, -1, -1) + expanded_backbone_out = { + "backbone_fpn": backbone_out["backbone_fpn"].copy(), + "vision_pos_enc": backbone_out["vision_pos_enc"].copy(), + } + for i, feat in enumerate(expanded_backbone_out["backbone_fpn"]): + feat = feat.expand(batch_size, -1, -1, -1) + expanded_backbone_out["backbone_fpn"][i] = feat + for i, pos in enumerate(expanded_backbone_out["vision_pos_enc"]): + pos = pos.expand(batch_size, -1, -1, -1) + expanded_backbone_out["vision_pos_enc"][i] = pos + + features = self._prepare_backbone_features(expanded_backbone_out) + features = (expanded_image,) + features + return features + + def _run_single_frame_inference( + self, + inference_state, + output_dict, + frame_idx, + batch_size, + is_init_cond_frame, + point_inputs, + mask_inputs, + reverse, + run_mem_encoder, + prev_sam_mask_logits=None, + use_prev_mem_frame=True, + ): + """Run tracking on a single frame based on current inputs and previous memory.""" + # Retrieve correct image features + ( + image, + _, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + ) = self._get_image_feature(inference_state, frame_idx, batch_size) + + # point and mask should not appear as input simultaneously on the same frame + assert point_inputs is None or mask_inputs is None + current_out = self.track_step( + frame_idx=frame_idx, + is_init_cond_frame=is_init_cond_frame, + current_vision_feats=current_vision_feats, + current_vision_pos_embeds=current_vision_pos_embeds, + feat_sizes=feat_sizes, + image=image, + point_inputs=point_inputs, + mask_inputs=mask_inputs, + output_dict=output_dict, + num_frames=inference_state["num_frames"], + track_in_reverse=reverse, + run_mem_encoder=run_mem_encoder, + prev_sam_mask_logits=prev_sam_mask_logits, + use_prev_mem_frame=use_prev_mem_frame, + ) + + # optionally offload the output to CPU memory to save GPU space + storage_device = inference_state["storage_device"] + maskmem_features = current_out["maskmem_features"] + if maskmem_features is not None: + maskmem_features = maskmem_features.to(torch.bfloat16) + maskmem_features = maskmem_features.to(storage_device, non_blocking=True) + pred_masks_gpu = current_out["pred_masks"] + pred_masks = pred_masks_gpu.to(storage_device, non_blocking=True) + # "maskmem_pos_enc" is the same across frames, so we only need to store one copy of it + maskmem_pos_enc = self._get_maskmem_pos_enc(inference_state, current_out) + # object pointer is a small tensor, so we always keep it on GPU memory for fast access + obj_ptr = current_out["obj_ptr"] + object_score_logits = current_out["object_score_logits"] + # make a compact version of this frame's output to reduce the state size + compact_current_out = { + "maskmem_features": maskmem_features, + "maskmem_pos_enc": maskmem_pos_enc, + "pred_masks": pred_masks, + "obj_ptr": obj_ptr, + "object_score_logits": object_score_logits, + } + if self.use_memory_selection: + compact_current_out["iou_score"] = current_out["iou_score"] + compact_current_out["eff_iou_score"] = current_out["eff_iou_score"] + return compact_current_out, pred_masks_gpu + + def _run_memory_encoder( + self, + inference_state, + frame_idx, + batch_size, + high_res_masks, + object_score_logits, + is_mask_from_pts, + ): + """ + Run the memory encoder on `high_res_masks`. This is usually after applying + non-overlapping constraints to object scores. Since their scores changed, their + memory also need to be computed again with the memory encoder. + """ + # Retrieve correct image features + image, _, current_vision_feats, _, feat_sizes = self._get_image_feature( + inference_state, frame_idx, batch_size + ) + maskmem_features, maskmem_pos_enc = self._encode_new_memory( + image=image, + current_vision_feats=current_vision_feats, + feat_sizes=feat_sizes, + pred_masks_high_res=high_res_masks, + object_score_logits=object_score_logits, + is_mask_from_pts=is_mask_from_pts, + ) + + # optionally offload the output to CPU memory to save GPU space + storage_device = inference_state["storage_device"] + maskmem_features = maskmem_features.to(torch.bfloat16) + maskmem_features = maskmem_features.to(storage_device, non_blocking=True) + # "maskmem_pos_enc" is the same across frames, so we only need to store one copy of it + maskmem_pos_enc = self._get_maskmem_pos_enc( + inference_state, {"maskmem_pos_enc": maskmem_pos_enc} + ) + return maskmem_features, maskmem_pos_enc + + def _get_maskmem_pos_enc(self, inference_state, current_out): + """ + `maskmem_pos_enc` is the same across frames and objects, so we cache it as + a constant in the inference session to reduce session storage size. + """ + model_constants = inference_state["constants"] + # "out_maskmem_pos_enc" should be either a list of tensors or None + out_maskmem_pos_enc = current_out["maskmem_pos_enc"] + if out_maskmem_pos_enc is not None: + if "maskmem_pos_enc" not in model_constants: + assert isinstance(out_maskmem_pos_enc, list) + # only take the slice for one object, since it's same across objects + maskmem_pos_enc = [x[0:1].clone() for x in out_maskmem_pos_enc] + model_constants["maskmem_pos_enc"] = maskmem_pos_enc + else: + maskmem_pos_enc = model_constants["maskmem_pos_enc"] + # expand the cached maskmem_pos_enc to the actual batch size + batch_size = out_maskmem_pos_enc[0].size(0) + expanded_maskmem_pos_enc = [ + x.expand(batch_size, -1, -1, -1) for x in maskmem_pos_enc + ] + else: + expanded_maskmem_pos_enc = None + return expanded_maskmem_pos_enc + + @torch.inference_mode() + def remove_object(self, inference_state, obj_id, strict=False, need_output=True): + """ + Remove an object id from the tracking state. If strict is True, we check whether + the object id actually exists and raise an error if it doesn't exist. + """ + old_obj_idx_to_rm = inference_state["obj_id_to_idx"].get(obj_id, None) + updated_frames = [] + # Check whether this object_id to remove actually exists and possibly raise an error. + if old_obj_idx_to_rm is None: + if not strict: + return inference_state["obj_ids"], updated_frames + raise RuntimeError( + f"Cannot remove object id {obj_id} as it doesn't exist. " + f"All existing object ids: {inference_state['obj_ids']}." + ) + + # If this is the only remaining object id, we simply reset the state. + if len(inference_state["obj_id_to_idx"]) == 1: + self.clear_all_points_in_video(inference_state) + return inference_state["obj_ids"], updated_frames + + # There are still remaining objects after removing this object id. In this case, + # we need to delete the object storage from inference state tensors. + # Step 0: clear the input on those frames where this object id has point or mask input + # (note that this step is required as it might downgrade conditioning frames to + # non-conditioning ones) + obj_input_frames_inds = set() + obj_input_frames_inds.update( + inference_state["point_inputs_per_obj"][old_obj_idx_to_rm] + ) + obj_input_frames_inds.update( + inference_state["mask_inputs_per_obj"][old_obj_idx_to_rm] + ) + for frame_idx in obj_input_frames_inds: + self.clear_all_points_in_frame( + inference_state, frame_idx, obj_id, need_output=False + ) + + # Step 1: Update the object id mapping (note that it must be done after Step 0, + # since Step 0 still requires the old object id mappings in inference_state) + old_obj_ids = inference_state["obj_ids"] + old_obj_inds = list(range(len(old_obj_ids))) + remain_old_obj_inds = old_obj_inds.copy() + remain_old_obj_inds.remove(old_obj_idx_to_rm) + new_obj_ids = [old_obj_ids[old_idx] for old_idx in remain_old_obj_inds] + new_obj_inds = list(range(len(new_obj_ids))) + # build new mappings + old_idx_to_new_idx = dict(zip(remain_old_obj_inds, new_obj_inds)) + inference_state["obj_id_to_idx"] = dict(zip(new_obj_ids, new_obj_inds)) + inference_state["obj_idx_to_id"] = dict(zip(new_obj_inds, new_obj_ids)) + inference_state["obj_ids"] = new_obj_ids + + # Step 2: For per-object tensor storage, we shift their obj_idx in the dict keys. + # (note that "consolidated_frame_inds" doesn't need to be updated in this step as + # it's already handled in Step 0) + def _map_keys(container): + new_kvs = [] + for k in old_obj_inds: + v = container.pop(k) + if k in old_idx_to_new_idx: + new_kvs.append((old_idx_to_new_idx[k], v)) + container.update(new_kvs) + + _map_keys(inference_state["point_inputs_per_obj"]) + _map_keys(inference_state["mask_inputs_per_obj"]) + _map_keys(inference_state["output_dict_per_obj"]) + _map_keys(inference_state["temp_output_dict_per_obj"]) + + # Step 3: For packed tensor storage, we index the remaining ids and rebuild the per-object slices. + def _slice_state(output_dict, storage_key): + for frame_idx, out in output_dict[storage_key].items(): + out["maskmem_features"] = out["maskmem_features"][remain_old_obj_inds] + out["maskmem_pos_enc"] = [ + x[remain_old_obj_inds] for x in out["maskmem_pos_enc"] + ] + # "maskmem_pos_enc" is the same across frames, so we only need to store one copy of it + out["maskmem_pos_enc"] = self._get_maskmem_pos_enc(inference_state, out) + out["pred_masks"] = out["pred_masks"][remain_old_obj_inds] + out["obj_ptr"] = out["obj_ptr"][remain_old_obj_inds] + out["object_score_logits"] = out["object_score_logits"][ + remain_old_obj_inds + ] + if self.use_memory_selection: + out["iou_score"] = out["iou_score"][remain_old_obj_inds] + out["eff_iou_score"] = self.cal_mem_score( + out["object_score_logits"], out["iou_score"] + ) # recalculate the memory frame score + # also update the per-object slices + self._add_output_per_object( + inference_state, frame_idx, out, storage_key + ) + + _slice_state(inference_state["output_dict"], "cond_frame_outputs") + _slice_state(inference_state["output_dict"], "non_cond_frame_outputs") + + # Step 4: Further collect the outputs on those frames in `obj_input_frames_inds`, which + # could show an updated mask for objects previously occluded by the object being removed + if need_output: + temp_output_dict_per_obj = inference_state["temp_output_dict_per_obj"] + for frame_idx in obj_input_frames_inds: + is_cond = any( + frame_idx in obj_temp_output_dict["cond_frame_outputs"] + for obj_temp_output_dict in temp_output_dict_per_obj.values() + ) + consolidated_out = self._consolidate_temp_output_across_obj( + inference_state, + frame_idx, + is_cond=is_cond, + run_mem_encoder=False, + consolidate_at_video_res=True, + ) + _, video_res_masks = self._get_orig_video_res_output( + inference_state, consolidated_out["pred_masks_video_res"] + ) + updated_frames.append((frame_idx, video_res_masks)) + + return inference_state["obj_ids"], updated_frames + + def _clear_non_cond_mem_around_input(self, inference_state, frame_idx): + """ + Remove the non-conditioning memory around the input frame. When users provide + correction clicks, the surrounding frames' non-conditioning memories can still + contain outdated object appearance information and could confuse the model. + + This method clears those non-conditioning memories surrounding the interacted + frame to avoid giving the model both old and new information about the object. + """ + r = self.memory_temporal_stride_for_eval + frame_idx_begin = frame_idx - r * self.num_maskmem + frame_idx_end = frame_idx + r * self.num_maskmem + batch_size = self._get_obj_num(inference_state) + for obj_idx in range(batch_size): + obj_output_dict = inference_state["output_dict_per_obj"][obj_idx] + non_cond_frame_outputs = obj_output_dict["non_cond_frame_outputs"] + for t in range(frame_idx_begin, frame_idx_end + 1): + non_cond_frame_outputs.pop(t, None) + + def _suppress_shrinked_masks( + self, pred_masks, new_pred_masks, shrink_threshold=0.3 + ): + area_before = (pred_masks > 0).sum(dim=(-1, -2)) + area_after = (new_pred_masks > 0).sum(dim=(-1, -2)) + area_before = torch.clamp(area_before, min=1.0) + area_ratio = area_after / area_before + keep = area_ratio >= shrink_threshold + keep_mask = keep[..., None, None].expand_as(pred_masks) + pred_masks_after = torch.where( + keep_mask, pred_masks, torch.clamp(pred_masks, max=-10.0) + ) + return pred_masks_after + + def _suppress_object_pw_area_shrinkage(self, pred_masks): + """ + This function suppresses masks that shrink in area after applying pixelwise non-overlapping constriants. + Note that the final output can still be overlapping. + """ + # Apply pixel-wise non-overlapping constraint based on mask scores + pixel_level_non_overlapping_masks = super()._apply_non_overlapping_constraints( + pred_masks + ) + # Fully suppress masks with high shrinkage (probably noisy) based on the pixel wise non-overlapping constraints + # NOTE: The output of this function can be a no op if none of the masks shrinked by a large factor. + pred_masks = self._suppress_shrinked_masks( + pred_masks, pixel_level_non_overlapping_masks + ) + return pred_masks + + def _apply_object_wise_non_overlapping_constraints( + self, pred_masks, obj_scores, background_value=-10.0 + ): + """ + Applies non-overlapping constraints object wise (i.e. only one object can claim the overlapping region) + """ + # Replace pixel scores with object scores + pred_masks_single_score = torch.where( + pred_masks > 0, obj_scores[..., None, None], background_value + ) + # Apply pixel-wise non-overlapping constraint based on mask scores + pixel_level_non_overlapping_masks = super()._apply_non_overlapping_constraints( + pred_masks_single_score + ) + # Replace object scores with pixel scores. Note, that now only one object can claim the overlapping region + pred_masks = torch.where( + pixel_level_non_overlapping_masks > 0, + pred_masks, + torch.clamp(pred_masks, max=background_value), + ) + return pred_masks diff --git a/src/nn/segearth_ov3/sam3/model/sam3_video_base.py b/src/nn/segearth_ov3/sam3/model/sam3_video_base.py new file mode 100644 index 0000000..e61969f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_video_base.py @@ -0,0 +1,1767 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import datetime +import logging +import math +import os +from collections import defaultdict +from copy import deepcopy +from enum import Enum +from typing import Any, Dict, List, Set + +import numpy as np +import numpy.typing as npt +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from sam3 import perflib +from sam3.logger import get_logger +from sam3.model.box_ops import fast_diag_box_iou +from sam3.model.data_misc import BatchedDatapoint +from sam3.model.sam3_tracker_utils import fill_holes_in_mask_scores, mask_to_box +from sam3.perflib.masks_ops import mask_iou +from sam3.train.masks_ops import rle_encode +from torch import nn, Tensor + +logger = get_logger(__name__) + + +class MaskletConfirmationStatus(Enum): + UNCONFIRMED = 1 # newly added masklet, not confirmed by any detection yet + CONFIRMED = 2 # confirmed by at least one detection + + +class Sam3VideoBase(nn.Module): + def __init__( + self, + detector: nn.Module, + tracker: nn.Module, + # prob threshold for detection outputs -- only keep detections above this threshold + # enters NMS and det-to-track matching + score_threshold_detection=0.5, + # IoU threshold for detection NMS + det_nms_thresh=0.0, + # IoU threshold for det-to-track matching -- a detection is considered "matched" to a tracklet it + # overlaps with a tracklet above this threshold -- it is often a loose threshold like 0.1 + assoc_iou_thresh=0.5, + # IoU threshold for det-to-track matching, which is used to determine whether a masklet is "unmatched" + # by any detections -- it is often a stricter threshold like 0.5 + trk_assoc_iou_thresh=0.5, + # prob threshold for a detection to be added as a new object + new_det_thresh=0.0, + # hotstart parameters: we hold off the outputs for `hotstart_delay` frames and + # 1) remove those tracklets unmatched by any detections based on `hotstart_unmatch_thresh` + # 2) remove those tracklets overlapping with one another based on `hotstart_dup_thresh` + hotstart_delay=0, + hotstart_unmatch_thresh=3, + hotstart_dup_thresh=3, + # Whether to suppress masks only within hotstart. If False, we can suppress masks even if they start before hotstart period. + suppress_unmatched_only_within_hotstart=True, + init_trk_keep_alive=0, + max_trk_keep_alive=8, + min_trk_keep_alive=-4, + # Threshold for suppressing overlapping objects based on recent occlusion + suppress_overlapping_based_on_recent_occlusion_threshold=0.0, + decrease_trk_keep_alive_for_empty_masklets=False, + o2o_matching_masklets_enable=False, # Enable hungarian matching to match existing masklets + suppress_det_close_to_boundary=False, + fill_hole_area=16, + # The maximum number of objects (masklets) to track across all GPUs (for no limit, set it to -1) + max_num_objects=-1, + recondition_every_nth_frame=-1, + # masket confirmation status (to suppress unconfirmed masklets) + masklet_confirmation_enable=False, + # a masklet is confirmed after being consecutively detected and matched for + # `masklet_confirmation_consecutive_det_thresh` + masklet_confirmation_consecutive_det_thresh=3, + # bbox heuristic parameters + reconstruction_bbox_iou_thresh=0.0, + reconstruction_bbox_det_score=0.0, + ): + super().__init__() + self.detector = detector + self.tracker = tracker + self.score_threshold_detection = score_threshold_detection + self.det_nms_thresh = det_nms_thresh + self.assoc_iou_thresh = assoc_iou_thresh + self.trk_assoc_iou_thresh = trk_assoc_iou_thresh + self.new_det_thresh = new_det_thresh + + # hotstart parameters + if hotstart_delay > 0: + assert hotstart_unmatch_thresh <= hotstart_delay + assert hotstart_dup_thresh <= hotstart_delay + self.hotstart_delay = hotstart_delay + self.hotstart_unmatch_thresh = hotstart_unmatch_thresh + self.hotstart_dup_thresh = hotstart_dup_thresh + self.suppress_unmatched_only_within_hotstart = ( + suppress_unmatched_only_within_hotstart + ) + self.init_trk_keep_alive = init_trk_keep_alive + self.max_trk_keep_alive = max_trk_keep_alive + self.min_trk_keep_alive = min_trk_keep_alive + self.suppress_overlapping_based_on_recent_occlusion_threshold = ( + suppress_overlapping_based_on_recent_occlusion_threshold + ) + self.suppress_det_close_to_boundary = suppress_det_close_to_boundary + self.decrease_trk_keep_alive_for_empty_masklets = ( + decrease_trk_keep_alive_for_empty_masklets + ) + self.o2o_matching_masklets_enable = o2o_matching_masklets_enable + self.fill_hole_area = fill_hole_area + self.eval() + self.rank = int(os.getenv("RANK", "0")) + self.world_size = int(os.getenv("WORLD_SIZE", "1")) + self._dist_pg_cpu = None # CPU process group (lazy-initialized on first use) + + # the maximum object number + if max_num_objects > 0: + num_obj_for_compile = math.ceil(max_num_objects / self.world_size) + else: + max_num_objects = 10000 # no limit + num_obj_for_compile = 16 + logger.info(f"setting {max_num_objects=} and {num_obj_for_compile=}") + self.max_num_objects = max_num_objects + self.num_obj_for_compile = num_obj_for_compile + self.recondition_every_nth_frame = recondition_every_nth_frame + self.masklet_confirmation_enable = masklet_confirmation_enable + self.masklet_confirmation_consecutive_det_thresh = ( + masklet_confirmation_consecutive_det_thresh + ) + self.reconstruction_bbox_iou_thresh = reconstruction_bbox_iou_thresh + self.reconstruction_bbox_det_score = reconstruction_bbox_det_score + + @property + def device(self): + self._device = getattr(self, "_device", None) or next(self.parameters()).device + return self._device + + def _init_dist_pg_cpu(self): + # a short 3-min timeout to quickly detect any synchronization failures + timeout_sec = int(os.getenv("SAM3_COLLECTIVE_OP_TIMEOUT_SEC", "180")) + timeout = datetime.timedelta(seconds=timeout_sec) + self._dist_pg_cpu = dist.new_group(backend="gloo", timeout=timeout) + + def broadcast_python_obj_cpu(self, python_obj_list, src): + if self._dist_pg_cpu is None: + self._init_dist_pg_cpu() + dist.broadcast_object_list(python_obj_list, src=src, group=self._dist_pg_cpu) + + def _det_track_one_frame( + self, + frame_idx: int, + num_frames: int, + reverse: bool, + input_batch: BatchedDatapoint, + geometric_prompt: Any, + tracker_states_local: List[Any], + tracker_metadata_prev: Dict[str, Any], + feature_cache: Dict, + orig_vid_height: int, + orig_vid_width: int, + is_image_only: bool = False, + allow_new_detections: bool = True, + ): + """ + This function handles one-step inference for the DenseTracking model in an SPMD manner. + At a high-level, all GPUs execute the same function calls as if it's done on a single GPU, + while under the hood, some function calls involve distributed computation based on sharded + SAM2 states. + + - `input_batch` contains image and other inputs on the entire video; it should be identical across GPUs + - `tracker_states_local` holds the local masklet information in this GPU shard + - `tracker_metadata_prev` manages the metadata for SAM2 objects, such as which masklet is hold on which GPUs + it contains both global and local masklet information + """ + + # Step 1: run backbone and detector in a distributed manner -- this is done via Sam3ImageOnVideoMultiGPU, + # a MultiGPU model (assigned to `self.detector`) that shards frames in a round-robin manner. + # It returns a "det_out" dict for `frame_idx` and fills SAM2 backbone features for `frame_idx` + # into `feature_cache`. Despite its distributed inference under the hood, the results would be + # the same as if it is running backbone and detector for every frame on a single GPU. + det_out = self.run_backbone_and_detection( + frame_idx=frame_idx, + num_frames=num_frames, + reverse=reverse, + input_batch=input_batch, + geometric_prompt=geometric_prompt, + feature_cache=feature_cache, + allow_new_detections=allow_new_detections, + ) + + # Step 2: each GPU propagates its local SAM2 states to get the SAM2 prediction masks. + # the returned `tracker_low_res_masks_global` contains the concatenated masklet predictions + # gathered from all GPUs (as if they are propagated on a single GPU). Note that this step only + # runs the SAM2 propagation step, but doesn't encode new memory for the predicted masks; + # we defer memory encoding to `run_tracker_update_execution_phase` after resolving all heuristics. + if tracker_metadata_prev == {}: + # initialize masklet metadata if it's uninitialized (empty dict) + tracker_metadata_prev.update(self._initialize_metadata()) + tracker_low_res_masks_global, tracker_obj_scores_global = ( + self.run_tracker_propagation( + frame_idx=frame_idx, + num_frames=num_frames, + reverse=reverse, + tracker_states_local=tracker_states_local, + tracker_metadata_prev=tracker_metadata_prev, + ) + ) + + # Step 3: based on detection outputs and the propagated SAM2 prediction masks, we make plans + # for SAM2 masklet updates (i.e. which objects to add and remove, how to load-balance them, etc). + # We also run SAM2 memory encoder globally in this step to resolve non-overlapping constraints. + # **This step should involve all the heuristics needed for any updates.** Most of the update + # planning will be done on the master rank (GPU 0) and the resulting plan `tracker_update_plan` is + # broadcasted to other GPUs (to be executed in a distributed manner). This step also generates the + # new masklet metadata `tracker_metadata_new` (based on its previous version `tracker_metadata_prev`). + tracker_update_plan, tracker_metadata_new = ( + self.run_tracker_update_planning_phase( + frame_idx=frame_idx, + num_frames=num_frames, + reverse=reverse, + det_out=det_out, + tracker_low_res_masks_global=tracker_low_res_masks_global, + tracker_obj_scores_global=tracker_obj_scores_global, + tracker_metadata_prev=tracker_metadata_prev, + tracker_states_local=tracker_states_local, + is_image_only=is_image_only, + ) + ) + + # Get reconditioning info from the update plan + reconditioned_obj_ids = tracker_update_plan.get("reconditioned_obj_ids", set()) + det_to_matched_trk_obj_ids = tracker_update_plan.get( + "det_to_matched_trk_obj_ids", {} + ) + + # Step 4: based on `tracker_update_plan`, each GPU executes the update w.r.t. its local SAM2 inference states + tracker_states_local_new = self.run_tracker_update_execution_phase( + frame_idx=frame_idx, + num_frames=num_frames, + reverse=reverse, + det_out=det_out, + tracker_states_local=tracker_states_local, + tracker_update_plan=tracker_update_plan, + orig_vid_height=orig_vid_height, + orig_vid_width=orig_vid_width, + feature_cache=feature_cache, + ) + + # Step 5: finally, build the outputs for this frame (it only needs to be done on GPU 0 since + # only GPU 0 will send outputs to the server). + if self.rank == 0: + obj_id_to_mask = self.build_outputs( + frame_idx=frame_idx, + num_frames=num_frames, + reverse=reverse, + det_out=det_out, + tracker_low_res_masks_global=tracker_low_res_masks_global, + tracker_obj_scores_global=tracker_obj_scores_global, + tracker_metadata_prev=tracker_metadata_prev, + tracker_update_plan=tracker_update_plan, + orig_vid_height=orig_vid_height, + orig_vid_width=orig_vid_width, + reconditioned_obj_ids=reconditioned_obj_ids, + det_to_matched_trk_obj_ids=det_to_matched_trk_obj_ids, + ) + obj_id_to_score = tracker_metadata_new["obj_id_to_score"] + else: + obj_id_to_mask, obj_id_to_score = {}, {} # dummy outputs on other GPUs + # a few statistics for the current frame as a part of the output + frame_stats = { + "num_obj_tracked": np.sum(tracker_metadata_new["num_obj_per_gpu"]), + "num_obj_dropped": tracker_update_plan["num_obj_dropped_due_to_limit"], + } + # add tracker scores to metadata, it should be fired for frames except the first frame + if tracker_obj_scores_global.shape[0] > 0: + # Convert tracker_obj_scores_global to sigmoid scores before updating + tracker_obj_scores_global = tracker_obj_scores_global.sigmoid().tolist() + tracker_obj_ids = tracker_metadata_prev["obj_ids_all_gpu"] + tracker_metadata_new["obj_id_to_tracker_score_frame_wise"][ + frame_idx + ].update(dict(zip(tracker_obj_ids, tracker_obj_scores_global))) + return ( + obj_id_to_mask, # a dict: obj_id --> output mask + obj_id_to_score, # a dict: obj_id --> output score (prob) + tracker_states_local_new, + tracker_metadata_new, + frame_stats, + tracker_obj_scores_global, # a dict: obj_id --> tracker frame-level scores + ) + + def _suppress_detections_close_to_boundary(self, boxes, margin=0.025): + """ + Suppress detections too close to image edges (for normalized boxes). + + boxes: (N, 4) in xyxy format, normalized [0,1] + margin: fraction of image + """ + x_min, y_min, x_max, y_max = boxes.unbind(-1) + x_c = (x_min + x_max) / 2 + y_c = (y_min + y_max) / 2 + keep = ( + (x_c > margin) + & (x_c < 1.0 - margin) + & (y_c > margin) + & (y_c < 1.0 - margin) + ) + + return keep + + def run_backbone_and_detection( + self, + frame_idx: int, + num_frames: int, + input_batch: BatchedDatapoint, + geometric_prompt: Any, + feature_cache: Dict, + reverse: bool, + allow_new_detections: bool, + ): + # Step 1: if text feature is not cached in `feature_cache`, compute and cache it + text_batch_key = tuple(input_batch.find_text_batch) + if "text" not in feature_cache or text_batch_key not in feature_cache["text"]: + text_outputs = self.detector.backbone.forward_text( + input_batch.find_text_batch, device=self.device + ) + # note: we only cache the text feature of the most recent prompt + feature_cache["text"] = {text_batch_key: text_outputs} + else: + text_outputs = feature_cache["text"][text_batch_key] + + # Step 2: run backbone, detector, and post-processing with NMS + if "multigpu_buffer" not in feature_cache: + # "multigpu_buffer" is a buffer cache used by `self.detector` and it needs + # to be passed to `forward_video_grounding_multigpu` for every call + feature_cache["multigpu_buffer"] = {} + + # Extract max_frame_num_to_track from feature_cache if available + tracking_bounds = feature_cache.get("tracking_bounds", {}) + max_frame_num_to_track = tracking_bounds.get("max_frame_num_to_track") + start_frame_idx = tracking_bounds.get("propagate_in_video_start_frame_idx") + + sam3_image_out, _ = self.detector.forward_video_grounding_multigpu( + backbone_out={ + "img_batch_all_stages": input_batch.img_batch, + **text_outputs, + }, + find_inputs=input_batch.find_inputs, + geometric_prompt=geometric_prompt, + frame_idx=frame_idx, + num_frames=num_frames, + multigpu_buffer=feature_cache["multigpu_buffer"], + track_in_reverse=reverse, + # also get the SAM2 backbone features + return_tracker_backbone_feats=True, + # run NMS as a part of distributed computation + run_nms=self.det_nms_thresh > 0.0, + nms_prob_thresh=self.score_threshold_detection, + nms_iou_thresh=self.det_nms_thresh, + # pass max_frame_num_to_track to respect tracking limits + max_frame_num_to_track=max_frame_num_to_track, + propagate_in_video_start_frame_idx=start_frame_idx, + ) + # note: detections in `sam3_image_out` has already gone through NMS + pred_probs = sam3_image_out["pred_logits"].squeeze(-1).sigmoid() + if not allow_new_detections: + pred_probs = pred_probs - 1e8 # make sure no detections are kept + pred_boxes_xyxy = sam3_image_out["pred_boxes_xyxy"] + pred_masks = sam3_image_out["pred_masks"] + # get the positive detection outputs above threshold + pos_pred_idx = torch.where(pred_probs > self.score_threshold_detection) + det_out = { + "bbox": pred_boxes_xyxy[pos_pred_idx[0], pos_pred_idx[1]], + "mask": pred_masks[pos_pred_idx[0], pos_pred_idx[1]], + "scores": pred_probs[pos_pred_idx[0], pos_pred_idx[1]], + } + + # Step 3: build SAM2 backbone features and store them in `feature_cache` + backbone_cache = {} + sam_mask_decoder = self.tracker.sam_mask_decoder + tracker_backbone_fpn = [ + sam_mask_decoder.conv_s0(sam3_image_out["tracker_backbone_fpn_0"]), + sam_mask_decoder.conv_s1(sam3_image_out["tracker_backbone_fpn_1"]), + sam3_image_out["tracker_backbone_fpn_2"], # fpn_2 doesn't need conv + ] + tracker_backbone_out = { + "vision_features": tracker_backbone_fpn[-1], # top-level feature + "vision_pos_enc": sam3_image_out["tracker_backbone_pos_enc"], + "backbone_fpn": tracker_backbone_fpn, + } + backbone_cache["tracker_backbone_out"] = tracker_backbone_out + feature_cache[frame_idx] = ( + input_batch.img_batch[frame_idx], + backbone_cache, + ) + # remove from `feature_cache` old features to save GPU memory + feature_cache.pop(frame_idx - 1 if not reverse else frame_idx + 1, None) + return det_out + + def run_tracker_propagation( + self, + frame_idx: int, + num_frames: int, + reverse: bool, + tracker_states_local: List[Any], + tracker_metadata_prev: Dict[str, npt.NDArray], + ): + # Step 1: propagate the local SAM2 states to get the current frame's prediction + # `low_res_masks_local` of the existing masklets on this GPU + # - obj_ids_local: List[int] -- list of object IDs + # - low_res_masks_local: Tensor -- (num_local_obj, H_mask, W_mask) + obj_ids_local, low_res_masks_local, obj_scores_local = ( + self._propogate_tracker_one_frame_local_gpu( + tracker_states_local, frame_idx=frame_idx, reverse=reverse + ) + ) + + assert np.all( + obj_ids_local == tracker_metadata_prev["obj_ids_per_gpu"][self.rank] + ), "{} != {}".format( + obj_ids_local, tracker_metadata_prev["obj_ids_per_gpu"][self.rank] + ) + + # Step 2: all-gather `low_res_masks_local` into `low_res_masks_global` + # - low_res_masks_global: Tensor -- (num_global_obj, H_mask, W_mask) + _, H_mask, W_mask = low_res_masks_local.shape + if self.world_size > 1: + # `low_res_masks_local` and `obj_scores_local` need to be contiguous and float32 + # (they could be non-contiguous due to slicing and/or bfloat16 due to autocast) + low_res_masks_local = low_res_masks_local.float().contiguous() + obj_scores_local = obj_scores_local.float().contiguous() + num_obj_this_gpu = tracker_metadata_prev["num_obj_per_gpu"][self.rank] + assert low_res_masks_local.size(0) == num_obj_this_gpu + assert obj_scores_local.size(0) == num_obj_this_gpu + low_res_masks_peers = [ + low_res_masks_local.new_empty(num_obj, H_mask, W_mask) + for num_obj in tracker_metadata_prev["num_obj_per_gpu"] + ] + obj_scores_peers = [ + obj_scores_local.new_empty(num_obj) + for num_obj in tracker_metadata_prev["num_obj_per_gpu"] + ] + dist.all_gather(low_res_masks_peers, low_res_masks_local) + dist.all_gather(obj_scores_peers, obj_scores_local) + low_res_masks_global = torch.cat(low_res_masks_peers, dim=0) + obj_scores_global = torch.cat(obj_scores_peers, dim=0) + else: + low_res_masks_global = low_res_masks_local + obj_scores_global = obj_scores_local + return low_res_masks_global, obj_scores_global + + def _recondition_masklets( + self, + frame_idx, + det_out: Dict[str, Tensor], + trk_id_to_max_iou_high_conf_det: List[int], + tracker_states_local: List[Any], + tracker_metadata: Dict[str, npt.NDArray], + tracker_obj_scores_global: Tensor, + ): + # Recondition the masklets based on the new detections + for trk_obj_id, det_idx in trk_id_to_max_iou_high_conf_det.items(): + new_mask = det_out["mask"][det_idx : det_idx + 1] + input_mask_res = self.tracker.input_mask_size + new_mask_binary = ( + F.interpolate( + new_mask.unsqueeze(1), + size=(input_mask_res, input_mask_res), + mode="bilinear", + align_corners=False, + ).squeeze(1)[0] + > 0 + ) + HIGH_CONF_THRESH = 0.8 + reconditioned_states_idx = set() + obj_idx = np.where(tracker_metadata["obj_ids_all_gpu"] == trk_obj_id)[ + 0 + ].item() + obj_score = tracker_obj_scores_global[obj_idx] + for state_idx, inference_state in enumerate(tracker_states_local): + if ( + trk_obj_id in inference_state["obj_ids"] + # NOTE: Goal of this condition is to avoid reconditioning masks that are occluded/low qualiy. + # Unfortunately, these can get reconditioned anyway due to batching. We should consider removing these heuristics. + and obj_score > HIGH_CONF_THRESH + ): + logger.debug( + f"Adding new mask for track {trk_obj_id} at frame {frame_idx}. Objects {inference_state['obj_ids']} are all reconditioned." + ) + self.tracker.add_new_mask( + inference_state=inference_state, + frame_idx=frame_idx, + obj_id=trk_obj_id, + mask=new_mask_binary, + ) + reconditioned_states_idx.add(state_idx) + + for idx in reconditioned_states_idx: + self.tracker.propagate_in_video_preflight( + tracker_states_local[idx], run_mem_encoder=True + ) + return tracker_states_local + + def run_tracker_update_planning_phase( + self, + frame_idx: int, + num_frames: int, + reverse: bool, + det_out: Dict[str, Tensor], + tracker_low_res_masks_global: Tensor, + tracker_obj_scores_global: Tensor, + tracker_metadata_prev: Dict[str, npt.NDArray], + tracker_states_local: List[Any], + is_image_only: bool = False, + ): + # initialize new metadata from previous metadata (its values will be updated later) + tracker_metadata_new = { + "obj_ids_per_gpu": deepcopy(tracker_metadata_prev["obj_ids_per_gpu"]), + "obj_ids_all_gpu": None, # will be filled later + "num_obj_per_gpu": deepcopy(tracker_metadata_prev["num_obj_per_gpu"]), + "obj_id_to_score": deepcopy(tracker_metadata_prev["obj_id_to_score"]), + "obj_id_to_tracker_score_frame_wise": deepcopy( + tracker_metadata_prev["obj_id_to_tracker_score_frame_wise"] + ), + "obj_id_to_last_occluded": {}, # will be filled later + "max_obj_id": deepcopy(tracker_metadata_prev["max_obj_id"]), + } + + # Initialize reconditioned_obj_ids early to avoid UnboundLocalError + reconditioned_obj_ids = set() + + # Step 1: make the update plan and resolve heuristics on GPU 0 + det_mask_preds: Tensor = det_out["mask"] # low-res mask logits + det_scores_np: npt.NDArray = det_out["scores"].float().cpu().numpy() + det_bbox_xyxy: Tensor = det_out["bbox"] + if self.rank == 0: + # a) match detector and tracker masks and find new objects + ( + new_det_fa_inds, + unmatched_trk_obj_ids, + det_to_matched_trk_obj_ids, + trk_id_to_max_iou_high_conf_det, + empty_trk_obj_ids, + ) = self._associate_det_trk( + det_masks=det_mask_preds, + det_scores_np=det_scores_np, + trk_masks=tracker_low_res_masks_global, + trk_obj_ids=tracker_metadata_prev["obj_ids_all_gpu"], + ) + if self.suppress_det_close_to_boundary: + keep = self._suppress_detections_close_to_boundary( + det_bbox_xyxy[new_det_fa_inds] + ) + new_det_fa_inds = new_det_fa_inds[keep.cpu().numpy()] + + # check whether we've hit the maximum number of objects we can track (and if so, drop some detections) + prev_obj_num = np.sum(tracker_metadata_prev["num_obj_per_gpu"]) + new_det_num = len(new_det_fa_inds) + num_obj_dropped_due_to_limit = 0 + if not is_image_only and prev_obj_num + new_det_num > self.max_num_objects: + logger.warning( + f"hitting {self.max_num_objects=} with {new_det_num=} and {prev_obj_num=}" + ) + new_det_num_to_keep = self.max_num_objects - prev_obj_num + num_obj_dropped_due_to_limit = new_det_num - new_det_num_to_keep + new_det_fa_inds = self._drop_new_det_with_obj_limit( + new_det_fa_inds, det_scores_np, new_det_num_to_keep + ) + assert len(new_det_fa_inds) == new_det_num_to_keep + new_det_num = len(new_det_fa_inds) + + # assign object IDs to new detections and decide which GPU to place them + new_det_start_obj_id = tracker_metadata_prev["max_obj_id"] + 1 + new_det_obj_ids = new_det_start_obj_id + np.arange(new_det_num) + prev_workload_per_gpu = tracker_metadata_prev["num_obj_per_gpu"] + new_det_gpu_ids = self._assign_new_det_to_gpus( + new_det_num=new_det_num, + prev_workload_per_gpu=prev_workload_per_gpu, + ) + + # b) handle hotstart heuristics to remove objects + # here `rank0_metadata` contains metadata stored on (and only accessible to) GPU 0; + # we avoid broadcasting them to other GPUs to save communication cost, assuming + # that `rank0_metadata` is not needed by other GPUs + rank0_metadata_new = deepcopy(tracker_metadata_prev["rank0_metadata"]) + if not hasattr(self, "_warm_up_complete") or self._warm_up_complete: + obj_ids_newly_removed, rank0_metadata_new = self._process_hotstart( + frame_idx=frame_idx, + num_frames=num_frames, + reverse=reverse, + det_to_matched_trk_obj_ids=det_to_matched_trk_obj_ids, + new_det_obj_ids=new_det_obj_ids, + empty_trk_obj_ids=empty_trk_obj_ids, + unmatched_trk_obj_ids=unmatched_trk_obj_ids, + rank0_metadata=rank0_metadata_new, + tracker_metadata=tracker_metadata_prev, + ) + else: + # if warm-up is not complete, we don't remove any objects + obj_ids_newly_removed = set() + tracker_metadata_new["rank0_metadata"] = rank0_metadata_new + + # Step 2: broadcast the update plan to other GPUs + NUM_BROADCAST_ITEMS = 9 + if self.rank == 0 and self.world_size > 1: + # `num_obj_per_gpu_on_rank0` is used for metadata consistency check on other GPUs + # (it's a small array with length==self.world_size, so broadcasting it is cheap) + num_obj_per_gpu_on_rank0 = tracker_metadata_prev["num_obj_per_gpu"] + update_plan = [ + new_det_fa_inds, + new_det_obj_ids, + new_det_gpu_ids, + num_obj_per_gpu_on_rank0, + unmatched_trk_obj_ids, + det_to_matched_trk_obj_ids, + obj_ids_newly_removed, + num_obj_dropped_due_to_limit, + trk_id_to_max_iou_high_conf_det, + ] + assert ( + len(update_plan) == NUM_BROADCAST_ITEMS + ), f"Manually update NUM_BROADCAST_ITEMS to be: {len(update_plan)}" + self.broadcast_python_obj_cpu(update_plan, src=0) + elif self.rank > 0 and self.world_size > 1: + update_plan = [ + None + ] * NUM_BROADCAST_ITEMS # other ranks receive the plan from rank 0 + self.broadcast_python_obj_cpu(update_plan, src=0) + ( + new_det_fa_inds, + new_det_obj_ids, + new_det_gpu_ids, + num_obj_per_gpu_on_rank0, + unmatched_trk_obj_ids, + det_to_matched_trk_obj_ids, + obj_ids_newly_removed, + num_obj_dropped_due_to_limit, + trk_id_to_max_iou_high_conf_det, + ) = update_plan + # metadata consistency check: verify that the received `num_obj_per_gpu_on_rank0` is consistent with the local metadata + # it's critical that all GPUs agree on the previous number of objects (otherwise the inference might hang or fail silently) + if not np.all( + num_obj_per_gpu_on_rank0 == tracker_metadata_prev["num_obj_per_gpu"] + ): + raise RuntimeError( + f"{self.rank=} received {num_obj_per_gpu_on_rank0=}, which is inconsistent with local record " + f"{tracker_metadata_prev['num_obj_per_gpu']=}. There's likely a bug in update planning or execution." + ) + + # `tracker_update_plan` should be identical on all GPUs after broadcasting + tracker_update_plan = { + "new_det_fa_inds": new_det_fa_inds, # npt.NDArray + "new_det_obj_ids": new_det_obj_ids, # npt.NDArray + "new_det_gpu_ids": new_det_gpu_ids, # npt.NDArray + "unmatched_trk_obj_ids": unmatched_trk_obj_ids, # npt.NDArray + "det_to_matched_trk_obj_ids": det_to_matched_trk_obj_ids, # dict + "obj_ids_newly_removed": obj_ids_newly_removed, # set + "num_obj_dropped_due_to_limit": num_obj_dropped_due_to_limit, # int + "trk_id_to_max_iou_high_conf_det": trk_id_to_max_iou_high_conf_det, # dict + "reconditioned_obj_ids": reconditioned_obj_ids, # set + } + + # Step 3 (optional): recondition masklets based on high-confidence detections before memory encoding + # NOTE: Running this in execution phase (after memory encoding) can lead to suboptimal results + should_recondition_iou = False + + # Evaluate tracklets for reconditioning based on bbox IoU mismatch with detections + if ( + self.reconstruction_bbox_iou_thresh > 0 + and len(trk_id_to_max_iou_high_conf_det) > 0 + ): + for trk_obj_id, det_idx in trk_id_to_max_iou_high_conf_det.items(): + det_box = det_out["bbox"][det_idx] + det_score = det_out["scores"][det_idx] + + try: + trk_idx = list(tracker_metadata_prev["obj_ids_all_gpu"]).index( + trk_obj_id + ) + except ValueError: + continue # Skip if tracklet not found + + tracker_mask = tracker_low_res_masks_global[trk_idx] + mask_binary = tracker_mask > 0 + mask_area = mask_binary.sum().item() + + if mask_area == 0: + continue # Skip tracklets with zero mask area + + # Get bounding box from SAM2 mask and convert to normalized coordinates + tracker_box_pixels = ( + mask_to_box(mask_binary.unsqueeze(0).unsqueeze(0)) + .squeeze(0) + .squeeze(0) + ) + mask_height, mask_width = tracker_mask.shape[-2:] + tracker_box_normalized = torch.tensor( + [ + tracker_box_pixels[0] / mask_width, + tracker_box_pixels[1] / mask_height, + tracker_box_pixels[2] / mask_width, + tracker_box_pixels[3] / mask_height, + ], + device=tracker_box_pixels.device, + ) + + # Compute IoU between detection and SAM2 tracklet bounding boxes + det_box_batch = det_box.unsqueeze(0) + tracker_box_batch = tracker_box_normalized.unsqueeze(0) + iou = fast_diag_box_iou(det_box_batch, tracker_box_batch)[0] + + if ( + iou < self.reconstruction_bbox_iou_thresh + and det_score >= self.reconstruction_bbox_det_score + ): + should_recondition_iou = True + reconditioned_obj_ids.add(trk_obj_id) + + should_recondition_periodic = ( + self.recondition_every_nth_frame > 0 + and frame_idx % self.recondition_every_nth_frame == 0 + and len(trk_id_to_max_iou_high_conf_det) > 0 + ) + + # Recondition if periodic or IoU condition met + if should_recondition_periodic or should_recondition_iou: + self._recondition_masklets( + frame_idx, + det_out, + trk_id_to_max_iou_high_conf_det, + tracker_states_local, + tracker_metadata_prev, + tracker_obj_scores_global, + ) + + # Step 4: Run SAM2 memory encoder on the current frame's prediction masks + # This is done on all GPUs + batch_size = tracker_low_res_masks_global.size(0) + if batch_size > 0: + if not hasattr(self, "_warm_up_complete") or self._warm_up_complete: + if self.suppress_overlapping_based_on_recent_occlusion_threshold > 0.0: + # NOTE: tracker_low_res_masks_global is updated in-place then returned + tracker_low_res_masks_global = ( + self._suppress_overlapping_based_on_recent_occlusion( + frame_idx, + tracker_low_res_masks_global, + tracker_metadata_prev, + tracker_metadata_new, + obj_ids_newly_removed, + reverse, + ) + ) + + self._tracker_update_memories( + tracker_states_local, + frame_idx, + tracker_metadata=tracker_metadata_prev, + low_res_masks=tracker_low_res_masks_global, + ) + + # Step 4: update the SAM2 metadata based on the update plan + # note: except for "rank0_metadata" (that is only available on GPU 0), + # the updated `tracker_metadata_new` should be identical on all GPUs + for rank in range(self.world_size): + new_det_obj_ids_this_gpu = new_det_obj_ids[new_det_gpu_ids == rank] + updated_obj_ids_this_gpu = tracker_metadata_new["obj_ids_per_gpu"][rank] + if len(new_det_obj_ids_this_gpu) > 0: + updated_obj_ids_this_gpu = np.concatenate( + [updated_obj_ids_this_gpu, new_det_obj_ids_this_gpu] + ) + if len(obj_ids_newly_removed) > 0: + is_removed = np.isin( + updated_obj_ids_this_gpu, list(obj_ids_newly_removed) + ) + updated_obj_ids_this_gpu = updated_obj_ids_this_gpu[~is_removed] + tracker_metadata_new["obj_ids_per_gpu"][rank] = updated_obj_ids_this_gpu + tracker_metadata_new["num_obj_per_gpu"][rank] = len( + updated_obj_ids_this_gpu + ) + tracker_metadata_new["obj_ids_all_gpu"] = np.concatenate( + tracker_metadata_new["obj_ids_per_gpu"] + ) + # update object scores and the maximum object ID assigned so far + if len(new_det_obj_ids) > 0: + tracker_metadata_new["obj_id_to_score"].update( + zip(new_det_obj_ids, det_scores_np[new_det_fa_inds]) + ) + # tracker scores are not available for new objects, use det score instead. + tracker_metadata_new["obj_id_to_tracker_score_frame_wise"][ + frame_idx + ].update(zip(new_det_obj_ids, det_scores_np[new_det_fa_inds])) + tracker_metadata_new["max_obj_id"] = max( + tracker_metadata_new["max_obj_id"], + np.max(new_det_obj_ids), + ) + # for removed objects, we set their scores to a very low value (-1e4) but still + # keep them in "obj_id_to_score" (it's easier to handle outputs this way) + for obj_id in obj_ids_newly_removed: + tracker_metadata_new["obj_id_to_score"][obj_id] = -1e4 + tracker_metadata_new["obj_id_to_tracker_score_frame_wise"][frame_idx][ + obj_id + ] = -1e4 + tracker_metadata_new["obj_id_to_last_occluded"].pop(obj_id, None) + # check that "rank0_metadata" is in tracker_metadata_new if and only if it's GPU 0 + assert ("rank0_metadata" in tracker_metadata_new) == (self.rank == 0) + if self.rank == 0 and self.masklet_confirmation_enable: + rank0_metadata = self.update_masklet_confirmation_status( + rank0_metadata=tracker_metadata_new["rank0_metadata"], + obj_ids_all_gpu_prev=tracker_metadata_prev["obj_ids_all_gpu"], + obj_ids_all_gpu_updated=tracker_metadata_new["obj_ids_all_gpu"], + det_to_matched_trk_obj_ids=det_to_matched_trk_obj_ids, + new_det_obj_ids=new_det_obj_ids, + ) + tracker_metadata_new["rank0_metadata"] = rank0_metadata + + return tracker_update_plan, tracker_metadata_new + + def _suppress_overlapping_based_on_recent_occlusion( + self, + frame_idx: int, + tracker_low_res_masks_global: Tensor, + tracker_metadata_prev: Dict[str, Any], + tracker_metadata_new: Dict[str, Any], + obj_ids_newly_removed: Set[int], + reverse: bool = False, + ): + """ + Suppress overlapping masks based on the most recent occlusion information. If an object is removed by hotstart, we always suppress it if it overlaps with any other object. + Args: + frame_idx (int): The current frame index. + tracker_low_res_masks_global (Tensor): The low-resolution masks for the current frame. + tracker_metadata_prev (Dict[str, Any]): The metadata from the previous frame. + tracker_metadata_new (Dict[str, Any]): The metadata for the current frame. + obj_ids_newly_removed (Set[int]): The object IDs that have been removed. + Return: + Tensor: The updated low-resolution masks with some objects suppressed. + """ + obj_ids_global = tracker_metadata_prev["obj_ids_all_gpu"] + binary_tracker_low_res_masks_global = tracker_low_res_masks_global > 0 + batch_size = tracker_low_res_masks_global.size(0) + if batch_size > 0: + assert ( + len(obj_ids_global) == batch_size + ), f"Mismatch in number of objects: {len(obj_ids_global)} vs {batch_size}" + NEVER_OCCLUDED = -1 + ALWAYS_OCCLUDED = 100000 # This value should be larger than any possible frame index, indicates that the object was removed by hotstart logic + last_occluded_prev = torch.cat( + [ + tracker_metadata_prev["obj_id_to_last_occluded"].get( + obj_id, + torch.full( + (1,), + fill_value=( + NEVER_OCCLUDED + if obj_id not in obj_ids_newly_removed + else ALWAYS_OCCLUDED + ), + device=binary_tracker_low_res_masks_global.device, + dtype=torch.long, + ), + ) + for obj_id in obj_ids_global + ], + dim=0, + ) + to_suppress = self._get_objects_to_suppress_based_on_most_recently_occluded( + binary_tracker_low_res_masks_global, + last_occluded_prev, + obj_ids_global, + frame_idx, + reverse, + ) + + # Update metadata with occlusion information + is_obj_occluded = ~(binary_tracker_low_res_masks_global.any(dim=(-1, -2))) + is_obj_occluded_or_suppressed = is_obj_occluded | to_suppress + last_occluded_new = last_occluded_prev.clone() + last_occluded_new[is_obj_occluded_or_suppressed] = frame_idx + # Slice out the last occluded frame for each object + tracker_metadata_new["obj_id_to_last_occluded"] = { + obj_id: last_occluded_new[obj_idx : obj_idx + 1] + for obj_idx, obj_id in enumerate(obj_ids_global) + } + + # Zero out suppressed masks before memory encoding + NO_OBJ_LOGIT = -10 + tracker_low_res_masks_global[to_suppress] = NO_OBJ_LOGIT + + return tracker_low_res_masks_global + + def run_tracker_update_execution_phase( + self, + frame_idx: int, + num_frames: int, + reverse: bool, + det_out: Dict[str, Tensor], + tracker_states_local: List[Any], + tracker_update_plan: Dict[str, npt.NDArray], + orig_vid_height: int, + orig_vid_width: int, + feature_cache: Dict, + ): + # initialize tracking scores with detection scores + new_det_fa_inds: npt.NDArray = tracker_update_plan["new_det_fa_inds"] + new_det_obj_ids: npt.NDArray = tracker_update_plan["new_det_obj_ids"] + new_det_gpu_ids: npt.NDArray = tracker_update_plan["new_det_gpu_ids"] + is_on_this_gpu: npt.NDArray = new_det_gpu_ids == self.rank + new_det_obj_ids_local: npt.NDArray = new_det_obj_ids[is_on_this_gpu] + new_det_fa_inds_local: npt.NDArray = new_det_fa_inds[is_on_this_gpu] + obj_ids_newly_removed: Set[int] = tracker_update_plan["obj_ids_newly_removed"] + + # Step 1: add new objects from the detector to SAM2 inference states + if len(new_det_fa_inds_local) > 0: + new_det_fa_inds_local_t = torch.from_numpy(new_det_fa_inds_local) + new_det_masks: Tensor = det_out["mask"][new_det_fa_inds_local_t] + # initialize SAM2 with new object masks + tracker_states_local = self._tracker_add_new_objects( + frame_idx=frame_idx, + num_frames=num_frames, + new_obj_ids=new_det_obj_ids_local, + new_obj_masks=new_det_masks, + tracker_states_local=tracker_states_local, + orig_vid_height=orig_vid_height, + orig_vid_width=orig_vid_width, + feature_cache=feature_cache, + ) + + # Step 2: remove from SAM2 inference states those objects removed by heuristics + if len(obj_ids_newly_removed) > 0: + self._tracker_remove_objects(tracker_states_local, obj_ids_newly_removed) + + return tracker_states_local + + def build_outputs( + self, + frame_idx: int, + num_frames: int, + reverse: bool, + det_out: Dict[str, Tensor], + tracker_low_res_masks_global: Tensor, + tracker_obj_scores_global: Tensor, + tracker_metadata_prev: Dict[str, npt.NDArray], + tracker_update_plan: Dict[str, npt.NDArray], + orig_vid_height: int, + orig_vid_width: int, + reconditioned_obj_ids: set = None, + det_to_matched_trk_obj_ids: dict = None, + ): + new_det_fa_inds: npt.NDArray = tracker_update_plan["new_det_fa_inds"] + new_det_obj_ids: npt.NDArray = tracker_update_plan["new_det_obj_ids"] + obj_id_to_mask = {} # obj_id --> output mask tensor + + # Part 1: masks from previous SAM2 propagation + existing_masklet_obj_ids = tracker_metadata_prev["obj_ids_all_gpu"] + existing_masklet_video_res_masks = F.interpolate( + tracker_low_res_masks_global.unsqueeze(1), + size=(orig_vid_height, orig_vid_width), + mode="bilinear", + align_corners=False, + ) # (num_obj, 1, H_video, W_video) + existing_masklet_binary = existing_masklet_video_res_masks > 0 + assert len(existing_masklet_obj_ids) == len(existing_masklet_binary) + for obj_id, mask in zip(existing_masklet_obj_ids, existing_masklet_binary): + obj_id_to_mask[obj_id] = mask # (1, H_video, W_video) + + # Part 2: masks from new detections + new_det_fa_inds_t = torch.from_numpy(new_det_fa_inds) + new_det_low_res_masks = det_out["mask"][new_det_fa_inds_t].unsqueeze(1) + new_det_low_res_masks = fill_holes_in_mask_scores( + new_det_low_res_masks, + max_area=self.fill_hole_area, + fill_holes=True, + remove_sprinkles=True, + ) + new_masklet_video_res_masks = F.interpolate( + new_det_low_res_masks, + size=(orig_vid_height, orig_vid_width), + mode="bilinear", + align_corners=False, + ) # (num_obj, 1, H_video, W_video) + + new_masklet_binary = new_masklet_video_res_masks > 0 + assert len(new_det_obj_ids) == len(new_masklet_video_res_masks) + for obj_id, mask in zip(new_det_obj_ids, new_masklet_binary): + obj_id_to_mask[obj_id] = mask # (1, H_video, W_video) + + # Part 3: Override masks for reconditioned objects using detection masks + if reconditioned_obj_ids is not None and len(reconditioned_obj_ids) > 0: + trk_id_to_max_iou_high_conf_det = tracker_update_plan.get( + "trk_id_to_max_iou_high_conf_det", {} + ) + + for obj_id in reconditioned_obj_ids: + det_idx = trk_id_to_max_iou_high_conf_det.get(obj_id) + + if det_idx is not None: + det_mask = det_out["mask"][det_idx] + det_mask = det_mask.unsqueeze(0).unsqueeze(0) + det_mask_resized = ( + F.interpolate( + det_mask.float(), + size=(orig_vid_height, orig_vid_width), + mode="bilinear", + align_corners=False, + ) + > 0 + ) + + det_mask_final = det_mask_resized.squeeze(0) + obj_id_to_mask[obj_id] = det_mask_final + + return obj_id_to_mask + + def _get_objects_to_suppress_based_on_most_recently_occluded( + self, + binary_low_res_masks: Tensor, + last_occluded: List[int], + obj_ids: List[int], + frame_idx: int = None, + reverse: bool = False, + ): + # Suppress overlapping masks for objects that were most recently occluded + assert ( + binary_low_res_masks.dtype == torch.bool + ), f"Expected boolean tensor, got {binary_low_res_masks.dtype}" + to_suppress = torch.zeros( + binary_low_res_masks.size(0), + device=binary_low_res_masks.device, + dtype=torch.bool, + ) + if len(obj_ids) <= 1: + return to_suppress + + iou = mask_iou(binary_low_res_masks, binary_low_res_masks) # [N,N] + + # Create masks for upper triangular matrix (i < j) and IoU threshold + mask_iou_thresh = ( + iou >= self.suppress_overlapping_based_on_recent_occlusion_threshold + ) + overlapping_pairs = torch.triu(mask_iou_thresh, diagonal=1) # [N,N] + + last_occ_expanded_i = last_occluded.unsqueeze(1) # (N, 1) + last_occ_expanded_j = last_occluded.unsqueeze(0) # (1, N) + # Suppress most recently occluded + cmp_op = torch.gt if not reverse else torch.lt + suppress_i_mask = ( + overlapping_pairs + & cmp_op( + last_occ_expanded_i, last_occ_expanded_j + ) # (last_occ_expanded_i > last_occ_expanded_j) + & ( + last_occ_expanded_j > -1 + ) # j can suppress i only if i was previously occluded + ) + suppress_j_mask = ( + overlapping_pairs + & cmp_op(last_occ_expanded_j, last_occ_expanded_i) + & ( + last_occ_expanded_i > -1 + ) # i can suppress j only if j was previously occluded + ) + # Apply suppression + to_suppress = suppress_i_mask.any(dim=1) | suppress_j_mask.any(dim=0) + + # Log for debugging + if ( + self.rank == 0 + and logger.isEnabledFor(logging.DEBUG) + and frame_idx is not None + ): + suppress_i_mask = suppress_i_mask.cpu().numpy() + suppress_j_mask = suppress_j_mask.cpu().numpy() + last_occluded = last_occluded.cpu().numpy() + + # Find all suppression pairs without using torch.where + batch_size = suppress_i_mask.shape[0] + + # Log i-suppression cases (where i gets suppressed in favor of j) + for i in range(batch_size): + for j in range(batch_size): + if suppress_i_mask[i, j]: + logger.debug( + f"{frame_idx=}: Suppressing obj {obj_ids[i]} last occluded {last_occluded[i]} in favor of {obj_ids[j]} last occluded {last_occluded[j]}" + ) + + # Log j-suppression cases (where j gets suppressed in favor of i) + for i in range(batch_size): + for j in range(batch_size): + if suppress_j_mask[i, j]: + logger.debug( + f"{frame_idx=}: Suppressing obj {obj_ids[j]} last occluded {last_occluded[j]} in favor of {obj_ids[i]} last occluded {last_occluded[i]}" + ) + + return to_suppress + + def _propogate_tracker_one_frame_local_gpu( + self, + inference_states: List[Any], + frame_idx: int, + reverse: bool, + # by default, we disable memory encoding until we gather all outputs + run_mem_encoder: bool = False, + ): + """ + inference_states: List of inference states, each state corresponds to a different set of objects. + """ + obj_ids_local = [] + low_res_masks_list = [] + obj_scores_list = [] + for inference_state in inference_states: + if len(inference_state["obj_ids"]) == 0: + continue # skip propagation on empty inference states + + # propagate one frame + num_frames_propagated = 0 + for out in self.tracker.propagate_in_video( + inference_state, + start_frame_idx=frame_idx, + # end_frame_idx = start_frame_idx + max_frame_num_to_track + # (i.e. propagating 1 frame since end_frame_idx is inclusive) + max_frame_num_to_track=0, + reverse=reverse, + tqdm_disable=True, + run_mem_encoder=run_mem_encoder, + ): + out_frame_idx, out_obj_ids, out_low_res_masks, _, out_obj_scores = out + num_frames_propagated += 1 + + # only 1 frames should be propagated + assert ( + num_frames_propagated == 1 and out_frame_idx == frame_idx + ), f"num_frames_propagated: {num_frames_propagated}, out_frame_idx: {out_frame_idx}, frame_idx: {frame_idx}" + assert isinstance(out_obj_ids, list) + obj_ids_local.extend(out_obj_ids) + low_res_masks_list.append(out_low_res_masks.squeeze(1)) + obj_scores_list.append(out_obj_scores.squeeze(1)) + + # concatenate the output masklets from all local inference states + H_mask = W_mask = self.tracker.low_res_mask_size + if len(low_res_masks_list) > 0: + low_res_masks_local = torch.cat(low_res_masks_list, dim=0) + obj_scores_local = torch.cat(obj_scores_list, dim=0) + assert low_res_masks_local.shape[1:] == (H_mask, W_mask) + + # Apply hole filling to the masks + low_res_masks_local = fill_holes_in_mask_scores( + low_res_masks_local.unsqueeze(1), + max_area=self.fill_hole_area, + fill_holes=True, + remove_sprinkles=True, + ) + low_res_masks_local = low_res_masks_local.squeeze(1) + else: + low_res_masks_local = torch.zeros(0, H_mask, W_mask, device=self.device) + obj_scores_local = torch.zeros(0, device=self.device) + + return obj_ids_local, low_res_masks_local, obj_scores_local + + def _associate_det_trk( + self, + det_masks: Tensor, + det_scores_np: npt.NDArray, + trk_masks: Tensor, + trk_obj_ids: npt.NDArray, + ): + """ + Match detections on the current frame with the existing masklets. + + Args: + - det_masks: (N, H, W) tensor of predicted masks + - det_scores_np: (N,) array of detection scores + - trk_masks: (M, H, W) tensor of track masks + - trk_obj_ids: (M,) array of object IDs corresponding to trk_masks + + Returns: + - new_det_fa_inds: array of new object indices. + - unmatched_trk_obj_ids: array of existing masklet object IDs that are not matched + to any detections on this frame (for unmatched, we only count masklets with >0 area) + - det_to_matched_trk_obj_ids: dict[int, npt.NDArray]: mapping from detector's detection indices + to the list of matched tracklet object IDs + - empty_trk_obj_ids: array of existing masklet object IDs with zero area in SAM2 prediction + """ + iou_threshold = self.assoc_iou_thresh + iou_threshold_trk = self.trk_assoc_iou_thresh + new_det_thresh = self.new_det_thresh + + assert det_masks.is_floating_point(), "float tensor expected (do not binarize)" + assert trk_masks.is_floating_point(), "float tensor expected (do not binarize)" + assert ( + trk_masks.size(0) == len(trk_obj_ids) + ), f"trk_masks and trk_obj_ids should have the same length, {trk_masks.size(0)} vs {len(trk_obj_ids)}" + if trk_masks.size(0) == 0: + # all detections are new + new_det_fa_inds = np.arange(det_masks.size(0)) + unmatched_trk_obj_ids = np.array([], np.int64) + empty_trk_obj_ids = np.array([], np.int64) + det_to_matched_trk_obj_ids = {} + trk_id_to_max_iou_high_conf_det = {} + return ( + new_det_fa_inds, + unmatched_trk_obj_ids, + det_to_matched_trk_obj_ids, + trk_id_to_max_iou_high_conf_det, + empty_trk_obj_ids, + ) + elif det_masks.size(0) == 0: + # all previous tracklets are unmatched if they have a non-zero area + new_det_fa_inds = np.array([], np.int64) + trk_is_nonempty = (trk_masks > 0).any(dim=(1, 2)).cpu().numpy() + unmatched_trk_obj_ids = trk_obj_ids[trk_is_nonempty] + empty_trk_obj_ids = trk_obj_ids[~trk_is_nonempty] + det_to_matched_trk_obj_ids = {} + trk_id_to_max_iou_high_conf_det = {} + return ( + new_det_fa_inds, + unmatched_trk_obj_ids, + det_to_matched_trk_obj_ids, + trk_id_to_max_iou_high_conf_det, + empty_trk_obj_ids, + ) + + if det_masks.shape[-2:] != trk_masks.shape[-2:]: + # resize to the smaller size to save GPU memory + if np.prod(det_masks.shape[-2:]) < np.prod(trk_masks.shape[-2:]): + trk_masks = F.interpolate( + trk_masks.unsqueeze(1), + size=det_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + else: + # resize detections to track size + det_masks = F.interpolate( + det_masks.unsqueeze(1), + size=trk_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + + det_masks_binary = det_masks > 0 + trk_masks_binary = trk_masks > 0 + ious = mask_iou(det_masks_binary, trk_masks_binary) # (N, M) + + ious_np = ious.cpu().numpy() + if self.o2o_matching_masklets_enable: + from scipy.optimize import linear_sum_assignment + + # Hungarian matching for tracks (one-to-one: each track matches at most one detection) + cost_matrix = 1 - ious_np # Hungarian solves for minimum cost + row_ind, col_ind = linear_sum_assignment(cost_matrix) + trk_is_matched = np.zeros(trk_masks.size(0), dtype=bool) + for d, t in zip(row_ind, col_ind): + if ious_np[d, t] >= iou_threshold_trk: + trk_is_matched[t] = True + else: + trk_is_matched = (ious_np >= iou_threshold_trk).any(axis=0) + # Non-empty tracks not matched by Hungarian assignment above threshold are unmatched + trk_is_nonempty = trk_masks_binary.any(dim=(1, 2)).cpu().numpy() + trk_is_unmatched = np.logical_and(trk_is_nonempty, ~trk_is_matched) + unmatched_trk_obj_ids = trk_obj_ids[trk_is_unmatched] + # also record masklets that have zero area in SAM 2 prediction + empty_trk_obj_ids = trk_obj_ids[~trk_is_nonempty] + + # For detections: allow many tracks to match to the same detection (many-to-one) + # So, a detection is 'new' if it does not match any track above threshold + is_new_det = np.logical_and( + det_scores_np >= new_det_thresh, + np.logical_not(np.any(ious_np >= iou_threshold, axis=1)), + ) + new_det_fa_inds = np.nonzero(is_new_det)[0] + + # for each detection, which tracks it matched to (above threshold) + det_to_matched_trk_obj_ids = {} + trk_id_to_max_iou_high_conf_det = {} # trk id --> exactly one detection idx + HIGH_CONF_THRESH = 0.8 + HIGH_IOU_THRESH = 0.8 + det_to_max_iou_trk_idx = np.argmax(ious_np, axis=1) + det_is_high_conf = (det_scores_np >= HIGH_CONF_THRESH) & ~is_new_det + det_is_high_iou = np.max(ious_np, axis=1) >= HIGH_IOU_THRESH + det_is_high_conf_and_iou = set( + np.nonzero(det_is_high_conf & det_is_high_iou)[0] + ) + for d in range(det_masks.size(0)): + det_to_matched_trk_obj_ids[d] = trk_obj_ids[ious_np[d, :] >= iou_threshold] + if d in det_is_high_conf_and_iou: + trk_obj_id = trk_obj_ids[det_to_max_iou_trk_idx[d]].item() + trk_id_to_max_iou_high_conf_det[trk_obj_id] = d + + return ( + new_det_fa_inds, + unmatched_trk_obj_ids, + det_to_matched_trk_obj_ids, + trk_id_to_max_iou_high_conf_det, + empty_trk_obj_ids, + ) + + def _assign_new_det_to_gpus(self, new_det_num, prev_workload_per_gpu): + """Distribute the new objects to the GPUs with the least workload.""" + workload_per_gpu: npt.NDArray = prev_workload_per_gpu.copy() + new_det_gpu_ids = np.zeros(new_det_num, np.int64) + + # assign the objects one by one + for i in range(len(new_det_gpu_ids)): + # find the GPU with the least workload + min_gpu = np.argmin(workload_per_gpu) + new_det_gpu_ids[i] = min_gpu + workload_per_gpu[min_gpu] += 1 + return new_det_gpu_ids + + def _process_hotstart( + self, + frame_idx: int, + num_frames: int, + reverse: bool, + det_to_matched_trk_obj_ids: Dict[int, npt.NDArray], + new_det_obj_ids: npt.NDArray, + empty_trk_obj_ids: npt.NDArray, + unmatched_trk_obj_ids: npt.NDArray, + rank0_metadata: Dict[str, Any], + tracker_metadata: Dict[str, Any], + ): + """Handle hotstart heuristics to remove unmatched or duplicated objects.""" + # obj_id --> first frame index where the object was detected + obj_first_frame_idx = rank0_metadata["obj_first_frame_idx"] + # obj_id --> [mismatched frame indices] + unmatched_frame_inds = rank0_metadata["unmatched_frame_inds"] + trk_keep_alive = rank0_metadata["trk_keep_alive"] + # (first_appear_obj_id, obj_id) --> [overlap frame indices] + overlap_pair_to_frame_inds = rank0_metadata["overlap_pair_to_frame_inds"] + # removed_obj_ids: object IDs that are suppressed via hot-start + removed_obj_ids = rank0_metadata["removed_obj_ids"] + suppressed_obj_ids = rank0_metadata["suppressed_obj_ids"][frame_idx] + + obj_ids_newly_removed = set() # object IDs to be newly removed on this frame + hotstart_diff = ( + frame_idx - self.hotstart_delay + if not reverse + else frame_idx + self.hotstart_delay + ) + + # Step 1: log the frame index where each object ID first appears + for obj_id in new_det_obj_ids: + if obj_id not in obj_first_frame_idx: + obj_first_frame_idx[obj_id] = frame_idx + assert obj_id not in trk_keep_alive + trk_keep_alive[obj_id] = self.init_trk_keep_alive + + matched_trks = set() + # We use the det-->tracks list to check for matched objects. Otherwise, we need to compute areas to decide whether they're occluded + for matched_trks_per_det in det_to_matched_trk_obj_ids.values(): + matched_trks.update(matched_trks_per_det) + for obj_id in matched_trks: + # NOTE: To minimize number of configurable params, we use the hotstart_unmatch_thresh to set the max value of trk_keep_alive + trk_keep_alive[obj_id] = min( + self.max_trk_keep_alive, trk_keep_alive[obj_id] + 1 + ) + for obj_id in unmatched_trk_obj_ids: + unmatched_frame_inds[obj_id].append(frame_idx) + # NOTE: To minimize number of configurable params, we use the hotstart_unmatch_thresh to set the min value of trk_keep_alive + # The max keep alive is 2x the min, means the model prefers to keep the prediction rather than suppress it if it was matched long enough. + trk_keep_alive[obj_id] = max( + self.min_trk_keep_alive, trk_keep_alive[obj_id] - 1 + ) + if self.decrease_trk_keep_alive_for_empty_masklets: + for obj_id in empty_trk_obj_ids: + # NOTE: To minimize number of configurable params, we use the hotstart_unmatch_thresh to set the min value of trk_keep_alive + trk_keep_alive[obj_id] = max( + self.min_trk_keep_alive, trk_keep_alive[obj_id] - 1 + ) + + # Step 2: removed tracks that has not matched with detections for `hotstart_unmatch_thresh` frames with hotstart period + # a) add unmatched frame indices for each existing object ID + # note that `unmatched_trk_obj_ids` contains those frames where the SAM2 output mask + # doesn't match any detection; it excludes those frames where SAM2 gives an empty mask + # b) remove a masklet if it first appears after `hotstart_diff` and is unmatched for more + # than `self.hotstart_unmatch_thresh` frames + for obj_id, frame_indices in unmatched_frame_inds.items(): + if obj_id in removed_obj_ids or obj_id in obj_ids_newly_removed: + continue # skip if the object is already removed + if len(frame_indices) >= self.hotstart_unmatch_thresh: + is_within_hotstart = ( + obj_first_frame_idx[obj_id] > hotstart_diff and not reverse + ) or (obj_first_frame_idx[obj_id] < hotstart_diff and reverse) + if is_within_hotstart: + obj_ids_newly_removed.add(obj_id) + logger.debug( + f"Removing object {obj_id} at frame {frame_idx} " + f"since it is unmatched for frames: {frame_indices}" + ) + if ( + trk_keep_alive[obj_id] <= 0 # Object has not been matched for too long + and not self.suppress_unmatched_only_within_hotstart + and obj_id not in removed_obj_ids + and obj_id not in obj_ids_newly_removed + ): + logger.debug( + f"Suppressing object {obj_id} at frame {frame_idx}, due to being unmatched" + ) + suppressed_obj_ids.add(obj_id) + + # Step 3: removed tracks that overlaps with another track for `hotstart_dup_thresh` frames + # a) find overlaps tracks -- we consider overlap if they match to the same detection + for _, matched_trk_obj_ids in det_to_matched_trk_obj_ids.items(): + if len(matched_trk_obj_ids) < 2: + continue # only count detections that are matched to multiple (>=2) masklets + # if there are multiple matched track ids, we need to find the one that appeared first; + # these later appearing ids may be removed since they may be considered as duplicates + first_appear_obj_id = ( + min(matched_trk_obj_ids, key=lambda x: obj_first_frame_idx[x]) + if not reverse + else max(matched_trk_obj_ids, key=lambda x: obj_first_frame_idx[x]) + ) + for obj_id in matched_trk_obj_ids: + if obj_id != first_appear_obj_id: + key = (first_appear_obj_id, obj_id) + overlap_pair_to_frame_inds[key].append(frame_idx) + + # b) remove a masklet if it first appears after `hotstart_diff` and it overlaps with another + # masklet (that appears earlier) for more than `self.hotstart_dup_thresh` frames + for (first_obj_id, obj_id), frame_indices in overlap_pair_to_frame_inds.items(): + if obj_id in removed_obj_ids or obj_id in obj_ids_newly_removed: + continue # skip if the object is already removed + if (obj_first_frame_idx[obj_id] > hotstart_diff and not reverse) or ( + obj_first_frame_idx[obj_id] < hotstart_diff and reverse + ): + if len(frame_indices) >= self.hotstart_dup_thresh: + obj_ids_newly_removed.add(obj_id) + logger.debug( + f"Removing object {obj_id} at frame {frame_idx} " + f"since it overlaps with another track {first_obj_id} at frames: {frame_indices}" + ) + + removed_obj_ids.update(obj_ids_newly_removed) + return obj_ids_newly_removed, rank0_metadata + + def _tracker_update_memories( + self, + tracker_inference_states: List[Any], + frame_idx: int, + tracker_metadata: Dict[str, Any], + low_res_masks: Tensor, + ): + """ + Run Sam2 memory encoder, enforcing non-overlapping constraints globally. + """ + if len(tracker_inference_states) == 0: + return + # Avoid an extra interpolation step by directly interpolating to `interpol_size` + high_res_H, high_res_W = ( + self.tracker.maskmem_backbone.mask_downsampler.interpol_size + ) + # NOTE: inspect this part if we observe OOMs in the demo + high_res_masks = F.interpolate( + low_res_masks.unsqueeze(1), + size=(high_res_H, high_res_W), + mode="bilinear", + align_corners=False, + ) + # We first apply non-overlapping constraints before memory encoding. This may include some suppression heuristics. + if not hasattr(self, "_warm_up_complete") or self._warm_up_complete: + high_res_masks = self.tracker._suppress_object_pw_area_shrinkage( + high_res_masks + ) + # Instead of gathering the predicted object scores, we use mask areas as a proxy. + object_score_logits = torch.where( + (high_res_masks > 0).any(dim=(-1, -2)), 10.0, -10.0 + ) + + # Run the memory encoder on local slices for each GPU + start_idx_gpu = sum(tracker_metadata["num_obj_per_gpu"][: self.rank]) + start_idx_state = start_idx_gpu + for tracker_state in tracker_inference_states: + num_obj_per_state = len(tracker_state["obj_ids"]) + if num_obj_per_state == 0: + continue + # Get the local high-res masks and object score logits for this inference state + end_idx_state = start_idx_state + num_obj_per_state + local_high_res_masks = high_res_masks[start_idx_state:end_idx_state] + local_object_score_logits = object_score_logits[ + start_idx_state:end_idx_state + ] + local_batch_size = local_high_res_masks.size(0) + # Run Sam2 memory encoder. Note that we do not re-enforce the non-overlapping constraint as it is turned off by default + + encoded_mem = self.tracker._run_memory_encoder( + tracker_state, + frame_idx, + local_batch_size, + local_high_res_masks, + local_object_score_logits, + is_mask_from_pts=False, + ) + local_maskmem_features, local_maskmem_pos_enc = encoded_mem + # Store encoded memories in the local inference state + output_dict = tracker_state["output_dict"] + for storage_key in ["cond_frame_outputs", "non_cond_frame_outputs"]: + if frame_idx not in output_dict[storage_key]: + continue + output_dict[storage_key][frame_idx]["maskmem_features"] = ( + local_maskmem_features + ) + output_dict[storage_key][frame_idx]["maskmem_pos_enc"] = [ + pos for pos in local_maskmem_pos_enc + ] + # for batched inference state, we also need to add per-object + # memory slides to support instance interactivity + self.tracker._add_output_per_object( + inference_state=tracker_state, + frame_idx=frame_idx, + current_out=output_dict[storage_key][frame_idx], + storage_key=storage_key, + ) + start_idx_state += num_obj_per_state + + def _tracker_add_new_objects( + self, + frame_idx: int, + num_frames: int, + new_obj_ids: List[int], + new_obj_masks: Tensor, + tracker_states_local: List[Any], + orig_vid_height: int, + orig_vid_width: int, + feature_cache: Dict, + ): + """Add a new object to SAM2 inference states.""" + prev_tracker_state = ( + tracker_states_local[0] if len(tracker_states_local) > 0 else None + ) + + # prepare inference_state + # batch objects that first appear on the same frame together + # Clear inference state. Keep the cached image features if available. + new_tracker_state = self.tracker.init_state( + cached_features=feature_cache, + video_height=orig_vid_height, + video_width=orig_vid_width, + num_frames=num_frames, + ) + new_tracker_state["backbone_out"] = ( + prev_tracker_state.get("backbone_out", None) + if prev_tracker_state is not None + else None + ) + + assert len(new_obj_ids) == new_obj_masks.size(0) + assert new_obj_masks.is_floating_point() + input_mask_res = self.tracker.input_mask_size + new_obj_masks = F.interpolate( + new_obj_masks.unsqueeze(1), + size=(input_mask_res, input_mask_res), + mode="bilinear", + align_corners=False, + ).squeeze(1) + new_obj_masks = new_obj_masks > 0 + + # add object one by one + for new_obj_id, new_mask in zip(new_obj_ids, new_obj_masks): + self.tracker.add_new_mask( + inference_state=new_tracker_state, + frame_idx=frame_idx, + obj_id=new_obj_id, + mask=new_mask, + add_mask_to_memory=True, + ) + # NOTE: we skip enforcing the non-overlapping constraint **globally** when adding new objects. + self.tracker.propagate_in_video_preflight( + new_tracker_state, run_mem_encoder=True + ) + tracker_states_local.append(new_tracker_state) + return tracker_states_local + + def _tracker_remove_object(self, tracker_states_local: List[Any], obj_id: int): + """ + Remove an object from SAM2 inference states. This would remove the object from + all frames in the video. + """ + tracker_states_local_before_removal = tracker_states_local.copy() + tracker_states_local.clear() + for tracker_inference_state in tracker_states_local_before_removal: + # we try to remove `obj_id` on every inference state with `strict=False` + # it will not do anything if an inference state doesn't contain `obj_id` + new_obj_ids, _ = self.tracker.remove_object( + tracker_inference_state, obj_id, strict=False, need_output=False + ) + # only keep an inference state if it's non-empty after object removal + if len(new_obj_ids) > 0: + tracker_states_local.append(tracker_inference_state) + + def _tracker_remove_objects( + self, tracker_states_local: List[Any], obj_ids: list[int] + ): + """ + Remove an object from SAM2 inference states. This would remove the object from + all frames in the video. + """ + for obj_id in obj_ids: + self._tracker_remove_object(tracker_states_local, obj_id) + + def _initialize_metadata(self): + """Initialize metadata for the masklets.""" + tracker_metadata = { + "obj_ids_per_gpu": [np.array([], np.int64) for _ in range(self.world_size)], + "obj_ids_all_gpu": np.array([], np.int64), + "num_obj_per_gpu": np.zeros(self.world_size, np.int64), + "max_obj_id": -1, + "obj_id_to_score": {}, + "obj_id_to_tracker_score_frame_wise": defaultdict(dict), + "obj_id_to_last_occluded": {}, + } + if self.rank == 0: + # "rank0_metadata" contains metadata that is only stored on (and accessible to) GPU 0 + # - obj_first_frame_idx: obj_id --> first frame index where the object was detected + # - unmatched_frame_inds: obj_id --> [mismatched frame indices] + # - overlap_pair_to_frame_inds: (first_appear_obj_id, obj_id) --> [overlap frame indices] + # - removed_obj_ids: object IDs that are suppressed via hot-start + rank0_metadata = { + "obj_first_frame_idx": {}, + "unmatched_frame_inds": defaultdict(list), + "trk_keep_alive": defaultdict( + int + ), # This is used only for object suppression not for removal + "overlap_pair_to_frame_inds": defaultdict(list), + "removed_obj_ids": set(), + "suppressed_obj_ids": defaultdict( + set + ), # frame_idx --> set of objects with suppressed outputs, but still continue to be tracked + } + if self.masklet_confirmation_enable: + # all the following are npt.NDArray with the same shape as `obj_ids_all_gpu` + rank0_metadata["masklet_confirmation"] = { + # "status" is the confirmation status of each masklet (in `MaskletConfirmationStatus`) + "status": np.array([], np.int64), + # "consecutive_det_num" is the number of consecutive frames where the masklet is + # detected by the detector (with a matched detection) + "consecutive_det_num": np.array([], np.int64), + } + tracker_metadata["rank0_metadata"] = rank0_metadata + + return tracker_metadata + + def update_masklet_confirmation_status( + self, + rank0_metadata: Dict[str, Any], + obj_ids_all_gpu_prev: npt.NDArray, + obj_ids_all_gpu_updated: npt.NDArray, + det_to_matched_trk_obj_ids: Dict[int, npt.NDArray], + new_det_obj_ids: npt.NDArray, + ): + confirmation_data = rank0_metadata["masklet_confirmation"] + + # a) first, expand "confirmation_data" to include new masklets added in this frame + status_prev = confirmation_data["status"] + consecutive_det_num_prev = confirmation_data["consecutive_det_num"] + assert ( + status_prev.shape == obj_ids_all_gpu_prev.shape + ), f"Got {status_prev.shape} vs {obj_ids_all_gpu_prev.shape}" + + obj_id_to_updated_idx = { + obj_id: idx for idx, obj_id in enumerate(obj_ids_all_gpu_updated) + } + prev_elem_is_in_updated = np.isin(obj_ids_all_gpu_prev, obj_ids_all_gpu_updated) + prev_elem_obj_ids_in_updated = obj_ids_all_gpu_prev[prev_elem_is_in_updated] + prev_elem_inds_in_updated = np.array( + [obj_id_to_updated_idx[obj_id] for obj_id in prev_elem_obj_ids_in_updated], + dtype=np.int64, + ) + # newly added masklets are initialized to "UNCONFIRMED" status + unconfirmed_val = MaskletConfirmationStatus.UNCONFIRMED.value + status = np.full_like(obj_ids_all_gpu_updated, fill_value=unconfirmed_val) + status[prev_elem_inds_in_updated] = status_prev[prev_elem_is_in_updated] + consecutive_det_num = np.zeros_like(obj_ids_all_gpu_updated) + consecutive_det_num[prev_elem_inds_in_updated] = consecutive_det_num_prev[ + prev_elem_is_in_updated + ] + + # b) update the confirmation status of all masklets based on the current frame + # b.1) update "consecutive_det_num" + # "is_matched": whether a masklet is matched to a detection on this frame + is_matched = np.isin(obj_ids_all_gpu_updated, new_det_obj_ids) + for matched_trk_obj_ids in det_to_matched_trk_obj_ids.values(): + is_matched |= np.isin(obj_ids_all_gpu_updated, matched_trk_obj_ids) + consecutive_det_num = np.where(is_matched, consecutive_det_num + 1, 0) + + # b.2) update "status" + change_to_confirmed = ( + consecutive_det_num >= self.masklet_confirmation_consecutive_det_thresh + ) + status[change_to_confirmed] = MaskletConfirmationStatus.CONFIRMED.value + + confirmation_data["status"] = status + confirmation_data["consecutive_det_num"] = consecutive_det_num + return rank0_metadata + + def forward(self, input: BatchedDatapoint, is_inference: bool = False): + raise NotImplementedError("Evaluation outside demo is not implemented yet") + + def _load_checkpoint(self, ckpt_path: str, strict: bool = True): + sd = torch.load(ckpt_path, map_location="cpu", weights_only=True)["model"] + missing_keys, unexpected_keys = self.load_state_dict(sd, strict=strict) + if len(missing_keys) > 0 or len(unexpected_keys) > 0: + logger.warning(f"Loaded ckpt with {missing_keys=}, {unexpected_keys=}") + else: + logger.info("Loaded ckpt successfully without missing or unexpected keys") + + def prep_for_evaluator(self, video_frames, tracking_res, scores_labels): + """This method is only used for benchmark eval (not used in the demo).""" + num_frames = len(video_frames) + w, h = video_frames[0].size + zero_mask = torch.zeros((1, h, w), dtype=torch.bool) + object_ids = list(scores_labels.keys()) + preds = {"scores": [], "labels": [], "boxes": [], "masks_rle": []} + for oid in object_ids: + o_masks = [] + o_score = scores_labels[oid][0].item() + o_label = scores_labels[oid][1] + for frame_idx in range(num_frames): + if frame_idx not in tracking_res: + o_masks.append(zero_mask) + else: + o_masks.append(tracking_res[frame_idx].get(oid, zero_mask)) + + o_masks = torch.cat(o_masks, dim=0) # (n_frames, H, W) + preds["scores"].append(o_score) + preds["labels"].append(o_label) + preds["boxes"].append(mask_to_box(o_masks.unsqueeze(1)).squeeze()) + preds["masks_rle"].append(rle_encode(o_masks, return_areas=True)) + + preds["boxes"] = ( + torch.stack(preds["boxes"], dim=0) + if len(preds["boxes"]) > 0 + else torch.empty( + (0, num_frames, 4), dtype=torch.float32, device=self.device + ) + ) + preds["scores"] = ( + torch.tensor(preds["scores"], device=self.device) + if len(preds["scores"]) > 0 + else torch.empty((0,), device=self.device) + ) + preds["per_frame_scores"] = preds["scores"] + preds["labels"] = ( + torch.tensor(preds["labels"], device=self.device) + if len(preds["labels"]) > 0 + else torch.empty((0,), device=self.device) + ) + return preds + + def _encode_prompt(self, **kwargs): + return self.detector._encode_prompt(**kwargs) + + def _drop_new_det_with_obj_limit(self, new_det_fa_inds, det_scores_np, num_to_keep): + """ + Drop a few new detections based on the maximum number of objects. We drop new objects based + on their detection scores, keeping the high-scoring ones and dropping the low-scoring ones. + """ + assert 0 <= num_to_keep <= len(new_det_fa_inds) + if num_to_keep == 0: + return np.array([], np.int64) # keep none + if num_to_keep == len(new_det_fa_inds): + return new_det_fa_inds # keep all + + # keep the top-scoring detections + score_order = np.argsort(det_scores_np[new_det_fa_inds])[::-1] + new_det_fa_inds = new_det_fa_inds[score_order[:num_to_keep]] + return new_det_fa_inds diff --git a/src/nn/segearth_ov3/sam3/model/sam3_video_inference.py b/src/nn/segearth_ov3/sam3/model/sam3_video_inference.py new file mode 100644 index 0000000..7fb87d0 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_video_inference.py @@ -0,0 +1,1709 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +from collections import defaultdict + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from sam3 import perflib +from sam3.logger import get_logger +from sam3.model.act_ckpt_utils import clone_output_wrapper +from sam3.model.box_ops import box_xywh_to_cxcywh, box_xyxy_to_xywh +from sam3.model.data_misc import BatchedDatapoint, convert_my_tensors, FindStage +from sam3.model.geometry_encoders import Prompt +from sam3.model.io_utils import IMAGE_EXTS, load_resource_as_video_frames +from sam3.model.sam3_tracker_utils import fill_holes_in_mask_scores +from sam3.model.sam3_video_base import MaskletConfirmationStatus, Sam3VideoBase +from sam3.model.utils.misc import copy_data_to_device +from sam3.perflib.compile import compile_wrapper, shape_logging_wrapper +from sam3.perflib.masks_ops import masks_to_boxes as perf_masks_to_boxes +from torchvision.ops import masks_to_boxes +from tqdm.auto import tqdm + +logger = get_logger(__name__) + + +class Sam3VideoInference(Sam3VideoBase): + TEXT_ID_FOR_TEXT = 0 + TEXT_ID_FOR_VISUAL = 1 + + def __init__( + self, + image_size=1008, + image_mean=(0.5, 0.5, 0.5), + image_std=(0.5, 0.5, 0.5), + compile_model=False, + **kwargs, + ): + """ + hotstart_delay: int, the delay (in #frames) before the model starts to yield output, 0 to disable hotstart delay. + hotstart_unmatch_thresh: int, remove the object if it has this many unmatched frames within its hotstart_delay period. + If `hotstart_delay` is set to 0, this parameter is ignored. + hotstart_dup_thresh: int, remove the object if it has overlapped with another object this many frames within its hotstart_delay period. + """ + super().__init__(**kwargs) + self.image_size = image_size + self.image_mean = image_mean + self.image_std = image_std + self.compile_model = compile_model + + @torch.inference_mode() + def init_state( + self, + resource_path, + offload_video_to_cpu=False, + async_loading_frames=False, + video_loader_type="cv2", + ): + """Initialize an inference state from `resource_path` (an image or a video).""" + images, orig_height, orig_width = load_resource_as_video_frames( + resource_path=resource_path, + image_size=self.image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=self.image_mean, + img_std=self.image_std, + async_loading_frames=async_loading_frames, + video_loader_type=video_loader_type, + ) + inference_state = {} + inference_state["image_size"] = self.image_size + inference_state["num_frames"] = len(images) + # the original video height and width, used for resizing final output scores + inference_state["orig_height"] = orig_height + inference_state["orig_width"] = orig_width + # values that don't change across frames (so we only need to hold one copy of them) + inference_state["constants"] = {} + # inputs on each frame + self._construct_initial_input_batch(inference_state, images) + # initialize extra states + inference_state["tracker_inference_states"] = [] + inference_state["tracker_metadata"] = {} + inference_state["feature_cache"] = {} + inference_state["cached_frame_outputs"] = {} + inference_state["action_history"] = [] # for logging user actions + inference_state["is_image_only"] = is_image_type(resource_path) + return inference_state + + @torch.inference_mode() + def reset_state(self, inference_state): + """Revert `inference_state` to what it was right after initialization.""" + inference_state["input_batch"].find_text_batch[0] = "" + inference_state["text_prompt"] = None + for t in range(inference_state["num_frames"]): + inference_state["input_batch"].find_inputs[t].text_ids[...] = 0 + # constructing an output list in inference state (we start with an empty list) + inference_state["previous_stages_out"][t] = None + inference_state["per_frame_raw_point_input"][t] = None + inference_state["per_frame_raw_box_input"][t] = None + inference_state["per_frame_visual_prompt"][t] = None + inference_state["per_frame_geometric_prompt"][t] = None + inference_state["per_frame_cur_step"][t] = 0 + + inference_state["visual_prompt_embed"] = None + inference_state["visual_prompt_mask"] = None + inference_state["tracker_inference_states"].clear() + inference_state["tracker_metadata"].clear() + inference_state["feature_cache"].clear() + inference_state["cached_frame_outputs"].clear() + inference_state["action_history"].clear() # for logging user actions + + def _construct_initial_input_batch(self, inference_state, images): + """Construct an initial `BatchedDatapoint` instance as input.""" + # 1) img_batch + num_frames = len(images) + device = self.device + + # 2) find_text_batch + # "" will be replaced by the actual text prompt when adding prompts + find_text_batch = ["", "visual"] + + # 3) find_inputs + input_box_embedding_dim = 258 # historical default + input_points_embedding_dim = 257 # historical default + stages = [ + FindStage( + img_ids=[stage_id], + text_ids=[0], + input_boxes=[torch.zeros(input_box_embedding_dim)], + input_boxes_mask=[torch.empty(0, dtype=torch.bool)], + input_boxes_label=[torch.empty(0, dtype=torch.long)], + input_points=[torch.empty(0, input_points_embedding_dim)], + input_points_mask=[torch.empty(0)], + object_ids=[], + ) + for stage_id in range(num_frames) + ] + for i in range(len(stages)): + stages[i] = convert_my_tensors(stages[i]) + + # construct the final `BatchedDatapoint` and cast to GPU + input_batch = BatchedDatapoint( + img_batch=images, + find_text_batch=find_text_batch, + find_inputs=stages, + find_targets=[None] * num_frames, + find_metadatas=[None] * num_frames, + ) + input_batch = copy_data_to_device(input_batch, device, non_blocking=True) + inference_state["input_batch"] = input_batch + + # construct the placeholder interactive prompts and tracking queries + bs = 1 + inference_state["constants"]["empty_geometric_prompt"] = Prompt( + box_embeddings=torch.zeros(0, bs, 4, device=device), + box_mask=torch.zeros(bs, 0, device=device, dtype=torch.bool), + box_labels=torch.zeros(0, bs, device=device, dtype=torch.long), + point_embeddings=torch.zeros(0, bs, 2, device=device), + point_mask=torch.zeros(bs, 0, device=device, dtype=torch.bool), + point_labels=torch.zeros(0, bs, device=device, dtype=torch.long), + ) + + # constructing an output list in inference state (we start with an empty list) + inference_state["previous_stages_out"] = [None] * num_frames + inference_state["text_prompt"] = None + inference_state["per_frame_raw_point_input"] = [None] * num_frames + inference_state["per_frame_raw_box_input"] = [None] * num_frames + inference_state["per_frame_visual_prompt"] = [None] * num_frames + inference_state["per_frame_geometric_prompt"] = [None] * num_frames + inference_state["per_frame_cur_step"] = [0] * num_frames + + # placeholders for cached outputs + # (note: currently, a single visual prompt embedding is shared for all frames) + inference_state["visual_prompt_embed"] = None + inference_state["visual_prompt_mask"] = None + + def _get_visual_prompt(self, inference_state, frame_idx, boxes_cxcywh, box_labels): + """ + Handle the case of visual prompt. Currently, in the inference API we do not + explicitly distinguish between initial box as visual prompt vs subsequent boxes + or boxes after inference for refinement. + """ + # If the frame hasn't had any inference results before (prompting or propagation), + # we treat the first added box prompt as a visual prompt; otherwise, we treat + # the first box just as a refinement prompt. + is_new_visual_prompt = ( + inference_state["per_frame_visual_prompt"][frame_idx] is None + and inference_state["previous_stages_out"][frame_idx] is None + ) + if is_new_visual_prompt: + if boxes_cxcywh.size(0) != 1: + raise RuntimeError( + "visual prompts (box as an initial prompt) should only have one box, " + f"but got {boxes_cxcywh.shape=}" + ) + if not box_labels.item(): + logging.warning("A negative box is added as a visual prompt.") + # take the first box prompt as a visual prompt + device = self.device + new_visual_prompt = Prompt( + box_embeddings=boxes_cxcywh[None, 0:1, :].to(device), # (seq, bs, 4) + box_mask=None, + box_labels=box_labels[None, 0:1].to(device), # (seq, bs) + point_embeddings=None, + point_mask=None, + point_labels=None, + ) + inference_state["per_frame_visual_prompt"][frame_idx] = new_visual_prompt + else: + new_visual_prompt = None + + # `boxes_cxcywh` and `box_labels` contains all the raw box inputs added so far + # strip any visual prompt from the input boxes (for geometric prompt encoding) + if inference_state["per_frame_visual_prompt"][frame_idx] is not None: + boxes_cxcywh = boxes_cxcywh[1:] + box_labels = box_labels[1:] + + return boxes_cxcywh, box_labels, new_visual_prompt + + def _get_processing_order( + self, inference_state, start_frame_idx, max_frame_num_to_track, reverse + ): + num_frames = inference_state["num_frames"] + previous_stages_out = inference_state["previous_stages_out"] + if all(out is None for out in previous_stages_out) and start_frame_idx is None: + raise RuntimeError( + "No prompts are received on any frames. Please add prompt on at least one frame before propagation." + ) + # set start index, end index, and processing order + if start_frame_idx is None: + # default: start from the earliest frame with input points + start_frame_idx = min( + t for t, out in enumerate(previous_stages_out) if out is not None + ) + if max_frame_num_to_track is None: + # default: track all the frames in the video + max_frame_num_to_track = num_frames + if reverse: + end_frame_idx = start_frame_idx - max_frame_num_to_track + end_frame_idx = max(end_frame_idx, 0) + processing_order = range(start_frame_idx - 1, end_frame_idx - 1, -1) + else: + end_frame_idx = start_frame_idx + max_frame_num_to_track + end_frame_idx = min(end_frame_idx, num_frames - 1) + processing_order = range(start_frame_idx, end_frame_idx + 1) + return processing_order, end_frame_idx + + @torch.inference_mode() + def propagate_in_video( + self, + inference_state, + start_frame_idx=None, + max_frame_num_to_track=None, + reverse=False, + ): + """ + Propagate the prompts to get grounding results for the entire video. This method + is a generator and yields inference outputs for all frames in the range specified + by `start_frame_idx`, `max_frame_num_to_track`, and `reverse`. + """ + # compile the model (it's a no-op if the model is already compiled) + # note that it's intentionally added to `self.propagate_in_video`, so that the first + # `self.add_prompt` call will be done in eager mode to fill in the decoder buffers + # such as positional encoding cache) + self._compile_model() + + processing_order, end_frame_idx = self._get_processing_order( + inference_state, + start_frame_idx, + max_frame_num_to_track, + reverse=reverse, + ) + + # Store max_frame_num_to_track in feature_cache for downstream methods + inference_state["feature_cache"]["tracking_bounds"] = { + "max_frame_num_to_track": max_frame_num_to_track, + "propagate_in_video_start_frame_idx": start_frame_idx, + } + + hotstart_buffer = [] + hotstart_removed_obj_ids = set() + # when deciding whether to output a masklet on `yield_frame_idx`, we check whether the object is confirmed + # in a future frame (`unconfirmed_frame_delay` frames after the current frame). For example, if we require + # an object to be detected in 3 consecutive frames to be confirmed, then we look 2 frames in the future -- + # e.g., we output an object on frame 4 only if it becomes confirmed on frame 6. + unconfirmed_status_delay = self.masklet_confirmation_consecutive_det_thresh - 1 + unconfirmed_obj_ids_per_frame = {} # frame_idx -> hidden_obj_ids + for frame_idx in tqdm( + processing_order, desc="propagate_in_video", disable=self.rank > 0 + ): + out = self._run_single_frame_inference(inference_state, frame_idx, reverse) + + if self.hotstart_delay > 0: + # accumulate the outputs for the first `hotstart_delay` frames + hotstart_buffer.append([frame_idx, out]) + # update the object IDs removed by hotstart so that we don't output them + if self.rank == 0: + hotstart_removed_obj_ids.update(out["removed_obj_ids"]) + unconfirmed_obj_ids = out.get("unconfirmed_obj_ids", None) + if unconfirmed_obj_ids is not None: + unconfirmed_obj_ids_per_frame[frame_idx] = unconfirmed_obj_ids + + if frame_idx == end_frame_idx: + # we reached the end of propagation -- yield all frames in the buffer + yield_list = hotstart_buffer + hotstart_buffer = [] + elif len(hotstart_buffer) >= self.hotstart_delay: + # we have enough frames -- yield and remove the first (oldest) frame from the buffer + yield_list = hotstart_buffer[:1] + hotstart_buffer = hotstart_buffer[1:] + else: + # not enough frames yet -- skip yielding + yield_list = [] + else: + yield_list = [(frame_idx, out)] # output the current frame + + for yield_frame_idx, yield_out in yield_list: + # post-process the output and yield it + if self.rank == 0: + suppressed_obj_ids = yield_out["suppressed_obj_ids"] + unconfirmed_status_frame_idx = ( + yield_frame_idx + unconfirmed_status_delay + if not reverse + else yield_frame_idx - unconfirmed_status_delay + ) + + # Clamp the frame index to stay within video bounds + num_frames = inference_state["num_frames"] + unconfirmed_status_frame_idx = max( + 0, min(unconfirmed_status_frame_idx, num_frames - 1) + ) + + unconfirmed_obj_ids = unconfirmed_obj_ids_per_frame.get( + unconfirmed_status_frame_idx, None + ) + postprocessed_out = self._postprocess_output( + inference_state, + yield_out, + hotstart_removed_obj_ids, + suppressed_obj_ids, + unconfirmed_obj_ids, + ) + + self._cache_frame_outputs( + inference_state, + yield_frame_idx, + yield_out["obj_id_to_mask"], + suppressed_obj_ids=suppressed_obj_ids, + removed_obj_ids=hotstart_removed_obj_ids, + unconfirmed_obj_ids=unconfirmed_obj_ids, + ) + else: + postprocessed_out = None # no output on other GPUs + yield yield_frame_idx, postprocessed_out + + def _run_single_frame_inference(self, inference_state, frame_idx, reverse): + """ + Perform inference on a single frame and get its inference results. This would + also update `inference_state`. + """ + # prepare inputs + input_batch = inference_state["input_batch"] + tracker_states_local = inference_state["tracker_inference_states"] + has_text_prompt = inference_state["text_prompt"] is not None + has_geometric_prompt = ( + inference_state["per_frame_geometric_prompt"][frame_idx] is not None + ) + # run inference for the current frame + ( + obj_id_to_mask, + obj_id_to_score, + tracker_states_local_new, + tracker_metadata_new, + frame_stats, + _, + ) = self._det_track_one_frame( + frame_idx=frame_idx, + num_frames=inference_state["num_frames"], + reverse=reverse, + input_batch=input_batch, + geometric_prompt=( + inference_state["constants"]["empty_geometric_prompt"] + if not has_geometric_prompt + else inference_state["per_frame_geometric_prompt"][frame_idx] + ), + tracker_states_local=tracker_states_local, + tracker_metadata_prev=inference_state["tracker_metadata"], + feature_cache=inference_state["feature_cache"], + orig_vid_height=inference_state["orig_height"], + orig_vid_width=inference_state["orig_width"], + is_image_only=inference_state["is_image_only"], + allow_new_detections=has_text_prompt or has_geometric_prompt, + ) + # update inference state + inference_state["tracker_inference_states"] = tracker_states_local_new + inference_state["tracker_metadata"] = tracker_metadata_new + # use a dummy string in "previous_stages_out" to indicate this frame has outputs + inference_state["previous_stages_out"][frame_idx] = "_THIS_FRAME_HAS_OUTPUTS_" + + if self.rank == 0: + self._cache_frame_outputs(inference_state, frame_idx, obj_id_to_mask) + + out = { + "obj_id_to_mask": obj_id_to_mask, + "obj_id_to_score": obj_id_to_score, # first frame detection score + "obj_id_to_tracker_score": tracker_metadata_new[ + "obj_id_to_tracker_score_frame_wise" + ][frame_idx], + } + # removed_obj_ids is only needed on rank 0 to handle hotstart delay buffer + if self.rank == 0: + rank0_metadata = tracker_metadata_new["rank0_metadata"] + removed_obj_ids = rank0_metadata["removed_obj_ids"] + out["removed_obj_ids"] = removed_obj_ids + out["suppressed_obj_ids"] = rank0_metadata["suppressed_obj_ids"][frame_idx] + out["frame_stats"] = frame_stats + if self.masklet_confirmation_enable: + status = rank0_metadata["masklet_confirmation"]["status"] + is_unconfirmed = status == MaskletConfirmationStatus.UNCONFIRMED.value + out["unconfirmed_obj_ids"] = tracker_metadata_new["obj_ids_all_gpu"][ + is_unconfirmed + ].tolist() + else: + out["unconfirmed_obj_ids"] = [] + + return out + + def _postprocess_output( + self, + inference_state, + out, + removed_obj_ids=None, + suppressed_obj_ids=None, + unconfirmed_obj_ids=None, + ): + obj_id_to_mask = out["obj_id_to_mask"] # low res masks + curr_obj_ids = sorted(obj_id_to_mask.keys()) + H_video, W_video = inference_state["orig_height"], inference_state["orig_width"] + if len(curr_obj_ids) == 0: + out_obj_ids = torch.zeros(0, dtype=torch.int64) + out_probs = torch.zeros(0, dtype=torch.float32) + out_binary_masks = torch.zeros(0, H_video, W_video, dtype=torch.bool) + out_boxes_xywh = torch.zeros(0, 4, dtype=torch.float32) + else: + out_obj_ids = torch.tensor(curr_obj_ids, dtype=torch.int64) + out_probs = torch.tensor( + [out["obj_id_to_score"][obj_id] for obj_id in curr_obj_ids] + ) + out_tracker_probs = torch.tensor( + [ + ( + out["obj_id_to_tracker_score"][obj_id] + if obj_id in out["obj_id_to_tracker_score"] + else 0.0 + ) + for obj_id in curr_obj_ids + ] + ) + out_binary_masks = torch.cat( + [obj_id_to_mask[obj_id] for obj_id in curr_obj_ids], dim=0 + ) + + assert out_binary_masks.dtype == torch.bool + keep = out_binary_masks.any(dim=(1, 2)).cpu() # remove masks with 0 areas + # hide outputs for those object IDs in `obj_ids_to_hide` + obj_ids_to_hide = [] + if suppressed_obj_ids is not None: + obj_ids_to_hide.extend(suppressed_obj_ids) + if removed_obj_ids is not None: + obj_ids_to_hide.extend(removed_obj_ids) + if unconfirmed_obj_ids is not None: + obj_ids_to_hide.extend(unconfirmed_obj_ids) + if len(obj_ids_to_hide) > 0: + obj_ids_to_hide_t = torch.tensor(obj_ids_to_hide, dtype=torch.int64) + keep &= ~torch.isin(out_obj_ids, obj_ids_to_hide_t) + + # slice those valid entries from the original outputs + keep_idx = torch.nonzero(keep, as_tuple=True)[0] + keep_idx_gpu = keep_idx.pin_memory().to( + device=out_binary_masks.device, non_blocking=True + ) + + out_obj_ids = torch.index_select(out_obj_ids, 0, keep_idx) + out_probs = torch.index_select(out_probs, 0, keep_idx) + out_tracker_probs = torch.index_select(out_tracker_probs, 0, keep_idx) + out_binary_masks = torch.index_select(out_binary_masks, 0, keep_idx_gpu) + + if perflib.is_enabled: + out_boxes_xyxy = perf_masks_to_boxes( + out_binary_masks, out_obj_ids.tolist() + ) + else: + out_boxes_xyxy = masks_to_boxes(out_binary_masks) + + out_boxes_xywh = box_xyxy_to_xywh(out_boxes_xyxy) # convert to xywh format + # normalize boxes + out_boxes_xywh[..., 0] /= W_video + out_boxes_xywh[..., 1] /= H_video + out_boxes_xywh[..., 2] /= W_video + out_boxes_xywh[..., 3] /= H_video + + # apply non-overlapping constraints on the existing masklets + if out_binary_masks.shape[0] > 1: + assert len(out_binary_masks) == len(out_tracker_probs) + out_binary_masks = ( + self.tracker._apply_object_wise_non_overlapping_constraints( + out_binary_masks.unsqueeze(1), + out_tracker_probs.unsqueeze(1).to(out_binary_masks.device), + background_value=0, + ).squeeze(1) + ) > 0 + + outputs = { + "out_obj_ids": out_obj_ids.cpu().numpy(), + "out_probs": out_probs.cpu().numpy(), + "out_boxes_xywh": out_boxes_xywh.cpu().numpy(), + "out_binary_masks": out_binary_masks.cpu().numpy(), + "frame_stats": out.get("frame_stats", None), + } + return outputs + + def _cache_frame_outputs( + self, + inference_state, + frame_idx, + obj_id_to_mask, + suppressed_obj_ids=None, + removed_obj_ids=None, + unconfirmed_obj_ids=None, + ): + # Filter out suppressed, removed, and unconfirmed objects from the cache + filtered_obj_id_to_mask = obj_id_to_mask.copy() + + objects_to_exclude = set() + if suppressed_obj_ids is not None: + objects_to_exclude.update(suppressed_obj_ids) + if removed_obj_ids is not None: + objects_to_exclude.update(removed_obj_ids) + if unconfirmed_obj_ids is not None: + objects_to_exclude.update(unconfirmed_obj_ids) + + if objects_to_exclude: + for obj_id in objects_to_exclude: + if obj_id in filtered_obj_id_to_mask: + del filtered_obj_id_to_mask[obj_id] + + inference_state["cached_frame_outputs"][frame_idx] = filtered_obj_id_to_mask + + def _build_tracker_output( + self, inference_state, frame_idx, refined_obj_id_to_mask=None + ): + assert ( + "cached_frame_outputs" in inference_state + and frame_idx in inference_state["cached_frame_outputs"] + ), "No cached outputs found. Ensure normal propagation has run first to populate the cache." + cached_outputs = inference_state["cached_frame_outputs"][frame_idx] + + obj_id_to_mask = cached_outputs.copy() + + # Update with refined masks if provided + if refined_obj_id_to_mask is not None: + for obj_id, refined_mask in refined_obj_id_to_mask.items(): + assert ( + refined_mask is not None + ), f"Refined mask data must be provided for obj_id {obj_id}" + obj_id_to_mask[obj_id] = refined_mask + + return obj_id_to_mask + + def _compile_model(self): + """Compile the SAM model with torch.compile for speedup.""" + is_compiled = getattr(self, "_model_is_compiled", False) + if is_compiled or not self.compile_model: + return + + import torch._dynamo + + # a larger cache size to hold varying number of shapes for torch.compile + # see https://github.com/pytorch/pytorch/blob/v2.5.1/torch/_dynamo/config.py#L42-L49 + torch._dynamo.config.cache_size_limit = 128 + torch._dynamo.config.accumulated_cache_size_limit = 2048 + torch._dynamo.config.capture_scalar_outputs = True + torch._dynamo.config.suppress_errors = True + + # Compile module components + # skip compilation of `_encode_prompt` since it sometimes tiggger SymInt errors + # self._encode_prompt = clone_output_wrapper( + # torch.compile(self._encode_prompt, fullgraph=True, mode="max-autotune") + # ) + + ## Compile SAM3 model components + self.detector.backbone.vision_backbone.forward = clone_output_wrapper( + torch.compile( + self.detector.backbone.vision_backbone.forward, + fullgraph=True, + mode="max-autotune", + ) + ) + self.detector.transformer.encoder.forward = clone_output_wrapper( + torch.compile( + self.detector.transformer.encoder.forward, + fullgraph=True, + mode="max-autotune", + ) + ) + self.detector.transformer.decoder.forward = clone_output_wrapper( + torch.compile( + self.detector.transformer.decoder.forward, + fullgraph=True, + mode="max-autotune", + dynamic=False, + ) + ) + + self.detector.segmentation_head.forward = clone_output_wrapper( + torch.compile( + self.detector.segmentation_head.forward, + fullgraph=True, + mode="max-autotune", + ) + ) + + ## Compile Tracker model components + self.tracker.maskmem_backbone.forward = compile_wrapper( + self.tracker.maskmem_backbone.forward, + mode="max-autotune", + fullgraph=True, + dynamic=False, + ) + + self.tracker.transformer.encoder.forward = shape_logging_wrapper( + compile_wrapper( + self.tracker.transformer.encoder.forward, + mode="max-autotune-no-cudagraphs", + fullgraph=True, + dynamic=True, + ), + keep_kwargs=["src", "src_pos", "prompt", "prompt_pos"], + ) + + self.tracker.sam_mask_decoder.forward = compile_wrapper( + self.tracker.sam_mask_decoder.forward, + mode="max-autotune", + fullgraph=True, + dynamic=False, # Accuracy regression on True + ) + + self._model_is_compiled = True + + def _warm_up_vg_propagation(self, inference_state, start_frame_idx=0): + # use different tracking score thresholds for each round to simulate different number of output objects + num_objects_list = range(self.num_obj_for_compile + 1) + new_det_score_thresh_list = [0.3, 0.5, 0.7] + num_rounds = len(new_det_score_thresh_list) + orig_new_det_thresh = self.new_det_thresh + + for i, thresh in enumerate(new_det_score_thresh_list): + self.new_det_thresh = thresh + for num_objects in num_objects_list: + logger.info(f"{i+1}/{num_rounds} warming up model compilation") + self.add_prompt( + inference_state, frame_idx=start_frame_idx, text_str="cat" + ) + logger.info( + f"{i+1}/{num_rounds} warming up model compilation -- simulating {num_objects}/{self.num_obj_for_compile} objects" + ) + inference_state = self.add_fake_objects_to_inference_state( + inference_state, num_objects, frame_idx=start_frame_idx + ) + inference_state["tracker_metadata"]["rank0_metadata"].update( + { + "masklet_confirmation": { + "status": np.zeros(num_objects, dtype=np.int64), + "consecutive_det_num": np.zeros( + num_objects, dtype=np.int64 + ), + } + } + ) + for _ in self.propagate_in_video( + inference_state, start_frame_idx, reverse=False + ): + pass + for _ in self.propagate_in_video( + inference_state, start_frame_idx, reverse=True + ): + pass + self.reset_state(inference_state) + logger.info( + f"{i+1}/{num_rounds} warming up model compilation -- completed round {i+1} out of {num_rounds}" + ) + + # Warm up Tracker memory encoder with varying input shapes + num_iters = 3 + feat_size = self.tracker.sam_image_embedding_size**2 # 72 * 72 = 5184 + hidden_dim = self.tracker.hidden_dim # 256 + mem_dim = self.tracker.mem_dim # 64 + for _ in tqdm(range(num_iters)): + for b in range(1, self.num_obj_for_compile + 1): + for i in range( + 1, + self.tracker.max_cond_frames_in_attn + self.tracker.num_maskmem, + ): + for j in range( + self.tracker.max_cond_frames_in_attn + + self.tracker.max_obj_ptrs_in_encoder + ): + num_obj_ptr_tokens = (hidden_dim // mem_dim) * j + src = torch.randn(feat_size, b, hidden_dim, device=self.device) + src_pos = torch.randn( + feat_size, b, hidden_dim, device=self.device + ) + prompt = torch.randn( + feat_size * i + num_obj_ptr_tokens, + b, + mem_dim, + device=self.device, + ) + prompt_pos = torch.randn( + feat_size * i + num_obj_ptr_tokens, + b, + mem_dim, + device=self.device, + ) + + self.tracker.transformer.encoder.forward( + src=src, + src_pos=src_pos, + prompt=prompt, + prompt_pos=prompt_pos, + num_obj_ptr_tokens=num_obj_ptr_tokens, + ) + + self.new_det_thresh = orig_new_det_thresh + return inference_state + + def add_fake_objects_to_inference_state( + self, inference_state, num_objects, frame_idx + ): + new_det_obj_ids_local = np.arange(num_objects) + high_res_H, high_res_W = ( + self.tracker.maskmem_backbone.mask_downsampler.interpol_size + ) + new_det_masks = torch.ones( + len(new_det_obj_ids_local), high_res_H, high_res_W + ).to(self.device) + + inference_state["tracker_inference_states"] = self._tracker_add_new_objects( + frame_idx=frame_idx, + num_frames=inference_state["num_frames"], + new_obj_ids=new_det_obj_ids_local, + new_obj_masks=new_det_masks, + tracker_states_local=inference_state["tracker_inference_states"], + orig_vid_height=inference_state["orig_height"], + orig_vid_width=inference_state["orig_width"], + feature_cache=inference_state["feature_cache"], + ) + + # Synthesize obj_id_to_mask data for cached_frame_outputs to support _build_tracker_output during warmup + obj_id_to_mask = {} + if num_objects > 0: + H_video = inference_state["orig_height"] + W_video = inference_state["orig_width"] + + video_res_masks = F.interpolate( + new_det_masks.unsqueeze(1), # Add channel dimension for interpolation + size=(H_video, W_video), + mode="bilinear", + align_corners=False, + ) # (num_objects, 1, H_video, W_video) + for i, obj_id in enumerate(new_det_obj_ids_local): + obj_id_to_mask[obj_id] = (video_res_masks[i] > 0.0).to(torch.bool) + if self.rank == 0: + for fidx in range(inference_state["num_frames"]): + self._cache_frame_outputs(inference_state, fidx, obj_id_to_mask) + + inference_state["tracker_metadata"].update( + { + "obj_ids_per_gpu": [np.arange(num_objects)], + "obj_ids_all_gpu": np.arange(num_objects), # Same as 1 GPU + "num_obj_per_gpu": [num_objects], + "obj_id_to_score": {i: 1.0 for i in range(num_objects)}, + "max_obj_id": num_objects, + "rank0_metadata": { + "masklet_confirmation": { + "status": np.zeros(num_objects, dtype=np.int64), + "consecutive_det_num": np.zeros(num_objects, dtype=np.int64), + }, + "removed_obj_ids": set(), + "suppressed_obj_ids": defaultdict(set), + }, + } + ) + return inference_state + + @torch.inference_mode() + @torch.autocast(device_type="cuda", dtype=torch.bfloat16) + def warm_up_compilation(self): + """ + Warm up the model by running a dummy inference to compile the model. This is + useful to avoid the compilation overhead in the first inference call. + """ + if not self.compile_model: + return + self._warm_up_complete = False + if self.device.type != "cuda": + raise RuntimeError( + f"The model must be on CUDA for warm-up compilation, got {self.device=}." + ) + + # temporally set to single GPU temporarily for warm-up compilation + orig_rank = self.rank + orig_world_size = self.world_size + self.rank = self.detector.rank = 0 + self.world_size = self.detector.world_size = 1 + orig_recondition_every_nth_frame = self.recondition_every_nth_frame + # self.recondition_every_nth_frame = 2 + + # Get a random video + inference_state = self.init_state(resource_path="") + start_frame_idx = 0 + + # Run basic propagation warm-up + inference_state = self._warm_up_vg_propagation(inference_state, start_frame_idx) + + logger.info("Warm-up compilation completed.") + + # revert to the original GPU and rank + self.rank = self.detector.rank = orig_rank + self.world_size = self.detector.world_size = orig_world_size + self.recondition_every_nth_frame = orig_recondition_every_nth_frame + self._warm_up_complete = True + self.tracker.transformer.encoder.forward.set_logging(True) + + @torch.inference_mode() + def add_prompt( + self, + inference_state, + frame_idx, + text_str=None, + boxes_xywh=None, + box_labels=None, + ): + """ + Add text, point or box prompts on a single frame. This method returns the inference + outputs only on the prompted frame. + + Note that text prompts are NOT associated with a particular frame (i.e. they apply + to all frames). However, we only run inference on the frame specified in `frame_idx`. + """ + logger.debug("Running add_prompt on frame %d", frame_idx) + + num_frames = inference_state["num_frames"] + assert ( + text_str is not None or boxes_xywh is not None + ), "at least one type of prompt (text, boxes) must be provided" + assert ( + 0 <= frame_idx < num_frames + ), f"{frame_idx=} is out of range for a total of {num_frames} frames" + + # since it's a semantic prompt, we start over + self.reset_state(inference_state) + + # 1) add text prompt + if text_str is not None and text_str != "visual": + inference_state["text_prompt"] = text_str + inference_state["input_batch"].find_text_batch[0] = text_str + text_id = self.TEXT_ID_FOR_TEXT + else: + inference_state["text_prompt"] = None + inference_state["input_batch"].find_text_batch[0] = "" + text_id = self.TEXT_ID_FOR_VISUAL + for t in range(inference_state["num_frames"]): + inference_state["input_batch"].find_inputs[t].text_ids[...] = text_id + + # 2) handle box prompt + assert (boxes_xywh is not None) == (box_labels is not None) + if boxes_xywh is not None: + boxes_xywh = torch.as_tensor(boxes_xywh, dtype=torch.float32) + box_labels = torch.as_tensor(box_labels, dtype=torch.long) + # input boxes are expected to be [xmin, ymin, width, height] format + # in normalized coordinates of range 0~1, similar to FA + assert boxes_xywh.dim() == 2 + assert boxes_xywh.size(0) > 0 and boxes_xywh.size(-1) == 4 + assert box_labels.dim() == 1 and box_labels.size(0) == boxes_xywh.size(0) + boxes_cxcywh = box_xywh_to_cxcywh(boxes_xywh) + assert (boxes_xywh >= 0).all().item() and (boxes_xywh <= 1).all().item() + assert (boxes_cxcywh >= 0).all().item() and (boxes_cxcywh <= 1).all().item() + + new_box_input = boxes_cxcywh, box_labels + inference_state["per_frame_raw_box_input"][frame_idx] = new_box_input + + # handle the case of visual prompt (also added as an input box from the UI) + boxes_cxcywh, box_labels, geometric_prompt = self._get_visual_prompt( + inference_state, frame_idx, boxes_cxcywh, box_labels + ) + + inference_state["per_frame_geometric_prompt"][frame_idx] = geometric_prompt + + out = self._run_single_frame_inference( + inference_state, frame_idx, reverse=False + ) + return frame_idx, self._postprocess_output(inference_state, out) + + @torch.autocast(device_type="cuda", dtype=torch.bfloat16) + def forward(self, input: BatchedDatapoint, is_inference: bool = False): + """This method is only used for benchmark eval (not used in the demo).""" + # set the model to single GPU for benchmark evaluation (to be compatible with trainer) + orig_rank = self.rank + orig_world_size = self.world_size + self.rank = self.detector.rank = 0 + self.world_size = self.detector.world_size = 1 + + # get data + text_prompt_ids = input.find_metadatas[0].original_category_id + text_prompt_list = input.find_text_batch + + # loop over txt prompts + tracking_res = defaultdict(dict) # frame_idx --> {obj_id: mask} + scores_labels = defaultdict(tuple) # obj_id --> (score, text_prompt_id) + inference_state = self.init_state(resource_path=input.raw_images) + for prompt_id, prompt in zip(text_prompt_ids, text_prompt_list): + self.add_prompt(inference_state, frame_idx=0, text_str=prompt) + start_obj_id = max(scores_labels.keys(), default=-1) + 1 # prev max + 1 + + # propagate the prompts + obj_ids_this_prompt = set() + for frame_idx, out in self.propagate_in_video( + inference_state, + start_frame_idx=0, + max_frame_num_to_track=inference_state["num_frames"], + reverse=False, + ): + current_frame_res = tracking_res[frame_idx] + for obj_id, mask in zip(out["out_obj_ids"], out["out_binary_masks"]): + mask_tensor = torch.tensor(mask[None], dtype=torch.bool) + current_frame_res[obj_id + start_obj_id] = mask_tensor + obj_ids_this_prompt.update(current_frame_res.keys()) + + obj_id_to_score = inference_state["tracker_metadata"]["obj_id_to_score"] + for obj_id, score in obj_id_to_score.items(): + if obj_id + start_obj_id in obj_ids_this_prompt: + score_tensor = torch.tensor(score, dtype=torch.float32) + scores_labels[obj_id + start_obj_id] = (score_tensor, prompt_id) + + self.reset_state(inference_state) + + video_id = input.find_metadatas[0].original_image_id[0].cpu().item() + preds = self.prep_for_evaluator(input.raw_images, tracking_res, scores_labels) + + # revert the model to the original GPU and rank + self.rank = self.detector.rank = orig_rank + self.world_size = self.detector.world_size = orig_world_size + return {video_id: preds} + + def back_convert(self, targets): + # Needed for retraining compatibility with trainer + return targets + + +class Sam3VideoInferenceWithInstanceInteractivity(Sam3VideoInference): + def __init__( + self, + use_prev_mem_frame=False, + use_stateless_refinement=False, + refinement_detector_cond_frame_removal_window=16, + **kwargs, + ): + """ + use_prev_mem_frame: bool, whether to condition on previous memory frames for adding points + use_stateless_refinement: bool, whether to enable stateless refinement behavior + refinement_detector_cond_frame_removal_window: int, we remove a detector conditioning frame if it + is within this many frames of a user refined frame. Set to a large value (e.g. 10000) to + always remove detector conditioning frames if there is any user refinement in the video. + """ + super().__init__(**kwargs) + self.use_prev_mem_frame = use_prev_mem_frame + self.use_stateless_refinement = use_stateless_refinement + self.refinement_detector_cond_frame_removal_window = ( + refinement_detector_cond_frame_removal_window + ) + + def _init_new_tracker_state(self, inference_state): + return self.tracker.init_state( + cached_features=inference_state["feature_cache"], + video_height=inference_state["orig_height"], + video_width=inference_state["orig_width"], + num_frames=inference_state["num_frames"], + ) + + @torch.inference_mode() + def propagate_in_video( + self, + inference_state, + start_frame_idx=None, + max_frame_num_to_track=None, + reverse=False, + ): + # step 1: check which type of propagation to run, should be the same for all GPUs. + propagation_type, obj_ids = self.parse_action_history_for_propagation( + inference_state + ) + self.add_action_history( + inference_state, + action_type=propagation_type, + obj_ids=obj_ids, + frame_idx=start_frame_idx, + ) + + # step 2: run full VG propagation + if propagation_type == "propagation_full": + logger.debug(f"Running full VG propagation (reverse={reverse}).") + yield from super().propagate_in_video( + inference_state, + start_frame_idx=start_frame_idx, + max_frame_num_to_track=max_frame_num_to_track, + reverse=reverse, + ) + return + + # step 3: run Tracker partial propagation or direct fetch existing predictions + assert propagation_type in ["propagation_partial", "propagation_fetch"] + logger.debug( + f"Running Tracker propagation for objects {obj_ids} and merging it with existing VG predictions (reverse={reverse})." + if propagation_type == "propagation_partial" + else f"Fetching existing VG predictions without running any propagation (reverse={reverse})." + ) + processing_order, _ = self._get_processing_order( + inference_state, + start_frame_idx=start_frame_idx, + max_frame_num_to_track=max_frame_num_to_track, + reverse=reverse, + ) + + tracker_metadata = inference_state["tracker_metadata"] + + # if fetch just return from output + if propagation_type == "propagation_fetch": + for frame_idx in tqdm(processing_order): + if self.rank == 0: + obj_id_to_mask = inference_state["cached_frame_outputs"].get( + frame_idx, {} + ) + # post processing - remove suppressed obj_ids + obj_id_to_score = tracker_metadata["obj_id_to_score"] + suppressed_obj_ids = tracker_metadata["rank0_metadata"][ + "suppressed_obj_ids" + ][frame_idx] + obj_id_to_tracker_score = tracker_metadata[ + "obj_id_to_tracker_score_frame_wise" + ][frame_idx] + + out = { + "obj_id_to_mask": obj_id_to_mask, + "obj_id_to_score": obj_id_to_score, + "obj_id_to_tracker_score": obj_id_to_tracker_score, + } + yield ( + frame_idx, + self._postprocess_output( + inference_state, out, suppressed_obj_ids=suppressed_obj_ids + ), + ) + else: + yield frame_idx, None + + return + + # get Tracker inference states containing selected obj_ids + if propagation_type == "propagation_partial": + # can be empty for GPUs where objects are not in their inference states + tracker_states_local = self._get_tracker_inference_states_by_obj_ids( + inference_state, obj_ids + ) + for tracker_state in tracker_states_local: + self.tracker.propagate_in_video_preflight( + tracker_state, run_mem_encoder=True + ) + + for frame_idx in tqdm(processing_order): + # run Tracker propagation + if propagation_type == "propagation_partial": + self._prepare_backbone_feats(inference_state, frame_idx, reverse) + obj_ids_local, low_res_masks_local, tracker_scores_local = ( + self._propogate_tracker_one_frame_local_gpu( + tracker_states_local, + frame_idx=frame_idx, + reverse=reverse, + run_mem_encoder=True, + ) + ) + + # broadcast refined object tracker scores and masks to all GPUs + # handle multiple objects that can be located on different GPUs + refined_obj_data = {} # obj_id -> (score, mask_video_res) + + # Collect data for objects on this GPU + local_obj_data = {} + for obj_id in obj_ids: + obj_rank = self._get_gpu_id_by_obj_id(inference_state, obj_id) + if self.rank == obj_rank and obj_id in obj_ids_local: + refined_obj_idx = obj_ids_local.index(obj_id) + refined_mask_low_res = low_res_masks_local[ + refined_obj_idx + ] # (H_low_res, W_low_res) + refined_score = tracker_scores_local[refined_obj_idx] + + # Keep low resolution for broadcasting to reduce communication cost + local_obj_data[obj_id] = (refined_score, refined_mask_low_res) + + # Broadcast data from each GPU that has refined objects + if self.world_size > 1: + for obj_id in obj_ids: + obj_rank = self._get_gpu_id_by_obj_id(inference_state, obj_id) + if self.rank == obj_rank: + # This GPU has the object, broadcast its data + data_to_broadcast = local_obj_data.get(obj_id, None) + data_list = [ + (data_to_broadcast[0].cpu(), data_to_broadcast[1].cpu()) + ] + self.broadcast_python_obj_cpu(data_list, src=obj_rank) + if data_to_broadcast is not None: + refined_obj_data[obj_id] = data_to_broadcast + elif self.rank != obj_rank: + # This GPU doesn't have the object, receive data + data_list = [None] + self.broadcast_python_obj_cpu(data_list, src=obj_rank) + refined_obj_data[obj_id] = ( + data_list[0][0].to(self.device), + data_list[0][1].to(self.device), + ) + else: + # Single GPU case + refined_obj_data = local_obj_data + + # Update Tracker scores for all refined objects + for obj_id, (refined_score, _) in refined_obj_data.items(): + tracker_metadata["obj_id_to_tracker_score_frame_wise"][ + frame_idx + ].update({obj_id: refined_score.item()}) + + if self.rank == 0: + # get predictions from Tracker inference states, it includes the original + # VG predictions and the refined predictions from interactivity. + + # Prepare refined masks dictionary - upscale to video resolution after broadcast + refined_obj_id_to_mask = {} + for obj_id, (_, refined_mask_low_res) in refined_obj_data.items(): + refined_mask_video_res = ( + self._convert_low_res_mask_to_video_res( + refined_mask_low_res, inference_state + ) + ) # (1, H_video, W_video) bool + refined_obj_id_to_mask[obj_id] = refined_mask_video_res + + obj_id_to_mask = self._build_tracker_output( + inference_state, frame_idx, refined_obj_id_to_mask + ) + out = { + "obj_id_to_mask": obj_id_to_mask, + "obj_id_to_score": tracker_metadata["obj_id_to_score"], + "obj_id_to_tracker_score": tracker_metadata[ + "obj_id_to_tracker_score_frame_wise" + ][frame_idx], + } + suppressed_obj_ids = tracker_metadata["rank0_metadata"][ + "suppressed_obj_ids" + ][frame_idx] + self._cache_frame_outputs( + inference_state, + frame_idx, + obj_id_to_mask, + suppressed_obj_ids=suppressed_obj_ids, + ) + suppressed_obj_ids = tracker_metadata["rank0_metadata"][ + "suppressed_obj_ids" + ][frame_idx] + yield ( + frame_idx, + self._postprocess_output( + inference_state, out, suppressed_obj_ids=suppressed_obj_ids + ), + ) + else: + yield frame_idx, None + + def add_action_history( + self, inference_state, action_type, frame_idx=None, obj_ids=None + ): + """ + action_history is used to automatically decide what to do during propagation. + action_type: one of ["add", "remove", "refine"] + ["propagation_full", "propagation_partial", "propagation_fetch"] + """ + instance_actions = ["add", "remove", "refine"] + propagation_actions = [ + "propagation_full", + "propagation_partial", + "propagation_fetch", + ] + assert ( + action_type in instance_actions + propagation_actions + ), f"Invalid action type: {action_type}, must be one of {instance_actions + propagation_actions}" + action = { + "type": action_type, + "frame_idx": frame_idx, + "obj_ids": obj_ids, + } + inference_state["action_history"].append(action) + + def _has_object_been_refined(self, inference_state, obj_id): + action_history = inference_state["action_history"] + for action in action_history: + if action["type"] in ["add", "refine"] and action.get("obj_ids"): + if obj_id in action["obj_ids"]: + return True + return False + + def parse_action_history_for_propagation(self, inference_state): + """ + Parse the actions in history before the last propagation and prepare for the next propagation. + We support multiple actions (add/remove/refine) between two propagations. If we had an action + history similar to this ["propagate", "add", "refine", "remove", "add"], the next propagation + would remove the removed object, and also propagate the two added/refined objects. + + Returns: + propagation_type: one of ["propagation_full", "propagation_partial", "propagation_fetch"] + - "propagation_full": run VG propagation for all objects + - "propagation_partial": run Tracker propagation for selected objects, useful for add/refine actions + - "propagation_fetch": fetch existing VG predictions without running any propagation + obj_ids: list of object ids to run Tracker propagation on if propagation_type is "propagation_partial". + """ + action_history = inference_state["action_history"] + if len(action_history) == 0: + # we run propagation for the first time + return "propagation_full", None + + if "propagation" in action_history[-1]["type"]: + if action_history[-1]["type"] in ["propagation_fetch"]: + # last propagation is direct fetch, we fetch existing predictions + return "propagation_fetch", None + elif action_history[-1]["type"] in [ + "propagation_partial", + "propagation_full", + ]: + # we do fetch prediction if we have already run propagation twice or we have run + # propagation once and it is from the first frame or last frame. + if ( + len(action_history) > 1 + and action_history[-2]["type"] + in ["propagation_partial", "propagation_full"] + ) or action_history[-1]["frame_idx"] in [ + 0, + inference_state["num_frames"] - 1, + ]: + # we have run both forward and backward partial/full propagation + return "propagation_fetch", None + else: + # we have run partial/full forward or backward propagation once, need run it for the rest of the frames + return action_history[-1]["type"], action_history[-1]["obj_ids"] + + # parse actions since last propagation + obj_ids = [] + for action in action_history[::-1]: + if "propagation" in action["type"]: + # we reached the last propagation action, stop parsing + break + if action["type"] in ["add", "refine"]: + obj_ids.extend(action["obj_ids"]) + # else action["type"] == "remove": noop + obj_ids = list(set(obj_ids)) if len(obj_ids) > 0 else None + propagation_type = ( + "propagation_partial" if obj_ids is not None else "propagation_fetch" + ) + return propagation_type, obj_ids + + def remove_object(self, inference_state, obj_id, is_user_action=False): + """ + We try to remove object from tracker states on every GPU, it will do nothing + for states without this object. + """ + obj_rank = self._get_gpu_id_by_obj_id(inference_state, obj_id) + assert obj_rank is not None, f"Object {obj_id} not found in any GPU." + + tracker_states_local = inference_state["tracker_inference_states"] + if self.rank == obj_rank: + self._tracker_remove_object(tracker_states_local, obj_id) + + if is_user_action: + self.add_action_history( + inference_state, action_type="remove", obj_ids=[obj_id] + ) + + # update metadata + tracker_metadata = inference_state["tracker_metadata"] + _obj_ids = tracker_metadata["obj_ids_per_gpu"][obj_rank] + tracker_metadata["obj_ids_per_gpu"][obj_rank] = _obj_ids[_obj_ids != obj_id] + tracker_metadata["num_obj_per_gpu"][obj_rank] = len( + tracker_metadata["obj_ids_per_gpu"][obj_rank] + ) + tracker_metadata["obj_ids_all_gpu"] = np.concatenate( + tracker_metadata["obj_ids_per_gpu"] + ) + tracker_metadata["obj_id_to_score"].pop(obj_id, None) + # tracker_metadata["max_obj_id"] # we do not reuse the object id, so we do not update it here + + # Clean up cached frame outputs to remove references to the deleted object + if "cached_frame_outputs" in inference_state: + for frame_idx in inference_state["cached_frame_outputs"]: + frame_cache = inference_state["cached_frame_outputs"][frame_idx] + if obj_id in frame_cache: + del frame_cache[obj_id] + + def _get_gpu_id_by_obj_id(self, inference_state, obj_id): + """ + Locate GPU ID for a given object. + """ + obj_ids_per_gpu = inference_state["tracker_metadata"]["obj_ids_per_gpu"] + for rank, obj_ids in enumerate(obj_ids_per_gpu): + if obj_id in obj_ids: + return rank + return None # object not found in any GPU + + def _get_tracker_inference_states_by_obj_ids(self, inference_state, obj_ids): + """ + Get the Tracker inference states that contain the given object ids. + This is used to run partial Tracker propagation on a single object/bucket. + Possibly multiple or zero states can be returned. + """ + states = [ + state + for state in inference_state["tracker_inference_states"] + if set(obj_ids) & set(state["obj_ids"]) + ] + return states + + def _prepare_backbone_feats(self, inference_state, frame_idx, reverse): + input_batch = inference_state["input_batch"] + feature_cache = inference_state["feature_cache"] + num_frames = inference_state["num_frames"] + geometric_prompt = ( + inference_state["constants"]["empty_geometric_prompt"] + if inference_state["per_frame_geometric_prompt"][frame_idx] is None + else inference_state["per_frame_geometric_prompt"][frame_idx] + ) + _ = self.run_backbone_and_detection( + frame_idx=frame_idx, + num_frames=num_frames, + input_batch=input_batch, + geometric_prompt=geometric_prompt, + feature_cache=feature_cache, + reverse=reverse, + allow_new_detections=True, + ) + + @torch.inference_mode() + def add_prompt( + self, + inference_state, + frame_idx, + text_str=None, + boxes_xywh=None, + box_labels=None, + points=None, + point_labels=None, + obj_id=None, + rel_coordinates=True, + ): + if points is not None: + # Tracker instance prompts + assert ( + text_str is None and boxes_xywh is None + ), "When points are provided, text_str and boxes_xywh must be None." + assert ( + obj_id is not None + ), "When points are provided, obj_id must be provided." + return self.add_tracker_new_points( + inference_state, + frame_idx, + obj_id=obj_id, + points=points, + labels=point_labels, + rel_coordinates=rel_coordinates, + use_prev_mem_frame=self.use_prev_mem_frame, + ) + else: + # SAM3 prompts + return super().add_prompt( + inference_state, + frame_idx, + text_str=text_str, + boxes_xywh=boxes_xywh, + box_labels=box_labels, + ) + + @torch.inference_mode() + def add_tracker_new_points( + self, + inference_state, + frame_idx, + obj_id, + points, + labels, + rel_coordinates=True, + use_prev_mem_frame=False, + ): + """Add a new point prompt to Tracker. Suppporting instance refinement to existing + objects by passing existing obj_id or adding a new object by passing a new obj_id. + use_prev_mem_frame=False to disable cross attention to previous memory frames. + Every GPU returns the same results, and results should contain all masks including + these masks not refined or not added by the current user points. + """ + assert obj_id is not None, "obj_id must be provided to add new points" + tracker_metadata = inference_state["tracker_metadata"] + if tracker_metadata == {}: + # initialize masklet metadata if it's uninitialized (empty dict) + tracker_metadata.update(self._initialize_metadata()) + + obj_rank = self._get_gpu_id_by_obj_id(inference_state, obj_id) + + # prepare feature + self._prepare_backbone_feats(inference_state, frame_idx, reverse=False) + + object_has_been_refined = self._has_object_been_refined(inference_state, obj_id) + if ( + obj_rank is not None + and self.use_stateless_refinement + and not object_has_been_refined + ): + # The first time we start refinement on the object, we remove it. + logger.debug( + f"[rank={self.rank}] Removing object {obj_id} before refinement." + ) + self.remove_object(inference_state, obj_id, is_user_action=False) + obj_rank = None + + if obj_rank is None: + # new object, we assign it a GPU and create a new inference state if limit allows + num_prev_obj = np.sum(tracker_metadata["num_obj_per_gpu"]) + if num_prev_obj >= self.max_num_objects: + logger.warning( + f"add_tracker_new_points: cannot add a new object as we are already tracking {num_prev_obj=} " + f"masklets (under {self.max_num_objects=})" + ) + obj_ids = [] + H_low_res = W_low_res = self.tracker.low_res_mask_size + H_video_res = inference_state["orig_height"] + W_video_res = inference_state["orig_width"] + low_res_masks = torch.zeros(0, 1, H_low_res, W_low_res) + video_res_masks = torch.zeros(0, 1, H_video_res, W_video_res) + return frame_idx, obj_ids, low_res_masks, video_res_masks + + new_det_gpu_ids = self._assign_new_det_to_gpus( + new_det_num=1, + prev_workload_per_gpu=tracker_metadata["num_obj_per_gpu"], + ) + obj_rank = new_det_gpu_ids[0] + + # get tracker inference state for the new object + if self.rank == obj_rank: + # for batched inference, we create a new inference state + tracker_state = self._init_new_tracker_state(inference_state) + inference_state["tracker_inference_states"].append(tracker_state) + + # update metadata + tracker_metadata["obj_ids_per_gpu"][obj_rank] = np.concatenate( + [ + tracker_metadata["obj_ids_per_gpu"][obj_rank], + np.array([obj_id], dtype=np.int64), + ] + ) + tracker_metadata["num_obj_per_gpu"][obj_rank] = len( + tracker_metadata["obj_ids_per_gpu"][obj_rank] + ) + tracker_metadata["obj_ids_all_gpu"] = np.concatenate( + tracker_metadata["obj_ids_per_gpu"] + ) + tracker_metadata["max_obj_id"] = max(tracker_metadata["max_obj_id"], obj_id) + + logger.debug( + f"[rank={self.rank}] Adding new object with id {obj_id} at frame {frame_idx}." + ) + self.add_action_history( + inference_state, "add", frame_idx=frame_idx, obj_ids=[obj_id] + ) + else: + # existing object, for refinement + if self.rank == obj_rank: + tracker_states = self._get_tracker_inference_states_by_obj_ids( + inference_state, [obj_id] + ) + assert ( + len(tracker_states) == 1 + ), f"[rank={self.rank}] Multiple Tracker inference states found for the same object id." + tracker_state = tracker_states[0] + + # log + logger.debug( + f"[rank={self.rank}] Refining existing object with id {obj_id} at frame {frame_idx}." + ) + self.add_action_history( + inference_state, "refine", frame_idx=frame_idx, obj_ids=[obj_id] + ) + + # assign higher score to added/refined object + tracker_metadata["obj_id_to_score"][obj_id] = 1.0 + tracker_metadata["obj_id_to_tracker_score_frame_wise"][frame_idx][obj_id] = 1.0 + + if self.rank == 0: + rank0_metadata = tracker_metadata.get("rank0_metadata", {}) + + if "removed_obj_ids" in rank0_metadata: + rank0_metadata["removed_obj_ids"].discard(obj_id) + + if "suppressed_obj_ids" in rank0_metadata: + for frame_id in rank0_metadata["suppressed_obj_ids"]: + rank0_metadata["suppressed_obj_ids"][frame_id].discard(obj_id) + + if "masklet_confirmation" in rank0_metadata: + obj_ids_all_gpu = tracker_metadata["obj_ids_all_gpu"] + obj_indices = np.where(obj_ids_all_gpu == obj_id)[0] + if len(obj_indices) > 0: + obj_idx = obj_indices[0] + if obj_idx < len(rank0_metadata["masklet_confirmation"]["status"]): + rank0_metadata["masklet_confirmation"]["status"][obj_idx] = 1 + rank0_metadata["masklet_confirmation"]["consecutive_det_num"][ + obj_idx + ] = self.masklet_confirmation_consecutive_det_thresh + + if self.rank == obj_rank: + frame_idx, obj_ids, low_res_masks, video_res_masks = ( + self.tracker.add_new_points( + inference_state=tracker_state, + frame_idx=frame_idx, + obj_id=obj_id, + points=points, + labels=labels, + clear_old_points=True, + rel_coordinates=rel_coordinates, + use_prev_mem_frame=use_prev_mem_frame, + ) + ) + + if video_res_masks is not None and len(video_res_masks) > 0: + video_res_masks = fill_holes_in_mask_scores( + video_res_masks, # shape (N, 1, H_video, W_video) + max_area=self.fill_hole_area, + fill_holes=True, + remove_sprinkles=True, + ) + + # Since the mem encoder has already run for the current input points? + self.tracker.propagate_in_video_preflight( + tracker_state, run_mem_encoder=True + ) + # Clear detector conditioning frames when user clicks are received to allow + # model updating masks on these frames. It is a noop if user is refining on the + # detector conditioning frames or adding new objects. + self.clear_detector_added_cond_frame_in_tracker( + tracker_state, obj_id, frame_idx + ) + + # fetch results from states and gather across GPUs + # Use optimized caching approach to avoid reprocessing unmodified objects + if self.rank == obj_rank and len(obj_ids) > 0: + new_mask_data = (video_res_masks[obj_ids.index(obj_id)] > 0.0).to( + torch.bool + ) + else: + new_mask_data = None + # Broadcast the new mask data across all ranks for consistency + if self.world_size > 1: + data_list = [new_mask_data.cpu() if new_mask_data is not None else None] + self.broadcast_python_obj_cpu(data_list, src=obj_rank) + new_mask_data = data_list[0].to(self.device) + + if self.rank == 0: + obj_id_to_mask = self._build_tracker_output( + inference_state, + frame_idx, + {obj_id: new_mask_data} if new_mask_data is not None else None, + ) + # post processing - remove suppressed obj_ids + obj_id_to_score = tracker_metadata["obj_id_to_score"] + suppressed_obj_ids = tracker_metadata["rank0_metadata"][ + "suppressed_obj_ids" + ][frame_idx] + obj_id_to_tracker_score = tracker_metadata[ + "obj_id_to_tracker_score_frame_wise" + ][frame_idx] + + out = { + "obj_id_to_mask": obj_id_to_mask, + "obj_id_to_score": obj_id_to_score, + "obj_id_to_tracker_score": obj_id_to_tracker_score, + } + self._cache_frame_outputs( + inference_state, + frame_idx, + obj_id_to_mask, + suppressed_obj_ids=suppressed_obj_ids, + ) + return frame_idx, self._postprocess_output( + inference_state, out, suppressed_obj_ids=suppressed_obj_ids + ) + else: + return frame_idx, None # no output on other GPUs + + def _gather_obj_id_to_mask_across_gpus(self, inference_state, obj_id_to_mask_local): + """Gather obj_id_to_mask from all GPUs. Optionally resize the masks to the video resolution.""" + tracker_metadata = inference_state["tracker_metadata"] + + # concatenate the output masklets from all local inference states + H_mask = W_mask = self.tracker.low_res_mask_size + obj_ids_local = tracker_metadata["obj_ids_per_gpu"][self.rank] + low_res_masks_local = [] + for obj_id in obj_ids_local: + if obj_id in obj_id_to_mask_local: + low_res_masks_local.append(obj_id_to_mask_local[obj_id]) + else: + low_res_masks_local.append( + torch.full((H_mask, W_mask), -1024.0, device=self.device) + ) + if len(low_res_masks_local) > 0: + low_res_masks_local = torch.stack(low_res_masks_local, dim=0) # (N, H, W) + assert low_res_masks_local.shape[1:] == (H_mask, W_mask) + else: + low_res_masks_local = torch.zeros(0, H_mask, W_mask, device=self.device) + + # all-gather `low_res_masks_local` into `low_res_masks_global` + # - low_res_masks_global: Tensor -- (num_global_obj, H_mask, W_mask) + if self.world_size > 1: + low_res_masks_local = low_res_masks_local.float().contiguous() + low_res_masks_peers = [ + low_res_masks_local.new_empty(num_obj, H_mask, W_mask) + for num_obj in tracker_metadata["num_obj_per_gpu"] + ] + dist.all_gather(low_res_masks_peers, low_res_masks_local) + low_res_masks_global = torch.cat(low_res_masks_peers, dim=0) + else: + low_res_masks_global = low_res_masks_local + return low_res_masks_global + + def _convert_low_res_mask_to_video_res(self, low_res_mask, inference_state): + """ + Convert a low-res mask to video resolution, matching the format expected by _build_tracker_output. + + Args: + low_res_mask: Tensor of shape (H_low_res, W_low_res) + inference_state: Contains video dimensions + + Returns: + video_res_mask: Tensor of shape (1, H_video, W_video) bool + """ + if low_res_mask is None: + return None + + # Convert to 3D for interpolation: (H_low_res, W_low_res) -> (1, H_low_res, W_low_res) + low_res_mask_3d = low_res_mask.unsqueeze(0).unsqueeze(0) + + # Get video dimensions + H_video = inference_state["orig_height"] + W_video = inference_state["orig_width"] + + video_res_mask = F.interpolate( + low_res_mask_3d.float(), + size=(H_video, W_video), + mode="bilinear", + align_corners=False, + ) # (1, H_video, W_video) + + # Convert to boolean - already in the right shape! + return (video_res_mask.squeeze(0) > 0.0).to(torch.bool) + + def clear_detector_added_cond_frame_in_tracker( + self, tracker_state, obj_id, refined_frame_idx + ): + """Clear detector added conditioning frame if it is within a predefined window + of the refined frame. This allow model to update masks on these frames.""" + obj_idx = self.tracker._obj_id_to_idx(tracker_state, obj_id) + + mask_only_cond_frame_indices = [] + window = self.refinement_detector_cond_frame_removal_window + for frame_idx in tracker_state["mask_inputs_per_obj"][obj_idx]: + if frame_idx not in tracker_state["point_inputs_per_obj"][obj_idx]: + # clear conditioning frames within a window of the refined frame + if abs(frame_idx - refined_frame_idx) <= window: + mask_only_cond_frame_indices.append(frame_idx) + + # clear + if len(mask_only_cond_frame_indices) > 0: + for frame_idx in mask_only_cond_frame_indices: + # obj_ids_on_this_frame is essentially all obj_ids in the state + # since they are bucket batched + obj_ids_on_this_frame = tracker_state["obj_id_to_idx"].keys() + for obj_id2 in obj_ids_on_this_frame: + self.tracker.clear_all_points_in_frame( + tracker_state, frame_idx, obj_id2, need_output=False + ) + logger.debug( + f"Cleared detector mask only conditioning frames ({mask_only_cond_frame_indices}) in Tracker." + ) + return + + +def is_image_type(resource_path: str) -> bool: + if isinstance(resource_path, list): + return len(resource_path) == 1 + return resource_path.lower().endswith(tuple(IMAGE_EXTS)) diff --git a/src/nn/segearth_ov3/sam3/model/sam3_video_predictor.py b/src/nn/segearth_ov3/sam3/model/sam3_video_predictor.py new file mode 100644 index 0000000..c639e1d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/sam3_video_predictor.py @@ -0,0 +1,521 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import datetime +import gc +import multiprocessing as mp +import os +import queue +import socket +import sys +import time +import uuid +from contextlib import closing +from typing import List, Optional + +import psutil +import torch + +from sam3.logger import get_logger + +logger = get_logger(__name__) + + +class Sam3VideoPredictor: + # a global dictionary that holds all inference states for this model (key is session_id) + _ALL_INFERENCE_STATES = {} + + def __init__( + self, + checkpoint_path=None, + bpe_path=None, + has_presence_token=True, + geo_encoder_use_img_cross_attn=True, + strict_state_dict_loading=True, + async_loading_frames=False, + video_loader_type="cv2", + apply_temporal_disambiguation: bool = True, + ): + self.async_loading_frames = async_loading_frames + self.video_loader_type = video_loader_type + from sam3.model_builder import build_sam3_video_model + + self.model = ( + build_sam3_video_model( + checkpoint_path=checkpoint_path, + bpe_path=bpe_path, + has_presence_token=has_presence_token, + geo_encoder_use_img_cross_attn=geo_encoder_use_img_cross_attn, + strict_state_dict_loading=strict_state_dict_loading, + apply_temporal_disambiguation=apply_temporal_disambiguation, + ) + .cuda() + .eval() + ) + + @torch.inference_mode() + def handle_request(self, request): + """Dispatch a request based on its type.""" + request_type = request["type"] + if request_type == "start_session": + return self.start_session( + resource_path=request["resource_path"], + session_id=request.get("session_id", None), + ) + elif request_type == "add_prompt": + return self.add_prompt( + session_id=request["session_id"], + frame_idx=request["frame_index"], + text=request.get("text", None), + points=request.get("points", None), + point_labels=request.get("point_labels", None), + bounding_boxes=request.get("bounding_boxes", None), + bounding_box_labels=request.get("bounding_box_labels", None), + obj_id=request.get("obj_id", None), + ) + elif request_type == "remove_object": + return self.remove_object( + session_id=request["session_id"], + obj_id=request["obj_id"], + is_user_action=request.get("is_user_action", True), + ) + elif request_type == "reset_session": + return self.reset_session(session_id=request["session_id"]) + elif request_type == "close_session": + return self.close_session(session_id=request["session_id"]) + else: + raise RuntimeError(f"invalid request type: {request_type}") + + @torch.inference_mode() + def handle_stream_request(self, request): + """Dispatch a stream request based on its type.""" + request_type = request["type"] + if request_type == "propagate_in_video": + yield from self.propagate_in_video( + session_id=request["session_id"], + propagation_direction=request.get("propagation_direction", "both"), + start_frame_idx=request.get("start_frame_index", None), + max_frame_num_to_track=request.get("max_frame_num_to_track", None), + ) + else: + raise RuntimeError(f"invalid request type: {request_type}") + + def start_session(self, resource_path, session_id=None): + """ + Start a new inference session on an image or a video. Here `resource_path` + can be either a path to an image file (for image inference) or an MP4 file + or directory with JPEG video frames (for video inference). + + If `session_id` is defined, it will be used as identifier for the + session. If it is not defined, the start_session function will create + a session id and return it. + """ + # get an initial inference_state from the model + inference_state = self.model.init_state( + resource_path=resource_path, + async_loading_frames=self.async_loading_frames, + video_loader_type=self.video_loader_type, + ) + if not session_id: + session_id = str(uuid.uuid4()) + self._ALL_INFERENCE_STATES[session_id] = { + "state": inference_state, + "session_id": session_id, + "start_time": time.time(), + } + logger.debug( + f"started new session {session_id}; {self._get_session_stats()}; " + f"{self._get_torch_and_gpu_properties()}" + ) + return {"session_id": session_id} + + def add_prompt( + self, + session_id: str, + frame_idx: int, + text: Optional[str] = None, + points: Optional[List[List[float]]] = None, + point_labels: Optional[List[int]] = None, + bounding_boxes: Optional[List[List[float]]] = None, + bounding_box_labels: Optional[List[int]] = None, + obj_id: Optional[int] = None, + ): + """Add text, box and/or point prompt on a specific video frame.""" + logger.debug( + f"add prompt on frame {frame_idx} in session {session_id}: " + f"{text=}, {points=}, {point_labels=}, " + f"{bounding_boxes=}, {bounding_box_labels=}" + ) + session = self._get_session(session_id) + inference_state = session["state"] + + frame_idx, outputs = self.model.add_prompt( + inference_state=inference_state, + frame_idx=frame_idx, + text_str=text, + points=points, + point_labels=point_labels, + boxes_xywh=bounding_boxes, + box_labels=bounding_box_labels, + obj_id=obj_id, + ) + return {"frame_index": frame_idx, "outputs": outputs} + + def remove_object( + self, + session_id: str, + obj_id: int, + is_user_action: bool = True, + ): + """Remove an object from tracking.""" + logger.debug( + f"remove object {obj_id} in session {session_id}: " f"{is_user_action=}" + ) + session = self._get_session(session_id) + inference_state = session["state"] + + self.model.remove_object( + inference_state=inference_state, + obj_id=obj_id, + is_user_action=is_user_action, + ) + return {"is_success": True} + + def propagate_in_video( + self, + session_id, + propagation_direction, + start_frame_idx, + max_frame_num_to_track, + ): + """Propagate the added prompts to get grounding results on all video frames.""" + logger.debug( + f"propagate in video in session {session_id}: " + f"{propagation_direction=}, {start_frame_idx=}, {max_frame_num_to_track=}" + ) + try: + session = self._get_session(session_id) + inference_state = session["state"] + if propagation_direction not in ["both", "forward", "backward"]: + raise ValueError( + f"invalid propagation direction: {propagation_direction}" + ) + + # First doing the forward propagation + if propagation_direction in ["both", "forward"]: + for frame_idx, outputs in self.model.propagate_in_video( + inference_state=inference_state, + start_frame_idx=start_frame_idx, + max_frame_num_to_track=max_frame_num_to_track, + reverse=False, + ): + yield {"frame_index": frame_idx, "outputs": outputs} + # Then doing the backward propagation (reverse in time) + if propagation_direction in ["both", "backward"]: + for frame_idx, outputs in self.model.propagate_in_video( + inference_state=inference_state, + start_frame_idx=start_frame_idx, + max_frame_num_to_track=max_frame_num_to_track, + reverse=True, + ): + yield {"frame_index": frame_idx, "outputs": outputs} + finally: + # Log upon completion (so that e.g. we can see if two propagations happen in parallel). + # Using `finally` here to log even when the tracking is aborted with GeneratorExit. + logger.debug( + f"propagation ended in session {session_id}; {self._get_session_stats()}" + ) + + def reset_session(self, session_id): + """Reset the session to its initial state (as when it's initial opened).""" + logger.debug(f"reset session {session_id}") + session = self._get_session(session_id) + inference_state = session["state"] + self.model.reset_state(inference_state) + return {"is_success": True} + + def close_session(self, session_id): + """ + Close a session. This method is idempotent and can be called multiple + times on the same "session_id". + """ + session = self._ALL_INFERENCE_STATES.pop(session_id, None) + if session is None: + logger.warning( + f"cannot close session {session_id} as it does not exist (it might have expired); " + f"{self._get_session_stats()}" + ) + else: + del session + gc.collect() + logger.info(f"removed session {session_id}; {self._get_session_stats()}") + return {"is_success": True} + + def _get_session(self, session_id): + session = self._ALL_INFERENCE_STATES.get(session_id, None) + if session is None: + raise RuntimeError( + f"Cannot find session {session_id}; it might have expired" + ) + return session + + def _get_session_stats(self): + """Get a statistics string for live sessions and their GPU usage.""" + # print both the session ids and their video frame numbers + live_session_strs = [ + f"'{session_id}' ({session['state']['num_frames']} frames)" + for session_id, session in self._ALL_INFERENCE_STATES.items() + ] + session_stats_str = ( + f"live sessions: [{', '.join(live_session_strs)}], GPU memory: " + f"{torch.cuda.memory_allocated() // 1024**2} MiB used and " + f"{torch.cuda.memory_reserved() // 1024**2} MiB reserved" + f" (max over time: {torch.cuda.max_memory_allocated() // 1024**2} MiB used " + f"and {torch.cuda.max_memory_reserved() // 1024**2} MiB reserved)" + ) + return session_stats_str + + def _get_torch_and_gpu_properties(self): + """Get a string for PyTorch and GPU properties (for logging and debugging).""" + torch_and_gpu_str = ( + f"torch: {torch.__version__} with CUDA arch {torch.cuda.get_arch_list()}, " + f"GPU device: {torch.cuda.get_device_properties(torch.cuda.current_device())}" + ) + return torch_and_gpu_str + + def shutdown(self): + """Shutdown the predictor and clear all sessions.""" + self._ALL_INFERENCE_STATES.clear() + + +class Sam3VideoPredictorMultiGPU(Sam3VideoPredictor): + def __init__(self, *model_args, gpus_to_use=None, **model_kwargs): + if gpus_to_use is None: + # if not specified, use only the current GPU by default + gpus_to_use = [torch.cuda.current_device()] + + IS_MAIN_PROCESS = os.getenv("IS_MAIN_PROCESS", "1") == "1" + if IS_MAIN_PROCESS: + gpus_to_use = sorted(set(gpus_to_use)) + logger.info(f"using the following GPU IDs: {gpus_to_use}") + assert len(gpus_to_use) > 0 and all(isinstance(i, int) for i in gpus_to_use) + assert all(0 <= i < torch.cuda.device_count() for i in gpus_to_use) + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = f"{self._find_free_port()}" + os.environ["RANK"] = "0" + os.environ["WORLD_SIZE"] = f"{len(gpus_to_use)}" + + self.gpus_to_use = gpus_to_use + self.rank = int(os.environ["RANK"]) + self.world_size = int(os.environ["WORLD_SIZE"]) + self.rank_str = f"rank={self.rank} with world_size={self.world_size}" + self.device = torch.device(f"cuda:{self.gpus_to_use[self.rank]}") + torch.cuda.set_device(self.device) + self.has_shutdown = False + if self.rank == 0: + logger.info("\n\n\n\t*** START loading model on all ranks ***\n\n") + + logger.info(f"loading model on {self.rank_str} -- this could take a while ...") + super().__init__(*model_args, **model_kwargs) + logger.info(f"loading model on {self.rank_str} -- DONE locally") + + if self.world_size > 1 and self.rank == 0: + # start the worker processes *after* the model is loaded in the main process + # so that the main process can run torch.compile and fill the cache first + self._start_worker_processes(*model_args, **model_kwargs) + for rank in range(1, self.world_size): + self.command_queues[rank].put(("start_nccl_process_group", None)) + self._start_nccl_process_group() + + if self.rank == 0: + logger.info("\n\n\n\t*** DONE loading model on all ranks ***\n\n") + + @torch.inference_mode() + def handle_request(self, request): + """Dispatch a request based on its type.""" + if self.has_shutdown: + raise RuntimeError( + "cannot handle request after the predictor has shutdown; please create a new predictor" + ) + + # when starting a session, we need to create a session id before dispatching + # the request to the workers + if request["type"] == "start_session" and request.get("session_id") is None: + request["session_id"] = str(uuid.uuid4()) + # dispatch the request to all worker processes + if self.world_size > 1 and self.rank == 0: + for rank in range(1, self.world_size): + self.command_queues[rank].put((request, False)) + + response = super().handle_request(request) + + if self.world_size > 1: + torch.distributed.barrier() # wait for all ranks to finish + return response + + @torch.inference_mode() + def handle_stream_request(self, request): + """Dispatch a stream request based on its type.""" + if self.has_shutdown: + raise RuntimeError( + "cannot handle request after the predictor has shutdown; please create a new predictor" + ) + + # dispatch the request to all worker processes + if self.world_size > 1 and self.rank == 0: + for rank in range(1, self.world_size): + self.command_queues[rank].put((request, True)) + + yield from super().handle_stream_request(request) + + if self.world_size > 1: + torch.distributed.barrier() # wait for all ranks to finish + + def _start_worker_processes(self, *model_args, **model_kwargs): + """Start worker processes for handling model inference.""" + world_size = self.world_size + logger.info(f"spawning {world_size - 1} worker processes") + # Use "spawn" (instead of "fork") for different PyTorch or CUDA context + mp_ctx = mp.get_context("spawn") + self.command_queues = {rank: mp_ctx.Queue() for rank in range(1, world_size)} + self.result_queues = {rank: mp_ctx.Queue() for rank in range(1, world_size)} + parent_pid = os.getpid() + for rank in range(1, world_size): + # set the environment variables for each worker process + os.environ["IS_MAIN_PROCESS"] = "0" # mark this as a worker process + os.environ["RANK"] = f"{rank}" + worker_process = mp_ctx.Process( + target=Sam3VideoPredictorMultiGPU._worker_process_command_loop, + args=( + rank, + world_size, + self.command_queues[rank], + self.result_queues[rank], + model_args, + model_kwargs, + self.gpus_to_use, + parent_pid, + ), + daemon=True, + ) + worker_process.start() + # revert the environment variables for the main process + os.environ["IS_MAIN_PROCESS"] = "1" + os.environ["RANK"] = "0" + # wait for all the worker processes to load the model and collect their PIDs + self.worker_pids = {} + for rank in range(1, self.world_size): + # a large timeout to cover potentially long model loading time due to compilation + _, worker_pid = self.result_queues[rank].get(timeout=7200) + self.worker_pids[rank] = worker_pid + logger.info(f"spawned {world_size - 1} worker processes") + + def _start_nccl_process_group(self): + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + if world_size == 1: + return + + logger.debug(f"starting NCCL process group on {rank=} with {world_size=}") + assert not torch.distributed.is_initialized() + # use the "env://" init method with environment variables set in start_worker_processes + # a short 3-min timeout to quickly detect any synchronization failures + timeout_sec = int(os.getenv("SAM3_COLLECTIVE_OP_TIMEOUT_SEC", "180")) + timeout = datetime.timedelta(seconds=timeout_sec) + torch.distributed.init_process_group( + backend="nccl", + init_method="env://", + timeout=timeout, + device_id=self.device, + ) + # warm-up the NCCL process group by running a dummy all-reduce + tensor = torch.ones(1024, 1024).cuda() + torch.distributed.all_reduce(tensor) + logger.debug(f"started NCCL process group on {rank=} with {world_size=}") + + def _find_free_port(self) -> int: + """ + Find a free port (a random free port from 1024 to 65535 will be selected) + https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number) + """ + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: + s.bind(("", 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + return s.getsockname()[1] + + @staticmethod + def _worker_process_command_loop( + rank, + world_size, + command_queue, + result_queue, + model_args, + model_kwargs, + gpus_to_use, + parent_pid, + ): + """ + The command loop for each worker process. It listens to commands from the main process + and executes them using the model. + """ + logger.info(f"starting worker process {rank=} with {world_size=}") + # verify that the environment variables are set correctly + assert int(os.environ["IS_MAIN_PROCESS"]) == 0 + assert int(os.environ["RANK"]) == rank + assert int(os.environ["WORLD_SIZE"]) == world_size + # load the model in this worker process + predictor = Sam3VideoPredictorMultiGPU( + *model_args, gpus_to_use=gpus_to_use, **model_kwargs + ) + logger.info(f"started worker {rank=} with {world_size=}") + # return the worker process id to the main process for bookkeeping + worker_pid = os.getpid() + result_queue.put(("load_model", worker_pid)) + + # wait for the command to start the NCCL process group + request_type, _ = command_queue.get(timeout=7200) + assert request_type == "start_nccl_process_group" + predictor._start_nccl_process_group() + + # keep listening to commands from the main process + while True: + try: + request, is_stream_request = command_queue.get(timeout=5.0) + if request == "shutdown": + logger.info(f"worker {rank=} shutting down") + torch.distributed.destroy_process_group() + result_queue.put(("shutdown", True)) # acknowledge the shutdown + sys.exit(0) + + logger.debug(f"worker {rank=} received request {request['type']=}") + if is_stream_request: + for _ in predictor.handle_stream_request(request): + pass # handle stream requests in a generator fashion + else: + predictor.handle_request(request) + except queue.Empty: + # Usually Python's multiprocessing module will shutdown all the daemon worker + # processes when the main process exits gracefully. However, the user may kill + # the main process using SIGKILL and thereby leaving no chance for the main process + # to clean up its daemon child processes. So here we manually check whether the + # parent process still exists (every 5 sec as in `command_queue.get` timeout). + if not psutil.pid_exists(parent_pid): + logger.info( + f"stopping worker {rank=} as its parent process has exited" + ) + sys.exit(1) + except Exception as e: + logger.error(f"worker {rank=} exception: {e}", exc_info=True) + + def shutdown(self): + """Shutdown all worker processes.""" + if self.rank == 0 and self.world_size > 1: + logger.info(f"shutting down {self.world_size - 1} worker processes") + for rank in range(1, self.world_size): + self.command_queues[rank].put(("shutdown", False)) + torch.distributed.destroy_process_group() + for rank in range(1, self.world_size): + self.result_queues[rank].get() # wait for the worker to acknowledge + logger.info(f"shut down {self.world_size - 1} worker processes") + self.has_shutdown = True + + super().shutdown() diff --git a/src/nn/segearth_ov3/sam3/model/text_encoder_ve.py b/src/nn/segearth_ov3/sam3/model/text_encoder_ve.py new file mode 100644 index 0000000..b1cf145 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/text_encoder_ve.py @@ -0,0 +1,328 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from collections import OrderedDict +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + +from .model_misc import LayerScale + + +class ResidualAttentionBlock(nn.Module): + def __init__( + self, + d_model: int, + n_head: int, + mlp_ratio: float = 4.0, + ls_init_value: Optional[float] = None, + act_layer: Callable[[], nn.Module] = nn.GELU, + norm_layer: Callable[[int], nn.Module] = nn.LayerNorm, + ): + super().__init__() + # Attention + self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True) + + # LayerNorm, LayerScale + self.ln_1 = norm_layer(d_model) + self.ln_2 = norm_layer(d_model) + + self.ls_1 = ( + LayerScale(d_model, ls_init_value) + if ls_init_value is not None + else nn.Identity() + ) + self.ls_2 = ( + LayerScale(d_model, ls_init_value) + if ls_init_value is not None + else nn.Identity() + ) + + # MLP + mlp_width = int(d_model * mlp_ratio) + self.mlp = nn.Sequential( + OrderedDict( + [ + ("c_fc", nn.Linear(d_model, mlp_width)), + ("gelu", act_layer()), + ("c_proj", nn.Linear(mlp_width, d_model)), + ] + ) + ) + + def attention( + self, + q_x: torch.Tensor, + k_x: Optional[torch.Tensor] = None, + v_x: Optional[torch.Tensor] = None, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + k_x = k_x if k_x is not None else q_x + v_x = v_x if v_x is not None else q_x + if attn_mask is not None: + # Leave boolean masks as is + if not attn_mask.dtype == torch.bool: + attn_mask = attn_mask.to(q_x.dtype) + + return self.attn(q_x, k_x, v_x, need_weights=False, attn_mask=attn_mask)[0] + + def forward( + self, + q_x: torch.Tensor, + k_x: Optional[torch.Tensor] = None, + v_x: Optional[torch.Tensor] = None, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + k_x = ( + self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None + ) + v_x = ( + self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None + ) + x = q_x + self.ls_1( + self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask) + ) + x = x + self.ls_2(self.mlp(self.ln_2(x))) + return x + + +class Transformer(nn.Module): + def __init__( + self, + width: int, + layers: int, + heads: int, + mlp_ratio: float = 4.0, + ls_init_value: Optional[float] = None, + act_layer: Callable[[], nn.Module] = nn.GELU, + norm_layer: Callable[[int], nn.Module] = nn.LayerNorm, + compile_mode: Optional[str] = None, + use_act_checkpoint: bool = False, + ): + super().__init__() + self.width = width + self.layers = layers + self.grad_checkpointing = use_act_checkpoint + self.resblocks = nn.ModuleList( + [ + ResidualAttentionBlock( + width, + heads, + mlp_ratio, + ls_init_value=ls_init_value, + act_layer=act_layer, + norm_layer=norm_layer, + ) + for _ in range(layers) + ] + ) + + if compile_mode is not None: + self.forward = torch.compile( + self.forward, mode=compile_mode, fullgraph=True + ) + if self.grad_checkpointing: + torch._dynamo.config.optimize_ddp = False + + def forward( + self, + x: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + for _, r in enumerate(self.resblocks): + if ( + self.grad_checkpointing + and not torch.jit.is_scripting() + and self.training + ): + x = checkpoint(r, x, None, None, attn_mask, use_reentrant=False) + else: + x = r( + x, + attn_mask=attn_mask, + ) + return x + + +def text_global_pool( + x: torch.Tensor, text: Optional[torch.Tensor] = None, pool_type: str = "argmax" +) -> Tuple[torch.Tensor, torch.Tensor]: + if pool_type == "first": + pooled, tokens = x[:, 0], x[:, 1:] + elif pool_type == "last": + pooled, tokens = x[:, -1], x[:, :-1] + elif pool_type == "argmax": + # take features from the eot embedding (eot_token is the highest number in each sequence) + assert text is not None + pooled, tokens = x[torch.arange(x.shape[0]), text.argmax(dim=-1)], x + else: + pooled = tokens = x + return pooled, tokens + + +class TextTransformer(nn.Module): + def __init__( + self, + context_length: int = 77, + vocab_size: int = 49408, + width: int = 512, + heads: int = 8, + layers: int = 12, + mlp_ratio: float = 4.0, + ls_init_value: Optional[float] = None, + output_dim: int = 512, + no_causal_mask: bool = False, + pool_type: str = "none", # no pooling + proj_bias: bool = False, + act_layer: Callable = nn.GELU, + norm_layer: Callable = nn.LayerNorm, + output_tokens: bool = False, + use_ln_post: bool = True, + compile_mode: Optional[str] = None, + use_act_checkpoint: bool = False, + ): + super().__init__() + assert pool_type in ("first", "last", "argmax", "none") + self.output_tokens = output_tokens + self.num_pos = self.context_length = context_length + self.vocab_size = vocab_size + self.width = width + self.output_dim = output_dim + self.heads = heads + self.pool_type = pool_type + + self.token_embedding = nn.Embedding(self.vocab_size, width) + self.positional_embedding = nn.Parameter(torch.empty(self.num_pos, width)) + self.transformer = Transformer( + width=width, + layers=layers, + heads=heads, + mlp_ratio=mlp_ratio, + ls_init_value=ls_init_value, + act_layer=act_layer, + norm_layer=norm_layer, + compile_mode=compile_mode, + use_act_checkpoint=use_act_checkpoint, + ) + self.ln_final = norm_layer(width) if use_ln_post else nn.Identity() + if no_causal_mask: + self.attn_mask = None + else: + self.register_buffer( + "attn_mask", self.build_causal_mask(), persistent=False + ) + if proj_bias: + self.text_projection = nn.Linear(width, output_dim) + else: + self.text_projection = nn.Parameter(torch.empty(width, output_dim)) + + def build_causal_mask(self) -> torch.Tensor: + # lazily create causal attention mask, with full attention between the tokens + # pytorch uses additive attention mask; fill with -inf + mask = torch.empty(self.num_pos, self.num_pos) + mask.fill_(float("-inf")) + mask.triu_(1) # zero out the lower diagonal + return mask + + def forward( + self, text: torch.Tensor + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + seq_len = text.shape[1] + x = self.token_embedding(text) # [batch_size, n_ctx, d_model] + + attn_mask = self.attn_mask + if attn_mask is not None: + attn_mask = attn_mask[:seq_len, :seq_len] + + x = x + self.positional_embedding[:seq_len] + x = self.transformer(x, attn_mask=attn_mask) + + x = self.ln_final(x) + pooled, tokens = text_global_pool(x, text, pool_type=self.pool_type) + if self.text_projection is not None: + if isinstance(self.text_projection, nn.Linear): + pooled = self.text_projection(pooled) + else: + pooled = pooled @ self.text_projection + if self.output_tokens: + return pooled, tokens + return pooled + + +class VETextEncoder(nn.Module): + def __init__( + self, + d_model: int, + tokenizer: Callable, + width: int = 1024, + heads: int = 16, + layers: int = 24, + context_length: int = 32, + vocab_size: int = 49408, + use_ln_post: bool = True, + compile_mode: Optional[str] = None, + use_act_checkpoint: bool = True, + ): + super().__init__() + self.context_length = context_length + self.use_ln_post = use_ln_post + self.tokenizer = tokenizer + + self.encoder = TextTransformer( + context_length=self.context_length, + vocab_size=vocab_size, + width=width, + heads=heads, + layers=layers, + # we want the tokens, not just the pooled output + output_tokens=True, + use_ln_post=use_ln_post, + compile_mode=compile_mode, + use_act_checkpoint=use_act_checkpoint, + ) + self.resizer = nn.Linear(self.encoder.width, d_model) + + def forward( + self, + text: Union[List[str], Tuple[torch.Tensor, torch.Tensor, dict]], + input_boxes: Optional[List] = None, + device: torch.device = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if isinstance(text[0], str): + # no use case for this + assert input_boxes is None or len(input_boxes) == 0, "not supported" + + # Encode the text + tokenized = self.tokenizer(text, context_length=self.context_length).to( + device + ) # [b, seq_len] + text_attention_mask = (tokenized != 0).bool() + + # manually embed the tokens + inputs_embeds = self.encoder.token_embedding( + tokenized + ) # [b, seq_len, d=1024] + _, text_memory = self.encoder(tokenized) # [b, seq_len, d=1024] + + assert text_memory.shape[1] == inputs_embeds.shape[1] + # Invert attention mask because its the opposite in pytorch transformer + text_attention_mask = text_attention_mask.ne(1) + # Transpose memory because pytorch's attention expects sequence first + text_memory = text_memory.transpose(0, 1) + # Resize the encoder hidden states to be of the same d_model as the decoder + text_memory_resized = self.resizer(text_memory) + else: + # The text is already encoded, use as is. + text_attention_mask, text_memory_resized, tokenized = text + inputs_embeds = tokenized["inputs_embeds"] + assert ( + input_boxes is None or len(input_boxes) == 0 + ), "Can't replace boxes in text if it's already encoded" + + # Note that the input_embeds are returned in pytorch's convention (sequence first) + return ( + text_attention_mask, + text_memory_resized, + inputs_embeds.transpose(0, 1), + ) diff --git a/src/nn/segearth_ov3/sam3/model/tokenizer_ve.py b/src/nn/segearth_ov3/sam3/model/tokenizer_ve.py new file mode 100644 index 0000000..ef42773 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/tokenizer_ve.py @@ -0,0 +1,253 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Text Tokenizer. + +Copied and lightly adapted from VE repo, which in turn copied +from open_clip and openAI CLIP. +""" + +import gzip +import html +import io +import os +import string +from functools import lru_cache +from typing import List, Optional, Union + +import ftfy +import regex as re +import torch +from iopath.common.file_io import g_pathmgr + + +# https://stackoverflow.com/q/62691279 +os.environ["TOKENIZERS_PARALLELISM"] = "false" +DEFAULT_CONTEXT_LENGTH = 77 + + +@lru_cache() +def bytes_to_unicode(): + """ + Returns list of utf-8 byte and a corresponding list of unicode strings. + The reversible bpe codes work on unicode strings. + This means you need a large # of unicode characters in your vocab if you want to avoid UNKs. + When you're at something like a 10B token dataset you end up needing around 5K for decent coverage. + This is a significant percentage of your normal, say, 32K bpe vocab. + To avoid that, we want lookup tables between utf-8 bytes and unicode strings. + And avoids mapping to whitespace/control characters the bpe code barfs on. + """ + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + """Return set of symbol pairs in a word. + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def _clean_canonicalize(x): + # basic, remove whitespace, remove punctuation, lower case + return canonicalize_text(basic_clean(x)) + + +def _clean_lower(x): + # basic, remove whitespace, lower case + return whitespace_clean(basic_clean(x)).lower() + + +def _clean_whitespace(x): + # basic, remove whitespace + return whitespace_clean(basic_clean(x)) + + +def get_clean_fn(type: str): + if type == "canonicalize": + return _clean_canonicalize + elif type == "lower": + return _clean_lower + elif type == "whitespace": + return _clean_whitespace + else: + assert False, f"Invalid clean function ({type})." + + +def canonicalize_text(text, *, keep_punctuation_exact_string=None): + """Returns canonicalized `text` (lowercase and punctuation removed). + From: https://github.com/google-research/big_vision/blob/53f18caf27a9419231bbf08d3388b07671616d3d/big_vision/evaluators/proj/image_text/prompt_engineering.py#L94 + Args: + text: string to be canonicalized. + keep_punctuation_exact_string: If provided, then this exact string kept. + For example providing '{}' will keep any occurrences of '{}' (but will + still remove '{' and '}' that appear separately). + """ + text = text.replace("_", " ") + if keep_punctuation_exact_string: + text = keep_punctuation_exact_string.join( + part.translate(str.maketrans("", "", string.punctuation)) + for part in text.split(keep_punctuation_exact_string) + ) + else: + text = text.translate(str.maketrans("", "", string.punctuation)) + text = text.lower() + text = re.sub(r"\s+", " ", text) + return text.strip() + + +class SimpleTokenizer(object): + def __init__( + self, + bpe_path: Union[str, os.PathLike], + additional_special_tokens: Optional[List[str]] = None, + context_length: Optional[int] = DEFAULT_CONTEXT_LENGTH, + clean: str = "lower", + ): + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + with g_pathmgr.open(bpe_path, "rb") as fh: + bpe_bytes = io.BytesIO(fh.read()) + merges = gzip.open(bpe_bytes).read().decode("utf-8").split("\n") + # merges = gzip.open(bpe_path).read().decode("utf-8").split("\n") + merges = merges[1 : 49152 - 256 - 2 + 1] + merges = [tuple(merge.split()) for merge in merges] + vocab = list(bytes_to_unicode().values()) + vocab = vocab + [v + "" for v in vocab] + for merge in merges: + vocab.append("".join(merge)) + special_tokens = ["", ""] + if additional_special_tokens: + special_tokens += additional_special_tokens + vocab.extend(special_tokens) + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.decoder = {v: k for k, v in self.encoder.items()} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {t: t for t in special_tokens} + special = "|".join(special_tokens) + self.pat = re.compile( + special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", + re.IGNORECASE, + ) + self.vocab_size = len(self.encoder) + self.all_special_ids = [self.encoder[t] for t in special_tokens] + self.sot_token_id = self.all_special_ids[0] + self.eot_token_id = self.all_special_ids[1] + self.context_length = context_length + self.clean_fn = get_clean_fn(clean) + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + (token[-1] + "",) + pairs = get_pairs(word) + if not pairs: + return token + "" + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + new_word.extend(word[i:j]) + i = j + except: + new_word.extend(word[i:]) + break + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = " ".join(word) + self.cache[token] = word + return word + + def encode(self, text): + bpe_tokens = [] + text = self.clean_fn(text) + for token in re.findall(self.pat, text): + token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) + bpe_tokens.extend( + self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" ") + ) + return bpe_tokens + + def decode(self, tokens): + text = "".join([self.decoder[token] for token in tokens]) + text = ( + bytearray([self.byte_decoder[c] for c in text]) + .decode("utf-8", errors="replace") + .replace("", " ") + ) + return text + + def __call__( + self, texts: Union[str, List[str]], context_length: Optional[int] = None + ) -> torch.LongTensor: + """Returns the tokenized representation of given input string(s) + Parameters + ---------- + texts : Union[str, List[str]] + An input string or a list of input strings to tokenize + context_length : int + The context length to use; all CLIP models use 77 as the context length + Returns + ------- + A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length] + """ + if isinstance(texts, str): + texts = [texts] + context_length = context_length or self.context_length + assert context_length, "Please set a valid context length" + all_tokens = [ + [self.sot_token_id] + self.encode(text) + [self.eot_token_id] + for text in texts + ] + result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + tokens = tokens[:context_length] # Truncate + tokens[-1] = self.eot_token_id + result[i, : len(tokens)] = torch.tensor(tokens) + return result diff --git a/src/nn/segearth_ov3/sam3/model/utils/__init__.py b/src/nn/segearth_ov3/sam3/model/utils/__init__.py new file mode 100644 index 0000000..47d9858 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/src/nn/segearth_ov3/sam3/model/utils/misc.py b/src/nn/segearth_ov3/sam3/model/utils/misc.py new file mode 100644 index 0000000..4f07206 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/utils/misc.py @@ -0,0 +1,77 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from collections import defaultdict +from dataclasses import fields, is_dataclass +from typing import Any, Mapping, Protocol, runtime_checkable + +import torch + + +def _is_named_tuple(x) -> bool: + return isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields") + + +@runtime_checkable +class _CopyableData(Protocol): + def to(self, device: torch.device, *args: Any, **kwargs: Any): + """Copy data to the specified device""" + ... + + +def copy_data_to_device(data, device: torch.device, *args: Any, **kwargs: Any): + """Function that recursively copies data to a torch.device. + + Args: + data: The data to copy to device + device: The device to which the data should be copied + args: positional arguments that will be passed to the `to` call + kwargs: keyword arguments that will be passed to the `to` call + + Returns: + The data on the correct device + """ + + if _is_named_tuple(data): + return type(data)( + **copy_data_to_device(data._asdict(), device, *args, **kwargs) + ) + elif isinstance(data, (list, tuple)): + return type(data)(copy_data_to_device(e, device, *args, **kwargs) for e in data) + elif isinstance(data, defaultdict): + return type(data)( + data.default_factory, + { + k: copy_data_to_device(v, device, *args, **kwargs) + for k, v in data.items() + }, + ) + elif isinstance(data, Mapping): + return type(data)( + { + k: copy_data_to_device(v, device, *args, **kwargs) + for k, v in data.items() + } + ) + elif is_dataclass(data) and not isinstance(data, type): + new_data_class = type(data)( + **{ + field.name: copy_data_to_device( + getattr(data, field.name), device, *args, **kwargs + ) + for field in fields(data) + if field.init + } + ) + for field in fields(data): + if not field.init: + setattr( + new_data_class, + field.name, + copy_data_to_device( + getattr(data, field.name), device, *args, **kwargs + ), + ) + return new_data_class + elif isinstance(data, _CopyableData): + return data.to(device, *args, **kwargs) + return data diff --git a/src/nn/segearth_ov3/sam3/model/utils/sam1_utils.py b/src/nn/segearth_ov3/sam3/model/utils/sam1_utils.py new file mode 100644 index 0000000..18f0d04 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/utils/sam1_utils.py @@ -0,0 +1,119 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import warnings + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.transforms import Normalize, Resize, ToTensor + + +# Adapted from https://github.com/facebookresearch/sam2/blob/main/sam2/utils/transforms.py +class SAM2Transforms(nn.Module): + def __init__( + self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0 + ): + """ + Transforms for SAM2. + """ + super().__init__() + self.resolution = resolution + self.mask_threshold = mask_threshold + self.max_hole_area = max_hole_area + self.max_sprinkle_area = max_sprinkle_area + self.mean = [0.5, 0.5, 0.5] + self.std = [0.5, 0.5, 0.5] + self.to_tensor = ToTensor() + self.transforms = torch.jit.script( + nn.Sequential( + Resize((self.resolution, self.resolution)), + Normalize(self.mean, self.std), + ) + ) + + def __call__(self, x): + x = self.to_tensor(x) + return self.transforms(x) + + def forward_batch(self, img_list): + img_batch = [self.transforms(self.to_tensor(img)) for img in img_list] + img_batch = torch.stack(img_batch, dim=0) + return img_batch + + def transform_coords( + self, coords: torch.Tensor, normalize=False, orig_hw=None + ) -> torch.Tensor: + """ + Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates, + If the coords are in absolute image coordinates, normalize should be set to True and original image size is required. + + Returns + Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model. + """ + if normalize: + assert orig_hw is not None + h, w = orig_hw + coords = coords.clone() + coords[..., 0] = coords[..., 0] / w + coords[..., 1] = coords[..., 1] / h + + coords = coords * self.resolution # unnormalize coords + return coords + + def transform_boxes( + self, boxes: torch.Tensor, normalize=False, orig_hw=None + ) -> torch.Tensor: + """ + Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates, + if the coords are in absolute image coordinates, normalize should be set to True and original image size is required. + """ + boxes = self.transform_coords(boxes.reshape(-1, 2, 2), normalize, orig_hw) + return boxes + + def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor: + """ + Perform PostProcessing on output masks. + """ + masks = masks.float() + input_masks = masks + mask_flat = masks.flatten(0, 1).unsqueeze(1) # flatten as 1-channel image + try: + from sam3.perflib.connected_components import connected_components + + if self.max_hole_area > 0: + # Holes are those connected components in background with area <= self.fill_hole_area + # (background regions are those with mask scores <= self.mask_threshold) + labels, areas = connected_components( + (mask_flat <= self.mask_threshold).to(torch.uint8) + ) + is_hole = (labels > 0) & (areas <= self.max_hole_area) + is_hole = is_hole.reshape_as(masks) + # We fill holes with a small positive mask score (10.0) to change them to foreground. + masks = torch.where(is_hole, self.mask_threshold + 10.0, masks) + + if self.max_sprinkle_area > 0: + labels, areas = connected_components( + (mask_flat > self.mask_threshold).to(torch.uint8) + ) + is_hole = (labels > 0) & (areas <= self.max_sprinkle_area) + is_hole = is_hole.reshape_as(masks) + # We fill holes with negative mask score (-10.0) to change them to background. + masks = torch.where(is_hole, self.mask_threshold - 10.0, masks) + except Exception as e: + # Skip the post-processing step if the CUDA kernel fails + warnings.warn( + f"{e}\n\nSkipping the post-processing step due to the error above. You can " + "still use SAM 3 and it's OK to ignore the error above, although some post-processing " + "functionality may be limited (which doesn't affect the results in most cases; see " + "https://github.com/facebookresearch/sam3/blob/main/INSTALL.md).", + category=UserWarning, + stacklevel=2, + ) + masks = input_masks + + masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False) + return masks diff --git a/src/nn/segearth_ov3/sam3/model/utils/sam2_utils.py b/src/nn/segearth_ov3/sam3/model/utils/sam2_utils.py new file mode 100644 index 0000000..cc3d12e --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/utils/sam2_utils.py @@ -0,0 +1,233 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import os +from threading import Thread + +import numpy as np +import torch +from PIL import Image +from tqdm import tqdm + + +def _load_img_as_tensor(img_path, image_size): + img_pil = Image.open(img_path) + img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size))) + if img_np.dtype == np.uint8: # np.uint8 is expected for JPEG images + img_np = img_np / 255.0 + else: + raise RuntimeError(f"Unknown image dtype: {img_np.dtype} on {img_path}") + img = torch.from_numpy(img_np).permute(2, 0, 1) + video_width, video_height = img_pil.size # the original video size + return img, video_height, video_width + + +class AsyncVideoFrameLoader: + """ + A list of video frames to be load asynchronously without blocking session start. + """ + + def __init__( + self, + img_paths, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + compute_device, + ): + self.img_paths = img_paths + self.image_size = image_size + self.offload_video_to_cpu = offload_video_to_cpu + self.img_mean = img_mean + self.img_std = img_std + # items in `self.images` will be loaded asynchronously + self.images = [None] * len(img_paths) + # catch and raise any exceptions in the async loading thread + self.exception = None + # video_height and video_width be filled when loading the first image + self.video_height = None + self.video_width = None + self.compute_device = compute_device + + # load the first frame to fill video_height and video_width and also + # to cache it (since it's most likely where the user will click) + self.__getitem__(0) + + # load the rest of frames asynchronously without blocking the session start + def _load_frames(): + try: + for n in tqdm(range(len(self.images)), desc="frame loading (JPEG)"): + self.__getitem__(n) + except Exception as e: + self.exception = e + + self.thread = Thread(target=_load_frames, daemon=True) + self.thread.start() + + def __getitem__(self, index): + if self.exception is not None: + raise RuntimeError("Failure in frame loading thread") from self.exception + + img = self.images[index] + if img is not None: + return img + + img, video_height, video_width = _load_img_as_tensor( + self.img_paths[index], self.image_size + ) + self.video_height = video_height + self.video_width = video_width + # normalize by mean and std + img -= self.img_mean + img /= self.img_std + if not self.offload_video_to_cpu: + img = img.to(self.compute_device, non_blocking=True) + self.images[index] = img + return img + + def __len__(self): + return len(self.images) + + +def load_video_frames( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.485, 0.456, 0.406), + img_std=(0.229, 0.224, 0.225), + async_loading_frames=False, + compute_device=torch.device("cuda"), +): + """ + Load the video frames from video_path. The frames are resized to image_size as in + the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo. + """ + is_bytes = isinstance(video_path, bytes) + is_str = isinstance(video_path, str) + is_mp4_path = is_str and os.path.splitext(video_path)[-1] in [".mp4", ".MP4"] + if is_bytes or is_mp4_path: + return load_video_frames_from_video_file( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + compute_device=compute_device, + ) + elif is_str and os.path.isdir(video_path): + return load_video_frames_from_jpg_images( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + async_loading_frames=async_loading_frames, + compute_device=compute_device, + ) + else: + raise NotImplementedError( + "Only MP4 video and JPEG folder are supported at this moment" + ) + + +def load_video_frames_from_jpg_images( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.485, 0.456, 0.406), + img_std=(0.229, 0.224, 0.225), + async_loading_frames=False, + compute_device=torch.device("cuda"), +): + """ + Load the video frames from a directory of JPEG files (".jpg" format). + + The frames are resized to image_size x image_size and are loaded to GPU if + `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`. + + You can load a frame asynchronously by setting `async_loading_frames` to `True`. + """ + if isinstance(video_path, str) and os.path.isdir(video_path): + jpg_folder = video_path + else: + raise NotImplementedError( + "Only JPEG frames are supported at this moment. For video files, you may use " + "ffmpeg (https://ffmpeg.org/) to extract frames into a folder of JPEG files, such as \n" + "```\n" + "ffmpeg -i .mp4 -q:v 2 -start_number 0 /'%05d.jpg'\n" + "```\n" + "where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks " + "ffmpeg to start the JPEG file from 00000.jpg." + ) + + frame_names = [ + p + for p in os.listdir(jpg_folder) + if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"] + ] + frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) + num_frames = len(frame_names) + if num_frames == 0: + raise RuntimeError(f"no images found in {jpg_folder}") + img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names] + img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None] + + if async_loading_frames: + lazy_images = AsyncVideoFrameLoader( + img_paths, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + compute_device, + ) + return lazy_images, lazy_images.video_height, lazy_images.video_width + + images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32) + for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")): + images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size) + if not offload_video_to_cpu: + images = images.to(compute_device) + img_mean = img_mean.to(compute_device) + img_std = img_std.to(compute_device) + # normalize by mean and std + images -= img_mean + images /= img_std + return images, video_height, video_width + + +def load_video_frames_from_video_file( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.485, 0.456, 0.406), + img_std=(0.229, 0.224, 0.225), + compute_device=torch.device("cuda"), +): + """Load the video frames from a video file.""" + import decord + + img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None] + # Get the original video height and width + decord.bridge.set_bridge("torch") + video_height, video_width, _ = decord.VideoReader(video_path).next().shape + # Iterate over all frames in the video + images = [] + for frame in decord.VideoReader(video_path, width=image_size, height=image_size): + images.append(frame.permute(2, 0, 1)) + + images = torch.stack(images, dim=0).float() / 255.0 + if not offload_video_to_cpu: + images = images.to(compute_device) + img_mean = img_mean.to(compute_device) + img_std = img_std.to(compute_device) + # normalize by mean and std + images -= img_mean + images /= img_std + return images, video_height, video_width diff --git a/src/nn/segearth_ov3/sam3/model/vitdet.py b/src/nn/segearth_ov3/sam3/model/vitdet.py new file mode 100644 index 0000000..aa56664 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/vitdet.py @@ -0,0 +1,879 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +ViTDet backbone adapted from Detectron2. +This module implements Vision Transformer (ViT) backbone for object detection. + +Rope embedding code adopted from: +1. https://github.com/meta-llama/codellama/blob/main/llama/model.py +2. https://github.com/naver-ai/rope-vit +3. https://github.com/lucidrains/rotary-embedding-torch +""" + +import math +from functools import partial +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint + +try: + from timm.layers import DropPath, Mlp, trunc_normal_ +except ModuleNotFoundError: + # compatibility for older timm versions + from timm.models.layers import DropPath, Mlp, trunc_normal_ +from torch import Tensor + +from .model_misc import LayerScale + + +def init_t_xy( + end_x: int, end_y: int, scale: float = 1.0, offset: int = 0 +) -> Tuple[torch.Tensor, torch.Tensor]: + t = torch.arange(end_x * end_y, dtype=torch.float32) + t_x = (t % end_x).float() + t_y = torch.div(t, end_x, rounding_mode="floor").float() + return t_x * scale + offset, t_y * scale + offset + + +def compute_axial_cis( + dim: int, + end_x: int, + end_y: int, + theta: float = 10000.0, + scale_pos: float = 1.0, + offset: int = 0, +) -> torch.Tensor: + freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) + freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) + + t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset) + freqs_x = torch.outer(t_x, freqs_x) + freqs_y = torch.outer(t_y, freqs_y) + freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) + freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) + return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) + + +def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + ndim = x.ndim + assert 0 <= 1 < ndim + assert freqs_cis.shape == (x.shape[-2], x.shape[-1]) + shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def apply_rotary_enc( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: torch.Tensor, + repeat_freqs_k: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) + xk_ = ( + torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) + if xk.shape[-2] != 0 + else None + ) + freqs_cis = reshape_for_broadcast(freqs_cis, xq_) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) + if xk_ is None: + # no keys to rotate, due to dropout + return xq_out.type_as(xq).to(xq.device), xk + # repeat freqs along seq_len dim to match k seq_len + if repeat_freqs_k: + r = xk_.shape[-2] // xq_.shape[-2] + freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) + return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device) + + +def window_partition(x: Tensor, window_size: int) -> Tuple[Tensor, Tuple[int, int]]: + """ + Partition into non-overlapping windows with padding if needed. + Args: + x (tensor): input tokens with [B, H, W, C]. + window_size (int): window size. + Returns: + windows: windows after partition with [B * num_windows, window_size, window_size, C]. + (Hp, Wp): padded height and width before partition + """ + B, H, W, C = x.shape + + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + if pad_h > 0 or pad_w > 0: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + Hp, Wp = H + pad_h, W + pad_w + + x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C) + return windows, (Hp, Wp) + + +def window_unpartition( + windows: Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int] +) -> Tensor: + """ + Window unpartition into original sequences and removing padding. + Args: + x (tensor): input tokens with [B * num_windows, window_size, window_size, C]. + window_size (int): window size. + pad_hw (Tuple): padded height and width (Hp, Wp). + hw (Tuple): original height and width (H, W) before padding. + Returns: + x: unpartitioned sequences with [B, H, W, C]. + """ + Hp, Wp = pad_hw + H, W = hw + B = windows.shape[0] // (Hp * Wp // window_size // window_size) + x = windows.reshape( + B, Hp // window_size, Wp // window_size, window_size, window_size, -1 + ) + x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1) + + if Hp > H or Wp > W: + x = x[:, :H, :W, :] + return x + + +def get_rel_pos(q_size: int, k_size: int, rel_pos: Tensor) -> Tensor: + """ + Get relative positional embeddings according to the relative positions of + query and key sizes. + Args: + q_size (int): size of query q. + k_size (int): size of key k. + rel_pos (Tensor): relative position embeddings (L, C). + Returns: + Extracted positional embeddings according to relative positions. + """ + max_rel_dist = int(2 * max(q_size, k_size) - 1) + # Interpolate rel pos if needed. + if rel_pos.shape[0] != max_rel_dist: + # Interpolate rel pos. + rel_pos_resized = F.interpolate( + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), + size=max_rel_dist, + mode="linear", + align_corners=False, + ) + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) + else: + rel_pos_resized = rel_pos + + # Scale the coords with short length if shapes for q and k are different. + q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) + k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + + return rel_pos_resized[relative_coords.long()] + + +def get_abs_pos( + abs_pos: Tensor, + has_cls_token: bool, + hw: Tuple[int, int], + retain_cls_token: bool = False, + tiling: bool = False, +) -> Tensor: + """ + Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token + dimension for the original embeddings. + Args: + abs_pos (Tensor): absolute positional embeddings with (1, num_position, C). + has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token. + hw (Tuple): size of input image tokens. + retain_cls_token: whether to retain the cls_token + tiling: whether to tile the embeddings, *instead* of interpolation (a la abs_win) + Returns: + Absolute positional embeddings after processing with shape (1, H, W, C), + if retain_cls_token is False, otherwise (1, 1+H*W, C) + """ + if retain_cls_token: + assert has_cls_token + + h, w = hw + if has_cls_token: + cls_pos = abs_pos[:, :1] + abs_pos = abs_pos[:, 1:] + + xy_num = abs_pos.shape[1] + size = int(math.sqrt(xy_num)) + assert size * size == xy_num + + if size != h or size != w: + new_abs_pos = abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2) + if tiling: + new_abs_pos = new_abs_pos.tile( + [1, 1] + [x // y + 1 for x, y in zip((h, w), new_abs_pos.shape[2:])] + )[:, :, :h, :w] + else: + new_abs_pos = F.interpolate( + new_abs_pos, + size=(h, w), + mode="bicubic", + align_corners=False, + ) + + if not retain_cls_token: + return new_abs_pos.permute(0, 2, 3, 1) + else: + # add cls_token back, flatten spatial dims + assert has_cls_token + return torch.cat( + [cls_pos, new_abs_pos.permute(0, 2, 3, 1).reshape(1, h * w, -1)], + dim=1, + ) + + else: + if not retain_cls_token: + return abs_pos.reshape(1, h, w, -1) + else: + assert has_cls_token + return torch.cat([cls_pos, abs_pos], dim=1) + + +def concat_rel_pos( + q: Tensor, + k: Tensor, + q_hw: Tuple[int, int], + k_hw: Tuple[int, int], + rel_pos_h: Tensor, + rel_pos_w: Tensor, + rescale: bool = False, + relative_coords: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor]: + """ + Concatenate rel pos coeffs to the q & k tensors, so that qk^T is now + effectively including rel pos biases. + Args: + q (Tensor): q tensor with shape (B, L_q, C). + k (Tensor): k tensor with shape (B, L_k, C). + q_hw, k_hw: These are spatial size of q & k tensors. + rel_pos_h, rel_pos_w: These are relative pos embeddings/params of height, width. + rescale (bool): whether to rescale. e.g. for use when using sdpa, pytorch will + scale by the wrong factor due to the concat. + Returns: + q, k: But, padded so that qk^T accounts for rel pos biases + """ + q_h, q_w = q_hw + k_h, k_w = k_hw + + assert (q_h == q_w) and (k_h == k_w), "only square inputs supported" + + if relative_coords is not None: + Rh = rel_pos_h[relative_coords] + Rw = rel_pos_w[relative_coords] + else: + Rh = get_rel_pos(q_h, k_h, rel_pos_h) + Rw = get_rel_pos(q_w, k_w, rel_pos_w) + + B, _, dim = q.shape + r_q = q.reshape(B, q_h, q_w, dim) + + old_scale = dim**0.5 + new_scale = (dim + k_h + k_w) ** 0.5 if rescale else old_scale # for sdpa + # attn will be divided by new_scale, but we want to divide q by old_scale + scale_ratio = new_scale / old_scale + + rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) * new_scale # (B, q_h, q_w, k_h) + rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) * new_scale # (B, q_h, q_w, k_w) + + eye_h = torch.eye(k_h, dtype=q.dtype, device=q.device) + eye_w = torch.eye(k_w, dtype=q.dtype, device=q.device) + + eye_h = eye_h.view(1, k_h, 1, k_h).expand([B, k_h, k_w, k_h]) + eye_w = eye_w.view(1, 1, k_w, k_w).expand([B, k_h, k_w, k_w]) + + q = torch.cat([r_q * scale_ratio, rel_h, rel_w], dim=-1).view(B, q_h * q_w, -1) + k = torch.cat([k.view(B, k_h, k_w, -1), eye_h, eye_w], dim=-1).view( + B, k_h * k_w, -1 + ) + + return q, k + + +class PatchEmbed(nn.Module): + """ + Image to Patch Embedding. + """ + + def __init__( + self, + kernel_size: Tuple[int, int] = (16, 16), + stride: Tuple[int, int] = (16, 16), + padding: Tuple[int, int] = (0, 0), + in_chans: int = 3, + embed_dim: int = 768, + bias: bool = True, + ): + """ + Args: + kernel_size (Tuple): kernel size of the projection layer. + stride (Tuple): stride of the projection layer. + padding (Tuple): padding size of the projection layer. + in_chans (int): Number of input image channels. + embed_dim (int): embed_dim (int): Patch embedding dimension. + """ + super().__init__() + + self.proj = nn.Conv2d( + in_chans, + embed_dim, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=bias, + ) + + def forward(self, x: Tensor) -> Tensor: + x = self.proj(x) + # B C H W -> B H W C + x = x.permute(0, 2, 3, 1) + return x + + +class Attention(nn.Module): + """Multi-head Attention block with relative position embeddings and 2d-rope.""" + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + input_size: Optional[Tuple[int, int]] = None, + cls_token: bool = False, + use_rope: bool = False, + rope_theta: float = 10000.0, + rope_pt_size: Optional[Tuple[int, int]] = None, + rope_interp: bool = False, + ): + """ + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + qkv_bias (bool: If True, add a learnable bias to query, key, value. + rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + input_size (int or None): Input resolution for calculating the relative positional + parameter size or rope size. + attn_type: Type of attention operation, e.g. "vanilla", "vanilla-xformer". + cls_token: whether a cls_token is present. + use_rope: whether to use rope 2d (indep of use_rel_pos, as it can be used together) + rope_theta: control frequencies of rope + rope_pt_size: size of rope in previous stage of training, needed for interpolation or tiling + rope_interp: whether to interpolate (or extrapolate) rope to match input size + """ + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.cls_token = cls_token + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.proj = nn.Linear(dim, dim) + + # rel_pos embeddings and rope + self.use_rel_pos = use_rel_pos + self.input_size = input_size + + self.use_rope = use_rope + self.rope_theta = rope_theta + self.rope_pt_size = rope_pt_size + self.rope_interp = rope_interp + + # init rel_pos embeddings and rope + self._setup_rel_pos(rel_pos_zero_init) + self._setup_rope_freqs() + + def _setup_rel_pos(self, rel_pos_zero_init: bool = True) -> None: + if not self.use_rel_pos: + self.rel_pos_h = None + self.rel_pos_w = None + return + + assert self.input_size is not None + assert self.cls_token is False, "not supported" + # initialize relative positional embeddings + self.rel_pos_h = nn.Parameter( + torch.zeros(2 * self.input_size[0] - 1, self.head_dim) + ) + self.rel_pos_w = nn.Parameter( + torch.zeros(2 * self.input_size[1] - 1, self.head_dim) + ) + + if not rel_pos_zero_init: + trunc_normal_(self.rel_pos_h, std=0.02) + trunc_normal_(self.rel_pos_w, std=0.02) + + # Precompute the relative coords + H, W = self.input_size + q_coords = torch.arange(H)[:, None] + k_coords = torch.arange(W)[None, :] + relative_coords = (q_coords - k_coords) + (H - 1) + self.register_buffer("relative_coords", relative_coords.long()) + + def _setup_rope_freqs(self) -> None: + if not self.use_rope: + self.freqs_cis = None + return + + assert self.input_size is not None + # determine rope input size + if self.rope_pt_size is None: + self.rope_pt_size = self.input_size + + # initialize 2d rope freqs + self.compute_cis = partial( + compute_axial_cis, + dim=self.head_dim, + theta=self.rope_theta, + ) + + # interpolate rope + scale_pos = 1.0 + if self.rope_interp: + scale_pos = self.rope_pt_size[0] / self.input_size[0] + # get scaled freqs_cis + freqs_cis = self.compute_cis( + end_x=self.input_size[0], + end_y=self.input_size[1], + scale_pos=scale_pos, + ) + if self.cls_token: + t = torch.zeros( + self.head_dim // 2, + dtype=torch.float32, + device=freqs_cis.device, + ) + cls_freqs_cis = torch.polar(torch.ones_like(t), t)[None, :] + freqs_cis = torch.cat([cls_freqs_cis, freqs_cis], dim=0) + + self.register_buffer("freqs_cis", freqs_cis) + + def _apply_rope(self, q, k) -> Tuple[Tensor, Tensor]: + if not self.use_rope: + return q, k + + assert self.freqs_cis is not None + return apply_rotary_enc(q, k, freqs_cis=self.freqs_cis) + + def forward(self, x: Tensor) -> Tensor: + s = 1 if self.cls_token else 0 # used to exclude cls_token + if x.ndim == 4: + B, H, W, _ = x.shape + assert s == 0 # no cls_token + L = H * W + ndim = 4 + else: + assert x.ndim == 3 + B, L, _ = x.shape + ndim = 3 + H = W = math.sqrt(L - s) + + # qkv with shape (3, B, nHead, L, C) + qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, -1) + # q, k, v with shape (B, nHead, L, C) + q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + + # handle rope and rel pos embeddings + q, k = self._apply_rope(q, k) + if self.use_rel_pos: + q, k = concat_rel_pos( + q.flatten(0, 1), + k.flatten(0, 1), + (H, W), + x.shape[1:3], + self.rel_pos_h, + self.rel_pos_w, + rescale=True, + relative_coords=self.relative_coords, + ) + + # sdpa expects [B, nheads, H*W, C] so we transpose back + q = q.reshape(B, self.num_heads, H * W, -1) + k = k.reshape(B, self.num_heads, H * W, -1) + + x = F.scaled_dot_product_attention(q, k, v) + + if ndim == 4: + x = ( + x.view(B, self.num_heads, H, W, -1) + .permute(0, 2, 3, 1, 4) + .reshape(B, H, W, -1) + ) + else: + x = x.view(B, self.num_heads, L, -1).permute(0, 2, 1, 3).reshape(B, L, -1) + + x = self.proj(x) + + return x + + +class Block(nn.Module): + """Transformer blocks with support of window attention""" + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + drop_path: float = 0.0, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + act_layer: Callable[..., nn.Module] = nn.GELU, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + window_size: int = 0, + input_size: Optional[Tuple[int, int]] = None, + use_rope: bool = False, + rope_pt_size: Optional[Tuple[int, int]] = None, + rope_tiled: bool = False, + rope_interp: bool = False, + use_ve_rope: bool = False, + cls_token: bool = False, + dropout: float = 0.0, + init_values: Optional[float] = None, + ): + """ + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + drop_path (float): Stochastic depth rate. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. If it equals 0, then not + use window attention. + input_size (int or None): Input resolution for calculating the relative positional + parameter size. + dropout (float): Dropout rate. + cls_token: whether a cls_token is present. + use_rope: whether to use rope 2d (indep of use_rel_pos, as it can be used together) + rope_pt_size: size of rope in previous stage of training, needed for interpolation or tiling + rope_interp: whether to interpolate (or extrapolate) rope to match target input size, + expected to specify source size as rope_pt_size. + """ + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + use_rel_pos=use_rel_pos, + rel_pos_zero_init=rel_pos_zero_init, + input_size=input_size if window_size == 0 else (window_size, window_size), + use_rope=use_rope, + rope_pt_size=rope_pt_size, + rope_interp=rope_interp, + cls_token=cls_token, + ) + self.ls1 = ( + LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + self.mlp = Mlp( + in_features=dim, + hidden_features=int(dim * mlp_ratio), + act_layer=act_layer, + drop=(dropout, 0.0), + ) + self.ls2 = ( + LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + ) + self.dropout = nn.Dropout(dropout) + self.window_size = window_size + + def forward(self, x: Tensor) -> Tensor: + shortcut = x + x = self.norm1(x) + # Window partition + if self.window_size > 0: + H, W = x.shape[1], x.shape[2] + x, pad_hw = window_partition(x, self.window_size) + + x = self.ls1(self.attn(x)) + # Reverse window partition + if self.window_size > 0: + x = window_unpartition(x, self.window_size, pad_hw, (H, W)) + + x = shortcut + self.dropout(self.drop_path(x)) + x = x + self.dropout(self.drop_path(self.ls2(self.mlp(self.norm2(x))))) + + return x + + +class ViT(nn.Module): + """ + This module implements Vision Transformer (ViT) backbone in :paper:`vitdet`. + "Exploring Plain Vision Transformer Backbones for Object Detection", + https://arxiv.org/abs/2203.16527 + """ + + def __init__( + self, + img_size: int = 1024, + patch_size: int = 16, + in_chans: int = 3, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + drop_path_rate: float = 0.0, + norm_layer: Union[Callable[..., nn.Module], str] = "LayerNorm", + act_layer: Callable[..., nn.Module] = nn.GELU, + use_abs_pos: bool = True, + tile_abs_pos: bool = True, + rel_pos_blocks: Union[Tuple[int, ...], bool] = (2, 5, 8, 11), + rel_pos_zero_init: bool = True, + window_size: int = 14, + global_att_blocks: Tuple[int, ...] = (2, 5, 8, 11), + use_rope: bool = False, + rope_pt_size: Optional[int] = None, + use_interp_rope: bool = False, + pretrain_img_size: int = 224, + pretrain_use_cls_token: bool = True, + retain_cls_token: bool = True, + dropout: float = 0.0, + return_interm_layers: bool = False, + init_values: Optional[float] = None, # for layerscale + ln_pre: bool = False, + ln_post: bool = False, + bias_patch_embed: bool = True, + compile_mode: Optional[str] = None, + use_act_checkpoint: bool = True, + ): + """ + Args: + img_size (int): Input image size. Only relevant for rel pos or rope. + patch_size (int): Patch size. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + depth (int): Depth of ViT. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + drop_path_rate (float): Stochastic depth rate. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_abs_pos (bool): If True, use absolute positional embeddings. + tile_abs_pos (bool): If True, tile absolute positional embeddings instead of interpolation. + rel_pos_blocks (list): Blocks which have rel pos embeddings. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. + global_att_blocks (list): Indexes for blocks using global attention (other blocks use window attention). + use_rope (bool): whether to use rope 2d (indep of rel_pos_blocks, as it can be used together). + rope_pt_size (int): size of rope in previous stage of training, needed for interpolation or tiling. + use_interp_rope: whether to interpolate (or extrapolate) rope to match target input size, + expected to specify source size as rope_pt_size. + use_act_checkpoint (bool): If True, use activation checkpointing. + pretrain_img_size (int): input image size for pretraining models. + pretrain_use_cls_token (bool): If True, pretraining models use class token. + retain_cls_token: whether cls_token should be retained. + dropout (float): Dropout rate. Applied in residual blocks of attn, mlp and inside the mlp. + + return_interm_layers (bool): Whether to return intermediate layers (all global attention blocks). + init_values: layer scale init, None for no layer scale. + + ln_pre (bool): If True, apply layer norm before transformer blocks. + ln_post (bool): If True, apply layer norm after transformer blocks. + bias_patch_embed (bool): bias in conv for patch embed? + compile_mode (str): mode to compile the forward + """ + super().__init__() + self.pretrain_use_cls_token = pretrain_use_cls_token + + window_block_indexes = [i for i in range(depth) if i not in global_att_blocks] + self.full_attn_ids = list(global_att_blocks) + self.rel_pos_blocks = [False] * depth + if isinstance(rel_pos_blocks, bool) and rel_pos_blocks: + self.rel_pos_blocks = [True] * depth + else: + for i in rel_pos_blocks: + self.rel_pos_blocks[i] = True + + self.retain_cls_token = retain_cls_token + if self.retain_cls_token: + assert pretrain_use_cls_token + assert ( + len(window_block_indexes) == 0 + ), "windowing not supported with cls token" + + assert sum(self.rel_pos_blocks) == 0, "rel pos not supported with cls token" + + scale = embed_dim**-0.5 + self.class_embedding = nn.Parameter(scale * torch.randn(1, 1, embed_dim)) + + if isinstance(norm_layer, str): + norm_layer = partial(getattr(nn, norm_layer), eps=1e-5) + + self.patch_embed = PatchEmbed( + kernel_size=(patch_size, patch_size), + stride=(patch_size, patch_size), + in_chans=in_chans, + embed_dim=embed_dim, + bias=bias_patch_embed, + ) + + # Handle absolute positional embedding + self.tile_abs_pos = tile_abs_pos + self.use_abs_pos = use_abs_pos + if self.tile_abs_pos: + assert self.use_abs_pos + + if self.use_abs_pos: + # Initialize absolute positional embedding with pretrain image size. + num_patches = (pretrain_img_size // patch_size) * ( + pretrain_img_size // patch_size + ) + num_positions = (num_patches + 1) if pretrain_use_cls_token else num_patches + self.pos_embed = nn.Parameter(torch.zeros(1, num_positions, embed_dim)) + else: + self.pos_embed = None + + # stochastic depth decay rule + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] + + self.blocks = nn.ModuleList() + cur_stage = 1 + for i in range(depth): + block = Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + use_rel_pos=self.rel_pos_blocks[i], + rel_pos_zero_init=rel_pos_zero_init, + window_size=window_size if i in window_block_indexes else 0, + input_size=(img_size // patch_size, img_size // patch_size), + use_rope=use_rope, + rope_pt_size=( + (window_size, window_size) + if rope_pt_size is None + else (rope_pt_size, rope_pt_size) + ), + rope_interp=use_interp_rope, + cls_token=self.retain_cls_token, + dropout=dropout, + init_values=init_values, + ) + + if i not in window_block_indexes: + cur_stage += 1 + + self.use_act_checkpoint = use_act_checkpoint + + self.blocks.append(block) + + self.return_interm_layers = return_interm_layers + self.channel_list = ( + [embed_dim] * len(self.full_attn_ids) + if return_interm_layers + else [embed_dim] + ) + + if self.pos_embed is not None: + trunc_normal_(self.pos_embed, std=0.02) + + self.ln_pre = norm_layer(embed_dim) if ln_pre else nn.Identity() + self.ln_post = norm_layer(embed_dim) if ln_post else nn.Identity() + + self.apply(self._init_weights) + + if compile_mode is not None: + self.forward = torch.compile( + self.forward, mode=compile_mode, fullgraph=True + ) + if self.use_act_checkpoint and self.training: + torch._dynamo.config.optimize_ddp = False + + def _init_weights(self, m: nn.Module) -> None: + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + x = self.patch_embed(x) + h, w = x.shape[1], x.shape[2] + + s = 0 + if self.retain_cls_token: + # If cls_token is retained, we don't + # maintain spatial shape + x = torch.cat([self.class_embedding, x.flatten(1, 2)], dim=1) + s = 1 + + if self.pos_embed is not None: + x = x + get_abs_pos( + self.pos_embed, + self.pretrain_use_cls_token, + (h, w), + self.retain_cls_token, + tiling=self.tile_abs_pos, + ) + + x = self.ln_pre(x) + + outputs = [] + for i, blk in enumerate(self.blocks): + if self.use_act_checkpoint and self.training: + x = checkpoint.checkpoint(blk, x, use_reentrant=False) + else: + x = blk(x) + if (i == self.full_attn_ids[-1]) or ( + self.return_interm_layers and i in self.full_attn_ids + ): + if i == self.full_attn_ids[-1]: + x = self.ln_post(x) + + feats = x[:, s:] + if feats.ndim == 4: + feats = feats.permute(0, 3, 1, 2) + else: + assert feats.ndim == 3 + h = w = math.sqrt(feats.shape[1]) + feats = feats.reshape( + feats.shape[0], h, w, feats.shape[-1] + ).permute(0, 3, 1, 2) + + outputs.append(feats) + + return outputs + + def get_layer_id(self, layer_name: str) -> int: + # https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33 + num_layers = self.get_num_layers() + + if layer_name.find("rel_pos") != -1: + return num_layers + 1 + elif layer_name.find("ln_pre") != -1: + return 0 + elif layer_name.find("pos_embed") != -1 or layer_name.find("cls_token") != -1: + return 0 + elif layer_name.find("patch_embed") != -1: + return 0 + elif layer_name.find("blocks") != -1: + return int(layer_name.split("blocks")[1].split(".")[1]) + 1 + else: + return num_layers + 1 + + def get_num_layers(self) -> int: + return len(self.blocks) diff --git a/src/nn/segearth_ov3/sam3/model/vl_combiner.py b/src/nn/segearth_ov3/sam3/model/vl_combiner.py new file mode 100644 index 0000000..a53f821 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model/vl_combiner.py @@ -0,0 +1,179 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Provides utility to combine a vision backbone with a language backbone.""" + +from copy import copy +from typing import List, Optional + +import torch +import torch.nn as nn + +from torch.nn.attention import sdpa_kernel, SDPBackend + +from .act_ckpt_utils import activation_ckpt_wrapper +from .necks import Sam3DualViTDetNeck + + +class SAM3VLBackbone(nn.Module): + """This backbone combines a vision backbone and a language backbone without fusion. + As such it is more of a convenience wrapper to handle the two backbones together. + + It adds support for activation checkpointing and compilation. + """ + + def __init__( + self, + visual: Sam3DualViTDetNeck, + text, + compile_visual: bool = False, + act_ckpt_whole_vision_backbone: bool = False, + act_ckpt_whole_language_backbone: bool = False, + scalp=0, + ): + """Initialize the backbone combiner. + + :param visual: The vision backbone to use + :param text: The text encoder to use + """ + super().__init__() + self.vision_backbone: Sam3DualViTDetNeck = ( + torch.compile(visual) if compile_visual else visual + ) + self.language_backbone = text + self.scalp = scalp + # allow running activation checkpointing on the entire vision and language backbones + self.act_ckpt_whole_vision_backbone = act_ckpt_whole_vision_backbone + self.act_ckpt_whole_language_backbone = act_ckpt_whole_language_backbone + + def forward( + self, + samples: torch.Tensor, + captions: List[str], + input_boxes: Optional[torch.Tensor] = None, + additional_text: Optional[List[str]] = None, + ): + """Forward pass of the backbone combiner. + + :param samples: The input images + :param captions: The input captions + :param input_boxes: If the text contains place-holders for boxes, this + parameter contains the tensor containing their spatial features + :param additional_text: This can be used to encode some additional text + (different from the captions) in the same forward of the backbone + :return: Output dictionary with the following keys: + - vision_features: The output of the vision backbone + - language_features: The output of the language backbone + - language_mask: The attention mask of the language backbone + - vision_pos_enc: The positional encoding of the vision backbone + - (optional) additional_text_features: The output of the language + backbone for the additional text + - (optional) additional_text_mask: The attention mask of the + language backbone for the additional text + """ + output = self.forward_image(samples) + device = output["vision_features"].device + output.update(self.forward_text(captions, input_boxes, additional_text, device)) + return output + + def forward_image(self, samples: torch.Tensor): + return activation_ckpt_wrapper(self._forward_image_no_act_ckpt)( + samples=samples, + act_ckpt_enable=self.act_ckpt_whole_vision_backbone and self.training, + ) + + def _forward_image_no_act_ckpt(self, samples): + # Forward through backbone + sam3_features, sam3_pos, sam2_features, sam2_pos = self.vision_backbone.forward( + samples + ) + # sam3_features: [1, 256, 288, 288], [1, 256, 144, 144], [1, 256, 72, 72], [1, 256, 32, 32] + if self.scalp > 0: + # Discard the lowest resolution features + sam3_features, sam3_pos = ( + sam3_features[: -self.scalp], + sam3_pos[: -self.scalp], + ) + # sam3_features: [1, 256, 144, 144], [1, 256, 72, 72], [1, 256, 32, 32] + # sam3_pos: [1, 256, 144, 144], [1, 256, 72, 72], [1, 256, 32, 32] + if sam2_features is not None and sam2_pos is not None: + sam2_features, sam2_pos = ( + sam2_features[: -self.scalp], + sam2_pos[: -self.scalp], + ) + + sam2_output = None + + if sam2_features is not None and sam2_pos is not None: + sam2_src = sam2_features[-1] + sam2_output = { + "vision_features": sam2_src, + "vision_pos_enc": sam2_pos, + "backbone_fpn": sam2_features, + } + + sam3_src = sam3_features[-1] # [1, 256, 72, 72] + output = { + "vision_features": sam3_src, + "vision_pos_enc": sam3_pos, + "backbone_fpn": sam3_features, + "sam2_backbone_out": sam2_output, + } + + return output + + def forward_text( + self, captions, input_boxes=None, additional_text=None, device="cuda" + ): + return activation_ckpt_wrapper(self._forward_text_no_ack_ckpt)( + captions=captions, + input_boxes=input_boxes, + additional_text=additional_text, + device=device, + act_ckpt_enable=self.act_ckpt_whole_language_backbone and self.training, + ) + + def _forward_text_no_ack_ckpt( + self, + captions, + input_boxes=None, + additional_text=None, + device="cuda", + ): + output = {} + + # Forward through text_encoder + text_to_encode = copy(captions) + if additional_text is not None: + # if there are additional_text, we piggy-back them into this forward. + # They'll be used later for output alignment + text_to_encode += additional_text + + sdpa_context = sdpa_kernel( + [ + SDPBackend.MATH, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.FLASH_ATTENTION, + ] + ) + + with sdpa_context: + text_attention_mask, text_memory, text_embeds = self.language_backbone( + text_to_encode, input_boxes, device=device + ) + + if additional_text is not None: + output["additional_text_features"] = text_memory[:, -len(additional_text) :] + output["additional_text_mask"] = text_attention_mask[ + -len(additional_text) : + ] + + text_memory = text_memory[:, : len(captions)] + text_attention_mask = text_attention_mask[: len(captions)] + text_embeds = text_embeds[:, : len(captions)] + output["language_features"] = text_memory + output["language_mask"] = text_attention_mask + output["language_embeds"] = ( + text_embeds # Text embeddings before forward to the encoder + ) + + return output diff --git a/src/nn/segearth_ov3/sam3/model_builder.py b/src/nn/segearth_ov3/sam3/model_builder.py new file mode 100644 index 0000000..058bbec --- /dev/null +++ b/src/nn/segearth_ov3/sam3/model_builder.py @@ -0,0 +1,793 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os +from typing import Optional + +import torch +import torch.nn as nn +from huggingface_hub import hf_hub_download +from iopath.common.file_io import g_pathmgr +from sam3.model.decoder import ( + TransformerDecoder, + TransformerDecoderLayer, + TransformerDecoderLayerv2, + TransformerEncoderCrossAttention, +) +from sam3.model.encoder import TransformerEncoderFusion, TransformerEncoderLayer +from sam3.model.geometry_encoders import SequenceGeometryEncoder +from sam3.model.maskformer_segmentation import PixelDecoder, UniversalSegmentationHead +from sam3.model.memory import ( + CXBlock, + SimpleFuser, + SimpleMaskDownSampler, + SimpleMaskEncoder, +) +from sam3.model.model_misc import ( + DotProductScoring, + MLP, + MultiheadAttentionWrapper as MultiheadAttention, + TransformerWrapper, +) +from sam3.model.necks import Sam3DualViTDetNeck +from sam3.model.position_encoding import PositionEmbeddingSine +from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor +from sam3.model.sam3_image import Sam3Image, Sam3ImageOnVideoMultiGPU +from sam3.model.sam3_tracking_predictor import Sam3TrackerPredictor +from sam3.model.sam3_video_inference import Sam3VideoInferenceWithInstanceInteractivity +from sam3.model.sam3_video_predictor import Sam3VideoPredictorMultiGPU +from sam3.model.text_encoder_ve import VETextEncoder +from sam3.model.tokenizer_ve import SimpleTokenizer +from sam3.model.vitdet import ViT +from sam3.model.vl_combiner import SAM3VLBackbone +from sam3.sam.transformer import RoPEAttention + + +# Setup TensorFloat-32 for Ampere GPUs if available +def _setup_tf32() -> None: + """Enable TensorFloat-32 for Ampere GPUs if available.""" + if torch.cuda.is_available(): + device_props = torch.cuda.get_device_properties(0) + if device_props.major >= 8: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + +_setup_tf32() + + +def _create_position_encoding(precompute_resolution=None): + """Create position encoding for visual backbone.""" + return PositionEmbeddingSine( + num_pos_feats=256, + normalize=True, + scale=None, + temperature=10000, + precompute_resolution=precompute_resolution, + ) + + +def _create_vit_backbone(compile_mode=None): + """Create ViT backbone for visual feature extraction.""" + return ViT( + img_size=1008, + pretrain_img_size=336, + patch_size=14, + embed_dim=1024, + depth=32, + num_heads=16, + mlp_ratio=4.625, + norm_layer="LayerNorm", + drop_path_rate=0.1, + qkv_bias=True, + use_abs_pos=True, + tile_abs_pos=True, + global_att_blocks=(7, 15, 23, 31), + rel_pos_blocks=(), + use_rope=True, + use_interp_rope=True, + window_size=24, + pretrain_use_cls_token=True, + retain_cls_token=False, + ln_pre=True, + ln_post=False, + return_interm_layers=False, + bias_patch_embed=False, + compile_mode=compile_mode, + ) + + +def _create_vit_neck(position_encoding, vit_backbone, enable_inst_interactivity=False): + """Create ViT neck for feature pyramid.""" + return Sam3DualViTDetNeck( + position_encoding=position_encoding, + d_model=256, + scale_factors=[4.0, 2.0, 1.0, 0.5], + trunk=vit_backbone, + add_sam2_neck=enable_inst_interactivity, + ) + + +def _create_vl_backbone(vit_neck, text_encoder): + """Create visual-language backbone.""" + return SAM3VLBackbone(visual=vit_neck, text=text_encoder, scalp=1) + + +def _create_transformer_encoder() -> TransformerEncoderFusion: + """Create transformer encoder with its layer.""" + encoder_layer = TransformerEncoderLayer( + activation="relu", + d_model=256, + dim_feedforward=2048, + dropout=0.1, + pos_enc_at_attn=True, + pos_enc_at_cross_attn_keys=False, + pos_enc_at_cross_attn_queries=False, + pre_norm=True, + self_attention=MultiheadAttention( + num_heads=8, + dropout=0.1, + embed_dim=256, + batch_first=True, + ), + cross_attention=MultiheadAttention( + num_heads=8, + dropout=0.1, + embed_dim=256, + batch_first=True, + ), + ) + + encoder = TransformerEncoderFusion( + layer=encoder_layer, + num_layers=6, + d_model=256, + num_feature_levels=1, + frozen=False, + use_act_checkpoint=True, + add_pooled_text_to_img_feat=False, + pool_text_with_mask=True, + ) + return encoder + + +def _create_transformer_decoder() -> TransformerDecoder: + """Create transformer decoder with its layer.""" + decoder_layer = TransformerDecoderLayer( + activation="relu", + d_model=256, + dim_feedforward=2048, + dropout=0.1, + cross_attention=MultiheadAttention( + num_heads=8, + dropout=0.1, + embed_dim=256, + ), + n_heads=8, + use_text_cross_attention=True, + ) + + decoder = TransformerDecoder( + layer=decoder_layer, + num_layers=6, + num_queries=200, + return_intermediate=True, + box_refine=True, + num_o2m_queries=0, + dac=True, + boxRPB="log", + d_model=256, + frozen=False, + interaction_layer=None, + dac_use_selfatt_ln=True, + resolution=1008, + stride=14, + use_act_checkpoint=True, + presence_token=True, + ) + return decoder + + +def _create_dot_product_scoring(): + """Create dot product scoring module.""" + prompt_mlp = MLP( + input_dim=256, + hidden_dim=2048, + output_dim=256, + num_layers=2, + dropout=0.1, + residual=True, + out_norm=nn.LayerNorm(256), + ) + return DotProductScoring(d_model=256, d_proj=256, prompt_mlp=prompt_mlp) + + +def _create_segmentation_head(compile_mode=None): + """Create segmentation head with pixel decoder.""" + pixel_decoder = PixelDecoder( + num_upsampling_stages=3, + interpolation_mode="nearest", + hidden_dim=256, + compile_mode=compile_mode, + ) + + cross_attend_prompt = MultiheadAttention( + num_heads=8, + dropout=0, + embed_dim=256, + ) + + segmentation_head = UniversalSegmentationHead( + hidden_dim=256, + upsampling_stages=3, + aux_masks=False, + presence_head=False, + dot_product_scorer=None, + act_ckpt=True, + cross_attend_prompt=cross_attend_prompt, + pixel_decoder=pixel_decoder, + ) + return segmentation_head + + +def _create_geometry_encoder(): + """Create geometry encoder with all its components.""" + # Create position encoding for geometry encoder + geo_pos_enc = _create_position_encoding() + # Create CX block for fuser + cx_block = CXBlock( + dim=256, + kernel_size=7, + padding=3, + layer_scale_init_value=1.0e-06, + use_dwconv=True, + ) + # Create geometry encoder layer + geo_layer = TransformerEncoderLayer( + activation="relu", + d_model=256, + dim_feedforward=2048, + dropout=0.1, + pos_enc_at_attn=False, + pre_norm=True, + self_attention=MultiheadAttention( + num_heads=8, + dropout=0.1, + embed_dim=256, + batch_first=False, + ), + pos_enc_at_cross_attn_queries=False, + pos_enc_at_cross_attn_keys=True, + cross_attention=MultiheadAttention( + num_heads=8, + dropout=0.1, + embed_dim=256, + batch_first=False, + ), + ) + + # Create geometry encoder + input_geometry_encoder = SequenceGeometryEncoder( + pos_enc=geo_pos_enc, + encode_boxes_as_points=False, + points_direct_project=True, + points_pool=True, + points_pos_enc=True, + boxes_direct_project=True, + boxes_pool=True, + boxes_pos_enc=True, + d_model=256, + num_layers=3, + layer=geo_layer, + use_act_ckpt=True, + add_cls=True, + add_post_encode_proj=True, + ) + return input_geometry_encoder + + +def _create_sam3_model( + backbone, + transformer, + input_geometry_encoder, + segmentation_head, + dot_prod_scoring, + inst_interactive_predictor, + eval_mode, +): + """Create the SAM3 image model.""" + common_params = { + "backbone": backbone, + "transformer": transformer, + "input_geometry_encoder": input_geometry_encoder, + "segmentation_head": segmentation_head, + "num_feature_levels": 1, + "o2m_mask_predict": True, + "dot_prod_scoring": dot_prod_scoring, + "use_instance_query": False, + "multimask_output": True, + "inst_interactive_predictor": inst_interactive_predictor, + } + + matcher = None + if not eval_mode: + from sam3.train.matcher import BinaryHungarianMatcherV2 + + matcher = BinaryHungarianMatcherV2( + focal=True, + cost_class=2.0, + cost_bbox=5.0, + cost_giou=2.0, + alpha=0.25, + gamma=2, + stable=False, + ) + common_params["matcher"] = matcher + model = Sam3Image(**common_params) + + return model + + +def _create_tracker_maskmem_backbone(): + """Create the SAM3 Tracker memory encoder.""" + # Position encoding for mask memory backbone + position_encoding = PositionEmbeddingSine( + num_pos_feats=64, + normalize=True, + scale=None, + temperature=10000, + precompute_resolution=1008, + ) + + # Mask processing components + mask_downsampler = SimpleMaskDownSampler( + kernel_size=3, stride=2, padding=1, interpol_size=[1152, 1152] + ) + + cx_block_layer = CXBlock( + dim=256, + kernel_size=7, + padding=3, + layer_scale_init_value=1.0e-06, + use_dwconv=True, + ) + + fuser = SimpleFuser(layer=cx_block_layer, num_layers=2) + + maskmem_backbone = SimpleMaskEncoder( + out_dim=64, + position_encoding=position_encoding, + mask_downsampler=mask_downsampler, + fuser=fuser, + ) + + return maskmem_backbone + + +def _create_tracker_transformer(): + """Create the SAM3 Tracker transformer components.""" + # Self attention + self_attention = RoPEAttention( + embedding_dim=256, + num_heads=1, + downsample_rate=1, + dropout=0.1, + rope_theta=10000.0, + feat_sizes=[72, 72], + use_fa3=False, + use_rope_real=False, + ) + + # Cross attention + cross_attention = RoPEAttention( + embedding_dim=256, + num_heads=1, + downsample_rate=1, + dropout=0.1, + kv_in_dim=64, + rope_theta=10000.0, + feat_sizes=[72, 72], + rope_k_repeat=True, + use_fa3=False, + use_rope_real=False, + ) + + # Encoder layer + encoder_layer = TransformerDecoderLayerv2( + cross_attention_first=False, + activation="relu", + dim_feedforward=2048, + dropout=0.1, + pos_enc_at_attn=False, + pre_norm=True, + self_attention=self_attention, + d_model=256, + pos_enc_at_cross_attn_keys=True, + pos_enc_at_cross_attn_queries=False, + cross_attention=cross_attention, + ) + + # Encoder + encoder = TransformerEncoderCrossAttention( + remove_cross_attention_layers=[], + batch_first=True, + d_model=256, + frozen=False, + pos_enc_at_input=True, + layer=encoder_layer, + num_layers=4, + use_act_checkpoint=False, + ) + + # Transformer wrapper + transformer = TransformerWrapper( + encoder=encoder, + decoder=None, + d_model=256, + ) + + return transformer + + +def build_tracker( + apply_temporal_disambiguation: bool, with_backbone: bool = False, compile_mode=None +) -> Sam3TrackerPredictor: + """ + Build the SAM3 Tracker module for video tracking. + + Returns: + Sam3TrackerPredictor: Wrapped SAM3 Tracker module + """ + + # Create model components + maskmem_backbone = _create_tracker_maskmem_backbone() + transformer = _create_tracker_transformer() + backbone = None + if with_backbone: + vision_backbone = _create_vision_backbone(compile_mode=compile_mode) + backbone = SAM3VLBackbone(scalp=1, visual=vision_backbone, text=None) + # Create the Tracker module + model = Sam3TrackerPredictor( + image_size=1008, + num_maskmem=7, + backbone=backbone, + backbone_stride=14, + transformer=transformer, + maskmem_backbone=maskmem_backbone, + # SAM parameters + multimask_output_in_sam=True, + # Evaluation + forward_backbone_per_frame_for_eval=True, + trim_past_non_cond_mem_for_eval=False, + # Multimask + multimask_output_for_tracking=True, + multimask_min_pt_num=0, + multimask_max_pt_num=1, + # Additional settings + always_start_from_first_ann_frame=False, + # Mask overlap + non_overlap_masks_for_mem_enc=False, + non_overlap_masks_for_output=False, + max_cond_frames_in_attn=4, + offload_output_to_cpu_for_eval=False, + # SAM decoder settings + sam_mask_decoder_extra_args={ + "dynamic_multimask_via_stability": True, + "dynamic_multimask_stability_delta": 0.05, + "dynamic_multimask_stability_thresh": 0.98, + }, + clear_non_cond_mem_around_input=True, + fill_hole_area=0, + use_memory_selection=apply_temporal_disambiguation, + ) + + return model + + +def _create_text_encoder(bpe_path: str) -> VETextEncoder: + """Create SAM3 text encoder.""" + tokenizer = SimpleTokenizer(bpe_path=bpe_path) + return VETextEncoder( + tokenizer=tokenizer, + d_model=256, + width=1024, + heads=16, + layers=24, + ) + + +def _create_vision_backbone( + compile_mode=None, enable_inst_interactivity=True +) -> Sam3DualViTDetNeck: + """Create SAM3 visual backbone with ViT and neck.""" + # Position encoding + position_encoding = _create_position_encoding(precompute_resolution=1008) + # ViT backbone + vit_backbone: ViT = _create_vit_backbone(compile_mode=compile_mode) + vit_neck: Sam3DualViTDetNeck = _create_vit_neck( + position_encoding, + vit_backbone, + enable_inst_interactivity=enable_inst_interactivity, + ) + # Visual neck + return vit_neck + + +def _create_sam3_transformer(has_presence_token: bool = True) -> TransformerWrapper: + """Create SAM3 transformer encoder and decoder.""" + encoder: TransformerEncoderFusion = _create_transformer_encoder() + decoder: TransformerDecoder = _create_transformer_decoder() + + return TransformerWrapper(encoder=encoder, decoder=decoder, d_model=256) + + +def _load_checkpoint(model, checkpoint_path): + """Load model checkpoint from file.""" + with g_pathmgr.open(checkpoint_path, "rb") as f: + ckpt = torch.load(f, map_location="cpu", weights_only=True) + if "model" in ckpt and isinstance(ckpt["model"], dict): + ckpt = ckpt["model"] + sam3_image_ckpt = { + k.replace("detector.", ""): v for k, v in ckpt.items() if "detector" in k + } + if model.inst_interactive_predictor is not None: + sam3_image_ckpt.update( + { + k.replace("tracker.", "inst_interactive_predictor.model."): v + for k, v in ckpt.items() + if "tracker" in k + } + ) + missing_keys, _ = model.load_state_dict(sam3_image_ckpt, strict=False) + if len(missing_keys) > 0: + print( + f"loaded {checkpoint_path} and found " + f"missing and/or unexpected keys:\n{missing_keys=}" + ) + + +def _setup_device_and_mode(model, device, eval_mode): + """Setup model device and evaluation mode.""" + if device == "cuda": + model = model.cuda() + if eval_mode: + model.eval() + return model + + +def build_sam3_image_model( + bpe_path=None, + device="cuda" if torch.cuda.is_available() else "cpu", + eval_mode=True, + checkpoint_path=None, + load_from_HF=True, + enable_segmentation=True, + enable_inst_interactivity=False, + compile=False, +): + """ + Build SAM3 image model + + Args: + bpe_path: Path to the BPE tokenizer vocabulary + device: Device to load the model on ('cuda' or 'cpu') + eval_mode: Whether to set the model to evaluation mode + checkpoint_path: Optional path to model checkpoint + enable_segmentation: Whether to enable segmentation head + enable_inst_interactivity: Whether to enable instance interactivity (SAM 1 task) + compile_mode: To enable compilation, set to "default" + + Returns: + A SAM3 image model + """ + if bpe_path is None: + bpe_path = os.path.join( + os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz" + ) + # Create visual components + compile_mode = "default" if compile else None + vision_encoder = _create_vision_backbone( + compile_mode=compile_mode, enable_inst_interactivity=enable_inst_interactivity + ) + + # Create text components + text_encoder = _create_text_encoder(bpe_path) + + # Create visual-language backbone + backbone = _create_vl_backbone(vision_encoder, text_encoder) + + # Create transformer components + transformer = _create_sam3_transformer() + + # Create dot product scoring + dot_prod_scoring = _create_dot_product_scoring() + + # Create segmentation head if enabled + segmentation_head = ( + _create_segmentation_head(compile_mode=compile_mode) + if enable_segmentation + else None + ) + + # Create geometry encoder + input_geometry_encoder = _create_geometry_encoder() + if enable_inst_interactivity: + sam3_pvs_base = build_tracker(apply_temporal_disambiguation=False) + inst_predictor = SAM3InteractiveImagePredictor(sam3_pvs_base) + else: + inst_predictor = None + # Create the SAM3 model + model = _create_sam3_model( + backbone, + transformer, + input_geometry_encoder, + segmentation_head, + dot_prod_scoring, + inst_predictor, + eval_mode, + ) + if load_from_HF and checkpoint_path is None: + checkpoint_path = download_ckpt_from_hf() + # Load checkpoint if provided + if checkpoint_path is not None: + _load_checkpoint(model, checkpoint_path) + + # Setup device and mode + model = _setup_device_and_mode(model, device, eval_mode) + + return model + + +def download_ckpt_from_hf(): + SAM3_MODEL_ID = "facebook/sam3" + SAM3_CKPT_NAME = "sam3.pt" + SAM3_CFG_NAME = "config.json" + _ = hf_hub_download(repo_id=SAM3_MODEL_ID, filename=SAM3_CFG_NAME) + checkpoint_path = hf_hub_download(repo_id=SAM3_MODEL_ID, filename=SAM3_CKPT_NAME) + return checkpoint_path + + +def build_sam3_video_model( + checkpoint_path: Optional[str] = None, + load_from_HF=True, + bpe_path: Optional[str] = None, + has_presence_token: bool = True, + geo_encoder_use_img_cross_attn: bool = True, + strict_state_dict_loading: bool = True, + apply_temporal_disambiguation: bool = True, + device="cuda" if torch.cuda.is_available() else "cpu", + compile=False, +) -> Sam3VideoInferenceWithInstanceInteractivity: + """ + Build SAM3 dense tracking model. + + Args: + checkpoint_path: Optional path to checkpoint file + bpe_path: Path to the BPE tokenizer file + + Returns: + Sam3VideoInferenceWithInstanceInteractivity: The instantiated dense tracking model + """ + if bpe_path is None: + bpe_path = os.path.join( + os.path.dirname(__file__), "..", "assets", "bpe_simple_vocab_16e6.txt.gz" + ) + + # Build Tracker module + tracker = build_tracker(apply_temporal_disambiguation=apply_temporal_disambiguation) + + # Build Detector components + visual_neck = _create_vision_backbone() + text_encoder = _create_text_encoder(bpe_path) + backbone = SAM3VLBackbone(scalp=1, visual=visual_neck, text=text_encoder) + transformer = _create_sam3_transformer(has_presence_token=has_presence_token) + segmentation_head: UniversalSegmentationHead = _create_segmentation_head() + input_geometry_encoder = _create_geometry_encoder() + + # Create main dot product scoring + main_dot_prod_mlp = MLP( + input_dim=256, + hidden_dim=2048, + output_dim=256, + num_layers=2, + dropout=0.1, + residual=True, + out_norm=nn.LayerNorm(256), + ) + main_dot_prod_scoring = DotProductScoring( + d_model=256, d_proj=256, prompt_mlp=main_dot_prod_mlp + ) + + # Build Detector module + detector = Sam3ImageOnVideoMultiGPU( + num_feature_levels=1, + backbone=backbone, + transformer=transformer, + segmentation_head=segmentation_head, + semantic_segmentation_head=None, + input_geometry_encoder=input_geometry_encoder, + use_early_fusion=True, + use_dot_prod_scoring=True, + dot_prod_scoring=main_dot_prod_scoring, + supervise_joint_box_scores=has_presence_token, + ) + + # Build the main SAM3 video model + if apply_temporal_disambiguation: + model = Sam3VideoInferenceWithInstanceInteractivity( + detector=detector, + tracker=tracker, + score_threshold_detection=0.5, + assoc_iou_thresh=0.1, + det_nms_thresh=0.1, + new_det_thresh=0.7, + hotstart_delay=15, + hotstart_unmatch_thresh=8, + hotstart_dup_thresh=8, + suppress_unmatched_only_within_hotstart=True, + min_trk_keep_alive=-1, + max_trk_keep_alive=30, + init_trk_keep_alive=30, + suppress_overlapping_based_on_recent_occlusion_threshold=0.7, + suppress_det_close_to_boundary=False, + fill_hole_area=16, + recondition_every_nth_frame=16, + masklet_confirmation_enable=False, + decrease_trk_keep_alive_for_empty_masklets=False, + image_size=1008, + image_mean=(0.5, 0.5, 0.5), + image_std=(0.5, 0.5, 0.5), + compile_model=compile, + ) + else: + # a version without any heuristics for ablation studies + model = Sam3VideoInferenceWithInstanceInteractivity( + detector=detector, + tracker=tracker, + score_threshold_detection=0.5, + assoc_iou_thresh=0.1, + det_nms_thresh=0.1, + new_det_thresh=0.7, + hotstart_delay=0, + hotstart_unmatch_thresh=0, + hotstart_dup_thresh=0, + suppress_unmatched_only_within_hotstart=True, + min_trk_keep_alive=-1, + max_trk_keep_alive=30, + init_trk_keep_alive=30, + suppress_overlapping_based_on_recent_occlusion_threshold=0.7, + suppress_det_close_to_boundary=False, + fill_hole_area=16, + recondition_every_nth_frame=0, + masklet_confirmation_enable=False, + decrease_trk_keep_alive_for_empty_masklets=False, + image_size=1008, + image_mean=(0.5, 0.5, 0.5), + image_std=(0.5, 0.5, 0.5), + compile_model=compile, + ) + + # Load checkpoint if provided + if load_from_HF and checkpoint_path is None: + checkpoint_path = download_ckpt_from_hf() + if checkpoint_path is not None: + with g_pathmgr.open(checkpoint_path, "rb") as f: + ckpt = torch.load(f, map_location="cpu", weights_only=True) + if "model" in ckpt and isinstance(ckpt["model"], dict): + ckpt = ckpt["model"] + + missing_keys, unexpected_keys = model.load_state_dict( + ckpt, strict=strict_state_dict_loading + ) + if missing_keys: + print(f"Missing keys: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys: {unexpected_keys}") + + model.to(device=device) + return model + + +def build_sam3_video_predictor(*model_args, gpus_to_use=None, **model_kwargs): + return Sam3VideoPredictorMultiGPU( + *model_args, gpus_to_use=gpus_to_use, **model_kwargs + ) diff --git a/src/nn/segearth_ov3/sam3/perflib/__init__.py b/src/nn/segearth_ov3/sam3/perflib/__init__.py new file mode 100644 index 0000000..5c3823b --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os + +is_enabled = False +if os.getenv("USE_PERFLIB", "1") == "1": + # print("Enabled the use of perflib.\n", end="") + is_enabled = True diff --git a/src/nn/segearth_ov3/sam3/perflib/associate_det_trk.py b/src/nn/segearth_ov3/sam3/perflib/associate_det_trk.py new file mode 100644 index 0000000..508ae81 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/associate_det_trk.py @@ -0,0 +1,137 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from collections import defaultdict + +import torch +import torch.nn.functional as F +from sam3.perflib.masks_ops import mask_iou +from scipy.optimize import linear_sum_assignment + + +def associate_det_trk( + det_masks, + track_masks, + iou_threshold=0.5, + iou_threshold_trk=0.5, + det_scores=None, + new_det_thresh=0.0, +): + """ + Optimized implementation of detection <-> track association that minimizes DtoH syncs. + + Args: + det_masks: (N, H, W) tensor of predicted masks + track_masks: (M, H, W) tensor of track masks + + Returns: + new_det_indices: list of indices in det_masks considered 'new' + unmatched_trk_indices: list of indices in track_masks considered 'unmatched' + """ + with torch.autograd.profiler.record_function("perflib: associate_det_trk"): + assert isinstance(det_masks, torch.Tensor), "det_masks should be a tensor" + assert isinstance(track_masks, torch.Tensor), "track_masks should be a tensor" + if det_masks.size(0) == 0 or track_masks.size(0) == 0: + return list(range(det_masks.size(0))), [], {}, {} # all detections are new + + if list(det_masks.shape[-2:]) != list(track_masks.shape[-2:]): + # resize to the smaller size to save GPU memory + if torch.numel(det_masks[-2:]) < torch.numel(track_masks[-2:]): + track_masks = ( + F.interpolate( + track_masks.unsqueeze(1).float(), + size=det_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + > 0 + ) + else: + # resize detections to track size + det_masks = ( + F.interpolate( + det_masks.unsqueeze(1).float(), + size=track_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + > 0 + ) + + det_masks = det_masks > 0 + track_masks = track_masks > 0 + + iou = mask_iou(det_masks, track_masks) # (N, M) + igeit = iou >= iou_threshold + igeit_any_dim_1 = igeit.any(dim=1) + igeit_trk = iou >= iou_threshold_trk + + iou_list = iou.cpu().numpy().tolist() + igeit_list = igeit.cpu().numpy().tolist() + igeit_any_dim_1_list = igeit_any_dim_1.cpu().numpy().tolist() + igeit_trk_list = igeit_trk.cpu().numpy().tolist() + + det_scores_list = ( + det_scores + if det_scores is None + else det_scores.cpu().float().numpy().tolist() + ) + + # Hungarian matching for tracks (one-to-one: each track matches at most one detection) + # For detections: allow many tracks to match to the same detection (many-to-one) + + # If either is empty, return all detections as new + if det_masks.size(0) == 0 or track_masks.size(0) == 0: + return list(range(det_masks.size(0))), [], {} + + # Hungarian matching: maximize IoU for tracks + cost_matrix = 1 - iou.cpu().numpy() # Hungarian solves for minimum cost + row_ind, col_ind = linear_sum_assignment(cost_matrix) + + def branchy_hungarian_better_uses_the_cpu( + cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks + ): + matched_trk = set() + matched_det = set() + matched_det_scores = {} # track index -> [det_score, det_score * iou] det score of matched detection mask + for d, t in zip(row_ind, col_ind): + matched_det_scores[t] = [ + det_scores_list[d], + det_scores_list[d] * iou_list[d][t], + ] + if igeit_trk_list[d][t]: + matched_trk.add(t) + matched_det.add(d) + + # Tracks not matched by Hungarian assignment above threshold are unmatched + unmatched_trk_indices = [ + t for t in range(track_masks.size(0)) if t not in matched_trk + ] + + # For detections: allow many tracks to match to the same detection (many-to-one) + # So, a detection is 'new' if it does not match any track above threshold + assert track_masks.size(0) == igeit.size( + 1 + ) # Needed for loop optimizaiton below + new_det_indices = [] + for d in range(det_masks.size(0)): + if not igeit_any_dim_1_list[d]: + if det_scores is not None and det_scores[d] >= new_det_thresh: + new_det_indices.append(d) + + # for each detection, which tracks it matched to (above threshold) + det_to_matched_trk = defaultdict(list) + for d in range(det_masks.size(0)): + for t in range(track_masks.size(0)): + if igeit_list[d][t]: + det_to_matched_trk[d].append(t) + + return ( + new_det_indices, + unmatched_trk_indices, + det_to_matched_trk, + matched_det_scores, + ) + + return (branchy_hungarian_better_uses_the_cpu)( + cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks + ) diff --git a/src/nn/segearth_ov3/sam3/perflib/compile.py b/src/nn/segearth_ov3/sam3/perflib/compile.py new file mode 100644 index 0000000..f427aa7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/compile.py @@ -0,0 +1,99 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +def recursive_fn_factory(fn): + def recursive_fn(b): + if isinstance(b, dict): + return {k: recursive_fn(b[k]) for k in b} + if isinstance(b, list): + return [recursive_fn(t) for t in b] + if isinstance(b, tuple): + return tuple(recursive_fn(t) for t in b) + if isinstance(b, torch.Tensor): + return fn(b) + # Yes, writing out an explicit white list of + # trivial types is tedious, but so are bugs that + # come from not applying fn, when expected to have + # applied it. + if b is None: + return b + trivial_types = [bool, int] + for t in trivial_types: + if isinstance(b, t): + return b + raise TypeError(f"Unexpected type {type(b)}") + + return recursive_fn + + +recursive_contiguous = recursive_fn_factory(lambda x: x.contiguous()) +recursive_clone = recursive_fn_factory(torch.clone) + + +def compile_wrapper( + fn, *, mode="max-autotune", fullgraph=True, dynamic=False, name=None +): + compiled_fn = torch.compile(fn, mode=mode, fullgraph=fullgraph, dynamic=dynamic) + + def compiled_fn_wrapper(*args, **kwargs): + with torch.autograd.profiler.record_function( + f"compiled {fn}" if name is None else name + ): + cont_args = recursive_contiguous(args) + cont_kwargs = recursive_contiguous(kwargs) + result = compiled_fn(*cont_args, **cont_kwargs) + cloned_result = recursive_clone(result) + return cloned_result + + return compiled_fn_wrapper + + +def shape_logging_wrapper(fn, keep_kwargs, enable_logging=False): + """ + Wraps a function and prints the shapes of all tensor inputs. + Only prints when a new combination of shapes is seen. + Thread-safe. + + Args: + fn: Function to wrap + enable_logging: Boolean flag to enable/disable logging + """ + seen_shapes = set() + + def get_shape(obj): + if isinstance(obj, torch.Tensor): + return obj.shape + elif isinstance(obj, (list, tuple)): + if len(obj) > 1: + return tuple(get_shape(x) for x in obj) + return get_shape(obj[0]) + elif isinstance(obj, dict): + return tuple(sorted((k, get_shape(v)) for k, v in obj.items())) + else: + return type(obj).__name__ + + def wrapper(*args, **kwargs): + shapes = tuple(get_shape(arg) for arg in args) + tuple( + (k, get_shape(v)) + for k, v in kwargs.items() + if isinstance(v, (torch.Tensor, list)) + and (len(keep_kwargs) > 0 and k in keep_kwargs) + ) + if shapes not in seen_shapes: + seen_shapes.add(shapes) + if enable_logging: + print(f"[ShapeLogger] New input shapes for {fn.__qualname__}: {shapes}") + return fn(*args, **kwargs) + + # Allow toggling the flag at runtime + wrapper.enable_logging = enable_logging + + def set_logging(enabled=False): + nonlocal enable_logging + enable_logging = enabled + wrapper.enable_logging = enable_logging + + wrapper.set_logging = set_logging + return wrapper diff --git a/src/nn/segearth_ov3/sam3/perflib/connected_components.py b/src/nn/segearth_ov3/sam3/perflib/connected_components.py new file mode 100644 index 0000000..c96932a --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/connected_components.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import logging + +import torch + +try: + from cc_torch import get_connected_components + + HAS_CC_TORCH = True +except ImportError: + logging.debug( + "cc_torch not found. Consider installing for better performance. Command line:" + " pip install git+https://github.com/ronghanghu/cc_torch.git" + ) + HAS_CC_TORCH = False + + +def connected_components_cpu_single(values: torch.Tensor): + assert values.dim() == 2 + from skimage.measure import label + + labels, num = label(values.cpu().numpy(), return_num=True) + labels = torch.from_numpy(labels) + counts = torch.zeros_like(labels) + for i in range(1, num + 1): + cur_mask = labels == i + cur_count = cur_mask.sum() + counts[cur_mask] = cur_count + return labels, counts + + +def connected_components_cpu(input_tensor: torch.Tensor): + out_shape = input_tensor.shape + if input_tensor.dim() == 4 and input_tensor.shape[1] == 1: + input_tensor = input_tensor.squeeze(1) + else: + assert ( + input_tensor.dim() == 3 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + batch_size = input_tensor.shape[0] + labels_list = [] + counts_list = [] + for b in range(batch_size): + labels, counts = connected_components_cpu_single(input_tensor[b]) + labels_list.append(labels) + counts_list.append(counts) + labels_tensor = torch.stack(labels_list, dim=0).to(input_tensor.device) + counts_tensor = torch.stack(counts_list, dim=0).to(input_tensor.device) + return labels_tensor.view(out_shape), counts_tensor.view(out_shape) + + +def connected_components(input_tensor: torch.Tensor): + """ + Computes connected components labeling on a batch of 2D tensors, using the best available backend. + + Args: + input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Both tensors have the same shape as input_tensor. + - A tensor with dense labels. Background is 0. + - A tensor with the size of the connected component for each pixel. + """ + if input_tensor.dim() == 3: + input_tensor = input_tensor.unsqueeze(1) + + assert ( + input_tensor.dim() == 4 and input_tensor.shape[1] == 1 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + if input_tensor.is_cuda: + if HAS_CC_TORCH: + return get_connected_components(input_tensor.to(torch.uint8)) + else: + # triton fallback + from sam3.perflib.triton.connected_components import ( + connected_components_triton, + ) + + return connected_components_triton(input_tensor) + + # CPU fallback + return connected_components_cpu(input_tensor) diff --git a/src/nn/segearth_ov3/sam3/perflib/fa3.py b/src/nn/segearth_ov3/sam3/perflib/fa3.py new file mode 100644 index 0000000..8f8c9bd --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/fa3.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +@torch.library.custom_op("flash::flash_attn_func", mutates_args=()) +def flash_attn_func_op( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor +) -> torch.Tensor: + from flash_attn_interface import flash_attn_func as fa3 + + return fa3(q, k, v) + + +def flash_attn_func(q, k, v): + dtype = torch.float8_e4m3fn + return flash_attn_func_op(q.to(dtype), k.to(dtype), v.to(dtype)).to(q.dtype) + + +@flash_attn_func_op.register_fake +def _(q, k, v, **kwargs): + # two outputs: + # 1. output: (batch, seq_len, num_heads, head_dim) + # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32 + # output needs to be bfloat16, not float8! + meta_q = torch.empty_like(q, dtype=torch.bfloat16).contiguous() + return meta_q diff --git a/src/nn/segearth_ov3/sam3/perflib/masks_ops.py b/src/nn/segearth_ov3/sam3/perflib/masks_ops.py new file mode 100644 index 0000000..48299d5 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/masks_ops.py @@ -0,0 +1,69 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +def masks_to_boxes(masks: torch.Tensor, obj_ids: list[int]): + with torch.autograd.profiler.record_function("perflib: masks_to_boxes"): + # Sanity check based on callsite for replacement + assert masks.shape[0] == len(obj_ids) + assert masks.dim() == 3 + + # Based on torchvision masks_to_boxes + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device, dtype=torch.float) + + N, H, W = masks.shape + device = masks.device + y = torch.arange(H, device=device).view(1, H) + x = torch.arange(W, device=device).view(1, W) + + masks_with_obj = masks != 0 # N, H, W + masks_with_obj_x = masks_with_obj.amax( + dim=1 + ) # N, H (which columns have objects) + masks_with_obj_y = masks_with_obj.amax(dim=2) # N, W (which rows have objects) + masks_without_obj_x = ~masks_with_obj_x + masks_without_obj_y = ~masks_with_obj_y + + bounding_boxes_0 = torch.amin( + (masks_without_obj_x * W) + (masks_with_obj_x * x), dim=1 + ) + bounding_boxes_1 = torch.amin( + (masks_without_obj_y * H) + (masks_with_obj_y * y), dim=1 + ) + bounding_boxes_2 = torch.amax(masks_with_obj_x * x, dim=1) + bounding_boxes_3 = torch.amax(masks_with_obj_y * y, dim=1) + + bounding_boxes = torch.stack( + [bounding_boxes_0, bounding_boxes_1, bounding_boxes_2, bounding_boxes_3], + dim=1, + ).to(dtype=torch.float) + assert bounding_boxes.shape == (N, 4) + assert bounding_boxes.device == masks.device + assert bounding_boxes.dtype == torch.float + return bounding_boxes + + +def mask_iou(pred_masks: torch.Tensor, gt_masks: torch.Tensor) -> torch.Tensor: + """ + Compute the IoU (Intersection over Union) between predicted masks and ground truth masks. + Args: + - pred_masks: (N, H, W) bool Tensor, containing binary predicted segmentation masks + - gt_masks: (M, H, W) bool Tensor, containing binary ground truth segmentation masks + Returns: + - ious: (N, M) float Tensor, containing IoUs for each pair of predicted and ground truth masks + """ + assert pred_masks.dtype == gt_masks.dtype == torch.bool + N, H, W = pred_masks.shape + M, _, _ = gt_masks.shape + + # Flatten masks: (N, 1, H*W) and (1, M, H*W) + pred_flat = pred_masks.view(N, 1, H * W) + gt_flat = gt_masks.view(1, M, H * W) + + # Compute intersection and union: (N, M) + intersection = (pred_flat & gt_flat).sum(dim=2).float() + union = (pred_flat | gt_flat).sum(dim=2).float() + ious = intersection / union.clamp(min=1) + return ious # shape: (N, M) diff --git a/src/nn/segearth_ov3/sam3/perflib/nms.py b/src/nn/segearth_ov3/sam3/perflib/nms.py new file mode 100644 index 0000000..b3efc59 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/nms.py @@ -0,0 +1,91 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging + +import numpy as np +import torch + +from sam3.perflib.masks_ops import mask_iou + + +try: + from torch_generic_nms import generic_nms as generic_nms_cuda + + GENERIC_NMS_AVAILABLE = True +except ImportError: + logging.debug( + "Falling back to triton or CPU mask NMS implementation -- please install `torch_generic_nms` via\n\t" + 'pip uninstall -y torch_generic_nms; TORCH_CUDA_ARCH_LIST="8.0 9.0" pip install git+https://github.com/ronghanghu/torch_generic_nms' + ) + GENERIC_NMS_AVAILABLE = False + + +def nms_masks( + pred_probs: torch.Tensor, + pred_masks: torch.Tensor, + prob_threshold: float, + iou_threshold: float, +) -> torch.Tensor: + """ + Args: + - pred_probs: (num_det,) float Tensor, containing the score (probability) of each detection + - pred_masks: (num_det, H_mask, W_mask) float Tensor, containing the binary segmentation mask of each detection + - prob_threshold: float, score threshold to prefilter detections (NMS is performed on detections above threshold) + - iou_threshold: float, mask IoU threshold for NMS + + Returns: + - keep: (num_det,) bool Tensor, indicating whether each detection is kept after score thresholding + NMS + """ + # prefilter the detections with prob_threshold ("valid" are those above prob_threshold) + is_valid = pred_probs > prob_threshold # (num_det,) + probs = pred_probs[is_valid] # (num_valid,) + masks_binary = pred_masks[is_valid] > 0 # (num_valid, H_mask, W_mask) + if probs.numel() == 0: + return is_valid # no valid detection, return empty keep mask + + ious = mask_iou(masks_binary, masks_binary) # (num_valid, num_valid) + kept_inds = generic_nms(ious, probs, iou_threshold) + + # valid_inds are the indices among `probs` of valid detections before NMS (or -1 for invalid) + valid_inds = torch.where(is_valid, is_valid.cumsum(dim=0) - 1, -1) # (num_det,) + keep = torch.isin(valid_inds, kept_inds) # (num_det,) + return keep + + +def generic_nms( + ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5 +) -> torch.Tensor: + """A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix.""" + + assert ious.dim() == 2 and ious.size(0) == ious.size(1) + assert scores.dim() == 1 and scores.size(0) == ious.size(0) + + if ious.is_cuda: + if GENERIC_NMS_AVAILABLE: + return generic_nms_cuda(ious, scores, iou_threshold, use_iou_matrix=True) + else: + from sam3.perflib.triton.nms import nms_triton + + return nms_triton(ious, scores, iou_threshold) + + return generic_nms_cpu(ious, scores, iou_threshold) + + +def generic_nms_cpu( + ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5 +) -> torch.Tensor: + """ + A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix. (CPU implementation + based on https://github.com/jwyang/faster-rcnn.pytorch/blob/master/lib/model/nms/nms_cpu.py) + """ + ious_np = ious.float().detach().cpu().numpy() + scores_np = scores.float().detach().cpu().numpy() + order = scores_np.argsort()[::-1] + kept_inds = [] + while order.size > 0: + i = order.item(0) + kept_inds.append(i) + inds = np.where(ious_np[i, order[1:]] <= iou_threshold)[0] + order = order[inds + 1] + + return torch.tensor(kept_inds, dtype=torch.int64, device=scores.device) diff --git a/src/nn/segearth_ov3/sam3/perflib/tests/assets/masks.tiff b/src/nn/segearth_ov3/sam3/perflib/tests/assets/masks.tiff new file mode 100644 index 0000000..7a8efc6 Binary files /dev/null and b/src/nn/segearth_ov3/sam3/perflib/tests/assets/masks.tiff differ diff --git a/src/nn/segearth_ov3/sam3/perflib/tests/tests.py b/src/nn/segearth_ov3/sam3/perflib/tests/tests.py new file mode 100644 index 0000000..0fb88ad --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/tests/tests.py @@ -0,0 +1,59 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os + +import numpy as np +import pytest +import torch +from PIL import Image +from sam3.perflib.masks_ops import masks_to_boxes + + +class TestMasksToBoxes: + def test_masks_box(self): + def masks_box_check(masks, expected, atol=1e-4): + out = masks_to_boxes(masks, [1 for _ in range(masks.shape[0])]) + assert out.dtype == torch.float + print("out: ", out) + print("expected: ", expected) + torch.testing.assert_close( + out, expected, rtol=0.0, check_dtype=True, atol=atol + ) + + # Check for int type boxes. + def _get_image(): + assets_directory = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "assets" + ) + mask_path = os.path.join(assets_directory, "masks.tiff") + image = Image.open(mask_path) + return image + + def _create_masks(image, masks): + for index in range(image.n_frames): + image.seek(index) + frame = np.array(image) + masks[index] = torch.tensor(frame) + + return masks + + expected = torch.tensor( + [ + [127, 2, 165, 40], + [2, 50, 44, 92], + [56, 63, 98, 100], + [139, 68, 175, 104], + [160, 112, 198, 145], + [49, 138, 99, 182], + [108, 148, 152, 213], + ], + dtype=torch.float, + ) + + image = _get_image() + for dtype in [torch.float16, torch.float32, torch.float64]: + masks = torch.zeros( + (image.n_frames, image.height, image.width), dtype=dtype + ) + masks = _create_masks(image, masks) + masks_box_check(masks, expected) diff --git a/src/nn/segearth_ov3/sam3/perflib/triton/connected_components.py b/src/nn/segearth_ov3/sam3/perflib/triton/connected_components.py new file mode 100644 index 0000000..253ca9d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/triton/connected_components.py @@ -0,0 +1,468 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import math + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _any_combine(a, b): + return a | b + + +@triton.jit +def tl_any(a, dim=0): + return tl.reduce(a, dim, _any_combine) + + +# ============================================================================== +# ## Phase 1: Initialization Kernel +# ============================================================================== +# Each foreground pixel (value > 0) gets a unique label equal to its +# linear index. Background pixels (value == 0) get a sentinel label of -1. +# Note that the indexing is done across batch boundaries for simplicity +# (i.e., the first pixel of image 1 gets label H*W, etc.) + + +@triton.jit +def _init_labels_kernel( + input_ptr, labels_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + input_values = tl.load(input_ptr + offsets, mask=mask, other=0) + + indices = tl.where((input_values != 0), offsets, -1) + tl.store(labels_ptr + offsets, indices, mask=mask) + + +# ============================================================================== +# ## Phase 2: Local merging +# ============================================================================== +# Each pixel tries to merge with its 8-connected neighbors (up, down, left, right) +# if they have the same value. This is done using a disjoint-set union operation. + + +@triton.jit +def find(labels_ptr, indices, mask): + current_pids = indices + + # 'is_done' tracks lanes that have finished their work. + # A lane is initially "done" if it's not active (mask is False). + is_done = ~mask + + # Loop as long as there is at least one lane that is NOT done. + while tl_any(~is_done): + # The work_mask is for lanes that are still active and seeking their root. + work_mask = ~is_done + parents = tl.load(labels_ptr + current_pids, mask=work_mask, other=-1) + # A lane is now done if its parent is itself (it's a root) + # or if it hits a -1 sentinel (a safe exit condition). + is_root = parents == current_pids + is_sentinel = parents == -1 + is_done |= is_root | is_sentinel + + # For lanes that are not yet done, update their pid to their parent to continue traversal. + current_pids = tl.where(is_done, current_pids, parents) + # We could add the following line to do path compression, but experimentally it's slower + # tl.atomic_min(labels_ptr + indices, current_pids, mask=mask) + return current_pids + + +@triton.jit +def union(labels_ptr, a, b, process_mask): + # This function implements a disjoint-set union + # As an invariant, we use the fact that the roots have the lower id. That helps parallelization + # However, that is not sufficient by itself. Suppose two threads want to do union(0,2) and union(1,2) at the same time + # Then if we do a naive atomic_min, 0 and 1 will compete to be the new parent of 2 and min(0, 1) will win. + # However, 1 still needs to be merged with the new {0, 2} component. + # To ensure that merge is also done, we need to detect whether the merge was successful, and if not retry until it is + + current_a = a + current_b = b + + final_root = a + # A mask to track which lanes have successfully completed their union. + done_mask = ~process_mask # tl.zeros_like(a) == 1 # Init with all False + + while tl_any(~done_mask): + # Define the mask for lanes that still need work in this iteration + work_mask = process_mask & ~done_mask + + # Find the roots for the current a and b values in the active lanes + root_a = find(labels_ptr, current_a, work_mask) + tl.debug_barrier() + root_b = find(labels_ptr, current_b, work_mask) + + # 7. Merge logic + # If roots are already the same, the sets are already merged. Mark as done. + are_equal = root_a == root_b + final_root = tl.where(are_equal & work_mask & ~done_mask, root_a, final_root) + done_mask |= are_equal & work_mask + + # Define masks for the two merge cases (a < b or b < a) + a_is_smaller = root_a < root_b + + # Case 1: root_a < root_b. Attempt to set parent[root_b] = root_a + merge_mask_a_smaller = work_mask & a_is_smaller & ~are_equal + ptr_b = labels_ptr + root_b + old_val_b = tl.atomic_min(ptr_b, root_a, mask=merge_mask_a_smaller) + + # A lane is done if its atomic op was successful (old value was what we expected) + success_b = old_val_b == root_b + final_root = tl.where(success_b & work_mask & ~done_mask, root_a, final_root) + done_mask |= success_b & merge_mask_a_smaller + + # *** Crucial Retry Logic *** + # If the update failed (old_val_b != root_b), another thread interfered. + # We update `current_b` to this new root (`old_val_b`) and will retry in the next loop iteration. + current_b = tl.where(success_b | ~merge_mask_a_smaller, current_b, old_val_b) + + # Case 2: root_b < root_a. Attempt to set parent[root_a] = root_b + merge_mask_b_smaller = work_mask & ~a_is_smaller & ~are_equal + ptr_a = labels_ptr + root_a + old_val_a = tl.atomic_min(ptr_a, root_b, mask=merge_mask_b_smaller) + + success_a = old_val_a == root_a + final_root = tl.where(success_a & work_mask & ~done_mask, root_b, final_root) + done_mask |= success_a & merge_mask_b_smaller + + # *** Crucial Retry Logic *** + # Similarly, update `current_a` if the atomic operation failed. + current_a = tl.where(success_a | ~merge_mask_b_smaller, current_a, old_val_a) + + return final_root + + +@triton.jit +def _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + dx: tl.constexpr, + dy: tl.constexpr, +): + # Helper functions to compute merge with a specific neighbor offset (dx, dy) + + neighbor_h = offsets_h + dy + neighbor_w = offsets_w + dx + # Proper bounds checking: all four bounds must be satisfied + mask_n = ( + mask_2d + & (neighbor_h[:, None] >= 0) + & (neighbor_h[:, None] < H) + & (neighbor_w[None, :] >= 0) + & (neighbor_w[None, :] < W) + ) + + offsets_neighbor = neighbor_h[:, None] * W + neighbor_w[None, :] + neighbor_values = tl.load( + input_ptr + base_offset + offsets_neighbor, mask=mask_n, other=-1 + ) + + mask_n = tl.ravel(mask_n) + neighbor_labels = tl.load( + labels_ptr + tl.ravel(base_offset + offsets_neighbor), mask=mask_n, other=-1 + ) + + to_merge = ( + mask_n & (neighbor_labels != -1) & tl.ravel(current_values == neighbor_values) + ) + valid_write = valid_current & to_merge + + # returns new parents for the pixels that were merged (otherwise keeps current labels) + return tl.where( + valid_write, + union(labels_ptr, current_labels, neighbor_labels, valid_write), + current_labels, + ) + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 16}, num_stages=1, num_warps=2 + ), + triton.Config( + {"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 32}, num_stages=2, num_warps=4 + ), + ], + key=["H", "W"], + restore_value=["labels_ptr"], +) +@triton.jit +def _local_prop_kernel( + labels_ptr, + input_ptr, + H: tl.constexpr, + W: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_W: tl.constexpr, +): + # This is the meat of the Phase 2 to do local merging + # It will be launched with a 2D grid: + # - dim 0: batch index + # - dim 1: block index over HxW image (2D tiling) + pid_b = tl.program_id(0) + pid_hw = tl.program_id(1) + + # Calculate offsets for the core block + offsets_h = (pid_hw // tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_H + tl.arange( + 0, BLOCK_SIZE_H + ) + offsets_w = (pid_hw % tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_W + tl.arange( + 0, BLOCK_SIZE_W + ) + + base_offset = pid_b * H * W + offsets_2d = offsets_h[:, None] * W + offsets_w[None, :] + mask_2d = (offsets_h[:, None] < H) & (offsets_w[None, :] < W) + mask_1d = tl.ravel(mask_2d) + + # Load the current labels for the block - these are parent pointers + current_labels = tl.load( + labels_ptr + tl.ravel(base_offset + offsets_2d), mask=mask_1d, other=-1 + ) + current_values = tl.load( + input_ptr + base_offset + offsets_2d, mask=mask_2d, other=-1 + ) + valid_current = mask_1d & (current_labels != -1) + + # Horizontal merge + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + 0, + ) + # Vertical merge + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + 0, + -1, + ) + # Diagonal merges + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + -1, + ) + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + 1, + ) + + # This actually does some path compression, in a lightweight but beneficial way + tl.atomic_min( + labels_ptr + tl.ravel(base_offset + offsets_2d), current_labels, mask=mask_1d + ) + + +# ============================================================================== +# ## Phase 3: Pointer Jumping Kernel +# ============================================================================== +# This kernel performs pointer jumping to ensure that all pixels point directly to their root labels. +# This is done in a loop until convergence. + + +@triton.jit +def _pointer_jump_kernel( + labels_in_ptr, labels_out_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + """ + Pointer jumping kernel with double buffering to avoid race conditions. + Reads from labels_in_ptr and writes to labels_out_ptr. + """ + # This kernel is launched with a 1D grid, and does not care about batching explicitly. + # By construction, the labels are global indices across the batch, and we never perform + # cross-batch merges, so this is safe. + + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load current labels from input buffer + current_labels = tl.load(labels_in_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (current_labels != -1) + + # A mask to track which lanes have successfully completed their union. + done_mask = ~valid_mask + while tl_any(~(done_mask | ~valid_mask)): + parent_labels = tl.load( + labels_in_ptr + current_labels, mask=valid_mask, other=-1 + ) + + are_equal = current_labels == parent_labels + done_mask |= are_equal & valid_mask + + current_labels = tl.where( + ~done_mask, tl.minimum(current_labels, parent_labels), current_labels + ) + + # Write to output buffer (safe because we're not reading from it) + tl.store(labels_out_ptr + offsets, current_labels, mask=mask) + + +# ============================================================================== +# ## Phase 4: Kernels for Computing Component Sizes +# ============================================================================== + + +# Step 4.1: Count occurrences of each root label using atomic adds. +@triton.jit +def _count_labels_kernel(labels_ptr, sizes_ptr, numel, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load the final, converged labels + labels = tl.load(labels_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (labels != -1) + + # Atomically increment the counter for each label. This builds a histogram. + tl.atomic_add(sizes_ptr + labels, 1, mask=valid_mask) + + +# Step 4.2: Broadcast the computed sizes back to the output tensor. +@triton.jit +def _broadcast_sizes_kernel( + labels_ptr, sizes_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load the final labels + labels = tl.load(labels_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (labels != -1) + + # Look up the size for each label from the histogram + component_sizes = tl.load(sizes_ptr + labels, mask=valid_mask, other=0) + + # Write the size to the final output tensor. Background pixels get size 0. + tl.store(out_ptr + offsets, component_sizes, mask=mask) + + +def connected_components_triton(input_tensor: torch.Tensor): + """ + Computes connected components labeling on a batch of 2D integer tensors using Triton. + + Args: + input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted + + Returns: + Tuple[torch.Tensor, int]: A tuple containing: + - A BxHxW output tensor with dense labels. Background is 0. + - A BxHxW tensor with the size of the connected component for each pixel. + """ + assert ( + input_tensor.is_cuda and input_tensor.is_contiguous() + ), "Input tensor must be a contiguous CUDA tensor." + out_shape = input_tensor.shape + if input_tensor.dim() == 4 and input_tensor.shape[1] == 1: + input_tensor = input_tensor.squeeze(1) + else: + assert ( + input_tensor.dim() == 3 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + B, H, W = input_tensor.shape + numel = B * H * W + device = input_tensor.device + + # --- Allocate Tensors --- + labels = torch.empty_like(input_tensor, dtype=torch.int32) + output = torch.empty_like(input_tensor, dtype=torch.int32) + + # --- Phase 1 --- + BLOCK_SIZE = 256 + grid_init = (triton.cdiv(numel, BLOCK_SIZE),) + _init_labels_kernel[grid_init]( + input_tensor, + labels, + numel, + BLOCK_SIZE=BLOCK_SIZE, + ) + + # --- Phase 2 --- + grid_local_prop = lambda meta: ( + B, + triton.cdiv(H, meta["BLOCK_SIZE_H"]) * triton.cdiv(W, meta["BLOCK_SIZE_W"]), + ) + _local_prop_kernel[grid_local_prop](labels, input_tensor, H, W) + + # --- Phase 3 --- + BLOCK_SIZE = 256 + grid_jump = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),) + _pointer_jump_kernel[grid_jump](labels, output, numel, BLOCK_SIZE=BLOCK_SIZE) + + # --- Phase 4 --- + # Allocate tensor to store the final output sizes + component_sizes_out = torch.empty_like(input_tensor, dtype=torch.int32) + + # Allocate a temporary 1D tensor to act as the histogram + # Size is numel because labels can be up to numel-1 + sizes_histogram = torch.zeros(numel, dtype=torch.int32, device=device) + + # 4.1: Count the occurrences of each label + grid_count = (triton.cdiv(numel, BLOCK_SIZE),) + _count_labels_kernel[grid_count]( + output, sizes_histogram, numel, BLOCK_SIZE=BLOCK_SIZE + ) + + # 2.2: Broadcast the counts to the final output tensor + grid_broadcast = (triton.cdiv(numel, BLOCK_SIZE),) + _broadcast_sizes_kernel[grid_broadcast]( + output, sizes_histogram, component_sizes_out, numel, BLOCK_SIZE=BLOCK_SIZE + ) + return output.view(out_shape) + 1, component_sizes_out.view(out_shape) diff --git a/src/nn/segearth_ov3/sam3/perflib/triton/nms.py b/src/nn/segearth_ov3/sam3/perflib/triton/nms.py new file mode 100644 index 0000000..ed800a1 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/perflib/triton/nms.py @@ -0,0 +1,124 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +# Adapted from https://github.com/stackav-oss/conch/blob/main/conch/kernels/vision/nms.py + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({"cxpr_block_size": 128}), + triton.Config({"cxpr_block_size": 256}), + triton.Config({"cxpr_block_size": 512}), + triton.Config({"cxpr_block_size": 1024}), + triton.Config({"cxpr_block_size": 2048}), + triton.Config({"cxpr_block_size": 4096}), + triton.Config({"cxpr_block_size": 8192}), + ], + key=["num_boxes"], +) +@triton.jit +def _nms_suppression_kernel( + # Tensors + iou_mask_ptr: tl.tensor, # [N, N] + keep_mask_ptr: tl.tensor, # [N] + # Scalars + num_boxes: tl.int32, + # Strides + iou_mask_stride: tl.int32, + # Constexprs + cxpr_block_size: tl.constexpr, +) -> None: + """NMS suppression kernel. + + Args: + iou_mask_ptr: Pointer to precomputed IoU mask, shape: (N, N). + keep_mask_ptr: Pointer to keep mask tensor, shape: (N,). + num_boxes: Number of boxes. + iou_mask_stride: Stride for IoU mask tensor. + cxpr_block_size: Block size for processing. + """ + # Sequential NMS: for each box in sorted order, suppress later boxes + for current_box_idx in range(num_boxes - 1): + # Check if current box is still kept + is_kept = tl.load(keep_mask_ptr + current_box_idx) + if is_kept: + # IoU mask row offset for the current box + # Because the IoU mask is sorted by score, we will only consider boxes that come after the current box. + # This means we only need to read the upper triangular part of the IoU mask. + iou_row_offset = current_box_idx * iou_mask_stride + + # Only process boxes that come after the current box + next_box_idx = current_box_idx + 1 + remaining_boxes = num_boxes - next_box_idx + + # Iterate blockwise through the columns + for block_idx in range(tl.cdiv(remaining_boxes, cxpr_block_size)): + # Masked load of indices for the target boxes in the current block + block_start = next_box_idx + block_idx * cxpr_block_size + target_box_offsets = block_start + tl.arange(0, cxpr_block_size) + target_box_mask = target_box_offsets < num_boxes + + # Suppress boxes with lower scores that have high IoU + suppression_mask = tl.load( + iou_mask_ptr + iou_row_offset + target_box_offsets, + mask=target_box_mask, + other=False, + ) + suppression_mask = tl.cast(suppression_mask, tl.int1) + + # Conditionally store suppression result for high-IoU boxes + tl.store( + keep_mask_ptr + target_box_offsets, False, mask=suppression_mask + ) + + # Potential race condition: we need to ensure all threads complete the store before the next + # iteration otherwise we may load stale data for whether or not a box has been suppressed. + tl.debug_barrier() + + +def nms_triton( + ious: torch.Tensor, + scores: torch.Tensor, + iou_threshold: float, +) -> torch.Tensor: + """Perform NMS given the iou matrix, the scores and the iou threshold + + Args: + ious: Pairwise IoU tensor of shape (N, N). + scores: Scores tensor of shape (N,). + iou_threshold: IoU threshold for suppression. + + Returns: + Tensor: Indices of kept boxes, sorted by decreasing score. + """ + assert scores.dim() == 1, "Scores must be 1D" + iou_mask = ious > iou_threshold + assert iou_mask.dim() == 2 + assert iou_mask.shape[0] == iou_mask.shape[1] == scores.shape[0] + assert iou_mask.device == scores.device + assert iou_mask.dtype == torch.bool + + num_boxes = scores.size(0) + keep_mask = torch.ones(len(scores), device=scores.device, dtype=torch.bool) + + # Sort boxes by scores in descending order + _, sorted_indices = torch.sort(scores, dim=0, stable=True, descending=True) + iou_mask = iou_mask[sorted_indices][:, sorted_indices].contiguous() + + # For the suppression stage, we need to process sequentially, but we'll still take + # advantage of parallelism by processing in blocks in one program. + stage2_grid = (1,) + _nms_suppression_kernel[stage2_grid]( + # Tensors + iou_mask_ptr=iou_mask, + keep_mask_ptr=keep_mask, + # Scalars + num_boxes=num_boxes, + # Strides + iou_mask_stride=iou_mask.stride(0), + ) + # Extract indices of kept boxes + return sorted_indices[keep_mask] diff --git a/src/nn/segearth_ov3/sam3/sam/__init__.py b/src/nn/segearth_ov3/sam3/sam/__init__.py new file mode 100644 index 0000000..8b35da6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/sam/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +from .mask_decoder import MaskDecoder +from .prompt_encoder import PromptEncoder +from .transformer import TwoWayTransformer diff --git a/src/nn/segearth_ov3/sam3/sam/common.py b/src/nn/segearth_ov3/sam3/sam/common.py new file mode 100644 index 0000000..b6d1587 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/sam/common.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Type + +import torch +import torch.nn as nn + + +class MLPBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + mlp_dim: int, + act: Type[nn.Module] = nn.GELU, + ) -> None: + super().__init__() + self.lin1 = nn.Linear(embedding_dim, mlp_dim) + self.lin2 = nn.Linear(mlp_dim, embedding_dim) + self.act = act() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lin2(self.act(self.lin1(x))) + + +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa +# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa +class LayerNorm2d(nn.Module): + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x diff --git a/src/nn/segearth_ov3/sam3/sam/mask_decoder.py b/src/nn/segearth_ov3/sam3/sam/mask_decoder.py new file mode 100644 index 0000000..b4ac397 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/sam/mask_decoder.py @@ -0,0 +1,319 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import List, Optional, Tuple, Type + +import torch +from torch import nn +from torch.nn import functional as F + +from .common import LayerNorm2d + + +class MaskDecoder(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + use_high_res_features: bool = False, + iou_prediction_use_sigmoid=False, + dynamic_multimask_via_stability=False, + dynamic_multimask_stability_delta=0.05, + dynamic_multimask_stability_thresh=0.98, + pred_obj_scores: bool = False, + pred_obj_scores_mlp: bool = False, + use_multimask_token_for_obj_ptr: bool = False, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.pred_obj_scores = pred_obj_scores + if self.pred_obj_scores: + self.obj_score_token = nn.Embedding(1, transformer_dim) + self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d( + transformer_dim, transformer_dim // 4, kernel_size=2, stride=2 + ), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d( + transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2 + ), + activation(), + ) + self.use_high_res_features = use_high_res_features + if use_high_res_features: + self.conv_s0 = nn.Conv2d( + transformer_dim, transformer_dim // 8, kernel_size=1, stride=1 + ) + self.conv_s1 = nn.Conv2d( + transformer_dim, transformer_dim // 4, kernel_size=1, stride=1 + ) + + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, + iou_head_hidden_dim, + self.num_mask_tokens, + iou_head_depth, + sigmoid_output=iou_prediction_use_sigmoid, + ) + if self.pred_obj_scores: + self.pred_obj_score_head = nn.Linear(transformer_dim, 1) + if pred_obj_scores_mlp: + self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3) + + # When outputting a single mask, optionally we can dynamically fall back to the best + # multimask output token if the single mask output token gives low stability scores. + self.dynamic_multimask_via_stability = dynamic_multimask_via_stability + self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta + self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + repeat_image: bool, + high_res_features: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + torch.Tensor: batched SAM token for mask output + """ + masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + repeat_image=repeat_image, + high_res_features=high_res_features, + ) + + # Select the correct mask or masks for output + if multimask_output: + masks = masks[:, 1:, :, :] + iou_pred = iou_pred[:, 1:] + elif self.dynamic_multimask_via_stability and not self.training: + masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred) + else: + masks = masks[:, 0:1, :, :] + iou_pred = iou_pred[:, 0:1] + + if multimask_output and self.use_multimask_token_for_obj_ptr: + sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape + else: + # Take the mask output token. Here we *always* use the token for single mask output. + # At test time, even if we track after 1-click (and using multimask_output=True), + # we still take the single mask token here. The rationale is that we always track + # after multiple clicks during training, so the past tokens seen during training + # are always the single mask token (and we'll let it be the object-memory token). + sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape + + # Prepare output + return masks, iou_pred, sam_tokens_out, object_score_logits + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + repeat_image: bool, + high_res_features: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + s = 0 + if self.pred_obj_scores: + output_tokens = torch.cat( + [ + self.obj_score_token.weight, + self.iou_token.weight, + self.mask_tokens.weight, + ], + dim=0, + ) + s = 1 + else: + output_tokens = torch.cat( + [self.iou_token.weight, self.mask_tokens.weight], dim=0 + ) + output_tokens = output_tokens.unsqueeze(0).expand( + sparse_prompt_embeddings.size(0), -1, -1 + ) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + if repeat_image: + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + else: + assert image_embeddings.shape[0] == tokens.shape[0] + src = image_embeddings + src = src + dense_prompt_embeddings + assert ( + image_pe.size(0) == 1 + ), "image_pe should have size 1 in batch dim (from `get_dense_pe()`)" + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, s, :] + mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + if not self.use_high_res_features: + upscaled_embedding = self.output_upscaling(src) + else: + dc1, ln1, act1, dc2, act2 = self.output_upscaling + feat_s0, feat_s1 = high_res_features + upscaled_embedding = act1(ln1(dc1(src) + feat_s1)) + upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0) + + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + hyper_in_list.append( + self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]) + ) + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding.shape + masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) + + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + if self.pred_obj_scores: + assert s == 1 + object_score_logits = self.pred_obj_score_head(hs[:, 0, :]) + else: + # Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1 + object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1) + + return masks, iou_pred, mask_tokens_out, object_score_logits + + def _get_stability_scores(self, mask_logits): + """ + Compute stability scores of the mask logits based on the IoU between upper and + lower thresholds. + """ + mask_logits = mask_logits.flatten(-2) + stability_delta = self.dynamic_multimask_stability_delta + area_i = torch.sum(mask_logits > stability_delta, dim=-1).float() + area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float() + stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0) + return stability_scores + + def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores): + """ + When outputting a single mask, if the stability score from the current single-mask + output (based on output token 0) falls below a threshold, we instead select from + multi-mask outputs (based on output token 1~3) the mask with the highest predicted + IoU score. This is intended to ensure a valid mask for both clicking and tracking. + """ + # The best mask from multimask output tokens (1~3) + multimask_logits = all_mask_logits[:, 1:, :, :] + multimask_iou_scores = all_iou_scores[:, 1:] + best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1) + batch_inds = torch.arange( + multimask_iou_scores.size(0), device=all_iou_scores.device + ) + best_multimask_logits = multimask_logits[batch_inds, best_scores_inds] + best_multimask_logits = best_multimask_logits.unsqueeze(1) + best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds] + best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1) + + # The mask from singlemask output token 0 and its stability score + singlemask_logits = all_mask_logits[:, 0:1, :, :] + singlemask_iou_scores = all_iou_scores[:, 0:1] + stability_scores = self._get_stability_scores(singlemask_logits) + is_stable = stability_scores >= self.dynamic_multimask_stability_thresh + + # Dynamically fall back to best multimask output upon low stability scores. + mask_logits_out = torch.where( + is_stable[..., None, None].expand_as(singlemask_logits), + singlemask_logits, + best_multimask_logits, + ) + iou_scores_out = torch.where( + is_stable.expand_as(singlemask_iou_scores), + singlemask_iou_scores, + best_multimask_iou_scores, + ) + return mask_logits_out, iou_scores_out + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x diff --git a/src/nn/segearth_ov3/sam3/sam/prompt_encoder.py b/src/nn/segearth_ov3/sam3/sam/prompt_encoder.py new file mode 100644 index 0000000..145ea9f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/sam/prompt_encoder.py @@ -0,0 +1,243 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Any, Optional, Tuple, Type + +import numpy as np +import torch +from torch import nn + +from .common import LayerNorm2d + + +class PromptEncoder(nn.Module): + def __init__( + self, + embed_dim: int, + image_embedding_size: Tuple[int, int], + input_image_size: Tuple[int, int], + mask_in_chans: int, + activation: Type[nn.Module] = nn.GELU, + ) -> None: + """ + Encodes prompts for input to SAM's mask decoder. + + Arguments: + embed_dim (int): The prompts' embedding dimension + image_embedding_size (tuple(int, int)): The spatial size of the + image embedding, as (H, W). + input_image_size (int): The padded size of the image as input + to the image encoder, as (H, W). + mask_in_chans (int): The number of hidden channels used for + encoding input masks. + activation (nn.Module): The activation to use when encoding + input masks. + """ + super().__init__() + self.embed_dim = embed_dim + self.input_image_size = input_image_size + self.image_embedding_size = image_embedding_size + self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) + + self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners + point_embeddings = [ + nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings) + ] + self.point_embeddings = nn.ModuleList(point_embeddings) + self.not_a_point_embed = nn.Embedding(1, embed_dim) + + self.mask_input_size = ( + 4 * image_embedding_size[0], + 4 * image_embedding_size[1], + ) + self.mask_downscaling = nn.Sequential( + nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans // 4), + activation(), + nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans), + activation(), + nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), + ) + self.no_mask_embed = nn.Embedding(1, embed_dim) + + def get_dense_pe(self) -> torch.Tensor: + """ + Returns the positional encoding used to encode point prompts, + applied to a dense set of points the shape of the image encoding. + + Returns: + torch.Tensor: Positional encoding with shape + 1x(embed_dim)x(embedding_h)x(embedding_w) + """ + return self.pe_layer(self.image_embedding_size).unsqueeze(0) + + def _embed_points( + self, + points: torch.Tensor, + labels: torch.Tensor, + pad: bool, + ) -> torch.Tensor: + """Embeds point prompts.""" + points = points + 0.5 # Shift to center of pixel + if pad: + padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device) + padding_label = -torch.ones((labels.shape[0], 1), device=labels.device) + points = torch.cat([points, padding_point], dim=1) + labels = torch.cat([labels, padding_label], dim=1) + point_embedding = self.pe_layer.forward_with_coords( + points, self.input_image_size + ) + + point_embedding = torch.where( + (labels == -1).unsqueeze(-1), + torch.zeros_like(point_embedding) + self.not_a_point_embed.weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 0).unsqueeze(-1), + point_embedding + self.point_embeddings[0].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 1).unsqueeze(-1), + point_embedding + self.point_embeddings[1].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 2).unsqueeze(-1), + point_embedding + self.point_embeddings[2].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 3).unsqueeze(-1), + point_embedding + self.point_embeddings[3].weight, + point_embedding, + ) + return point_embedding + + def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: + """Embeds box prompts.""" + boxes = boxes + 0.5 # Shift to center of pixel + coords = boxes.reshape(-1, 2, 2) + corner_embedding = self.pe_layer.forward_with_coords( + coords, self.input_image_size + ) + corner_embedding[:, 0, :] += self.point_embeddings[2].weight + corner_embedding[:, 1, :] += self.point_embeddings[3].weight + return corner_embedding + + def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor: + """Embeds mask inputs.""" + mask_embedding = self.mask_downscaling(masks) + return mask_embedding + + def _get_batch_size( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> int: + """ + Gets the batch size of the output given the batch size of the input prompts. + """ + if points is not None: + return points[0].shape[0] + elif boxes is not None: + return boxes.shape[0] + elif masks is not None: + return masks.shape[0] + else: + return 1 + + def _get_device(self) -> torch.device: + return self.point_embeddings[0].weight.device + + def forward( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Embeds different types of prompts, returning both sparse and dense + embeddings. + + Arguments: + points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates + and labels to embed. + boxes (torch.Tensor or none): boxes to embed + masks (torch.Tensor or none): masks to embed + + Returns: + torch.Tensor: sparse embeddings for the points and boxes, with shape + BxNx(embed_dim), where N is determined by the number of input points + and boxes. + torch.Tensor: dense embeddings for the masks, in the shape + Bx(embed_dim)x(embed_H)x(embed_W) + """ + bs = self._get_batch_size(points, boxes, masks) + sparse_embeddings = torch.empty( + (bs, 0, self.embed_dim), device=self._get_device() + ) + if points is not None: + coords, labels = points + point_embeddings = self._embed_points(coords, labels, pad=(boxes is None)) + sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1) + if boxes is not None: + box_embeddings = self._embed_boxes(boxes) + sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1) + + if masks is not None: + dense_embeddings = self._embed_masks(masks) + else: + dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( + bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] + ) + + return sparse_embeddings, dense_embeddings + + +class PositionEmbeddingRandom(nn.Module): + """ + Positional encoding using random spatial frequencies. + """ + + def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: + super().__init__() + if scale is None or scale <= 0.0: + scale = 1.0 + self.register_buffer( + "positional_encoding_gaussian_matrix", + scale * torch.randn((2, num_pos_feats)), + ) + + def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: + """Positionally encode points that are normalized to [0,1].""" + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coords = 2 * coords - 1 + coords = coords @ self.positional_encoding_gaussian_matrix + coords = 2 * np.pi * coords + # outputs d_1 x ... x d_n x C shape + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + def forward(self, size: Tuple[int, int]) -> torch.Tensor: + """Generate positional encoding for a grid of the specified size.""" + h, w = size + device: Any = self.positional_encoding_gaussian_matrix.device + grid = torch.ones((h, w), device=device, dtype=torch.float32) + y_embed = grid.cumsum(dim=0) - 0.5 + x_embed = grid.cumsum(dim=1) - 0.5 + y_embed = y_embed / h + x_embed = x_embed / w + + pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) + return pe.permute(2, 0, 1) # C x H x W + + def forward_with_coords( + self, coords_input: torch.Tensor, image_size: Tuple[int, int] + ) -> torch.Tensor: + """Positionally encode points that are not normalized to [0,1].""" + coords = coords_input.clone() + coords[:, :, 0] = coords[:, :, 0] / image_size[1] + coords[:, :, 1] = coords[:, :, 1] / image_size[0] + return self._pe_encoding(coords.to(torch.float)) # B x N x C diff --git a/src/nn/segearth_ov3/sam3/sam/rope.py b/src/nn/segearth_ov3/sam3/sam/rope.py new file mode 100644 index 0000000..2db01b6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/sam/rope.py @@ -0,0 +1,161 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Adapted from: +1. https://github.com/meta-llama/codellama/blob/main/llama/model.py +2. https://github.com/naver-ai/rope-vit +3. https://github.com/lucidrains/rotary-embedding-torch +""" + +from typing import Optional + +import torch +from einops import rearrange, repeat +from torch import broadcast_tensors, nn + + +def init_t_xy(end_x: int, end_y: int, scale: float = 1.0, offset: int = 0, device=None): + t = torch.arange(end_x * end_y, dtype=torch.float32, device=device) + t_x = (t % end_x).float() + t_y = torch.div(t, end_x, rounding_mode="floor").float() + return t_x * scale + offset, t_y * scale + offset + + +def compute_axial_cis( + dim: int, + end_x: int, + end_y: int, + theta: float = 10000.0, + scale_pos: float = 1.0, + offset: int = 0, + device=None, +): + freqs_x = 1.0 / ( + theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim) + ) + freqs_y = 1.0 / ( + theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim) + ) + + t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset, device=device) + freqs_x = torch.outer(t_x, freqs_x) + freqs_y = torch.outer(t_y, freqs_y) + freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) + freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) + return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) + + +def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): + ndim = x.ndim + assert 0 <= 1 < ndim + assert freqs_cis.shape == (x.shape[-2], x.shape[-1]) + shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def apply_rotary_enc( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: torch.Tensor, + repeat_freqs_k: bool = False, +): + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) + xk_ = ( + torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) + if xk.shape[-2] != 0 + else None + ) + freqs_cis = reshape_for_broadcast(freqs_cis, xq_) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) + if xk_ is None: + # no keys to rotate, due to dropout + return xq_out.type_as(xq).to(xq.device), xk + # repeat freqs along seq_len dim to match k seq_len + if repeat_freqs_k: + r = xk_.shape[-2] // xq_.shape[-2] + freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) + return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device) + + +def complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag): + # Compute the real part of the product + real_part = xq_real * freqs_cis_real - xq_imag * freqs_cis_imag + # Compute the imaginary part of the product + imag_part = xq_real * freqs_cis_imag + xq_imag * freqs_cis_real + # Stack the real and imaginary parts along the last dimension + return torch.stack([real_part, imag_part], dim=-1) + + +def apply_rotary_enc_real( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis_real: torch.Tensor, + freqs_cis_imag: torch.Tensor, + repeat_freqs_k: bool = False, +): + assert xk is not None + assert xk.shape[-2] != 0 + + xq_real = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 0] + xq_imag = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 1] + xk_real = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 0] + xk_imag = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 1] + freqs_cis_real = reshape_for_broadcast(freqs_cis_real, xq_real) + freqs_cis_imag = reshape_for_broadcast(freqs_cis_imag, xq_imag) + xq_out = complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag).flatten(3) + if repeat_freqs_k: + r = xk_real.shape[-2] // xq_real.shape[-2] + freqs_cis_real = freqs_cis_real.repeat(*([1] * (freqs_cis_real.ndim - 2)), r, 1) + freqs_cis_imag = freqs_cis_imag.repeat(*([1] * (freqs_cis_imag.ndim - 2)), r, 1) + xk_out = complex_mult(xk_real, xk_imag, freqs_cis_real, freqs_cis_imag).flatten(3) + # xq_out = torch.view_as_real(torch.complex(xq_real, xq_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3) + # xk_out = torch.view_as_real(torch.compelx(xk_real, xk_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3) + return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device) + + +# rotary embedding helper functions +def broadcat(tensors, dim=-1): + broadcasted_tensors = broadcast_tensors(*tensors) + return torch.cat(broadcasted_tensors, dim=dim) + + +def rotate_half(x: torch.Tensor): + x = rearrange(x, "... (d r) -> ... d r", r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, "... d r -> ... (d r)") + + +class VisionRotaryEmbeddingVE(nn.Module): + def __init__( + self, + dim: int, + seq_len: int, + pt_seq_len: Optional[int] = None, + theta: float = 10000.0, + offset: int = 1, # specific to VE + ): + super().__init__() + + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + scale = 1.0 + if pt_seq_len is not None: + scale = pt_seq_len / seq_len + + # offset of +1 following VE - even though for the + # attention op only differences matter + t = torch.arange(seq_len) * scale + offset + + freqs = torch.einsum("..., f -> ... f", t, freqs) + freqs = repeat(freqs, "... n -> ... (n r)", r=2) + + freqs = broadcat((freqs[None, :, :], freqs[:, None, :]), dim=-1) + freqs_cos = freqs.cos().view(-1, freqs.shape[-1]) + freqs_sin = freqs.sin().view(-1, freqs.shape[-1]) + + self.register_buffer("freqs_cos", freqs_cos) + self.register_buffer("freqs_sin", freqs_sin) + + def forward(self, t: torch.Tensor): + return t * self.freqs_cos + rotate_half(t) * self.freqs_sin diff --git a/src/nn/segearth_ov3/sam3/sam/transformer.py b/src/nn/segearth_ov3/sam3/sam/transformer.py new file mode 100644 index 0000000..3e96c28 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/sam/transformer.py @@ -0,0 +1,358 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math +from functools import partial +from typing import Tuple, Type + +import torch +import torch.nn.functional as F + +from sam3.sam.rope import apply_rotary_enc, apply_rotary_enc_real, compute_axial_cis +from torch import nn, Tensor + +from .common import MLPBlock + + +class TwoWayTransformer(nn.Module): + def __init__( + self, + depth: int, + embedding_dim: int, + num_heads: int, + mlp_dim: int, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + ) -> None: + """ + A transformer decoder that attends to an input image using + queries whose positional embedding is supplied. + + Args: + depth (int): number of layers in the transformer + embedding_dim (int): the channel dimension for the input embeddings + num_heads (int): the number of heads for multihead attention. Must + divide embedding_dim + mlp_dim (int): the channel dimension internal to the MLP block + activation (nn.Module): the activation to use in the MLP block + """ + super().__init__() + self.depth = depth + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.layers = nn.ModuleList() + + for i in range(depth): + self.layers.append( + TwoWayAttentionBlock( + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_dim=mlp_dim, + activation=activation, + attention_downsample_rate=attention_downsample_rate, + skip_first_layer_pe=(i == 0), + ) + ) + + self.final_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm_final_attn = nn.LayerNorm(embedding_dim) + + def forward( + self, + image_embedding: Tensor, + image_pe: Tensor, + point_embedding: Tensor, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + image_embedding (torch.Tensor): image to attend to. Should be shape + B x embedding_dim x h x w for any h and w. + image_pe (torch.Tensor): the positional encoding to add to the image. Must + have the same shape as image_embedding. + point_embedding (torch.Tensor): the embedding to add to the query points. + Must have shape B x N_points x embedding_dim for any N_points. + + Returns: + torch.Tensor: the processed point_embedding + torch.Tensor: the processed image_embedding + """ + # BxCxHxW -> BxHWxC == B x N_image_tokens x C + bs, c, h, w = image_embedding.shape + image_embedding = image_embedding.flatten(2).permute(0, 2, 1) + image_pe = image_pe.flatten(2).permute(0, 2, 1) + + # Prepare queries + queries = point_embedding + keys = image_embedding + + # Apply transformer blocks and final layernorm + for layer in self.layers: + queries, keys = layer( + queries=queries, + keys=keys, + query_pe=point_embedding, + key_pe=image_pe, + ) + + # Apply the final attention layer from the points to the image + q = queries + point_embedding + k = keys + image_pe + attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm_final_attn(queries) + + return queries, keys + + +class TwoWayAttentionBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + num_heads: int, + mlp_dim: int = 2048, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + skip_first_layer_pe: bool = False, + ) -> None: + """ + A transformer block with four layers: (1) self-attention of sparse + inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp + block on sparse inputs, and (4) cross attention of dense inputs to sparse + inputs. + + Arguments: + embedding_dim (int): the channel dimension of the embeddings + num_heads (int): the number of heads in the attention layers + mlp_dim (int): the hidden dimension of the mlp block + activation (nn.Module): the activation of the mlp block + skip_first_layer_pe (bool): skip the PE on the first layer + """ + super().__init__() + self.self_attn = Attention(embedding_dim, num_heads) + self.norm1 = nn.LayerNorm(embedding_dim) + + self.cross_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm2 = nn.LayerNorm(embedding_dim) + + self.mlp = MLPBlock(embedding_dim, mlp_dim, activation) + self.norm3 = nn.LayerNorm(embedding_dim) + + self.norm4 = nn.LayerNorm(embedding_dim) + self.cross_attn_image_to_token = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + + self.skip_first_layer_pe = skip_first_layer_pe + + def forward( + self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor + ) -> Tuple[Tensor, Tensor]: + # Self attention block + if self.skip_first_layer_pe: + queries = self.self_attn(q=queries, k=queries, v=queries) + else: + q = queries + query_pe + attn_out = self.self_attn(q=q, k=q, v=queries) + queries = queries + attn_out + queries = self.norm1(queries) + + # Cross attention block, tokens attending to image embedding + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm2(queries) + + # MLP block + mlp_out = self.mlp(queries) + queries = queries + mlp_out + queries = self.norm3(queries) + + # Cross attention block, image embedding attending to tokens + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) + keys = keys + attn_out + keys = self.norm4(keys) + + return queries, keys + + +class Attention(nn.Module): + """ + An attention layer that allows for downscaling the size of the embedding + after projection to queries, keys, and values. + """ + + def __init__( + self, + embedding_dim: int, + num_heads: int, + downsample_rate: int = 1, + dropout: float = 0.0, + kv_in_dim: int = None, + use_fa3: bool = False, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim + self.internal_dim = embedding_dim // downsample_rate + self.num_heads = num_heads + self.use_fa3 = use_fa3 + assert ( + self.internal_dim % num_heads == 0 + ), "num_heads must divide embedding_dim." + + self.q_proj = nn.Linear(embedding_dim, self.internal_dim) + self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim) + self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim) + self.out_proj = nn.Linear(self.internal_dim, embedding_dim) + + self.dropout_p = dropout + + def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: + b, n, c = x.shape + x = x.reshape(b, n, num_heads, c // num_heads) + return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head + + def _recombine_heads(self, x: Tensor) -> Tensor: + b, n_heads, n_tokens, c_per_head = x.shape + x = x.transpose(1, 2) + return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + dropout_p = self.dropout_p if self.training else 0.0 + # Attention + # with torch.backends.cuda.sdp_kernel( + # enable_flash=USE_FLASH_ATTN, + # # if Flash attention kernel is off, then math kernel needs to be enabled + # enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON, + # enable_mem_efficient=OLD_GPU, + # ): + # Let's trust the dispatcher.... + if self.use_fa3: + from sam3.perflib.fa3 import flash_attn_func + + assert dropout_p == 0.0 + out = flash_attn_func( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + else: + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out + + +class RoPEAttention(Attention): + """Attention with rotary position encoding.""" + + def __init__( + self, + *args, + rope_theta=10000.0, + # whether to repeat q rope to match k length + # this is needed for cross-attention to memories + rope_k_repeat=False, + feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution + use_rope_real=False, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.use_rope_real = use_rope_real + self.compute_cis = partial( + compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta + ) + device = torch.device("cuda") if torch.cuda.is_available() else None + self.freqs_cis = self.compute_cis( + end_x=feat_sizes[0], end_y=feat_sizes[1], device=device + ) + if self.use_rope_real: + self.freqs_cis_real = self.freqs_cis.real + self.freqs_cis_imag = self.freqs_cis.imag + self.rope_k_repeat = rope_k_repeat + + def forward( + self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0 + ) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + # Apply rotary position encoding + w = h = math.sqrt(q.shape[-2]) + if self.freqs_cis.shape[0] != q.shape[-2]: + self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device) + self.freqs_cis_real = self.freqs_cis.real + self.freqs_cis_imag = self.freqs_cis.imag + if q.shape[-2] != k.shape[-2]: + assert self.rope_k_repeat + + num_k_rope = k.size(-2) - num_k_exclude_rope + if self.use_rope_real: + q, k[:, :, :num_k_rope] = apply_rotary_enc_real( + q, + k[:, :, :num_k_rope], + freqs_cis_real=self.freqs_cis_real, + freqs_cis_imag=self.freqs_cis_imag, + repeat_freqs_k=self.rope_k_repeat, + ) + else: + q, k[:, :, :num_k_rope] = apply_rotary_enc( + q, + k[:, :, :num_k_rope], + self.freqs_cis, + repeat_freqs_k=self.rope_k_repeat, + ) + + dropout_p = self.dropout_p if self.training else 0.0 + # Attention + # with torch.backends.cuda.sdp_kernel( + # enable_flash=USE_FLASH_ATTN, + # # if Flash attention kernel is off, then math kernel needs to be enabled + # enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON, + # enable_mem_efficient=OLD_GPU, + # ): + # Let's trust the dispatcher.... + if self.use_fa3: + from sam3.perflib.fa3 import flash_attn_func + + assert dropout_p == 0.0 + out = flash_attn_func( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + else: + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out diff --git a/src/nn/segearth_ov3/sam3/train/__init__.py b/src/nn/segearth_ov3/sam3/train/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/train/configs/eval_base.yaml b/src/nn/segearth_ov3/sam3/train/configs/eval_base.yaml new file mode 100644 index 0000000..20890d8 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/eval_base.yaml @@ -0,0 +1,279 @@ +# @package _global_ +defaults: + - _self_ + +# This config is the base configuration for all evaluations. Amongst other things, it defines: +# - the model +# - the image transforms +# - the post processors +# - cluster configuration (only relevant for slurm-based evals, ignored otherwise) +# +# Most of the parameters should be kept as-is. The main modifications you may want to make are: +# - the cluster configuration, to adjust partitions/qos to your system +# - the flag gather_pred_via_filesys if you ram is tight +# - num_val_workers if your number of cores is small (should be roughly number of cores / number of gpus) +# - the paths below + + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + # If you leave the checkpoint path to null, the model will be downloaded from hugging-face. Otherwise provide a path + checkpoint_path: null + # the experiments will be subfolders of this + base_experiment_log_dir: + + # base path to the annotation folder for gold (refer to the readmes on how to download) + base_annotation_path: + + # base path to the annotation folder for silver (refer to the readmes on how to download) + base_annotation_path_silver: + + # path to the metaclip images, used for SA-Co gold (refer to the readme for instructions). Can be null if you don't intend on evaluating on this dataset. + metaclip_img_path: + + # path to the sa1b images, used for SA-Co gold (refer to the readme for instructions). Can be null if you don't intend on evaluating on this dataset. + sa1b_img_path: + + # path to the SA-Co/silver images + silver_img_path: + + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + + use_presence_eval: True + + base_val_transform: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + ######## transforms for validation (begin) ######## + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: False + ######## transforms for validation (end) ######## + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + loss: null + + # Model parameters + d_model: 256 + input_box_embedding_dim: ${add:${scratch.d_model},2} + + # Box processing + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 #infinite detections + use_original_ids: false + use_original_sizes_box: false + use_presence: ${scratch.use_presence_eval} + + box_postprocessor_thresholded: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 #infinite detections + use_original_ids: false + use_original_sizes_box: false + detection_threshold: 0.3 + use_presence: ${scratch.use_presence_eval} + + mask_postprocessor_thresholded: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 #infinite detections + iou_type: "segm" + use_original_ids: false + use_original_sizes_box: false + use_original_sizes_mask: true + convert_mask_to_rle: True + detection_threshold: 0.3 + use_presence: ${scratch.use_presence_eval} + + # Image processing parameters + resolution: 1008 + max_ann_per_img: 200 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + train_batch_size: 1 + val_batch_size: 1 + num_train_workers: 0 + num_val_workers: 10 # change this depending on the number of cpu cores available + max_data_epochs: 20 + target_epoch_size: 1500 + hybrid_repeats: 1 + context_length: 2 + + # All reduce - this controls how the predictions are sent back to node 0. + # If you have a lot of ram, CPU gather is faster. Otherwise, we provide a fallback through filesystem (eg NFS) + # Switch to true if you get cpu ooms during gather. + gather_pred_via_filesys: false + + # Learning rate and scheduler parameters (unused for eval) + lr_scale: 0.1 + lr_transformer: ${times:8e-4,${scratch.lr_scale}} + lr_vision_backbone: ${times:2.5e-4,${scratch.lr_scale}} + lr_language_backbone: ${times:5e-5,${scratch.lr_scale}} + lrd_vision_backbone: 0.9 # (lower for in-domain adn higher for ood) + wd: 0.1 + scheduler_timescale: 20 + scheduler_warmup: 20 + scheduler_cooldown: 20 + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: null + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: true + enable_segmentation: true # Warning: Enable this if using segmentation. + checkpoint_path: ${paths.checkpoint_path} + + meters: + val: null + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + optimizer: + _target_: torch.optim.AdamW + + gradient_clip: + _target_: sam3.train.optim.optimizer.GradientClipper + max_norm: 0.1 + norm_type: 2 + + param_group_modifiers: + - _target_: sam3.train.optim.optimizer.layer_decay_param_modifier + _partial_: True + layer_decay_value: ${scratch.lrd_vision_backbone} + apply_to: 'backbone.vision_backbone.trunk' + overrides: + - pattern: '*pos_embed*' + value: 1.0 + + options: + lr: + - scheduler: # transformer and class_embed + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_transformer} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_vision_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.vision_backbone.*' + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_language_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.language_backbone.*' + + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: ${scratch.wd} + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.0 + param_names: + - '*bias*' + module_cls_names: ['torch.nn.LayerNorm'] + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 4 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + + +submitit: + account: null # Add your SLURM account if use_cluster == 1 + partition: null + qos: null # Add your QoS if use_cluster == 1 + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_attributes.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_attributes.yaml new file mode 100644 index 0000000..8646b69 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_attributes.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_attributes/ + coco_gt: ${paths.base_annotation_path}/gold_attributes_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_attributes_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_attributes_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_attributes_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.metaclip_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_attributes + + meters: + val: + gold_attributes: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_attributes + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_crowded.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_crowded.yaml new file mode 100644 index 0000000..fef74a6 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_crowded.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_crowded/ + coco_gt: ${paths.base_annotation_path}/gold_crowded_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_crowded_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_crowded_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_crowded_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.metaclip_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_crowded + + meters: + val: + gold_crowded: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_crowded + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml new file mode 100644 index 0000000..b08c4a4 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_fg_food/ + coco_gt: ${paths.base_annotation_path}/gold_fg_food_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_fg_food_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_fg_food_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_fg_food_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.metaclip_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_fg_food + + meters: + val: + gold_fg_food: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_fg_food + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml new file mode 100644 index 0000000..89a93be --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_fg_sports_equipment/ + coco_gt: ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.metaclip_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_fg_sports_equipment + + meters: + val: + gold_fg_sports_equipment: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_fg_sports_equipment + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml new file mode 100644 index 0000000..e9c276f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_metaclip_nps/ + coco_gt: ${paths.base_annotation_path}/gold_metaclip_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_metaclip_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_metaclip_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_metaclip_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.metaclip_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_metaclip_nps + + meters: + val: + gold_metaclip_nps: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_metaclip_nps + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml new file mode 100644 index 0000000..52c87ee --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_sa1b_nps/ + coco_gt: ${paths.base_annotation_path}/gold_sa1b_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_sa1b_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_sa1b_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_sa1b_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.sa1b_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_sa1b_nps + + meters: + val: + gold_sa1b_nps: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_sa1b_nps + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml new file mode 100644 index 0000000..6304954 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml @@ -0,0 +1,66 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/gold_wiki_common/ + coco_gt: ${paths.base_annotation_path}/gold_wiki_common_merged_a_release_test.json + coco_gts: + - ${paths.base_annotation_path}/gold_wiki_common_merged_a_release_test.json + - ${paths.base_annotation_path}/gold_wiki_common_merged_b_release_test.json + - ${paths.base_annotation_path}/gold_wiki_common_merged_c_release_test.json + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.metaclip_img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: gold_wiki_common + + meters: + val: + gold_wiki_common: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/gold_wiki_common + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gts} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_and_visual.yaml b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_and_visual.yaml new file mode 100644 index 0000000..51e93b4 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_and_visual.yaml @@ -0,0 +1,255 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} + +paths: + odinw_data_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + +supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} +# Validation transforms pipeline +val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + - _target_: sam3.train.transforms.filter_query_transforms.TextQueryToVisual + keep_text_queries: true # Note: set this to false if you only want visual + probability: 1.0 # always + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: True + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + # Image processing parameters + resolution: 1008 + # Normalization parameters + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + val_batch_size: 2 + num_val_workers: 0 + gather_pred_via_filesys: false + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + max_epochs: 1 + accelerator: cuda + seed_value: 123 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + prompts: ${odinw35_prompts.${supercategory_tuple.name}} + include_negatives: true + category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! + _partial_: true + img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} + ann_file: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + transforms: ${val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: 1 + dict_key: odinw35 + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: true # Set to false if training + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + odinw35: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${supercategory_tuple.name} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + tide: False + iou_type: "bbox" + positive_split: true + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${supercategory_tuple.name} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + + job_array: + num_tasks: 13 + task_index: 0 + +# ============================================================================ +# ODinW13 Supercategories +# ============================================================================ + +all_odinw_supercategories: + - name: AerialMaritimeDrone_large + val: + img_folder: AerialMaritimeDrone/large/test/ + json: AerialMaritimeDrone/large/test/annotations_without_background.json + - name: Aquarium + val: + img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ + json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json + - name: CottontailRabbits + val: + img_folder: CottontailRabbits/test/ + json: CottontailRabbits/test/annotations_without_background.json + - name: EgoHands_generic + val: + img_folder: EgoHands/generic/test/ + json: EgoHands/generic/test/annotations_without_background.json + - name: NorthAmericaMushrooms + val: + img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ + json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json + - name: Packages + val: + img_folder: Packages/Raw/test/ + json: Packages/Raw/test/annotations_without_background.json + - name: PascalVOC + val: + img_folder: PascalVOC/valid/ + json: PascalVOC/valid/annotations_without_background.json + - name: Raccoon + val: + img_folder: Raccoon/Raccoon.v2-raw.coco/test/ + json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json + - name: ShellfishOpenImages + val: + img_folder: ShellfishOpenImages/raw/test/ + json: ShellfishOpenImages/raw/test/annotations_without_background.json + - name: VehiclesOpenImages + val: + img_folder: VehiclesOpenImages/416x416/test/ + json: VehiclesOpenImages/416x416/test/annotations_without_background.json + - name: pistols + val: + img_folder: pistols/export/ + json: pistols/export/test_annotations_without_background.json + - name: pothole + val: + img_folder: pothole/test/ + json: pothole/test/annotations_without_background.json + - name: thermalDogsAndPeople + val: + img_folder: thermalDogsAndPeople/test/ + json: thermalDogsAndPeople/test/annotations_without_background.json + + +odinw35_prompts: + AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, + {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", + "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, + {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' + Aquarium: null + CottontailRabbits: null + EgoHands_generic: null + NorthAmericaMushrooms: '[{''id'': 1, ''name'': + ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' + Packages: null + PascalVOC: null + Raccoon: null + ShellfishOpenImages: null + VehiclesOpenImages: null + pistols: null + pothole: null + thermalDogsAndPeople: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only.yaml b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only.yaml new file mode 100644 index 0000000..e28fa5d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only.yaml @@ -0,0 +1,253 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} + +paths: + odinw_data_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + + +supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} +# Validation transforms pipeline +val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: True + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + # Image processing parameters + resolution: 1008 + # Normalization parameters + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + val_batch_size: 2 + num_val_workers: 0 + gather_pred_via_filesys: false + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + max_epochs: 1 + accelerator: cuda + seed_value: 123 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + prompts: ${odinw35_prompts.${supercategory_tuple.name}} + include_negatives: true + category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! + _partial_: true + img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} + ann_file: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + transforms: ${val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: 1 + dict_key: odinw35 + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: true # Set to false if training + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + odinw35: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/odinw/${supercategory_tuple.name} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + tide: False + iou_type: "bbox" + positive_split: False + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${supercategory_tuple.name} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + + job_array: + num_tasks: 13 + task_index: 0 + +# ============================================================================ +# ODinW13 Supercategories +# ============================================================================ + +all_odinw_supercategories: + - name: AerialMaritimeDrone_large + val: + img_folder: AerialMaritimeDrone/large/test/ + json: AerialMaritimeDrone/large/test/annotations_without_background.json + - name: Aquarium + val: + img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ + json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json + - name: CottontailRabbits + val: + img_folder: CottontailRabbits/test/ + json: CottontailRabbits/test/annotations_without_background.json + - name: EgoHands_generic + val: + img_folder: EgoHands/generic/test/ + json: EgoHands/generic/test/annotations_without_background.json + - name: NorthAmericaMushrooms + val: + img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ + json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json + - name: Packages + val: + img_folder: Packages/Raw/test/ + json: Packages/Raw/test/annotations_without_background.json + - name: PascalVOC + val: + img_folder: PascalVOC/valid/ + json: PascalVOC/valid/annotations_without_background.json + - name: Raccoon + val: + img_folder: Raccoon/Raccoon.v2-raw.coco/test/ + json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json + - name: ShellfishOpenImages + val: + img_folder: ShellfishOpenImages/raw/test/ + json: ShellfishOpenImages/raw/test/annotations_without_background.json + - name: VehiclesOpenImages + val: + img_folder: VehiclesOpenImages/416x416/test/ + json: VehiclesOpenImages/416x416/test/annotations_without_background.json + - name: pistols + val: + img_folder: pistols/export/ + json: pistols/export/test_annotations_without_background.json + - name: pothole + val: + img_folder: pothole/test/ + json: pothole/test/annotations_without_background.json + - name: thermalDogsAndPeople + val: + img_folder: thermalDogsAndPeople/test/ + json: thermalDogsAndPeople/test/annotations_without_background.json + + +odinw35_prompts: + AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, + {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", + "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, + {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' + Aquarium: null + CottontailRabbits: null + EgoHands_generic: null + NorthAmericaMushrooms: '[{''id'': 1, ''name'': + ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' + Packages: null + PascalVOC: null + Raccoon: null + ShellfishOpenImages: null + VehiclesOpenImages: null + pistols: null + pothole: null + thermalDogsAndPeople: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only_positive.yaml b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only_positive.yaml new file mode 100644 index 0000000..9a86a52 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only_positive.yaml @@ -0,0 +1,253 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} + +paths: + odinw_data_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + + +supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} +# Validation transforms pipeline +val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: True + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + # Image processing parameters + resolution: 1008 + # Normalization parameters + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + val_batch_size: 2 + num_val_workers: 0 + gather_pred_via_filesys: false + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + max_epochs: 1 + accelerator: cuda + seed_value: 123 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + prompts: ${odinw35_prompts.${supercategory_tuple.name}} + include_negatives: true + category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! + _partial_: true + img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} + ann_file: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + transforms: ${val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: 1 + dict_key: odinw35 + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: true # Set to false if training + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + odinw35: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${supercategory_tuple.name} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + tide: False + iou_type: "bbox" + positive_split: true + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${supercategory_tuple.name} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + + job_array: + num_tasks: 13 + task_index: 0 + +# ============================================================================ +# ODinW13 Supercategories +# ============================================================================ + +all_odinw_supercategories: + - name: AerialMaritimeDrone_large + val: + img_folder: AerialMaritimeDrone/large/test/ + json: AerialMaritimeDrone/large/test/annotations_without_background.json + - name: Aquarium + val: + img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ + json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json + - name: CottontailRabbits + val: + img_folder: CottontailRabbits/test/ + json: CottontailRabbits/test/annotations_without_background.json + - name: EgoHands_generic + val: + img_folder: EgoHands/generic/test/ + json: EgoHands/generic/test/annotations_without_background.json + - name: NorthAmericaMushrooms + val: + img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ + json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json + - name: Packages + val: + img_folder: Packages/Raw/test/ + json: Packages/Raw/test/annotations_without_background.json + - name: PascalVOC + val: + img_folder: PascalVOC/valid/ + json: PascalVOC/valid/annotations_without_background.json + - name: Raccoon + val: + img_folder: Raccoon/Raccoon.v2-raw.coco/test/ + json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json + - name: ShellfishOpenImages + val: + img_folder: ShellfishOpenImages/raw/test/ + json: ShellfishOpenImages/raw/test/annotations_without_background.json + - name: VehiclesOpenImages + val: + img_folder: VehiclesOpenImages/416x416/test/ + json: VehiclesOpenImages/416x416/test/annotations_without_background.json + - name: pistols + val: + img_folder: pistols/export/ + json: pistols/export/test_annotations_without_background.json + - name: pothole + val: + img_folder: pothole/test/ + json: pothole/test/annotations_without_background.json + - name: thermalDogsAndPeople + val: + img_folder: thermalDogsAndPeople/test/ + json: thermalDogsAndPeople/test/annotations_without_background.json + + +odinw35_prompts: + AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, + {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", + "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, + {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' + Aquarium: null + CottontailRabbits: null + EgoHands_generic: null + NorthAmericaMushrooms: '[{''id'': 1, ''name'': + ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' + Packages: null + PascalVOC: null + Raccoon: null + ShellfishOpenImages: null + VehiclesOpenImages: null + pistols: null + pothole: null + thermalDogsAndPeople: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only_train.yaml b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only_train.yaml new file mode 100644 index 0000000..eb03cdf --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_text_only_train.yaml @@ -0,0 +1,591 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} + +paths: + odinw_data_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + + +odinw_train: + train_file: fewshot_train_shot10_seed300 + num_images: null + supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} + # Training transforms pipeline + train_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterCrowds + - _target_: sam3.train.transforms.point_sampling.RandomizeInputBbox + box_noise_std: 0.1 + box_noise_max: 20 + - _target_: sam3.train.transforms.segmentation.DecodeRle + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: + _target_: sam3.train.transforms.basic.get_random_resize_scales + size: ${scratch.resolution} + min_size: 480 + rounded: false + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: ${scratch.consistent_transform} + - _target_: sam3.train.transforms.basic_for_api.PadToSizeAPI + size: ${scratch.resolution} + consistent_transform: ${scratch.consistent_transform} + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.train_norm_mean} + std: ${scratch.train_norm_std} + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterFindQueriesWithTooManyOut + max_num_objects: ${scratch.max_ann_per_img} + + # Validation transforms pipeline + val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # loss config (no mask loss) + loss: + _target_: sam3.train.loss.sam3_loss.Sam3LossWrapper + matcher: ${scratch.matcher} + o2m_weight: 2.0 + o2m_matcher: + _target_: sam3.train.matcher.BinaryOneToManyMatcher + alpha: 0.3 + threshold: 0.4 + topk: 4 + use_o2m_matcher_on_o2m_aux: ${scratch.use_o2m_matcher_on_o2m_aux} + loss_fns_find: + - _target_: sam3.train.loss.loss_fns.Boxes + weight_dict: + loss_bbox: 5.0 + loss_giou: 2.0 + - _target_: sam3.train.loss.loss_fns.IABCEMdetr + weak_loss: False + weight_dict: + loss_ce: ${scratch.loss_ce_weight} # Change + presence_loss: ${scratch.presence_weight} # Change + pos_weight: ${scratch.iabce_pos_weight} + alpha: ${scratch.iabce_alpha} + gamma: 2 + use_presence: True # Change + pos_focal: ${scratch.iabce_pos_focal} + pad_n_queries: ${scratch.num_queries} + pad_scale_pos: ${scratch.instance_query_loss_pad_scale_pos} + + loss_fn_semantic_seg: null + scale_by_find_batch_size: ${scratch.scale_by_find_batch_size} + + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: False + use_act_checkpoint_geo_encoder: True + input_geometry_encoder: + _target_: sam3.model.geometry_encoders.SequenceGeometryEncoder + pos_enc: ${scratch.pos_embed} + encode_boxes_as_points: False + points_direct_project: True + points_pool: True + points_pos_enc: True + boxes_direct_project: True + boxes_pool: True + boxes_pos_enc: True + d_model: ${scratch.d_model} + num_layers: 3 + use_act_ckpt: ${scratch.use_act_checkpoint_geo_encoder} + layer: + _target_: sam3.model.encoder.TransformerEncoderLayer + activation: "relu" + d_model: ${scratch.d_model} + dim_feedforward: 2048 + dropout: ${scratch.encoder_dropout} + pos_enc_at_attn: false + pre_norm: True + pos_enc_at_cross_attn_queries: false + pos_enc_at_cross_attn_keys: true + self_attention: + _target_: sam3.model.attention.MultiheadAttention + attn_type: Vanilla + num_heads: 8 + dropout: ${scratch.encoder_dropout} + embed_dim: ${scratch.d_model} + batch_first: False + cross_attention: + _target_: sam3.model.attention.MultiheadAttention + attn_type: Vanilla + num_heads: 8 + dropout: ${scratch.encoder_dropout} + embed_dim: ${scratch.d_model} + batch_first: False + add_cls: true + add_post_encode_proj: True + + boxRPB: "log" + dac: True + use_early_fusion: true + o2m_mask: false + num_feature_levels: 1 # > 1 not implemented + encoder_dropout: 0.1 + decoder_dropout: 0.1 + + tokenizer_ve: + _target_: sam3.model.tokenizer_ve.SimpleTokenizer + bpe_path: ${paths.bpe_path} + + + freeze_text_tower: False + freeze_image_tower: NoFreeze + vis_backbone_dp: 0.0 + # Activation checkpointing (Save memory) + use_act_checkpoint_vision_backbone: True + use_act_checkpoint_text_backbone: True + use_act_checkpoint_encoder: True + use_act_checkpoint_decoder: True + + loss: null + # Loss parameters + num_queries: 200 + presence_weight: 20.0 + loss_ce_weight: 20.0 + iabce_pos_weight: 5.0 + iabce_pos_focal: false + iabce_alpha: 0.25 + instance_query_loss_pad_scale_pos: 1.0 + use_o2m_matcher_on_o2m_aux: false + + # Model parameters + use_instance_query: true + d_model: 256 + pos_embed: + _target_: sam3.model.position_encoding.PositionEmbeddingSine + num_pos_feats: ${scratch.d_model} + normalize: true + scale: null + temperature: 10000 + + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + + # Matcher configuration + matcher: + _target_: sam3.train.matcher.BinaryHungarianMatcherV2 + focal: true + cost_class: 2.0 + cost_bbox: 5.0 + cost_giou: 2.0 + alpha: 0.25 + gamma: 2 + stable: False + scale_by_find_batch_size: True + + # Image processing parameters + resolution: 1008 + consistent_transform: False + max_ann_per_img: 200 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + train_batch_size: 1 + val_batch_size: 1 + num_train_workers: 0 + num_val_workers: 0 + max_data_epochs: 40 + target_epoch_size: 1500 + hybrid_repeats: 1 + context_length: 2 + gather_pred_via_filesys: false + + # Learning rate and scheduler parameters + lr_scale: 0.1 + lr_transformer: ${times:8e-4,${scratch.lr_scale}} + lr_vision_backbone: ${times:2.5e-4,${scratch.lr_scale}} + lr_language_backbone: ${times:5e-5,${scratch.lr_scale}} + lrd_vision_backbone: 0.9 + wd: 0.1 + scheduler_timescale: 20 + scheduler_warmup: 20 + scheduler_cooldown: 20 + + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + # _target_: sam3.train.trainer.Trainer + # skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: train + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: ${odinw_train.loss} + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + limit_ids: ${odinw_train.num_images} + transforms: ${odinw_train.train_transforms} + load_segmentation: ${scratch.enable_segmentation} + max_ann_per_img: 500000 + multiplier: 1 + max_train_queries: 50000 + max_val_queries: 50000 + training: true + use_caching: False + img_folder: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.train.img_folder} + ann_file: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.train.json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + prompts: ${odinw35_prompts.${odinw_train.supercategory_tuple.name}} #${odinw_train.supercategory_tuple.name) + _partial_: true + shuffle: True + batch_size: ${scratch.train_batch_size} + num_workers: ${scratch.num_train_workers} + pin_memory: False + drop_last: True + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: all + with_seg_masks: ${scratch.enable_segmentation} + + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + load_segmentation: ${scratch.enable_segmentation} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + prompts: ${odinw35_prompts.${odinw_train.supercategory_tuple.name}} + include_negatives: true + category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! + _partial_: true + img_folder: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.val.img_folder} + ann_file: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.val.json} + transforms: ${odinw_train.val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: 1 + dict_key: odinw35 + with_seg_masks: ${scratch.enable_segmentation} + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: false # Set to false if training + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + odinw35: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/odinw/${odinw_train.supercategory_tuple.name} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.val.json} + tide: False + iou_type: "bbox" + positive_split: False + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + optimizer: + _target_: torch.optim.AdamW + + gradient_clip: + _target_: sam3.train.optim.optimizer.GradientClipper + max_norm: 0.1 + norm_type: 2 + + param_group_modifiers: + - _target_: sam3.train.optim.optimizer.layer_decay_param_modifier + _partial_: True + layer_decay_value: ${scratch.lrd_vision_backbone} + apply_to: 'backbone.vision_backbone.trunk' + overrides: + - pattern: '*pos_embed*' + value: 1.0 + + options: + lr: + - scheduler: # transformer and class_embed + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_transformer} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_vision_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.vision_backbone.*' + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_language_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.language_backbone.*' + + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: ${scratch.wd} + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.0 + param_names: + - '*bias*' + module_cls_names: ['torch.nn.LayerNorm'] + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${odinw_train.supercategory_tuple.name} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: null #${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + + # task_index: 2 + # Uncomment for job array configuration + job_array: + num_tasks: 13 + task_index: 0 + + +# ============================================================================ +# ODinW13 Supercategories +# ============================================================================ + +all_odinw_supercategories: + - name: AerialMaritimeDrone_large + val: + img_folder: AerialMaritimeDrone/large/test/ + json: AerialMaritimeDrone/large/test/annotations_without_background.json + train: + img_folder: AerialMaritimeDrone/large/train/ + json: AerialMaritimeDrone/large/train/${odinw_train.train_file}.json + - name: Aquarium + val: + img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ + json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json + train: + img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/train/ + json: Aquarium/Aquarium Combined.v2-raw-1024.coco/train/${odinw_train.train_file}.json + - name: CottontailRabbits + val: + img_folder: CottontailRabbits/test/ + json: CottontailRabbits/test/annotations_without_background.json + train: + img_folder: CottontailRabbits/train/ + json: CottontailRabbits/train/${odinw_train.train_file}.json + - name: EgoHands_generic + val: + img_folder: EgoHands/generic/test/ + json: EgoHands/generic/test/annotations_without_background.json + train: + img_folder: EgoHands/generic/train/ + json: EgoHands/generic/train/${odinw_train.train_file}.json + - name: NorthAmericaMushrooms + val: + img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ + json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json + train: + img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/train/ + json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/train/${odinw_train.train_file}.json + - name: Packages + val: + img_folder: Packages/Raw/test/ + json: Packages/Raw/test/annotations_without_background.json + train: + img_folder: Packages/Raw/train/ + json: Packages/Raw/train/${odinw_train.train_file}.json + - name: PascalVOC + val: + img_folder: PascalVOC/valid/ + json: PascalVOC/valid/annotations_without_background.json + train: + img_folder: PascalVOC/train/ + json: PascalVOC/train/${odinw_train.train_file}.json + - name: Raccoon + val: + img_folder: Raccoon/Raccoon.v2-raw.coco/test/ + json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json + train: + img_folder: Raccoon/Raccoon.v2-raw.coco/train/ + json: Raccoon/Raccoon.v2-raw.coco/train/${odinw_train.train_file}.json + - name: ShellfishOpenImages + val: + img_folder: ShellfishOpenImages/raw/test/ + json: ShellfishOpenImages/raw/test/annotations_without_background.json + train: + img_folder: ShellfishOpenImages/raw/train/ + json: ShellfishOpenImages/raw/train/${odinw_train.train_file}.json + - name: VehiclesOpenImages + val: + img_folder: VehiclesOpenImages/416x416/test/ + json: VehiclesOpenImages/416x416/test/annotations_without_background.json + train: + img_folder: VehiclesOpenImages/416x416/train/ + json: VehiclesOpenImages/416x416/train/${odinw_train.train_file}.json + - name: pistols + val: + img_folder: pistols/export/ + json: pistols/export/test_annotations_without_background.json + train: + img_folder: pistols/export/ + json: pistols/export/${odinw_train.train_file}.json + - name: pothole + val: + img_folder: pothole/test/ + json: pothole/test/annotations_without_background.json + train: + img_folder: pothole/train/ + json: pothole/train/${odinw_train.train_file}.json + - name: thermalDogsAndPeople + val: + img_folder: thermalDogsAndPeople/test/ + json: thermalDogsAndPeople/test/annotations_without_background.json + train: + img_folder: thermalDogsAndPeople/train/ + json: thermalDogsAndPeople/train/${odinw_train.train_file}.json + + +odinw35_prompts: + AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, + {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", + "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, + {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' + Aquarium: null + CottontailRabbits: null + EgoHands_generic: null + NorthAmericaMushrooms: '[{''id'': 1, ''name'': + ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' + Packages: null + PascalVOC: null + Raccoon: null + ShellfishOpenImages: null + VehiclesOpenImages: null + pistols: null + pothole: null + thermalDogsAndPeople: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_visual_only.yaml b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_visual_only.yaml new file mode 100644 index 0000000..e724f2d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/odinw13/odinw_visual_only.yaml @@ -0,0 +1,256 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} + +paths: + odinw_data_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + + +supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} +# Validation transforms pipeline +val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + - _target_: sam3.train.transforms.filter_query_transforms.TextQueryToVisual + keep_text_queries: false # Note: set this to false if you only want visual + probability: 1.0 # always + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: True + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + # Image processing parameters + resolution: 1008 + # Normalization parameters + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + val_batch_size: 2 + num_val_workers: 0 + gather_pred_via_filesys: false + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + max_epochs: 1 + accelerator: cuda + seed_value: 123 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + prompts: ${odinw35_prompts.${supercategory_tuple.name}} + include_negatives: true + category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! + _partial_: true + img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} + ann_file: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + transforms: ${val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: 1 + dict_key: odinw35 + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: true # Set to false if training + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + odinw35: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${supercategory_tuple.name} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: + _target_: sam3.eval.coco_reindex.reindex_coco_to_temp + input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} + tide: False + iou_type: "bbox" + positive_split: true + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${supercategory_tuple.name} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + + job_array: + num_tasks: 13 + task_index: 0 + +# ============================================================================ +# ODinW13 Supercategories +# ============================================================================ + +all_odinw_supercategories: + - name: AerialMaritimeDrone_large + val: + img_folder: AerialMaritimeDrone/large/test/ + json: AerialMaritimeDrone/large/test/annotations_without_background.json + - name: Aquarium + val: + img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ + json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json + - name: CottontailRabbits + val: + img_folder: CottontailRabbits/test/ + json: CottontailRabbits/test/annotations_without_background.json + - name: EgoHands_generic + val: + img_folder: EgoHands/generic/test/ + json: EgoHands/generic/test/annotations_without_background.json + - name: NorthAmericaMushrooms + val: + img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ + json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json + - name: Packages + val: + img_folder: Packages/Raw/test/ + json: Packages/Raw/test/annotations_without_background.json + - name: PascalVOC + val: + img_folder: PascalVOC/valid/ + json: PascalVOC/valid/annotations_without_background.json + - name: Raccoon + val: + img_folder: Raccoon/Raccoon.v2-raw.coco/test/ + json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json + - name: ShellfishOpenImages + val: + img_folder: ShellfishOpenImages/raw/test/ + json: ShellfishOpenImages/raw/test/annotations_without_background.json + - name: VehiclesOpenImages + val: + img_folder: VehiclesOpenImages/416x416/test/ + json: VehiclesOpenImages/416x416/test/annotations_without_background.json + - name: pistols + val: + img_folder: pistols/export/ + json: pistols/export/test_annotations_without_background.json + - name: pothole + val: + img_folder: pothole/test/ + json: pothole/test/annotations_without_background.json + - name: thermalDogsAndPeople + val: + img_folder: thermalDogsAndPeople/test/ + json: thermalDogsAndPeople/test/annotations_without_background.json + + +odinw35_prompts: + AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, + {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", + "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, + {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' + Aquarium: null + CottontailRabbits: null + EgoHands_generic: null + NorthAmericaMushrooms: '[{''id'': 1, ''name'': + ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' + Packages: null + PascalVOC: null + Raccoon: null + ShellfishOpenImages: null + VehiclesOpenImages: null + pistols: null + pothole: null + thermalDogsAndPeople: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/roboflow_v100/roboflow_v100_eval.yaml b/src/nn/segearth_ov3/sam3/train/configs/roboflow_v100/roboflow_v100_eval.yaml new file mode 100644 index 0000000..361e622 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/roboflow_v100/roboflow_v100_eval.yaml @@ -0,0 +1,539 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + roboflow_vl_100_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + +# Roboflow dataset configuration +roboflow_train: + num_images: 100 # Note: This is the number of images used for training. If null, all images are used. + supercategory: ${all_roboflow_supercategories.${string:${submitit.job_array.task_index}}} + + # Training transforms pipeline + train_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterCrowds + - _target_: sam3.train.transforms.point_sampling.RandomizeInputBbox + box_noise_std: 0.1 + box_noise_max: 20 + - _target_: sam3.train.transforms.segmentation.DecodeRle + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: + _target_: sam3.train.transforms.basic.get_random_resize_scales + size: ${scratch.resolution} + min_size: 480 + rounded: false + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: ${scratch.consistent_transform} + - _target_: sam3.train.transforms.basic_for_api.PadToSizeAPI + size: ${scratch.resolution} + consistent_transform: ${scratch.consistent_transform} + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.train_norm_mean} + std: ${scratch.train_norm_std} + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterFindQueriesWithTooManyOut + max_num_objects: ${scratch.max_ann_per_img} + + # Validation transforms pipeline + val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.train_norm_mean} + std: ${scratch.train_norm_std} + + # loss config (no mask loss) + loss: + _target_: sam3.train.loss.sam3_loss.Sam3LossWrapper + matcher: ${scratch.matcher} + o2m_weight: 2.0 + o2m_matcher: + _target_: sam3.train.matcher.BinaryOneToManyMatcher + alpha: 0.3 + threshold: 0.4 + topk: 4 + use_o2m_matcher_on_o2m_aux: false # Another option is true + loss_fns_find: + - _target_: sam3.train.loss.loss_fns.Boxes + weight_dict: + loss_bbox: 5.0 + loss_giou: 2.0 + - _target_: sam3.train.loss.loss_fns.IABCEMdetr + weak_loss: False + weight_dict: + loss_ce: 20.0 # Another option is 100.0 + presence_loss: 20.0 + pos_weight: 10.0 # Another option is 5.0 + alpha: 0.25 + gamma: 2 + use_presence: True # Change + pos_focal: false + pad_n_queries: 200 + pad_scale_pos: 1.0 + + loss_fn_semantic_seg: null + scale_by_find_batch_size: ${scratch.scale_by_find_batch_size} + + + # NOTE: Loss to be used for training in case of segmentation + # loss: + # _target_: sam3.train.loss.sam3_loss.Sam3LossWrapper + # matcher: ${scratch.matcher} + # o2m_weight: 2.0 + # o2m_matcher: + # _target_: sam3.train.matcher.BinaryOneToManyMatcher + # alpha: 0.3 + # threshold: 0.4 + # topk: 4 + # use_o2m_matcher_on_o2m_aux: false + # loss_fns_find: + # - _target_: sam3.train.loss.loss_fns.Boxes + # weight_dict: + # loss_bbox: 5.0 + # loss_giou: 2.0 + # - _target_: sam3.train.loss.loss_fns.IABCEMdetr + # weak_loss: False + # weight_dict: + # loss_ce: 20.0 # Another option is 100.0 + # presence_loss: 20.0 + # pos_weight: 10.0 # Another option is 5.0 + # alpha: 0.25 + # gamma: 2 + # use_presence: True # Change + # pos_focal: false + # pad_n_queries: 200 + # pad_scale_pos: 1.0 + # - _target_: sam3.train.loss.loss_fns.Masks + # focal_alpha: 0.25 + # focal_gamma: 2.0 + # weight_dict: + # loss_mask: 200.0 + # loss_dice: 10.0 + # compute_aux: false + # loss_fn_semantic_seg: + # _target_: sam3.losses.loss_fns.SemanticSegCriterion + # presence_head: True + # presence_loss: False # Change + # focal: True + # focal_alpha: 0.6 + # focal_gamma: 2.0 + # downsample: False + # weight_dict: + # loss_semantic_seg: 20.0 + # loss_semantic_presence: 1.0 + # loss_semantic_dice: 30.0 + # scale_by_find_batch_size: ${scratch.scale_by_find_batch_size} + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: False # NOTE: This is the number of queries used for segmentation + # Model parameters + d_model: 256 + pos_embed: + _target_: sam3.model.position_encoding.PositionEmbeddingSine + num_pos_feats: ${scratch.d_model} + normalize: true + scale: null + temperature: 10000 + + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + # Matcher configuration + matcher: + _target_: sam3.train.matcher.BinaryHungarianMatcherV2 + focal: true # with `focal: true` it is equivalent to BinaryFocalHungarianMatcher + cost_class: 2.0 + cost_bbox: 5.0 + cost_giou: 2.0 + alpha: 0.25 + gamma: 2 + stable: False + scale_by_find_batch_size: True + + # Image processing parameters + resolution: 1008 + consistent_transform: False + max_ann_per_img: 200 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + num_train_workers: 10 + num_val_workers: 0 + max_data_epochs: 20 + target_epoch_size: 1500 + hybrid_repeats: 1 + context_length: 2 + gather_pred_via_filesys: false + + # Learning rate and scheduler parameters + lr_scale: 0.1 + lr_transformer: ${times:8e-4,${scratch.lr_scale}} + lr_vision_backbone: ${times:2.5e-4,${scratch.lr_scale}} + lr_language_backbone: ${times:5e-5,${scratch.lr_scale}} + lrd_vision_backbone: 0.9 + wd: 0.1 + scheduler_timescale: 20 + scheduler_warmup: 20 + scheduler_cooldown: 20 + + val_batch_size: 1 + collate_fn_val: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: roboflow100 + with_seg_masks: ${scratch.enable_segmentation} # Note: Set this to true if using segmentation masks! + + gradient_accumulation_steps: 1 + train_batch_size: 1 + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: all + with_seg_masks: ${scratch.enable_segmentation} # Note: Set this to true if using segmentation masks! + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: 20 + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + gradient_accumulation_steps: ${scratch.gradient_accumulation_steps} + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: ${roboflow_train.loss} + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + limit_ids: ${roboflow_train.num_images} + transforms: ${roboflow_train.train_transforms} + load_segmentation: ${scratch.enable_segmentation} + max_ann_per_img: 500000 + multiplier: 1 + max_train_queries: 50000 + max_val_queries: 50000 + training: true + use_caching: False + img_folder: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/train/ + ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/train/_annotations.coco.json + + shuffle: True + batch_size: ${scratch.train_batch_size} + num_workers: ${scratch.num_train_workers} + pin_memory: True + drop_last: True + collate_fn: ${scratch.collate_fn} + + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + load_segmentation: ${scratch.enable_segmentation} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + include_negatives: true + category_chunk_size: 2 # Note: You can increase this based on the memory of your GPU. + _partial_: true + img_folder: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/ + ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json + transforms: ${roboflow_train.val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: ${scratch.collate_fn_val} + + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: true + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + roboflow100: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${roboflow_train.supercategory} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json + tide: False + iou_type: "bbox" + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + optimizer: + _target_: torch.optim.AdamW + + gradient_clip: + _target_: sam3.train.optim.optimizer.GradientClipper + max_norm: 0.1 + norm_type: 2 + + param_group_modifiers: + - _target_: sam3.train.optim.optimizer.layer_decay_param_modifier + _partial_: True + layer_decay_value: ${scratch.lrd_vision_backbone} + apply_to: 'backbone.vision_backbone.trunk' + overrides: + - pattern: '*pos_embed*' + value: 1.0 + + options: + lr: + - scheduler: # transformer and class_embed + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_transformer} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_vision_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.vision_backbone.*' + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_language_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.language_backbone.*' + + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: ${scratch.wd} + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.0 + param_names: + - '*bias*' + module_cls_names: ['torch.nn.LayerNorm'] + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${roboflow_train.supercategory} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + # Uncomment for job array configuration + job_array: + num_tasks: 100 + task_index: 0 + +# ============================================================================ +# Available Roboflow Supercategories (for reference) +# ============================================================================ + +all_roboflow_supercategories: + - -grccs + - zebrasatasturias + - cod-mw-warzone + - canalstenosis + - label-printing-defect-version-2 + - new-defects-in-wood + - orionproducts + - aquarium-combined + - varroa-mites-detection--test-set + - clashroyalechardetector + - stomata-cells + - halo-infinite-angel-videogame + - pig-detection + - urine-analysis1 + - aerial-sheep + - orgharvest + - actions + - mahjong + - liver-disease + - needle-base-tip-min-max + - wheel-defect-detection + - aircraft-turnaround-dataset + - xray + - wildfire-smoke + - spinefrxnormalvindr + - ufba-425 + - speech-bubbles-detection + - train + - pill + - truck-movement + - car-logo-detection + - inbreast + - sea-cucumbers-new-tiles + - uavdet-small + - penguin-finder-seg + - aerial-airport + - bibdetection + - taco-trash-annotations-in-context + - bees + - recode-waste + - screwdetectclassification + - wine-labels + - aerial-cows + - into-the-vale + - gwhd2021 + - lacrosse-object-detection + - defect-detection + - dataconvert + - x-ray-id + - ball + - tube + - 2024-frc + - crystal-clean-brain-tumors-mri-dataset + - grapes-5 + - human-detection-in-floods + - buoy-onboarding + - apoce-aerial-photographs-for-object-detection-of-construction-equipment + - l10ul502 + - floating-waste + - deeppcb + - ism-band-packet-detection + - weeds4 + - invoice-processing + - thermal-cheetah + - tomatoes-2 + - marine-sharks + - peixos-fish + - sssod + - aerial-pool + - countingpills + - asphaltdistressdetection + - roboflow-trained-dataset + - everdaynew + - underwater-objects + - soda-bottles + - dentalai + - jellyfish + - deepfruits + - activity-diagrams + - circuit-voltages + - all-elements + - macro-segmentation + - exploratorium-daphnia + - signatures + - conveyor-t-shirts + - fruitjes + - grass-weeds + - infraredimageofpowerequipment + - 13-lkc01 + - wb-prova + - flir-camera-objects + - paper-parts + - football-player-detection + - trail-camera + - smd-components + - water-meter + - nih-xray + - the-dreidel-project + - electric-pylon-detection-in-rsi + - cable-damage diff --git a/src/nn/segearth_ov3/sam3/train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml b/src/nn/segearth_ov3/sam3/train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml new file mode 100644 index 0000000..6b95f62 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml @@ -0,0 +1,539 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + roboflow_vl_100_root: + experiment_log_dir: + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + +# Roboflow dataset configuration +roboflow_train: + num_images: 100 # Note: This is the number of images used for training. If null, all images are used. + supercategory: ${all_roboflow_supercategories.${string:${submitit.job_array.task_index}}} + + # Training transforms pipeline + train_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterCrowds + - _target_: sam3.train.transforms.point_sampling.RandomizeInputBbox + box_noise_std: 0.1 + box_noise_max: 20 + - _target_: sam3.train.transforms.segmentation.DecodeRle + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: + _target_: sam3.train.transforms.basic.get_random_resize_scales + size: ${scratch.resolution} + min_size: 480 + rounded: false + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: ${scratch.consistent_transform} + - _target_: sam3.train.transforms.basic_for_api.PadToSizeAPI + size: ${scratch.resolution} + consistent_transform: ${scratch.consistent_transform} + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.train_norm_mean} + std: ${scratch.train_norm_std} + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets + - _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries + query_filter: + _target_: sam3.train.transforms.filter_query_transforms.FilterFindQueriesWithTooManyOut + max_num_objects: ${scratch.max_ann_per_img} + + # Validation transforms pipeline + val_transforms: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} + max_size: + _target_: sam3.train.transforms.basic.get_random_resize_max_size + size: ${scratch.resolution} + square: true + consistent_transform: False + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.train_norm_mean} + std: ${scratch.train_norm_std} + + # loss config (no mask loss) + loss: + _target_: sam3.train.loss.sam3_loss.Sam3LossWrapper + matcher: ${scratch.matcher} + o2m_weight: 2.0 + o2m_matcher: + _target_: sam3.train.matcher.BinaryOneToManyMatcher + alpha: 0.3 + threshold: 0.4 + topk: 4 + use_o2m_matcher_on_o2m_aux: false # Another option is true + loss_fns_find: + - _target_: sam3.train.loss.loss_fns.Boxes + weight_dict: + loss_bbox: 5.0 + loss_giou: 2.0 + - _target_: sam3.train.loss.loss_fns.IABCEMdetr + weak_loss: False + weight_dict: + loss_ce: 20.0 # Another option is 100.0 + presence_loss: 20.0 + pos_weight: 10.0 # Another option is 5.0 + alpha: 0.25 + gamma: 2 + use_presence: True # Change + pos_focal: false + pad_n_queries: 200 + pad_scale_pos: 1.0 + + loss_fn_semantic_seg: null + scale_by_find_batch_size: ${scratch.scale_by_find_batch_size} + + + # NOTE: Loss to be used for training in case of segmentation + # loss: + # _target_: sam3.train.loss.sam3_loss.Sam3LossWrapper + # matcher: ${scratch.matcher} + # o2m_weight: 2.0 + # o2m_matcher: + # _target_: sam3.train.matcher.BinaryOneToManyMatcher + # alpha: 0.3 + # threshold: 0.4 + # topk: 4 + # use_o2m_matcher_on_o2m_aux: false + # loss_fns_find: + # - _target_: sam3.train.loss.loss_fns.Boxes + # weight_dict: + # loss_bbox: 5.0 + # loss_giou: 2.0 + # - _target_: sam3.train.loss.loss_fns.IABCEMdetr + # weak_loss: False + # weight_dict: + # loss_ce: 20.0 # Another option is 100.0 + # presence_loss: 20.0 + # pos_weight: 10.0 # Another option is 5.0 + # alpha: 0.25 + # gamma: 2 + # use_presence: True # Change + # pos_focal: false + # pad_n_queries: 200 + # pad_scale_pos: 1.0 + # - _target_: sam3.train.loss.loss_fns.Masks + # focal_alpha: 0.25 + # focal_gamma: 2.0 + # weight_dict: + # loss_mask: 200.0 + # loss_dice: 10.0 + # compute_aux: false + # loss_fn_semantic_seg: + # _target_: sam3.losses.loss_fns.SemanticSegCriterion + # presence_head: True + # presence_loss: False # Change + # focal: True + # focal_alpha: 0.6 + # focal_gamma: 2.0 + # downsample: False + # weight_dict: + # loss_semantic_seg: 20.0 + # loss_semantic_presence: 1.0 + # loss_semantic_dice: 30.0 + # scale_by_find_batch_size: ${scratch.scale_by_find_batch_size} + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + enable_segmentation: False # NOTE: This is the number of queries used for segmentation + # Model parameters + d_model: 256 + pos_embed: + _target_: sam3.model.position_encoding.PositionEmbeddingSine + num_pos_feats: ${scratch.d_model} + normalize: true + scale: null + temperature: 10000 + + # Box processing + use_presence_eval: True + original_box_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessImage + max_dets_per_img: -1 # infinite detections + use_original_ids: true + use_original_sizes_box: true + use_presence: ${scratch.use_presence_eval} + + # Matcher configuration + matcher: + _target_: sam3.train.matcher.BinaryHungarianMatcherV2 + focal: true # with `focal: true` it is equivalent to BinaryFocalHungarianMatcher + cost_class: 2.0 + cost_bbox: 5.0 + cost_giou: 2.0 + alpha: 0.25 + gamma: 2 + stable: False + scale_by_find_batch_size: True + + # Image processing parameters + resolution: 1008 + consistent_transform: False + max_ann_per_img: 200 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + # Training parameters + num_train_workers: 10 + num_val_workers: 0 + max_data_epochs: 20 + target_epoch_size: 1500 + hybrid_repeats: 1 + context_length: 2 + gather_pred_via_filesys: false + + # Learning rate and scheduler parameters + lr_scale: 0.1 + lr_transformer: ${times:8e-4,${scratch.lr_scale}} + lr_vision_backbone: ${times:2.5e-4,${scratch.lr_scale}} + lr_language_backbone: ${times:5e-5,${scratch.lr_scale}} + lrd_vision_backbone: 0.9 + wd: 0.1 + scheduler_timescale: 20 + scheduler_warmup: 20 + scheduler_cooldown: 20 + + val_batch_size: 1 + collate_fn_val: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: roboflow100 + with_seg_masks: ${scratch.enable_segmentation} # Note: Set this to true if using segmentation masks! + + gradient_accumulation_steps: 1 + train_batch_size: 1 + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: all + with_seg_masks: ${scratch.enable_segmentation} # Note: Set this to true if using segmentation masks! + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: 20 + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: train + gradient_accumulation_steps: ${scratch.gradient_accumulation_steps} + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: ${roboflow_train.loss} + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + limit_ids: ${roboflow_train.num_images} + transforms: ${roboflow_train.train_transforms} + load_segmentation: ${scratch.enable_segmentation} + max_ann_per_img: 500000 + multiplier: 1 + max_train_queries: 50000 + max_val_queries: 50000 + training: true + use_caching: False + img_folder: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/train/ + ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/train/_annotations.coco.json + + shuffle: True + batch_size: ${scratch.train_batch_size} + num_workers: ${scratch.num_train_workers} + pin_memory: True + drop_last: True + collate_fn: ${scratch.collate_fn} + + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + load_segmentation: ${scratch.enable_segmentation} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON + include_negatives: true + category_chunk_size: 2 # Note: You can increase this based on the memory of your GPU. + _partial_: true + img_folder: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/ + ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json + transforms: ${roboflow_train.val_transforms} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: ${scratch.collate_fn_val} + + + model: + _target_: sam3.model_builder.build_sam3_image_model + bpe_path: ${paths.bpe_path} + device: cpus + eval_mode: false + enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. + + meters: + val: + roboflow100: + detection: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "bbox" + dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${roboflow_train.supercategory} + merge_predictions: True + postprocessor: ${scratch.original_box_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 100 + pred_file_evaluators: + - _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators + gt_path: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json + tide: False + iou_type: "bbox" + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + optimizer: + _target_: torch.optim.AdamW + + gradient_clip: + _target_: sam3.train.optim.optimizer.GradientClipper + max_norm: 0.1 + norm_type: 2 + + param_group_modifiers: + - _target_: sam3.train.optim.optimizer.layer_decay_param_modifier + _partial_: True + layer_decay_value: ${scratch.lrd_vision_backbone} + apply_to: 'backbone.vision_backbone.trunk' + overrides: + - pattern: '*pos_embed*' + value: 1.0 + + options: + lr: + - scheduler: # transformer and class_embed + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_transformer} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_vision_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.vision_backbone.*' + - scheduler: + _target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler + base_lr: ${scratch.lr_language_backbone} + timescale: ${scratch.scheduler_timescale} + warmup_steps: ${scratch.scheduler_warmup} + cooldown_steps: ${scratch.scheduler_cooldown} + param_names: + - 'backbone.language_backbone.*' + + weight_decay: + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: ${scratch.wd} + - scheduler: + _target_: fvcore.common.param_scheduler.ConstantParamScheduler + value: 0.0 + param_names: + - '*bias*' + module_cls_names: ['torch.nn.LayerNorm'] + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/${roboflow_train.supercategory} + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 1 + gpus_per_node: 2 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null + # Uncomment for job array configuration + job_array: + num_tasks: 100 + task_index: 0 + +# ============================================================================ +# Available Roboflow Supercategories (for reference) +# ============================================================================ + +all_roboflow_supercategories: + - -grccs + - zebrasatasturias + - cod-mw-warzone + - canalstenosis + - label-printing-defect-version-2 + - new-defects-in-wood + - orionproducts + - aquarium-combined + - varroa-mites-detection--test-set + - clashroyalechardetector + - stomata-cells + - halo-infinite-angel-videogame + - pig-detection + - urine-analysis1 + - aerial-sheep + - orgharvest + - actions + - mahjong + - liver-disease + - needle-base-tip-min-max + - wheel-defect-detection + - aircraft-turnaround-dataset + - xray + - wildfire-smoke + - spinefrxnormalvindr + - ufba-425 + - speech-bubbles-detection + - train + - pill + - truck-movement + - car-logo-detection + - inbreast + - sea-cucumbers-new-tiles + - uavdet-small + - penguin-finder-seg + - aerial-airport + - bibdetection + - taco-trash-annotations-in-context + - bees + - recode-waste + - screwdetectclassification + - wine-labels + - aerial-cows + - into-the-vale + - gwhd2021 + - lacrosse-object-detection + - defect-detection + - dataconvert + - x-ray-id + - ball + - tube + - 2024-frc + - crystal-clean-brain-tumors-mri-dataset + - grapes-5 + - human-detection-in-floods + - buoy-onboarding + - apoce-aerial-photographs-for-object-detection-of-construction-equipment + - l10ul502 + - floating-waste + - deeppcb + - ism-band-packet-detection + - weeds4 + - invoice-processing + - thermal-cheetah + - tomatoes-2 + - marine-sharks + - peixos-fish + - sssod + - aerial-pool + - countingpills + - asphaltdistressdetection + - roboflow-trained-dataset + - everdaynew + - underwater-objects + - soda-bottles + - dentalai + - jellyfish + - deepfruits + - activity-diagrams + - circuit-voltages + - all-elements + - macro-segmentation + - exploratorium-daphnia + - signatures + - conveyor-t-shirts + - fruitjes + - grass-weeds + - infraredimageofpowerequipment + - 13-lkc01 + - wb-prova + - flir-camera-objects + - paper-parts + - football-player-detection + - trail-camera + - smd-components + - water-meter + - nih-xray + - the-dreidel-project + - electric-pylon-detection-in-rsi + - cable-damage diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_test.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_test.yaml new file mode 100644 index 0000000..b5bed47 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_test.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_sav_test + experiment_log_dir: + ytvis_json: /saco_veval_sav_test.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: True + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml new file mode 100644 index 0000000..abc3289 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_sav_test + experiment_log_dir: + ytvis_json: /saco_veval_sav_test.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: False + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_val.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_val.yaml new file mode 100644 index 0000000..25c6e60 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_val.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_sav_val + experiment_log_dir: + ytvis_json: /saco_veval_sav_val.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: True + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml new file mode 100644 index 0000000..9a89eba --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_sav_val + experiment_log_dir: + ytvis_json: /saco_veval_sav_val.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: False + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml new file mode 100644 index 0000000..fdc8185 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_smartglasses_test + experiment_log_dir: + ytvis_json: /saco_veval_smartglasses_test.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: True + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml new file mode 100644 index 0000000..2d6150e --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_smartglasses_test + experiment_log_dir: + ytvis_json: /saco_veval_smartglasses_test.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: False + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml new file mode 100644 index 0000000..a9bffda --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_smartglasses_val + experiment_log_dir: + ytvis_json: /saco_veval_smartglasses_val.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: True + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml new file mode 100644 index 0000000..e1f6443 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_smartglasses_val + experiment_log_dir: + ytvis_json: /saco_veval_smartglasses_val.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: False + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_test.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_test.yaml new file mode 100644 index 0000000..71e5034 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_test.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_yt1b_test + experiment_log_dir: + ytvis_json: /saco_veval_yt1b_test.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: True + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml new file mode 100644 index 0000000..f8df6ae --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_yt1b_test + experiment_log_dir: + ytvis_json: /saco_veval_yt1b_test.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: False + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_val.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_val.yaml new file mode 100644 index 0000000..5816952 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_val.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_yt1b_val + experiment_log_dir: + ytvis_json: /saco_veval_yt1b_val.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: True + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml new file mode 100644 index 0000000..374e5e2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml @@ -0,0 +1,174 @@ +# @package _global_ +defaults: + - _self_ + +# ============================================================================ +# Paths Configuration (Chage this to your own paths) +# ============================================================================ +paths: + + dump_file_name: saco_veval_yt1b_val + experiment_log_dir: + ytvis_json: /saco_veval_yt1b_val.json + ytvis_dir : + bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz + num_videos: null + +# ============================================================================ +# Different helper parameters and functions +# ============================================================================ +scratch: + vid_mask_postprocessor: + _target_: sam3.eval.postprocessors.PostProcessNullOp + + use_presence_eval: True + + video_transforms_val: + - _target_: sam3.train.transforms.basic_for_api.ComposeAPI + transforms: + - _target_: sam3.train.transforms.segmentation.DecodeRle + # resize the image to 1024x1024 resolution + - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI + sizes: ${scratch.resolution} # originally `resolution: 1024` + square: true + consistent_transform: true + - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI + - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI + mean: ${scratch.val_norm_mean} + std: ${scratch.val_norm_std} + + # Model parameters + d_model: 256 + + # Image processing parameters + resolution: 1008 + + # Normalization parameters + train_norm_mean: [0.5, 0.5, 0.5] + train_norm_std: [0.5, 0.5, 0.5] + val_norm_mean: [0.5, 0.5, 0.5] + val_norm_std: [0.5, 0.5, 0.5] + + val_batch_size: 1 + num_val_workers: 0 + max_data_epochs: 20 + hybrid_repeats: 1 + gather_pred_via_filesys: false + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + _target_: sam3.train.trainer.Trainer + skip_saving_ckpts: true + empty_gpu_mem_cache_after_eval: True + skip_first_val: True + max_epochs: ${scratch.max_data_epochs} + accelerator: cuda + seed_value: 123 + val_epoch_freq: 10 + mode: val + + distributed: + backend: nccl + find_unused_parameters: True + gradient_as_bucket_view: True + + loss: + all: + _target_: sam3.train.loss.sam3_loss.DummyLoss + default: + _target_: sam3.train.loss.sam3_loss.DummyLoss + + data: + train: null + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset + limit_ids: ${paths.num_videos} + img_folder: ${paths.ytvis_dir} + ann_file: ${paths.ytvis_json} + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP + _partial_: true + + transforms: ${scratch.video_transforms_val} + max_ann_per_img: 100000 # filtered in transforms + max_val_queries: 100000 + multiplier: 1 + load_segmentation: true + training: false + + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: True + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: ytvis_val + with_seg_masks: true + + + model: + _target_: sam3.model_builder.build_sam3_video_model + bpe_path: ${paths.bpe_path} + has_presence_token: True + geo_encoder_use_img_cross_attn: True + apply_temporal_disambiguation: False + + meters: + val: + ytvis_val: + pred_file: # key + _target_: sam3.eval.ytvis_eval.YTVISResultsWriter + dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json + postprocessor: ${scratch.vid_mask_postprocessor} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + + optim: + amp: + enabled: True + amp_dtype: bfloat16 + + + checkpoint: + save_dir: ${launcher.experiment_log_dir}/checkpoints + save_freq: 0 # 0 only last checkpoint is saved. + + + logging: + tensorboard_writer: + _target_: sam3.train.utils.logger.make_tensorboard_logger + log_dir: ${launcher.experiment_log_dir}/tensorboard + flush_secs: 120 + should_log: True + wandb_writer: null + log_dir: ${launcher.experiment_log_dir}/logs/ + log_freq: 10 + +# ============================================================================ +# Launcher and Submitit Configuration +# ============================================================================ + +launcher: + num_nodes: 8 + gpus_per_node: 8 + experiment_log_dir: ${paths.experiment_log_dir} + multiprocessing_context: forkserver + +submitit: + account: null + partition: null + qos: null + timeout_hour: 72 + use_cluster: True + cpus_per_task: 10 + port_range: [10000, 65000] + constraint: null diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml new file mode 100644 index 0000000..e5587cf --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_bdd100k/ + coco_gt: ${paths.base_annotation_path_silver}/silver_bdd100k_merged_test.json + img_path: ${paths.silver_img_path}/bdd100k/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_bdd100k + + meters: + val: + silver_bdd100k: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_bdd100k + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_droid.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_droid.yaml new file mode 100644 index 0000000..c0d6234 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_droid.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_droid/ + coco_gt: ${paths.base_annotation_path_silver}/silver_droid_merged_test.json + img_path: ${paths.silver_img_path}/droid/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_droid + + meters: + val: + silver_droid: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_droid + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml new file mode 100644 index 0000000..d5a036d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_ego4d/ + coco_gt: ${paths.base_annotation_path_silver}/silver_ego4d_merged_test.json + img_path: ${paths.silver_img_path}/ego4d/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_ego4d + + meters: + val: + silver_ego4d: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_ego4d + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml new file mode 100644 index 0000000..b15d0c8 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_fathomnet/ + coco_gt: ${paths.base_annotation_path_silver}/silver_fathomnet_test.json + img_path: ${paths.silver_img_path}/fathomnet/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_fathomnet + + meters: + val: + silver_fathomnet: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_fathomnet + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml new file mode 100644 index 0000000..5158ff5 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_food_rec/ + coco_gt: ${paths.base_annotation_path_silver}/silver_food_rec_merged_test.json + img_path: ${paths.silver_img_path}/food_rec/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_food_rec + + meters: + val: + silver_food_rec: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_food_rec + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_geode.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_geode.yaml new file mode 100644 index 0000000..08f159f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_geode.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_geode/ + coco_gt: ${paths.base_annotation_path_silver}/silver_geode_merged_test.json + img_path: ${paths.silver_img_path}/geode/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_geode + + meters: + val: + silver_geode: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_geode + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml new file mode 100644 index 0000000..1d56d97 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_inaturalist/ + coco_gt: ${paths.base_annotation_path_silver}/silver_inaturalist_merged_test.json + img_path: ${paths.silver_img_path}/inaturalist/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_inaturalist + + meters: + val: + silver_inaturalist: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_inaturalist + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_nga.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_nga.yaml new file mode 100644 index 0000000..b2de0af --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_nga.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_nga_art/ + coco_gt: ${paths.base_annotation_path_silver}/silver_nga_art_merged_test.json + img_path: ${paths.silver_img_path}/nga/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_nga_art + + meters: + val: + silver_nga_art: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_nga_art + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_sav.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_sav.yaml new file mode 100644 index 0000000..7ebbb0f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_sav.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_sav/ + coco_gt: ${paths.base_annotation_path_silver}/silver_sav_merged_test.json + img_path: ${paths.silver_img_path}/sav/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_sav + + meters: + val: + silver_sav: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_sav + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml new file mode 100644 index 0000000..901bd3a --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml @@ -0,0 +1,64 @@ +# @package _global_ +defaults: + - /configs/eval_base.yaml + - _self_ + +# ============================================================================ +# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct +# ============================================================================ +paths: + experiment_log_dir: ${paths.base_experiment_log_dir}/silver_yt1b/ + coco_gt: ${paths.base_annotation_path_silver}/silver_yt1b_merged_test.json + img_path: ${paths.silver_img_path}/yt1b/ + + + +# ============================================================================ +# Trainer Configuration +# ============================================================================ + +trainer: + data: + val: + _target_: sam3.train.data.torch_dataset.TorchDataset + dataset: + _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset + coco_json_loader: + _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP + _partial_: true + img_folder: ${paths.img_path} + ann_file: ${paths.coco_gt} + transforms: ${scratch.base_val_transform} + max_ann_per_img: 100000 + multiplier: 1 + training: false + + shuffle: False + batch_size: ${scratch.val_batch_size} + num_workers: ${scratch.num_val_workers} + pin_memory: False + drop_last: False + collate_fn: + _target_: sam3.train.data.collator.collate_fn_api + _partial_: true + repeats: ${scratch.hybrid_repeats} + dict_key: silver_yt1b + + meters: + val: + silver_yt1b: # this key matches the "dict_key" in the dataloader's collate function + cgf1: + _target_: sam3.eval.coco_writer.PredictionDumper + iou_type: "segm" + dump_dir: ${launcher.experiment_log_dir}/dumps/silver_yt1b + merge_predictions: True + postprocessor: ${scratch.mask_postprocessor_thresholded} + gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} + maxdets: 1000000 # no limit + pred_file_evaluators: + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "bbox" + - _target_: sam3.eval.cgf1_eval.CGF1Evaluator + gt_path: ${paths.coco_gt} + iou_type: "segm" diff --git a/src/nn/segearth_ov3/sam3/train/loss/__init__.py b/src/nn/segearth_ov3/sam3/train/loss/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/loss/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/train/loss/loss_fns.py b/src/nn/segearth_ov3/sam3/train/loss/loss_fns.py new file mode 100644 index 0000000..3b61d7f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/loss/loss_fns.py @@ -0,0 +1,1319 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +import warnings + +import torch +import torch.distributed +import torch.nn.functional as F +import torchmetrics + +from sam3.model import box_ops + +from sam3.model.data_misc import interpolate + +from sam3.train.loss.sigmoid_focal_loss import ( + triton_sigmoid_focal_loss, + triton_sigmoid_focal_loss_reduce, +) +from torch import nn + +from .mask_sampling import ( + calculate_uncertainty, + get_uncertain_point_coords_with_randomness, + point_sample, +) + + +CORE_LOSS_KEY = "core_loss" + + +def instance_masks_to_semantic_masks( + instance_masks: torch.Tensor, num_instances: torch.Tensor +) -> torch.Tensor: + """This function converts instance masks to semantic masks. + It accepts a collapsed batch of instances masks (ie all instance masks are concatenated in a single tensor) and + the number of instances in each image of the batch. + It returns a mask with the same spatial dimensions as the input instance masks, where for each batch element the + semantic mask is the union of all the instance masks in the batch element. + + If for a given batch element there are no instances (ie num_instances[i]==0), the corresponding semantic mask will be a tensor of zeros. + + Args: + instance_masks (torch.Tensor): A tensor of shape (N, H, W) where N is the number of instances in the batch. + num_instances (torch.Tensor): A tensor of shape (B,) where B is the batch size. It contains the number of instances + in each image of the batch. + + Returns: + torch.Tensor: A tensor of shape (B, H, W) where B is the batch size and H, W are the spatial dimensions of the + input instance masks. + """ + if num_instances.sum() == 0: + # all negative batch, create a tensor of zeros (B, 1, 1) + return num_instances.unsqueeze(-1).unsqueeze(-1) + + masks_per_query = torch.split(instance_masks, num_instances.tolist()) + + return torch.stack([torch.any(masks, dim=0) for masks in masks_per_query], dim=0) + + +@torch.no_grad() +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + if target.numel() == 0: + return [torch.zeros([], device=output.device)] + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +def dice_loss(inputs, targets, num_boxes, loss_on_multimask=False, reduce=True): + """ + Compute the DICE loss, similar to generalized IOU for masks + Args: + inputs: A float tensor of arbitrary shape. + The predictions for each example. + targets: A float tensor with the same shape as inputs. Stores the binary + classification label for each element in inputs + (0 for the negative class and 1 for the positive class). + """ + try: + loss = _dice_loss(inputs, targets, num_boxes, loss_on_multimask, reduce) + except torch.OutOfMemoryError: + logging.error("GPU OOM, computing dice loss on CPU") + # try to recover from GPU OOM by moving tensors to CPU and computing loss there + orig_device = inputs.device + inputs = inputs.cpu() + targets = targets.cpu() + if isinstance(num_boxes, torch.Tensor): + num_boxes = num_boxes.cpu() + loss = _dice_loss(inputs, targets, num_boxes, loss_on_multimask, reduce) + loss = loss.to(orig_device) + + return loss + + +def _dice_loss(inputs, targets, num_boxes, loss_on_multimask=False, reduce=True): + inputs = inputs.sigmoid() + if loss_on_multimask: + # inputs and targets are [N, M, H, W] where M corresponds to multiple predicted masks + assert inputs.dim() == 4 and targets.dim() == 4 + # flatten spatial dimension while keeping multimask channel dimension + inputs = inputs.flatten(2) + targets = targets.flatten(2) + numerator = 2 * (inputs * targets).sum(-1) + else: + inputs = inputs.flatten(1) + numerator = 2 * (inputs * targets).sum(1) + denominator = inputs.sum(-1) + targets.sum(-1) + loss = 1 - (numerator + 1) / (denominator + 1) + if loss_on_multimask: + return loss / num_boxes + if not reduce: + return loss + return loss.sum() / num_boxes + + +def sigmoid_focal_loss( + inputs, + targets, + num_boxes, + alpha: float = 0.25, + gamma: float = 2, + loss_on_multimask=False, + reduce=True, + triton=True, +): + """ + Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. + Args: + inputs: A float tensor of arbitrary shape. + The predictions for each example. + targets: A float tensor with the same shape as inputs. Stores the binary + classification label for each element in inputs + (0 for the negative class and 1 for the positive class). + alpha: (optional) Weighting factor in range (0,1) to balance + positive vs negative examples. Default = -1 (no weighting). + gamma: Exponent of the modulating factor (1 - p_t) to + balance easy vs hard examples. + Returns: + Loss tensor + """ + if not (0 <= alpha <= 1) and triton: + raise RuntimeError(f"Alpha should be in [0,1], got {alpha}") + if triton: + if reduce and not loss_on_multimask: + loss = triton_sigmoid_focal_loss_reduce(inputs, targets, alpha, gamma) + return loss / (num_boxes * inputs.shape[1]) + + loss = triton_sigmoid_focal_loss(inputs, targets, alpha, gamma) + else: + prob = inputs.sigmoid() + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = prob * targets + (1 - prob) * (1 - targets) + loss = ce_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * loss + + if not reduce: + return loss + + if loss_on_multimask: + # loss is [N, M, H, W] where M corresponds to multiple predicted masks + assert loss.dim() == 4 + return loss.flatten(2).mean(-1) / num_boxes # average over spatial dims + return loss.mean(1).sum() / num_boxes + + +def iou_loss( + inputs, targets, pred_ious, num_boxes, loss_on_multimask=False, use_l1_loss=False +): + """MSE loss between predicted IoUs and actual IoUs between inputs and targets.""" + assert inputs.dim() == 4 and targets.dim() == 4 + pred_mask = inputs.flatten(2) > 0 + gt_mask = targets.flatten(2) > 0 + area_i = torch.sum(pred_mask & gt_mask, dim=-1).float() + area_u = torch.sum(pred_mask | gt_mask, dim=-1).float() + actual_ious = area_i / torch.clamp(area_u, min=1.0) + + if use_l1_loss: + loss = F.l1_loss(pred_ious, actual_ious, reduction="none") + else: + loss = F.mse_loss(pred_ious, actual_ious, reduction="none") + if loss_on_multimask: + return loss / num_boxes + return loss.sum() / num_boxes + + +@torch.jit.script +def _contrastive_align(logits, positive_map): + positive_logits = -logits.masked_fill(~positive_map, 0) + negative_logits = logits # .masked_fill(positive_map, -1000000) + + boxes_with_pos = positive_map.any(2) + pos_term = positive_logits.sum(2) + neg_term = negative_logits.logsumexp(2) + + nb_pos = positive_map.sum(2) + 1e-6 + + box_to_token_loss = ( + (pos_term / nb_pos + neg_term).masked_fill(~boxes_with_pos, 0).sum() + ) + + tokens_with_pos = positive_map.any(1) + pos_term = positive_logits.sum(1) + neg_term = negative_logits.logsumexp(1) + + nb_pos = positive_map.sum(1) + 1e-6 + + tokens_to_boxes_loss = ( + (pos_term / nb_pos + neg_term).masked_fill(~tokens_with_pos, 0).sum() + ) + return (box_to_token_loss + tokens_to_boxes_loss) / 2 + + +def _get_src_permutation_idx(indices): + # permute predictions following indices + batch_idx = torch.cat( + [torch.full_like(src, i) for i, (src, _) in enumerate(indices)] + ) + src_idx = torch.cat([src for (src, _) in indices]) + return batch_idx, src_idx + + +class LossWithWeights(nn.Module): + def __init__(self, weight_dict, compute_aux, supports_o2m_loss=True): + super().__init__() + # weights for each computed loss key (those losses not in weight_dict + # will not be aggregated in the final reduced core loss) + self.weight_dict = weight_dict if weight_dict is not None else {} + # whether this loss will be applied on auxiliary outputs + self.compute_aux = compute_aux + self.supports_o2m_loss = supports_o2m_loss + self.target_keys = [] + + def forward(self, *args, is_aux=False, **kwargs): + if is_aux and not self.compute_aux: + return {CORE_LOSS_KEY: 0.0} + losses = self.get_loss(*args, **kwargs) + losses[CORE_LOSS_KEY] = self.reduce_loss(losses) + return losses + + def get_loss(self, **kwargs): + raise NotImplementedError() + + def reduce_loss(self, losses): + reduced_loss = 0.0 + for loss_key, weight in self.weight_dict.items(): + if loss_key not in losses: + raise ValueError(f"{type(self)} doesn't compute {loss_key}") + if weight != 0: + reduced_loss += losses[loss_key] * weight + + return reduced_loss + + +class IABCEMdetr(LossWithWeights): + def __init__( + self, + pos_weight, + weight_dict=None, + compute_aux=True, + gamma=0, + weak_loss=True, + alpha=0.25, + pad_n_queries=None, + pad_scale_pos=1.0, + use_separate_loss_for_det_and_trk=False, + num_det_queries=None, + det_exhaustive_loss_scale_pos=1.0, + det_exhaustive_loss_scale_neg=1.0, + det_non_exhaustive_loss_scale_pos=1.0, + det_non_exhaustive_loss_scale_neg=1.0, + trk_loss_scale_pos=1.0, + trk_loss_scale_neg=1.0, + no_loss_for_fp_propagation=False, + apply_loss_to_det_queries_in_video_grounding=True, + use_presence=False, + use_presence_semgseg=False, # If True, use presence scores from the semgseg head. + presence_alpha=0.5, + presence_gamma=0.0, + pos_focal: bool = False, # for box scores, use focal loss for positives as well + ): + super().__init__(weight_dict, compute_aux) + self.pos_weight = pos_weight + self.gamma = gamma + self.weak_loss = weak_loss + self.alpha = alpha + self.target_keys.append("boxes_xyxy") + self.no_loss_for_fp_propagation = no_loss_for_fp_propagation + if self.weak_loss: + self.target_keys.append("is_exhaustive") + # NOTE: This is hacky solution to have the same CE loss scale across datasets where the model might predict different number of object queries for different tasks. + # If not None, we assume there are a total pad_n_queries object queries. + # For example, if the model predicts only 1 object query and pad_n_queries=100, we pad the predictions with 99 zero preds. + # Currently this only affects the BCE loss and not the F1 score. + self.pad_n_queries = pad_n_queries + self.pad_scale_pos = pad_scale_pos + if self.pad_scale_pos != 1.0: + assert self.pad_n_queries is not None + # whether to use presence scores + self.use_presence = use_presence + self.use_presence_semgseg = use_presence_semgseg + if self.use_presence_semgseg: + assert self.use_presence + self.presence_alpha = presence_alpha + self.presence_gamma = presence_gamma + self.pos_focal = pos_focal + + # Decoupled loss for detection and tracking queries + self.apply_loss_to_det_queries_in_video_grounding = ( + apply_loss_to_det_queries_in_video_grounding + ) + self.use_separate_loss_for_det_and_trk = use_separate_loss_for_det_and_trk + if num_det_queries is not None: + logging.warning("note: it's not needed to set num_det_queries anymore") + if self.use_separate_loss_for_det_and_trk: + assert not self.weak_loss, "Do not use weak_loss in this case -- set separate loss for detection and tracking queries instead" + self.det_exhaustive_loss_scale_pos = det_exhaustive_loss_scale_pos + self.det_exhaustive_loss_scale_neg = det_exhaustive_loss_scale_neg + self.det_non_exhaustive_loss_scale_pos = det_non_exhaustive_loss_scale_pos + self.det_non_exhaustive_loss_scale_neg = det_non_exhaustive_loss_scale_neg + self.trk_loss_scale_pos = trk_loss_scale_pos + self.trk_loss_scale_neg = trk_loss_scale_neg + else: + assert ( + det_exhaustive_loss_scale_pos == 1.0 + and det_exhaustive_loss_scale_neg == 1.0 + and det_non_exhaustive_loss_scale_pos == 1.0 + and det_non_exhaustive_loss_scale_neg == 1.0 + and trk_loss_scale_pos == 1.0 + and trk_loss_scale_neg == 1.0 + ), "If not using separate loss for detection and tracking queries, separate detection and tracking loss scales should all be 1.0" + + def get_loss(self, outputs, targets, indices, num_boxes): + assert len(outputs["pred_logits"].shape) > 2, "Incorrect predicted logits shape" + assert outputs["pred_logits"].shape[-1] == 1, "Incorrect predicted logits shape" + src_logits = outputs["pred_logits"].squeeze(-1) + prob = src_logits.sigmoid() + + with torch.no_grad(): + target_classes = torch.full( + src_logits.shape[:2], + 0, + dtype=torch.float, + device=src_logits.device, + ) + target_classes[(indices[0], indices[1])] = 1 + src_boxes_xyxy = outputs["pred_boxes_xyxy"][(indices[0], indices[1])] + target_boxes_giou = ( + targets["boxes_xyxy"][indices[2]] + if indices[2] is not None + else targets["boxes_xyxy"] + ) + + iou = box_ops.fast_diag_box_iou(src_boxes_xyxy, target_boxes_giou) + t = prob[(indices[0], indices[1])] ** self.alpha * iou ** (1 - self.alpha) + t = torch.clamp(t, 0.01).detach() + positive_target_classes = target_classes.clone() + positive_target_classes[(indices[0], indices[1])] = t + + # Soft loss on positives + if self.pos_focal: + loss_bce = sigmoid_focal_loss( + src_logits.contiguous(), + positive_target_classes, + num_boxes=1, + alpha=0.5, + gamma=self.gamma, + reduce=False, + ) + else: + loss_bce = F.binary_cross_entropy_with_logits( + src_logits, positive_target_classes, reduction="none" + ) + loss_bce = loss_bce * target_classes * self.pos_weight + + if ( + self.pad_n_queries is not None + and isinstance(self.pad_n_queries, int) + and loss_bce.size(1) < self.pad_n_queries + ): + loss_bce = loss_bce * self.pad_scale_pos + # Negatives + loss_bce = loss_bce + F.binary_cross_entropy_with_logits( + src_logits, target_classes, reduction="none" + ) * (1 - target_classes) * (prob**self.gamma) + + # Optionally, not applying IABCEMdetr loss to detection queries in video. + is_video_grounding = outputs.get("is_video_grounding_batch", False) + if is_video_grounding and not self.apply_loss_to_det_queries_in_video_grounding: + Q_det = outputs["Q_det"] + loss_bce[:, :Q_det] *= 0.0 + presence_loss = torch.tensor(0.0, device=src_logits.device) + presence_dec_acc = torch.tensor(0.0, device=src_logits.device) + if self.use_presence: + # no classifiction loss for individual tokens if no target gt + # cannot directly use targets["num_boxes"] to check if some + # GT box exists as there may be dummy boxes for "invisible objects" + # in video grounding data + + gt_padded_object_ids = targets["object_ids_padded"] # (B, H) + gt_padded_boxes = targets["boxes_padded"] # (B, H, 4) shape, CxCyWH + gt_padded_is_visible = ( + (gt_padded_object_ids >= 0) + & (gt_padded_boxes[..., 2] > 0) # width > 0 + & (gt_padded_boxes[..., 3] > 0) # height > 0 + ) + keep_loss = (gt_padded_is_visible.sum(dim=-1)[..., None] != 0).float() + + loss_bce = loss_bce * keep_loss + + if self.use_presence_semgseg: + # no loss here, has it's own separate loss computation + assert "presence_logit_dec" not in outputs + elif "presence_logit_dec" in outputs: + presence_logits = outputs["presence_logit_dec"].view_as(keep_loss) + bs = presence_logits.shape[0] + presence_loss = sigmoid_focal_loss( + presence_logits, + keep_loss, + # not num_boxes, but we'll use it to normalize by bs + num_boxes=bs, + alpha=self.presence_alpha, + gamma=self.presence_gamma, + ) + pred = (presence_logits.sigmoid() > 0.5).float() + presence_dec_acc = (pred == keep_loss).float().mean() + else: + # for o2m, nothing to do + pass + + if self.weak_loss: + assert not self.use_separate_loss_for_det_and_trk, "Do not use weak_loss in this case -- set separate loss for detection and tracking queries instead" + + # nullify the negative loss for the non-exhaustive classes + assert loss_bce.shape[0] == targets["is_exhaustive"].shape[0] + assert targets["is_exhaustive"].ndim == 1 + + loss_mask = (~targets["is_exhaustive"]).view(-1, 1).expand_as(loss_bce) + # restrict the mask to the negative supervision + loss_mask = loss_mask & (target_classes < 0.5) + loss_mask = ~loss_mask + # Mask the loss + loss_bce = loss_bce * loss_mask.float() + # Average + loss_bce = loss_bce.sum() / (loss_mask.sum() + 1e-6) + else: + # apply separate loss weights to detection and tracking queries + if self.use_separate_loss_for_det_and_trk: + Q_det = outputs["Q_det"] + assert loss_bce.size(1) >= Q_det + is_positive = target_classes > 0.5 + is_positive_det = is_positive[:, :Q_det] + is_positive_trk = is_positive[:, Q_det:] + assert loss_bce.size(0) == targets["is_exhaustive"].size(0) + is_exhaustive = targets["is_exhaustive"].unsqueeze(1).bool() + loss_scales = torch.zeros_like(loss_bce) + # detection query loss weights + loss_scales[:, :Q_det] = ( + (is_exhaustive & is_positive_det).float() + * self.det_exhaustive_loss_scale_pos + + (is_exhaustive & ~is_positive_det).float() + * self.det_exhaustive_loss_scale_neg + + (~is_exhaustive & is_positive_det).float() + * self.det_non_exhaustive_loss_scale_pos + + (~is_exhaustive & ~is_positive_det).float() + * self.det_non_exhaustive_loss_scale_neg + ) + # tracking query weights + loss_scales[:, Q_det:] = ( + is_positive_trk.float() * self.trk_loss_scale_pos + + (~is_positive_trk).float() * self.trk_loss_scale_neg + ) + # apply the loss weights + + # if the id is -2 means it is a fp propagation , we don't apply the loss to them + if self.no_loss_for_fp_propagation: + is_original_queries = outputs["pred_old_obj_ids"] != -2 + loss_scales *= (is_exhaustive | is_original_queries).float() + + loss_bce = loss_bce * loss_scales + + if self.pad_n_queries is None or loss_bce.size(1) >= self.pad_n_queries: + loss_bce = loss_bce.mean() + else: + assert isinstance(self.pad_n_queries, int) + assert ( + loss_bce.size(1) < self.pad_n_queries + ), f"The number of predictions is more than the expected total after padding. Got {loss_bce.size(1)} predictions." + loss_bce = loss_bce.sum() / (self.pad_n_queries * loss_bce.size(0)) + + bce_f1 = torchmetrics.functional.f1_score( + src_logits.sigmoid().flatten(), + target=target_classes.flatten().long(), + task="binary", + ) + + losses = { + "loss_ce": loss_bce, + "ce_f1": bce_f1, + "presence_loss": presence_loss, + "presence_dec_acc": presence_dec_acc, + } + return losses + + +class Boxes(LossWithWeights): + def __init__( + self, + weight_dict=None, + compute_aux=True, + apply_loss_to_det_queries_in_video_grounding=True, + ): + super().__init__(weight_dict, compute_aux) + self.apply_loss_to_det_queries_in_video_grounding = ( + apply_loss_to_det_queries_in_video_grounding + ) + self.target_keys.extend(["boxes", "boxes_xyxy"]) + + def get_loss(self, outputs, targets, indices, num_boxes): + """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss + targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4] + The target boxes are expected in format (center_x, center_y, h, w), normalized by the image size. + """ + # Optionally, not applying Boxes loss to detection queries in video. + is_video_grounding = outputs.get("is_video_grounding_batch", False) + if is_video_grounding and not self.apply_loss_to_det_queries_in_video_grounding: + indices = _keep_only_trk_queries_in_match_inds( + indices, Q_det=outputs["Q_det"] + ) + + assert "pred_boxes" in outputs + # idx = self._get_src_permutation_idx(indices) + src_boxes = outputs["pred_boxes"][(indices[0], indices[1])] + src_boxes_xyxy = outputs["pred_boxes_xyxy"][(indices[0], indices[1])] + target_boxes = ( + targets["boxes"] if indices[2] is None else targets["boxes"][indices[2]] + ) + target_boxes_giou = ( + targets["boxes_xyxy"] + if indices[2] is None + else targets["boxes_xyxy"][indices[2]] + ) + + loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction="none") + + losses = {} + losses["loss_bbox"] = loss_bbox.sum() / num_boxes + + loss_giou = 1 - box_ops.fast_diag_generalized_box_iou( + src_boxes_xyxy, target_boxes_giou + ) + losses["loss_giou"] = loss_giou.sum() / num_boxes + return losses + + +class Masks(LossWithWeights): + def __init__( + self, + weight_dict=None, + compute_aux=False, + focal_alpha=0.25, + focal_gamma=2, + num_sample_points=None, + oversample_ratio=None, + importance_sample_ratio=None, + apply_loss_to_det_queries_in_video_grounding=True, + ): + super().__init__(weight_dict, compute_aux) + if compute_aux: + warnings.warn("Masks loss usually shouldn't be applied to aux outputs") + self.focal_alpha = focal_alpha + self.focal_gamma = focal_gamma + self.num_sample_points = num_sample_points + self.oversample_ratio = oversample_ratio + self.importance_sample_ratio = importance_sample_ratio + self.apply_loss_to_det_queries_in_video_grounding = ( + apply_loss_to_det_queries_in_video_grounding + ) + self.target_keys.extend(["masks", "is_valid_mask"]) + + def _sampled_loss(self, src_masks, target_masks, num_boxes): + assert len(src_masks.shape) == 3 and len(target_masks.shape) == 3 + src_masks = src_masks[:, None] + target_masks = target_masks[:, None] + with torch.no_grad(): + # Sample point_coords + point_coords = get_uncertain_point_coords_with_randomness( + src_masks, + calculate_uncertainty, + self.num_sample_points, + self.oversample_ratio, + self.importance_sample_ratio, + ) + + # get GT labels + sampled_target_masks = point_sample( + target_masks, + point_coords, + align_corners=False, + ).squeeze(1) + + sampled_src_masks = point_sample( + src_masks, + point_coords, + align_corners=False, + ).squeeze(1) + + losses = { + "loss_mask": sigmoid_focal_loss( + sampled_src_masks, + sampled_target_masks, + num_boxes, + alpha=self.focal_alpha, + gamma=self.focal_gamma, + ), + "loss_dice": dice_loss(sampled_src_masks, sampled_target_masks, num_boxes), + } + # Not needed for backward + del src_masks + del target_masks + + return losses + + def get_loss(self, outputs, targets, indices, num_boxes): + """Compute the losses related to the masks: the focal loss and the dice loss. + targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w] + """ + assert "pred_masks" in outputs + assert "is_valid_mask" in targets + # Optionally, not applying Masks loss to detection queries in video. + is_video_grounding = outputs.get("is_video_grounding_batch", False) + if is_video_grounding and not self.apply_loss_to_det_queries_in_video_grounding: + indices = _keep_only_trk_queries_in_match_inds( + indices, Q_det=outputs["Q_det"] + ) + + src_masks = outputs["pred_masks"] + + # Dataset doesn't have segmentation masks + if targets["masks"] is None: + return { + "loss_mask": torch.tensor(0.0, device=src_masks.device), + "loss_dice": torch.tensor(0.0, device=src_masks.device), + } + + target_masks = ( + targets["masks"] if indices[2] is None else targets["masks"][indices[2]] + ) + target_masks = target_masks.to(src_masks) + keep = ( + targets["is_valid_mask"] + if indices[2] is None + else targets["is_valid_mask"][indices[2]] + ) + + src_masks = src_masks[(indices[0], indices[1])] + + # Remove invalid masks from loss + src_masks = src_masks[keep] + target_masks = target_masks[keep] + + if self.num_sample_points is not None: + # Compute loss on sampled points for the Mask + losses = self._sampled_loss(src_masks, target_masks, num_boxes) + + else: + # upsample predictions to the target size + if target_masks.shape[0] == 0 and src_masks.shape[0] == 0: + src_masks = src_masks.flatten(1) + target_masks = target_masks.reshape(src_masks.shape) + else: + if len(src_masks.shape) == 3: + src_masks = src_masks[:, None] + if src_masks.dtype == torch.bfloat16: + # Bilinear interpolation does not support bf16 + src_masks = src_masks.to(dtype=torch.float32) + src_masks = interpolate( + src_masks, + size=target_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ) + src_masks = src_masks[:, 0].flatten(1) + target_masks = target_masks.flatten(1) + + losses = { + "loss_mask": sigmoid_focal_loss( + src_masks, + target_masks, + num_boxes, + alpha=self.focal_alpha, + gamma=self.focal_gamma, + ), + "loss_dice": dice_loss(src_masks, target_masks, num_boxes), + } + + return losses + + +# class MultiStepIteractiveMasks(LossWithWeights): +# def __init__( +# self, +# weight_dict=None, +# compute_aux=False, +# focal_alpha=0.25, +# focal_gamma=2, +# ): +# warnings.warn( +# "MultiStepIteractiveMasks is deprecated. Please use MultiStepMultiMasksAndIous", +# DeprecationWarning, +# ) +# super().__init__(weight_dict, compute_aux) +# self.focal_alpha = focal_alpha +# self.focal_gamma = focal_gamma +# self.target_keys.extend(["masks"]) + +# def get_loss(self, outputs, targets, indices, num_boxes): +# """Compute the losses related to the masks: the focal loss and the dice loss. +# targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w] + +# Unlike `Masks`, here the "multistep_pred_masks" can have multiple channels, each +# corresponding to one iterative prediction step in SAM-style training. We treat each +# channel as a mask prediction and sum the loss across channels. +# """ +# src_masks = outputs["multistep_pred_masks"] +# target_masks = targets["masks"] +# assert src_masks.size(0) == target_masks.size(0) +# assert src_masks.dim() == 4 +# assert target_masks.dim() == 3 + +# # tile target_masks according to the number of +# # channels `src_masks`. +# num_steps = src_masks.size(1) +# target_masks = target_masks.unsqueeze(1).to(src_masks.dtype) +# if num_steps > 1: +# target_masks = target_masks.repeat(1, num_steps, 1, 1) + +# # resize `src_masks` to target mask resolution +# if src_masks.shape != target_masks.shape: +# src_masks = interpolate( +# src_masks, +# size=target_masks.shape[-2:], +# mode="bilinear", +# align_corners=False, +# ) +# assert src_masks.shape == target_masks.shape + +# # flatten the multiple steps in to the batch dimension +# src_masks = src_masks.flatten(0, 1).flatten(1) +# target_masks = target_masks.flatten(0, 1).flatten(1) +# losses = { +# "loss_mask": sigmoid_focal_loss( +# src_masks, +# target_masks, +# num_boxes, +# alpha=self.focal_alpha, +# gamma=self.focal_gamma, +# ), +# "loss_dice": dice_loss(src_masks, target_masks, num_boxes), +# } + +# return losses + + +# class MultiStepMultiMasksAndIous(LossWithWeights): +# def __init__( +# self, +# weight_dict=None, +# compute_aux=False, +# focal_alpha=0.25, +# focal_gamma=2, +# # if True, back-prop on all predicted ious +# # not just the one with lowest loss_combo +# supervise_all_iou=False, +# # Less slack vs MSE loss in [-1, 1] error range +# iou_use_l1_loss=False, +# # Settings for obj score prediction +# pred_obj_scores=False, +# focal_gamma_obj_score=0.0, +# focal_alpha_obj_score=-1, +# ): +# super().__init__(weight_dict, compute_aux) +# self.focal_alpha = focal_alpha +# self.focal_gamma = focal_gamma +# self.target_keys.extend(["masks"]) +# assert "loss_mask" in self.weight_dict +# assert "loss_dice" in self.weight_dict +# assert "loss_iou" in self.weight_dict +# if "loss_class" not in self.weight_dict: +# self.weight_dict["loss_class"] = 0.0 +# self.focal_alpha_obj_score = focal_alpha_obj_score +# self.focal_gamma_obj_score = focal_gamma_obj_score +# self.supervise_all_iou = supervise_all_iou +# self.iou_use_l1_loss = iou_use_l1_loss +# self.pred_obj_scores = pred_obj_scores + +# def get_loss(self, outputs, targets, indices, num_boxes): +# """ +# Compute the losses related to the masks: the focal loss and the dice loss. +# and also the MSE loss between predicted IoUs and actual IoUs. + +# Here "multistep_pred_multimasks_high_res" is a list of multimasks (tensors +# of shape [N, M, H, W], where M could be 1 or larger, corresponding to +# one or multiple predicted masks from a click. + +# We back-propagate focal, dice and iou losses only on the prediction channel +# with the lowest focal+dice loss between predicted mask and ground-truth. +# """ + +# target_masks = targets["masks"].unsqueeze(1).float() +# assert target_masks.dim() == 4 # [N, 1, H, W] +# src_masks_list = outputs["multistep_pred_multimasks_high_res"] +# ious_list = outputs["multistep_pred_ious"] +# object_score_logits_list = outputs["multistep_object_score_logits"] + +# assert len(src_masks_list) == len(ious_list) +# assert len(object_score_logits_list) == len(ious_list) + +# # Remove invalid masks from loss +# keep = targets["is_valid_mask"] +# target_masks = target_masks[keep] + +# # accumulate the loss over prediction steps +# losses = {"loss_mask": 0, "loss_dice": 0, "loss_iou": 0, "loss_class": 0} +# for src_masks, ious, object_score_logits in zip( +# src_masks_list, ious_list, object_score_logits_list +# ): +# object_score_logits = object_score_logits[keep] +# ious = ious[keep] +# src_masks = src_masks[keep] +# self._update_losses( +# losses, src_masks, target_masks, ious, num_boxes, object_score_logits +# ) +# return losses + +# def _update_losses( +# self, losses, src_masks, target_masks, ious, num_boxes, object_score_logits +# ): +# target_masks = target_masks.expand_as(src_masks) +# # get focal, dice and iou loss on all output masks in a prediction step +# loss_multimask = sigmoid_focal_loss( +# src_masks, +# target_masks, +# num_boxes, +# alpha=self.focal_alpha, +# gamma=self.focal_gamma, +# loss_on_multimask=True, +# triton=False, # only use triton if alpha > 0 +# ) +# loss_multidice = dice_loss( +# src_masks, target_masks, num_boxes, loss_on_multimask=True +# ) +# if not self.pred_obj_scores: +# loss_class = torch.tensor( +# 0.0, dtype=loss_multimask.dtype, device=loss_multimask.device +# ) +# target_obj = torch.ones( +# loss_multimask.shape[0], +# 1, +# dtype=loss_multimask.dtype, +# device=loss_multimask.device, +# ) +# else: +# target_obj = torch.any((target_masks[:, 0] > 0).flatten(1), dim=-1)[ +# ..., None +# ].float() +# loss_class = sigmoid_focal_loss( +# object_score_logits, +# target_obj, +# num_boxes, +# alpha=self.focal_alpha_obj_score, +# gamma=self.focal_gamma_obj_score, +# triton=False, +# ) + +# loss_multiiou = iou_loss( +# src_masks, +# target_masks, +# ious, +# num_boxes, +# loss_on_multimask=True, +# use_l1_loss=self.iou_use_l1_loss, +# ) +# assert loss_multimask.dim() == 2 +# assert loss_multidice.dim() == 2 +# assert loss_multiiou.dim() == 2 +# if loss_multimask.size(1) > 1: +# # take the mask indices with the smallest focal + dice loss for back propagation +# loss_combo = ( +# loss_multimask * self.weight_dict["loss_mask"] +# + loss_multidice * self.weight_dict["loss_dice"] +# ) +# best_loss_inds = torch.argmin(loss_combo, dim=-1) +# batch_inds = torch.arange(loss_combo.size(0), device=loss_combo.device) +# loss_mask = loss_multimask[batch_inds, best_loss_inds].unsqueeze(1) +# loss_dice = loss_multidice[batch_inds, best_loss_inds].unsqueeze(1) +# # calculate the iou prediction and slot losses only in the index +# # with the minimum loss for each mask (to be consistent w/ SAM) +# if self.supervise_all_iou: +# loss_iou = loss_multiiou.mean(dim=-1).unsqueeze(1) +# else: +# loss_iou = loss_multiiou[batch_inds, best_loss_inds].unsqueeze(1) +# else: +# loss_mask = loss_multimask +# loss_dice = loss_multidice +# loss_iou = loss_multiiou + +# # backprop focal, dice and iou loss only if obj present +# loss_mask = loss_mask * target_obj +# loss_dice = loss_dice * target_obj +# loss_iou = loss_iou * target_obj + +# # sum over batch dimension (note that the losses are already divided by num_boxes) +# losses["loss_mask"] += loss_mask.sum() +# losses["loss_dice"] += loss_dice.sum() +# losses["loss_iou"] += loss_iou.sum() +# losses["loss_class"] += loss_class + + +# class TextCriterion(LossWithWeights): +# def __init__( +# self, +# pad_token, +# max_seq_len=100, +# weight_dict=None, +# compute_aux=False, +# ): +# super().__init__(weight_dict, compute_aux) +# self.pad_token = pad_token +# self.max_seq_len = max_seq_len +# self.in_lengths = None + +# def get_loss(self, outputs, **kwargs): +# nb_tokens = outputs["captioning_tokenized_target"].input_ids.numel() +# bs, seq_len = outputs["captioning_tokenized_target"].input_ids.shape +# ce = F.cross_entropy( +# outputs["captioning_pred_text"].flatten(0, -2), +# outputs["captioning_tokenized_target"].input_ids.flatten(), +# ignore_index=self.pad_token, +# reduction="sum", +# ) + +# not_pad = ( +# outputs["captioning_tokenized_target"] +# .input_ids.reshape(-1) +# .ne(self.pad_token) +# ) + +# if nb_tokens > 0: +# nb_non_pad = not_pad.numel() +# ce = ce / nb_non_pad + +# preds = outputs["captioning_pred_text"].flatten(0, -2).argmax(-1)[not_pad] +# targets = outputs["captioning_tokenized_target"].input_ids.flatten()[not_pad] +# correct = preds == targets +# correct = correct.sum() / (correct.numel() + 1e-5) + +# correct_sequence_level = torch.all( +# ( +# outputs["captioning_pred_text"] +# .flatten(0, -2) +# .argmax(-1) +# .reshape(bs, seq_len) +# == outputs["captioning_tokenized_target"].input_ids +# ) +# | (~not_pad).view(bs, seq_len), +# dim=1, +# ) +# seq_level_acc = correct_sequence_level.float().mean() + +# return {"loss_text": ce, "text_acc": correct, "text_seq_acc": seq_level_acc} + + +def segment_miou(source, target): + """Compute the mean IoU between two sets of masks""" + assert source.shape == target.shape, "The two masks must have the same shape" + assert source.ndim == 3, "The masks must be 3D" + + valid_targets = (target.sum(dim=(1, 2)) > 0).sum() + if valid_targets == 0: + return torch.tensor(1.0, device=source.device) + intersection = (source.bool() & target.bool()).sum(dim=(1, 2)) + union = (source.bool() | target.bool()).sum(dim=(1, 2)) + iou = intersection / (union + 1e-8) + return iou.sum() / valid_targets + + +class SemanticSegCriterion(LossWithWeights): + def __init__( + self, + weight_dict, + focal: bool = False, + focal_alpha: float = 0.6, + focal_gamma: float = 1.6, + downsample: bool = True, + presence_head: bool = False, + # Option to turn off presence loss, if some other component + # is already doing it, e.g. decoder - in which case, + # we could still set presence_head to True so that + # losses are not propogated to masks when there is no GT mask + presence_loss: bool = True, + ): + super().__init__(weight_dict, False) + self.focal = focal + self.focal_alpha = focal_alpha + self.focal_gamma = focal_gamma + self.downsample = downsample + self.presence_head = presence_head + self.presence_loss = presence_loss + + def get_loss(self, out_dict, targets): + outputs = out_dict["semantic_seg"] + presence_logit = out_dict["presence_logit"] + if ( + "semantic_masks" in targets + and targets["semantic_masks"] is not None + and targets["semantic_masks"].size(0) > 0 + ): + semantic_targets = targets["semantic_masks"] + with torch.no_grad(): + if self.downsample: + # downsample targets to the size of predictions + size = outputs.shape[-2:] + semantic_targets = ( + F.interpolate( + semantic_targets.float().unsqueeze(1), + size=size, + mode="bilinear", + align_corners=False, + ) + .squeeze(1) + .bool() + ) + else: + with torch.no_grad(): + if self.downsample: + # downsample targets to the size of predictions + size = outputs.shape[-2:] + segments = ( + F.interpolate( + targets["masks"].float().unsqueeze(1), + size=size, + mode="bilinear", + align_corners=False, + ) + .squeeze(1) + .bool() + ) + else: + segments = targets["masks"].bool() + + # the annotations are for instance segmentation, so we merge them to get semantic segmentation + semantic_targets = instance_masks_to_semantic_masks( + segments, targets["num_boxes"] + ) + + if not self.downsample: + # upsample predictions to the target size + size = semantic_targets.shape[-2:] + outputs = F.interpolate( + outputs.float(), + size=size, + mode="bilinear", + align_corners=False, + ) + + if self.focal: + loss = sigmoid_focal_loss( + outputs.squeeze(1).flatten(-2), + semantic_targets.float().flatten(-2), + num_boxes=len(semantic_targets), + alpha=self.focal_alpha, + gamma=self.focal_gamma, + reduce=not self.presence_head, + ) + if self.presence_head: + loss = loss.mean(1) + else: + loss = F.binary_cross_entropy_with_logits( + outputs.squeeze(1), + semantic_targets.float(), + reduction="none" if self.presence_head else "mean", + ) + if self.presence_head: + loss = loss.flatten(1).mean(1) + + loss_dice = dice_loss( + outputs.squeeze(1).flatten(1), + semantic_targets.flatten(1), + len(semantic_targets), + reduce=not self.presence_head, + ) + + miou = segment_miou(outputs.sigmoid().squeeze(1) > 0.5, semantic_targets) + + loss_dict = {} + + if self.presence_head: + presence_target = semantic_targets.flatten(1).any(-1) + if self.presence_loss: + loss_presence = F.binary_cross_entropy_with_logits( + presence_logit.flatten(), + presence_target.float(), + ) + presence_acc = ( + ((presence_logit.flatten().sigmoid() > 0.5) == presence_target) + .float() + .mean() + ) + else: + # Dummy values + loss_presence = torch.tensor(0.0, device=loss.device) + # Whichever component is computing the presence loss, + # should also track presence_acc + presence_acc = torch.tensor(0.0, device=loss.device) + + loss_dict["loss_semantic_presence"] = loss_presence + loss_dict["presence_acc"] = presence_acc + + # reduce the other losses, skipping the negative ones + bs = loss.shape[0] + assert presence_target.numel() == bs + + mask = presence_target + nb_valid = presence_target.sum().item() + + loss = (loss * mask.float()).sum() / (nb_valid + 1e-6) + loss_dice = (loss_dice * mask.float()).sum() / (nb_valid + 1e-6) + + loss_dict.update( + { + "loss_semantic_seg": loss, + "loss_semantic_dice": loss_dice, + "miou_semantic_seg": miou, + } + ) + + return loss_dict + + +class Det2TrkAssoc(LossWithWeights): + def __init__( + self, + weight_dict, + use_fp_loss=False, + fp_loss_on_exhaustive_only=True, + treat_fp_as_new_obj=False, + ): + super().__init__(weight_dict, compute_aux=False) + self.use_fp_loss = use_fp_loss + self.fp_loss_on_exhaustive_only = fp_loss_on_exhaustive_only + self.treat_fp_as_new_obj = treat_fp_as_new_obj + if self.use_fp_loss: + self.target_keys.append("is_exhaustive") + + def get_loss(self, outputs, targets, indices, num_boxes): + det2trk_assoc_logits = outputs["det2trk_assoc_logits"] + device = det2trk_assoc_logits.device + B, Q_det, Q_trk_plus_2 = det2trk_assoc_logits.shape + assert Q_trk_plus_2 >= 2 + Q_trk = Q_trk_plus_2 - 2 + + # We only apply association losses to those detection queries that either match + # a GT instance or have score > 0 (i.e. those TP, FN and FP detection queries) + matched_object_ids = outputs["matched_object_ids"] + assert matched_object_ids.shape == (B, Q_det + Q_trk) + matched_obj_ids_det = matched_object_ids[:, :Q_det] + matched_obj_ids_trk = matched_object_ids[:, Q_det:] + det_is_matched_to_gt = matched_obj_ids_det >= 0 + trk_is_matched_to_gt = matched_obj_ids_trk >= 0 + + # note: -1 label is ignored in the (softmax) cross_entropy loss below + det2trk_assoc_labels = -torch.ones(B, Q_det, dtype=torch.long, device=device) + # a) If a detection query is matched to a same object ID as a tracking query, + # we assign it the index of the tracking query as a label + det_is_same_obj_id_as_trk = ( + det_is_matched_to_gt[:, :, None] + & trk_is_matched_to_gt[:, None, :] + & (matched_obj_ids_det[:, :, None] == matched_obj_ids_trk[:, None, :]) + ) + batch_idx, det_idx, trk_idx = det_is_same_obj_id_as_trk.nonzero(as_tuple=True) + det2trk_assoc_labels[batch_idx, det_idx] = trk_idx + + # b) If a detection query is matched to GT but not to any tracking query, + # we assign it a "new_object" label + det_is_new_obj = det_is_matched_to_gt & ~det_is_same_obj_id_as_trk.any(dim=-1) + det2trk_assoc_labels[det_is_new_obj] = Q_trk + + # c) If a detection query is not matched to GT but have score > 0, + # we assign it a "false_positive" label + if self.use_fp_loss: + det_is_above_thresh = outputs["pred_logits"][:, :Q_det].squeeze(2) > 0 + det_is_fp = ~det_is_matched_to_gt & det_is_above_thresh + if self.treat_fp_as_new_obj: + det2trk_assoc_labels[det_is_fp] = Q_trk + else: + if self.fp_loss_on_exhaustive_only: + # only count FP detections on batches that are exhaustively annotated + det_is_fp &= targets["is_exhaustive"].unsqueeze(1).bool() + det2trk_assoc_labels[det_is_fp] = Q_trk + 1 + + # softmax cross-entropy loss for detection-to-tracking association + loss_det2trk_assoc = F.cross_entropy( + input=det2trk_assoc_logits.flatten(0, 1), # (B * Q_det, Q_trk + 2) + target=det2trk_assoc_labels.flatten(0, 1), # (B * Q_det) + ignore_index=-1, + reduction="none", + ).view(B, Q_det) + # skip det2trk assocation loss on frames w/o any (non-padding) tracking queries + frame_has_valid_trk = trk_is_matched_to_gt.any(dim=-1, keepdims=True) # (B, 1) + loss_det2trk_assoc = loss_det2trk_assoc * frame_has_valid_trk.float() + + loss_det2trk_assoc = loss_det2trk_assoc.sum() / (B * num_boxes) + return {"loss_det2trk_assoc": loss_det2trk_assoc} + + +class TrackingByDetectionAssoc(LossWithWeights): + def __init__(self, weight_dict): + super().__init__(weight_dict, compute_aux=False, supports_o2m_loss=False) + assert "loss_det2trk_assoc" in self.weight_dict + assert "loss_trk2det_assoc" in self.weight_dict + + def get_loss(self, outputs, targets, indices, num_boxes): + # Part A: gather object id matching between detection and tracking + det2trk_assoc_logits = outputs["det2trk_assoc_logits"] # (B, Q_det+1, Q_trk+1) + B, Q_det_plus_1, Q_trk_plus_1 = det2trk_assoc_logits.shape + assert Q_det_plus_1 >= 1 and Q_trk_plus_1 >= 1 + Q_det = Q_det_plus_1 - 1 + Q_trk = Q_trk_plus_1 - 1 + device = det2trk_assoc_logits.device + + matched_obj_ids_det = outputs["matched_object_ids"] + assert matched_obj_ids_det.shape == (B, Q_det) + det_is_matched_to_gt = matched_obj_ids_det >= 0 + matched_obj_ids_trk = outputs["prev_trk_object_ids"] + assert matched_obj_ids_trk.shape == (B, Q_trk) + trk_is_matched_to_gt = matched_obj_ids_trk >= 0 + frame_has_valid_trk = trk_is_matched_to_gt.any(dim=-1, keepdims=True) # (B, 1) + + # check whether a detection object is the same as a tracking object + det_is_same_obj_id_as_trk = ( + det_is_matched_to_gt[:, :, None] + & trk_is_matched_to_gt[:, None, :] + & (matched_obj_ids_det[:, :, None] == matched_obj_ids_trk[:, None, :]) + ) # (B, Q_det, Q_trk) + # there should be at most one match for each detection and each previous tracked object + torch._assert_async(torch.all(det_is_same_obj_id_as_trk.sum(dim=2) <= 1)) + torch._assert_async(torch.all(det_is_same_obj_id_as_trk.sum(dim=1) <= 1)) + batch_idx, det_idx, trk_idx = det_is_same_obj_id_as_trk.nonzero(as_tuple=True) + + # Part B: Detection-to-tracking association loss + # assign detection-to-tracking labels (note: -1 label is ignored in the loss below) + det2trk_assoc_labels = -torch.ones(B, Q_det, dtype=torch.long, device=device) + det2trk_assoc_labels[batch_idx, det_idx] = trk_idx + # if a detection is matched to GT but not to any tracking, assign it a "new-object" label + det_is_new_obj = det_is_matched_to_gt & ~det_is_same_obj_id_as_trk.any(dim=2) + det2trk_assoc_labels[det_is_new_obj] = Q_trk # "Q_trk" label is "new-object" + + # softmax cross-entropy loss for detection-to-tracking association + loss_det2trk_assoc = F.cross_entropy( + input=det2trk_assoc_logits[:, :-1].flatten(0, 1), # (B*Q_det, Q_trk+1) + target=det2trk_assoc_labels.flatten(0, 1), # (B*Q_det) + ignore_index=-1, + reduction="none", + ).view(B, Q_det) + # skip det2trk assocation loss on frames w/o any (non-padding) tracking queries + loss_det2trk_assoc = loss_det2trk_assoc * frame_has_valid_trk.float() + loss_det2trk_assoc = loss_det2trk_assoc.sum() / (B * num_boxes) + loss_dict = {"loss_det2trk_assoc": loss_det2trk_assoc} + + # Part C: tracking-to-detection association loss + trk2det_assoc_logits = det2trk_assoc_logits.transpose(1, 2) + assert trk2det_assoc_logits.shape == (B, Q_trk + 1, Q_det + 1) + # assign tracking-to-detection labels (note: -1 label is ignored in the loss below) + trk2det_assoc_labels = -torch.ones(B, Q_trk, dtype=torch.long, device=device) + trk2det_assoc_labels[batch_idx, trk_idx] = det_idx + # if a tracking is matched to GT but not to any detection, assign it a "occluded" label + trk_is_occluded = trk_is_matched_to_gt & ~det_is_same_obj_id_as_trk.any(dim=1) + trk2det_assoc_labels[trk_is_occluded] = Q_det # "Q_det" label is "occluded" + + # softmax cross-entropy loss for tracking-to-detection association + loss_trk2det_assoc = F.cross_entropy( + input=trk2det_assoc_logits[:, :-1].flatten(0, 1), # (B*Q_trk, Q_det+1) + target=trk2det_assoc_labels.flatten(0, 1), # (B*Q_trk) + ignore_index=-1, + reduction="none", + ).view(B, Q_trk) + # skip trk2det association loss on frames w/o any (non-padding) tracking queries + loss_trk2det_assoc = loss_trk2det_assoc * frame_has_valid_trk.float() + loss_trk2det_assoc = loss_trk2det_assoc.sum() / (B * num_boxes) + loss_dict["loss_trk2det_assoc"] = loss_trk2det_assoc + + return loss_dict + + +def _keep_only_trk_queries_in_match_inds(inds, Q_det): + """Keep only the tracking query indices in the indices tuple""" + batch_idx, src_idx, tgt_idx = inds + if batch_idx.numel() == 0: + return (batch_idx, src_idx, tgt_idx) # empty indices, nothing to filter + + # keep only the tracking query indices + is_trk_query = src_idx >= Q_det + batch_idx_trk = batch_idx[is_trk_query] + src_idx_trk = src_idx[is_trk_query] + tgt_idx_trk = tgt_idx[is_trk_query] if tgt_idx is not None else None + return (batch_idx_trk, src_idx_trk, tgt_idx_trk) diff --git a/src/nn/segearth_ov3/sam3/train/loss/mask_sampling.py b/src/nn/segearth_ov3/sam3/train/loss/mask_sampling.py new file mode 100644 index 0000000..aeba3fe --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/loss/mask_sampling.py @@ -0,0 +1,113 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Callable + +import torch +from torch.nn import functional as F + + +# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py +def point_sample(input, point_coords, **kwargs): + """ + A wrapper around :function:`torch.nn.functional.grid_sample` to support 3D point_coords tensors. + Unlike :function:`torch.nn.functional.grid_sample` it assumes `point_coords` to lie inside + [0, 1] x [0, 1] square. + + Args: + input (Tensor): A tensor of shape (N, C, H, W) that contains features map on a H x W grid. + point_coords (Tensor): A tensor of shape (N, P, 2) or (N, Hgrid, Wgrid, 2) that contains + [0, 1] x [0, 1] normalized point coordinates. + + Returns: + output (Tensor): A tensor of shape (N, C, P) or (N, C, Hgrid, Wgrid) that contains + features for points in `point_coords`. The features are obtained via bilinear + interplation from `input` the same way as :function:`torch.nn.functional.grid_sample`. + """ + add_dim = False + if point_coords.dim() == 3: + add_dim = True + point_coords = point_coords.unsqueeze(2) + normalized_point_coords = 2.0 * point_coords - 1.0 # Normalize to [-1,1] + output = F.grid_sample(input, normalized_point_coords, **kwargs) + if add_dim: + output = output.squeeze(3) + return output + + +# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py +def get_uncertain_point_coords_with_randomness( + logits: torch.Tensor, + uncertainty_func: Callable, + num_points: int, + oversample_ratio: int, + importance_sample_ratio: float, +) -> torch.Tensor: + """ + Sample points in [0, 1] x [0, 1] coordinate space based on their uncertainty. The unceratinties + are calculated for each point using 'uncertainty_func' function that takes point's logit + prediction as input. + See PointRend paper for details. + + Args: + logits (Tensor): A tensor of shape (N, C, Hmask, Wmask) or (N, 1, Hmask, Wmask) for + class-specific or class-agnostic prediction. + uncertainty_func: A function that takes a Tensor of shape (N, C, P) or (N, 1, P) that + contains logit predictions for P points and returns their uncertainties as a Tensor of + shape (N, 1, P). + num_points (int): The number of points P to sample. + oversample_ratio (int): Oversampling parameter. + importance_sample_ratio (float): Ratio of points that are sampled via importnace sampling. + + Returns: + point_coords (Tensor): A tensor of shape (N, P, 2) that contains the coordinates of P + sampled points. + """ + assert oversample_ratio >= 1 + assert importance_sample_ratio <= 1 and importance_sample_ratio >= 0 + num_boxes = logits.shape[0] + num_sampled = int(num_points * oversample_ratio) + point_coords = torch.rand(num_boxes, num_sampled, 2, device=logits.device) + point_logits = point_sample(logits, point_coords, align_corners=False) + # It is crucial to calculate uncertainty based on the sampled prediction value for the points. + # Calculating uncertainties of the predictions first and sampling them for points leads + # to incorrect results. + # To illustrate this: assume uncertainty_func(logits)=-abs(logits), a sampled point between + # two predictions with -1 and 1 logits has 0 logits, and therefore 0 uncertainty value. + # However, if we calculate uncertainties for the predictions first, + # both will have -1 uncertainty, and the sampled point will get -1 uncertainty. + point_uncertainties = uncertainty_func(point_logits) + num_uncertain_points = int(importance_sample_ratio * num_points) + num_random_points = num_points - num_uncertain_points + idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] + # Flatten the indices + shift = num_sampled * torch.arange( + num_boxes, dtype=torch.long, device=logits.device + ) + idx += shift[:, None] + point_coords = point_coords.view(-1, 2)[idx.view(-1), :].view( + num_boxes, num_uncertain_points, 2 + ) + if num_random_points > 0: + point_coords = torch.cat( + [ + point_coords, + torch.rand(num_boxes, num_random_points, 2, device=logits.device), + ], + dim=1, + ) + return point_coords + + +# Adapted from https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/criterion.py +def calculate_uncertainty(logits: torch.Tensor) -> torch.Tensor: + """ + Estimates uncerainty as L1 distance between 0.0 and the logit prediction. + Args: + logits (Tensor): A tensor of shape (R, 1, ...) for class-agnostic + predicted masks + Returns: + scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with + the most uncertain locations having the highest uncertainty score. + """ + assert logits.shape[1] == 1 + return -(torch.abs(logits)) diff --git a/src/nn/segearth_ov3/sam3/train/loss/sam3_loss.py b/src/nn/segearth_ov3/sam3/train/loss/sam3_loss.py new file mode 100644 index 0000000..7ef59dc --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/loss/sam3_loss.py @@ -0,0 +1,203 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + +from sam3.model.model_misc import SAM3Output + +from sam3.train.utils.distributed import get_world_size + +from .loss_fns import CORE_LOSS_KEY, Det2TrkAssoc, Masks + + +class DummyLoss(torch.nn.Module): + """A dummy loss that always returns 0 (as a placeholder for eval)""" + + def __init__( + self, + core_loss_key: str = CORE_LOSS_KEY, + device: str = "cuda", + **kwargs, + ): + super().__init__() + self.core_loss_key = core_loss_key + self.device = torch.device(device) + + def forward(self, *args, **kwargs): + return {self.core_loss_key: torch.tensor(0.0, device=self.device)} + + def accumulate(self, out_dict): + """ + Called by iterative losses. + """ + if self.core_loss_key not in out_dict: + out_dict[self.core_loss_key] = torch.tensor(0.0, device=self.device) + return out_dict + + +class Sam3LossWrapper(torch.nn.Module): + def __init__( + self, + loss_fns_find, + normalization="global", + matcher=None, + o2m_matcher=None, + o2m_weight=1.0, + use_o2m_matcher_on_o2m_aux=True, + loss_fn_semantic_seg=None, + normalize_by_valid_object_num=False, + normalize_by_stage_num=False, + scale_by_find_batch_size=False, + ): + super().__init__() + self.loss_fns_find = loss_fns_find + assert normalization in ["global", "local", "none"] + self.normalization = normalization + self.normalize_by_valid_object_num = normalize_by_valid_object_num + self.normalize_by_stage_num = normalize_by_stage_num + self.matcher = matcher + self.o2m_matcher = o2m_matcher + self.o2m_weight = o2m_weight + # whether to use the o2m matcher on the o2m queries in auxiliary outputs + self.use_o2m_matcher_on_o2m_aux = use_o2m_matcher_on_o2m_aux + self.loss_fn_semantic_seg = loss_fn_semantic_seg + self.scale_by_find_batch_size = scale_by_find_batch_size + + def _get_num_boxes(self, targets): + # the average number of target boxes for loss normalization + if self.normalize_by_valid_object_num: + # valid boxes are those with non-zero height and width + # (while padded invisible boxes are ) + boxes_hw = targets["boxes"].view(-1, 4) # cx, cy, w, h + num_boxes = (boxes_hw[:, 2:] > 0).all(dim=-1).sum().float() + else: + num_boxes = targets["num_boxes"].sum().float() + if self.normalization == "global": + torch.distributed.all_reduce(num_boxes) + num_boxes = torch.clamp(num_boxes / get_world_size(), min=1) + elif self.normalization == "local": + num_boxes = torch.clamp(num_boxes, min=1) + elif self.normalization == "none": + num_boxes = 1 + return num_boxes + + def compute_loss(self, nested_out, targets): + num_boxes = self._get_num_boxes(targets) + o2m_out_is_valid = nested_out.get("o2m_out_is_valid", None) + o2m_target_is_valid_padded = nested_out.get("o2m_target_is_valid_padded", None) + + # Get a list of outputs, including auxiliary and first stage outputs + output_list = [(nested_out, "", False)] # (out, suffix, is_aux) + if "aux_outputs" in nested_out: + output_list.extend( + (aux_out, f"_aux_{i}", True) + for i, aux_out in enumerate(nested_out["aux_outputs"]) + ) + if "first_stage" in nested_out: + output_list.append((nested_out["first_stage"], "_fs", True)) + + # Compute all the requested losses + losses = {} + total_core_loss = 0.0 + for out, suffix, is_aux in output_list: + # o2o matcher indices need to be computed by the model (as the video model requires + # a specific way of matching free and locked indices beyond just calling the matcher) + indices = out["indices"] + has_o2m_out = "pred_logits_o2m" in out + if has_o2m_out: + o2m_out = { + k[: -len("_o2m")]: v for k, v in out.items() if k.endswith("_o2m") + } + # o2m targets are the same as the o2o targets (assuming repeat=1) + o2m_targets = targets + if self.use_o2m_matcher_on_o2m_aux or not is_aux: + o2m_indices = self.o2m_matcher( + o2m_out, + o2m_targets, + out_is_valid=o2m_out_is_valid, + target_is_valid_padded=o2m_target_is_valid_padded, + ) + else: + o2m_indices = self.matcher( + o2m_out, + o2m_targets, + out_is_valid=o2m_out_is_valid, + target_is_valid_padded=o2m_target_is_valid_padded, + ) + + for loss_fn in self.loss_fns_find: + l_dict = loss_fn( + outputs=out, + targets=targets, + indices=indices, + num_boxes=num_boxes, + is_aux=is_aux, + ) + total_core_loss += l_dict.pop(CORE_LOSS_KEY) + losses.update({f"{k}{suffix}": v for k, v in l_dict.items()}) + + compute_o2m_loss = has_o2m_out + # a special handling to allow turning off mask loss in o2m + # (to be compatible with the original implementation) + if isinstance(loss_fn, Masks): + compute_o2m_loss = compute_o2m_loss and "pred_masks" in o2m_out + if isinstance(loss_fn, Det2TrkAssoc): + compute_o2m_loss = False # Det2TrkAssoc does not support o2m + if compute_o2m_loss: + l_dict = loss_fn( + outputs=o2m_out, + targets=o2m_targets, + indices=o2m_indices, + num_boxes=num_boxes, + is_aux=is_aux, + ) + for k in l_dict: + l_dict[k] *= self.o2m_weight + total_core_loss += l_dict.pop(CORE_LOSS_KEY) + losses.update({f"{k}{suffix}_o2m": v for k, v in l_dict.items()}) + + losses[CORE_LOSS_KEY] = total_core_loss + return losses + + def forward(self, find_stages: SAM3Output, find_targets): + if find_stages.loss_stages is not None: + find_targets = [find_targets[i] for i in find_stages.loss_stages] + with SAM3Output.iteration_mode( + find_stages, iter_mode=SAM3Output.IterMode.ALL_STEPS_PER_STAGE + ) as find_stages: + assert len(find_stages) == len(find_targets) + total_losses = {} + for stage_outputs, stage_targets in zip(find_stages, find_targets): + stage_targets = [stage_targets] * len(stage_outputs) + # If there are multiple steps within a stage, compute the loss for all of them (e.g. interactivity) + for outputs, targets in zip(stage_outputs, stage_targets): + cur_losses = self.compute_loss(outputs, targets) + + if self.loss_fn_semantic_seg is not None: + cur_losses_semantic = self.loss_fn_semantic_seg( + outputs, targets + ) + cur_losses[CORE_LOSS_KEY] += cur_losses_semantic.pop( + CORE_LOSS_KEY + ) + # make sure the semantic losses don't overlap with the find losses + assert set(cur_losses).isdisjoint(set(cur_losses_semantic)) + cur_losses.update(cur_losses_semantic) + + # Optionally, normalize the loss by the number of find stages (training video frames) so that + # image batches and video batches have similar loss scales. (Otherwise video batches would + # have a much higher loss scale due to summing the losses over all the find stages.) + if self.normalize_by_stage_num: + cur_losses[CORE_LOSS_KEY] /= len(find_stages) + + if self.scale_by_find_batch_size: + bs = targets["num_boxes"].shape[0] + # sqrt scaling based on the "effective" batch size + cur_losses[CORE_LOSS_KEY] *= bs**0.5 + + for k, v in cur_losses.items(): + if k not in total_losses: + total_losses[k] = v + else: + total_losses[k] += v + + return total_losses diff --git a/src/nn/segearth_ov3/sam3/train/loss/sigmoid_focal_loss.py b/src/nn/segearth_ov3/sam3/train/loss/sigmoid_focal_loss.py new file mode 100644 index 0000000..15e6db4 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/loss/sigmoid_focal_loss.py @@ -0,0 +1,321 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Triton kernel for faster and memory efficient sigmoid focal loss""" + +import torch +import triton +import triton.language as tl +from torch._inductor.runtime.triton_helpers import libdevice + +""" + +The sigmoid focal loss is defined as: + + prob = inputs.sigmoid() + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = prob * targets + (1 - prob) * (1 - targets) + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * ce_loss * ((1 - p_t) ** gamma) + +Where alpha and gamma are scalar parameters, inputs are the logits, targets the float targets. + +We implement two versions of the sigmoid focal loss: with and without sum reduction. +The latter is implemented with built-in reduction to avoid materializing wrt the output of the loss. +This can help save a bit of peak memory. + +The reduction version is implemented using somewhat of a hack. Pytorch's generated kernels usually do the point-wise operation in a first kernel, and implement the reduction another kernel launched on a grid of size 1, where the reduction happens as a for loop in the triton kernel. +Since we want to fuse those two kernels, that is not a good idea: we'd have to launch the overall kernel on a grid of size 1, which is obviously inefficient. +On the other hand, typical CUDA algorithms for reduction (eg reduction tree) are hard to implement in triton due to the lack of thread sync primitives. +We settle for a version that abuses triton's atomic_add: we can have all threads simply add to the same location. +In practice, this is not good, since it creates a massive bottleneck on the semaphore for that single memory location. So instead, we create M reduction locations. Each thread will simply write to thread_id%M. The python code can finally sum over the M reductions. +M = 32 works fine in benchmarking tests. The forward is a tiny bit slower compared to the non-reduced kernel, but the backward breaks even due to one less memory allocation. +""" + + +@triton.jit +def _inner_focal_loss_fwd(inputs, targets, alpha, gamma): + inv_targets = 1 - targets + # Sigmoid + sig = tl.sigmoid(inputs) + + # Binary cross entropy with logits + # In practice, we want the following: + # bce_loss = -targets * tl.log(sig) - (1 - targets) * tl.log(1 - sig) + # However, the above is not numerically stable. + # We're also not directly taking the sum here, so the usual log-sum-exp trick doesn't apply + # The bce can be reformulated, after algebraic manipulation, to + # bce_loss = log(1 + exp(-x)) + x * (1-y) + # This is still not stable, because for large (-x) the exponential will blow up. + # We'll use the following alternate formulation: + # bce_loss = max(x, 0) - x * y + log(1 + exp(-abs(x))) + # Let's show that it's equivalent: + # Case x>=0: abs(x) = x , max(x, 0) = x + # so we get x - x * y + log(1 + exp(-x)) which is equivalent + # Case x<0: abs(x) = -x, max(x, 0) = 0 + # we have log(1 + exp(-abs(x))) = log(1 + exp(x)) = log(exp(x)(1 + exp(-x))) = x+log(1 + exp(-x)) + # plugging it in, we get + # 0 - x * y + x + log(1 + exp(-x)), which is also equivalent + # Note that this is stable because now the exponent are guaranteed to be below 0. + max_val = tl.clamp(inputs, min=0, max=1e9) + bce_loss = max_val - inputs * targets + tl.log(1 + tl.exp(-tl.abs(inputs))) + + # Modulating factor + p_t = sig * targets + (1 - sig) * inv_targets + mod_factor = libdevice.pow(1 - p_t, gamma) + + # Alpha factor + alpha_t = alpha * targets + (1 - alpha) * inv_targets + + # Final loss calculation + return alpha_t * mod_factor * bce_loss + + +# Non-reduced version +@triton.jit +def sigmoid_focal_loss_fwd_kernel( + inputs_ptr, + targets_ptr, + loss_ptr, + alpha: float, + gamma: float, + n_elements: int, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offset = block_start + tl.arange(0, BLOCK_SIZE) + mask = offset < n_elements + + # Load data + inputs = tl.load(inputs_ptr + offset, mask=mask).to(tl.float32) + targets = tl.load(targets_ptr + offset, mask=mask) + + final_loss = _inner_focal_loss_fwd(inputs, targets, alpha, gamma) + + # Store result + tl.store(loss_ptr + offset, final_loss, mask=mask) + + +# version with reduction +@triton.jit +def sigmoid_focal_loss_fwd_kernel_reduce( + inputs_ptr, + targets_ptr, + loss_ptr, + alpha: float, + gamma: float, + n_elements: int, + BLOCK_SIZE: tl.constexpr, + REDUCE_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + reduce_loc = pid % REDUCE_SIZE + offset = block_start + tl.arange(0, BLOCK_SIZE) + mask = offset < n_elements + # Load data + inputs = tl.load(inputs_ptr + offset, mask=mask).to(tl.float32) + targets = tl.load(targets_ptr + offset, mask=mask) + + final_loss = _inner_focal_loss_fwd(inputs, targets, alpha, gamma) * mask + + fl = tl.sum(final_loss) + + # Store result + tl.atomic_add(loss_ptr + reduce_loc, fl) + + +@triton.jit +def _inner_focal_loss_bwd(inputs, targets, alpha, gamma): + inv_targets = 1 - targets + + # Recompute forward + max_val = tl.clamp(inputs, min=0, max=1e9) + bce_loss = max_val - inputs * targets + tl.log(1 + tl.exp(-tl.abs(inputs))) + + # Sigmoid + sig = tl.sigmoid(inputs) + inv_sig = 1 - sig + + # Modulating factor + p_t = sig * targets + inv_sig * inv_targets + tmp = libdevice.pow(1 - p_t, gamma - 1) + mod_factor = tmp * (1 - p_t) + + # Alpha factor + alpha_t = alpha * targets + (1 - alpha) * inv_targets + + # Now computing the derivatives + d_pt = (2 * targets - 1) * sig * inv_sig + d_mod_factor = -gamma * d_pt * tmp + + d_bce_loss = sig - targets + + return alpha_t * (d_bce_loss * mod_factor + d_mod_factor * bce_loss) + + +@triton.jit +def sigmoid_focal_loss_bwd_kernel( + inputs_ptr, + targets_ptr, + grad_inputs_ptr, + grad_out_ptr, + alpha: float, + gamma: float, + n_elements: int, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offset = block_start + tl.arange(0, BLOCK_SIZE) + mask = offset < n_elements + input_ptrs = inputs_ptr + offset + target_ptrs = targets_ptr + offset + grad_input_ptrs = grad_inputs_ptr + offset + grad_out_ptrs = grad_out_ptr + offset + # Load data + inputs = tl.load(input_ptrs, mask=mask).to(tl.float32) + targets = tl.load(target_ptrs, mask=mask) + grad_out = tl.load(grad_out_ptrs, mask=mask) + d_loss = grad_out * _inner_focal_loss_bwd(inputs, targets, alpha, gamma) + tl.store(grad_input_ptrs, d_loss, mask=mask) + + +@triton.jit +def sigmoid_focal_loss_bwd_kernel_reduce( + inputs_ptr, + targets_ptr, + grad_inputs_ptr, + grad_out_ptr, + alpha: float, + gamma: float, + n_elements: int, + BLOCK_SIZE: tl.constexpr, +): + # The only difference is that the gradient is now a single scalar + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offset = block_start + tl.arange(0, BLOCK_SIZE) + mask = offset < n_elements + input_ptrs = inputs_ptr + offset + target_ptrs = targets_ptr + offset + grad_input_ptrs = grad_inputs_ptr + offset + # Load data + inputs = tl.load(input_ptrs, mask=mask).to(tl.float32) + targets = tl.load(target_ptrs, mask=mask) + grad_out = tl.load(grad_out_ptr) + d_loss = grad_out * _inner_focal_loss_bwd(inputs, targets, alpha, gamma) + tl.store(grad_input_ptrs, d_loss, mask=mask) + + +class SigmoidFocalLoss(torch.autograd.Function): + BLOCK_SIZE = 256 + + @staticmethod + def forward(ctx, inputs, targets, alpha=0.25, gamma=2): + n_elements = inputs.numel() + assert targets.numel() == n_elements + input_shape = inputs.shape + inputs = inputs.view(-1).contiguous() + targets = targets.view(-1).contiguous() + loss = torch.empty(inputs.shape, dtype=torch.float32, device=inputs.device) + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + sigmoid_focal_loss_fwd_kernel[grid]( + inputs, targets, loss, alpha, gamma, n_elements, SigmoidFocalLoss.BLOCK_SIZE + ) + ctx.save_for_backward(inputs.view(input_shape), targets.view(input_shape)) + ctx.alpha = alpha + ctx.gamma = gamma + return loss.view(input_shape) + + @staticmethod + def backward(ctx, grad_output): + inputs, targets = ctx.saved_tensors + alpha = ctx.alpha + gamma = ctx.gamma + n_elements = inputs.numel() + input_shape = inputs.shape + grad_inputs = torch.empty( + inputs.shape, dtype=grad_output.dtype, device=grad_output.device + ) + inputs_ptr = inputs.view(-1).contiguous() + targets_ptr = targets.view(-1).contiguous() + grad_output_ptr = grad_output.view(-1).contiguous() + grad_inputs_ptr = grad_inputs + assert grad_output.numel() == n_elements + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + sigmoid_focal_loss_bwd_kernel[grid]( + inputs_ptr, + targets_ptr, + grad_inputs_ptr, + grad_output_ptr, + alpha, + gamma, + n_elements, + SigmoidFocalLoss.BLOCK_SIZE, + ) + return grad_inputs.view(input_shape), None, None, None + + +triton_sigmoid_focal_loss = SigmoidFocalLoss.apply + + +class SigmoidFocalLossReduced(torch.autograd.Function): + BLOCK_SIZE = 256 + REDUCE_SIZE = 32 + + @staticmethod + def forward(ctx, inputs, targets, alpha=0.25, gamma=2): + n_elements = inputs.numel() + input_shape = inputs.shape + inputs = inputs.view(-1).contiguous() + targets = targets.view(-1).contiguous() + loss = torch.zeros( + SigmoidFocalLossReduced.REDUCE_SIZE, + device=inputs.device, + dtype=torch.float32, + ) + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + sigmoid_focal_loss_fwd_kernel_reduce[grid]( + inputs, + targets, + loss, + alpha, + gamma, + n_elements, + SigmoidFocalLossReduced.BLOCK_SIZE, + SigmoidFocalLossReduced.REDUCE_SIZE, + ) + ctx.save_for_backward(inputs.view(input_shape), targets.view(input_shape)) + ctx.alpha = alpha + ctx.gamma = gamma + return loss.sum() + + @staticmethod + def backward(ctx, grad_output): + inputs, targets = ctx.saved_tensors + alpha = ctx.alpha + gamma = ctx.gamma + n_elements = inputs.numel() + input_shape = inputs.shape + grad_inputs = torch.empty( + inputs.shape, dtype=grad_output.dtype, device=grad_output.device + ) + inputs_ptr = inputs.view(-1).contiguous() + targets_ptr = targets.reshape(-1).contiguous() + assert grad_output.numel() == 1 + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + sigmoid_focal_loss_bwd_kernel_reduce[grid]( + inputs_ptr, + targets_ptr, + grad_inputs, + grad_output, + alpha, + gamma, + n_elements, + SigmoidFocalLossReduced.BLOCK_SIZE, + ) + return grad_inputs.view(input_shape), None, None, None + + +triton_sigmoid_focal_loss_reduce = SigmoidFocalLossReduced.apply diff --git a/src/nn/segearth_ov3/sam3/train/masks_ops.py b/src/nn/segearth_ov3/sam3/train/masks_ops.py new file mode 100644 index 0000000..f9d2fd7 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/masks_ops.py @@ -0,0 +1,272 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +"""Utilities for masks manipulation""" + +import numpy as np +import pycocotools.mask as maskUtils +import torch +from pycocotools import mask as mask_util + + +def instance_masks_to_semantic_masks( + instance_masks: torch.Tensor, num_instances: torch.Tensor +) -> torch.Tensor: + """This function converts instance masks to semantic masks. + It accepts a collapsed batch of instances masks (ie all instance masks are concatenated in a single tensor) and + the number of instances in each image of the batch. + It returns a mask with the same spatial dimensions as the input instance masks, where for each batch element the + semantic mask is the union of all the instance masks in the batch element. + + If for a given batch element there are no instances (ie num_instances[i]==0), the corresponding semantic mask will be a tensor of zeros. + + Args: + instance_masks (torch.Tensor): A tensor of shape (N, H, W) where N is the number of instances in the batch. + num_instances (torch.Tensor): A tensor of shape (B,) where B is the batch size. It contains the number of instances + in each image of the batch. + + Returns: + torch.Tensor: A tensor of shape (B, H, W) where B is the batch size and H, W are the spatial dimensions of the + input instance masks. + """ + + masks_per_query = torch.split(instance_masks, num_instances.tolist()) + + return torch.stack([torch.any(masks, dim=0) for masks in masks_per_query], dim=0) + + +def mask_intersection(masks1, masks2, block_size=16): + """Compute the intersection of two sets of masks, without blowing the memory""" + + assert masks1.shape[1:] == masks2.shape[1:] + assert masks1.dtype == torch.bool and masks2.dtype == torch.bool + + result = torch.zeros( + masks1.shape[0], masks2.shape[0], device=masks1.device, dtype=torch.long + ) + for i in range(0, masks1.shape[0], block_size): + for j in range(0, masks2.shape[0], block_size): + intersection = ( + (masks1[i : i + block_size, None] * masks2[None, j : j + block_size]) + .flatten(-2) + .sum(-1) + ) + result[i : i + block_size, j : j + block_size] = intersection + return result + + +def mask_iom(masks1, masks2): + """ + Similar to IoU, except the denominator is the area of the smallest mask + """ + assert masks1.shape[1:] == masks2.shape[1:] + assert masks1.dtype == torch.bool and masks2.dtype == torch.bool + + # intersection = (masks1[:, None] * masks2[None]).flatten(-2).sum(-1) + intersection = mask_intersection(masks1, masks2) + area1 = masks1.flatten(-2).sum(-1) + area2 = masks2.flatten(-2).sum(-1) + min_area = torch.min(area1[:, None], area2[None, :]) + return intersection / (min_area + 1e-8) + + +def compute_boundary(seg): + """ + Adapted from https://github.com/JonathonLuiten/TrackEval/blob/master/trackeval/metrics/j_and_f.py#L148 + Return a 1pix wide boundary of the given mask + """ + assert seg.ndim >= 2 + e = torch.zeros_like(seg) + s = torch.zeros_like(seg) + se = torch.zeros_like(seg) + + e[..., :, :-1] = seg[..., :, 1:] + s[..., :-1, :] = seg[..., 1:, :] + se[..., :-1, :-1] = seg[..., 1:, 1:] + + b = seg ^ e | seg ^ s | seg ^ se + b[..., -1, :] = seg[..., -1, :] ^ e[..., -1, :] + b[..., :, -1] = seg[..., :, -1] ^ s[..., :, -1] + b[..., -1, -1] = 0 + return b + + +def dilation(mask, kernel_size): + """ + Implements the dilation operation. If the input is on cpu, we call the cv2 version. + Otherwise, we implement it using a convolution + + The kernel is assumed to be a square kernel + + """ + + assert mask.ndim == 3 + kernel_size = int(kernel_size) + assert ( + kernel_size % 2 == 1 + ), f"Dilation expects a odd kernel size, got {kernel_size}" + + if mask.is_cuda: + m = mask.unsqueeze(1).to(torch.float16) + k = torch.ones(1, 1, kernel_size, 1, dtype=m.dtype, device=m.device) + + result = torch.nn.functional.conv2d(m, k, padding="same") + result = torch.nn.functional.conv2d(result, k.transpose(-1, -2), padding="same") + return result.view_as(mask) > 0 + + all_masks = mask.view(-1, mask.size(-2), mask.size(-1)).numpy().astype(np.uint8) + kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8) + + import cv2 + + processed = [torch.from_numpy(cv2.dilate(m, kernel)) for m in all_masks] + return torch.stack(processed).view_as(mask).to(mask) + + +def compute_F_measure( + gt_boundary_rle, gt_dilated_boundary_rle, dt_boundary_rle, dt_dilated_boundary_rle +): + """Adapted from https://github.com/JonathonLuiten/TrackEval/blob/master/trackeval/metrics/j_and_f.py#L207 + + Assumes the boundary and dilated boundaries have already been computed and converted to RLE + """ + gt_match = maskUtils.merge([gt_boundary_rle, dt_dilated_boundary_rle], True) + dt_match = maskUtils.merge([dt_boundary_rle, gt_dilated_boundary_rle], True) + + n_dt = maskUtils.area(dt_boundary_rle) + n_gt = maskUtils.area(gt_boundary_rle) + # % Compute precision and recall + if n_dt == 0 and n_gt > 0: + precision = 1 + recall = 0 + elif n_dt > 0 and n_gt == 0: + precision = 0 + recall = 1 + elif n_dt == 0 and n_gt == 0: + precision = 1 + recall = 1 + else: + precision = maskUtils.area(dt_match) / float(n_dt) + recall = maskUtils.area(gt_match) / float(n_gt) + + # Compute F measure + if precision + recall == 0: + f_val = 0 + else: + f_val = 2 * precision * recall / (precision + recall) + + return f_val + + +@torch.no_grad() +def rle_encode(orig_mask, return_areas=False): + """Encodes a collection of masks in RLE format + + This function emulates the behavior of the COCO API's encode function, but + is executed partially on the GPU for faster execution. + + Args: + mask (torch.Tensor): A mask of shape (N, H, W) with dtype=torch.bool + return_areas (bool): If True, add the areas of the masks as a part of + the RLE output dict under the "area" key. Default is False. + + Returns: + str: The RLE encoded masks + """ + assert orig_mask.ndim == 3, "Mask must be of shape (N, H, W)" + assert orig_mask.dtype == torch.bool, "Mask must have dtype=torch.bool" + + if orig_mask.numel() == 0: + return [] + + # First, transpose the spatial dimensions. + # This is necessary because the COCO API uses Fortran order + mask = orig_mask.transpose(1, 2) + + # Flatten the mask + flat_mask = mask.reshape(mask.shape[0], -1) + if return_areas: + mask_areas = flat_mask.sum(-1).tolist() + # Find the indices where the mask changes + differences = torch.ones( + mask.shape[0], flat_mask.shape[1] + 1, device=mask.device, dtype=torch.bool + ) + differences[:, 1:-1] = flat_mask[:, :-1] != flat_mask[:, 1:] + differences[:, 0] = flat_mask[:, 0] + _, change_indices = torch.where(differences) + + try: + boundaries = torch.cumsum(differences.sum(-1), 0).cpu() + except RuntimeError as _: + boundaries = torch.cumsum(differences.cpu().sum(-1), 0) + + change_indices_clone = change_indices.clone() + # First pass computes the RLEs on GPU, in a flatten format + for i in range(mask.shape[0]): + # Get the change indices for this batch item + beg = 0 if i == 0 else boundaries[i - 1].item() + end = boundaries[i].item() + change_indices[beg + 1 : end] -= change_indices_clone[beg : end - 1] + + # Now we can split the RLES of each batch item, and convert them to strings + # No more gpu at this point + change_indices = change_indices.tolist() + + batch_rles = [] + # Process each mask in the batch separately + for i in range(mask.shape[0]): + beg = 0 if i == 0 else boundaries[i - 1].item() + end = boundaries[i].item() + run_lengths = change_indices[beg:end] + + uncompressed_rle = {"counts": run_lengths, "size": list(orig_mask.shape[1:])} + h, w = uncompressed_rle["size"] + rle = mask_util.frPyObjects(uncompressed_rle, h, w) + rle["counts"] = rle["counts"].decode("utf-8") + if return_areas: + rle["area"] = mask_areas[i] + batch_rles.append(rle) + + return batch_rles + + +def robust_rle_encode(masks): + """Encodes a collection of masks in RLE format. Uses the gpu version fist, falls back to the cpu version if it fails""" + + assert masks.ndim == 3, "Mask must be of shape (N, H, W)" + assert masks.dtype == torch.bool, "Mask must have dtype=torch.bool" + + try: + return rle_encode(masks) + except RuntimeError as _: + masks = masks.cpu().numpy() + rles = [ + mask_util.encode( + np.array(mask[:, :, np.newaxis], dtype=np.uint8, order="F") + )[0] + for mask in masks + ] + for rle in rles: + rle["counts"] = rle["counts"].decode("utf-8") + return rles + + +def ann_to_rle(segm, im_info): + """Convert annotation which can be polygons, uncompressed RLE to RLE. + Args: + ann (dict) : annotation object + Returns: + ann (rle) + """ + h, w = im_info["height"], im_info["width"] + if isinstance(segm, list): + # polygon -- a single object might consist of multiple parts + # we merge all parts into one mask rle code + rles = mask_util.frPyObjects(segm, h, w) + rle = mask_util.merge(rles) + elif isinstance(segm["counts"], list): + # uncompressed RLE + rle = mask_util.frPyObjects(segm, h, w) + else: + # rle + rle = segm + return rle diff --git a/src/nn/segearth_ov3/sam3/train/matcher.py b/src/nn/segearth_ov3/sam3/train/matcher.py new file mode 100644 index 0000000..b0b8d62 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/matcher.py @@ -0,0 +1,806 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Modules to compute the matching cost and solve the corresponding LSAP. +""" + +import numpy as np +import torch + +from sam3.model.box_ops import box_cxcywh_to_xyxy, box_iou, generalized_box_iou +from scipy.optimize import linear_sum_assignment +from torch import nn + + +def _do_matching(cost, repeats=1, return_tgt_indices=False, do_filtering=False): + if repeats > 1: + cost = np.tile(cost, (1, repeats)) + + i, j = linear_sum_assignment(cost) + if do_filtering: + # filter out invalid entries (i.e. those with cost > 1e8) + valid_thresh = 1e8 + valid_ijs = [(ii, jj) for ii, jj in zip(i, j) if cost[ii, jj] < valid_thresh] + i, j = zip(*valid_ijs) if len(valid_ijs) > 0 else ([], []) + i, j = np.array(i, dtype=np.int64), np.array(j, dtype=np.int64) + if return_tgt_indices: + return i, j + order = np.argsort(j) + return i[order] + + +class HungarianMatcher(nn.Module): + """This class computes an assignment between the targets and the predictions of the network + + For efficiency reasons, the targets don't include the no_object. Because of this, in general, + there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, + while the others are un-matched (and thus treated as non-objects). + """ + + def __init__( + self, + cost_class: float = 1, + cost_bbox: float = 1, + cost_giou: float = 1, + focal_loss: bool = False, + focal_alpha: float = 0.25, + focal_gamma: float = 2, + ): + """Creates the matcher + + Params: + cost_class: This is the relative weight of the classification error in the matching cost + cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost + cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost + """ + super().__init__() + self.cost_class = cost_class + self.cost_bbox = cost_bbox + self.cost_giou = cost_giou + self.norm = nn.Sigmoid() if focal_loss else nn.Softmax(-1) + assert ( + cost_class != 0 or cost_bbox != 0 or cost_giou != 0 + ), "all costs cant be 0" + self.focal_loss = focal_loss + self.focal_alpha = focal_alpha + self.focal_gamma = focal_gamma + + @torch.no_grad() + def forward(self, outputs, batched_targets): + """Performs the matching + + Params: + outputs: This is a dict that contains at least these entries: + "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits + "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates + + targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: + "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth + objects in the target) containing the class labels + "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates + + Returns: + A list of size batch_size, containing tuples of (index_i, index_j) where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected targets (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_target_boxes) + """ + bs, num_queries = outputs["pred_logits"].shape[:2] + + # We flatten to compute the cost matrices in a batch + out_prob = self.norm( + outputs["pred_logits"].flatten(0, 1) + ) # [batch_size * num_queries, num_classes] + out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] + + # Also concat the target labels and boxes + tgt_bbox = batched_targets["boxes"] + + if "positive_map" in batched_targets: + # In this case we have a multi-hot target + positive_map = batched_targets["positive_map"] + assert len(tgt_bbox) == len(positive_map) + + if self.focal_loss: + positive_map = positive_map > 1e-4 + alpha = self.focal_alpha + gamma = self.focal_gamma + neg_cost_class = ( + (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log()) + ) + pos_cost_class = ( + alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log()) + ) + cost_class = ( + (pos_cost_class - neg_cost_class).unsqueeze(1) + * positive_map.unsqueeze(0) + ).sum(-1) + else: + # Compute the soft-cross entropy between the predicted token alignment and the GT one for each box + cost_class = -(out_prob.unsqueeze(1) * positive_map.unsqueeze(0)).sum( + -1 + ) + else: + # In this case we are doing a "standard" cross entropy + tgt_ids = batched_targets["labels"] + assert len(tgt_bbox) == len(tgt_ids) + + if self.focal_loss: + alpha = self.focal_alpha + gamma = self.focal_gamma + neg_cost_class = ( + (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log()) + ) + pos_cost_class = ( + alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log()) + ) + cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] + else: + # Compute the classification cost. Contrary to the loss, we don't use the NLL, + # but approximate it in 1 - proba[target class]. + # The 1 is a constant that doesn't change the matching, it can be omitted. + cost_class = -out_prob[:, tgt_ids] + + # Compute the L1 cost between boxes + cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) + assert cost_class.shape == cost_bbox.shape + + # Compute the giou cost betwen boxes + cost_giou = -generalized_box_iou( + box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) + ) + + # Final cost matrix + C = ( + self.cost_bbox * cost_bbox + + self.cost_class * cost_class + + self.cost_giou * cost_giou + ) + C = C.view(bs, num_queries, -1).cpu().numpy() + + sizes = torch.cumsum(batched_targets["num_boxes"], -1)[:-1] + costs = [c[i] for i, c in enumerate(np.split(C, sizes.cpu().numpy(), axis=-1))] + indices = [_do_matching(c) for c in costs] + batch_idx = torch.as_tensor( + sum([[i] * len(src) for i, src in enumerate(indices)], []), dtype=torch.long + ) + src_idx = torch.from_numpy(np.concatenate(indices)).long() + return batch_idx, src_idx + + +class BinaryHungarianMatcher(nn.Module): + """This class computes an assignment between the targets and the predictions of the network + + For efficiency reasons, the targets don't include the no_object. Because of this, in general, + there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, + while the others are un-matched (and thus treated as non-objects). + """ + + def __init__( + self, + cost_class: float = 1, + cost_bbox: float = 1, + cost_giou: float = 1, + ): + """Creates the matcher + + Params: + cost_class: This is the relative weight of the classification error in the matching cost + cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost + cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost + """ + super().__init__() + self.cost_class = cost_class + self.cost_bbox = cost_bbox + self.cost_giou = cost_giou + self.norm = nn.Sigmoid() + assert ( + cost_class != 0 or cost_bbox != 0 or cost_giou != 0 + ), "all costs cant be 0" + + @torch.no_grad() + def forward(self, outputs, batched_targets, repeats=0, repeat_batch=1): + """Performs the matching + + Params: + outputs: This is a dict that contains at least these entries: + "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits + "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates + + targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: + "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth + objects in the target) containing the class labels + "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates + + Returns: + A list of size batch_size, containing tuples of (index_i, index_j) where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected targets (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_target_boxes) + """ + if repeat_batch != 1: + raise NotImplementedError("please use BinaryHungarianMatcherV2 instead") + + bs, num_queries = outputs["pred_logits"].shape[:2] + + # We flatten to compute the cost matrices in a batch + out_prob = self.norm(outputs["pred_logits"].flatten(0, 1)).squeeze( + -1 + ) # [batch_size * num_queries] + out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] + + # Also concat the target labels and boxes + tgt_bbox = batched_targets["boxes"] + + # Compute the L1 cost between boxes + cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) + + cost_class = -out_prob.unsqueeze(-1).expand_as(cost_bbox) + + assert cost_class.shape == cost_bbox.shape + + # Compute the giou cost betwen boxes + cost_giou = -generalized_box_iou( + box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) + ) + + # Final cost matrix + C = ( + self.cost_bbox * cost_bbox + + self.cost_class * cost_class + + self.cost_giou * cost_giou + ) + C = C.view(bs, num_queries, -1).cpu().numpy() + + sizes = torch.cumsum(batched_targets["num_boxes"], -1)[:-1] + costs = [c[i] for i, c in enumerate(np.split(C, sizes.cpu().numpy(), axis=-1))] + return_tgt_indices = False + for c in costs: + n_targ = c.shape[1] + if repeats > 1: + n_targ *= repeats + if c.shape[0] < n_targ: + return_tgt_indices = True + break + if return_tgt_indices: + indices, tgt_indices = zip( + *( + _do_matching( + c, repeats=repeats, return_tgt_indices=return_tgt_indices + ) + for c in costs + ) + ) + tgt_indices = list(tgt_indices) + for i in range(1, len(tgt_indices)): + tgt_indices[i] += sizes[i - 1].item() + tgt_idx = torch.from_numpy(np.concatenate(tgt_indices)).long() + else: + indices = [_do_matching(c, repeats=repeats) for c in costs] + tgt_idx = None + + batch_idx = torch.as_tensor( + sum([[i] * len(src) for i, src in enumerate(indices)], []), dtype=torch.long + ) + src_idx = torch.from_numpy(np.concatenate(indices)).long() + return batch_idx, src_idx, tgt_idx + + +class BinaryFocalHungarianMatcher(nn.Module): + """This class computes an assignment between the targets and the predictions of the network + + For efficiency reasons, the targets don't include the no_object. Because of this, in general, + there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, + while the others are un-matched (and thus treated as non-objects). + """ + + def __init__( + self, + cost_class: float = 1, + cost_bbox: float = 1, + cost_giou: float = 1, + alpha: float = 0.25, + gamma: float = 2.0, + stable: bool = False, + ): + """Creates the matcher + + Params: + cost_class: This is the relative weight of the classification error in the matching cost + cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost + cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost + """ + super().__init__() + self.cost_class = cost_class + self.cost_bbox = cost_bbox + self.cost_giou = cost_giou + self.norm = nn.Sigmoid() + self.alpha = alpha + self.gamma = gamma + self.stable = stable + assert ( + cost_class != 0 or cost_bbox != 0 or cost_giou != 0 + ), "all costs cant be 0" + + @torch.no_grad() + def forward(self, outputs, batched_targets, repeats=1, repeat_batch=1): + """Performs the matching + + Params: + outputs: This is a dict that contains at least these entries: + "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits + "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates + + targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: + "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth + objects in the target) containing the class labels + "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates + + Returns: + A list of size batch_size, containing tuples of (index_i, index_j) where: + - index_i is the indices of the selected predictions (in order) + - index_j is the indices of the corresponding selected targets (in order) + For each batch element, it holds: + len(index_i) = len(index_j) = min(num_queries, num_target_boxes) + """ + if repeat_batch != 1: + raise NotImplementedError("please use BinaryHungarianMatcherV2 instead") + + bs, num_queries = outputs["pred_logits"].shape[:2] + + # We flatten to compute the cost matrices in a batch + out_score = outputs["pred_logits"].flatten(0, 1).squeeze(-1) + out_prob = self.norm(out_score) # [batch_size * num_queries] + out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] + + # Also concat the target labels and boxes + tgt_bbox = batched_targets["boxes"] + + # Compute the L1 cost between boxes + cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) + + # Compute the giou cost betwen boxes + cost_giou = -generalized_box_iou( + box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) + ) + + # cost_class = -out_prob.unsqueeze(-1).expand_as(cost_bbox) + if self.stable: + rescaled_giou = (-cost_giou + 1) / 2 + out_prob = out_prob.unsqueeze(-1).expand_as(cost_bbox) * rescaled_giou + cost_class = -self.alpha * (1 - out_prob) ** self.gamma * torch.log( + out_prob + ) + (1 - self.alpha) * out_prob**self.gamma * torch.log(1 - out_prob) + else: + # directly computing log sigmoid (more numerically stable) + log_out_prob = torch.nn.functional.logsigmoid(out_score) + log_one_minus_out_prob = torch.nn.functional.logsigmoid(-out_score) + cost_class = ( + -self.alpha * (1 - out_prob) ** self.gamma * log_out_prob + + (1 - self.alpha) * out_prob**self.gamma * log_one_minus_out_prob + ) + if not self.stable: + cost_class = cost_class.unsqueeze(-1).expand_as(cost_bbox) + + assert cost_class.shape == cost_bbox.shape + + # Final cost matrix + C = ( + self.cost_bbox * cost_bbox + + self.cost_class * cost_class + + self.cost_giou * cost_giou + ) + C = C.view(bs, num_queries, -1).cpu().numpy() + + sizes = torch.cumsum(batched_targets["num_boxes"], -1)[:-1] + costs = [c[i] for i, c in enumerate(np.split(C, sizes.cpu().numpy(), axis=-1))] + return_tgt_indices = False + for c in costs: + n_targ = c.shape[1] + if repeats > 1: + n_targ *= repeats + if c.shape[0] < n_targ: + return_tgt_indices = True + break + if return_tgt_indices: + indices, tgt_indices = zip( + *( + _do_matching( + c, repeats=repeats, return_tgt_indices=return_tgt_indices + ) + for c in costs + ) + ) + tgt_indices = list(tgt_indices) + for i in range(1, len(tgt_indices)): + tgt_indices[i] += sizes[i - 1].item() + tgt_idx = torch.from_numpy(np.concatenate(tgt_indices)).long() + else: + indices = [_do_matching(c, repeats=repeats) for c in costs] + tgt_idx = None + + batch_idx = torch.as_tensor( + sum([[i] * len(src) for i, src in enumerate(indices)], []), dtype=torch.long + ) + src_idx = torch.from_numpy(np.concatenate(indices)).long() + return batch_idx, src_idx, tgt_idx + + +class BinaryHungarianMatcherV2(nn.Module): + """ + This class computes an assignment between the targets and the predictions + of the network + + For efficiency reasons, the targets don't include the no_object. Because of + this, in general, there are more predictions than targets. In this case, we + do a 1-to-1 matching of the best predictions, while the others are + un-matched (and thus treated as non-objects). + + This is a more efficient implementation of BinaryHungarianMatcher. + """ + + def __init__( + self, + cost_class: float = 1, + cost_bbox: float = 1, + cost_giou: float = 1, + focal: bool = False, + alpha: float = 0.25, + gamma: float = 2.0, + stable: bool = False, + remove_samples_with_0_gt: bool = True, + ): + """ + Creates the matcher + + Params: + - cost_class: Relative weight of the classification error in the + matching cost + - cost_bbox: Relative weight of the L1 error of the bounding box + coordinates in the matching cost + - cost_giou: This is the relative weight of the giou loss of the + bounding box in the matching cost + """ + super().__init__() + self.cost_class = cost_class + self.cost_bbox = cost_bbox + self.cost_giou = cost_giou + self.norm = nn.Sigmoid() + assert ( + cost_class != 0 or cost_bbox != 0 or cost_giou != 0 + ), "all costs cant be 0" + self.focal = focal + if focal: + self.alpha = alpha + self.gamma = gamma + self.stable = stable + self.remove_samples_with_0_gt = remove_samples_with_0_gt + + @torch.no_grad() + def forward( + self, + outputs, + batched_targets, + repeats=1, + repeat_batch=1, + out_is_valid=None, + target_is_valid_padded=None, + ): + """ + Performs the matching. The inputs and outputs are the same as + BinaryHungarianMatcher.forward, except for the optional cached_padded + flag and the optional "_boxes_padded" entry of batched_targets. + + Inputs: + - outputs: A dict with the following keys: + - "pred_logits": Tensor of shape (batch_size, num_queries, 1) with + classification logits + - "pred_boxes": Tensor of shape (batch_size, num_queries, 4) with + predicted box coordinates in cxcywh format. + - batched_targets: A dict of targets. There may be a variable number of + targets per batch entry; suppose that there are T_b targets for batch + entry 0 <= b < batch_size. It should have the following keys: + - "boxes": Tensor of shape (sum_b T_b, 4) giving ground-truth boxes + in cxcywh format for all batch entries packed into a single tensor + - "num_boxes": int64 Tensor of shape (batch_size,) giving the number + of ground-truth boxes per batch entry: num_boxes[b] = T_b + - "_boxes_padded": Tensor of shape (batch_size, max_b T_b, 4) giving + a padded version of ground-truth boxes. If this is not present then + it will be computed from batched_targets["boxes"] instead, but + caching it here can improve performance for repeated calls with the + same targets. + - out_is_valid: If not None, it should be a boolean tensor of shape + (batch_size, num_queries) indicating which predictions are valid. + Invalid predictions are ignored during matching and won't appear in + the output indices. + - target_is_valid_padded: If not None, it should be a boolean tensor of + shape (batch_size, max_num_gt_boxes) in padded format indicating + which GT boxes are valid. Invalid GT boxes are ignored during matching + and won't appear in the output indices. + + Returns: + A list of size batch_size, containing tuples of (idx_i, idx_j): + - idx_i is the indices of the selected predictions (in order) + - idx_j is the indices of the corresponding selected targets + (in order) + For each batch element, it holds: + len(index_i) = len(index_j) + = min(num_queries, num_target_boxes) + """ + _, num_queries = outputs["pred_logits"].shape[:2] + + out_score = outputs["pred_logits"].squeeze(-1) # (B, Q) + out_bbox = outputs["pred_boxes"] # (B, Q, 4)) + + device = out_score.device + + num_boxes = batched_targets["num_boxes"].cpu() + # Get a padded version of target boxes (as precomputed in the collator). + # It should work for both repeat==1 (o2o) and repeat>1 (o2m) matching. + tgt_bbox = batched_targets["boxes_padded"] + if self.remove_samples_with_0_gt: + # keep only samples w/ at least 1 GT box in targets (num_boxes and tgt_bbox) + batch_keep = num_boxes > 0 + num_boxes = num_boxes[batch_keep] + tgt_bbox = tgt_bbox[batch_keep] + if target_is_valid_padded is not None: + target_is_valid_padded = target_is_valid_padded[batch_keep] + # Repeat the targets (for the case of batched aux outputs in the matcher) + if repeat_batch > 1: + # In this case, out_prob and out_bbox will be a concatenation of + # both final and auxiliary outputs, so we also repeat the targets + num_boxes = num_boxes.repeat(repeat_batch) + tgt_bbox = tgt_bbox.repeat(repeat_batch, 1, 1) + if target_is_valid_padded is not None: + target_is_valid_padded = target_is_valid_padded.repeat(repeat_batch, 1) + + # keep only samples w/ at least 1 GT box in outputs + if self.remove_samples_with_0_gt: + if repeat_batch > 1: + batch_keep = batch_keep.repeat(repeat_batch) + out_score = out_score[batch_keep] + out_bbox = out_bbox[batch_keep] + if out_is_valid is not None: + out_is_valid = out_is_valid[batch_keep] + assert out_bbox.shape[0] == tgt_bbox.shape[0] + assert out_bbox.shape[0] == num_boxes.shape[0] + + # Compute the L1 cost between boxes + cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) + + # Compute the giou cost betwen boxes + cost_giou = -generalized_box_iou( + box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) + ) + + out_prob = self.norm(out_score) + if not self.focal: + cost_class = -out_prob.unsqueeze(-1).expand_as(cost_bbox) + else: + if self.stable: + rescaled_giou = (-cost_giou + 1) / 2 + out_prob = out_prob.unsqueeze(-1).expand_as(cost_bbox) * rescaled_giou + cost_class = -self.alpha * (1 - out_prob) ** self.gamma * torch.log( + out_prob + ) + (1 - self.alpha) * out_prob**self.gamma * torch.log(1 - out_prob) + else: + # directly computing log sigmoid (more numerically stable) + log_out_prob = torch.nn.functional.logsigmoid(out_score) + log_one_minus_out_prob = torch.nn.functional.logsigmoid(-out_score) + cost_class = ( + -self.alpha * (1 - out_prob) ** self.gamma * log_out_prob + + (1 - self.alpha) * out_prob**self.gamma * log_one_minus_out_prob + ) + if not self.stable: + cost_class = cost_class.unsqueeze(-1).expand_as(cost_bbox) + + assert cost_class.shape == cost_bbox.shape + + # Final cost matrix + C = ( + self.cost_bbox * cost_bbox + + self.cost_class * cost_class + + self.cost_giou * cost_giou + ) + # assign a very high cost (1e9) to invalid outputs and targets, so that we can + # filter them out (in `_do_matching`) from bipartite matching results + do_filtering = out_is_valid is not None or target_is_valid_padded is not None + if out_is_valid is not None: + C = torch.where(out_is_valid[:, :, None], C, 1e9) + if target_is_valid_padded is not None: + C = torch.where(target_is_valid_padded[:, None, :], C, 1e9) + C = C.cpu().numpy() + costs = [C[i, :, :s] for i, s in enumerate(num_boxes.tolist())] + return_tgt_indices = ( + do_filtering or torch.any(num_queries < num_boxes * max(repeats, 1)).item() + ) + if len(costs) == 0: + # We have size 0 in the batch dimension, so we return empty matching indices + # (note that this can happen due to `remove_samples_with_0_gt=True` even if + # the original input batch size is not 0, when all queries have empty GTs). + indices = [] + tgt_idx = torch.zeros(0).long().to(device) if return_tgt_indices else None + elif return_tgt_indices: + indices, tgt_indices = zip( + *( + _do_matching( + c, + repeats=repeats, + return_tgt_indices=return_tgt_indices, + do_filtering=do_filtering, + ) + for c in costs + ) + ) + tgt_indices = list(tgt_indices) + sizes = torch.cumsum(num_boxes, -1)[:-1] + for i in range(1, len(tgt_indices)): + tgt_indices[i] += sizes[i - 1].item() + tgt_idx = torch.from_numpy(np.concatenate(tgt_indices)).long().to(device) + else: + indices = [ + _do_matching(c, repeats=repeats, do_filtering=do_filtering) + for c in costs + ] + tgt_idx = None + + if self.remove_samples_with_0_gt: + kept_inds = batch_keep.nonzero().squeeze(1) + batch_idx = torch.as_tensor( + sum([[kept_inds[i]] * len(src) for i, src in enumerate(indices)], []), + dtype=torch.long, + device=device, + ) + else: + batch_idx = torch.as_tensor( + sum([[i] * len(src) for i, src in enumerate(indices)], []), + dtype=torch.long, + device=device, + ) + + # indices could be an empty list (since we remove samples w/ 0 GT boxes) + if len(indices) > 0: + src_idx = torch.from_numpy(np.concatenate(indices)).long().to(device) + else: + src_idx = torch.empty(0, dtype=torch.long, device=device) + return batch_idx, src_idx, tgt_idx + + +class BinaryOneToManyMatcher(nn.Module): + """ + This class computes a greedy assignment between the targets and the predictions of the network. + In this formulation, several predictions can be assigned to each target, but each prediction can be assigned to + at most one target. + + See DAC-Detr for details + """ + + def __init__( + self, + alpha: float = 0.3, + threshold: float = 0.4, + topk: int = 6, + ): + """ + Creates the matcher + + Params: + alpha: relative balancing between classification and localization + threshold: threshold used to select positive predictions + topk: number of top scoring predictions to consider + """ + super().__init__() + self.norm = nn.Sigmoid() + self.alpha = alpha + self.threshold = threshold + self.topk = topk + + @torch.no_grad() + def forward( + self, + outputs, + batched_targets, + repeats=1, + repeat_batch=1, + out_is_valid=None, + target_is_valid_padded=None, + ): + """ + Performs the matching. The inputs and outputs are the same as + BinaryHungarianMatcher.forward + + Inputs: + - outputs: A dict with the following keys: + - "pred_logits": Tensor of shape (batch_size, num_queries, 1) with + classification logits + - "pred_boxes": Tensor of shape (batch_size, num_queries, 4) with + predicted box coordinates in cxcywh format. + - batched_targets: A dict of targets. There may be a variable number of + targets per batch entry; suppose that there are T_b targets for batch + entry 0 <= b < batch_size. It should have the following keys: + - "num_boxes": int64 Tensor of shape (batch_size,) giving the number + of ground-truth boxes per batch entry: num_boxes[b] = T_b + - "_boxes_padded": Tensor of shape (batch_size, max_b T_b, 4) giving + a padded version of ground-truth boxes. If this is not present then + it will be computed from batched_targets["boxes"] instead, but + caching it here can improve performance for repeated calls with the + same targets. + - out_is_valid: If not None, it should be a boolean tensor of shape + (batch_size, num_queries) indicating which predictions are valid. + Invalid predictions are ignored during matching and won't appear in + the output indices. + - target_is_valid_padded: If not None, it should be a boolean tensor of + shape (batch_size, max_num_gt_boxes) in padded format indicating + which GT boxes are valid. Invalid GT boxes are ignored during matching + and won't appear in the output indices. + Returns: + A list of size batch_size, containing tuples of (idx_i, idx_j): + - idx_i is the indices of the selected predictions (in order) + - idx_j is the indices of the corresponding selected targets + (in order) + For each batch element, it holds: + len(index_i) = len(index_j) + = min(num_queries, num_target_boxes) + """ + assert repeats <= 1 and repeat_batch <= 1 + bs, num_queries = outputs["pred_logits"].shape[:2] + + out_prob = self.norm(outputs["pred_logits"]).squeeze(-1) # (B, Q) + out_bbox = outputs["pred_boxes"] # (B, Q, 4)) + + num_boxes = batched_targets["num_boxes"] + + # Get a padded version of target boxes (as precomputed in the collator). + tgt_bbox = batched_targets["boxes_padded"] + assert len(tgt_bbox) == bs + num_targets = tgt_bbox.shape[1] + if num_targets == 0: + return ( + torch.empty(0, dtype=torch.long, device=out_prob.device), + torch.empty(0, dtype=torch.long, device=out_prob.device), + torch.empty(0, dtype=torch.long, device=out_prob.device), + ) + + iou, _ = box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) + + assert iou.shape == (bs, num_queries, num_targets) + + # Final cost matrix (higher is better in `C`; this is unlike the case + # of BinaryHungarianMatcherV2 above where lower is better in its `C`) + C = self.alpha * out_prob.unsqueeze(-1) + (1 - self.alpha) * iou + if out_is_valid is not None: + C = torch.where(out_is_valid[:, :, None], C, -1e9) + if target_is_valid_padded is not None: + C = torch.where(target_is_valid_padded[:, None, :], C, -1e9) + + # Selecting topk predictions + matches = C > torch.quantile( + C, 1 - self.topk / num_queries, dim=1, keepdim=True + ) + + # Selecting predictions above threshold + matches = matches & (C > self.threshold) + if out_is_valid is not None: + matches = matches & out_is_valid[:, :, None] + if target_is_valid_padded is not None: + matches = matches & target_is_valid_padded[:, None, :] + + # Removing padding + matches = matches & ( + torch.arange(0, num_targets, device=num_boxes.device)[None] + < num_boxes[:, None] + ).unsqueeze(1) + + batch_idx, src_idx, tgt_idx = torch.nonzero(matches, as_tuple=True) + + cum_num_boxes = torch.cat( + [ + torch.zeros(1, dtype=num_boxes.dtype, device=num_boxes.device), + num_boxes.cumsum(-1)[:-1], + ] + ) + tgt_idx += cum_num_boxes[batch_idx] + + return batch_idx, src_idx, tgt_idx diff --git a/src/nn/segearth_ov3/sam3/train/nms_helper.py b/src/nn/segearth_ov3/sam3/train/nms_helper.py new file mode 100644 index 0000000..cd5b6dc --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/nms_helper.py @@ -0,0 +1,306 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import warnings +from typing import Dict, List + +import numpy as np + +# Check if Numba is available +HAS_NUMBA = False +try: + import numba as nb + + HAS_NUMBA = True +except ImportError: + warnings.warn( + "Numba not found. Using slower pure Python implementations.", UserWarning + ) + + +# -------------------- Helper Functions -------------------- +def is_zero_box(bbox: list) -> bool: + """Check if bounding box is invalid""" + if bbox is None: + return True + return all(x <= 0 for x in bbox[:4]) or len(bbox) < 4 + + +def convert_bbox_format(bbox: list) -> List[float]: + """Convert bbox from (x,y,w,h) to (x1,y1,x2,y2)""" + x, y, w, h = bbox + return [x, y, x + w, y + h] + + +# -------------------- Track-level NMS -------------------- +def process_track_level_nms(video_groups: Dict, nms_threshold: float) -> Dict: + """Apply track-level NMS to all videos""" + for video_id, tracks in video_groups.items(): + track_detections = [] + + # Process tracks + for track_idx, track in enumerate(tracks): + if not track["bboxes"]: + continue + + converted_bboxes = [] + valid_frames = [] + for bbox in track["bboxes"]: + if bbox and not is_zero_box(bbox): + converted_bboxes.append(convert_bbox_format(bbox)) + valid_frames.append(True) + else: + converted_bboxes.append([np.nan] * 4) + valid_frames.append(False) + + if any(valid_frames): + track_detections.append( + { + "track_idx": track_idx, + "bboxes": np.array(converted_bboxes, dtype=np.float32), + "score": track["score"], + } + ) + + # Apply NMS + if track_detections: + scores = np.array([d["score"] for d in track_detections], dtype=np.float32) + keep = apply_track_nms(track_detections, scores, nms_threshold) + + # Suppress non-kept tracks + for idx, track in enumerate(track_detections): + if idx not in keep: + tracks[track["track_idx"]]["bboxes"] = [None] * len(track["bboxes"]) + + return video_groups + + +# -------------------- Frame-level NMS -------------------- +def process_frame_level_nms(video_groups: Dict, nms_threshold: float) -> Dict: + """Apply frame-level NMS to all videos""" + for video_id, tracks in video_groups.items(): + if not tracks: + continue + + num_frames = len(tracks[0]["bboxes"]) + + for frame_idx in range(num_frames): + frame_detections = [] + + # Collect valid detections + for track_idx, track in enumerate(tracks): + bbox = track["bboxes"][frame_idx] + if bbox and not is_zero_box(bbox): + frame_detections.append( + { + "track_idx": track_idx, + "bbox": np.array( + convert_bbox_format(bbox), dtype=np.float32 + ), + "score": track["score"], + } + ) + + # Apply NMS + if frame_detections: + bboxes = np.stack([d["bbox"] for d in frame_detections]) + scores = np.array( + [d["score"] for d in frame_detections], dtype=np.float32 + ) + keep = apply_frame_nms(bboxes, scores, nms_threshold) + + # Suppress non-kept detections + for i, d in enumerate(frame_detections): + if i not in keep: + tracks[d["track_idx"]]["bboxes"][frame_idx] = None + + return video_groups + + +# Track-level NMS helpers ------------------------------------------------------ +def compute_track_iou_matrix( + bboxes_stacked: np.ndarray, valid_masks: np.ndarray, areas: np.ndarray +) -> np.ndarray: + """IoU matrix computation for track-level NMS with fallback to pure Python""" + num_tracks = bboxes_stacked.shape[0] + iou_matrix = np.zeros((num_tracks, num_tracks), dtype=np.float32) + if HAS_NUMBA: + iou_matrix = _compute_track_iou_matrix_numba(bboxes_stacked, valid_masks, areas) + else: + # Pure Python implementation + for i in range(num_tracks): + for j in range(i + 1, num_tracks): + valid_ij = valid_masks[i] & valid_masks[j] + if not valid_ij.any(): + continue + bboxes_i = bboxes_stacked[i, valid_ij] + bboxes_j = bboxes_stacked[j, valid_ij] + area_i = areas[i, valid_ij] + area_j = areas[j, valid_ij] + inter_total = 0.0 + union_total = 0.0 + for k in range(bboxes_i.shape[0]): + x1 = max(bboxes_i[k, 0], bboxes_j[k, 0]) + y1 = max(bboxes_i[k, 1], bboxes_j[k, 1]) + x2 = min(bboxes_i[k, 2], bboxes_j[k, 2]) + y2 = min(bboxes_i[k, 3], bboxes_j[k, 3]) + inter = max(0, x2 - x1) * max(0, y2 - y1) + union = area_i[k] + area_j[k] - inter + inter_total += inter + union_total += union + if union_total > 0: + iou_matrix[i, j] = inter_total / union_total + iou_matrix[j, i] = iou_matrix[i, j] + return iou_matrix + + +if HAS_NUMBA: + + @nb.jit(nopython=True, parallel=True) + def _compute_track_iou_matrix_numba(bboxes_stacked, valid_masks, areas): + """Numba-optimized IoU matrix computation for track-level NMS""" + num_tracks = bboxes_stacked.shape[0] + iou_matrix = np.zeros((num_tracks, num_tracks), dtype=np.float32) + for i in nb.prange(num_tracks): + for j in range(i + 1, num_tracks): + valid_ij = valid_masks[i] & valid_masks[j] + if not valid_ij.any(): + continue + bboxes_i = bboxes_stacked[i, valid_ij] + bboxes_j = bboxes_stacked[j, valid_ij] + area_i = areas[i, valid_ij] + area_j = areas[j, valid_ij] + inter_total = 0.0 + union_total = 0.0 + for k in range(bboxes_i.shape[0]): + x1 = max(bboxes_i[k, 0], bboxes_j[k, 0]) + y1 = max(bboxes_i[k, 1], bboxes_j[k, 1]) + x2 = min(bboxes_i[k, 2], bboxes_j[k, 2]) + y2 = min(bboxes_i[k, 3], bboxes_j[k, 3]) + inter = max(0, x2 - x1) * max(0, y2 - y1) + union = area_i[k] + area_j[k] - inter + inter_total += inter + union_total += union + if union_total > 0: + iou_matrix[i, j] = inter_total / union_total + iou_matrix[j, i] = iou_matrix[i, j] + return iou_matrix + + +def apply_track_nms( + track_detections: List[dict], scores: np.ndarray, nms_threshold: float +) -> List[int]: + """Vectorized track-level NMS implementation""" + if not track_detections: + return [] + bboxes_stacked = np.stack([d["bboxes"] for d in track_detections], axis=0) + valid_masks = ~np.isnan(bboxes_stacked).any(axis=2) + areas = (bboxes_stacked[:, :, 2] - bboxes_stacked[:, :, 0]) * ( + bboxes_stacked[:, :, 3] - bboxes_stacked[:, :, 1] + ) + areas[~valid_masks] = 0 + iou_matrix = compute_track_iou_matrix(bboxes_stacked, valid_masks, areas) + keep = [] + order = np.argsort(-scores) + suppress = np.zeros(len(track_detections), dtype=bool) + for i in range(len(order)): + if not suppress[order[i]]: + keep.append(order[i]) + suppress[order[i:]] = suppress[order[i:]] | ( + iou_matrix[order[i], order[i:]] >= nms_threshold + ) + return keep + + +# Frame-level NMS helpers ------------------------------------------------------ +def compute_frame_ious(bbox: np.ndarray, bboxes: np.ndarray) -> np.ndarray: + """IoU computation for frame-level NMS with fallback to pure Python""" + if HAS_NUMBA: + return _compute_frame_ious_numba(bbox, bboxes) + else: + # Pure Python implementation + ious = np.zeros(len(bboxes), dtype=np.float32) + for i in range(len(bboxes)): + x1 = max(bbox[0], bboxes[i, 0]) + y1 = max(bbox[1], bboxes[i, 1]) + x2 = min(bbox[2], bboxes[i, 2]) + y2 = min(bbox[3], bboxes[i, 3]) + + inter = max(0, x2 - x1) * max(0, y2 - y1) + area1 = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) + area2 = (bboxes[i, 2] - bboxes[i, 0]) * (bboxes[i, 3] - bboxes[i, 1]) + union = area1 + area2 - inter + + ious[i] = inter / union if union > 0 else 0.0 + return ious + + +if HAS_NUMBA: + + @nb.jit(nopython=True, parallel=True) + def _compute_frame_ious_numba(bbox, bboxes): + """Numba-optimized IoU computation""" + ious = np.zeros(len(bboxes), dtype=np.float32) + for i in nb.prange(len(bboxes)): + x1 = max(bbox[0], bboxes[i, 0]) + y1 = max(bbox[1], bboxes[i, 1]) + x2 = min(bbox[2], bboxes[i, 2]) + y2 = min(bbox[3], bboxes[i, 3]) + + inter = max(0, x2 - x1) * max(0, y2 - y1) + area1 = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) + area2 = (bboxes[i, 2] - bboxes[i, 0]) * (bboxes[i, 3] - bboxes[i, 1]) + union = area1 + area2 - inter + + ious[i] = inter / union if union > 0 else 0.0 + return ious + + +def apply_frame_nms( + bboxes: np.ndarray, scores: np.ndarray, nms_threshold: float +) -> List[int]: + """Frame-level NMS implementation with fallback to pure Python""" + if HAS_NUMBA: + return _apply_frame_nms_numba(bboxes, scores, nms_threshold) + else: + # Pure Python implementation + order = np.argsort(-scores) + keep = [] + suppress = np.zeros(len(bboxes), dtype=bool) + + for i in range(len(order)): + if not suppress[order[i]]: + keep.append(order[i]) + current_bbox = bboxes[order[i]] + + remaining_bboxes = bboxes[order[i + 1 :]] + if len(remaining_bboxes) > 0: # Check if there are any remaining boxes + ious = compute_frame_ious(current_bbox, remaining_bboxes) + suppress[order[i + 1 :]] = suppress[order[i + 1 :]] | ( + ious >= nms_threshold + ) + + return keep + + +if HAS_NUMBA: + + @nb.jit(nopython=True) + def _apply_frame_nms_numba(bboxes, scores, nms_threshold): + """Numba-optimized NMS implementation""" + order = np.argsort(-scores) + keep = [] + suppress = np.zeros(len(bboxes), dtype=nb.boolean) + + for i in range(len(order)): + if not suppress[order[i]]: + keep.append(order[i]) + current_bbox = bboxes[order[i]] + + if i + 1 < len(order): # Check bounds + ious = _compute_frame_ious_numba( + current_bbox, bboxes[order[i + 1 :]] + ) + suppress[order[i + 1 :]] = suppress[order[i + 1 :]] | ( + ious >= nms_threshold + ) + + return keep diff --git a/src/nn/segearth_ov3/sam3/train/optim/__init__.py b/src/nn/segearth_ov3/sam3/train/optim/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/optim/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/train/optim/optimizer.py b/src/nn/segearth_ov3/sam3/train/optim/optimizer.py new file mode 100644 index 0000000..d401b98 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/optim/optimizer.py @@ -0,0 +1,498 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import fnmatch +import inspect +import itertools +import logging +import types +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Set, + Tuple, + Type, + Union, +) + +import hydra + +import torch +import torch.nn as nn +from omegaconf import DictConfig +from torch import Tensor + + +class Optimizer: + def __init__(self, optimizer, schedulers=None) -> None: + self.optimizer = optimizer + self.schedulers = schedulers + self._validate_optimizer_schedulers() + self.step_schedulers(0.0, 0) + + def _validate_optimizer_schedulers(self): + if self.schedulers is None: + return + for _, set_of_schedulers in enumerate(self.schedulers): + for option, _ in set_of_schedulers.items(): + assert option in self.optimizer.defaults, ( + "Optimizer option " + f"{option} not found in {self.optimizer}. Valid options are " + f"{self.optimizer.defaults.keys()}" + ) + + def step_schedulers(self, where: float, step: int) -> None: + if self.schedulers is None: + return + for i, param_group in enumerate(self.optimizer.param_groups): + for option, scheduler in self.schedulers[i].items(): + if "step" in inspect.signature(scheduler.__call__).parameters: + new_value = scheduler(step=step, where=where) + elif ( + hasattr(scheduler, "scheduler") + and "step" + in inspect.signature(scheduler.scheduler.__call__).parameters + ): + # To handle ValueScaler wrappers + new_value = scheduler(step=step, where=where) + else: + new_value = scheduler(where) + param_group[option] = new_value + + def step(self, where, step, closure=None): + self.step_schedulers(where, step) + return self.optimizer.step(closure) + + def zero_grad(self, *args, **kwargs): + return self.optimizer.zero_grad(*args, **kwargs) + + +def set_default_parameters( + scheduler_cfgs: List[DictConfig], all_parameter_names: Set[str] +) -> None: + """Set up the "default" scheduler with the right parameters. + + Args: + scheduler_cgfs: A list of scheduler configs, where each scheduler also + specifies which parameters it applies to, based on the names of parameters + or the class of the modules. At most one scheduler is allowed to skip this + specification, which is used as a "default" specification for any remaining + parameters. + all_parameter_names: Names of all the parameters to consider. + """ + constraints = [ + scheduler_cfg.parameter_names + for scheduler_cfg in scheduler_cfgs + if scheduler_cfg.parameter_names is not None + ] + if len(constraints) == 0: + default_params = set(all_parameter_names) + else: + default_params = all_parameter_names - set.union(*constraints) + default_count = 0 + for scheduler_cfg in scheduler_cfgs: + if scheduler_cfg.parameter_names is None: + scheduler_cfg.parameter_names = default_params + default_count += 1 + assert default_count <= 1, "Only one scheduler per option can be default" + if default_count == 0: + # No default scheduler specified, add a default, but without any scheduler + # for that option + scheduler_cfgs.append({"parameter_names": default_params}) + + +def name_constraints_to_parameters( + param_constraints: List[Set[str]], named_parameters: Dict[str, Tensor] +) -> List[torch.nn.Parameter]: + """Return parameters which match the intersection of parameter constraints. + + Note that this returns the parameters themselves, not their names. + + Args: + param_constraints: A list, with each element being a set of allowed parameters. + named_parameters: Mapping from a parameter name to the parameter itself. + + Returns: + A list containing the parameters which overlap with _each_ constraint set from + param_constraints. + """ + matching_names = set.intersection(*param_constraints) + return [value for name, value in named_parameters.items() if name in matching_names] + + +def map_scheduler_cfgs_to_param_groups( + all_scheduler_cfgs: Iterable[List[Dict]], + named_parameters: Dict[str, Tensor], +) -> Tuple[List[Dict[Any, Any]], List[Dict[str, List[torch.nn.Parameter]]]]: + """Produce parameter groups corresponding to all the scheduler configs. + + Takes all the scheduler configs, each of which applies to a specific optimizer + option (like "lr" or "weight_decay") and has a set of parameter names which it + applies to, and produces a final set of param groups where each param group + covers all the options which apply to a particular set of parameters. + + Args: + all_scheduler_cfgs: All the scheduler configs covering every option. + named_parameters: Mapping from a parameter name to the parameter itself. + Returns: + Tuple of lists of schedulers and param_groups, where schedulers[i] + applies to param_groups[i]. + """ + + scheduler_cfgs_per_param_group = itertools.product(*all_scheduler_cfgs) + schedulers = [] + param_groups = [] + for scheduler_cfgs in scheduler_cfgs_per_param_group: + param_constraints = [ + scheduler_cfg["parameter_names"] for scheduler_cfg in scheduler_cfgs + ] + matching_parameters = name_constraints_to_parameters( + param_constraints, named_parameters + ) + if len(matching_parameters) == 0: # If no overlap of parameters, skip + continue + schedulers_for_group = { + scheduler_cfg["option"]: scheduler_cfg["scheduler"] + for scheduler_cfg in scheduler_cfgs + if "option" in scheduler_cfg + } + schedulers.append(schedulers_for_group) + param_groups.append({"params": matching_parameters}) + return schedulers, param_groups + + +def validate_param_group_params(param_groups: List[Dict], model: nn.Module): + """Check that the param groups are non-overlapping and cover all the parameters. + + Args: + param_groups: List of all param groups + model: Model to validate against. The check ensures that all the model + parameters are part of param_groups + """ + for pg in param_groups: + # no param should be repeated within a group + assert len(pg["params"]) == len(set(pg["params"])) + parameters = [set(param_group["params"]) for param_group in param_groups] + model_parameters = {parameter for _, parameter in model.named_parameters()} + for p1, p2 in itertools.permutations(parameters, 2): + assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint" + assert set.union(*parameters) == model_parameters, ( + "Scheduler generated param_groups must include all parameters of the model." + f" Found {len(set.union(*parameters))} params whereas model has" + f" {len(model_parameters)} params" + ) + + +def unix_module_cls_pattern_to_parameter_names( + filter_module_cls_names: List[str], + module_cls_to_param_names: Dict[Type, str], +) -> Union[None, Set[str]]: + """Returns param names which pass the filters specified in filter_module_cls_names. + + Args: + filter_module_cls_names: A list of filter strings containing class names, like + ["torch.nn.LayerNorm", "torch.nn.BatchNorm2d"] + module_cls_to_param_names: Mapping from module classes to the parameter names + they contain. See `get_module_cls_to_param_names`. + """ + if filter_module_cls_names is None: + return set() + allowed_parameter_names = [] + for module_cls_name in filter_module_cls_names: + module_cls = hydra.utils.get_class(module_cls_name) + if module_cls not in module_cls_to_param_names: + raise AssertionError( + f"module_cls_name {module_cls_name} does not " + "match any classes in the model" + ) + matching_parameters = module_cls_to_param_names[module_cls] + assert ( + len(matching_parameters) > 0 + ), f"module_cls_name {module_cls_name} does not contain any parameters in the model" + logging.info( + f"Matches for module_cls_name [{module_cls_name}]: {matching_parameters} " + ) + allowed_parameter_names.append(matching_parameters) + return set.union(*allowed_parameter_names) + + +def unix_param_pattern_to_parameter_names( + filter_param_names: Optional[List[str]], + parameter_names: Dict[str, torch.Tensor], +) -> Union[None, Set[str]]: + """Returns param names which pass the filters specified in filter_param_names. + + Args: + filter_param_names: A list of unix-style filter strings with optional + wildcards, like ["block.2.*", "block.2.linear.weight"] + module_cls_to_param_names: Mapping from module classes to the parameter names + they contain. See `get_module_cls_to_param_names`. + """ + + if filter_param_names is None: + return set() + allowed_parameter_names = [] + for param_name in filter_param_names: + matching_parameters = set(fnmatch.filter(parameter_names, param_name)) + assert ( + len(matching_parameters) >= 1 + ), f"param_name {param_name} does not match any parameters in the model" + logging.info(f"Matches for param_name [{param_name}]: {matching_parameters}") + allowed_parameter_names.append(matching_parameters) + return set.union(*allowed_parameter_names) + + +def _unix_pattern_to_parameter_names( + scheduler_cfg: DictConfig, + parameter_names: Set[str], + module_cls_to_param_names: Dict[Type, str], +) -> Union[None, Set[str]]: + """Returns param names which pass the filters specified in scheduler_cfg. + + Args: + scheduler_cfg: The config for the scheduler + parameter_names: The set of all parameter names which will be filtered + """ + if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg: + return None + return unix_param_pattern_to_parameter_names( + scheduler_cfg.get("param_names"), parameter_names + ).union( + unix_module_cls_pattern_to_parameter_names( + scheduler_cfg.get("module_cls_names"), module_cls_to_param_names + ) + ) + + +def get_module_cls_to_param_names( + model: nn.Module, param_allowlist: Set[str] = None +) -> Dict[Type, str]: + """Produce a mapping from all the modules classes to the names of parames they own. + + Only counts a parameter as part of the immediate parent module, i.e. recursive + parents do not count. + + Args: + model: Model to iterate over + param_allowlist: If specified, only these param names will be processed + """ + + module_cls_to_params = {} + for module_name, module in model.named_modules(): + module_cls = type(module) + module_cls_to_params.setdefault(module_cls, set()) + for param_name, _ in module.named_parameters(recurse=False): + full_param_name = get_full_parameter_name(module_name, param_name) + if param_allowlist is None or full_param_name in param_allowlist: + module_cls_to_params[module_cls].add(full_param_name) + return module_cls_to_params + + +def construct_optimizer( + model: torch.nn.Module, + optimizer_conf: Any, + options_conf: Mapping[str, List] = None, + param_group_modifiers_conf: List[Callable] = None, + param_allowlist: Optional[Set[str]] = None, + validate_param_groups=True, +) -> Optimizer: + """ + Constructs a stochastic gradient descent or ADAM (or ADAMw) optimizer + with momentum. i.e, constructs a torch.optim.Optimizer with zero-weight decay + Batchnorm and/or no-update 1-D parameters support, based on the config. + + Supports wrapping the optimizer with Layer-wise Adaptive Rate Scaling + (LARS): https://arxiv.org/abs/1708.03888 + + Args: + model: model to perform stochastic gradient descent + optimization or ADAM optimization. + optimizer_conf: Hydra config consisting a partial torch optimizer like SGD or + ADAM, still missing the params argument which this function provides to + produce the final optimizer + param_group_modifiers_conf: Optional user specified functions which can modify + the final scheduler configs before the optimizer's param groups are built + param_allowlist: The parameters to optimize. Parameters which are not part of + this allowlist will be skipped. + validate_param_groups: If enabled, valides that the produced param_groups don't + overlap and cover all the model parameters. + """ + if param_allowlist is None: + param_allowlist = {name for name, _ in model.named_parameters()} + + named_parameters = { + name: param + for name, param in model.named_parameters() + if name in param_allowlist + } + + if not options_conf: + optimizer = hydra.utils.instantiate(optimizer_conf, named_parameters.values()) + return Optimizer(optimizer) + + all_parameter_names = { + name for name, _ in model.named_parameters() if name in param_allowlist + } + module_cls_to_all_param_names = get_module_cls_to_param_names( + model, param_allowlist + ) + + scheduler_cfgs_per_option = hydra.utils.instantiate(options_conf) + all_scheduler_cfgs = [] + for option, scheduler_cfgs in scheduler_cfgs_per_option.items(): + for config in scheduler_cfgs: + config.option = option + config.parameter_names = _unix_pattern_to_parameter_names( + config, all_parameter_names, module_cls_to_all_param_names + ) + set_default_parameters(scheduler_cfgs, all_parameter_names) + all_scheduler_cfgs.append(scheduler_cfgs) + + if param_group_modifiers_conf: + for custom_param_modifier in param_group_modifiers_conf: + custom_param_modifier = hydra.utils.instantiate(custom_param_modifier) + all_scheduler_cfgs = custom_param_modifier( + scheduler_cfgs=all_scheduler_cfgs, model=model + ) + schedulers, param_groups = map_scheduler_cfgs_to_param_groups( + all_scheduler_cfgs, named_parameters + ) + if validate_param_groups: + validate_param_group_params(param_groups, model) + optimizer = hydra.utils.instantiate(optimizer_conf, param_groups) + return Optimizer(optimizer, schedulers) + + +def get_full_parameter_name(module_name, param_name): + if module_name == "": + return param_name + return f"{module_name}.{param_name}" + + +class GradientClipper: + """ + Gradient clipping utils that works for DDP + """ + + def __init__(self, max_norm: float = 1.0, norm_type: int = 2): + assert isinstance(max_norm, (int, float)) or max_norm is None + self.max_norm = max_norm if max_norm is None else float(max_norm) + self.norm_type = norm_type + + def __call__(self, model: nn.Module): + if self.max_norm is None: + return # no-op + + nn.utils.clip_grad_norm_( + model.parameters(), max_norm=self.max_norm, norm_type=self.norm_type + ) + + +class ValueScaler: + def __init__(self, scheduler, mult_val: float): + self.scheduler = scheduler + self.mult_val = mult_val + + def __call__(self, *args, **kwargs): + val = self.scheduler(*args, **kwargs) + return val * self.mult_val + + +def rgetattr(obj, rattrs: str = None): + """ + Like getattr(), but supports dotted notation for nested objects. + rattrs is a str of form 'attr1.attr2', returns obj.attr1.attr2 + """ + if rattrs is None: + return obj + attrs = rattrs.split(".") + for attr in attrs: + obj = getattr(obj, attr) + return obj + + +def layer_decay_param_modifier( + scheduler_cfgs: List[List[Dict]], + model, + layer_decay_value: float, + layer_decay_min: Optional[float] = None, + apply_to: Optional[str] = None, + overrides: List[Dict] = (), +) -> List[List[Dict]]: + """ + Args + - scheduler_cfgs: a list of omegaconf.ListConfigs. + Each element in the list is a omegaconfg.DictConfig with the following structure + { + "scheduler": + "option": possible options are "lr", "weight_decay" etc. + "parameter_names": Set of str indicating param names that this scheduler applies to + } + - model: a model that implements a method `get_layer_id` that maps layer_name to an integer and + and a method get_num_layers. + Alternatively, use apply_to argument to select a specific component of the model. + - layer_decay_value: float + - layer_decay_min: min val for layer decay + - apply_to: optional arg to select which component of the model to apply the the layer decay modifier to + - overrides: to manually override lr for specific patterns. Is a list of dicts. Each dict, has keys "pattern", "value". + Returns + - scheduler_configs: same structure as the input, elements can be modified + """ + model = rgetattr(model, apply_to) + num_layers = model.get_num_layers() + 1 + layer_decays = [ + layer_decay_value ** (num_layers - i) for i in range(num_layers + 1) + ] + if layer_decay_min is not None: + layer_decays = [max(val, layer_decay_min) for val in layer_decays] + final_scheduler_cfgs = [] + # scheduler_cfgs is a list of lists + for scheduler_cfg_group in scheduler_cfgs: + curr_cfg_group = [] + # scheduler_cfg_group is a list of dictionaries + for scheduler_cfg in scheduler_cfg_group: + if scheduler_cfg["option"] != "lr": + curr_cfg_group.append(scheduler_cfg) + continue + # Need sorted so that the list of parameter names is deterministic and consistent + # across re-runs of this job. Else it was causing issues with loading the optimizer + # state during a job restart + parameter_names = sorted(scheduler_cfg["parameter_names"]) + + # Only want one cfg group per layer + layer_cfg_groups = {} + for param_name in parameter_names: + layer_id = num_layers + this_scale = layer_decays[layer_id] + if param_name.startswith(apply_to): + layer_id = model.get_layer_id(param_name) + this_scale = layer_decays[layer_id] + # Overrides + for override in overrides: + if fnmatch.fnmatchcase(param_name, override["pattern"]): + this_scale = float(override["value"]) + layer_id = override["pattern"] + break + + if layer_id not in layer_cfg_groups: + curr_param = { + "option": scheduler_cfg["option"], + "scheduler": ValueScaler( + scheduler_cfg["scheduler"], this_scale + ), + "parameter_names": {param_name}, + } + else: + curr_param = layer_cfg_groups[layer_id] + curr_param["parameter_names"].add(param_name) + layer_cfg_groups[layer_id] = curr_param + + for layer_cfg in layer_cfg_groups.values(): + curr_cfg_group.append(layer_cfg) + + final_scheduler_cfgs.append(curr_cfg_group) + return final_scheduler_cfgs diff --git a/src/nn/segearth_ov3/sam3/train/optim/schedulers.py b/src/nn/segearth_ov3/sam3/train/optim/schedulers.py new file mode 100644 index 0000000..59da840 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/optim/schedulers.py @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math + + +class InverseSquareRootParamScheduler: + def __init__( + self, + base_lr: float, + warmup_steps: int, + cooldown_steps: int, + timescale: int, + ): + self.base_lr = base_lr + self.warmup_steps = warmup_steps + self.cooldown_steps = cooldown_steps + self.timescale = timescale + + def __call__(self, step: int, where: float): + lr = self.base_lr + + if where > 0: + total_steps = step / where + progress = (step - self.warmup_steps) / float( + total_steps - self.warmup_steps + ) + progress = max(min(progress, 1), 0) + else: + progress = 0 + total_steps = 1 + + shift = self.timescale - self.warmup_steps + if self.warmup_steps < step: + lr = lr / math.sqrt((step + shift) / self.timescale) + + if self.warmup_steps: + lr = lr * min(1.0, step / self.warmup_steps) + if self.cooldown_steps: + lr = lr * min(1.0, (total_steps - step) / self.cooldown_steps) + + return lr diff --git a/src/nn/segearth_ov3/sam3/train/train.py b/src/nn/segearth_ov3/sam3/train/train.py new file mode 100644 index 0000000..b3e995c --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/train.py @@ -0,0 +1,339 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +import os +import random +import sys +import traceback +from argparse import ArgumentParser +from copy import deepcopy + +import submitit +import torch + +from hydra import compose, initialize_config_module +from hydra.utils import instantiate + +from iopath.common.file_io import g_pathmgr +from omegaconf import OmegaConf + +from sam3.train.utils.train_utils import makedir, register_omegaconf_resolvers +from tqdm import tqdm + + +os.environ["HYDRA_FULL_ERROR"] = "1" + + +class SlurmEvent: + QUEUED = "QUEUED" + START = "START" + FINISH = "FINISH" + JOB_ERROR = "JOB_ERROR" + SLURM_SIGNAL = "SLURM_SIGNAL" + + +def handle_custom_resolving(cfg): + # We'll resolve the config here, so we can catch mistakes early. + # However, we need to pass the un-resolved config to the launcher + # (because DVC resolving needs to be done on the node it will run on) + # First, do a copy without triggering resolving + cfg_resolved = OmegaConf.to_container(cfg, resolve=False) + cfg_resolved = OmegaConf.create(cfg_resolved) + return cfg_resolved + + +def single_proc_run(local_rank, main_port, cfg, world_size): + """Single GPU process""" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(main_port) + os.environ["RANK"] = str(local_rank) + os.environ["LOCAL_RANK"] = str(local_rank) + os.environ["WORLD_SIZE"] = str(world_size) + try: + register_omegaconf_resolvers() + except Exception as e: + logging.info(e) + + trainer = instantiate(cfg.trainer, _recursive_=False) + trainer.run() + + +def single_node_runner(cfg, main_port: int): + assert cfg.launcher.num_nodes == 1 + # assert cfg.launcher.gpus_per_node == 1 + num_proc = cfg.launcher.gpus_per_node + torch.multiprocessing.set_start_method( + "spawn" + ) # CUDA runtime does not support `fork` + if num_proc == 1: + # directly call single_proc so we can easily set breakpoints + # mp.spawn does not let us set breakpoints + single_proc_run(local_rank=0, main_port=main_port, cfg=cfg, world_size=num_proc) + else: + mp_runner = torch.multiprocessing.start_processes + args = (main_port, cfg, num_proc) + # Note: using "fork" below, "spawn" causes time and error regressions. Using + # spawn changes the default multiprocessing context to spawn, which doesn't + # interact well with the dataloaders (likely due to the use of OpenCV). + mp_runner(single_proc_run, args=args, nprocs=num_proc, start_method="spawn") + + +def format_exception(e: Exception, limit=20): + traceback_str = "".join(traceback.format_tb(e.__traceback__, limit=limit)) + return f"{type(e).__name__}: {e}\nTraceback:\n{traceback_str}" + + +class SubmititRunner(submitit.helpers.Checkpointable): + """A callable which is passed to submitit to launch the jobs.""" + + def __init__(self, port, cfg): + self.cfg = cfg + self.port = port + self.has_setup = False + + def run_trainer(self): + job_env = submitit.JobEnvironment() + # Need to add this again so the hydra.job.set_env PYTHONPATH + # is also set when launching jobs. + add_pythonpath_to_sys_path() + os.environ["MASTER_ADDR"] = job_env.hostnames[0] + os.environ["MASTER_PORT"] = str(self.port) + os.environ["RANK"] = str(job_env.global_rank) + os.environ["LOCAL_RANK"] = str(job_env.local_rank) + os.environ["WORLD_SIZE"] = str(job_env.num_tasks) + + register_omegaconf_resolvers() + cfg_resolved = OmegaConf.to_container(self.cfg, resolve=False) + cfg_resolved = OmegaConf.create(cfg_resolved) + + trainer = instantiate(cfg_resolved.trainer, _recursive_=False) + trainer.run() + + def __call__(self): + job_env = submitit.JobEnvironment() + self.setup_job_info(job_env.job_id, job_env.global_rank) + try: + self.run_trainer() + except Exception as e: + # Log the exception. Then raise it again (as what SubmititRunner currently does). + message = format_exception(e) + logging.error(message) + raise e + + def setup_job_info(self, job_id, rank): + """Set up slurm job info""" + self.job_info = { + "job_id": job_id, + "rank": rank, + "cluster": self.cfg.get("cluster", None), + "experiment_log_dir": self.cfg.launcher.experiment_log_dir, + } + + self.has_setup = True + + +def add_pythonpath_to_sys_path(): + if "PYTHONPATH" not in os.environ or not os.environ["PYTHONPATH"]: + return + sys.path = os.environ["PYTHONPATH"].split(":") + sys.path + + +def main(args) -> None: + cfg = compose(config_name=args.config) + if cfg.launcher.experiment_log_dir is None: + cfg.launcher.experiment_log_dir = os.path.join( + os.getcwd(), "sam3_logs", args.config + ) + print("###################### Train App Config ####################") + print(OmegaConf.to_yaml(cfg)) + print("############################################################") + + add_pythonpath_to_sys_path() + makedir(cfg.launcher.experiment_log_dir) + with g_pathmgr.open( + os.path.join(cfg.launcher.experiment_log_dir, "config.yaml"), "w" + ) as f: + f.write(OmegaConf.to_yaml(cfg)) + + cfg_resolved = OmegaConf.to_container(cfg, resolve=False) + cfg_resolved = OmegaConf.create(cfg_resolved) + + with g_pathmgr.open( + os.path.join(cfg.launcher.experiment_log_dir, "config_resolved.yaml"), "w" + ) as f: + f.write(OmegaConf.to_yaml(cfg_resolved, resolve=True)) + + submitit_conf = cfg.get("submitit", None) + assert submitit_conf is not None, "Missing submitit config" + + experiment_log_dir = cfg.launcher.experiment_log_dir + print(f"Experiment Log Dir:\n{experiment_log_dir}") + submitit_dir = os.path.join(experiment_log_dir, "submitit_logs") + + # Prioritize cmd line args + cfg.launcher.gpus_per_node = ( + args.num_gpus if args.num_gpus is not None else cfg.launcher.gpus_per_node + ) + cfg.launcher.num_nodes = ( + args.num_nodes if args.num_nodes is not None else cfg.launcher.num_nodes + ) + submitit_conf.use_cluster = ( + args.use_cluster if args.use_cluster is not None else submitit_conf.use_cluster + ) + if submitit_conf.use_cluster: + executor = submitit.AutoExecutor(folder=submitit_dir) + submitit_conf.partition = ( + args.partition + if args.partition is not None + else submitit_conf.get("partition", None) + ) + submitit_conf.account = ( + args.account + if args.account is not None + else submitit_conf.get("account", None) + ) + submitit_conf.qos = ( + args.qos if args.qos is not None else submitit_conf.get("qos", None) + ) + job_kwargs = { + "timeout_min": 60 * submitit_conf.timeout_hour, + "name": ( + submitit_conf.name if hasattr(submitit_conf, "name") else args.config + ), + "slurm_partition": submitit_conf.partition, + "gpus_per_node": cfg.launcher.gpus_per_node, + "tasks_per_node": cfg.launcher.gpus_per_node, # one task per GPU + "cpus_per_task": submitit_conf.cpus_per_task, + "nodes": cfg.launcher.num_nodes, + "slurm_additional_parameters": { + "exclude": " ".join(submitit_conf.get("exclude_nodes", [])), + }, + } + if "include_nodes" in submitit_conf: + assert ( + len(submitit_conf["include_nodes"]) >= cfg.launcher.num_nodes + ), "Not enough nodes" + job_kwargs["slurm_additional_parameters"]["nodelist"] = " ".join( + submitit_conf["include_nodes"] + ) + if submitit_conf.account is not None: + job_kwargs["slurm_additional_parameters"]["account"] = submitit_conf.account + if submitit_conf.qos is not None: + job_kwargs["slurm_additional_parameters"]["qos"] = submitit_conf.qos + + if submitit_conf.get("mem_gb", None) is not None: + job_kwargs["mem_gb"] = submitit_conf.mem_gb + elif submitit_conf.get("mem", None) is not None: + job_kwargs["slurm_mem"] = submitit_conf.mem + + if submitit_conf.get("constraints", None) is not None: + job_kwargs["slurm_constraint"] = submitit_conf.constraints + + if submitit_conf.get("comment", None) is not None: + job_kwargs["slurm_comment"] = submitit_conf.comment + + # Supports only cpu-bind option within srun_args. New options can be added here + if submitit_conf.get("srun_args", None) is not None: + job_kwargs["slurm_srun_args"] = [] + if submitit_conf.srun_args.get("cpu_bind", None) is not None: + job_kwargs["slurm_srun_args"].extend( + ["--cpu-bind", submitit_conf.srun_args.cpu_bind] + ) + + print("###################### SLURM Config ####################") + print(job_kwargs) + print("##########################################") + executor.update_parameters(**job_kwargs) + + if ( + "job_array" in submitit_conf + and submitit_conf.job_array.get("num_tasks", -1) > 0 + ): + num_tasks = submitit_conf.job_array.num_tasks + job_array_config_dir = os.path.join( + cfg.launcher.experiment_log_dir, "job_array_configs" + ) + makedir(job_array_config_dir) + + job_indices = range(num_tasks) + ports = random.sample( + range(submitit_conf.port_range[0], submitit_conf.port_range[1] + 1), + k=len(job_indices), + ) + + jobs_runners_configs = [] + with executor.batch(): + task_index = 0 + for indices, main_port in tqdm(zip(job_indices, ports)): + curr_cfg = deepcopy(cfg) + curr_cfg.submitit.job_array["task_index"] = task_index + curr_cfg_resolved = handle_custom_resolving(cfg) + runner = SubmititRunner(main_port, curr_cfg) + job = executor.submit(runner) + jobs_runners_configs.append( + (job, runner, curr_cfg, curr_cfg_resolved) + ) + task_index += 1 + + for job, runner, job_cfg, job_cfg_resolved in jobs_runners_configs: + print("Submitit Job ID:", job.job_id) + + # Save job specific config + job_array_config_file = os.path.join( + job_array_config_dir, "{}.config.yaml".format(job.job_id) + ) + with g_pathmgr.open(job_array_config_file, "w") as f: + f.write(OmegaConf.to_yaml(job_cfg)) + + job_array_config_resolved_file = os.path.join( + job_array_config_dir, "{}.config_resolved.yaml".format(job.job_id) + ) + with g_pathmgr.open(job_array_config_resolved_file, "w") as f: + f.write(OmegaConf.to_yaml(job_cfg_resolved, resolve=True)) + + runner.setup_job_info(job.job_id, rank=0) + # runner.log_event(event_type=SlurmEvent.QUEUED) + else: + main_port = random.randint( + submitit_conf.port_range[0], submitit_conf.port_range[1] + ) + runner = SubmititRunner(main_port, cfg) + job = executor.submit(runner) + print(f"Submitit Job ID: {job.job_id}") + runner.setup_job_info(job.job_id, rank=0) + + else: + cfg.launcher.num_nodes = 1 + main_port = random.randint( + submitit_conf.port_range[0], submitit_conf.port_range[1] + ) + single_node_runner(cfg, main_port) + + +if __name__ == "__main__": + initialize_config_module("sam3.train", version_base="1.2") + parser = ArgumentParser() + parser.add_argument( + "-c", + "--config", + required=True, + type=str, + help="path to config file (e.g. configs/roboflow_v100_full_ft_100_images.yaml)", + ) + parser.add_argument( + "--use-cluster", + type=int, + default=None, + help="whether to launch on a cluster, 0: run locally, 1: run on a cluster", + ) + parser.add_argument("--partition", type=str, default=None, help="SLURM partition") + parser.add_argument("--account", type=str, default=None, help="SLURM account") + parser.add_argument("--qos", type=str, default=None, help="SLURM qos") + parser.add_argument( + "--num-gpus", type=int, default=None, help="number of GPUS per node" + ) + parser.add_argument("--num-nodes", type=int, default=None, help="Number of nodes") + args = parser.parse_args() + args.use_cluster = bool(args.use_cluster) if args.use_cluster is not None else None + register_omegaconf_resolvers() + main(args) diff --git a/src/nn/segearth_ov3/sam3/train/trainer.py b/src/nn/segearth_ov3/sam3/train/trainer.py new file mode 100644 index 0000000..ac7c1b5 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/trainer.py @@ -0,0 +1,1193 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import contextlib +import fnmatch +import gc +import json +import logging +import math +import os +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional + +import numpy as np + +import torch +import torch.distributed as dist +import torch.nn as nn +from hydra.utils import instantiate +from iopath.common.file_io import g_pathmgr + +from sam3.model.data_misc import BatchedDatapoint +from sam3.model.model_misc import SAM3Output +from sam3.model.utils.misc import copy_data_to_device + +from sam3.train.optim.optimizer import construct_optimizer + +from sam3.train.utils.checkpoint_utils import ( + assert_skipped_parameters_are_frozen, + exclude_params_matching_unix_pattern, + load_state_dict_into_model, + with_check_parameter_frozen, +) + +from sam3.train.utils.distributed import all_reduce_max, barrier, get_rank + +from sam3.train.utils.logger import Logger, setup_logging +from sam3.train.utils.train_utils import ( + AverageMeter, + collect_dict_keys, + DurationMeter, + get_amp_type, + get_machine_local_and_dist_rank, + get_resume_checkpoint, + human_readable_time, + is_dist_avail_and_initialized, + log_env_variables, + makedir, + MemMeter, + Phase, + ProgressMeter, + set_seeds, + setup_distributed_backend, +) + + +CORE_LOSS_KEY = "core_loss" + + +def unwrap_ddp_if_wrapped(model): + if isinstance(model, torch.nn.parallel.DistributedDataParallel): + return model.module + return model + + +@dataclass +class OptimAMPConf: + enabled: bool = False + amp_dtype: str = "float16" + + +@dataclass +class OptimConf: + optimizer: torch.optim.Optimizer = None + options: Optional[Dict[str, Any]] = None + param_group_modifiers: Optional[List] = None + amp: Optional[Dict[str, Any]] = None + gradient_clip: Any = None + gradient_logger: Any = None + + def __post_init__(self): + # amp + if not isinstance(self.amp, OptimAMPConf): + if self.amp is None: + self.amp = {} + assert isinstance(self.amp, Mapping) + self.amp = OptimAMPConf(**self.amp) + + +@dataclass +class DistributedConf: + backend: Optional[str] = None # inferred from accelerator type + comms_dtype: Optional[str] = None + find_unused_parameters: bool = False + timeout_mins: int = 30 + gradient_as_bucket_view: bool = False # PyTorch DDP default is False + static_graph: bool = False # PyTorch DDP default is False + + +@dataclass +class CudaConf: + cudnn_deterministic: bool = False + cudnn_benchmark: bool = True + allow_tf32: bool = False + # if not None, `matmul_allow_tf32` key will override `allow_tf32` for matmul + matmul_allow_tf32: Optional[bool] = None + # if not None, `cudnn_allow_tf32` key will override `allow_tf32` for cudnn + cudnn_allow_tf32: Optional[bool] = None + + +@dataclass +class CheckpointConf: + save_dir: str + save_freq: int + save_list: List[int] = field(default_factory=list) + model_weight_initializer: Any = None + save_best_meters: List[str] = None + skip_saving_parameters: List[str] = field(default_factory=list) + initialize_after_preemption: Optional[bool] = None + # if not None, training will be resumed from this checkpoint + resume_from: Optional[str] = None + + def infer_missing(self): + if self.initialize_after_preemption is None: + with_skip_saving = len(self.skip_saving_parameters) > 0 + self.initialize_after_preemption = with_skip_saving + return self + + +@dataclass +class LoggingConf: + log_dir: str + log_freq: int # In iterations + tensorboard_writer: Any + log_level_primary: str = "INFO" + log_level_secondary: str = "ERROR" + log_scalar_frequency: int = 100 + log_visual_frequency: int = 100 + scalar_keys_to_log: Optional[Dict[str, Any]] = None + log_batch_stats: bool = False + wandb_writer: Optional[Any] = None + + +class Trainer: + """ + Trainer supporting the DDP training strategies. + """ + + EPSILON = 1e-8 + + def __init__( + self, + *, # the order of these args can change at any time, so they are keyword-only + data: Dict[str, Any], + model: Dict[str, Any], + logging: Dict[str, Any], + checkpoint: Dict[str, Any], + max_epochs: int, + mode: str = "train", + accelerator: str = "cuda", + seed_value: int = 123, + val_epoch_freq: int = 1, + distributed: Dict[str, bool] = None, + cuda: Dict[str, bool] = None, + env_variables: Optional[Dict[str, Any]] = None, + optim: Optional[Dict[str, Any]] = None, + optim_overrides: Optional[List[Dict[str, Any]]] = None, + meters: Optional[Dict[str, Any]] = None, + loss: Optional[Dict[str, Any]] = None, + skip_first_val: bool = False, + skip_saving_ckpts: bool = False, + empty_gpu_mem_cache_after_eval: bool = True, + gradient_accumulation_steps: int = 1, + ): + self._setup_env_variables(env_variables) + self._setup_timers() + + self.data_conf = data + self.model_conf = model + self.logging_conf = LoggingConf(**logging) + self.checkpoint_conf = CheckpointConf(**checkpoint).infer_missing() + self.max_epochs = max_epochs + self.mode = mode + self.val_epoch_freq = val_epoch_freq + self.optim_conf = OptimConf(**optim) if optim is not None else OptimConf() + self.meters_conf = meters + self.loss_conf = loss + self.gradient_accumulation_steps = gradient_accumulation_steps + distributed = DistributedConf(**distributed or {}) + cuda = CudaConf(**cuda or {}) + self.where = 0.0 + + self.skip_first_val = skip_first_val + self.skip_saving_ckpts = skip_saving_ckpts + self.empty_gpu_mem_cache_after_eval = empty_gpu_mem_cache_after_eval + + self._infer_distributed_backend_if_none(distributed, accelerator) + + self._setup_device(accelerator) + + self._setup_torch_dist_and_backend(cuda, distributed) + + makedir(self.logging_conf.log_dir) + setup_logging( + __name__, + output_dir=self.logging_conf.log_dir, + rank=self.rank, + log_level_primary=self.logging_conf.log_level_primary, + log_level_secondary=self.logging_conf.log_level_secondary, + ) + + set_seeds(seed_value, self.max_epochs, self.distributed_rank) + log_env_variables() + + assert ( + is_dist_avail_and_initialized() + ), "Torch distributed needs to be initialized before calling the trainer." + + self._setup_components() # Except Optimizer everything is setup here. + self._move_to_device() + self._construct_optimizers() + self._setup_dataloaders() + + self.time_elapsed_meter = DurationMeter("Time Elapsed", self.device, ":.2f") + + if self.checkpoint_conf.resume_from is not None: + assert os.path.exists( + self.checkpoint_conf.resume_from + ), f"The 'resume_from' checkpoint {self.checkpoint_conf.resume_from} does not exist!" + dst = os.path.join(self.checkpoint_conf.save_dir, "checkpoint.pt") + if self.distributed_rank == 0 and not os.path.exists(dst): + # Copy the "resume_from" checkpoint to the checkpoint folder + # if there is not a checkpoint to resume from already there + makedir(self.checkpoint_conf.save_dir) + g_pathmgr.copy(self.checkpoint_conf.resume_from, dst) + barrier() + + self.load_checkpoint() + self._setup_ddp_distributed_training(distributed, accelerator) + barrier() + + def _setup_timers(self): + """ + Initializes counters for elapsed time and eta. + """ + self.start_time = time.time() + self.ckpt_time_elapsed = 0 + self.est_epoch_time = dict.fromkeys([Phase.TRAIN, Phase.VAL], 0) + + def _get_meters(self, phase_filters=None): + if self.meters is None: + return {} + meters = {} + for phase, phase_meters in self.meters.items(): + if phase_filters is not None and phase not in phase_filters: + continue + for key, key_meters in phase_meters.items(): + if key_meters is None: + continue + for name, meter in key_meters.items(): + meters[f"{phase}_{key}/{name}"] = meter + return meters + + def _infer_distributed_backend_if_none(self, distributed_conf, accelerator): + if distributed_conf.backend is None: + distributed_conf.backend = "nccl" if accelerator == "cuda" else "gloo" + + def _setup_env_variables(self, env_variables_conf) -> None: + if env_variables_conf is not None: + for variable_name, value in env_variables_conf.items(): + os.environ[variable_name] = value + + def _setup_torch_dist_and_backend(self, cuda_conf, distributed_conf) -> None: + if torch.cuda.is_available(): + torch.backends.cudnn.deterministic = cuda_conf.cudnn_deterministic + torch.backends.cudnn.benchmark = cuda_conf.cudnn_benchmark + torch.backends.cuda.matmul.allow_tf32 = ( + cuda_conf.matmul_allow_tf32 + if cuda_conf.matmul_allow_tf32 is not None + else cuda_conf.allow_tf32 + ) + torch.backends.cudnn.allow_tf32 = ( + cuda_conf.cudnn_allow_tf32 + if cuda_conf.cudnn_allow_tf32 is not None + else cuda_conf.allow_tf32 + ) + + self.rank = setup_distributed_backend( + distributed_conf.backend, distributed_conf.timeout_mins + ) + + def _setup_device(self, accelerator): + self.local_rank, self.distributed_rank = get_machine_local_and_dist_rank() + if accelerator == "cuda": + self.device = torch.device("cuda", self.local_rank) + torch.cuda.set_device(self.local_rank) + elif accelerator == "cpu": + self.device = torch.device("cpu") + else: + raise ValueError(f"Unsupported accelerator: {accelerator}") + + def _setup_ddp_distributed_training(self, distributed_conf, accelerator): + assert isinstance(self.model, torch.nn.Module) + + self.model = nn.parallel.DistributedDataParallel( + self.model, + device_ids=[self.local_rank] if accelerator == "cuda" else [], + find_unused_parameters=distributed_conf.find_unused_parameters, + gradient_as_bucket_view=distributed_conf.gradient_as_bucket_view, + static_graph=distributed_conf.static_graph, + ) + if distributed_conf.comms_dtype is not None: # noqa + from torch.distributed.algorithms import ddp_comm_hooks + + amp_type = get_amp_type(distributed_conf.comms_dtype) + if amp_type == torch.bfloat16: + hook = ddp_comm_hooks.default_hooks.bf16_compress_hook + logging.info("Enabling bfloat16 grad communication") + else: + hook = ddp_comm_hooks.default_hooks.fp16_compress_hook + logging.info("Enabling fp16 grad communication") + process_group = None + self.model.register_comm_hook(process_group, hook) + + def _move_to_device(self): + logging.info( + f"Moving components to device {self.device} and local rank {self.local_rank}." + ) + + self.model.to(self.device) + + logging.info( + f"Done moving components to device {self.device} and local rank {self.local_rank}." + ) + + def save_checkpoint(self, epoch, checkpoint_names=None): + if self.skip_saving_ckpts: + logging.info( + "skip_saving_ckpts is set to True. So, no checkpoints have been saved." + ) + return + checkpoint_folder = self.checkpoint_conf.save_dir + makedir(checkpoint_folder) + if checkpoint_names is None: + checkpoint_names = ["checkpoint"] + if ( + self.checkpoint_conf.save_freq > 0 + and (int(epoch) % self.checkpoint_conf.save_freq == 0) + ) or int(epoch) in self.checkpoint_conf.save_list: + checkpoint_names.append(f"checkpoint_{int(epoch)}") + + checkpoint_paths = [] + for ckpt_name in checkpoint_names: + checkpoint_paths.append(os.path.join(checkpoint_folder, f"{ckpt_name}.pt")) + + state_dict = unwrap_ddp_if_wrapped(self.model).state_dict() + state_dict = exclude_params_matching_unix_pattern( + patterns=self.checkpoint_conf.skip_saving_parameters, state_dict=state_dict + ) + + checkpoint = { + "model": state_dict, + "optimizer": self.optim.optimizer.state_dict(), + "epoch": epoch, + "loss": self.loss.state_dict(), + "steps": self.steps, + "time_elapsed": self.time_elapsed_meter.val, + "best_meter_values": self.best_meter_values, + } + if self.optim_conf.amp.enabled: + checkpoint["scaler"] = self.scaler.state_dict() + + # DDP checkpoints are only saved on rank 0 (all workers are identical) + if self.distributed_rank != 0: + return + + for checkpoint_path in checkpoint_paths: + self._save_checkpoint(checkpoint, checkpoint_path) + + def _save_checkpoint(self, checkpoint, checkpoint_path): + """ + Save a checkpoint while guarding against the job being killed in the middle + of checkpoint saving (which corrupts the checkpoint file and ruins the + entire training since usually only the last checkpoint is kept per run). + + We first save the new checkpoint to a temp file (with a '.tmp' suffix), and + and move it to overwrite the old checkpoint_path. + """ + checkpoint_path_tmp = f"{checkpoint_path}.tmp" + with g_pathmgr.open(checkpoint_path_tmp, "wb") as f: + torch.save(checkpoint, f) + # after torch.save is completed, replace the old checkpoint with the new one + if g_pathmgr.exists(checkpoint_path): + # remove the old checkpoint_path file first (otherwise g_pathmgr.mv fails) + g_pathmgr.rm(checkpoint_path) + success = g_pathmgr.mv(checkpoint_path_tmp, checkpoint_path) + assert success + + def load_checkpoint(self): + ckpt_path = get_resume_checkpoint(self.checkpoint_conf.save_dir) + if ckpt_path is None: + self._init_model_state() + else: + if self.checkpoint_conf.initialize_after_preemption: + self._call_model_initializer() + self._load_resuming_checkpoint(ckpt_path) + + def _init_model_state(self): + # Checking that parameters that won't be saved are indeed frozen + # We do this check here before even saving the model to catch errors + # are early as possible and not at the end of the first epoch + assert_skipped_parameters_are_frozen( + patterns=self.checkpoint_conf.skip_saving_parameters, + model=self.model, + ) + + # Checking that parameters that won't be saved are initialized from + # within the model definition, unless `initialize_after_preemption` + # is explicitly set to `True`. If not, this is a bug, and after + # preemption, the `skip_saving_parameters` will have random values + allow_init_skip_parameters = self.checkpoint_conf.initialize_after_preemption + with with_check_parameter_frozen( + patterns=self.checkpoint_conf.skip_saving_parameters, + model=self.model, + disabled=allow_init_skip_parameters, + ): + self._call_model_initializer() + + def _call_model_initializer(self): + model_weight_initializer = instantiate( + self.checkpoint_conf.model_weight_initializer + ) + if model_weight_initializer is not None: + logging.info( + f"Loading pretrained checkpoint from {self.checkpoint_conf.model_weight_initializer}" + ) + self.model = model_weight_initializer(model=self.model) + + def _load_resuming_checkpoint(self, ckpt_path: str): + logging.info(f"Resuming training from {ckpt_path}") + + with g_pathmgr.open(ckpt_path, "rb") as f: + checkpoint = torch.load(f, map_location="cpu") + load_state_dict_into_model( + model=self.model, + state_dict=checkpoint["model"], + ignore_missing_keys=self.checkpoint_conf.skip_saving_parameters, + ) + + self.optim.optimizer.load_state_dict(checkpoint["optimizer"]) + self.loss.load_state_dict(checkpoint["loss"], strict=True) + self.epoch = checkpoint["epoch"] + self.steps = checkpoint["steps"] + self.ckpt_time_elapsed = checkpoint.get("time_elapsed") + + if self.optim_conf.amp.enabled and "scaler" in checkpoint: + self.scaler.load_state_dict(checkpoint["scaler"]) + + self.best_meter_values = checkpoint.get("best_meter_values", {}) + + if "train_dataset" in checkpoint and self.train_dataset is not None: + self.train_dataset.load_checkpoint_state(checkpoint["train_dataset"]) + + def is_intermediate_val_epoch(self, epoch): + skip_epoch = self.skip_first_val and epoch == 0 + return ( + epoch % self.val_epoch_freq == 0 + and epoch < self.max_epochs - 1 + and not skip_epoch + ) + + def _find_loss(self, key: str): + if key in self.loss: + return self.loss[key] + + assert key != "all", "Loss must be specified for key='all'" + assert ( + "default" in self.loss + ), f"Key {key} not found in losss, and no default provided" + return self.loss["default"] + + def _find_meter(self, phase: str, key: str): + if key in self.meters[phase]: + return self.meters[phase][key] + + for cand_key, meter in self.meters[phase].items(): + if fnmatch.fnmatch(key, cand_key): + return meter + return None + + def _step( + self, + batch: BatchedDatapoint, + model: nn.Module, + phase: str, + ): + key, batch = batch.popitem() + batch = copy_data_to_device(batch, self.device, non_blocking=True) + + find_stages = model(batch) + find_targets = [ + unwrap_ddp_if_wrapped(model).back_convert(x) for x in batch.find_targets + ] + batch_size = len(batch.img_batch) + loss = self._find_loss(key)(find_stages, find_targets) + + loss_str = f"Losses/{phase}_{key}_loss" + + loss_log_str = os.path.join("Step_Losses", loss_str) + + # loss contains multiple sub-components we wish to log + step_losses = {} + if isinstance(loss, dict): + step_losses.update( + {f"Losses/{phase}_{key}_{k}": v for k, v in loss.items()} + ) + loss = self._log_loss_detailed_and_return_core_loss( + loss, loss_log_str, self.steps[phase] + ) + + if self.steps[phase] % self.logging_conf.log_scalar_frequency == 0: + self.logger.log( + loss_log_str, + loss, + self.steps[phase], + ) + + self.steps[phase] += 1 + + ret_tuple = {loss_str: loss}, batch_size, step_losses + + if phase not in self.meters: + return ret_tuple + + meters_dict = self._find_meter(phase, key) + if meters_dict is None: + return ret_tuple + if meters_dict is not None: + for _, meter in meters_dict.items(): + meter.update( + find_stages=find_stages, + find_metadatas=batch.find_metadatas, + model=model, + batch=batch, + key=key, + ) + # Cleanup memory + if isinstance(find_stages, SAM3Output): + for fs in find_stages: + for k in list(fs.keys()): + del fs[k] + + return ret_tuple + + def run(self): + assert self.mode in ["train", "train_only", "val"] + if self.mode == "train": + if self.epoch > 0: + logging.info(f"Resuming training from epoch: {self.epoch}") + # resuming from a checkpoint + if self.is_intermediate_val_epoch(self.epoch - 1): + logging.info("Running previous val epoch") + self.epoch -= 1 + self.run_val() + self.epoch += 1 + self.run_train() + self.run_val() + elif self.mode == "val": + self.run_val() + elif self.mode == "train_only": + self.run_train() + + def _setup_dataloaders(self): + self.train_dataset = None + self.val_dataset = None + + if self.mode in ["train", "val"]: + self.val_dataset = instantiate(self.data_conf.get(Phase.VAL, None)) + + if self.mode in ["train", "train_only"]: + self.train_dataset = instantiate(self.data_conf.train) + + def run_train(self): + while self.epoch < self.max_epochs: + dataloader = self.train_dataset.get_loader(epoch=int(self.epoch)) + barrier() + outs = self.train_epoch(dataloader) + self.logger.log_dict(outs, self.epoch) # Logged only on rank 0 + + # log train to text file. + if self.distributed_rank == 0: + with g_pathmgr.open( + os.path.join(self.logging_conf.log_dir, "train_stats.json"), + "a", + ) as f: + f.write(json.dumps(outs) + "\n") + + # Save checkpoint before validating + self.save_checkpoint(self.epoch + 1) + + del dataloader + gc.collect() + + # Run val, not running on last epoch since will run after the + # loop anyway + if self.is_intermediate_val_epoch(self.epoch): + self.run_val() + if torch.cuda.is_available() and self.empty_gpu_mem_cache_after_eval: + # release memory buffers held by the model during eval (which typically + # involves a lot more frames in video grounding that during training) + torch.cuda.empty_cache() + + if self.distributed_rank == 0: + self.best_meter_values.update(self._get_trainer_state("train")) + with g_pathmgr.open( + os.path.join(self.logging_conf.log_dir, "best_stats.json"), + "a", + ) as f: + f.write(json.dumps(self.best_meter_values) + "\n") + + self.epoch += 1 + # epoch was incremented in the loop but the val step runs out of the loop + self.epoch -= 1 + + def run_val(self): + if not self.val_dataset: + return + + dataloader = self.val_dataset.get_loader(epoch=int(self.epoch)) + outs = self.val_epoch(dataloader, phase=Phase.VAL) + del dataloader + gc.collect() + self.logger.log_dict(outs, self.epoch) # Logged only on rank 0 + + if self.distributed_rank == 0: + with g_pathmgr.open( + os.path.join(self.logging_conf.log_dir, "val_stats.json"), + "a", + ) as f: + f.write(json.dumps(outs) + "\n") + + def val_epoch(self, val_loader, phase): + batch_time = AverageMeter("Batch Time", self.device, ":.2f") + data_time = AverageMeter("Data Time", self.device, ":.2f") + mem = MemMeter("Mem (GB)", self.device, ":.2f") + + iters_per_epoch = len(val_loader) + + curr_phases = [phase] + curr_models = [self.model] + + loss_names = [] + for p in curr_phases: + for key in self.loss.keys(): + loss_names.append(f"Losses/{p}_{key}_loss") + + loss_mts = OrderedDict( + [(name, AverageMeter(name, self.device, ":.2e")) for name in loss_names] + ) + extra_loss_mts = {} + + for model in curr_models: + model.eval() + if hasattr(unwrap_ddp_if_wrapped(model), "on_validation_epoch_start"): + unwrap_ddp_if_wrapped(model).on_validation_epoch_start() + + progress = ProgressMeter( + iters_per_epoch, + [batch_time, data_time, mem, self.time_elapsed_meter, *loss_mts.values()], + self._get_meters(curr_phases), + prefix="Val Epoch: [{}]".format(self.epoch), + ) + + end = time.time() + + for data_iter, batch in enumerate(val_loader): + # measure data loading time + data_time.update(time.time() - end) + + # batch = batch.to(self.device, non_blocking=True) + + # compute output + with torch.no_grad(): + with torch.amp.autocast( + device_type="cuda", + enabled=(self.optim_conf.amp.enabled if self.optim_conf else False), + dtype=( + get_amp_type(self.optim_conf.amp.amp_dtype) + if self.optim_conf + else None + ), + ): + for phase, model in zip(curr_phases, curr_models): + loss_dict, batch_size, extra_losses = self._step( + batch, + model, + phase, + ) + + assert len(loss_dict) == 1 + loss_key, loss = loss_dict.popitem() + + if loss_key in loss_mts: + loss_mts[loss_key].update(loss.item(), batch_size) + + for k, v in extra_losses.items(): + if k not in extra_loss_mts: + extra_loss_mts[k] = AverageMeter(k, self.device, ":.2e") + extra_loss_mts[k].update(v.item(), batch_size) + + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() + + self.time_elapsed_meter.update( + time.time() - self.start_time + self.ckpt_time_elapsed + ) + + if torch.cuda.is_available(): + mem.update(reset_peak_usage=True) + + if data_iter % self.logging_conf.log_freq == 0: + progress.display(data_iter) + + if data_iter % self.logging_conf.log_scalar_frequency == 0: + # Log progress meters. + for progress_meter in progress.meters: + self.logger.log( + os.path.join("Step_Stats", phase, progress_meter.name), + progress_meter.val, + self.steps[Phase.VAL], + ) + + if data_iter % 10 == 0: + dist.barrier() + + self.est_epoch_time[phase] = batch_time.avg * iters_per_epoch + self._log_timers(phase) + for model in curr_models: + if hasattr(unwrap_ddp_if_wrapped(model), "on_validation_epoch_end"): + unwrap_ddp_if_wrapped(model).on_validation_epoch_end() + + out_dict = self._log_meters_and_save_best_ckpts(curr_phases) + + for k, v in loss_mts.items(): + out_dict[k] = v.avg + for k, v in extra_loss_mts.items(): + out_dict[k] = v.avg + + for phase in curr_phases: + out_dict.update(self._get_trainer_state(phase)) + self._reset_meters(curr_phases) + logging.info(f"Meters: {out_dict}") + return out_dict + + def _get_trainer_state(self, phase): + return { + "Trainer/where": self.where, + "Trainer/epoch": self.epoch, + f"Trainer/steps_{phase}": self.steps[phase], + } + + def train_epoch(self, train_loader): + # Init stat meters + batch_time_meter = AverageMeter("Batch Time", self.device, ":.2f") + data_time_meter = AverageMeter("Data Time", self.device, ":.2f") + mem_meter = MemMeter("Mem (GB)", self.device, ":.2f") + data_times = [] + phase = Phase.TRAIN + + iters_per_epoch = len(train_loader) + + loss_names = [] + for batch_key in self.loss.keys(): + loss_names.append(f"Losses/{phase}_{batch_key}_loss") + + loss_mts = OrderedDict( + [(name, AverageMeter(name, self.device, ":.2e")) for name in loss_names] + ) + extra_loss_mts = {} + + progress = ProgressMeter( + iters_per_epoch, + [ + batch_time_meter, + data_time_meter, + mem_meter, + self.time_elapsed_meter, + *loss_mts.values(), + ], + self._get_meters([phase]), + prefix="Train Epoch: [{}]".format(self.epoch), + ) + + # Model training loop + self.model.train() + end = time.time() + + for data_iter, batch in enumerate(train_loader): + # measure data loading time + data_time_meter.update(time.time() - end) + data_times.append(data_time_meter.val) + # batch = batch.to( + # self.device, non_blocking=True + # ) # move tensors in a tensorclass + + try: + self._run_step(batch, phase, loss_mts, extra_loss_mts) + + # compute gradient and do optim step + exact_epoch = self.epoch + float(data_iter) / iters_per_epoch + self.where = float(exact_epoch) / self.max_epochs + assert self.where <= 1 + self.EPSILON + if self.where < 1.0: + self.optim.step_schedulers( + self.where, step=int(exact_epoch * iters_per_epoch) + ) + else: + logging.warning( + f"Skipping scheduler update since the training is at the end, i.e, {self.where} of [0,1]." + ) + + # Log schedulers + if data_iter % self.logging_conf.log_scalar_frequency == 0: + for j, param_group in enumerate(self.optim.optimizer.param_groups): + for option in self.optim.schedulers[j]: + optim_prefix = ( + "" + f"{j}_" + if len(self.optim.optimizer.param_groups) > 1 + else "" + ) + self.logger.log( + os.path.join("Optim", f"{optim_prefix}", option), + param_group[option], + self.steps[phase], + ) + + # Clipping gradients and detecting diverging gradients + if self.gradient_clipper is not None: + self.scaler.unscale_(self.optim.optimizer) + self.gradient_clipper(model=self.model) + + if self.gradient_logger is not None: + self.gradient_logger( + self.model, rank=self.distributed_rank, where=self.where + ) + + # Optimizer step: the scaler will make sure gradients are not + # applied if the gradients are infinite + self.scaler.step(self.optim.optimizer) + self.scaler.update() + + # measure elapsed time + batch_time_meter.update(time.time() - end) + end = time.time() + + self.time_elapsed_meter.update( + time.time() - self.start_time + self.ckpt_time_elapsed + ) + + mem_meter.update(reset_peak_usage=True) + if data_iter % self.logging_conf.log_freq == 0: + progress.display(data_iter) + + if data_iter % self.logging_conf.log_scalar_frequency == 0: + # Log progress meters. + for progress_meter in progress.meters: + self.logger.log( + os.path.join("Step_Stats", phase, progress_meter.name), + progress_meter.val, + self.steps[phase], + ) + + # Catching NaN/Inf errors in the loss + except FloatingPointError as e: + raise e + + self.est_epoch_time[Phase.TRAIN] = batch_time_meter.avg * iters_per_epoch + self._log_timers(Phase.TRAIN) + self._log_sync_data_times(Phase.TRAIN, data_times) + + out_dict = self._log_meters_and_save_best_ckpts([Phase.TRAIN]) + + for k, v in loss_mts.items(): + out_dict[k] = v.avg + for k, v in extra_loss_mts.items(): + out_dict[k] = v.avg + out_dict.update(self._get_trainer_state(phase)) + logging.info(f"Losses and meters: {out_dict}") + self._reset_meters([phase]) + return out_dict + + def _log_sync_data_times(self, phase, data_times): + data_times = all_reduce_max(torch.tensor(data_times)).tolist() + steps = range(self.steps[phase] - len(data_times), self.steps[phase]) + for step, data_time in zip(steps, data_times): + if step % self.logging_conf.log_scalar_frequency == 0: + self.logger.log( + os.path.join("Step_Stats", phase, "Data Time Synced"), + data_time, + step, + ) + + def _run_step( + self, + batch: BatchedDatapoint, + phase: str, + loss_mts: Dict[str, AverageMeter], + extra_loss_mts: Dict[str, AverageMeter], + raise_on_error: bool = True, + ): + """ + Run the forward / backward + """ + + # it's important to set grads to None, especially with Adam since 0 + # grads will also update a model even if the step doesn't produce + # gradients + self.optim.zero_grad(set_to_none=True) + + if self.gradient_accumulation_steps > 1: + assert isinstance( + batch, list + ), f"Expected a list of batches, got {type(batch)}" + assert ( + len(batch) == self.gradient_accumulation_steps + ), f"Expected {self.gradient_accumulation_steps} batches, got {len(batch)}" + accum_steps = len(batch) + else: + accum_steps = 1 + batch = [batch] + + for i, chunked_batch in enumerate(batch): + ddp_context = ( + self.model.no_sync() + if i < accum_steps - 1 + else contextlib.nullcontext() + ) + with ddp_context: + with torch.amp.autocast( + device_type="cuda", + enabled=self.optim_conf.amp.enabled, + dtype=get_amp_type(self.optim_conf.amp.amp_dtype), + ): + loss_dict, batch_size, extra_losses = self._step( + chunked_batch, + self.model, + phase, + ) + + assert len(loss_dict) == 1 + loss_key, loss = loss_dict.popitem() + + if not math.isfinite(loss.item()): + error_msg = f"Loss is {loss.item()}, attempting to stop training" + logging.error(error_msg) + if raise_on_error: + raise FloatingPointError(error_msg) + else: + return + + self.scaler.scale(loss).backward() + loss_mts[loss_key].update(loss.item(), batch_size) + for extra_loss_key, extra_loss in extra_losses.items(): + if extra_loss_key not in extra_loss_mts: + extra_loss_mts[extra_loss_key] = AverageMeter( + extra_loss_key, self.device, ":.2e" + ) + extra_loss_mts[extra_loss_key].update(extra_loss.item(), batch_size) + + def _log_meters_and_save_best_ckpts(self, phases: List[str]): + logging.info("Synchronizing meters") + out_dict = {} + checkpoint_save_keys = [] + for key, meter in self._get_meters(phases).items(): + meter_output = meter.compute_synced() + is_better_check = getattr(meter, "is_better", None) + + for meter_subkey, meter_value in meter_output.items(): + out_dict[os.path.join("Meters_train", key, meter_subkey)] = meter_value + + if is_better_check is None: + continue + + tracked_meter_key = os.path.join(key, meter_subkey) + if tracked_meter_key not in self.best_meter_values or is_better_check( + meter_value, + self.best_meter_values[tracked_meter_key], + ): + self.best_meter_values[tracked_meter_key] = meter_value + + if ( + self.checkpoint_conf.save_best_meters is not None + and key in self.checkpoint_conf.save_best_meters + ): + checkpoint_save_keys.append(tracked_meter_key.replace("/", "_")) + + if len(checkpoint_save_keys) > 0: + self.save_checkpoint(self.epoch + 1, checkpoint_save_keys) + + return out_dict + + def _log_timers(self, phase): + time_remaining = 0 + epochs_remaining = self.max_epochs - self.epoch - 1 + val_epochs_remaining = sum( + n % self.val_epoch_freq == 0 for n in range(self.epoch, self.max_epochs) + ) + + # Adding the guaranteed val run at the end if val_epoch_freq doesn't coincide with + # the end epoch. + if (self.max_epochs - 1) % self.val_epoch_freq != 0: + val_epochs_remaining += 1 + + # Remove the current val run from estimate + if phase == Phase.VAL: + val_epochs_remaining -= 1 + + time_remaining += ( + epochs_remaining * self.est_epoch_time[Phase.TRAIN] + + val_epochs_remaining * self.est_epoch_time[Phase.VAL] + ) + + self.logger.log( + os.path.join("Step_Stats", phase, self.time_elapsed_meter.name), + self.time_elapsed_meter.val, + self.steps[phase], + ) + + logging.info(f"Estimated time remaining: {human_readable_time(time_remaining)}") + + def _reset_meters(self, phases: str) -> None: + for meter in self._get_meters(phases).values(): + meter.reset() + + def _check_val_key_match(self, val_keys, phase): + if val_keys is not None: + # Check if there are any duplicates + assert len(val_keys) == len( + set(val_keys) + ), f"Duplicate keys in val datasets, keys: {val_keys}" + + # Check that the keys match the meter keys + if self.meters_conf is not None and phase in self.meters_conf: + assert set(val_keys) == set(self.meters_conf[phase].keys()), ( + f"Keys in val datasets do not match the keys in meters." + f"\nMissing in meters: {set(val_keys) - set(self.meters_conf[phase].keys())}" + f"\nMissing in val datasets: {set(self.meters_conf[phase].keys()) - set(val_keys)}" + ) + + if self.loss_conf is not None: + loss_keys = set(self.loss_conf.keys()) - set(["all"]) + if "default" not in loss_keys: + for k in val_keys: + assert ( + k in loss_keys + ), f"Error: key {k} is not defined in the losses, and no default is set" + + def _setup_components(self): + # Get the keys for all the val datasets, if any + val_phase = Phase.VAL + val_keys = None + if self.data_conf.get(val_phase, None) is not None: + val_keys = collect_dict_keys(self.data_conf[val_phase]) + # Additional checks on the sanity of the config for val datasets + self._check_val_key_match(val_keys, phase=val_phase) + + logging.info("Setting up components: Model, loss, optim, meters etc.") + self.epoch = 0 + self.steps = {Phase.TRAIN: 0, Phase.VAL: 0} + + self.logger = Logger(self.logging_conf) + + self.model = instantiate(self.model_conf, _convert_="all") + print_model_summary(self.model) + + self.loss = None + if self.loss_conf: + self.loss = { + key: el # wrap_base_loss(el) + for (key, el) in instantiate(self.loss_conf, _convert_="all").items() + } + self.loss = nn.ModuleDict(self.loss) + + self.meters = {} + self.best_meter_values = {} + if self.meters_conf: + self.meters = instantiate(self.meters_conf, _convert_="all") + + self.scaler = torch.amp.GradScaler( + self.device, + enabled=self.optim_conf.amp.enabled if self.optim_conf else False, + ) + + self.gradient_clipper = ( + instantiate(self.optim_conf.gradient_clip) if self.optim_conf else None + ) + self.gradient_logger = ( + instantiate(self.optim_conf.gradient_logger) if self.optim_conf else None + ) + + logging.info("Finished setting up components: Model, loss, optim, meters etc.") + + def _construct_optimizers(self): + self.optim = construct_optimizer( + self.model, + self.optim_conf.optimizer, + self.optim_conf.options, + self.optim_conf.param_group_modifiers, + ) + + def _log_loss_detailed_and_return_core_loss(self, loss, loss_str, step): + core_loss = loss.pop(CORE_LOSS_KEY) + if step % self.logging_conf.log_scalar_frequency == 0: + for k in loss: + log_str = os.path.join(loss_str, k) + self.logger.log(log_str, loss[k], step) + return core_loss + + +def print_model_summary(model: torch.nn.Module, log_dir: str = ""): + """ + Prints the model and the number of parameters in the model. + # Multiple packages provide this info in a nice table format + # However, they need us to provide an `input` (as they also write down the output sizes) + # Our models are complex, and a single input is restrictive. + # https://github.com/sksq96/pytorch-summary + # https://github.com/nmhkahn/torchsummaryX + """ + if get_rank() != 0: + return + param_kwargs = {} + trainable_parameters = sum( + p.numel() for p in model.parameters(**param_kwargs) if p.requires_grad + ) + total_parameters = sum(p.numel() for p in model.parameters(**param_kwargs)) + non_trainable_parameters = total_parameters - trainable_parameters + logging.info("==" * 10) + logging.info(f"Summary for model {type(model)}") + logging.info(f"Model is {model}") + logging.info(f"\tTotal parameters {get_human_readable_count(total_parameters)}") + logging.info( + f"\tTrainable parameters {get_human_readable_count(trainable_parameters)}" + ) + logging.info( + f"\tNon-Trainable parameters {get_human_readable_count(non_trainable_parameters)}" + ) + logging.info("==" * 10) + + if log_dir: + output_fpath = os.path.join(log_dir, "model.txt") + with g_pathmgr.open(output_fpath, "w") as f: + print(model, file=f) + + +PARAMETER_NUM_UNITS = [" ", "K", "M", "B", "T"] + + +def get_human_readable_count(number: int) -> str: + """ + Abbreviates an integer number with K, M, B, T for thousands, millions, + billions and trillions, respectively. + Examples: + >>> get_human_readable_count(123) + '123 ' + >>> get_human_readable_count(1234) # (one thousand) + '1.2 K' + >>> get_human_readable_count(2e6) # (two million) + '2.0 M' + >>> get_human_readable_count(3e9) # (three billion) + '3.0 B' + >>> get_human_readable_count(4e14) # (four hundred trillion) + '400 T' + >>> get_human_readable_count(5e15) # (more than trillion) + '5,000 T' + Args: + number: a positive integer number + Return: + A string formatted according to the pattern described above. + """ + assert number >= 0 + labels = PARAMETER_NUM_UNITS + num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1) + num_groups = int(np.ceil(num_digits / 3)) + num_groups = min(num_groups, len(labels)) # don't abbreviate beyond trillions + shift = -3 * (num_groups - 1) + number = number * (10**shift) + index = num_groups - 1 + if index < 1 or number >= 100: + return f"{int(number):,d} {labels[index]}" + else: + return f"{number:,.1f} {labels[index]}" diff --git a/src/nn/segearth_ov3/sam3/train/transforms/__init__.py b/src/nn/segearth_ov3/sam3/train/transforms/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/transforms/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/train/transforms/basic.py b/src/nn/segearth_ov3/sam3/train/transforms/basic.py new file mode 100644 index 0000000..3e8cf1f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/transforms/basic.py @@ -0,0 +1,455 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Transforms and data augmentation for both image + bbox. +""" + +import math +import random +from typing import Iterable + +import PIL +import torch +import torchvision.transforms as T +import torchvision.transforms.functional as F + +from sam3.model.box_ops import box_xyxy_to_cxcywh +from sam3.model.data_misc import interpolate + + +def crop(image, target, region): + cropped_image = F.crop(image, *region) + + target = target.copy() + i, j, h, w = region + + # should we do something wrt the original size? + target["size"] = torch.tensor([h, w]) + + fields = ["labels", "area", "iscrowd", "positive_map"] + + if "boxes" in target: + boxes = target["boxes"] + max_size = torch.as_tensor([w, h], dtype=torch.float32) + cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) + cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) + cropped_boxes = cropped_boxes.clamp(min=0) + area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) + target["boxes"] = cropped_boxes.reshape(-1, 4) + target["area"] = area + fields.append("boxes") + + if "input_boxes" in target: + boxes = target["input_boxes"] + max_size = torch.as_tensor([w, h], dtype=torch.float32) + cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) + cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) + cropped_boxes = cropped_boxes.clamp(min=0) + target["input_boxes"] = cropped_boxes.reshape(-1, 4) + + if "masks" in target: + # FIXME should we update the area here if there are no boxes? + target["masks"] = target["masks"][:, i : i + h, j : j + w] + fields.append("masks") + + # remove elements for which the boxes or masks that have zero area + if "boxes" in target or "masks" in target: + # favor boxes selection when defining which elements to keep + # this is compatible with previous implementation + if "boxes" in target: + cropped_boxes = target["boxes"].reshape(-1, 2, 2) + keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) + else: + keep = target["masks"].flatten(1).any(1) + + for field in fields: + if field in target: + target[field] = target[field][keep] + + return cropped_image, target + + +def hflip(image, target): + flipped_image = F.hflip(image) + + w, h = image.size + + target = target.copy() + if "boxes" in target: + boxes = target["boxes"] + boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( + [-1, 1, -1, 1] + ) + torch.as_tensor([w, 0, w, 0]) + target["boxes"] = boxes + + if "input_boxes" in target: + boxes = target["input_boxes"] + boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( + [-1, 1, -1, 1] + ) + torch.as_tensor([w, 0, w, 0]) + target["input_boxes"] = boxes + + if "masks" in target: + target["masks"] = target["masks"].flip(-1) + + if "text_input" in target: + text_input = ( + target["text_input"] + .replace("left", "[TMP]") + .replace("right", "left") + .replace("[TMP]", "right") + ) + target["text_input"] = text_input + + return flipped_image, target + + +def resize(image, target, size, max_size=None, square=False): + # size can be min_size (scalar) or (w, h) tuple + + def get_size_with_aspect_ratio(image_size, size, max_size=None): + w, h = image_size + if max_size is not None: + min_original_size = float(min((w, h))) + max_original_size = float(max((w, h))) + if max_original_size / min_original_size * size > max_size: + size = int(round(max_size * min_original_size / max_original_size)) + + if (w <= h and w == size) or (h <= w and h == size): + return (h, w) + + if w < h: + ow = size + oh = int(size * h / w) + else: + oh = size + ow = int(size * w / h) + + return (oh, ow) + + def get_size(image_size, size, max_size=None): + if isinstance(size, (list, tuple)): + return size[::-1] + else: + return get_size_with_aspect_ratio(image_size, size, max_size) + + if square: + size = size, size + else: + size = get_size(image.size, size, max_size) + rescaled_image = F.resize(image, size) + + if target is None: + return rescaled_image, None + + ratios = tuple( + float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size) + ) + ratio_width, ratio_height = ratios + + target = target.copy() + if "boxes" in target: + boxes = target["boxes"] + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32 + ) + target["boxes"] = scaled_boxes + if "input_boxes" in target: + boxes = target["input_boxes"] + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32 + ) + target["input_boxes"] = scaled_boxes + + if "area" in target: + area = target["area"] + scaled_area = area * (ratio_width * ratio_height) + target["area"] = scaled_area + + h, w = size + target["size"] = torch.tensor([h, w]) + + if "masks" in target: + target["masks"] = ( + interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0] + > 0.5 + ) + + return rescaled_image, target + + +def pad(image, target, padding): + if len(padding) == 2: + # assumes that we only pad on the bottom right corners + padded_image = F.pad(image, (0, 0, padding[0], padding[1])) + else: + # left, top, right, bottom + padded_image = F.pad(image, (padding[0], padding[1], padding[2], padding[3])) + if target is None: + return padded_image, None + target = target.copy() + + w, h = padded_image.size + + # should we do something wrt the original size? + target["size"] = torch.tensor([h, w]) + if "boxes" in target and len(padding) == 4: + boxes = target["boxes"] + boxes = boxes + torch.as_tensor( + [padding[0], padding[1], padding[0], padding[1]], dtype=torch.float32 + ) + target["boxes"] = boxes + + if "input_boxes" in target and len(padding) == 4: + boxes = target["input_boxes"] + boxes = boxes + torch.as_tensor( + [padding[0], padding[1], padding[0], padding[1]], dtype=torch.float32 + ) + target["input_boxes"] = boxes + + if "masks" in target: + if len(padding) == 2: + target["masks"] = torch.nn.functional.pad( + target["masks"], (0, padding[0], 0, padding[1]) + ) + else: + target["masks"] = torch.nn.functional.pad( + target["masks"], (padding[0], padding[2], padding[1], padding[3]) + ) + return padded_image, target + + +class RandomCrop: + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + region = T.RandomCrop.get_params(img, self.size) + return crop(img, target, region) + + +class RandomSizeCrop: + def __init__(self, min_size: int, max_size: int, respect_boxes: bool = False): + self.min_size = min_size + self.max_size = max_size + self.respect_boxes = respect_boxes # if True we can't crop a box out + + def __call__(self, img: PIL.Image.Image, target: dict): + init_boxes = len(target["boxes"]) + init_boxes_tensor = target["boxes"].clone() + if self.respect_boxes and init_boxes > 0: + minW, minH, maxW, maxH = ( + min(img.width, self.min_size), + min(img.width, self.min_size), + min(img.width, self.max_size), + min(img.height, self.max_size), + ) + minX, minY = ( + target["boxes"][:, 0].max().item() + 10.0, + target["boxes"][:, 1].max().item() + 10.0, + ) + minX = min(img.width, minX) + minY = min(img.height, minY) + maxX, maxY = ( + target["boxes"][:, 2].min().item() - 10, + target["boxes"][:, 3].min().item() - 10, + ) + maxX = max(0.0, maxX) + maxY = max(0.0, maxY) + minW = max(minW, minX - maxX) + minH = max(minH, minY - maxY) + w = random.uniform(minW, max(minW, maxW)) + h = random.uniform(minH, max(minH, maxH)) + if minX > maxX: + # i = random.uniform(max(0, minX - w + 1), max(maxX, max(0, minX - w + 1))) + i = random.uniform(max(0, minX - w), max(maxX, max(0, minX - w))) + else: + i = random.uniform( + max(0, minX - w + 1), max(maxX - 1, max(0, minX - w + 1)) + ) + if minY > maxY: + # j = random.uniform(max(0, minY - h + 1), max(maxY, max(0, minY - h + 1))) + j = random.uniform(max(0, minY - h), max(maxY, max(0, minY - h))) + else: + j = random.uniform( + max(0, minY - h + 1), max(maxY - 1, max(0, minY - h + 1)) + ) + result_img, result_target = crop(img, target, [j, i, h, w]) + assert ( + len(result_target["boxes"]) == init_boxes + ), f"img_w={img.width}\timg_h={img.height}\tminX={minX}\tminY={minY}\tmaxX={maxX}\tmaxY={maxY}\tminW={minW}\tminH={minH}\tmaxW={maxW}\tmaxH={maxH}\tw={w}\th={h}\ti={i}\tj={j}\tinit_boxes={init_boxes_tensor}\tresults={result_target['boxes']}" + + return result_img, result_target + else: + w = random.randint(self.min_size, min(img.width, self.max_size)) + h = random.randint(self.min_size, min(img.height, self.max_size)) + region = T.RandomCrop.get_params(img, (h, w)) + result_img, result_target = crop(img, target, region) + return result_img, result_target + + +class CenterCrop: + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + image_width, image_height = img.size + crop_height, crop_width = self.size + crop_top = int(round((image_height - crop_height) / 2.0)) + crop_left = int(round((image_width - crop_width) / 2.0)) + return crop(img, target, (crop_top, crop_left, crop_height, crop_width)) + + +class RandomHorizontalFlip: + def __init__(self, p=0.5): + self.p = p + + def __call__(self, img, target): + if random.random() < self.p: + return hflip(img, target) + return img, target + + +class RandomResize: + def __init__(self, sizes, max_size=None, square=False): + if isinstance(sizes, int): + sizes = (sizes,) + assert isinstance(sizes, Iterable) + self.sizes = list(sizes) + self.max_size = max_size + self.square = square + + def __call__(self, img, target=None): + size = random.choice(self.sizes) + return resize(img, target, size, self.max_size, square=self.square) + + +class RandomPad: + def __init__(self, max_pad): + self.max_pad = max_pad + + def __call__(self, img, target): + pad_x = random.randint(0, self.max_pad) + pad_y = random.randint(0, self.max_pad) + return pad(img, target, (pad_x, pad_y)) + + +class PadToSize: + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + w, h = img.size + pad_x = self.size - w + pad_y = self.size - h + assert pad_x >= 0 and pad_y >= 0 + pad_left = random.randint(0, pad_x) + pad_right = pad_x - pad_left + pad_top = random.randint(0, pad_y) + pad_bottom = pad_y - pad_top + return pad(img, target, (pad_left, pad_top, pad_right, pad_bottom)) + + +class Identity: + def __call__(self, img, target): + return img, target + + +class RandomSelect: + """ + Randomly selects between transforms1 and transforms2, + with probability p for transforms1 and (1 - p) for transforms2 + """ + + def __init__(self, transforms1=None, transforms2=None, p=0.5): + self.transforms1 = transforms1 or Identity() + self.transforms2 = transforms2 or Identity() + self.p = p + + def __call__(self, img, target): + if random.random() < self.p: + return self.transforms1(img, target) + return self.transforms2(img, target) + + +class ToTensor: + def __call__(self, img, target): + return F.to_tensor(img), target + + +class RandomErasing: + def __init__(self, *args, **kwargs): + self.eraser = T.RandomErasing(*args, **kwargs) + + def __call__(self, img, target): + return self.eraser(img), target + + +class Normalize: + def __init__(self, mean, std): + self.mean = mean + self.std = std + + def __call__(self, image, target=None): + image = F.normalize(image, mean=self.mean, std=self.std) + if target is None: + return image, None + target = target.copy() + h, w = image.shape[-2:] + if "boxes" in target: + boxes = target["boxes"] + boxes = box_xyxy_to_cxcywh(boxes) + boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) + target["boxes"] = boxes + if "input_boxes" in target: + boxes = target["input_boxes"] + boxes = box_xyxy_to_cxcywh(boxes) + boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) + target["input_boxes"] = boxes + return image, target + + +class RemoveDifficult: + def __init__(self, enabled=False): + self.remove_difficult = enabled + + def __call__(self, image, target=None): + if target is None: + return image, None + target = target.copy() + keep = ~target["iscrowd"].to(torch.bool) | (not self.remove_difficult) + if "boxes" in target: + target["boxes"] = target["boxes"][keep] + target["labels"] = target["labels"][keep] + target["iscrowd"] = target["iscrowd"][keep] + return image, target + + +class Compose: + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, image, target): + for t in self.transforms: + image, target = t(image, target) + return image, target + + def __repr__(self): + format_string = self.__class__.__name__ + "(" + for t in self.transforms: + format_string += "\n" + format_string += " {0}".format(t) + format_string += "\n)" + return format_string + + +def get_random_resize_scales(size, min_size, rounded): + stride = 128 if rounded else 32 + min_size = int(stride * math.ceil(min_size / stride)) + scales = list(range(min_size, size + 1, stride)) + return scales + + +def get_random_resize_max_size(size, ratio=5 / 3): + max_size = round(ratio * size) + return max_size diff --git a/src/nn/segearth_ov3/sam3/train/transforms/basic_for_api.py b/src/nn/segearth_ov3/sam3/train/transforms/basic_for_api.py new file mode 100644 index 0000000..e0ec2af --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/transforms/basic_for_api.py @@ -0,0 +1,1396 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Transforms and data augmentation for both image + bbox. +""" + +import logging + +import numbers +import random +from collections.abc import Sequence +from typing import Iterable + +import torch +import torchvision.transforms as T +import torchvision.transforms.functional as F +import torchvision.transforms.v2.functional as Fv2 + +from PIL import Image as PILImage + +from sam3.model.box_ops import box_xyxy_to_cxcywh, masks_to_boxes +from sam3.train.data.sam3_image_dataset import Datapoint +from torchvision.transforms import InterpolationMode + + +def crop( + datapoint, + index, + region, + v2=False, + check_validity=True, + check_input_validity=True, + recompute_box_from_mask=False, +): + if v2: + rtop, rleft, rheight, rwidth = (int(round(r)) for r in region) + datapoint.images[index].data = Fv2.crop( + datapoint.images[index].data, + top=rtop, + left=rleft, + height=rheight, + width=rwidth, + ) + else: + datapoint.images[index].data = F.crop(datapoint.images[index].data, *region) + + i, j, h, w = region + + # should we do something wrt the original size? + datapoint.images[index].size = (h, w) + + for obj in datapoint.images[index].objects: + # crop the mask + if obj.segment is not None: + obj.segment = F.crop(obj.segment, int(i), int(j), int(h), int(w)) + + # crop the bounding box + if recompute_box_from_mask and obj.segment is not None: + # here the boxes are still in XYXY format with absolute coordinates (they are + # converted to CxCyWH with relative coordinates in basic_for_api.NormalizeAPI) + obj.bbox, obj.area = get_bbox_xyxy_abs_coords_from_mask(obj.segment) + else: + if recompute_box_from_mask and obj.segment is None and obj.area > 0: + logging.warning( + "Cannot recompute bounding box from mask since `obj.segment` is None. " + "Falling back to directly cropping from the input bounding box." + ) + boxes = obj.bbox.view(1, 4) + max_size = torch.as_tensor([w, h], dtype=torch.float32) + cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) + cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) + cropped_boxes = cropped_boxes.clamp(min=0) + obj.area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) + obj.bbox = cropped_boxes.reshape(-1, 4) + + for query in datapoint.find_queries: + if query.semantic_target is not None: + query.semantic_target = F.crop( + query.semantic_target, int(i), int(j), int(h), int(w) + ) + if query.image_id == index and query.input_bbox is not None: + boxes = query.input_bbox + max_size = torch.as_tensor([w, h], dtype=torch.float32) + cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) + cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) + cropped_boxes = cropped_boxes.clamp(min=0) + + # cur_area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) + # if check_input_validity: + # assert ( + # (cur_area > 0).all().item() + # ), "Some input box got cropped out by the crop transform" + + query.input_bbox = cropped_boxes.reshape(-1, 4) + if query.image_id == index and query.input_points is not None: + print( + "Warning! Point cropping with this function may lead to unexpected results" + ) + points = query.input_points + # Unlike right-lower box edges, which are exclusive, the + # point must be in [0, length-1], hence the -1 + max_size = torch.as_tensor([w, h], dtype=torch.float32) - 1 + cropped_points = points - torch.as_tensor([j, i, 0], dtype=torch.float32) + cropped_points[:, :, :2] = torch.min(cropped_points[:, :, :2], max_size) + cropped_points[:, :, :2] = cropped_points[:, :, :2].clamp(min=0) + query.input_points = cropped_points + + if check_validity: + # Check that all boxes are still valid + for obj in datapoint.images[index].objects: + assert obj.area > 0, "Box {} has no area".format(obj.bbox) + + return datapoint + + +def hflip(datapoint, index): + datapoint.images[index].data = F.hflip(datapoint.images[index].data) + + w, h = datapoint.images[index].data.size + for obj in datapoint.images[index].objects: + boxes = obj.bbox.view(1, 4) + boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( + [-1, 1, -1, 1] + ) + torch.as_tensor([w, 0, w, 0]) + obj.bbox = boxes + if obj.segment is not None: + obj.segment = F.hflip(obj.segment) + + for query in datapoint.find_queries: + if query.semantic_target is not None: + query.semantic_target = F.hflip(query.semantic_target) + if query.image_id == index and query.input_bbox is not None: + boxes = query.input_bbox + boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( + [-1, 1, -1, 1] + ) + torch.as_tensor([w, 0, w, 0]) + query.input_bbox = boxes + if query.image_id == index and query.input_points is not None: + points = query.input_points + points = points * torch.as_tensor([-1, 1, 1]) + torch.as_tensor([w, 0, 0]) + query.input_points = points + return datapoint + + +def get_size_with_aspect_ratio(image_size, size, max_size=None): + w, h = image_size + if max_size is not None: + min_original_size = float(min((w, h))) + max_original_size = float(max((w, h))) + if max_original_size / min_original_size * size > max_size: + size = max_size * min_original_size / max_original_size + + if (w <= h and w == size) or (h <= w and h == size): + return (h, w) + + if w < h: + ow = int(round(size)) + oh = int(round(size * h / w)) + else: + oh = int(round(size)) + ow = int(round(size * w / h)) + + return (oh, ow) + + +def resize(datapoint, index, size, max_size=None, square=False, v2=False): + # size can be min_size (scalar) or (w, h) tuple + + def get_size(image_size, size, max_size=None): + if isinstance(size, (list, tuple)): + return size[::-1] + else: + return get_size_with_aspect_ratio(image_size, size, max_size) + + if square: + size = size, size + else: + cur_size = ( + datapoint.images[index].data.size()[-2:][::-1] + if v2 + else datapoint.images[index].data.size + ) + size = get_size(cur_size, size, max_size) + + old_size = ( + datapoint.images[index].data.size()[-2:][::-1] + if v2 + else datapoint.images[index].data.size + ) + if v2: + datapoint.images[index].data = Fv2.resize( + datapoint.images[index].data, size, antialias=True + ) + else: + datapoint.images[index].data = F.resize(datapoint.images[index].data, size) + + new_size = ( + datapoint.images[index].data.size()[-2:][::-1] + if v2 + else datapoint.images[index].data.size + ) + ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(new_size, old_size)) + ratio_width, ratio_height = ratios + + for obj in datapoint.images[index].objects: + boxes = obj.bbox.view(1, 4) + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32 + ) + obj.bbox = scaled_boxes + obj.area *= ratio_width * ratio_height + if obj.segment is not None: + obj.segment = F.resize(obj.segment[None, None], size).squeeze() + + for query in datapoint.find_queries: + if query.semantic_target is not None: + query.semantic_target = F.resize( + query.semantic_target[None, None], size + ).squeeze() + if query.image_id == index and query.input_bbox is not None: + boxes = query.input_bbox + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height], + dtype=torch.float32, + ) + query.input_bbox = scaled_boxes + if query.image_id == index and query.input_points is not None: + points = query.input_points + scaled_points = points * torch.as_tensor( + [ratio_width, ratio_height, 1], + dtype=torch.float32, + ) + query.input_points = scaled_points + + h, w = size + datapoint.images[index].size = (h, w) + return datapoint + + +def pad(datapoint, index, padding, v2=False): + old_h, old_w = datapoint.images[index].size + h, w = old_h, old_w + if len(padding) == 2: + # assumes that we only pad on the bottom right corners + if v2: + datapoint.images[index].data = Fv2.pad( + datapoint.images[index].data, (0, 0, padding[0], padding[1]) + ) + else: + datapoint.images[index].data = F.pad( + datapoint.images[index].data, (0, 0, padding[0], padding[1]) + ) + h += padding[1] + w += padding[0] + else: + if v2: + # left, top, right, bottom + datapoint.images[index].data = Fv2.pad( + datapoint.images[index].data, + (padding[0], padding[1], padding[2], padding[3]), + ) + else: + # left, top, right, bottom + datapoint.images[index].data = F.pad( + datapoint.images[index].data, + (padding[0], padding[1], padding[2], padding[3]), + ) + h += padding[1] + padding[3] + w += padding[0] + padding[2] + + datapoint.images[index].size = (h, w) + + for obj in datapoint.images[index].objects: + if len(padding) != 2: + obj.bbox += torch.as_tensor( + [padding[0], padding[1], padding[0], padding[1]], dtype=torch.float32 + ) + if obj.segment is not None: + if v2: + if len(padding) == 2: + obj.segment = Fv2.pad( + obj.segment[None], (0, 0, padding[0], padding[1]) + ).squeeze(0) + else: + obj.segment = Fv2.pad(obj.segment[None], tuple(padding)).squeeze(0) + else: + if len(padding) == 2: + obj.segment = F.pad(obj.segment, (0, 0, padding[0], padding[1])) + else: + obj.segment = F.pad(obj.segment, tuple(padding)) + + for query in datapoint.find_queries: + if query.semantic_target is not None: + if v2: + if len(padding) == 2: + query.semantic_target = Fv2.pad( + query.semantic_target[None, None], + (0, 0, padding[0], padding[1]), + ).squeeze() + else: + query.semantic_target = Fv2.pad( + query.semantic_target[None, None], tuple(padding) + ).squeeze() + else: + if len(padding) == 2: + query.semantic_target = F.pad( + query.semantic_target[None, None], + (0, 0, padding[0], padding[1]), + ).squeeze() + else: + query.semantic_target = F.pad( + query.semantic_target[None, None], tuple(padding) + ).squeeze() + if query.image_id == index and query.input_bbox is not None: + if len(padding) != 2: + query.input_bbox += torch.as_tensor( + [padding[0], padding[1], padding[0], padding[1]], + dtype=torch.float32, + ) + if query.image_id == index and query.input_points is not None: + if len(padding) != 2: + query.input_points += torch.as_tensor( + [padding[0], padding[1], 0], dtype=torch.float32 + ) + + return datapoint + + +class RandomSizeCropAPI: + def __init__( + self, + min_size: int, + max_size: int, + respect_boxes: bool, + consistent_transform: bool, + respect_input_boxes: bool = True, + v2: bool = False, + recompute_box_from_mask: bool = False, + ): + self.min_size = min_size + self.max_size = max_size + self.respect_boxes = respect_boxes # if True we can't crop a box out + self.respect_input_boxes = respect_input_boxes + self.consistent_transform = consistent_transform + self.v2 = v2 + self.recompute_box_from_mask = recompute_box_from_mask + + def _sample_no_respect_boxes(self, img): + w = random.randint(self.min_size, min(img.width, self.max_size)) + h = random.randint(self.min_size, min(img.height, self.max_size)) + return T.RandomCrop.get_params(img, (h, w)) + + def _sample_respect_boxes(self, img, boxes, points, min_box_size=10.0): + """ + Assure that no box or point is dropped via cropping, though portions + of boxes may be removed. + """ + if len(boxes) == 0 and len(points) == 0: + return self._sample_no_respect_boxes(img) + + if self.v2: + img_height, img_width = img.size()[-2:] + else: + img_width, img_height = img.size + + minW, minH, maxW, maxH = ( + min(img_width, self.min_size), + min(img_height, self.min_size), + min(img_width, self.max_size), + min(img_height, self.max_size), + ) + + # The crop box must extend one pixel beyond points to the bottom/right + # to assure the exclusive box contains the points. + minX = ( + torch.cat([boxes[:, 0] + min_box_size, points[:, 0] + 1], dim=0) + .max() + .item() + ) + minY = ( + torch.cat([boxes[:, 1] + min_box_size, points[:, 1] + 1], dim=0) + .max() + .item() + ) + minX = min(img_width, minX) + minY = min(img_height, minY) + maxX = torch.cat([boxes[:, 2] - min_box_size, points[:, 0]], dim=0).min().item() + maxY = torch.cat([boxes[:, 3] - min_box_size, points[:, 1]], dim=0).min().item() + maxX = max(0.0, maxX) + maxY = max(0.0, maxY) + minW = max(minW, minX - maxX) + minH = max(minH, minY - maxY) + w = random.uniform(minW, max(minW, maxW)) + h = random.uniform(minH, max(minH, maxH)) + if minX > maxX: + # i = random.uniform(max(0, minX - w + 1), max(maxX, max(0, minX - w + 1))) + i = random.uniform(max(0, minX - w), max(maxX, max(0, minX - w))) + else: + i = random.uniform( + max(0, minX - w + 1), max(maxX - 1, max(0, minX - w + 1)) + ) + if minY > maxY: + # j = random.uniform(max(0, minY - h + 1), max(maxY, max(0, minY - h + 1))) + j = random.uniform(max(0, minY - h), max(maxY, max(0, minY - h))) + else: + j = random.uniform( + max(0, minY - h + 1), max(maxY - 1, max(0, minY - h + 1)) + ) + + return [j, i, h, w] + + def __call__(self, datapoint, **kwargs): + if self.respect_boxes or self.respect_input_boxes: + if self.consistent_transform: + # Check that all the images are the same size + w, h = datapoint.images[0].data.size + for img in datapoint.images: + assert img.data.size == (w, h) + + all_boxes = [] + # Getting all boxes in all the images + if self.respect_boxes: + all_boxes += [ + obj.bbox.view(-1, 4) + for img in datapoint.images + for obj in img.objects + ] + # Get all the boxes in the find queries + if self.respect_input_boxes: + all_boxes += [ + q.input_bbox.view(-1, 4) + for q in datapoint.find_queries + if q.input_bbox is not None + ] + if all_boxes: + all_boxes = torch.cat(all_boxes, 0) + else: + all_boxes = torch.empty(0, 4) + + all_points = [ + q.input_points.view(-1, 3)[:, :2] + for q in datapoint.find_queries + if q.input_points is not None + ] + if all_points: + all_points = torch.cat(all_points, 0) + else: + all_points = torch.empty(0, 2) + + crop_param = self._sample_respect_boxes( + datapoint.images[0].data, all_boxes, all_points + ) + for i in range(len(datapoint.images)): + datapoint = crop( + datapoint, + i, + crop_param, + v2=self.v2, + check_validity=self.respect_boxes, + check_input_validity=self.respect_input_boxes, + recompute_box_from_mask=self.recompute_box_from_mask, + ) + return datapoint + else: + for i in range(len(datapoint.images)): + all_boxes = [] + # Get all boxes in the current image + if self.respect_boxes: + all_boxes += [ + obj.bbox.view(-1, 4) for obj in datapoint.images[i].objects + ] + # Get all the boxes in the find queries that correspond to this image + if self.respect_input_boxes: + all_boxes += [ + q.input_bbox.view(-1, 4) + for q in datapoint.find_queries + if q.image_id == i and q.input_bbox is not None + ] + if all_boxes: + all_boxes = torch.cat(all_boxes, 0) + else: + all_boxes = torch.empty(0, 4) + + all_points = [ + q.input_points.view(-1, 3)[:, :2] + for q in datapoint.find_queries + if q.input_points is not None + ] + if all_points: + all_points = torch.cat(all_points, 0) + else: + all_points = torch.empty(0, 2) + + crop_param = self._sample_respect_boxes( + datapoint.images[i].data, all_boxes, all_points + ) + datapoint = crop( + datapoint, + i, + crop_param, + v2=self.v2, + check_validity=self.respect_boxes, + check_input_validity=self.respect_input_boxes, + recompute_box_from_mask=self.recompute_box_from_mask, + ) + return datapoint + else: + if self.consistent_transform: + # Check that all the images are the same size + w, h = datapoint.images[0].data.size + for img in datapoint.images: + assert img.data.size == (w, h) + + crop_param = self._sample_no_respect_boxes(datapoint.images[0].data) + for i in range(len(datapoint.images)): + datapoint = crop( + datapoint, + i, + crop_param, + v2=self.v2, + check_validity=self.respect_boxes, + check_input_validity=self.respect_input_boxes, + recompute_box_from_mask=self.recompute_box_from_mask, + ) + return datapoint + else: + for i in range(len(datapoint.images)): + crop_param = self._sample_no_respect_boxes(datapoint.images[i].data) + datapoint = crop( + datapoint, + i, + crop_param, + v2=self.v2, + check_validity=self.respect_boxes, + check_input_validity=self.respect_input_boxes, + recompute_box_from_mask=self.recompute_box_from_mask, + ) + return datapoint + + +class CenterCropAPI: + def __init__(self, size, consistent_transform, recompute_box_from_mask=False): + self.size = size + self.consistent_transform = consistent_transform + self.recompute_box_from_mask = recompute_box_from_mask + + def _sample_crop(self, image_width, image_height): + crop_height, crop_width = self.size + crop_top = int(round((image_height - crop_height) / 2.0)) + crop_left = int(round((image_width - crop_width) / 2.0)) + return crop_top, crop_left, crop_height, crop_width + + def __call__(self, datapoint, **kwargs): + if self.consistent_transform: + # Check that all the images are the same size + w, h = datapoint.images[0].data.size + for img in datapoint.images: + assert img.size == (w, h) + + crop_top, crop_left, crop_height, crop_width = self._sample_crop(w, h) + for i in range(len(datapoint.images)): + datapoint = crop( + datapoint, + i, + (crop_top, crop_left, crop_height, crop_width), + recompute_box_from_mask=self.recompute_box_from_mask, + ) + return datapoint + + for i in range(len(datapoint.images)): + w, h = datapoint.images[i].data.size + crop_top, crop_left, crop_height, crop_width = self._sample_crop(w, h) + datapoint = crop( + datapoint, + i, + (crop_top, crop_left, crop_height, crop_width), + recompute_box_from_mask=self.recompute_box_from_mask, + ) + + return datapoint + + +class RandomHorizontalFlip: + def __init__(self, consistent_transform, p=0.5): + self.p = p + self.consistent_transform = consistent_transform + + def __call__(self, datapoint, **kwargs): + if self.consistent_transform: + if random.random() < self.p: + for i in range(len(datapoint.images)): + datapoint = hflip(datapoint, i) + return datapoint + for i in range(len(datapoint.images)): + if random.random() < self.p: + datapoint = hflip(datapoint, i) + return datapoint + + +class RandomResizeAPI: + def __init__( + self, sizes, consistent_transform, max_size=None, square=False, v2=False + ): + if isinstance(sizes, int): + sizes = (sizes,) + assert isinstance(sizes, Iterable) + self.sizes = list(sizes) + self.max_size = max_size + self.square = square + self.consistent_transform = consistent_transform + self.v2 = v2 + + def __call__(self, datapoint, **kwargs): + if self.consistent_transform: + size = random.choice(self.sizes) + for i in range(len(datapoint.images)): + datapoint = resize( + datapoint, i, size, self.max_size, square=self.square, v2=self.v2 + ) + return datapoint + for i in range(len(datapoint.images)): + size = random.choice(self.sizes) + datapoint = resize( + datapoint, i, size, self.max_size, square=self.square, v2=self.v2 + ) + return datapoint + + +class ScheduledRandomResizeAPI(RandomResizeAPI): + def __init__(self, size_scheduler, consistent_transform, square=False): + self.size_scheduler = size_scheduler + # Just a meaningful init value for super + params = self.size_scheduler(epoch_num=0) + sizes, max_size = params["sizes"], params["max_size"] + super().__init__(sizes, consistent_transform, max_size=max_size, square=square) + + def __call__(self, datapoint, **kwargs): + assert "epoch" in kwargs, "Param scheduler needs to know the current epoch" + params = self.size_scheduler(kwargs["epoch"]) + sizes, max_size = params["sizes"], params["max_size"] + self.sizes = sizes + self.max_size = max_size + datapoint = super(ScheduledRandomResizeAPI, self).__call__(datapoint, **kwargs) + return datapoint + + +class RandomPadAPI: + def __init__(self, max_pad, consistent_transform): + self.max_pad = max_pad + self.consistent_transform = consistent_transform + + def _sample_pad(self): + pad_x = random.randint(0, self.max_pad) + pad_y = random.randint(0, self.max_pad) + return pad_x, pad_y + + def __call__(self, datapoint, **kwargs): + if self.consistent_transform: + pad_x, pad_y = self._sample_pad() + for i in range(len(datapoint.images)): + datapoint = pad(datapoint, i, (pad_x, pad_y)) + return datapoint + + for i in range(len(datapoint.images)): + pad_x, pad_y = self._sample_pad() + datapoint = pad(datapoint, i, (pad_x, pad_y)) + return datapoint + + +class PadToSizeAPI: + def __init__(self, size, consistent_transform, bottom_right=False, v2=False): + self.size = size + self.consistent_transform = consistent_transform + self.v2 = v2 + self.bottom_right = bottom_right + + def _sample_pad(self, w, h): + pad_x = self.size - w + pad_y = self.size - h + assert pad_x >= 0 and pad_y >= 0 + pad_left = random.randint(0, pad_x) + pad_right = pad_x - pad_left + pad_top = random.randint(0, pad_y) + pad_bottom = pad_y - pad_top + return pad_left, pad_top, pad_right, pad_bottom + + def __call__(self, datapoint, **kwargs): + if self.consistent_transform: + # Check that all the images are the same size + w, h = datapoint.images[0].data.size + for img in datapoint.images: + assert img.size == (w, h) + if self.bottom_right: + pad_right = self.size - w + pad_bottom = self.size - h + padding = (pad_right, pad_bottom) + else: + padding = self._sample_pad(w, h) + for i in range(len(datapoint.images)): + datapoint = pad(datapoint, i, padding, v2=self.v2) + return datapoint + + for i, img in enumerate(datapoint.images): + w, h = img.data.size + if self.bottom_right: + pad_right = self.size - w + pad_bottom = self.size - h + padding = (pad_right, pad_bottom) + else: + padding = self._sample_pad(w, h) + datapoint = pad(datapoint, i, padding, v2=self.v2) + return datapoint + + +class RandomMosaicVideoAPI: + def __init__(self, prob=0.15, grid_h=2, grid_w=2, use_random_hflip=False): + self.prob = prob + self.grid_h = grid_h + self.grid_w = grid_w + self.use_random_hflip = use_random_hflip + + def __call__(self, datapoint, **kwargs): + if random.random() > self.prob: + return datapoint + + # select a random location to place the target mask in the mosaic + target_grid_y = random.randint(0, self.grid_h - 1) + target_grid_x = random.randint(0, self.grid_w - 1) + # whether to flip each grid in the mosaic horizontally + if self.use_random_hflip: + should_hflip = torch.rand(self.grid_h, self.grid_w) < 0.5 + else: + should_hflip = torch.zeros(self.grid_h, self.grid_w, dtype=torch.bool) + for i in range(len(datapoint.images)): + datapoint = random_mosaic_frame( + datapoint, + i, + grid_h=self.grid_h, + grid_w=self.grid_w, + target_grid_y=target_grid_y, + target_grid_x=target_grid_x, + should_hflip=should_hflip, + ) + + return datapoint + + +def random_mosaic_frame( + datapoint, + index, + grid_h, + grid_w, + target_grid_y, + target_grid_x, + should_hflip, +): + # Step 1: downsize the images and paste them into a mosaic + image_data = datapoint.images[index].data + is_pil = isinstance(image_data, PILImage.Image) + if is_pil: + H_im = image_data.height + W_im = image_data.width + image_data_output = PILImage.new("RGB", (W_im, H_im)) + else: + H_im = image_data.size(-2) + W_im = image_data.size(-1) + image_data_output = torch.zeros_like(image_data) + + downsize_cache = {} + for grid_y in range(grid_h): + for grid_x in range(grid_w): + y_offset_b = grid_y * H_im // grid_h + x_offset_b = grid_x * W_im // grid_w + y_offset_e = (grid_y + 1) * H_im // grid_h + x_offset_e = (grid_x + 1) * W_im // grid_w + H_im_downsize = y_offset_e - y_offset_b + W_im_downsize = x_offset_e - x_offset_b + + if (H_im_downsize, W_im_downsize) in downsize_cache: + image_data_downsize = downsize_cache[(H_im_downsize, W_im_downsize)] + else: + image_data_downsize = F.resize( + image_data, + size=(H_im_downsize, W_im_downsize), + interpolation=InterpolationMode.BILINEAR, + antialias=True, # antialiasing for downsizing + ) + downsize_cache[(H_im_downsize, W_im_downsize)] = image_data_downsize + if should_hflip[grid_y, grid_x].item(): + image_data_downsize = F.hflip(image_data_downsize) + + if is_pil: + image_data_output.paste(image_data_downsize, (x_offset_b, y_offset_b)) + else: + image_data_output[:, y_offset_b:y_offset_e, x_offset_b:x_offset_e] = ( + image_data_downsize + ) + + datapoint.images[index].data = image_data_output + + # Step 2: downsize the masks and paste them into the target grid of the mosaic + # (note that we don't scale input/target boxes since they are not used in TA) + for obj in datapoint.images[index].objects: + if obj.segment is None: + continue + assert obj.segment.shape == (H_im, W_im) and obj.segment.dtype == torch.uint8 + segment_output = torch.zeros_like(obj.segment) + + target_y_offset_b = target_grid_y * H_im // grid_h + target_x_offset_b = target_grid_x * W_im // grid_w + target_y_offset_e = (target_grid_y + 1) * H_im // grid_h + target_x_offset_e = (target_grid_x + 1) * W_im // grid_w + target_H_im_downsize = target_y_offset_e - target_y_offset_b + target_W_im_downsize = target_x_offset_e - target_x_offset_b + + segment_downsize = F.resize( + obj.segment[None, None], + size=(target_H_im_downsize, target_W_im_downsize), + interpolation=InterpolationMode.BILINEAR, + antialias=True, # antialiasing for downsizing + )[0, 0] + if should_hflip[target_grid_y, target_grid_x].item(): + segment_downsize = F.hflip(segment_downsize[None, None])[0, 0] + + segment_output[ + target_y_offset_b:target_y_offset_e, target_x_offset_b:target_x_offset_e + ] = segment_downsize + obj.segment = segment_output + + return datapoint + + +class ScheduledPadToSizeAPI(PadToSizeAPI): + def __init__(self, size_scheduler, consistent_transform): + self.size_scheduler = size_scheduler + size = self.size_scheduler(epoch_num=0)["sizes"] + super().__init__(size, consistent_transform) + + def __call__(self, datapoint, **kwargs): + assert "epoch" in kwargs, "Param scheduler needs to know the current epoch" + params = self.size_scheduler(kwargs["epoch"]) + self.size = params["resolution"] + return super(ScheduledPadToSizeAPI, self).__call__(datapoint, **kwargs) + + +class IdentityAPI: + def __call__(self, datapoint, **kwargs): + return datapoint + + +class RandomSelectAPI: + """ + Randomly selects between transforms1 and transforms2, + with probability p for transforms1 and (1 - p) for transforms2 + """ + + def __init__(self, transforms1=None, transforms2=None, p=0.5): + self.transforms1 = transforms1 or IdentityAPI() + self.transforms2 = transforms2 or IdentityAPI() + self.p = p + + def __call__(self, datapoint, **kwargs): + if random.random() < self.p: + return self.transforms1(datapoint, **kwargs) + return self.transforms2(datapoint, **kwargs) + + +class ToTensorAPI: + def __init__(self, v2=False): + self.v2 = v2 + + def __call__(self, datapoint: Datapoint, **kwargs): + for img in datapoint.images: + if self.v2: + img.data = Fv2.to_image_tensor(img.data) + # img.data = Fv2.to_dtype(img.data, torch.uint8, scale=True) + # img.data = Fv2.convert_image_dtype(img.data, torch.uint8) + else: + img.data = F.to_tensor(img.data) + return datapoint + + +class NormalizeAPI: + def __init__(self, mean, std, v2=False): + self.mean = mean + self.std = std + self.v2 = v2 + + def __call__(self, datapoint: Datapoint, **kwargs): + for img in datapoint.images: + if self.v2: + img.data = Fv2.convert_image_dtype(img.data, torch.float32) + img.data = Fv2.normalize(img.data, mean=self.mean, std=self.std) + else: + img.data = F.normalize(img.data, mean=self.mean, std=self.std) + for obj in img.objects: + boxes = obj.bbox + cur_h, cur_w = img.data.shape[-2:] + boxes = box_xyxy_to_cxcywh(boxes) + boxes = boxes / torch.tensor( + [cur_w, cur_h, cur_w, cur_h], dtype=torch.float32 + ) + obj.bbox = boxes + + for query in datapoint.find_queries: + if query.input_bbox is not None: + boxes = query.input_bbox + cur_h, cur_w = datapoint.images[query.image_id].data.shape[-2:] + boxes = box_xyxy_to_cxcywh(boxes) + boxes = boxes / torch.tensor( + [cur_w, cur_h, cur_w, cur_h], dtype=torch.float32 + ) + query.input_bbox = boxes + if query.input_points is not None: + points = query.input_points + cur_h, cur_w = datapoint.images[query.image_id].data.shape[-2:] + points = points / torch.tensor([cur_w, cur_h, 1.0], dtype=torch.float32) + query.input_points = points + + return datapoint + + +class ComposeAPI: + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, datapoint, **kwargs): + for t in self.transforms: + datapoint = t(datapoint, **kwargs) + return datapoint + + def __repr__(self): + format_string = self.__class__.__name__ + "(" + for t in self.transforms: + format_string += "\n" + format_string += " {0}".format(t) + format_string += "\n)" + return format_string + + +class RandomGrayscale: + def __init__(self, consistent_transform, p=0.5): + self.p = p + self.consistent_transform = consistent_transform + self.Grayscale = T.Grayscale(num_output_channels=3) + + def __call__(self, datapoint: Datapoint, **kwargs): + if self.consistent_transform: + if random.random() < self.p: + for img in datapoint.images: + img.data = self.Grayscale(img.data) + return datapoint + for img in datapoint.images: + if random.random() < self.p: + img.data = self.Grayscale(img.data) + return datapoint + + +class ColorJitter: + def __init__(self, consistent_transform, brightness, contrast, saturation, hue): + self.consistent_transform = consistent_transform + self.brightness = ( + brightness + if isinstance(brightness, list) + else [max(0, 1 - brightness), 1 + brightness] + ) + self.contrast = ( + contrast + if isinstance(contrast, list) + else [max(0, 1 - contrast), 1 + contrast] + ) + self.saturation = ( + saturation + if isinstance(saturation, list) + else [max(0, 1 - saturation), 1 + saturation] + ) + self.hue = hue if isinstance(hue, list) or hue is None else ([-hue, hue]) + + def __call__(self, datapoint: Datapoint, **kwargs): + if self.consistent_transform: + # Create a color jitter transformation params + ( + fn_idx, + brightness_factor, + contrast_factor, + saturation_factor, + hue_factor, + ) = T.ColorJitter.get_params( + self.brightness, self.contrast, self.saturation, self.hue + ) + for img in datapoint.images: + if not self.consistent_transform: + ( + fn_idx, + brightness_factor, + contrast_factor, + saturation_factor, + hue_factor, + ) = T.ColorJitter.get_params( + self.brightness, self.contrast, self.saturation, self.hue + ) + for fn_id in fn_idx: + if fn_id == 0 and brightness_factor is not None: + img.data = F.adjust_brightness(img.data, brightness_factor) + elif fn_id == 1 and contrast_factor is not None: + img.data = F.adjust_contrast(img.data, contrast_factor) + elif fn_id == 2 and saturation_factor is not None: + img.data = F.adjust_saturation(img.data, saturation_factor) + elif fn_id == 3 and hue_factor is not None: + img.data = F.adjust_hue(img.data, hue_factor) + return datapoint + + +class RandomAffine: + def __init__( + self, + degrees, + consistent_transform, + scale=None, + translate=None, + shear=None, + image_mean=(123, 116, 103), + log_warning=True, + num_tentatives=1, + image_interpolation="bicubic", + ): + """ + The mask is required for this transform. + if consistent_transform if True, then the same random affine is applied to all frames and masks. + """ + self.degrees = degrees if isinstance(degrees, list) else ([-degrees, degrees]) + self.scale = scale + self.shear = ( + shear if isinstance(shear, list) else ([-shear, shear] if shear else None) + ) + self.translate = translate + self.fill_img = image_mean + self.consistent_transform = consistent_transform + self.log_warning = log_warning + self.num_tentatives = num_tentatives + + if image_interpolation == "bicubic": + self.image_interpolation = InterpolationMode.BICUBIC + elif image_interpolation == "bilinear": + self.image_interpolation = InterpolationMode.BILINEAR + else: + raise NotImplementedError + + def __call__(self, datapoint: Datapoint, **kwargs): + for _tentative in range(self.num_tentatives): + res = self.transform_datapoint(datapoint) + if res is not None: + return res + + if self.log_warning: + logging.warning( + f"Skip RandomAffine for zero-area mask in first frame after {self.num_tentatives} tentatives" + ) + return datapoint + + def transform_datapoint(self, datapoint: Datapoint): + _, height, width = F.get_dimensions(datapoint.images[0].data) + img_size = [width, height] + + if self.consistent_transform: + # Create a random affine transformation + affine_params = T.RandomAffine.get_params( + degrees=self.degrees, + translate=self.translate, + scale_ranges=self.scale, + shears=self.shear, + img_size=img_size, + ) + + for img_idx, img in enumerate(datapoint.images): + this_masks = [ + obj.segment.unsqueeze(0) if obj.segment is not None else None + for obj in img.objects + ] + if not self.consistent_transform: + # if not consistent we create a new affine params for every frame&mask pair Create a random affine transformation + affine_params = T.RandomAffine.get_params( + degrees=self.degrees, + translate=self.translate, + scale_ranges=self.scale, + shears=self.shear, + img_size=img_size, + ) + + transformed_bboxes, transformed_masks = [], [] + for i in range(len(img.objects)): + if this_masks[i] is None: + transformed_masks.append(None) + # Dummy bbox for a dummy target + transformed_bboxes.append(torch.tensor([[0, 0, 0, 0]])) + else: + transformed_mask = F.affine( + this_masks[i], + *affine_params, + interpolation=InterpolationMode.NEAREST, + fill=0.0, + ) + if img_idx == 0 and transformed_mask.max() == 0: + # We are dealing with a video and the object is not visible in the first frame + # Return the datapoint without transformation + return None + transformed_bbox = masks_to_boxes(transformed_mask) + transformed_bboxes.append(transformed_bbox) + transformed_masks.append(transformed_mask.squeeze()) + + for i in range(len(img.objects)): + img.objects[i].bbox = transformed_bboxes[i] + img.objects[i].segment = transformed_masks[i] + + img.data = F.affine( + img.data, + *affine_params, + interpolation=self.image_interpolation, + fill=self.fill_img, + ) + return datapoint + + +class RandomResizedCrop: + def __init__( + self, + consistent_transform, + size, + scale=None, + ratio=None, + log_warning=True, + num_tentatives=4, + keep_aspect_ratio=False, + ): + """ + The mask is required for this transform. + if consistent_transform if True, then the same random resized crop is applied to all frames and masks. + """ + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + elif isinstance(size, Sequence) and len(size) == 1: + self.size = (size[0], size[0]) + elif len(size) != 2: + raise ValueError("Please provide only two dimensions (h, w) for size.") + else: + self.size = size + + self.scale = scale if scale is not None else (0.08, 1.0) + self.ratio = ratio if ratio is not None else (3.0 / 4.0, 4.0 / 3.0) + self.consistent_transform = consistent_transform + self.log_warning = log_warning + self.num_tentatives = num_tentatives + self.keep_aspect_ratio = keep_aspect_ratio + + def __call__(self, datapoint: Datapoint, **kwargs): + for _tentative in range(self.num_tentatives): + res = self.transform_datapoint(datapoint) + if res is not None: + return res + + if self.log_warning: + logging.warning( + f"Skip RandomResizeCrop for zero-area mask in first frame after {self.num_tentatives} tentatives" + ) + return datapoint + + def transform_datapoint(self, datapoint: Datapoint): + if self.keep_aspect_ratio: + original_size = datapoint.images[0].size + original_ratio = original_size[1] / original_size[0] + ratio = [r * original_ratio for r in self.ratio] + else: + ratio = self.ratio + + if self.consistent_transform: + # Create a random crop transformation + crop_params = T.RandomResizedCrop.get_params( + img=datapoint.images[0].data, + scale=self.scale, + ratio=ratio, + ) + + for img_idx, img in enumerate(datapoint.images): + if not self.consistent_transform: + # Create a random crop transformation + crop_params = T.RandomResizedCrop.get_params( + img=img.data, + scale=self.scale, + ratio=ratio, + ) + + this_masks = [ + obj.segment.unsqueeze(0) if obj.segment is not None else None + for obj in img.objects + ] + + transformed_bboxes, transformed_masks = [], [] + for i in range(len(img.objects)): + if this_masks[i] is None: + transformed_masks.append(None) + # Dummy bbox for a dummy target + transformed_bboxes.append(torch.tensor([[0, 0, 0, 0]])) + else: + transformed_mask = F.resized_crop( + this_masks[i], + *crop_params, + size=self.size, + interpolation=InterpolationMode.NEAREST, + ) + if img_idx == 0 and transformed_mask.max() == 0: + # We are dealing with a video and the object is not visible in the first frame + # Return the datapoint without transformation + return None + transformed_masks.append(transformed_mask.squeeze()) + transformed_bbox = masks_to_boxes(transformed_mask) + transformed_bboxes.append(transformed_bbox) + + # Set the new boxes and masks if all transformed masks and boxes are good. + for i in range(len(img.objects)): + img.objects[i].bbox = transformed_bboxes[i] + img.objects[i].segment = transformed_masks[i] + + img.data = F.resized_crop( + img.data, + *crop_params, + size=self.size, + interpolation=InterpolationMode.BILINEAR, + ) + return datapoint + + +class ResizeToMaxIfAbove: + # Resize datapoint image if one of its sides is larger that max_size + def __init__( + self, + max_size=None, + ): + self.max_size = max_size + + def __call__(self, datapoint: Datapoint, **kwargs): + _, height, width = F.get_dimensions(datapoint.images[0].data) + + if height <= self.max_size and width <= self.max_size: + # The original frames are small enough + return datapoint + elif height >= width: + new_height = self.max_size + new_width = int(round(self.max_size * width / height)) + else: + new_height = int(round(self.max_size * height / width)) + new_width = self.max_size + + size = new_height, new_width + + for index in range(len(datapoint.images)): + datapoint.images[index].data = F.resize(datapoint.images[index].data, size) + + for obj in datapoint.images[index].objects: + obj.segment = F.resize( + obj.segment[None, None], + size, + interpolation=InterpolationMode.NEAREST, + ).squeeze() + + h, w = size + datapoint.images[index].size = (h, w) + return datapoint + + +def get_bbox_xyxy_abs_coords_from_mask(mask): + """Get the bounding box (XYXY format w/ absolute coordinates) of a binary mask.""" + assert mask.dim() == 2 + rows = torch.any(mask, dim=1) + cols = torch.any(mask, dim=0) + row_inds = rows.nonzero().view(-1) + col_inds = cols.nonzero().view(-1) + if row_inds.numel() == 0: + # mask is empty + bbox = torch.zeros(1, 4, dtype=torch.float32) + bbox_area = 0.0 + else: + ymin, ymax = row_inds.min(), row_inds.max() + xmin, xmax = col_inds.min(), col_inds.max() + bbox = torch.tensor([xmin, ymin, xmax, ymax], dtype=torch.float32).view(1, 4) + bbox_area = float((ymax - ymin) * (xmax - xmin)) + return bbox, bbox_area + + +class MotionBlur: + def __init__(self, kernel_size=5, consistent_transform=True, p=0.5): + assert kernel_size % 2 == 1, "Kernel size must be odd." + self.kernel_size = kernel_size + self.consistent_transform = consistent_transform + self.p = p + + def __call__(self, datapoint: Datapoint, **kwargs): + if random.random() >= self.p: + return datapoint + if self.consistent_transform: + # Generate a single motion blur kernel for all images + kernel = self._generate_motion_blur_kernel() + for img in datapoint.images: + if not self.consistent_transform: + # Generate a new motion blur kernel for each image + kernel = self._generate_motion_blur_kernel() + img.data = self._apply_motion_blur(img.data, kernel) + + return datapoint + + def _generate_motion_blur_kernel(self): + kernel = torch.zeros((self.kernel_size, self.kernel_size)) + direction = random.choice(["horizontal", "vertical", "diagonal"]) + if direction == "horizontal": + kernel[self.kernel_size // 2, :] = 1.0 + elif direction == "vertical": + kernel[:, self.kernel_size // 2] = 1.0 + elif direction == "diagonal": + for i in range(self.kernel_size): + kernel[i, i] = 1.0 + kernel /= kernel.sum() + return kernel + + def _apply_motion_blur(self, image, kernel): + if isinstance(image, PILImage.Image): + image = F.to_tensor(image) + channels = image.shape[0] + kernel = kernel.to(image.device).unsqueeze(0).unsqueeze(0) + blurred_image = torch.nn.functional.conv2d( + image.unsqueeze(0), + kernel.repeat(channels, 1, 1, 1), + padding=self.kernel_size // 2, + groups=channels, + ) + return F.to_pil_image(blurred_image.squeeze(0)) + + +class LargeScaleJitter: + def __init__( + self, + scale_range=(0.1, 2.0), + aspect_ratio_range=(0.75, 1.33), + crop_size=(640, 640), + consistent_transform=True, + p=0.5, + ): + """ + Args:rack + scale_range (tuple): Range of scaling factors (min_scale, max_scale). + aspect_ratio_range (tuple): Range of aspect ratios (min_aspect_ratio, max_aspect_ratio). + crop_size (tuple): Target size of the cropped region (width, height). + consistent_transform (bool): Whether to apply the same transformation across all frames. + p (float): Probability of applying the transformation. + """ + self.scale_range = scale_range + self.aspect_ratio_range = aspect_ratio_range + self.crop_size = crop_size + self.consistent_transform = consistent_transform + self.p = p + + def __call__(self, datapoint: Datapoint, **kwargs): + if random.random() >= self.p: + return datapoint + + # Sample a single scale factor and aspect ratio for all frames + log_ratio = torch.log(torch.tensor(self.aspect_ratio_range)) + scale_factor = torch.empty(1).uniform_(*self.scale_range).item() + aspect_ratio = torch.exp( + torch.empty(1).uniform_(log_ratio[0], log_ratio[1]) + ).item() + + for idx, img in enumerate(datapoint.images): + if not self.consistent_transform: + # Sample a new scale factor and aspect ratio for each frame + log_ratio = torch.log(torch.tensor(self.aspect_ratio_range)) + scale_factor = torch.empty(1).uniform_(*self.scale_range).item() + aspect_ratio = torch.exp( + torch.empty(1).uniform_(log_ratio[0], log_ratio[1]) + ).item() + + # Compute the dimensions of the jittered crop + original_width, original_height = img.data.size + target_area = original_width * original_height * scale_factor + crop_width = int(round((target_area * aspect_ratio) ** 0.5)) + crop_height = int(round((target_area / aspect_ratio) ** 0.5)) + + # Randomly select the top-left corner of the crop + crop_x = random.randint(0, max(0, original_width - crop_width)) + crop_y = random.randint(0, max(0, original_height - crop_height)) + + # Extract the cropped region + datapoint = crop(datapoint, idx, (crop_x, crop_y, crop_width, crop_height)) + + # Resize the cropped region to the target crop size + datapoint = resize(datapoint, idx, self.crop_size) + + return datapoint diff --git a/src/nn/segearth_ov3/sam3/train/transforms/filter_query_transforms.py b/src/nn/segearth_ov3/sam3/train/transforms/filter_query_transforms.py new file mode 100644 index 0000000..2d6708f --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/transforms/filter_query_transforms.py @@ -0,0 +1,607 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +import random + +from collections import defaultdict +from typing import List, Optional, Union + +import torch + +from sam3.train.data.sam3_image_dataset import Datapoint, FindQuery, Object + + +class FilterDataPointQueries: + find_ids_to_filter: set = None + get_ids_to_filter: set = None + obj_ids_to_filter: set = None # stored as pairs (img_id, obj_id) + + def identify_queries_to_filter(self, datapoint: Datapoint) -> None: + """ + Compute set of query ids to keep, for both find and get queries + """ + raise NotImplementedError + + def _do_filter_query(self, query: Union[FindQuery], query_id: int): + assert self.find_ids_to_filter is not None + + return query_id in self.find_ids_to_filter + + +class FilterQueryWithText(FilterDataPointQueries): + """ + Filter all datapoints which have query text in a specified list of exluded terms + """ + + def __init__( + self, exclude_find_keys: List[str] = None, exclude_get_keys: List[str] = None + ): + self.find_filter_keys = exclude_find_keys if exclude_find_keys else [] + self.get_filter_keys = exclude_get_keys if exclude_get_keys else [] + + def identify_queries_to_filter(self, datapoint): + self.obj_ids_to_filter = set() + del_find_ids = [] + del_get_ids = [] + for i, f_q in enumerate(datapoint.find_queries): + if f_q.query_text in self.find_filter_keys: + del_find_ids.append(i) + + self.find_ids_to_filter = set(del_find_ids) + + +class KeepMaxNumFindQueries(FilterDataPointQueries): + def __init__( + self, max_num_find_queries: int, retain_positive_queries: bool = False + ): + self.max_num_find_queries = max_num_find_queries + self.retain_positive_queries = retain_positive_queries + + def identify_queries_to_filter(self, datapoint: Datapoint) -> None: + self.obj_ids_to_filter = set() + num_find_queries = len(datapoint.find_queries) + if num_find_queries <= self.max_num_find_queries: + self.find_ids_to_filter = set() # keep all find queries + return + + if not self.retain_positive_queries: + all_find_query_ids = list(range(num_find_queries)) + num_queries_to_filter = max(0, num_find_queries - self.max_num_find_queries) + query_ids_to_filter = random.sample( + all_find_query_ids, k=num_queries_to_filter + ) + else: + # keep up to max_num_find_queries postive find queries and fill + # the remaining slots (if any) with negative find queries + pos_find_ids, neg_find_ids = [], [] + for i, f_q in enumerate(datapoint.find_queries): + # Negative finds return an empty list of object_ids_output + if len(f_q.object_ids_output) == 0: + neg_find_ids.append(i) + else: + pos_find_ids.append(i) + + if len(pos_find_ids) >= self.max_num_find_queries: + # we have more positive find queries than `max_num_find_queries`, + # so we subsample postive find queries and remove all negative find queries + num_queries_to_filter = len(pos_find_ids) - self.max_num_find_queries + query_ids_to_filter = random.sample( + pos_find_ids, k=num_queries_to_filter + ) + query_ids_to_filter.extend(neg_find_ids) + else: + # we have fewer positive find queries than `max_num_find_queries` + # so we need to fill the remaining with negative find queries + num_queries_to_filter = num_find_queries - self.max_num_find_queries + query_ids_to_filter = random.sample( + neg_find_ids, k=num_queries_to_filter + ) + + assert len(query_ids_to_filter) == num_find_queries - self.max_num_find_queries + self.find_ids_to_filter = set(query_ids_to_filter) + + +class KeepMaxNumFindQueriesVideo(FilterDataPointQueries): + def __init__( + self, + video_mosaic_max_num_find_queries_per_frame: int, + retain_positive_queries: bool = False, + ): + self.video_mosaic_max_num_find_queries_per_frame = ( + video_mosaic_max_num_find_queries_per_frame + ) + self.retain_positive_queries = retain_positive_queries + + def identify_queries_to_filter(self, datapoint: Datapoint) -> None: + self.obj_ids_to_filter = set() + num_find_queries = len(datapoint.find_queries) + + findQueries_to_imageIds = defaultdict(list) + max_queries_per_frame = True + for i, f_q in enumerate(datapoint.find_queries): + findQueries_to_imageIds[f_q.image_id].append(i) + if ( + len(findQueries_to_imageIds[f_q.image_id]) + > self.video_mosaic_max_num_find_queries_per_frame + ): + max_queries_per_frame = False + + if max_queries_per_frame: + self.find_ids_to_filter = set() + return + + num_frames = len(findQueries_to_imageIds) + findQueries_0 = findQueries_to_imageIds[0] + num_find_queries_0 = len(findQueries_0) + max_num_find_queries_per_frame = ( + self.video_mosaic_max_num_find_queries_per_frame + ) + if not self.retain_positive_queries: + find_query_ids_0 = list(range(num_find_queries_0)) + num_queries_to_filter = max( + 0, num_find_queries_0 - max_num_find_queries_per_frame + ) + query_ids_to_filter_0 = random.sample( + find_query_ids_0, k=num_queries_to_filter + ) + else: + # keep up to max_num_find_queries postive find queries and fill + # the remaining slots (if any) with negative find queries + pos_find_ids_0, neg_find_ids_0 = [], [] + for i, f_q_id in enumerate(findQueries_0): + f_q = datapoint.find_queries[f_q_id] + # Negative finds return an empty list of object_ids_output + if len(f_q.object_ids_output) == 0: + neg_find_ids_0.append(i) + else: + pos_find_ids_0.append(i) + + if len(pos_find_ids_0) >= max_num_find_queries_per_frame: + # we have more positive find queries than `max_num_find_queries`, + # so we subsample postive find queries and remove all negative find queries + num_queries_to_filter = ( + len(pos_find_ids_0) - max_num_find_queries_per_frame + ) + query_ids_to_filter_0 = random.sample( + pos_find_ids_0, k=num_queries_to_filter + ) + query_ids_to_filter_0.extend(neg_find_ids_0) + else: + # we have fewer positive find queries than `max_num_find_queries` + # so we need to fill the remaining with negative find queries + num_queries_to_filter = ( + num_find_queries_0 - max_num_find_queries_per_frame + ) + query_ids_to_filter_0 = random.sample( + neg_find_ids_0, k=num_queries_to_filter + ) + + # get based on frame 0 all find queries from all the frames with the same indices as in frame 0 + query_ids_to_filter = [] + for i in range(num_frames): + findQueries_i = findQueries_to_imageIds[i] + query_ids_to_filter.extend( + [findQueries_i[j] for j in query_ids_to_filter_0] + ) + + assert ( + len(query_ids_to_filter) + == num_find_queries + - self.video_mosaic_max_num_find_queries_per_frame * num_frames + ) + self.find_ids_to_filter = set(query_ids_to_filter) + + +class KeepSemanticFindQueriesOnly(FilterDataPointQueries): + def identify_queries_to_filter(self, datapoint: Datapoint) -> None: + self.obj_ids_to_filter = set() + self.find_ids_to_filter = { + i for i, q in enumerate(datapoint.find_queries) if q.input_bbox is not None + } # filter (remove) geometric find queries (whose input_bbox is not None) + + # Keep all get queries which don't depend on filtered finds + + +class KeepUnaryFindQueriesOnly(FilterDataPointQueries): + def identify_queries_to_filter(self, datapoint: Datapoint) -> None: + self.obj_ids_to_filter = set() + self.find_ids_to_filter = set() + + # Keep all get queries which don't depend on filtered finds + + +class FilterZeroBoxQueries(FilterDataPointQueries): + """ + Filters all find queries which predict a box with zero area + """ + + @staticmethod + def _is_zero_area_object(obj: Object): + # Check if height or width of bounding box is zero + bbox = obj.bbox # Assume in XYXY format + height = bbox[..., 3].item() - bbox[..., 1].item() + width = bbox[..., 2].item() - bbox[..., 0].item() + + return height == 0 or width == 0 + + def identify_queries_to_filter(self, datapoint): + self.obj_ids_to_filter = set() + + # Find objects with zero area + # Assume only one image per datapoint + image_objects = datapoint.images[0].objects + exclude_objects = { + obj_id + for obj_id, obj in enumerate(image_objects) + if self._is_zero_area_object(obj) + } + + # If a query predicts an object with zero area, drop the whole find query + del_find_ids = [] + for i, f_q in enumerate(datapoint.find_queries): + f_q_objects = set(f_q.object_ids_output) + if len(exclude_objects.intersection(f_q_objects)) > 0: + del_find_ids.append(i) + + self.find_ids_to_filter = set(del_find_ids) + + +class FilterFindQueriesWithTooManyOut(FilterDataPointQueries): + """ + Filters all find queries which have more than a specified number of objects in the output + """ + + def __init__(self, max_num_objects: int): + self.max_num_objects = max_num_objects + + def identify_queries_to_filter(self, datapoint): + self.obj_ids_to_filter = set() + + # If a query predicts more than max_num_objects, drop the whole find query + del_find_ids = [] + for i, f_q in enumerate(datapoint.find_queries): + if len(f_q.object_ids_output) > self.max_num_objects: + del_find_ids.append(i) + + self.find_ids_to_filter = set(del_find_ids) + + +class FilterEmptyTargets(FilterDataPointQueries): + """ + Filters all targets which have zero area + """ + + def identify_queries_to_filter(self, datapoint): + self.obj_ids_to_filter = set() + + for img_id in range(len(datapoint.images)): + for obj_id, obj in enumerate(datapoint.images[img_id].objects): + if obj.area < 1e-6: + self.obj_ids_to_filter.add((img_id, obj_id)) + self.find_ids_to_filter = set() + + +class FilterNonExhaustiveFindQueries(FilterDataPointQueries): + """ + Filters all find queries which are non-exhaustive + """ + + def __init__(self, exhaustivity_type: str): + """ + Args: + exhaustivity_type: Can be "pixel" or "instance": + -pixel: filter queries where the union of all segments covers every pixel belonging to target class + -instance: filter queries where there are non-separable or non annotated instances + Note that instance exhaustivity implies pixel exhaustivity + """ + assert exhaustivity_type in ["pixel", "instance"] + self.exhaustivity_type = exhaustivity_type + + def identify_queries_to_filter(self, datapoint): + self.obj_ids_to_filter = set() + + # If a query predicts more than max_num_objects, drop the whole find query + del_find_ids = [] + for i, f_q in enumerate(datapoint.find_queries): + if self.exhaustivity_type == "instance": + if not f_q.is_exhaustive: + del_find_ids.append(i) + elif self.exhaustivity_type == "pixel": + if f_q.is_pixel_exhaustive is not None and not f_q.is_pixel_exhaustive: + del_find_ids.append(i) + else: + raise RuntimeError( + f"Unknown exhaustivity type {self.exhaustivity_type}" + ) + + self.find_ids_to_filter = set(del_find_ids) + + +class FilterInvalidGeometricQueries(FilterDataPointQueries): + """ + Filters geometric queries whose output got deleted (eg due to cropping) + """ + + def identify_queries_to_filter(self, datapoint): + self.obj_ids_to_filter = set() + + # If a query predicts more than max_num_objects, drop the whole find query + del_find_ids = [] + for i, f_q in enumerate(datapoint.find_queries): + if f_q.input_bbox is not None and f_q.query_text == "geometric": + if len(f_q.object_ids_output) == 0: + del_find_ids.append(i) + self.find_ids_to_filter = set(del_find_ids) + + +class FlexibleFilterFindGetQueries: + def __init__( + self, query_filter: FilterDataPointQueries, enabled: bool = True + ) -> None: + self.query_filter = query_filter + self.enabled = enabled + + def __call__(self, datapoint, **kwargs): + if not self.enabled: + return datapoint + + # Identify all queries to filter + self.query_filter.identify_queries_to_filter(datapoint=datapoint) + + del_find_ids = [] + del_get_ids = [] + for i, f_q in enumerate(datapoint.find_queries): + if self.query_filter._do_filter_query(f_q, i): + datapoint.find_queries[i] = None + del_find_ids.append(i) + + new_find_queries = [] + new_get_queries = [] + + find_old_to_new_map = {} + get_old_to_new_map = {} + + find_counter = 0 + get_counter = 0 + + for i, f_q in enumerate(datapoint.find_queries): + if f_q is not None: + find_old_to_new_map[i] = find_counter + find_counter += 1 + new_find_queries.append(f_q) + + start_with_zero_check = False + for n_f_q in new_find_queries: + if n_f_q.query_processing_order == 0: + start_with_zero_check = True + break + + if len(new_find_queries) == 0: + start_with_zero_check = True + + assert ( + start_with_zero_check + ), "Invalid Find queries, they need to start at query_processing_order = 0" + + datapoint.find_queries = new_find_queries + + if len(datapoint.find_queries) == 0: + print("Warning: No find queries left in datapoint, this is not allowed") + print("Filtering function:", self.query_filter) + print("Datapoint:", datapoint) + raise ValueError + + # The deletion may have removed intermediate steps, so we need to remap to make them contiguous again + all_stages = sorted( + list(set(q.query_processing_order for q in datapoint.find_queries)) + ) + stage_map = {qpo: i for i, qpo in enumerate(all_stages)} + for i in range(len(datapoint.find_queries)): + qpo = datapoint.find_queries[i].query_processing_order + datapoint.find_queries[i].query_processing_order = stage_map[qpo] + + # Final step, clear up objects that are not used anymore + for img_id in range(len(datapoint.images)): + all_objects_ids = set( + i + for find in datapoint.find_queries + for i in find.object_ids_output + if find.image_id == img_id + ) + unused_ids = ( + set(range(len(datapoint.images[img_id].objects))) - all_objects_ids + ) + for tgt_img_id, tgt_obj_id in self.query_filter.obj_ids_to_filter: + if tgt_img_id == img_id: + unused_ids.add(tgt_obj_id) + + if len(unused_ids) > 0: + old_objects = datapoint.images[img_id].objects + object_old_to_new_map = {} + new_objects = [] + for i, o in enumerate(old_objects): + if i not in unused_ids: + object_old_to_new_map[i] = len(new_objects) + new_objects.append(o) + + datapoint.images[img_id].objects = new_objects + + # Remap the outputs of the find queries + affected_find_queries_ids = set() + object_old_to_new_map_per_query = {} + for fid, find in enumerate(datapoint.find_queries): + if find.image_id == img_id: + old_object_ids_output = find.object_ids_output + object_old_to_new_map_per_query[fid] = {} + find.object_ids_output = [] + for oid, old_obj_id in enumerate(old_object_ids_output): + if old_obj_id not in unused_ids: + new_obj_id = object_old_to_new_map[old_obj_id] + find.object_ids_output.append(new_obj_id) + object_old_to_new_map_per_query[fid][oid] = ( + len(find.object_ids_output) - 1 + ) + affected_find_queries_ids.add(fid) + + # finally remove unused images + all_imgs_to_keep = set() + for f_q in datapoint.find_queries: + all_imgs_to_keep.add(f_q.image_id) + + old_img_id_to_new_img_id = {} + new_images = [] + for img_id, img in enumerate(datapoint.images): + if img_id in all_imgs_to_keep: + old_img_id_to_new_img_id[img_id] = len(new_images) + new_images.append(img) + datapoint.images = new_images + + for f_q in datapoint.find_queries: + f_q.image_id = old_img_id_to_new_img_id[f_q.image_id] + + return datapoint + + +class AddPrefixSuffixToFindText: + """ + Add prefix or suffix strings to find query text on the fly. + + If `condition_on_text` is True, the prefix or suffix strings are only added + to those find query text in `condition_text_list` (case-insensitive). + """ + + def __init__( + self, + prefix: Optional[str] = None, + suffix: Optional[str] = None, + condition_on_text: bool = False, + condition_text_list: Optional[List[str]] = None, + enabled: bool = True, + ) -> None: + self.prefix = prefix + self.suffix = suffix + self.condition_on_text = condition_on_text + if self.condition_on_text: + assert condition_text_list is not None + self.condition_text_set = {s.lower().strip() for s in condition_text_list} + self.enabled = enabled + if self.enabled: + logging.info( + f"AddPrefixSuffixToFindText: prefix={prefix}, suffix={suffix}, " + f"condition_on_text={condition_on_text}, condition_text_list={condition_text_list}" + ) + + def __call__(self, datapoint, **kwargs): + if not self.enabled: + return datapoint + + for find in datapoint.find_queries: + if find.query_text == "geometric": + # skip geometric find queries + continue + if ( + self.condition_on_text + and find.query_text.lower().strip() not in self.condition_text_set + ): + # if condition_on_text is True, skip those queries not in condition_text_set + continue + + # add prefix and/or suffix strings to the find query text + if self.prefix is not None: + find.query_text = self.prefix + find.query_text + if self.suffix is not None: + find.query_text = find.query_text + self.suffix + + return datapoint + + +class FilterCrowds(FilterDataPointQueries): + def identify_queries_to_filter(self, datapoint: Datapoint) -> None: + """ + Compute set of query ids to keep, for both find and get queries + """ + self.obj_ids_to_filter = set() + self.find_ids_to_filter = set() + # self.get_ids_to_filter = set() + for img_id, img in enumerate(datapoint.images): + for obj_id, obj in enumerate(img.objects): + if obj.is_crowd: + self.obj_ids_to_filter.add((img_id, obj_id)) + + +class TextQueryToVisual: + """ + Transform a test query to a visual query (with some proba), using any of the output targets as the prompt + """ + + def __init__(self, probability, keep_text_queries=False) -> None: + self.probability = probability + assert 0 <= probability <= 1 + self.keep_text_queries = keep_text_queries + + def __call__(self, datapoint: Datapoint, **kwargs): + for find in datapoint.find_queries: + if find.input_bbox is not None or find.input_points is not None: + # skip geometric find queries + continue + + if len(find.object_ids_output) == 0: + # Can't create a visual query, skip + continue + + if find.query_processing_order > 0: + # Second stage query, can't use + continue + + if random.random() > self.probability: + continue + + selected_vq_id = random.choice(find.object_ids_output) + img_id = find.image_id + + find.input_bbox = datapoint.images[img_id].objects[selected_vq_id].bbox + find.input_bbox_label = torch.ones(1, dtype=torch.bool) + if not self.keep_text_queries: + find.query_text = "visual" + + return datapoint + + +class RemoveInputBoxes: + """ + Remove input boxes from find queries + """ + + def __init__(self) -> None: + pass + + def __call__(self, datapoint: Datapoint, **kwargs): + for find in datapoint.find_queries: + if find.input_bbox is None: + continue + + if find.query_text == "geometric": + print("Warning: removing input box from geometric find query") + + find.input_bbox = None + return datapoint + + +class OverwriteTextQuery: + """ + With some probability, overwrite the text query with a custom text + """ + + def __init__(self, target_text, probability=1.0) -> None: + self.probability = probability + self.target_text = target_text + assert 0 <= probability <= 1 + + def __call__(self, datapoint: Datapoint, **kwargs): + for find in datapoint.find_queries: + if random.random() > self.probability: + continue + + find.query_text = self.target_text + + return datapoint diff --git a/src/nn/segearth_ov3/sam3/train/transforms/point_sampling.py b/src/nn/segearth_ov3/sam3/train/transforms/point_sampling.py new file mode 100644 index 0000000..e083fde --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/transforms/point_sampling.py @@ -0,0 +1,345 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import cv2 +import numpy as np +import torch +from PIL import Image as PILImage +from pycocotools import mask as mask_util + +from sam3.train.data.sam3_image_dataset import Datapoint +from torchvision.ops import masks_to_boxes + + +def sample_points_from_rle(rle, n_points, mode, box=None, normalize=True): + """ + Sample random points from a mask provided in COCO RLE format. 'mode' + 'mode' is in ["centered", "random_mask", "random_box"] + "centered": points are sampled farthest from the mask edges and each other + "random_mask": points are sampled uniformly from the mask + "random_box": points are sampled uniformly from the annotation's box + 'box' must be provided if 'mode' is "random_box". + If 'normalize' is true, points are in [0,1], relative to mask h,w. + """ + mask = np.ascontiguousarray(mask_util.decode(rle)) + points = sample_points_from_mask(mask, n_points, mode, box) + + if normalize: + h, w = mask.shape + norm = np.array([w, h, 1.0])[None, :] + points = points / norm + + return points + + +def sample_points_from_mask(mask, n_points, mode, box=None): + if mode == "centered": + points = center_positive_sample(mask, n_points) + elif mode == "random_mask": + points = uniform_positive_sample(mask, n_points) + elif mode == "random_box": + assert box is not None, "'random_box' mode requires a provided box." + points = uniform_sample_from_box(mask, box, n_points) + else: + raise ValueError(f"Unknown point sampling mode {mode}.") + return points + + +def uniform_positive_sample(mask, n_points): + """ + Samples positive points uniformly from the mask. Only integer pixel + values are sampled. + """ + # Sampling directly from the uncompressed RLE would be faster but is + # likely unnecessary. + mask_points = np.stack(np.nonzero(mask), axis=0).transpose(1, 0) + assert len(mask_points) > 0, "Can't sample positive points from an empty mask." + selected_idxs = np.random.randint(low=0, high=len(mask_points), size=n_points) + selected_points = mask_points[selected_idxs] + + selected_points = selected_points[:, ::-1] # (y, x) -> (x, y) + labels = np.ones((len(selected_points), 1)) + selected_points = np.concatenate([selected_points, labels], axis=1) + + return selected_points + + +def center_positive_sample(mask, n_points): + """ + Samples points farthest from mask edges (by distance transform) + and subsequent points also farthest from each other. Each new point + sampled is treated as an edge for future points. Edges of the image are + treated as edges of the mask. + """ + + # Pad mask by one pixel on each end to assure distance transform + # avoids edges + padded_mask = np.pad(mask, 1) + + points = [] + for _ in range(n_points): + assert np.max(mask) > 0, "Can't sample positive points from an empty mask." + dist = cv2.distanceTransform(padded_mask, cv2.DIST_L2, 0) + point = np.unravel_index(dist.argmax(), dist.shape) + # Mark selected point as background so next point avoids it + padded_mask[point[0], point[1]] = 0 + points.append(point[::-1]) # (y, x) -> (x, y) + + points = np.stack(points, axis=0) + points = points - 1 # Subtract left/top padding of 1 + labels = np.ones((len(points), 1)) + points = np.concatenate([points, labels], axis=1) + + return points + + +def uniform_sample_from_box(mask, box, n_points): + """ + Sample points uniformly from the provided box. The points' labels + are determined by the provided mask. Does not guarantee a positive + point is sampled. The box is assumed unnormalized in XYXY format. + Points are sampled at integer values. + """ + + # Since lower/right edges are exclusive, ceil can be applied to all edges + int_box = np.ceil(box) + + x = np.random.randint(low=int_box[0], high=int_box[2], size=n_points) + y = np.random.randint(low=int_box[1], high=int_box[3], size=n_points) + labels = mask[y, x] + points = np.stack([x, y, labels], axis=1) + + return points + + +def rescale_box_xyxy(box, factor, imsize=None): + """ + Rescale a box providing in unnormalized XYXY format, fixing the center. + If imsize is provided, clamp to the image. + """ + cx, cy = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2 + w, h = box[2] - box[0], box[3] - box[1] + + new_w, new_h = factor * w, factor * h + + new_x0, new_y0 = cx - new_w / 2, cy - new_h / 2 + new_x1, new_y1 = cx + new_w / 2, cy + new_h / 2 + + if imsize is not None: + new_x0 = max(min(new_x0, imsize[1]), 0) + new_x1 = max(min(new_x1, imsize[1]), 0) + new_y0 = max(min(new_y0, imsize[0]), 0) + new_y1 = max(min(new_y1, imsize[0]), 0) + + return [new_x0, new_y0, new_x1, new_y1] + + +def noise_box(box, im_size, box_noise_std, box_noise_max, min_box_area): + if box_noise_std <= 0.0: + return box + noise = box_noise_std * torch.randn(size=(4,)) + w, h = box[2] - box[0], box[3] - box[1] + scale_factor = torch.tensor([w, h, w, h]) + noise = noise * scale_factor + if box_noise_max is not None: + noise = torch.clamp(noise, -box_noise_max, box_noise_max) + input_box = box + noise + # Clamp to maximum image size + img_clamp = torch.tensor([im_size[1], im_size[0], im_size[1], im_size[0]]) + input_box = torch.maximum(input_box, torch.zeros_like(input_box)) + input_box = torch.minimum(input_box, img_clamp) + if (input_box[2] - input_box[0]) * (input_box[3] - input_box[1]) <= min_box_area: + return box + + return input_box + + +class RandomGeometricInputsAPI: + """ + For geometric queries, replaces the input box or points with a random + one sampled from the GT mask. Segments must be provided for objects + that are targets of geometric queries, and must be binary masks. Existing + point and box queries in the datapoint will be ignored and completely replaced. + Will sample points and boxes in XYXY format in absolute pixel space. + + Geometry queries are currently determined by taking any query whose + query text is a set value. + + Args: + num_points (int or (int, int)): how many points to sample. If a tuple, + sample a random number of points uniformly over the inclusive range. + box_chance (float): fraction of time a box is sampled. A box will replace + one sampled point. + box_noise_std (float): if greater than 0, add noise to the sampled boxes + with this std. Noise is relative to the length of the box side. + box_noise_max (int): if not none, truncate any box noise larger than this + in terms of absolute pixels. + resample_box_from_mask (bool): if True, any sampled box will be determined + by finding the extrema of the provided mask. If False, the bbox provided + in the target object will be used. + point_sample_mode (str): In ["centered", "random_mask", "random_box"], + controlling how points are sampled: + "centered": points are sampled farthest from the mask edges and each other + "random_mask": points are sampled uniformly from the mask + "random_box": points are sampled uniformly from the annotation's box + Note that "centered" may be too slow for on-line generation. + geometric_query_str (str): what string in query_text indicates a + geometry query. + minimum_box_area (float): sampled boxes with area this size or smaller after + noising will use the original box instead. It is the input's responsibility + to avoid original boxes that violate necessary area bounds. + concat_points (bool): if True, any sampled points will be added to existing + ones instead of replacing them. + + """ + + def __init__( + self, + num_points, + box_chance, + box_noise_std=0.0, + box_noise_max=None, + minimum_box_area=0.0, + resample_box_from_mask=False, + point_sample_mode="random_mask", + sample_box_scale_factor=1.0, + geometric_query_str="geometric", + concat_points=False, + ): + self.num_points = num_points + if not isinstance(self.num_points, int): + # Convert from inclusive range to exclusive range expected by torch + self.num_points[1] += 1 + self.num_points = tuple(self.num_points) + self.box_chance = box_chance + self.box_noise_std = box_noise_std + self.box_noise_max = box_noise_max + self.minimum_box_area = minimum_box_area + self.resample_box_from_mask = resample_box_from_mask + self.point_sample_mode = point_sample_mode + assert point_sample_mode in [ + "centered", + "random_mask", + "random_box", + ], "Unknown point sample mode." + self.geometric_query_str = geometric_query_str + self.concat_points = concat_points + self.sample_box_scale_factor = sample_box_scale_factor + + def _sample_num_points_and_if_box(self): + if isinstance(self.num_points, tuple): + n_points = torch.randint( + low=self.num_points[0], high=self.num_points[1], size=(1,) + ).item() + else: + n_points = self.num_points + if self.box_chance > 0.0: + use_box = torch.rand(size=(1,)).item() < self.box_chance + n_points -= int(use_box) # box stands in for one point + else: + use_box = False + return n_points, use_box + + def _get_original_box(self, target_object): + if not self.resample_box_from_mask: + return target_object.bbox + mask = target_object.segment + return masks_to_boxes(mask[None, :, :])[0] + + def _get_target_object(self, datapoint, query): + img = datapoint.images[query.image_id] + targets = query.object_ids_output + assert ( + len(targets) == 1 + ), "Geometric queries only support a single target object." + target_idx = targets[0] + return img.objects[target_idx] + + def __call__(self, datapoint, **kwargs): + for query in datapoint.find_queries: + if query.query_text != self.geometric_query_str: + continue + + target_object = self._get_target_object(datapoint, query) + n_points, use_box = self._sample_num_points_and_if_box() + box = self._get_original_box(target_object) + + mask = target_object.segment + if n_points > 0: + # FIXME: The conversion to numpy and back to reuse code + # is awkward, but this is all in the dataloader worker anyway + # on CPU and so I don't think it should matter. + if self.sample_box_scale_factor != 1.0: + sample_box = rescale_box_xyxy( + box.numpy(), self.sample_box_scale_factor, mask.shape + ) + else: + sample_box = box.numpy() + input_points = sample_points_from_mask( + mask.numpy(), + n_points, + self.point_sample_mode, + sample_box, + ) + input_points = torch.as_tensor(input_points) + input_points = input_points[None, :, :] + if self.concat_points and query.input_points is not None: + input_points = torch.cat([query.input_points, input_points], dim=1) + else: + input_points = query.input_points if self.concat_points else None + + if use_box: + w, h = datapoint.images[query.image_id].size + input_box = noise_box( + box, + (h, w), + box_noise_std=self.box_noise_std, + box_noise_max=self.box_noise_max, + min_box_area=self.minimum_box_area, + ) + input_box = input_box[None, :] + else: + input_box = query.input_bbox if self.concat_points else None + + query.input_points = input_points + query.input_bbox = input_box + + return datapoint + + +class RandomizeInputBbox: + """ + Simplified version of the geometric transform that only deals with input boxes + """ + + def __init__( + self, + box_noise_std=0.0, + box_noise_max=None, + minimum_box_area=0.0, + ): + self.box_noise_std = box_noise_std + self.box_noise_max = box_noise_max + self.minimum_box_area = minimum_box_area + + def __call__(self, datapoint: Datapoint, **kwargs): + for query in datapoint.find_queries: + if query.input_bbox is None: + continue + + img = datapoint.images[query.image_id].data + if isinstance(img, PILImage.Image): + w, h = img.size + else: + assert isinstance(img, torch.Tensor) + h, w = img.shape[-2:] + + for box_id in range(query.input_bbox.shape[0]): + query.input_bbox[box_id, :] = noise_box( + query.input_bbox[box_id, :].view(4), + (h, w), + box_noise_std=self.box_noise_std, + box_noise_max=self.box_noise_max, + min_box_area=self.minimum_box_area, + ).view(1, 4) + + return datapoint diff --git a/src/nn/segearth_ov3/sam3/train/transforms/segmentation.py b/src/nn/segearth_ov3/sam3/train/transforms/segmentation.py new file mode 100644 index 0000000..3e97db0 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/transforms/segmentation.py @@ -0,0 +1,157 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import numpy as np +import pycocotools.mask as mask_utils +import torch + +import torchvision.transforms.functional as F +from PIL import Image as PILImage + +from sam3.model.box_ops import masks_to_boxes + +from sam3.train.data.sam3_image_dataset import Datapoint + + +class InstanceToSemantic(object): + """Convert instance segmentation to semantic segmentation.""" + + def __init__(self, delete_instance=True, use_rle=False): + self.delete_instance = delete_instance + self.use_rle = use_rle + + def __call__(self, datapoint: Datapoint, **kwargs): + for fquery in datapoint.find_queries: + h, w = datapoint.images[fquery.image_id].size + + if self.use_rle: + all_segs = [ + datapoint.images[fquery.image_id].objects[obj_id].segment + for obj_id in fquery.object_ids_output + ] + if len(all_segs) > 0: + # we need to double check that all rles are the correct size + # Otherwise cocotools will fail silently to an empty [0,0] mask + for seg in all_segs: + assert seg["size"] == all_segs[0]["size"], ( + "Instance segments have inconsistent sizes. " + f"Found sizes {seg['size']} and {all_segs[0]['size']}" + ) + fquery.semantic_target = mask_utils.merge(all_segs) + else: + # There is no good way to create an empty RLE of the correct size + # We resort to converting an empty box to RLE + fquery.semantic_target = mask_utils.frPyObjects( + np.array([[0, 0, 0, 0]], dtype=np.float64), h, w + )[0] + + else: + # `semantic_target` is uint8 and remains uint8 throughout the transforms + # (it contains binary 0 and 1 values just like `segment` for each object) + fquery.semantic_target = torch.zeros((h, w), dtype=torch.uint8) + for obj_id in fquery.object_ids_output: + segment = datapoint.images[fquery.image_id].objects[obj_id].segment + if segment is not None: + assert ( + isinstance(segment, torch.Tensor) + and segment.dtype == torch.uint8 + ) + fquery.semantic_target |= segment + + if self.delete_instance: + for img in datapoint.images: + for obj in img.objects: + del obj.segment + obj.segment = None + + return datapoint + + +class RecomputeBoxesFromMasks: + """Recompute bounding boxes from masks.""" + + def __call__(self, datapoint: Datapoint, **kwargs): + for img in datapoint.images: + for obj in img.objects: + # Note: if the mask is empty, the bounding box will be undefined + # The empty targets should be subsequently filtered + obj.bbox = masks_to_boxes(obj.segment) + obj.area = obj.segment.sum().item() + + return datapoint + + +class DecodeRle: + """This transform decodes RLEs into binary segments. + Implementing it as a transforms allows lazy loading. Some transforms (eg query filters) + may be deleting masks, so decoding them from the beginning is wasteful. + + This transforms needs to be called before any kind of geometric manipulation + """ + + def __call__(self, datapoint: Datapoint, **kwargs): + imgId2size = {} + warning_shown = False + for imgId, img in enumerate(datapoint.images): + if isinstance(img.data, PILImage.Image): + img_w, img_h = img.data.size + elif isinstance(img.data, torch.Tensor): + img_w, img_h = img.data.shape[-2:] + else: + raise RuntimeError(f"Unexpected image type {type(img.data)}") + + imgId2size[imgId] = (img_h, img_w) + + for obj in img.objects: + if obj.segment is not None and not isinstance( + obj.segment, torch.Tensor + ): + if mask_utils.area(obj.segment) == 0: + print("Warning, empty mask found, approximating from box") + obj.segment = torch.zeros(img_h, img_w, dtype=torch.uint8) + x1, y1, x2, y2 = obj.bbox.int().tolist() + obj.segment[y1 : max(y2, y1 + 1), x1 : max(x1 + 1, x2)] = 1 + else: + obj.segment = mask_utils.decode(obj.segment) + # segment is uint8 and remains uint8 throughout the transforms + obj.segment = torch.tensor(obj.segment).to(torch.uint8) + + if list(obj.segment.shape) != [img_h, img_w]: + # Should not happen often, but adding for security + if not warning_shown: + print( + f"Warning expected instance segmentation size to be {[img_h, img_w]} but found {list(obj.segment.shape)}" + ) + # Printing only once per datapoint to avoid spam + warning_shown = True + + obj.segment = F.resize( + obj.segment[None], (img_h, img_w) + ).squeeze(0) + + assert list(obj.segment.shape) == [img_h, img_w] + + warning_shown = False + for query in datapoint.find_queries: + if query.semantic_target is not None and not isinstance( + query.semantic_target, torch.Tensor + ): + query.semantic_target = mask_utils.decode(query.semantic_target) + # segment is uint8 and remains uint8 throughout the transforms + query.semantic_target = torch.tensor(query.semantic_target).to( + torch.uint8 + ) + if tuple(query.semantic_target.shape) != imgId2size[query.image_id]: + if not warning_shown: + print( + f"Warning expected semantic segmentation size to be {imgId2size[query.image_id]} but found {tuple(query.semantic_target.shape)}" + ) + # Printing only once per datapoint to avoid spam + warning_shown = True + + query.semantic_target = F.resize( + query.semantic_target[None], imgId2size[query.image_id] + ).squeeze(0) + + assert tuple(query.semantic_target.shape) == imgId2size[query.image_id] + + return datapoint diff --git a/src/nn/segearth_ov3/sam3/train/utils/__init__.py b/src/nn/segearth_ov3/sam3/train/utils/__init__.py new file mode 100644 index 0000000..46d37d2 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/utils/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/src/nn/segearth_ov3/sam3/train/utils/checkpoint_utils.py b/src/nn/segearth_ov3/sam3/train/utils/checkpoint_utils.py new file mode 100644 index 0000000..7f2736a --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/utils/checkpoint_utils.py @@ -0,0 +1,358 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + + +import contextlib +import fnmatch +import logging +from typing import ( + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, + Set, + Tuple, + Union, +) + +import numpy as np +import torch +import torch.nn as nn +from iopath.common.file_io import g_pathmgr +from torch.jit._script import RecursiveScriptModule + + +def unix_pattern_to_parameter_names( + constraints: List[str], all_parameter_names: Sequence[str] +) -> Union[None, Set[str]]: + """ + Go through the list of parameter names and select those that match + any of the provided constraints + """ + parameter_names = [] + for param_name in constraints: + matching_parameters = set(fnmatch.filter(all_parameter_names, param_name)) + assert ( + len(matching_parameters) > 0 + ), f"param_names {param_name} don't match any param in the given names." + parameter_names.append(matching_parameters) + return set.union(*parameter_names) + + +def filter_params_matching_unix_pattern( + patterns: List[str], state_dict: Dict[str, torch.Tensor] +) -> Dict[str, torch.Tensor]: + """ + Remove from the state dictionary the parameters matching the provided unix patterns + + Args: + patterns: the list of unix patterns to exclude + state_dict: the dictionary to filter + + Returns: + A new state dictionary + """ + if len(patterns) == 0: + return {} + + all_keys = list(state_dict.keys()) + included_keys = unix_pattern_to_parameter_names(patterns, all_keys) + return {k: state_dict[k] for k in included_keys} + + +def exclude_params_matching_unix_pattern( + patterns: List[str], state_dict: Dict[str, torch.Tensor] +) -> Dict[str, torch.Tensor]: + """ + Remove from the state dictionary the parameters matching the provided unix patterns + + Args: + patterns: the list of unix patterns to exclude + state_dict: the dictionary to filter + + Returns: + A new state dictionary + """ + if len(patterns) == 0: + return state_dict + + all_keys = list(state_dict.keys()) + excluded_keys = unix_pattern_to_parameter_names(patterns, all_keys) + return {k: v for k, v in state_dict.items() if k not in excluded_keys} + + +def _get_state_dict_summary(state_dict: Dict[str, torch.Tensor]): + keys = [] + trace = [] + for k, v in state_dict.items(): + keys.append(k) + trace.append(v.sum().item()) + trace = np.array(trace)[np.argsort(keys)] + return trace + + +def assert_skipped_parameters_are_frozen(model: nn.Module, patterns: List[str]): + """ + Verifies that all the parameters matching the provided patterns + are frozen - this acts as a safeguard when ignoring parameter + when saving checkpoints - if the parameters are in fact trainable + """ + if not patterns: + return + + frozen_state_dict = filter_params_matching_unix_pattern( + patterns=patterns, state_dict=model.state_dict() + ) + non_frozen_keys = { + n + for n, p in model.named_parameters() + if n in frozen_state_dict and p.requires_grad + } + if non_frozen_keys: + raise ValueError( + f"Parameters excluded with `skip_saving_parameters` should be frozen: {non_frozen_keys}" + ) + + +@contextlib.contextmanager +def with_check_parameter_frozen( + model: nn.Module, patterns: List[str], disabled: bool = True +): + """ + Context manager that inspects a model surrounding a piece of code + and verifies if the model has been updated by this piece of code + + The function will raise an exception if the model has been updated + on at least one of the parameter that matches one of the pattern + + Args: + model: the model that might have been updated + patterns: for the parameters we want to observe + allowed: + """ + if not patterns or disabled: + yield + return + + frozen_state_dict = filter_params_matching_unix_pattern( + patterns=patterns, state_dict=model.state_dict() + ) + summary_before = _get_state_dict_summary(frozen_state_dict) + + yield + + frozen_state_dict = filter_params_matching_unix_pattern( + patterns=patterns, state_dict=model.state_dict() + ) + summary_after = _get_state_dict_summary(frozen_state_dict) + + if not np.allclose(summary_before, summary_after, atol=1e-6): + raise ValueError( + f""" + The `model_weight_initializer` has initialized parameters frozen with `skip_saving_parameters`. + You can resolve this error by either initializing those parameters from within the model definition + or using the flag `trainer.checkpoint.initialize_after_preemption` to True. + """ + ) + + +class CkptExcludeKernel: + """ + Removes the keys from the given model state_dict that match the key_pattern. + + Args: + key_pattern: Patterns used to select the keys in the state_dict + that are eligible for this kernel. + """ + + def __init__(self, key_pattern: List[str]): + self.key_pattern = key_pattern + + def __call__(self, state_dict: Dict): + """ + Args: + state_dict: A dictionary representing the given checkpoint's state dict. + """ + if len(self.key_pattern) == 0: + return state_dict + exclude_keys = unix_pattern_to_parameter_names( + self.key_pattern, state_dict.keys() + ) + return {k: v for k, v in state_dict.items() if k not in exclude_keys} + + +def load_checkpoint( + path_list: List[str], + pick_recursive_keys: Optional[List[str]] = None, + map_location: str = "cpu", +) -> Any: + """ + Loads a checkpoint from the specified path. + + Args: + path_list: A list of paths which contain the checkpoint. Each element + is tried (in order) until a file that exists is found. That file is then + used to read the checkpoint. + pick_recursive_keys: Picks sub dicts from the loaded checkpoint if not None. + For pick_recursive_keys = ["a", "b"], will return checkpoint_dict["a"]["b"] + map_location (str): a function, torch.device, string or a dict specifying how to + remap storage locations + + Returns: Model with the matchin pre-trained weights loaded. + """ + path_exists = False + for path in path_list: + if g_pathmgr.exists(path): + path_exists = True + break + + if not path_exists: + raise ValueError(f"No path exists in {path_list}") + + with g_pathmgr.open(path, "rb") as f: + checkpoint = torch.load(f, map_location=map_location) + + logging.info(f"Loaded checkpoint from {path}") + if pick_recursive_keys is not None: + for key in pick_recursive_keys: + checkpoint = checkpoint[key] + return checkpoint + + +def get_state_dict(checkpoint, ckpt_state_dict_keys): + if isinstance(checkpoint, RecursiveScriptModule): + # This is a torchscript JIT model + return checkpoint.state_dict() + pre_train_dict = checkpoint + for i, key in enumerate(ckpt_state_dict_keys): + if (isinstance(pre_train_dict, Mapping) and key not in pre_train_dict) or ( + isinstance(pre_train_dict, Sequence) and key >= len(pre_train_dict) + ): + key_str = ( + '["' + '"]["'.join(list(map(ckpt_state_dict_keys[:i], str))) + '"]' + ) + raise KeyError( + f"'{key}' not found in checkpoint{key_str} " + f"with keys: {pre_train_dict.keys()}" + ) + pre_train_dict = pre_train_dict[key] + return pre_train_dict + + +def load_checkpoint_and_apply_kernels( + checkpoint_path: str, + checkpoint_kernels: List[Callable] = None, + ckpt_state_dict_keys: Tuple[str] = ("state_dict",), + map_location: str = "cpu", +) -> nn.Module: + """ + Performs checkpoint loading with a variety of pre-processing kernel applied in + sequence. + + Args: + checkpoint_path (str): Path to the checkpoint. + checkpoint_kernels List(Callable): A list of checkpoint processing kernels + to apply in the specified order. Supported kernels include `CkptIncludeKernel`, + `CkptExcludeKernel`, etc. These kernels are applied in the + given order. + ckpt_state_dict_keys (str): Keys containing the model state dict. + map_location (str): a function, torch.device, string or a dict specifying how to + remap storage locations + + Returns: Model with the matchin pre-trained weights loaded. + """ + assert g_pathmgr.exists(checkpoint_path), "Checkpoint '{}' not found".format( + checkpoint_path + ) + + # Load the checkpoint on CPU to avoid GPU mem spike. + with g_pathmgr.open(checkpoint_path, "rb") as f: + checkpoint = torch.load(f, map_location=map_location) + + pre_train_dict = get_state_dict(checkpoint, ckpt_state_dict_keys) + + # Not logging into info etc since it's a huge log + logging.debug( + "Loaded Checkpoint State Dict pre-kernel application: %s" + % str(", ".join(list(pre_train_dict.keys()))) + ) + # Apply kernels + if checkpoint_kernels is not None: + for f in checkpoint_kernels: + pre_train_dict = f(state_dict=pre_train_dict) + + logging.debug( + "Loaded Checkpoint State Dict Post-kernel application %s" + % str(", ".join(list(pre_train_dict.keys()))) + ) + + return pre_train_dict + + +def check_load_state_dict_errors( + missing_keys, + unexpected_keys, + strict: bool, + ignore_missing_keys: List[str] = None, + ignore_unexpected_keys: List[str] = None, +): + if ignore_missing_keys is not None and len(ignore_missing_keys) > 0: + ignored_keys = unix_pattern_to_parameter_names( + ignore_missing_keys, missing_keys + ) + missing_keys = [key for key in missing_keys if key not in ignored_keys] + + if ignore_unexpected_keys is not None and len(ignore_unexpected_keys) > 0: + ignored_unexpected_keys = unix_pattern_to_parameter_names( + ignore_unexpected_keys, unexpected_keys + ) + unexpected_keys = [ + key for key in unexpected_keys if key not in ignored_unexpected_keys + ] + + err = "State key mismatch." + if unexpected_keys: + err += f" Unexpected keys: {unexpected_keys}." + if missing_keys: + err += f" Missing keys: {missing_keys}." + + if unexpected_keys or missing_keys: + logging.warning(err) + if unexpected_keys or strict: + raise KeyError(err) + + +def load_state_dict_into_model( + state_dict: Dict, + model: nn.Module, + strict: bool = True, + ignore_missing_keys: List[str] = None, + ignore_unexpected_keys: List[str] = None, + checkpoint_kernels: List[Callable] = None, +): + """ + Loads a state dict into the given model. + + Args: + state_dict: A dictionary containing the model's + state dict, or a subset if strict is False + model: Model to load the checkpoint weights into + strict: raise if the state_dict has missing state keys + ignore_missing_keys: unix pattern of keys to ignore + """ + # Apply kernels + if checkpoint_kernels is not None: + for f in checkpoint_kernels: + state_dict = f(state_dict=state_dict) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + check_load_state_dict_errors( + missing_keys, + unexpected_keys, + strict=strict, + ignore_missing_keys=ignore_missing_keys, + ignore_unexpected_keys=ignore_unexpected_keys, + ) + return model diff --git a/src/nn/segearth_ov3/sam3/train/utils/distributed.py b/src/nn/segearth_ov3/sam3/train/utils/distributed.py new file mode 100644 index 0000000..3c87a91 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/utils/distributed.py @@ -0,0 +1,585 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import datetime +import functools +import io +import logging +import os +import random +import tempfile +import time +from typing import Any, Callable, List, Tuple + +import torch +import torch.autograd as autograd +import torch.distributed as dist + + +# Default to GPU 0 +_cuda_device_index: int = 0 + +# Setting _cuda_device_index to -1 internally implies that we should use CPU +_CPU_DEVICE_INDEX = -1 +_PRIMARY_RANK = 0 + + +@functools.lru_cache() +def _get_global_gloo_group(): + """ + Return a process group based on gloo backend, containing all the ranks + The result is cached. + """ + + if dist.get_backend() == "nccl": + # Increase timeout from 1800 sec to 43200 sec (12 hr) to avoid some processes + # being much slower than others causing a timeout (which can happen in relation + # or LVIS class mAP evaluation). + timeout = 43200 + return dist.new_group( + backend="gloo", + timeout=datetime.timedelta(seconds=timeout), + ) + + return dist.group.WORLD + + +def is_main_process(): + """Return true if the current process is the main one""" + return get_rank() == 0 + + +def all_gather_via_filesys(data, filesys_save_dir=None, gather_to_rank_0_only=False): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors), similar to + `all_gather` above, but using filesystem instead of collective ops. + + If gather_to_rank_0_only is True, only rank 0 will load the gathered object list + (and other ranks will have an empty list). + """ + world_size = get_world_size() + if world_size == 1: + return [data] + + print("gathering via files") + cpu_group = _get_global_gloo_group() + + # if unspecified, we will save to the current python file dir + if filesys_save_dir is not None: + save_dir = filesys_save_dir + elif "EXP_DIR" in os.environ: + save_dir = os.environ["EXP_DIR"] + else: + # try the same directory where the code is stored + save_dir = filesys_save_dir or os.path.dirname(__file__) + save_dir = os.path.join(save_dir, "all_gather_via_filesys") + if is_main_process(): + os.makedirs(save_dir, exist_ok=True) + + # use a timestamp and salt to distinguish different all_gather + timestamp = int(time.time()) if is_main_process() else 0 + salt = random.randint(0, 2**31 - 1) if is_main_process() else 0 + # broadcast the timestamp and salt across ranks + # (all-reduce will do the broadcasting since only rank 0 is non-zero) + timestamp_and_salt = torch.tensor([timestamp, salt], dtype=torch.long) + dist.all_reduce(timestamp_and_salt, group=cpu_group) + timestamp, salt = timestamp_and_salt.tolist() + + # save the data to a file on the disk + rank_save = get_rank() + save_data_filename = f"data_to_gather_{timestamp}_{salt}_{rank_save}.pkl" + save_data_path = os.path.join(save_dir, save_data_filename) + assert not os.path.exists(save_data_path), f"{save_data_path} already exists" + torch.save(data, save_data_path) + dist.barrier(group=cpu_group) + + # read the data from the files + data_list = [] + if rank_save == 0 or not gather_to_rank_0_only: + for rank_load in range(world_size): + load_data_filename = f"data_to_gather_{timestamp}_{salt}_{rank_load}.pkl" + load_data_path = os.path.join(save_dir, load_data_filename) + assert os.path.exists(load_data_path), f"cannot read {save_data_path}" + data_list.append(torch.load(load_data_path, weights_only=False)) + dist.barrier(group=cpu_group) + + # delete the saved file + os.remove(save_data_path) + return data_list + + +def all_gather(data, force_cpu=False, force_filesys=False, filesys_save_dir=None): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + + world_size = get_world_size() + if world_size == 1: + return [data] + + if os.getenv("MDETR_FILESYS_REDUCE_RANK_0_ONLY") == "1": + return all_gather_via_filesys( + data, filesys_save_dir, gather_to_rank_0_only=True + ) + + if os.getenv("MDETR_FILESYS_REDUCE") == "1" or force_filesys: + return all_gather_via_filesys(data, filesys_save_dir) + + cpu_group = None + if os.getenv("MDETR_CPU_REDUCE") == "1" or force_cpu: + cpu_group = _get_global_gloo_group() + + buffer = io.BytesIO() + torch.save(data, buffer) + data_view = buffer.getbuffer() + device = "cuda" if cpu_group is None else "cpu" + tensor = torch.ByteTensor(data_view).to(device) + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device=device, dtype=torch.long) + size_list = [ + torch.tensor([0], device=device, dtype=torch.long) for _ in range(world_size) + ] + if cpu_group is None: + dist.all_gather(size_list, local_size) + else: + print("gathering on cpu") + dist.all_gather(size_list, local_size, group=cpu_group) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + assert isinstance(local_size.item(), int) + local_size = int(local_size.item()) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=device)) + if local_size != max_size: + padding = torch.empty( + size=(max_size - local_size,), dtype=torch.uint8, device=device + ) + tensor = torch.cat((tensor, padding), dim=0) + if cpu_group is None: + dist.all_gather(tensor_list, tensor) + else: + dist.all_gather(tensor_list, tensor, group=cpu_group) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + tensor = torch.split(tensor, [size, max_size - size], dim=0)[0] + buffer = io.BytesIO(tensor.cpu().numpy()) + obj = torch.load(buffer, weights_only=False) + data_list.append(obj) + + return data_list + + +def convert_to_distributed_tensor(tensor: torch.Tensor) -> Tuple[torch.Tensor, str]: + """ + For some backends, such as NCCL, communication only works if the + tensor is on the GPU. This helper function converts to the correct + device and returns the tensor + original device. + """ + orig_device = "cpu" if not tensor.is_cuda else "gpu" + if ( + torch.distributed.is_available() + and torch.distributed.get_backend() == torch.distributed.Backend.NCCL + and not tensor.is_cuda + ): + tensor = tensor.cuda() + return (tensor, orig_device) + + +def convert_to_normal_tensor(tensor: torch.Tensor, orig_device: str) -> torch.Tensor: + """ + For some backends, such as NCCL, communication only works if the + tensor is on the GPU. This converts the tensor back to original device. + """ + if tensor.is_cuda and orig_device == "cpu": + tensor = tensor.cpu() + return tensor + + +def is_distributed_training_run() -> bool: + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and (torch.distributed.get_world_size() > 1) + ) + + +def is_primary() -> bool: + """ + Returns True if this is rank 0 of a distributed training job OR if it is + a single trainer job. Otherwise False. + """ + return get_rank() == _PRIMARY_RANK + + +def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor: + """ + Wrapper over torch.distributed.all_reduce for performing mean reduction + of tensor over all processes. + """ + return all_reduce_op( + tensor, + torch.distributed.ReduceOp.SUM, + lambda t: t / torch.distributed.get_world_size(), + ) + + +def all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor: + """ + Wrapper over torch.distributed.all_reduce for performing sum + reduction of tensor over all processes in both distributed / + non-distributed scenarios. + """ + return all_reduce_op(tensor, torch.distributed.ReduceOp.SUM) + + +def all_reduce_min(tensor: torch.Tensor) -> torch.Tensor: + """ + Wrapper over torch.distributed.all_reduce for performing min + reduction of tensor over all processes in both distributed / + non-distributed scenarios. + """ + return all_reduce_op(tensor, torch.distributed.ReduceOp.MIN) + + +def all_reduce_max(tensor: torch.Tensor) -> torch.Tensor: + """ + Wrapper over torch.distributed.all_reduce for performing min + reduction of tensor over all processes in both distributed / + non-distributed scenarios. + """ + return all_reduce_op(tensor, torch.distributed.ReduceOp.MAX) + + +def all_reduce_op( + tensor: torch.Tensor, + op: torch.distributed.ReduceOp, + after_op_func: Callable[[torch.Tensor], torch.Tensor] = None, +) -> torch.Tensor: + """ + Wrapper over torch.distributed.all_reduce for performing + reduction of tensor over all processes in both distributed / + non-distributed scenarios. + """ + if is_distributed_training_run(): + tensor, orig_device = convert_to_distributed_tensor(tensor) + torch.distributed.all_reduce(tensor, op) + if after_op_func is not None: + tensor = after_op_func(tensor) + tensor = convert_to_normal_tensor(tensor, orig_device) + return tensor + + +def gather_tensors_from_all(tensor: torch.Tensor) -> List[torch.Tensor]: + """ + Wrapper over torch.distributed.all_gather for performing + 'gather' of 'tensor' over all processes in both distributed / + non-distributed scenarios. + """ + if tensor.ndim == 0: + # 0 dim tensors cannot be gathered. so unsqueeze + tensor = tensor.unsqueeze(0) + + if is_distributed_training_run(): + tensor, orig_device = convert_to_distributed_tensor(tensor) + gathered_tensors = [ + torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size()) + ] + torch.distributed.all_gather(gathered_tensors, tensor) + gathered_tensors = [ + convert_to_normal_tensor(_tensor, orig_device) + for _tensor in gathered_tensors + ] + else: + gathered_tensors = [tensor] + + return gathered_tensors + + +def gather_from_all(tensor: torch.Tensor) -> torch.Tensor: + gathered_tensors = gather_tensors_from_all(tensor) + gathered_tensor = torch.cat(gathered_tensors, 0) + return gathered_tensor + + +def broadcast(tensor: torch.Tensor, src: int = 0) -> torch.Tensor: + """ + Wrapper over torch.distributed.broadcast for broadcasting a tensor from the source + to all processes in both distributed / non-distributed scenarios. + """ + if is_distributed_training_run(): + tensor, orig_device = convert_to_distributed_tensor(tensor) + torch.distributed.broadcast(tensor, src) + tensor = convert_to_normal_tensor(tensor, orig_device) + return tensor + + +def barrier() -> None: + """ + Wrapper over torch.distributed.barrier, returns without waiting + if the distributed process group is not initialized instead of throwing error. + """ + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return + torch.distributed.barrier() + + +def get_world_size() -> int: + """ + Simple wrapper for correctly getting worldsize in both distributed + / non-distributed settings + """ + return ( + torch.distributed.get_world_size() + if torch.distributed.is_available() and torch.distributed.is_initialized() + else 1 + ) + + +def get_rank() -> int: + """ + Simple wrapper for correctly getting rank in both distributed + / non-distributed settings + """ + return ( + torch.distributed.get_rank() + if torch.distributed.is_available() and torch.distributed.is_initialized() + else 0 + ) + + +def get_primary_rank() -> int: + return _PRIMARY_RANK + + +def set_cuda_device_index(idx: int) -> None: + global _cuda_device_index + _cuda_device_index = idx + torch.cuda.set_device(_cuda_device_index) + + +def set_cpu_device() -> None: + global _cuda_device_index + _cuda_device_index = _CPU_DEVICE_INDEX + + +def get_cuda_device_index() -> int: + return _cuda_device_index + + +def init_distributed_data_parallel_model( + model: torch.nn.Module, + broadcast_buffers: bool = False, + find_unused_parameters: bool = True, + bucket_cap_mb: int = 25, +) -> torch.nn.parallel.DistributedDataParallel: + global _cuda_device_index + + if _cuda_device_index == _CPU_DEVICE_INDEX: + # CPU-only model, don't specify device + return torch.nn.parallel.DistributedDataParallel( + model, + broadcast_buffers=broadcast_buffers, + find_unused_parameters=find_unused_parameters, + bucket_cap_mb=bucket_cap_mb, + ) + else: + # GPU model + return torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[_cuda_device_index], + output_device=_cuda_device_index, + broadcast_buffers=broadcast_buffers, + find_unused_parameters=find_unused_parameters, + bucket_cap_mb=bucket_cap_mb, + ) + + +def broadcast_object(obj: Any, src: int = _PRIMARY_RANK, use_disk: bool = True) -> Any: + """Broadcast an object from a source to all workers. + + Args: + obj: Object to broadcast, must be serializable + src: Source rank for broadcast (default is primary) + use_disk: If enabled, removes redundant CPU memory copies by writing to + disk + """ + # Either broadcast from primary to the fleet (default), + # or use the src setting as the original rank + if get_rank() == src: + # Emit data + buffer = io.BytesIO() + torch.save(obj, buffer) + data_view = buffer.getbuffer() + length_tensor = torch.LongTensor([len(data_view)]) + length_tensor = broadcast(length_tensor, src=src) + data_tensor = torch.ByteTensor(data_view) + data_tensor = broadcast(data_tensor, src=src) + else: + # Fetch from the source + length_tensor = torch.LongTensor([0]) + length_tensor = broadcast(length_tensor, src=src) + data_tensor = torch.empty([length_tensor.item()], dtype=torch.uint8) + data_tensor = broadcast(data_tensor, src=src) + if use_disk: + with tempfile.TemporaryFile("r+b") as f: + f.write(data_tensor.numpy()) + # remove reference to the data tensor and hope that Python garbage + # collects it + del data_tensor + f.seek(0) + obj = torch.load(f, weights_only=False) + else: + buffer = io.BytesIO(data_tensor.numpy()) + obj = torch.load(buffer, weights_only=False) + return obj + + +def all_gather_tensor(tensor: torch.Tensor, world_size=None): + if world_size is None: + world_size = get_world_size() + # make contiguous because NCCL won't gather the tensor otherwise + assert tensor.is_contiguous(), f"{tensor.shape} is not contiguous!" + tensor, orig_device = convert_to_distributed_tensor(tensor) + tensor_all = [torch.ones_like(tensor) for _ in range(world_size)] + dist.all_gather(tensor_all, tensor, async_op=False) # performance opt + tensor_all = [ + convert_to_normal_tensor(tensor, orig_device) for tensor in tensor_all + ] + return tensor_all + + +def all_gather_batch(tensors: List[torch.Tensor]): + """ + Performs all_gather operation on the provided tensors. + """ + # Queue the gathered tensors + world_size = get_world_size() + # There is no need for reduction in the single-proc case + if world_size == 1: + return tensors + tensor_list = [] + output_tensor = [] + for tensor in tensors: + tensor_all = all_gather_tensor(tensor, world_size) + tensor_list.append(tensor_all) + + for tensor_all in tensor_list: + output_tensor.append(torch.cat(tensor_all, dim=0)) + return output_tensor + + +class GatherLayer(autograd.Function): + """ + Gather tensors from all workers with support for backward propagation: + This implementation does not cut the gradients as torch.distributed.all_gather does. + """ + + @staticmethod + def forward(ctx, x): + output = [torch.zeros_like(x) for _ in range(dist.get_world_size())] + dist.all_gather(output, x) + return tuple(output) + + @staticmethod + def backward(ctx, *grads): + all_gradients = torch.stack(grads) + dist.all_reduce(all_gradients) + return all_gradients[dist.get_rank()] + + +def all_gather_batch_with_grad(tensors): + """ + Performs all_gather operation on the provided tensors. + Graph remains connected for backward grad computation. + """ + # Queue the gathered tensors + world_size = get_world_size() + # There is no need for reduction in the single-proc case + if world_size == 1: + return tensors + tensor_list = [] + output_tensor = [] + + for tensor in tensors: + tensor_all = GatherLayer.apply(tensor) + tensor_list.append(tensor_all) + + for tensor_all in tensor_list: + output_tensor.append(torch.cat(tensor_all, dim=0)) + return output_tensor + + +def unwrap_ddp_if_wrapped(model): + if isinstance(model, torch.nn.parallel.DistributedDataParallel): + return model.module + return model + + +def create_new_process_group(group_size): + """ + Creates process groups of a gives `group_size` and returns + process group that current GPU participates in. + + `group_size` must divide the total number of GPUs (world_size). + + Modified from + https://github.com/NVIDIA/apex/blob/4e1ae43f7f7ac69113ef426dd15f37123f0a2ed3/apex/parallel/__init__.py#L60 + + Args: + group_size (int): number of GPU's to collaborate for sync bn + """ + + assert group_size > 0 + + world_size = torch.distributed.get_world_size() + if world_size <= 8: + if group_size > world_size: + logging.warning( + f"Requested group size [{group_size}] > world size [{world_size}]. " + "Assuming local debug run and capping it to world size." + ) + group_size = world_size + assert world_size >= group_size + assert world_size % group_size == 0 + + group = None + for group_num in range(world_size // group_size): + group_ids = range(group_num * group_size, (group_num + 1) * group_size) + cur_group = torch.distributed.new_group(ranks=group_ids) + if torch.distributed.get_rank() // group_size == group_num: + group = cur_group + # can not drop out and return here, every process must go through creation of all subgroups + + assert group is not None + return group + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def gather_to_rank_0_via_filesys(data, filesys_save_dir=None): + """ + Gather any picklable data to rank 0 via filesystem, using all_gather_via_filesys. + """ + return all_gather_via_filesys(data, filesys_save_dir, gather_to_rank_0_only=True) diff --git a/src/nn/segearth_ov3/sam3/train/utils/logger.py b/src/nn/segearth_ov3/sam3/train/utils/logger.py new file mode 100644 index 0000000..127f6c8 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/utils/logger.py @@ -0,0 +1,241 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import atexit +import functools +import logging +import sys +import uuid +from typing import Any, Dict, Optional, Union + +from hydra.utils import instantiate + +from iopath.common.file_io import g_pathmgr +from numpy import ndarray + +from sam3.train.utils.train_utils import get_machine_local_and_dist_rank, makedir +from torch import Tensor +from torch.utils.tensorboard import SummaryWriter + +Scalar = Union[Tensor, ndarray, int, float] + + +def make_tensorboard_logger(log_dir: str, **writer_kwargs: Any): + makedir(log_dir) + summary_writer_method = SummaryWriter + return TensorBoardLogger( + path=log_dir, summary_writer_method=summary_writer_method, **writer_kwargs + ) + + +class TensorBoardWriterWrapper: + """ + A wrapper around a SummaryWriter object. + """ + + def __init__( + self, + path: str, + *args: Any, + filename_suffix: str = None, + summary_writer_method: Any = SummaryWriter, + **kwargs: Any, + ) -> None: + """Create a new TensorBoard logger. + On construction, the logger creates a new events file that logs + will be written to. If the environment variable `RANK` is defined, + logger will only log if RANK = 0. + + NOTE: If using the logger with distributed training: + - This logger can call collective operations + - Logs will be written on rank 0 only + - Logger must be constructed synchronously *after* initializing distributed process group. + + Args: + path (str): path to write logs to + *args, **kwargs: Extra arguments to pass to SummaryWriter + """ + self._writer: Optional[SummaryWriter] = None + _, self._rank = get_machine_local_and_dist_rank() + self._path: str = path + if self._rank == 0: + logging.info( + f"TensorBoard SummaryWriter instantiated. Files will be stored in: {path}" + ) + self._writer = summary_writer_method( + log_dir=path, + *args, + filename_suffix=filename_suffix or str(uuid.uuid4()), + **kwargs, + ) + else: + logging.debug( + f"Not logging meters on this host because env RANK: {self._rank} != 0" + ) + atexit.register(self.close) + + @property + def writer(self) -> Optional[SummaryWriter]: + return self._writer + + @property + def path(self) -> str: + return self._path + + def flush(self) -> None: + """Writes pending logs to disk.""" + + if not self._writer: + return + + self._writer.flush() + + def close(self) -> None: + """Close writer, flushing pending logs to disk. + Logs cannot be written after `close` is called. + """ + + if not self._writer: + return + + self._writer.close() + self._writer = None + + +class TensorBoardLogger(TensorBoardWriterWrapper): + """ + A simple logger for TensorBoard. + """ + + def log_dict(self, payload: Dict[str, Scalar], step: int) -> None: + """Add multiple scalar values to TensorBoard. + + Args: + payload (dict): dictionary of tag name and scalar value + step (int, Optional): step value to record + """ + if not self._writer: + return + for k, v in payload.items(): + self.log(k, v, step) + + def log(self, name: str, data: Scalar, step: int) -> None: + """Add scalar data to TensorBoard. + + Args: + name (string): tag name used to group scalars + data (float/int/Tensor): scalar data to log + step (int, optional): step value to record + """ + if not self._writer: + return + self._writer.add_scalar(name, data, global_step=step, new_style=True) + + def log_hparams( + self, hparams: Dict[str, Scalar], meters: Dict[str, Scalar] + ) -> None: + """Add hyperparameter data to TensorBoard. + + Args: + hparams (dict): dictionary of hyperparameter names and corresponding values + meters (dict): dictionary of name of meter and corersponding values + """ + if not self._writer: + return + self._writer.add_hparams(hparams, meters) + + +class Logger: + """ + A logger class that can interface with multiple loggers. It now supports tensorboard only for simplicity, but you can extend it with your own logger. + """ + + def __init__(self, logging_conf): + # allow turning off TensorBoard with "should_log: false" in config + tb_config = logging_conf.tensorboard_writer + tb_should_log = tb_config and tb_config.pop("should_log", True) + self.tb_logger = instantiate(tb_config) if tb_should_log else None + + def log_dict(self, payload: Dict[str, Scalar], step: int) -> None: + if self.tb_logger: + self.tb_logger.log_dict(payload, step) + + def log(self, name: str, data: Scalar, step: int) -> None: + if self.tb_logger: + self.tb_logger.log(name, data, step) + + def log_hparams( + self, hparams: Dict[str, Scalar], meters: Dict[str, Scalar] + ) -> None: + if self.tb_logger: + self.tb_logger.log_hparams(hparams, meters) + + +# cache the opened file object, so that different calls to `setup_logger` +# with the same file name can safely write to the same file. +@functools.lru_cache(maxsize=None) +def _cached_log_stream(filename): + # we tune the buffering value so that the logs are updated + # frequently. + log_buffer_kb = 10 * 1024 # 10KB + io = g_pathmgr.open(filename, mode="a", buffering=log_buffer_kb) + atexit.register(io.close) + return io + + +def setup_logging( + name, + output_dir=None, + rank=0, + log_level_primary="INFO", + log_level_secondary="ERROR", +): + """ + Setup various logging streams: stdout and file handlers. + For file handlers, we only setup for the master gpu. + """ + # get the filename if we want to log to the file as well + log_filename = None + if output_dir: + makedir(output_dir) + if rank == 0: + log_filename = f"{output_dir}/log.txt" + + logger = logging.getLogger(name) + logger.setLevel(log_level_primary) + + # create formatter + FORMAT = "%(levelname)s %(asctime)s %(filename)s:%(lineno)4d: %(message)s" + formatter = logging.Formatter(FORMAT) + + # Cleanup any existing handlers + for h in logger.handlers: + logger.removeHandler(h) + logger.root.handlers = [] + + # setup the console handler + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + if rank == 0: + console_handler.setLevel(log_level_primary) + else: + console_handler.setLevel(log_level_secondary) + + # we log to file as well if user wants + if log_filename and rank == 0: + file_handler = logging.StreamHandler(_cached_log_stream(log_filename)) + file_handler.setLevel(log_level_primary) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + logging.root = logger + + +def shutdown_logging(): + """ + After training is done, we ensure to shut down all the logger streams. + """ + logging.info("Shutting down loggers...") + handlers = logging.root.handlers + for handler in handlers: + handler.close() diff --git a/src/nn/segearth_ov3/sam3/train/utils/train_utils.py b/src/nn/segearth_ov3/sam3/train/utils/train_utils.py new file mode 100644 index 0000000..43a8078 --- /dev/null +++ b/src/nn/segearth_ov3/sam3/train/utils/train_utils.py @@ -0,0 +1,285 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging +import math +import os +import random +import re +from datetime import timedelta +from typing import Optional + +import hydra + +import numpy as np +import omegaconf +import torch +import torch.distributed as dist +from iopath.common.file_io import g_pathmgr +from omegaconf import OmegaConf + + +def multiply_all(*args): + return np.prod(np.array(args)).item() + + +def collect_dict_keys(config): + """This function recursively iterates through a dataset configuration, and collect all the dict_key that are defined""" + val_keys = [] + # If the this config points to the collate function, then it has a key + if "_target_" in config and re.match(r".*collate_fn.*", config["_target_"]): + val_keys.append(config["dict_key"]) + else: + # Recursively proceed + for v in config.values(): + if isinstance(v, type(config)): + val_keys.extend(collect_dict_keys(v)) + elif isinstance(v, omegaconf.listconfig.ListConfig): + for item in v: + if isinstance(item, type(config)): + val_keys.extend(collect_dict_keys(item)) + return val_keys + + +class Phase: + TRAIN = "train" + VAL = "val" + + +def register_omegaconf_resolvers(): + OmegaConf.register_new_resolver("get_method", hydra.utils.get_method) + OmegaConf.register_new_resolver("get_class", hydra.utils.get_class) + OmegaConf.register_new_resolver("add", lambda x, y: x + y) + OmegaConf.register_new_resolver("times", multiply_all) + OmegaConf.register_new_resolver("divide", lambda x, y: x / y) + OmegaConf.register_new_resolver("pow", lambda x, y: x**y) + OmegaConf.register_new_resolver("subtract", lambda x, y: x - y) + OmegaConf.register_new_resolver("range", lambda x: list(range(x))) + OmegaConf.register_new_resolver("int", lambda x: int(x)) + OmegaConf.register_new_resolver("ceil_int", lambda x: int(math.ceil(x))) + OmegaConf.register_new_resolver("merge", lambda *x: OmegaConf.merge(*x)) + OmegaConf.register_new_resolver("string", lambda x: str(x)) + + +def setup_distributed_backend(backend, timeout_mins): + """ + Initialize torch.distributed and set the CUDA device. + Expects environment variables to be set as per + https://pytorch.org/docs/stable/distributed.html#environment-variable-initialization + along with the environ variable "LOCAL_RANK" which is used to set the CUDA device. + """ + # enable TORCH_NCCL_ASYNC_ERROR_HANDLING to ensure dist nccl ops time out after timeout_mins + # of waiting + os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1" + logging.info(f"Setting up torch.distributed with a timeout of {timeout_mins} mins") + dist.init_process_group(backend=backend, timeout=timedelta(minutes=timeout_mins)) + return dist.get_rank() + + +def get_machine_local_and_dist_rank(): + """ + Get the distributed and local rank of the current gpu. + """ + local_rank = int(os.environ.get("LOCAL_RANK", None)) + distributed_rank = int(os.environ.get("RANK", None)) + assert ( + local_rank is not None and distributed_rank is not None + ), "Please the set the RANK and LOCAL_RANK environment variables." + return local_rank, distributed_rank + + +def print_cfg(cfg): + """ + Supports printing both Hydra DictConfig and also the AttrDict config + """ + logging.info("Training with config:") + logging.info(OmegaConf.to_yaml(cfg)) + + +def set_seeds(seed_value, max_epochs, dist_rank): + """ + Set the python random, numpy and torch seed for each gpu. Also set the CUDA + seeds if the CUDA is available. This ensures deterministic nature of the training. + """ + # Since in the pytorch sampler, we increment the seed by 1 for every epoch. + seed_value = (seed_value + dist_rank) * max_epochs + logging.info(f"MACHINE SEED: {seed_value}") + random.seed(seed_value) + np.random.seed(seed_value) + torch.manual_seed(seed_value) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed_value) + + +def makedir(dir_path): + """ + Create the directory if it does not exist. + """ + is_success = False + try: + if not g_pathmgr.exists(dir_path): + g_pathmgr.mkdirs(dir_path) + is_success = True + except BaseException: + logging.info(f"Error creating directory: {dir_path}") + return is_success + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_amp_type(amp_type: Optional[str] = None): + if amp_type is None: + return None + assert amp_type in ["bfloat16", "float16"], "Invalid Amp type." + if amp_type == "bfloat16": + return torch.bfloat16 + else: + return torch.float16 + + +def log_env_variables(): + env_keys = sorted(list(os.environ.keys())) + st = "" + for k in env_keys: + v = os.environ[k] + st += f"{k}={v}\n" + logging.info("Logging ENV_VARIABLES") + logging.info(st) + + +class AverageMeter: + """Computes and stores the average and current value""" + + def __init__(self, name, device, fmt=":f"): + self.name = name + self.fmt = fmt + self.device = device + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + self._allow_updates = True + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + def __str__(self): + fmtstr = "{name}: {val" + self.fmt + "} ({avg" + self.fmt + "})" + return fmtstr.format(**self.__dict__) + + +class MemMeter: + """Computes and stores the current, avg, and max of peak Mem usage per iteration""" + + def __init__(self, name, device, fmt=":f"): + self.name = name + self.fmt = fmt + self.device = device + self.reset() + + def reset(self): + self.val = 0 # Per iteration max usage + self.avg = 0 # Avg per iteration max usage + self.peak = 0 # Peak usage for lifetime of program + self.sum = 0 + self.count = 0 + self._allow_updates = True + + def update(self, n=1, reset_peak_usage=True): + self.val = torch.cuda.max_memory_allocated() // 1e9 + self.sum += self.val * n + self.count += n + self.avg = self.sum / self.count + self.peak = max(self.peak, self.val) + if reset_peak_usage: + torch.cuda.reset_peak_memory_stats() + + def __str__(self): + fmtstr = ( + "{name}: {val" + + self.fmt + + "} ({avg" + + self.fmt + + "}/{peak" + + self.fmt + + "})" + ) + return fmtstr.format(**self.__dict__) + + +def human_readable_time(time_seconds): + time = int(time_seconds) + minutes, seconds = divmod(time, 60) + hours, minutes = divmod(minutes, 60) + days, hours = divmod(hours, 24) + return f"{days:02}d {hours:02}h {minutes:02}m" + + +class DurationMeter: + def __init__(self, name, device, fmt=":f"): + self.name = name + self.device = device + self.fmt = fmt + self.val = 0 + + def reset(self): + self.val = 0 + + def update(self, val): + self.val = val + + def add(self, val): + self.val += val + + def __str__(self): + return f"{self.name}: {human_readable_time(self.val)}" + + +class ProgressMeter: + def __init__(self, num_batches, meters, real_meters, prefix=""): + self.batch_fmtstr = self._get_batch_fmtstr(num_batches) + self.meters = meters + self.real_meters = real_meters + self.prefix = prefix + + def display(self, batch, enable_print=False): + entries = [self.prefix + self.batch_fmtstr.format(batch)] + entries += [str(meter) for meter in self.meters] + entries += [ + " | ".join( + [ + f"{os.path.join(name, subname)}: {val:.4f}" + for subname, val in meter.compute().items() + ] + ) + for name, meter in self.real_meters.items() + ] + logging.info(" | ".join(entries)) + if enable_print: + print(" | ".join(entries)) + + def _get_batch_fmtstr(self, num_batches): + num_digits = len(str(num_batches // 1)) + fmt = "{:" + str(num_digits) + "d}" + return "[" + fmt + "/" + fmt.format(num_batches) + "]" + + +def get_resume_checkpoint(checkpoint_save_dir): + if not g_pathmgr.isdir(checkpoint_save_dir): + return None + ckpt_file = os.path.join(checkpoint_save_dir, "checkpoint.pt") + if not g_pathmgr.isfile(ckpt_file): + return None + + return ckpt_file diff --git a/src/nn/segearth_ov3/sam3/visualization_utils.py b/src/nn/segearth_ov3/sam3/visualization_utils.py new file mode 100644 index 0000000..8f4112d --- /dev/null +++ b/src/nn/segearth_ov3/sam3/visualization_utils.py @@ -0,0 +1,941 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import json +import os +import subprocess +from pathlib import Path + +import cv2 +import matplotlib.patches as patches +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import pycocotools.mask as mask_utils +import torch +from matplotlib.colors import to_rgb +from PIL import Image +from skimage.color import lab2rgb, rgb2lab +from sklearn.cluster import KMeans +from torchvision.ops import masks_to_boxes +from tqdm import tqdm + + +def generate_colors(n_colors=256, n_samples=5000): + # Step 1: Random RGB samples + np.random.seed(42) + rgb = np.random.rand(n_samples, 3) + # Step 2: Convert to LAB for perceptual uniformity + # print(f"Converting {n_samples} RGB samples to LAB color space...") + lab = rgb2lab(rgb.reshape(1, -1, 3)).reshape(-1, 3) + # print("Conversion to LAB complete.") + # Step 3: k-means clustering in LAB + kmeans = KMeans(n_clusters=n_colors, n_init=10) + # print(f"Fitting KMeans with {n_colors} clusters on {n_samples} samples...") + kmeans.fit(lab) + # print("KMeans fitting complete.") + centers_lab = kmeans.cluster_centers_ + # Step 4: Convert LAB back to RGB + colors_rgb = lab2rgb(centers_lab.reshape(1, -1, 3)).reshape(-1, 3) + colors_rgb = np.clip(colors_rgb, 0, 1) + return colors_rgb + + +COLORS = generate_colors(n_colors=128, n_samples=5000) + + +def show_img_tensor(img_batch, vis_img_idx=0): + MEAN_IMG = np.array([0.485, 0.456, 0.406]) + STD_IMG = np.array([0.229, 0.224, 0.225]) + im_tensor = img_batch[vis_img_idx].detach().cpu() + assert im_tensor.dim() == 3 + im_tensor = im_tensor.numpy().transpose((1, 2, 0)) + im_tensor = (im_tensor * STD_IMG) + MEAN_IMG + im_tensor = np.clip(im_tensor, 0, 1) + plt.imshow(im_tensor) + + +def draw_box_on_image(image, box, color=(0, 255, 0)): + """ + Draws a rectangle on a given PIL image using the provided box coordinates in xywh format. + :param image: PIL.Image - The image on which to draw the rectangle. + :param box: tuple - A tuple (x, y, w, h) representing the top-left corner, width, and height of the rectangle. + :param color: tuple - A tuple (R, G, B) representing the color of the rectangle. Default is red. + :return: PIL.Image - The image with the rectangle drawn on it. + """ + # Ensure the image is in RGB mode + image = image.convert("RGB") + # Unpack the box coordinates + x, y, w, h = box + x, y, w, h = int(x), int(y), int(w), int(h) + # Get the pixel data + pixels = image.load() + # Draw the top and bottom edges + for i in range(x, x + w): + pixels[i, y] = color + pixels[i, y + h - 1] = color + pixels[i, y + 1] = color + pixels[i, y + h] = color + pixels[i, y - 1] = color + pixels[i, y + h - 2] = color + # Draw the left and right edges + for j in range(y, y + h): + pixels[x, j] = color + pixels[x + 1, j] = color + pixels[x - 1, j] = color + pixels[x + w - 1, j] = color + pixels[x + w, j] = color + pixels[x + w - 2, j] = color + return image + + +def plot_bbox( + img_height, + img_width, + box, + box_format="XYXY", + relative_coords=True, + color="r", + linestyle="solid", + text=None, + ax=None, +): + if box_format == "XYXY": + x, y, x2, y2 = box + w = x2 - x + h = y2 - y + elif box_format == "XYWH": + x, y, w, h = box + elif box_format == "CxCyWH": + cx, cy, w, h = box + x = cx - w / 2 + y = cy - h / 2 + else: + raise RuntimeError(f"Invalid box_format {box_format}") + + if relative_coords: + x *= img_width + w *= img_width + y *= img_height + h *= img_height + + if ax is None: + ax = plt.gca() + rect = patches.Rectangle( + (x, y), + w, + h, + linewidth=1.5, + edgecolor=color, + facecolor="none", + linestyle=linestyle, + ) + ax.add_patch(rect) + if text is not None: + facecolor = "w" + ax.text( + x, + y - 5, + text, + color=color, + weight="bold", + fontsize=8, + bbox={"facecolor": facecolor, "alpha": 0.75, "pad": 2}, + ) + + +def plot_mask(mask, color="r", ax=None): + im_h, im_w = mask.shape + mask_img = np.zeros((im_h, im_w, 4), dtype=np.float32) + mask_img[..., :3] = to_rgb(color) + mask_img[..., 3] = mask * 0.5 + # Use the provided ax or the current axis + if ax is None: + ax = plt.gca() + ax.imshow(mask_img) + + +def normalize_bbox(bbox_xywh, img_w, img_h): + # Assumes bbox_xywh is in XYWH format + if isinstance(bbox_xywh, list): + assert ( + len(bbox_xywh) == 4 + ), "bbox_xywh list must have 4 elements. Batching not support except for torch tensors." + normalized_bbox = bbox_xywh.copy() + normalized_bbox[0] /= img_w + normalized_bbox[1] /= img_h + normalized_bbox[2] /= img_w + normalized_bbox[3] /= img_h + else: + assert isinstance( + bbox_xywh, torch.Tensor + ), "Only torch tensors are supported for batching." + normalized_bbox = bbox_xywh.clone() + assert ( + normalized_bbox.size(-1) == 4 + ), "bbox_xywh tensor must have last dimension of size 4." + normalized_bbox[..., 0] /= img_w + normalized_bbox[..., 1] /= img_h + normalized_bbox[..., 2] /= img_w + normalized_bbox[..., 3] /= img_h + return normalized_bbox + + +def visualize_frame_output(frame_idx, video_frames, outputs, figsize=(12, 8)): + plt.figure(figsize=figsize) + plt.title(f"frame {frame_idx}") + img = load_frame(video_frames[frame_idx]) + img_H, img_W, _ = img.shape + plt.imshow(img) + for i in range(len(outputs["out_probs"])): + box_xywh = outputs["out_boxes_xywh"][i] + prob = outputs["out_probs"][i] + obj_id = outputs["out_obj_ids"][i] + binary_mask = outputs["out_binary_masks"][i] + color = COLORS[obj_id % len(COLORS)] + plot_bbox( + img_H, + img_W, + box_xywh, + text=f"(id={obj_id}, {prob=:.2f})", + box_format="XYWH", + color=color, + ) + plot_mask(binary_mask, color=color) + + +def visualize_formatted_frame_output( + frame_idx, + video_frames, + outputs_list, + titles=None, + points_list=None, + points_labels_list=None, + figsize=(12, 8), + title_suffix="", + prompt_info=None, +): + """Visualize up to three sets of segmentation masks on a video frame. + + Args: + frame_idx: Frame index to visualize + image_files: List of image file paths + outputs_list: List of {frame_idx: {obj_id: mask_tensor}} or single dict {obj_id: mask_tensor} + titles: List of titles for each set of outputs_list + points_list: Optional list of point coordinates + points_labels_list: Optional list of point labels + figsize: Figure size tuple + save: Whether to save the visualization to file + output_dir: Base output directory when saving + scenario_name: Scenario name for organizing saved files + title_suffix: Additional title suffix + prompt_info: Dictionary with prompt information (boxes, points, etc.) + """ + # Handle single output dict case + if isinstance(outputs_list, dict) and frame_idx in outputs_list: + # This is a single outputs dict with frame indices as keys + outputs_list = [outputs_list] + elif isinstance(outputs_list, dict) and not any( + isinstance(k, int) for k in outputs_list.keys() + ): + # This is a single frame's outputs {obj_id: mask} + single_frame_outputs = {frame_idx: outputs_list} + outputs_list = [single_frame_outputs] + + num_outputs = len(outputs_list) + if titles is None: + titles = [f"Set {i+1}" for i in range(num_outputs)] + assert ( + len(titles) == num_outputs + ), "length of `titles` should match that of `outputs_list` if not None." + + _, axes = plt.subplots(1, num_outputs, figsize=figsize) + if num_outputs == 1: + axes = [axes] # Make it iterable + + img = load_frame(video_frames[frame_idx]) + img_H, img_W, _ = img.shape + + for idx in range(num_outputs): + ax, outputs_set, ax_title = axes[idx], outputs_list[idx], titles[idx] + ax.set_title(f"Frame {frame_idx} - {ax_title}{title_suffix}") + ax.imshow(img) + + if frame_idx in outputs_set: + _outputs = outputs_set[frame_idx] + else: + print(f"Warning: Frame {frame_idx} not found in outputs_set") + continue + + if prompt_info and frame_idx == 0: # Show prompts on first frame + if "boxes" in prompt_info: + for box in prompt_info["boxes"]: + # box is in [x, y, w, h] normalized format + x, y, w, h = box + plot_bbox( + img_H, + img_W, + [x, y, x + w, y + h], # Convert to XYXY + box_format="XYXY", + relative_coords=True, + color="yellow", + linestyle="dashed", + text="PROMPT BOX", + ax=ax, + ) + + if "points" in prompt_info and "point_labels" in prompt_info: + points = np.array(prompt_info["points"]) + labels = np.array(prompt_info["point_labels"]) + # Convert normalized to pixel coordinates + points_pixel = points * np.array([img_W, img_H]) + + # Draw positive points (green stars) + pos_points = points_pixel[labels == 1] + if len(pos_points) > 0: + ax.scatter( + pos_points[:, 0], + pos_points[:, 1], + color="lime", + marker="*", + s=200, + edgecolor="white", + linewidth=2, + label="Positive Points", + zorder=10, + ) + + # Draw negative points (red stars) + neg_points = points_pixel[labels == 0] + if len(neg_points) > 0: + ax.scatter( + neg_points[:, 0], + neg_points[:, 1], + color="red", + marker="*", + s=200, + edgecolor="white", + linewidth=2, + label="Negative Points", + zorder=10, + ) + + objects_drawn = 0 + for obj_id, binary_mask in _outputs.items(): + mask_sum = ( + binary_mask.sum() + if hasattr(binary_mask, "sum") + else np.sum(binary_mask) + ) + + if mask_sum > 0: # Only draw if mask has content + # Convert to torch tensor if it's not already + if not isinstance(binary_mask, torch.Tensor): + binary_mask = torch.tensor(binary_mask) + + # Find bounding box from mask + if binary_mask.any(): + box_xyxy = masks_to_boxes(binary_mask.unsqueeze(0)).squeeze() + box_xyxy = normalize_bbox(box_xyxy, img_W, img_H) + else: + # Fallback: create a small box at center + box_xyxy = [0.45, 0.45, 0.55, 0.55] + + color = COLORS[obj_id % len(COLORS)] + + plot_bbox( + img_H, + img_W, + box_xyxy, + text=f"(id={obj_id})", + box_format="XYXY", + color=color, + ax=ax, + ) + + # Convert back to numpy for plotting + mask_np = ( + binary_mask.numpy() + if isinstance(binary_mask, torch.Tensor) + else binary_mask + ) + plot_mask(mask_np, color=color, ax=ax) + objects_drawn += 1 + + if objects_drawn == 0: + ax.text( + 0.5, + 0.5, + "No objects detected", + transform=ax.transAxes, + fontsize=16, + ha="center", + va="center", + color="red", + weight="bold", + ) + + # Draw additional points if provided + if points_list is not None and points_list[idx] is not None: + show_points( + points_list[idx], points_labels_list[idx], ax=ax, marker_size=200 + ) + + ax.axis("off") + + plt.tight_layout() + plt.show() + + +def render_masklet_frame(img, outputs, frame_idx=None, alpha=0.5): + """ + Overlays masklets and bounding boxes on a single image frame. + Args: + img: np.ndarray, shape (H, W, 3), uint8 or float32 in [0,255] or [0,1] + outputs: dict with keys: out_boxes_xywh, out_probs, out_obj_ids, out_binary_masks + frame_idx: int or None, for overlaying frame index text + alpha: float, mask overlay alpha + Returns: + overlay: np.ndarray, shape (H, W, 3), uint8 + """ + if img.dtype == np.float32 or img.max() <= 1.0: + img = (img * 255).astype(np.uint8) + img = img[..., :3] # drop alpha if present + height, width = img.shape[:2] + overlay = img.copy() + + for i in range(len(outputs["out_probs"])): + obj_id = outputs["out_obj_ids"][i] + color = COLORS[obj_id % len(COLORS)] + color255 = (color * 255).astype(np.uint8) + mask = outputs["out_binary_masks"][i] + if mask.shape != img.shape[:2]: + mask = cv2.resize( + mask.astype(np.float32), + (img.shape[1], img.shape[0]), + interpolation=cv2.INTER_NEAREST, + ) + mask_bool = mask > 0.5 + for c in range(3): + overlay[..., c][mask_bool] = ( + alpha * color255[c] + (1 - alpha) * overlay[..., c][mask_bool] + ).astype(np.uint8) + + # Draw bounding boxes and text + for i in range(len(outputs["out_probs"])): + box_xywh = outputs["out_boxes_xywh"][i] + obj_id = outputs["out_obj_ids"][i] + prob = outputs["out_probs"][i] + color = COLORS[obj_id % len(COLORS)] + color255 = tuple(int(x * 255) for x in color) + x, y, w, h = box_xywh + x1 = int(x * width) + y1 = int(y * height) + x2 = int((x + w) * width) + y2 = int((y + h) * height) + cv2.rectangle(overlay, (x1, y1), (x2, y2), color255, 2) + if prob is not None: + label = f"id={obj_id}, p={prob:.2f}" + else: + label = f"id={obj_id}" + cv2.putText( + overlay, + label, + (x1, max(y1 - 10, 0)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + color255, + 1, + cv2.LINE_AA, + ) + + # Overlay frame index at the top-left corner + if frame_idx is not None: + cv2.putText( + overlay, + f"Frame {frame_idx}", + (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, + 1.0, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + + return overlay + + +def save_masklet_video(video_frames, outputs, out_path, alpha=0.5, fps=10): + # Each outputs dict has keys: "out_boxes_xywh", "out_probs", "out_obj_ids", "out_binary_masks" + # video_frames: list of video frame data, same length as outputs_list + + # Read first frame to get size + first_img = load_frame(video_frames[0]) + height, width = first_img.shape[:2] + if first_img.dtype == np.float32 or first_img.max() <= 1.0: + first_img = (first_img * 255).astype(np.uint8) + # Use 'mp4v' for best compatibility with VSCode playback (.mp4 files) + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter("temp.mp4", fourcc, fps, (width, height)) + + outputs_list = [ + (video_frames[frame_idx], frame_idx, outputs[frame_idx]) + for frame_idx in sorted(outputs.keys()) + ] + + for frame, frame_idx, frame_outputs in tqdm(outputs_list): + img = load_frame(frame) + overlay = render_masklet_frame( + img, frame_outputs, frame_idx=frame_idx, alpha=alpha + ) + writer.write(cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)) + + writer.release() + + # Re-encode the video for VSCode compatibility using ffmpeg + subprocess.run(["ffmpeg", "-y", "-i", "temp.mp4", out_path]) + print(f"Re-encoded video saved to {out_path}") + + os.remove("temp.mp4") # Clean up temporary file + + +def save_masklet_image(frame, outputs, out_path, alpha=0.5, frame_idx=None): + """ + Save a single image with masklet overlays. + """ + img = load_frame(frame) + overlay = render_masklet_frame(img, outputs, frame_idx=frame_idx, alpha=alpha) + Image.fromarray(overlay).save(out_path) + print(f"Overlay image saved to {out_path}") + + +def prepare_masks_for_visualization(frame_to_output): + # frame_to_obj_masks --> {frame_idx: {'output_probs': np.array, `out_obj_ids`: np.array, `out_binary_masks`: np.array}} + for frame_idx, out in frame_to_output.items(): + _processed_out = {} + for idx, obj_id in enumerate(out["out_obj_ids"].tolist()): + if out["out_binary_masks"][idx].any(): + _processed_out[obj_id] = out["out_binary_masks"][idx] + frame_to_output[frame_idx] = _processed_out + return frame_to_output + + +def convert_coco_to_masklet_format( + annotations, img_info, is_prediction=False, score_threshold=0.5 +): + """ + Convert COCO format annotations to format expected by render_masklet_frame + """ + outputs = { + "out_boxes_xywh": [], + "out_probs": [], + "out_obj_ids": [], + "out_binary_masks": [], + } + + img_h, img_w = img_info["height"], img_info["width"] + + for idx, ann in enumerate(annotations): + # Get bounding box in relative XYWH format + if "bbox" in ann: + bbox = ann["bbox"] + if max(bbox) > 1.0: # Convert absolute to relative coordinates + bbox = [ + bbox[0] / img_w, + bbox[1] / img_h, + bbox[2] / img_w, + bbox[3] / img_h, + ] + else: + mask = mask_utils.decode(ann["segmentation"]) + rows = np.any(mask, axis=1) + cols = np.any(mask, axis=0) + if np.any(rows) and np.any(cols): + rmin, rmax = np.where(rows)[0][[0, -1]] + cmin, cmax = np.where(cols)[0][[0, -1]] + # Convert to relative XYWH + bbox = [ + cmin / img_w, + rmin / img_h, + (cmax - cmin + 1) / img_w, + (rmax - rmin + 1) / img_h, + ] + else: + bbox = [0, 0, 0, 0] + + outputs["out_boxes_xywh"].append(bbox) + + # Get probability/score + if is_prediction: + prob = ann["score"] + else: + prob = 1.0 # GT has no probability + outputs["out_probs"].append(prob) + + outputs["out_obj_ids"].append(idx) + mask = mask_utils.decode(ann["segmentation"]) + mask = (mask > score_threshold).astype(np.uint8) + + outputs["out_binary_masks"].append(mask) + + return outputs + + +def save_side_by_side_visualization(img, gt_anns, pred_anns, noun_phrase): + """ + Create side-by-side visualization of GT and predictions + """ + + # Create side-by-side visualization + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7)) + + main_title = f"Noun phrase: '{noun_phrase}'" + fig.suptitle(main_title, fontsize=16, fontweight="bold") + + gt_overlay = render_masklet_frame(img, gt_anns, alpha=0.5) + ax1.imshow(gt_overlay) + ax1.set_title("Ground Truth", fontsize=14, fontweight="bold") + ax1.axis("off") + + pred_overlay = render_masklet_frame(img, pred_anns, alpha=0.5) + ax2.imshow(pred_overlay) + ax2.set_title("Predictions", fontsize=14, fontweight="bold") + ax2.axis("off") + + plt.subplots_adjust(top=0.88) + plt.tight_layout() + + +def bitget(val, idx): + return (val >> idx) & 1 + + +def pascal_color_map(): + colormap = np.zeros((512, 3), dtype=int) + ind = np.arange(512, dtype=int) + for shift in reversed(list(range(8))): + for channel in range(3): + colormap[:, channel] |= bitget(ind, channel) << shift + ind >>= 3 + + return colormap.astype(np.uint8) + + +def draw_masks_to_frame( + frame: np.ndarray, masks: np.ndarray, colors: np.ndarray +) -> np.ndarray: + masked_frame = frame + for mask, color in zip(masks, colors): + curr_masked_frame = np.where(mask[..., None], color, masked_frame) + masked_frame = cv2.addWeighted(masked_frame, 0.75, curr_masked_frame, 0.25, 0) + + if int(cv2.__version__[0]) > 3: + contours, _ = cv2.findContours( + np.array(mask, dtype=np.uint8).copy(), + cv2.RETR_TREE, + cv2.CHAIN_APPROX_NONE, + ) + else: + _, contours, _ = cv2.findContours( + np.array(mask, dtype=np.uint8).copy(), + cv2.RETR_TREE, + cv2.CHAIN_APPROX_NONE, + ) + + cv2.drawContours( + masked_frame, contours, -1, (255, 255, 255), 7 + ) # White outer contour + cv2.drawContours( + masked_frame, contours, -1, (0, 0, 0), 5 + ) # Black middle contour + cv2.drawContours( + masked_frame, contours, -1, color.tolist(), 3 + ) # Original color inner contour + return masked_frame + + +def get_annot_df(file_path: str): + with open(file_path, "r") as f: + data = json.load(f) + + dfs = {} + + for k, v in data.items(): + if k in ("info", "licenses"): + dfs[k] = v + continue + df = pd.DataFrame(v) + dfs[k] = df + + return dfs + + +def get_annot_dfs(file_list: list[str]): + dfs = {} + for annot_file in tqdm(file_list): + dataset_name = Path(annot_file).stem + dfs[dataset_name] = get_annot_df(annot_file) + + return dfs + + +def get_media_dir(media_dir: str, dataset: str): + if dataset in ["saco_veval_sav_test", "saco_veval_sav_val"]: + return os.path.join(media_dir, "saco_sav", "JPEGImages_24fps") + elif dataset in ["saco_veval_yt1b_test", "saco_veval_yt1b_val"]: + return os.path.join(media_dir, "saco_yt1b", "JPEGImages_6fps") + elif dataset in ["saco_veval_smartglasses_test", "saco_veval_smartglasses_val"]: + return os.path.join(media_dir, "saco_sg", "JPEGImages_6fps") + elif dataset == "sa_fari_test": + return os.path.join(media_dir, "sa_fari", "JPEGImages_6fps") + else: + raise ValueError(f"Dataset {dataset} not found") + + +def get_all_annotations_for_frame( + dataset_df: pd.DataFrame, video_id: int, frame_idx: int, data_dir: str, dataset: str +): + media_dir = os.path.join(data_dir, "media") + + # Load the annotation and video data + annot_df = dataset_df["annotations"] + video_df = dataset_df["videos"] + + # Get the frame + video_df_current = video_df[video_df.id == video_id] + assert ( + len(video_df_current) == 1 + ), f"Expected 1 video row, got {len(video_df_current)}" + video_row = video_df_current.iloc[0] + file_name = video_row.file_names[frame_idx] + file_path = os.path.join( + get_media_dir(media_dir=media_dir, dataset=dataset), file_name + ) + frame = cv2.imread(file_path) + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # Get the masks and noun phrases annotated in this video in this frame + annot_df_current_video = annot_df[annot_df.video_id == video_id] + if len(annot_df_current_video) == 0: + print(f"No annotations found for video_id {video_id}") + return frame, None, None + else: + empty_mask = np.zeros(frame.shape[:2], dtype=np.uint8) + mask_np_pairs = annot_df_current_video.apply( + lambda row: ( + ( + mask_utils.decode(row.segmentations[frame_idx]) + if row.segmentations[frame_idx] + else empty_mask + ), + row.noun_phrase, + ), + axis=1, + ) + # sort based on noun_phrases + mask_np_pairs = sorted(mask_np_pairs, key=lambda x: x[1]) + masks, noun_phrases = zip(*mask_np_pairs) + + return frame, masks, noun_phrases + + +def visualize_prompt_overlay( + frame_idx, + video_frames, + title="Prompt Visualization", + text_prompt=None, + point_prompts=None, + point_labels=None, + bounding_boxes=None, + box_labels=None, + obj_id=None, +): + """Simple prompt visualization function""" + img = Image.fromarray(load_frame(video_frames[frame_idx])) + fig, ax = plt.subplots(1, figsize=(6, 4)) + ax.imshow(img) + + img_w, img_h = img.size + + if text_prompt: + ax.text( + 0.02, + 0.98, + f'Text: "{text_prompt}"', + transform=ax.transAxes, + fontsize=12, + color="white", + weight="bold", + bbox=dict(boxstyle="round,pad=0.3", facecolor="red", alpha=0.7), + verticalalignment="top", + ) + + if point_prompts: + for i, point in enumerate(point_prompts): + x, y = point + # Convert relative to absolute coordinates + x_img, y_img = x * img_w, y * img_h + + # Use different colors for positive/negative points + if point_labels and len(point_labels) > i: + color = "green" if point_labels[i] == 1 else "red" + marker = "o" if point_labels[i] == 1 else "x" + else: + color = "green" + marker = "o" + + ax.plot( + x_img, + y_img, + marker=marker, + color=color, + markersize=10, + markeredgewidth=2, + markeredgecolor="white", + ) + ax.text( + x_img + 5, + y_img - 5, + f"P{i+1}", + color=color, + fontsize=10, + weight="bold", + bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8), + ) + + if bounding_boxes: + for i, box in enumerate(bounding_boxes): + x, y, w, h = box + # Convert relative to absolute coordinates + x_img, y_img = x * img_w, y * img_h + w_img, h_img = w * img_w, h * img_h + + # Use different colors for positive/negative boxes + if box_labels and len(box_labels) > i: + color = "green" if box_labels[i] == 1 else "red" + else: + color = "green" + + rect = patches.Rectangle( + (x_img, y_img), + w_img, + h_img, + linewidth=2, + edgecolor=color, + facecolor="none", + ) + ax.add_patch(rect) + ax.text( + x_img, + y_img - 5, + f"B{i+1}", + color=color, + fontsize=10, + weight="bold", + bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8), + ) + + # Add object ID info if provided + if obj_id is not None: + ax.text( + 0.02, + 0.02, + f"Object ID: {obj_id}", + transform=ax.transAxes, + fontsize=10, + color="white", + weight="bold", + bbox=dict(boxstyle="round,pad=0.3", facecolor="blue", alpha=0.7), + verticalalignment="bottom", + ) + + ax.set_title(title) + ax.axis("off") + plt.tight_layout() + plt.show() + + +def plot_results(img, results): + plt.figure(figsize=(12, 8)) + plt.imshow(img) + nb_objects = len(results["scores"]) + print(f"found {nb_objects} object(s)") + for i in range(nb_objects): + color = COLORS[i % len(COLORS)] + plot_mask(results["masks"][i].squeeze(0).cpu(), color=color) + w, h = img.size + prob = results["scores"][i].item() + # plot_bbox( + # h, + # w, + # results["boxes"][i].cpu(), + # text=f"(id={i}, {prob=:.2f})", + # box_format="XYXY", + # color=color, + # relative_coords=False, + # ) + plt.savefig('./figs/test.png') + +def single_visualization(img, anns, title): + """ + Create a single image visualization with overlays. + """ + fig, ax = plt.subplots(figsize=(7, 7)) + fig.suptitle(title, fontsize=16, fontweight="bold") + overlay = render_masklet_frame(img, anns, alpha=0.5) + ax.imshow(overlay) + ax.axis("off") + plt.tight_layout() + + +def show_mask(mask, ax, obj_id=None, random_color=False): + if random_color: + color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) + else: + cmap = plt.get_cmap("tab10") + cmap_idx = 0 if obj_id is None else obj_id + color = np.array([*cmap(cmap_idx)[:3], 0.6]) + h, w = mask.shape[-2:] + mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) + ax.imshow(mask_image) + + +def show_box(box, ax): + x0, y0 = box[0], box[1] + w, h = box[2] - box[0], box[3] - box[1] + ax.add_patch( + plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2) + ) + + +def show_points(coords, labels, ax, marker_size=375): + pos_points = coords[labels == 1] + neg_points = coords[labels == 0] + ax.scatter( + pos_points[:, 0], + pos_points[:, 1], + color="green", + marker="*", + s=marker_size, + edgecolor="white", + linewidth=1.25, + ) + ax.scatter( + neg_points[:, 0], + neg_points[:, 1], + color="red", + marker="*", + s=marker_size, + edgecolor="white", + linewidth=1.25, + ) + + +def load_frame(frame): + if isinstance(frame, np.ndarray): + img = frame + elif isinstance(frame, Image.Image): + img = np.array(frame) + elif isinstance(frame, str) and os.path.isfile(frame): + img = plt.imread(frame) + else: + raise ValueError(f"Invalid video frame type: {type(frame)=}") + return img diff --git a/src/nn/segearth_ov3/segearthov3_segmentor.py b/src/nn/segearth_ov3/segearthov3_segmentor.py new file mode 100644 index 0000000..e37a185 --- /dev/null +++ b/src/nn/segearth_ov3/segearthov3_segmentor.py @@ -0,0 +1,417 @@ +"""SegEarth-OV3 segmentor — standalone, no mmseg/mmcv/mmengine dependencies.""" + +import os +import torch +from torch import nn +import torch.nn.functional as F +from dataclasses import dataclass, field +from typing import Any +from PIL import Image + +from sam3 import build_sam3_image_model +from sam3.model.sam3_image_processor import Sam3Processor + +_DEFAULT_BPE_PATH = os.path.join( + os.path.dirname(__file__), "sam3", "assets", "bpe_simple_vocab_16e6.txt.gz", +) + + +# --------------------------------------------------------------------------- +# Lightweight replacements for mmengine/mmseg data structures +# --------------------------------------------------------------------------- + +@dataclass +class PixelData: + """Minimal replacement for mmengine.structures.PixelData.""" + data: torch.Tensor = None + + +@dataclass +class SegDataSample: + """Minimal replacement for mmseg.structures.SegDataSample.""" + metainfo: dict = field(default_factory=dict) + pred_sem_seg: PixelData = field(default_factory=PixelData) + seg_logits: PixelData = field(default_factory=PixelData) + + def set_metainfo(self, meta: dict) -> None: + self.metainfo.update(meta) + + def set_data(self, data: dict) -> None: + for k, v in data.items(): + setattr(self, k, v) + + +# --------------------------------------------------------------------------- +# Segmentor +# --------------------------------------------------------------------------- + +class SegEarthOV3Segmentation(nn.Module): + def __init__(self, classname_path, + device=torch.device('cuda'), + prob_thd=0.0, + bg_idx=0, + slide_stride=0, + slide_crop=0, + confidence_threshold=0.5, + use_sem_seg=True, + use_presence_score=True, + use_transformer_decoder=True, + bpe_path=None, + checkpoint_path="weights/sam3/sam3.pt", + **kwargs): + super().__init__() + + self.device = device + # Initialize SAM3 model + if bpe_path is None: + bpe_path = _DEFAULT_BPE_PATH + model = build_sam3_image_model( + bpe_path=bpe_path, + checkpoint_path=checkpoint_path, + device=str(device), + ) + self.processor = Sam3Processor(model, confidence_threshold=confidence_threshold, device=device) + self.query_words, self.query_idx = get_cls_idx(classname_path) + self.num_cls = max(self.query_idx) + 1 + self.num_queries = len(self.query_idx) + self.query_idx = torch.Tensor(self.query_idx).to(torch.int64).to(device) + + self.prob_thd = prob_thd + self.bg_idx = bg_idx + self.slide_stride = slide_stride + self.slide_crop = slide_crop + self.confidence_threshold = confidence_threshold + self.use_sem_seg = use_sem_seg + self.use_presence_score = use_presence_score + self.use_transformer_decoder = use_transformer_decoder + + def _inference_single_view(self, image): + """Inference on a single PIL image or crop patch.""" + w, h = image.size + seg_logits = torch.zeros((self.num_queries, h, w), device=self.device) + + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + inference_state = self.processor.set_image(image) + + for query_idx, query_word in enumerate(self.query_words): + self.processor.reset_all_prompts(inference_state) + inference_state = self.processor.set_text_prompt(state=inference_state, prompt=query_word) + + if self.use_transformer_decoder: + if inference_state['masks_logits'].shape[0] > 0: + inst_len = inference_state['masks_logits'].shape[0] + for inst_id in range(inst_len): + instance_logits = inference_state['masks_logits'][inst_id].squeeze() + instance_score = inference_state['object_score'][inst_id] + + if instance_logits.shape != (h, w): + instance_logits = F.interpolate( + instance_logits.view(1, 1, *instance_logits.shape), + size=(h, w), + mode='bilinear', + align_corners=False + ).squeeze() + + seg_logits[query_idx] = torch.max(seg_logits[query_idx], instance_logits * instance_score) + + if self.use_sem_seg: + semantic_logits = inference_state['semantic_mask_logits'] + if semantic_logits.shape != (h, w): + semantic_logits = F.interpolate( + semantic_logits, + size=(h, w), + mode='bilinear', + align_corners=False + ).squeeze() + + seg_logits[query_idx] = torch.max(seg_logits[query_idx], semantic_logits) + + if self.use_presence_score: + seg_logits[query_idx] = seg_logits[query_idx] * inference_state["presence_score"] + + return seg_logits + + def slide_inference(self, image, stride, crop_size): + """Inference by sliding-window with overlap using PIL cropping.""" + w_img, h_img = image.size + + if isinstance(stride, int): + stride = (stride, stride) + if isinstance(crop_size, int): + crop_size = (crop_size, crop_size) + + h_stride, w_stride = stride + h_crop, w_crop = crop_size + + preds = torch.zeros((self.num_queries, h_img, w_img), device=self.device) + count_mat = torch.zeros((1, h_img, w_img), device=self.device) + + h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1 + w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1 + + for h_idx in range(h_grids): + for w_idx in range(w_grids): + y1 = h_idx * h_stride + x1 = w_idx * w_stride + y2 = min(y1 + h_crop, h_img) + x2 = min(x1 + w_crop, w_img) + + y1 = max(y2 - h_crop, 0) + x1 = max(x2 - w_crop, 0) + + crop_img = image.crop((x1, y1, x2, y2)) + crop_seg_logit = self._inference_single_view(crop_img) + + preds[:, y1:y2, x1:x2] += crop_seg_logit + count_mat[:, y1:y2, x1:x2] += 1 + + assert (count_mat == 0).sum() == 0, "Error: Sparse sliding window coverage." + + preds = preds / count_mat + return preds + + def predict(self, inputs, data_samples): + """Run segmentation prediction. + + Args: + inputs: [B, 3, H, W] tensor (placeholder — actual image read from img_path). + data_samples: list[SegDataSample] with metainfo containing + 'img_path' and 'ori_shape'. + + Returns: + list[SegDataSample] with pred_sem_seg and seg_logits filled. + """ + if data_samples is not None: + batch_img_metas = [ds.metainfo for ds in data_samples] + else: + batch_img_metas = [ + dict( + ori_shape=inputs.shape[2:], + img_shape=inputs.shape[2:], + pad_shape=inputs.shape[2:], + padding_size=[0, 0, 0, 0]) + ] * inputs.shape[0] + + for i, meta in enumerate(batch_img_metas): + image_path = meta.get('img_path') + image = Image.open(image_path).convert('RGB') + ori_shape = meta['ori_shape'] + + if self.slide_crop > 0 and (self.slide_crop < image.size[0] or self.slide_crop < image.size[1]): + seg_logits = self.slide_inference(image, self.slide_stride, self.slide_crop) + else: + seg_logits = self._inference_single_view(image) + + if seg_logits.shape[-2:] != ori_shape: + seg_logits = F.interpolate( + seg_logits.unsqueeze(0), + size=ori_shape, + mode='bilinear', + align_corners=False + ).squeeze(0) + + # Post-processing: merge synonym queries into classes. + if self.num_cls != self.num_queries: + seg_logits = seg_logits.unsqueeze(0) + cls_index = nn.functional.one_hot(self.query_idx) + cls_index = cls_index.T.view(self.num_cls, len(self.query_idx), 1, 1) + seg_logits = (seg_logits * cls_index).max(1)[0] + + seg_pred = torch.argmax(seg_logits, dim=0) + + # Apply probability threshold. + max_vals = seg_logits.max(0)[0] + seg_pred[max_vals < self.prob_thd] = self.bg_idx + + data_samples[i].set_data({ + 'seg_logits': PixelData(data=seg_logits), + 'pred_sem_seg': PixelData(data=seg_pred.unsqueeze(0)), + }) + + return data_samples + + @staticmethod + def _slice_backbone_out(backbone_out, idx): + """Slice batched backbone output to extract single-image features.""" + sliced = {} + for key, val in backbone_out.items(): + if key == "sam2_backbone_out": + sliced[key] = SegEarthOV3Segmentation._slice_backbone_out(val, idx) if val is not None else None + elif isinstance(val, torch.Tensor): + sliced[key] = val[idx:idx + 1] + elif isinstance(val, list): + sliced[key] = [t[idx:idx + 1] if isinstance(t, torch.Tensor) else t for t in val] + else: + sliced[key] = val + return sliced + + 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 = [] + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for query_word in self.query_words: + text_out = self.processor.model.backbone.forward_text( + [query_word], device=self.device, + ) + # Deep-clone tensors so they survive state mutations. + cached = {k: v.clone() if isinstance(v, torch.Tensor) else v + for k, v in text_out.items()} + self._text_cache.append(cached) + + def _apply_cached_text(self, state, query_idx): + """Inject cached text embeddings and run grounding decoder.""" + cached = self._text_cache[query_idx] + state["backbone_out"].update( + {k: v.clone() if isinstance(v, torch.Tensor) else v + for k, v in cached.items()} + ) + if "geometric_prompt" not in state: + state["geometric_prompt"] = self.processor.model._get_dummy_prompt() + return self.processor._forward_grounding(state) + + def predict_pil_batch(self, images: list[Image.Image], + ori_shapes: list[tuple[int, int]] | None = None): + """Batch prediction: single backbone pass for all images, then per-image grounding. + + Args: + images: list of PIL RGB images. + ori_shapes: list of (H, W) target output shapes. Defaults to each image's size. + + Returns: + list of (seg_pred, seg_logits) tuples. + """ + B = len(images) + if ori_shapes is None: + ori_shapes = [(img.size[1], img.size[0]) for img in images] + + # Cache text embeddings once (amortized across all batches). + self._cache_text_embeddings() + + results = [] + + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + # Single backbone pass for all images + batch_state = self.processor.set_image_batch(images) + + for i in range(B): + w, h = images[i].size + seg_logits = torch.zeros((self.num_queries, h, w), device=self.device) + + # Create single-image state by slicing backbone features + state = { + "original_height": batch_state["original_heights"][i], + "original_width": batch_state["original_widths"][i], + "backbone_out": self._slice_backbone_out(batch_state["backbone_out"], i), + } + + for query_idx in range(len(self.query_words)): + self.processor.reset_all_prompts(state) + state = self._apply_cached_text(state, query_idx) + + if self.use_transformer_decoder: + if state['masks_logits'].shape[0] > 0: + inst_len = state['masks_logits'].shape[0] + for inst_id in range(inst_len): + instance_logits = state['masks_logits'][inst_id].squeeze() + instance_score = state['object_score'][inst_id] + + if instance_logits.shape != (h, w): + instance_logits = F.interpolate( + instance_logits.view(1, 1, *instance_logits.shape), + size=(h, w), + mode='bilinear', + align_corners=False + ).squeeze() + + seg_logits[query_idx] = torch.max(seg_logits[query_idx], instance_logits * instance_score) + + if self.use_sem_seg: + semantic_logits = state['semantic_mask_logits'] + if semantic_logits.shape != (h, w): + semantic_logits = F.interpolate( + semantic_logits, + size=(h, w), + mode='bilinear', + align_corners=False + ).squeeze() + + seg_logits[query_idx] = torch.max(seg_logits[query_idx], semantic_logits) + + if self.use_presence_score: + seg_logits[query_idx] = seg_logits[query_idx] * state["presence_score"] + + # Post-processing + ori_shape = ori_shapes[i] + if seg_logits.shape[-2:] != ori_shape: + seg_logits = F.interpolate( + seg_logits.unsqueeze(0), size=ori_shape, + mode='bilinear', align_corners=False, + ).squeeze(0) + + if self.num_cls != self.num_queries: + seg_logits = seg_logits.unsqueeze(0) + cls_index = nn.functional.one_hot(self.query_idx) + cls_index = cls_index.T.view(self.num_cls, len(self.query_idx), 1, 1) + seg_logits = (seg_logits * cls_index).max(1)[0] + + seg_pred = torch.argmax(seg_logits, dim=0) + max_vals = seg_logits.max(0)[0] + seg_pred[max_vals < self.prob_thd] = self.bg_idx + + results.append((seg_pred.to(torch.uint8), seg_logits)) + + return results + + def predict_pil(self, image: Image.Image, ori_shape: tuple[int, int] | None = None): + """Convenience method: predict from a PIL image directly (no file path needed). + + Args: + image: PIL RGB image. + ori_shape: (H, W) target output shape. Defaults to image size. + + Returns: + seg_pred: [H, W] uint8 tensor of class indices. + seg_logits: [num_cls, H, W] float tensor. + """ + if ori_shape is None: + ori_shape = (image.size[1], image.size[0]) # PIL is (W, H) + + if self.slide_crop > 0 and (self.slide_crop < image.size[0] or self.slide_crop < image.size[1]): + seg_logits = self.slide_inference(image, self.slide_stride, self.slide_crop) + else: + seg_logits = self._inference_single_view(image) + + if seg_logits.shape[-2:] != ori_shape: + seg_logits = F.interpolate( + seg_logits.unsqueeze(0), size=ori_shape, + mode='bilinear', align_corners=False, + ).squeeze(0) + + if self.num_cls != self.num_queries: + seg_logits = seg_logits.unsqueeze(0) + cls_index = nn.functional.one_hot(self.query_idx) + cls_index = cls_index.T.view(self.num_cls, len(self.query_idx), 1, 1) + seg_logits = (seg_logits * cls_index).max(1)[0] + + seg_pred = torch.argmax(seg_logits, dim=0) + max_vals = seg_logits.max(0)[0] + seg_pred[max_vals < self.prob_thd] = self.bg_idx + + return seg_pred.to(torch.uint8), seg_logits + + +def get_cls_idx(path): + with open(path, 'r') as f: + name_sets = f.readlines() + num_cls = len(name_sets) + + class_names, class_indices = [], [] + for idx in range(num_cls): + names_i = name_sets[idx].split(',') + names_i = [i.strip() for i in names_i] + class_names += names_i + class_indices += [idx for _ in range(len(names_i))] + class_names = [item.replace('\n', '') for item in class_names] + return class_names, class_indices diff --git a/src/scripts/move_query_aux.py b/src/scripts/move_query_aux.py new file mode 100644 index 0000000..506fb57 --- /dev/null +++ b/src/scripts/move_query_aux.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Move existing query aux files (256x256) from World-UAV-aug to a backup folder. + +Preserves the full directory structure so files can be restored if needed. +Uses os.rename (same filesystem) — no data copying, only inode renames. + +Usage:: + + # Dry run (default) — only counts, moves nothing: + python -m src.scripts.move_query_aux + + # Actually move: + python -m src.scripts.move_query_aux --execute + + # Restore (move back): + python -m src.scripts.move_query_aux --restore --execute +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +from pathlib import Path + +logger = logging.getLogger(__name__) + +PARTITIONS = ("Country", "Terrain", "Rot") +AUX_SUFFIXES = ("_depth.png", "_edge.png", "_segm.png", "_chm.png") + + +def _is_query_aux(p: Path) -> bool: + """Check if file is a query aux map (inside a query/ subtree).""" + if p.suffix.lower() != ".png": + return False + if "query" not in p.parts: + return False + name = p.name + return any(name.endswith(s) for s in AUX_SUFFIXES) + + +def move_query_aux( + aux_root: Path, + backup_root: Path, + execute: bool = False, + restore: bool = False, +) -> None: + """Move query aux files between aux_root and backup_root.""" + + if restore: + src_root, dst_root = backup_root, aux_root + action = "RESTORE" + else: + src_root, dst_root = aux_root, backup_root + action = "BACKUP" + + if not src_root.exists(): + logger.error("Source root does not exist: %s", src_root) + return + + logger.info("=== %s query aux files ===", action) + logger.info("From: %s", src_root) + logger.info("To: %s", dst_root) + logger.info("Mode: %s", "EXECUTE" if execute else "DRY RUN") + + moved = 0 + skipped = 0 + total_bytes = 0 + dirs_created: set[str] = set() + + for partition in PARTITIONS: + pdir = src_root / partition + if not pdir.exists(): + continue + for p in sorted(pdir.rglob("*")): + if not p.is_file(): + continue + if not _is_query_aux(p): + continue + + rel = p.relative_to(src_root) + dst = dst_root / rel + + if dst.exists(): + skipped += 1 + continue + + total_bytes += p.stat().st_size + moved += 1 + + if execute: + dst_dir = str(dst.parent) + if dst_dir not in dirs_created: + dst.parent.mkdir(parents=True, exist_ok=True) + dirs_created.add(dst_dir) + os.rename(p, dst) + + if moved % 100_000 == 0: + logger.info(" ... %d files processed (%.1f GB)", moved, total_bytes / 1e9) + + logger.info("=== Done ===") + logger.info(" Files %s: %d (%.1f GB)", "moved" if execute else "to move", moved, total_bytes / 1e9) + if skipped: + logger.info(" Skipped (already exists at destination): %d", skipped) + if not execute: + logger.info(" Run with --execute to actually move files.") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Move query aux files to/from backup") + parser.add_argument("--aux-root", type=Path, + default=Path("/mnt/data1tb/cvgl_datasets/World-UAV-aug")) + parser.add_argument("--backup-root", type=Path, + default=Path("/mnt/data1tb/cvgl_datasets/World-UAV-aug-query-256-backup")) + parser.add_argument("--execute", action="store_true", + help="Actually move files (default is dry run)") + parser.add_argument("--restore", action="store_true", + help="Restore files from backup back to aux_root") + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)-7s %(message)s", + datefmt="%H:%M:%S", + ) + move_query_aux(args.aux_root, args.backup_root, args.execute, args.restore) + + +if __name__ == "__main__": + main() diff --git a/src/tests/__init__.py b/src/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/tests/conftest.py b/src/tests/conftest.py new file mode 100644 index 0000000..94b2a4b --- /dev/null +++ b/src/tests/conftest.py @@ -0,0 +1,111 @@ +"""Shared fixtures for the augmentation pipeline tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image + +# Ensure src package is importable. +_SRC_ROOT = Path(__file__).resolve().parent.parent +if str(_SRC_ROOT) not in sys.path: + sys.path.insert(0, str(_SRC_ROOT)) + +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.dataset import ImageRecord + + +# --------------------------------------------------------------------------- +# Config fixtures (no gin, plain defaults) +# --------------------------------------------------------------------------- + +@pytest.fixture() +def hw_conf() -> HardwareConfig: + import gin + gin.clear_config() + return HardwareConfig( + profile_name="test", total_ram_gb=8.0, reserve_gb=1.0, + use_fp16=False, batch_size=2, num_workers=0, + ) + + +@pytest.fixture() +def input_conf() -> InputConfig: + import gin + gin.clear_config() + return InputConfig(image_size=64) + + +@pytest.fixture() +def seg_conf() -> SegConfig: + import gin + gin.clear_config() + return SegConfig(prompts=["background", "building", "road"]) + + +@pytest.fixture() +def models_conf() -> ModelsConfig: + import gin + gin.clear_config() + return ModelsConfig() + + +@pytest.fixture() +def pipeline_conf(tmp_path: Path) -> PipelineConfig: + import gin + gin.clear_config() + return PipelineConfig( + input_root=str(tmp_path / "input"), + output_root=str(tmp_path / "output"), + stages=["depth", "edges", "segmentation", "chmv2"], + save_npy=True, + save_vis=True, + save_concat=False, + resume=False, + log_level="WARNING", + ) + + +# --------------------------------------------------------------------------- +# Synthetic data fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture() +def fake_image_dir(tmp_path: Path) -> Path: + """Create a minimal World-UAV-like directory with 4 dummy images.""" + img_dir = tmp_path / "input" / "scene01" / "query" + img_dir.mkdir(parents=True) + for i in range(4): + img = Image.fromarray( + np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8), + ) + img.save(img_dir / f"img_{i:03d}.png") + return tmp_path / "input" + + +@pytest.fixture() +def sample_batch() -> torch.Tensor: + """Random [B, 3, H, W] float32 [0, 1] batch.""" + return torch.rand(2, 3, 64, 64) + + +@pytest.fixture() +def sample_depth() -> torch.Tensor: + """Random [B, 1, H, W] float32 [0, 1] depth batch.""" + return torch.rand(2, 1, 64, 64) + + +@pytest.fixture() +def sample_records(fake_image_dir: Path, tmp_path: Path) -> list[ImageRecord]: + """ImageRecord list matching fake_image_dir.""" + from src.augmentor.dataset import discover_images, attach_output_dirs + records = discover_images(fake_image_dir) + return attach_output_dirs(records, tmp_path / "output") diff --git a/src/tests/test_config.py b/src/tests/test_config.py new file mode 100644 index 0000000..b2281a1 --- /dev/null +++ b/src/tests/test_config.py @@ -0,0 +1,60 @@ +"""Tests for configuration dataclasses and auto_batch_size.""" + +from __future__ import annotations + +from src.conf.hardware_conf import HardwareConfig +from src.conf.pipeline_conf import PipelineConfig +from src.conf.seg_conf import SegConfig +from src.conf.input_conf import InputConfig + + +class TestHardwareConfig: + def test_defaults(self, hw_conf: HardwareConfig) -> None: + assert hw_conf.available_gb == 7.0 + assert hw_conf.batch_size == 2 + assert hw_conf.use_fp16 is False + + def test_auto_batch_size_returns_manual_when_set(self, hw_conf: HardwareConfig) -> None: + bs = hw_conf.auto_batch_size(act_per_sample_mb=120) + assert bs == 2 # manual override + + def test_auto_batch_size_computes_when_none(self) -> None: + import gin + gin.clear_config() + cfg = HardwareConfig( + total_ram_gb=24.0, reserve_gb=2.0, + batch_size=None, # auto + ) + bs = cfg.auto_batch_size(act_per_sample_mb=120) + assert bs >= 1 + # Must be power of 2. + assert (bs & (bs - 1)) == 0 + + def test_auto_batch_size_minimum_1(self) -> None: + import gin + gin.clear_config() + cfg = HardwareConfig(total_ram_gb=1.0, reserve_gb=0.5, batch_size=None) + bs = cfg.auto_batch_size(act_per_sample_mb=1000) + assert bs >= 1 + + +class TestPipelineConfig: + def test_default_stages(self, pipeline_conf: PipelineConfig) -> None: + assert pipeline_conf.stages == ["depth", "edges", "segmentation", "chmv2"] + + def test_save_flags(self, pipeline_conf: PipelineConfig) -> None: + assert pipeline_conf.save_npy is True + assert pipeline_conf.save_vis is True + assert pipeline_conf.save_concat is False + + +class TestSegConfig: + def test_num_classes(self, seg_conf: SegConfig) -> None: + assert seg_conf.num_classes == 3 + assert seg_conf.prompts == ["background", "building", "road"] + + +class TestInputConfig: + def test_defaults(self, input_conf: InputConfig) -> None: + assert input_conf.image_size == 64 + assert len(input_conf.imagenet_mean) == 3 diff --git a/src/tests/test_dataset.py b/src/tests/test_dataset.py new file mode 100644 index 0000000..9d7a1cc --- /dev/null +++ b/src/tests/test_dataset.py @@ -0,0 +1,114 @@ +"""Tests for dataset discovery, filtering, and AugmentDataset.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from src.augmentor.dataset import ( + AugmentDataset, + ImageRecord, + attach_output_dirs, + discover_images, + filter_completed, + INCOMPLETE_SCENES, +) + + +class TestDiscoverImages: + def test_finds_all_images(self, fake_image_dir: Path) -> None: + records = discover_images(fake_image_dir) + assert len(records) == 4 + + def test_records_are_sorted(self, fake_image_dir: Path) -> None: + records = discover_images(fake_image_dir) + paths = [r.rel_path for r in records] + assert paths == sorted(paths) + + def test_skips_incomplete_scenes(self, tmp_path: Path) -> None: + scene = list(INCOMPLETE_SCENES)[0] + img_dir = tmp_path / scene + img_dir.mkdir(parents=True) + Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(img_dir / "a.png") + records = discover_images(tmp_path) + assert len(records) == 0 + + def test_skips_excluded_extensions(self, tmp_path: Path) -> None: + (tmp_path / "file.tiff").touch() + (tmp_path / "file.txt").touch() + (tmp_path / "good.png").write_bytes( + Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).tobytes() + ) + # Only actually valid images with correct extensions are found. + records = discover_images(tmp_path) + # good.png is not a valid PNG file (raw bytes), but discover only checks extension. + assert all(r.abs_path.suffix in {".png", ".jpg", ".jpeg", ".bmp"} for r in records) + + def test_subset_filter(self, tmp_path: Path) -> None: + sub = tmp_path / "Rot" / "scene" + sub.mkdir(parents=True) + Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(sub / "a.png") + (tmp_path / "Country").mkdir() + Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save( + tmp_path / "Country" / "b.png", + ) + records = discover_images(tmp_path, subset="Rot") + assert len(records) == 1 + assert "Rot" in records[0].rel_path + + def test_empty_root(self, tmp_path: Path) -> None: + records = discover_images(tmp_path / "nonexistent") + assert records == [] + + +class TestAttachOutputDirs: + def test_output_dir_structure(self, fake_image_dir: Path, tmp_path: Path) -> None: + records = discover_images(fake_image_dir) + out_root = tmp_path / "output" + attached = attach_output_dirs(records, out_root) + for r in attached: + assert str(r.output_dir).startswith(str(out_root)) + # output_dir is the parent directory (same structure, no per-image subfolder) + assert r.output_dir == out_root / Path(r.rel_path).parent + + +class TestFilterCompleted: + def test_all_pending_when_no_output(self, sample_records: list[ImageRecord]) -> None: + pending, skipped = filter_completed(sample_records, "depth") + assert len(pending) == len(sample_records) + assert skipped == 0 + + def test_skips_completed(self, sample_records: list[ImageRecord]) -> None: + r = sample_records[0] + r.output_dir.mkdir(parents=True, exist_ok=True) + np.save(r.output_dir / f"{r.stem}_depth.npy", np.zeros((1, 8, 8))) + pending, skipped = filter_completed(sample_records, "depth") + assert skipped == 1 + assert len(pending) == len(sample_records) - 1 + + def test_unknown_stage_returns_all(self, sample_records: list[ImageRecord]) -> None: + pending, skipped = filter_completed(sample_records, "unknown_stage") + assert len(pending) == len(sample_records) + assert skipped == 0 + + +class TestAugmentDataset: + def test_len(self, sample_records: list[ImageRecord]) -> None: + ds = AugmentDataset(sample_records, image_size=32) + assert len(ds) == len(sample_records) + + def test_getitem_shape(self, sample_records: list[ImageRecord]) -> None: + ds = AugmentDataset(sample_records, image_size=32) + item = ds[0] + assert item["image_raw"].shape == (3, 32, 32) + assert item["image_raw"].dtype == torch.float32 + assert 0.0 <= item["image_raw"].min() + assert item["image_raw"].max() <= 1.0 + + def test_getitem_keys(self, sample_records: list[ImageRecord]) -> None: + ds = AugmentDataset(sample_records, image_size=32) + item = ds[0] + assert set(item.keys()) == {"image_raw", "rel_path", "stem", "output_dir"} diff --git a/src/tests/test_inference.py b/src/tests/test_inference.py new file mode 100644 index 0000000..c19963e --- /dev/null +++ b/src/tests/test_inference.py @@ -0,0 +1,222 @@ +"""Tests for inference functions: depth, edges, segmentation.""" + +from __future__ import annotations + +from unittest.mock import MagicMock +from types import SimpleNamespace + +import numpy as np +import torch + +from src.augmentor.inference import ( + compute_edges_from_depth, + infer_chmv2_batch, + infer_depth_batch, + infer_segmentation_batch, + _SOBEL_X, + _SOBEL_Y, +) + + +# --------------------------------------------------------------------------- +# CHMv2 — mock model +# --------------------------------------------------------------------------- + +class TestInferChmv2Batch: + def _make_mock(self, B, H, W): + model = MagicMock() + model.parameters = MagicMock(return_value=iter([torch.nn.Parameter(torch.zeros(1))])) + + def forward(pixel_values=None): + bb = pixel_values.shape[0] + return SimpleNamespace(predicted_depth=torch.rand(bb, H, W)) + + model.side_effect = forward + + processor = MagicMock() + + def preprocess(images=None, return_tensors=None): + bb = len(images) + return {"pixel_values": torch.rand(bb, 3, H, W)} + + processor.side_effect = preprocess + + def post_process(outputs, target_sizes=None): + bb = outputs.predicted_depth.shape[0] + return [ + {"predicted_depth": torch.rand(target_sizes[i][0], target_sizes[i][1])} + for i in range(bb) + ] + + processor.post_process_depth_estimation = post_process + return model, processor + + def test_output_shape(self, sample_batch: torch.Tensor) -> None: + B, _, H, W = sample_batch.shape + model, processor = self._make_mock(B, H, W) + result = infer_chmv2_batch(model, processor, sample_batch, torch.device("cpu")) + assert result.shape == (B, 1, H, W) + + def test_output_normalized(self, sample_batch: torch.Tensor) -> None: + B, _, H, W = sample_batch.shape + model, processor = self._make_mock(B, H, W) + result = infer_chmv2_batch(model, processor, sample_batch, torch.device("cpu")) + assert result.min() >= 0.0 + assert result.max() <= 1.0 + 1e-6 + + +# --------------------------------------------------------------------------- +# Edges (Sobel) — pure computation, no model needed +# --------------------------------------------------------------------------- + +class TestComputeEdges: + def test_output_shape(self, sample_depth: torch.Tensor) -> None: + edges = compute_edges_from_depth(sample_depth) + assert edges.shape == sample_depth.shape + + def test_output_range(self, sample_depth: torch.Tensor) -> None: + edges = compute_edges_from_depth(sample_depth) + assert edges.min() >= 0.0 + assert edges.max() <= 1.0 + 1e-6 + + def test_flat_depth_gives_zero_edges(self) -> None: + flat = torch.ones(1, 1, 32, 32) * 0.5 + edges = compute_edges_from_depth(flat) + assert edges.max().item() < 1e-6 + + def test_single_image(self) -> None: + depth = torch.rand(1, 1, 16, 16) + edges = compute_edges_from_depth(depth) + assert edges.shape == (1, 1, 16, 16) + + def test_batch_consistency(self) -> None: + """Batched result matches per-image results.""" + depth = torch.rand(4, 1, 32, 32) + edges_batch = compute_edges_from_depth(depth) + for i in range(4): + edges_single = compute_edges_from_depth(depth[i : i + 1]) + torch.testing.assert_close(edges_batch[i], edges_single[0], atol=1e-5, rtol=1e-5) + + +class TestSobelKernelCache: + def test_kernels_are_module_level(self) -> None: + assert _SOBEL_X.shape == (1, 1, 3, 3) + assert _SOBEL_Y.shape == (1, 1, 3, 3) + assert _SOBEL_X.dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# Depth — mock model +# --------------------------------------------------------------------------- + +class TestInferDepthBatch: + def test_da3_api_path(self, sample_batch: torch.Tensor) -> None: + """Test DA3 branch with model.inference().""" + B, _, H, W = sample_batch.shape + + def fake_inference(image_list, process_res=None): + depth = np.random.rand(len(image_list), H, W).astype(np.float32) + return SimpleNamespace(depth=depth) + + model = MagicMock() + model.inference = fake_inference + + result = infer_depth_batch(model, sample_batch, torch.device("cpu")) + assert result.shape == (B, 1, H, W) + assert result.min() >= 0.0 + assert result.max() <= 1.0 + 1e-6 + + def test_transformers_api_path(self, sample_batch: torch.Tensor) -> None: + """Test DA V2 fallback with model(pixel_values=x).predicted_depth.""" + B, _, H, W = sample_batch.shape + + # Mock a transformers-style model — spec=[] ensures hasattr(model, "inference") is False. + mock_param = torch.nn.Parameter(torch.zeros(1)) # float32 + model = MagicMock(spec=[]) + model.parameters = MagicMock(return_value=iter([mock_param])) + + pred_depth = torch.rand(B, H, W) + model.return_value = SimpleNamespace(predicted_depth=pred_depth) + + result = infer_depth_batch(model, sample_batch, torch.device("cpu")) + assert result.shape == (B, 1, H, W) + assert result.min() >= 0.0 + assert result.max() <= 1.0 + 1e-6 + + def test_per_frame_normalization(self) -> None: + """Each frame should be independently normalized to [0, 1].""" + B, H, W = 2, 16, 16 + + def fake_inference(image_list, process_res=None): + depth = np.random.rand(len(image_list), H, W).astype(np.float32) * 100 + return SimpleNamespace(depth=depth) + + model = MagicMock() + model.inference = fake_inference + + batch = torch.rand(B, 3, H, W) + result = infer_depth_batch(model, batch, torch.device("cpu")) + for i in range(B): + assert result[i].min().item() < 0.01 + assert result[i].max().item() > 0.99 + + +# --------------------------------------------------------------------------- +# Segmentation — mock model +# --------------------------------------------------------------------------- + +class TestInferSegmentationBatch: + def test_segformer_path(self, sample_batch: torch.Tensor) -> None: + """Test SegFormer fallback branch.""" + B, _, H, W = sample_batch.shape + + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock() + model.parameters.return_value = iter([mock_param]) + + num_classes = 5 + logits = torch.randn(B, num_classes, H // 4, W // 4) + model.return_value = SimpleNamespace(logits=logits) + + processor = SimpleNamespace( + image_mean=[0.485, 0.456, 0.406], + image_std=[0.229, 0.224, 0.225], + ) + seg_config = {"type": "segformer", "processor": processor} + + result = infer_segmentation_batch(model, seg_config, sample_batch, torch.device("cpu")) + assert result.shape == (B, 1, H, W) + assert result.dtype == torch.uint8 + assert result.min() >= 0 + assert result.max() < num_classes + + def test_segearth_path(self, sample_batch: torch.Tensor) -> None: + """Test SegEarth-OV3 branch with mock pipeline.""" + B, _, H, W = sample_batch.shape + + def fake_predict(pil_img, text_prompts=None): + return np.random.randint(0, 3, (H, W), dtype=np.uint8) + + pipeline = MagicMock() + pipeline.predict = fake_predict + + seg_config = { + "type": "segearth-ov3", + "prompts": ["bg", "building", "road"], + } + + result = infer_segmentation_batch(pipeline, seg_config, sample_batch, torch.device("cpu")) + assert result.shape == (B, 1, H, W) + assert result.dtype == torch.uint8 + + def test_segearth_handles_failure(self, sample_batch: torch.Tensor) -> None: + """Failed prediction should produce zero tensor, not crash.""" + B, _, H, W = sample_batch.shape + + pipeline = MagicMock() + pipeline.predict.side_effect = RuntimeError("OOM") + + seg_config = {"type": "segearth-ov3", "prompts": ["bg"]} + result = infer_segmentation_batch(pipeline, seg_config, sample_batch, torch.device("cpu")) + assert result.shape == (B, 1, H, W) + assert (result == 0).all() diff --git a/src/tests/test_io_utils.py b/src/tests/test_io_utils.py new file mode 100644 index 0000000..e367224 --- /dev/null +++ b/src/tests/test_io_utils.py @@ -0,0 +1,191 @@ +"""Tests for I/O utilities: save functions, async pool, dtype correctness.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from src.augmentor.io_utils import ( + make_palette, + save_chmv2, + save_chmv2_async, + save_depth, + save_depth_async, + save_edges, + save_edges_async, + save_segmentation, + save_segmentation_async, + save_concat_6ch, + shutdown_io_pool, + _atomic_save_npy, +) + + +# --------------------------------------------------------------------------- +# Atomic save +# --------------------------------------------------------------------------- + +class TestAtomicSaveNpy: + def test_creates_file(self, tmp_path: Path) -> None: + arr = np.array([1.0, 2.0, 3.0]) + path = tmp_path / "test.npy" + _atomic_save_npy(arr, path) + assert path.exists() + loaded = np.load(path) + np.testing.assert_array_equal(loaded, arr) + + def test_overwrites_existing(self, tmp_path: Path) -> None: + path = tmp_path / "test.npy" + _atomic_save_npy(np.array([1.0]), path) + _atomic_save_npy(np.array([2.0]), path) + loaded = np.load(path) + assert loaded.item() == 2.0 + + def test_no_leftover_tmp_on_success(self, tmp_path: Path) -> None: + _atomic_save_npy(np.array([1.0]), tmp_path / "test.npy") + tmp_files = list(tmp_path.glob("*.tmp")) + assert len(tmp_files) == 0 + + +# --------------------------------------------------------------------------- +# save_depth +# --------------------------------------------------------------------------- + +class TestSaveDepth: + def test_saves_float16_npy(self, tmp_path: Path) -> None: + save_depth(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=True, save_vis=False) + arr = np.load(tmp_path / "img01_depth.npy") + assert arr.dtype == np.float16 + assert arr.shape == (1, 32, 32) + + def test_saves_vis_png(self, tmp_path: Path) -> None: + save_depth(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=False, save_vis=True) + assert (tmp_path / "img01_depth.png").exists() + assert not (tmp_path / "img01_depth.npy").exists() + + def test_npy_values_in_range(self, tmp_path: Path) -> None: + save_depth(torch.rand(1, 16, 16), tmp_path, "img01") + arr = np.load(tmp_path / "img01_depth.npy").astype(np.float32) + assert arr.min() >= 0.0 + assert arr.max() <= 1.0 + + +class TestSaveChmv2: + def test_saves_float16_npy(self, tmp_path: Path) -> None: + save_chmv2(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=True, save_vis=False) + arr = np.load(tmp_path / "img01_chm.npy") + assert arr.dtype == np.float16 + assert arr.shape == (1, 32, 32) + + def test_saves_vis_png(self, tmp_path: Path) -> None: + save_chmv2(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=False, save_vis=True) + assert (tmp_path / "img01_chm.png").exists() + assert not (tmp_path / "img01_chm.npy").exists() + + def test_npy_values_in_range(self, tmp_path: Path) -> None: + save_chmv2(torch.rand(1, 16, 16), tmp_path, "img01") + arr = np.load(tmp_path / "img01_chm.npy").astype(np.float32) + assert arr.min() >= 0.0 + assert arr.max() <= 1.0 + + +class TestSaveEdges: + def test_saves_float16_npy(self, tmp_path: Path) -> None: + save_edges(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=True, save_vis=False) + arr = np.load(tmp_path / "img01_edge.npy") + assert arr.dtype == np.float16 + + def test_saves_vis(self, tmp_path: Path) -> None: + save_edges(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=False, save_vis=True) + assert (tmp_path / "img01_edge.png").exists() + + +class TestSaveSegmentation: + def test_saves_uint8_npy(self, tmp_path: Path) -> None: + seg = torch.randint(0, 5, (1, 32, 32), dtype=torch.uint8) + save_segmentation(seg, tmp_path, "img01", save_npy=True, save_vis=False, num_classes=5) + arr = np.load(tmp_path / "img01_segm.npy") + assert arr.dtype == np.uint8 + assert arr.shape == (1, 32, 32) + + def test_saves_vis(self, tmp_path: Path) -> None: + seg = torch.randint(0, 3, (1, 32, 32), dtype=torch.uint8) + save_segmentation(seg, tmp_path, "img01", save_npy=False, save_vis=True, num_classes=3) + assert (tmp_path / "img01_segm.png").exists() + + def test_int_tensor_also_works(self, tmp_path: Path) -> None: + seg = torch.randint(0, 5, (1, 16, 16), dtype=torch.int64) + save_segmentation(seg, tmp_path, "img01", num_classes=5) + arr = np.load(tmp_path / "img01_segm.npy") + assert arr.dtype == np.uint8 + + +class TestAsyncSaves: + def test_depth_async(self, tmp_path: Path) -> None: + save_depth_async(torch.rand(1, 16, 16), tmp_path, "img01", save_npy=True, save_vis=False) + shutdown_io_pool() + assert (tmp_path / "img01_depth.npy").exists() + + def test_chmv2_async(self, tmp_path: Path) -> None: + save_chmv2_async(torch.rand(1, 16, 16), tmp_path, "img01", save_npy=True, save_vis=False) + shutdown_io_pool() + assert (tmp_path / "img01_chm.npy").exists() + + def test_edges_async(self, tmp_path: Path) -> None: + save_edges_async(torch.rand(1, 16, 16), tmp_path, "img01", save_npy=True, save_vis=False) + shutdown_io_pool() + assert (tmp_path / "img01_edge.npy").exists() + + def test_segmentation_async(self, tmp_path: Path) -> None: + seg = torch.randint(0, 3, (1, 16, 16), dtype=torch.uint8) + save_segmentation_async(seg, tmp_path, "img01", save_npy=True, save_vis=False, num_classes=3) + shutdown_io_pool() + assert (tmp_path / "img01_segm.npy").exists() + + def test_multiple_async_writes(self, tmp_path: Path) -> None: + for i in range(8): + save_depth_async(torch.rand(1, 8, 8), tmp_path, f"img_{i}", + save_npy=True, save_vis=False) + shutdown_io_pool() + for i in range(8): + assert (tmp_path / f"img_{i}_depth.npy").exists() + + +class TestSaveConcat6ch: + def test_concat_shape(self, tmp_path: Path) -> None: + H, W = 16, 16 + rgb = torch.rand(3, H, W) + np.save(tmp_path / "img01_depth.npy", np.random.rand(1, H, W).astype(np.float16)) + np.save(tmp_path / "img01_edge.npy", np.random.rand(1, H, W).astype(np.float16)) + np.save(tmp_path / "img01_segm.npy", np.random.randint(0, 5, (1, H, W)).astype(np.uint8)) + + save_concat_6ch(rgb, tmp_path, stem="img01", num_classes=5) + concat = np.load(tmp_path / "img01_concat.npy") + assert concat.shape == (6, H, W) + + def test_skips_when_missing(self, tmp_path: Path) -> None: + rgb = torch.rand(3, 8, 8) + save_concat_6ch(rgb, tmp_path, stem="img01", num_classes=5) + assert not (tmp_path / "img01_concat.npy").exists() + + +# --------------------------------------------------------------------------- +# Palette +# --------------------------------------------------------------------------- + +class TestMakePalette: + def test_shape(self) -> None: + pal = make_palette(10) + assert pal.shape == (10, 3) + assert pal.dtype == np.uint8 + + def test_background_is_black(self) -> None: + pal = make_palette(5) + np.testing.assert_array_equal(pal[0], [0, 0, 0]) + + def test_cache_returns_same(self) -> None: + p1 = make_palette(7) + p2 = make_palette(7) + assert p1 is p2 diff --git a/src/tests/test_models_and_benchmark.py b/src/tests/test_models_and_benchmark.py new file mode 100644 index 0000000..d375014 --- /dev/null +++ b/src/tests/test_models_and_benchmark.py @@ -0,0 +1,661 @@ +"""Tests for model loading, synthetic inference, and benchmarking. + +All models are mocked — no real weights needed. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader + +from src.augmentor.inference import ( + compute_edges_from_depth, + infer_chmv2_batch, + infer_depth_batch, + infer_segmentation_batch, +) +from src.utils.benchmark import ( + BenchmarkResult, + SyntheticImageDataset, + benchmark_chmv2, + benchmark_depth, + benchmark_edges, + benchmark_full_pipeline, + benchmark_segmentation, +) + +from src.conf.hardware_conf import HardwareConfig + +DEVICE = torch.device("cpu") + + +# --------------------------------------------------------------------------- +# Mock model factories +# --------------------------------------------------------------------------- + +def _mock_da3(H: int = 64, W: int = 64) -> MagicMock: + """Mock DA3 model with inference() method.""" + model = MagicMock() + + def inference(image_list, process_res=None): + depth = np.random.rand(len(image_list), H, W).astype(np.float32) + return SimpleNamespace(depth=depth) + + model.inference = inference + return model + + +def _mock_transformers_depth(H: int = 64, W: int = 64) -> MagicMock: + """Mock transformers-style depth model (DA V2 fallback).""" + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock(spec=[]) + model.parameters = MagicMock(return_value=iter([mock_param])) + + def forward(pixel_values=None): + B = pixel_values.shape[0] + return SimpleNamespace(predicted_depth=torch.rand(B, H, W)) + + model.return_value = None + model.side_effect = lambda *a, **kw: forward(**kw) + # spec=[] means hasattr(model, "inference") is False + return model + + +def _mock_segformer(num_classes: int = 5, H: int = 64, W: int = 64): + """Mock SegFormer model + config dict.""" + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock() + model.parameters = lambda: iter([mock_param]) + + def forward(pixel_values=None): + B = pixel_values.shape[0] + return SimpleNamespace(logits=torch.randn(B, num_classes, H // 4, W // 4)) + + model.side_effect = forward + processor = SimpleNamespace( + image_mean=[0.485, 0.456, 0.406], + image_std=[0.229, 0.224, 0.225], + ) + seg_config = {"type": "segformer", "processor": processor, "num_classes": num_classes} + return model, seg_config + + +def _mock_segearth(num_classes: int = 5, H: int = 64, W: int = 64): + """Mock SegEarth-OV3 model + config dict.""" + model = MagicMock() + + def predict(pil_img, text_prompts=None): + return np.random.randint(0, num_classes, (H, W), dtype=np.uint8) + + model.predict = predict + seg_config = { + "type": "segearth-ov3", + "prompts": [f"class_{i}" for i in range(num_classes)], + "num_classes": num_classes, + } + return model, seg_config + + +def _mock_chmv2(H: int = 64, W: int = 64) -> tuple[MagicMock, MagicMock]: + """Mock CHMv2 model + processor.""" + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock() + model.parameters = lambda: iter([mock_param]) + + def forward(pixel_values=None): + B = pixel_values.shape[0] + return SimpleNamespace(predicted_depth=torch.rand(B, H, W)) + + model.side_effect = forward + + processor = MagicMock() + + def preprocess(images=None, return_tensors=None): + B = len(images) + return {"pixel_values": torch.rand(B, 3, H, W)} + + processor.side_effect = preprocess + + def post_process(outputs, target_sizes=None): + B = outputs.predicted_depth.shape[0] + results = [] + for i in range(B): + th, tw = target_sizes[i] if target_sizes else (H, W) + results.append({"predicted_depth": torch.rand(th, tw)}) + return results + + processor.post_process_depth_estimation = post_process + return model, processor + + +# =========================================================================== +# Model loading tests +# =========================================================================== + +class TestModelLoading: + """Test that model loading functions work with mock imports.""" + + def test_da3_model_loads(self) -> None: + """DA3 mock model initializes and has inference method.""" + model = _mock_da3() + assert hasattr(model, "inference") + result = model.inference([np.zeros((64, 64, 3), dtype=np.uint8)]) + assert result.depth.shape == (1, 64, 64) + + def test_transformers_depth_model_loads(self) -> None: + """Transformers fallback mock model initializes.""" + model = _mock_transformers_depth() + assert not hasattr(model, "inference") + params = list(model.parameters()) + assert len(params) > 0 + + def test_segformer_model_loads(self) -> None: + """SegFormer mock loads with processor and config.""" + model, seg_config = _mock_segformer() + assert seg_config["type"] == "segformer" + assert "processor" in seg_config + assert seg_config["num_classes"] == 5 + + def test_chmv2_model_loads(self) -> None: + """CHMv2 mock loads with model + processor.""" + model, processor = _mock_chmv2() + assert callable(processor.post_process_depth_estimation) + + def test_segearth_model_loads(self) -> None: + """SegEarth mock loads with predict method.""" + model, seg_config = _mock_segearth() + assert hasattr(model, "predict") + assert seg_config["type"] == "segearth-ov3" + assert len(seg_config["prompts"]) == 5 + + +# =========================================================================== +# Synthetic inference tests +# =========================================================================== + +class TestDepthInference: + """Depth inference on synthetic data.""" + + @pytest.mark.parametrize("batch_size", [1, 2, 4]) + def test_da3_output_shape(self, batch_size: int) -> None: + H = W = 64 + model = _mock_da3(H, W) + images = torch.rand(batch_size, 3, H, W) + result = infer_depth_batch(model, images, DEVICE) + assert result.shape == (batch_size, 1, H, W) + + def test_da3_output_normalized(self) -> None: + model = _mock_da3(64, 64) + images = torch.rand(2, 3, 64, 64) + result = infer_depth_batch(model, images, DEVICE) + assert result.min() >= 0.0 + assert result.max() <= 1.0 + 1e-6 + + def test_transformers_fallback_output(self) -> None: + H = W = 32 + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock(spec=[]) + model.parameters = MagicMock(return_value=iter([mock_param])) + pred = torch.rand(2, H, W) + model.return_value = SimpleNamespace(predicted_depth=pred) + + images = torch.rand(2, 3, H, W) + result = infer_depth_batch(model, images, DEVICE) + assert result.shape == (2, 1, H, W) + assert result.min() >= 0.0 + assert result.max() <= 1.0 + 1e-6 + + @pytest.mark.parametrize("image_size", [32, 64, 128]) + def test_different_resolutions(self, image_size: int) -> None: + model = _mock_da3(image_size, image_size) + images = torch.rand(2, 3, image_size, image_size) + result = infer_depth_batch(model, images, DEVICE) + assert result.shape == (2, 1, image_size, image_size) + + +class TestChmv2Inference: + """CHMv2 inference on synthetic data.""" + + @pytest.mark.parametrize("batch_size", [1, 2, 4]) + def test_output_shape(self, batch_size: int) -> None: + H = W = 64 + model, processor = _mock_chmv2(H, W) + images = torch.rand(batch_size, 3, H, W) + result = infer_chmv2_batch(model, processor, images, DEVICE) + assert result.shape == (batch_size, 1, H, W) + + def test_output_normalized(self) -> None: + H = W = 64 + model, processor = _mock_chmv2(H, W) + images = torch.rand(2, 3, H, W) + result = infer_chmv2_batch(model, processor, images, DEVICE) + assert result.min() >= 0.0 + assert result.max() <= 1.0 + 1e-6 + + @pytest.mark.parametrize("image_size", [32, 64, 128]) + def test_different_resolutions(self, image_size: int) -> None: + model, processor = _mock_chmv2(image_size, image_size) + images = torch.rand(2, 3, image_size, image_size) + result = infer_chmv2_batch(model, processor, images, DEVICE) + assert result.shape == (2, 1, image_size, image_size) + + +class TestEdgesInference: + """Edges (Sobel) computation on synthetic data.""" + + @pytest.mark.parametrize("batch_size", [1, 2, 8]) + def test_output_shape(self, batch_size: int) -> None: + depth = torch.rand(batch_size, 1, 64, 64) + edges = compute_edges_from_depth(depth) + assert edges.shape == depth.shape + + def test_output_range(self) -> None: + depth = torch.rand(4, 1, 64, 64) + edges = compute_edges_from_depth(depth) + assert edges.min() >= 0.0 + assert edges.max() <= 1.0 + 1e-6 + + def test_flat_gives_zero(self) -> None: + depth = torch.ones(2, 1, 32, 32) * 0.5 + edges = compute_edges_from_depth(depth) + assert edges.max().item() < 1e-6 + + def test_deterministic(self) -> None: + depth = torch.rand(2, 1, 32, 32) + e1 = compute_edges_from_depth(depth) + e2 = compute_edges_from_depth(depth) + torch.testing.assert_close(e1, e2) + + +class TestSegmentationInference: + """Segmentation inference on synthetic data.""" + + @pytest.mark.parametrize("batch_size", [1, 2, 4]) + def test_segformer_output_shape(self, batch_size: int) -> None: + H = W = 64 + model, seg_config = _mock_segformer(5, H, W) + images = torch.rand(batch_size, 3, H, W) + result = infer_segmentation_batch(model, seg_config, images, DEVICE) + assert result.shape == (batch_size, 1, H, W) + assert result.dtype == torch.uint8 + + def test_segformer_class_range(self) -> None: + num_classes = 5 + model, seg_config = _mock_segformer(num_classes, 64, 64) + images = torch.rand(4, 3, 64, 64) + result = infer_segmentation_batch(model, seg_config, images, DEVICE) + assert result.min() >= 0 + assert result.max() < num_classes + + @pytest.mark.parametrize("batch_size", [1, 2]) + def test_segearth_output_shape(self, batch_size: int) -> None: + H = W = 64 + model, seg_config = _mock_segearth(5, H, W) + images = torch.rand(batch_size, 3, H, W) + result = infer_segmentation_batch(model, seg_config, images, DEVICE) + assert result.shape == (batch_size, 1, H, W) + assert result.dtype == torch.uint8 + + def test_segearth_failure_returns_zeros(self) -> None: + model = MagicMock() + model.predict.side_effect = RuntimeError("OOM") + seg_config = {"type": "segearth-ov3", "prompts": ["bg"]} + images = torch.rand(2, 3, 32, 32) + result = infer_segmentation_batch(model, seg_config, images, DEVICE) + assert result.shape == (2, 1, 32, 32) + assert (result == 0).all() + + +# =========================================================================== +# DataLoader integration tests +# =========================================================================== + +class TestSyntheticDataset: + """Test the synthetic dataset and DataLoader.""" + + def test_dataset_length(self) -> None: + ds = SyntheticImageDataset(num_images=50, image_size=64) + assert len(ds) == 50 + + def test_dataset_item_shape(self) -> None: + ds = SyntheticImageDataset(num_images=10, image_size=128) + item = ds[0] + assert item["image_raw"].shape == (3, 128, 128) + assert item["image_raw"].dtype == torch.float32 + assert item["image_raw"].min() >= 0.0 + assert item["image_raw"].max() <= 1.0 + + def test_dataloader_batching(self) -> None: + ds = SyntheticImageDataset(num_images=10, image_size=64) + loader = DataLoader(ds, batch_size=4, shuffle=False) + batches = list(loader) + assert len(batches) == 3 # 4 + 4 + 2 + assert batches[0]["image_raw"].shape == (4, 3, 64, 64) + assert batches[2]["image_raw"].shape == (2, 3, 64, 64) + + @pytest.mark.parametrize("num_images", [1, 7, 16]) + def test_various_counts(self, num_images: int) -> None: + ds = SyntheticImageDataset(num_images=num_images, image_size=32) + loader = DataLoader(ds, batch_size=4, shuffle=False) + total = sum(b["image_raw"].shape[0] for b in loader) + assert total == num_images + + def test_depth_inference_from_loader(self) -> None: + """Full path: DataLoader → inference → check output.""" + model = _mock_da3(64, 64) + ds = SyntheticImageDataset(num_images=8, image_size=64) + loader = DataLoader(ds, batch_size=4, shuffle=False) + all_depths = [] + for batch in loader: + depths = infer_depth_batch(model, batch["image_raw"], DEVICE) + all_depths.append(depths) + combined = torch.cat(all_depths, dim=0) + assert combined.shape == (8, 1, 64, 64) + + def test_segformer_inference_from_loader(self) -> None: + model, seg_config = _mock_segformer(5, 64, 64) + ds = SyntheticImageDataset(num_images=6, image_size=64) + loader = DataLoader(ds, batch_size=3, shuffle=False) + all_segs = [] + for batch in loader: + segs = infer_segmentation_batch(model, seg_config, batch["image_raw"], DEVICE) + all_segs.append(segs) + combined = torch.cat(all_segs, dim=0) + assert combined.shape == (6, 1, 64, 64) + assert combined.dtype == torch.uint8 + + +# =========================================================================== +# Benchmark result tests +# =========================================================================== + +class TestBenchmarkResult: + """Test BenchmarkResult calculations.""" + + def test_fps_calculation(self) -> None: + r = BenchmarkResult("test", num_images=100, total_time_s=2.0) + assert r.fps == pytest.approx(50.0) + + def test_avg_latency(self) -> None: + r = BenchmarkResult("test", num_images=100, total_time_s=2.0) + assert r.avg_latency_ms == pytest.approx(20.0) + + def test_median_batch_latency(self) -> None: + r = BenchmarkResult("test", num_images=100, total_time_s=2.0, + batch_times_s=[0.01, 0.02, 0.03, 0.04, 0.05]) + assert r.median_batch_latency_ms == pytest.approx(30.0) + + def test_p95_batch_latency(self) -> None: + times = [0.01] * 95 + [0.1] * 5 + r = BenchmarkResult("test", num_images=100, total_time_s=2.0, + batch_times_s=times) + assert r.p95_batch_latency_ms > r.median_batch_latency_ms # p95 > median + + def test_zero_images(self) -> None: + r = BenchmarkResult("test", num_images=0, total_time_s=0.0) + assert r.fps == 0.0 + assert r.avg_latency_ms == 0.0 + + def test_empty_batch_times(self) -> None: + r = BenchmarkResult("test", num_images=10, total_time_s=1.0) + assert r.median_batch_latency_ms == 0.0 + assert r.p95_batch_latency_ms == 0.0 + + +# =========================================================================== +# Benchmark function tests (small runs with mocks) +# =========================================================================== + +class TestBenchmarkDepth: + """Benchmark depth with mocked model.""" + + def test_runs_and_returns_result(self) -> None: + model = _mock_da3(32, 32) + result = benchmark_depth(model, DEVICE, num_images=8, + batch_size=4, image_size=32, warmup_batches=1) + assert isinstance(result, BenchmarkResult) + assert result.num_images == 8 + assert result.total_time_s > 0 + assert result.fps > 0 + assert len(result.batch_times_s) == 2 # 8 images / 4 batch + + def test_single_image(self) -> None: + model = _mock_da3(32, 32) + result = benchmark_depth(model, DEVICE, num_images=1, + batch_size=1, image_size=32, warmup_batches=0) + assert result.num_images == 1 + assert len(result.batch_times_s) == 1 + + +class TestBenchmarkChmv2: + """Benchmark CHMv2 with mocked model.""" + + def test_runs_and_returns_result(self) -> None: + model, processor = _mock_chmv2(32, 32) + result = benchmark_chmv2(model, processor, DEVICE, num_images=8, + batch_size=4, image_size=32, warmup_batches=1) + assert isinstance(result, BenchmarkResult) + assert result.num_images == 8 + assert result.total_time_s > 0 + assert result.fps > 0 + assert len(result.batch_times_s) == 2 + + def test_single_image(self) -> None: + model, processor = _mock_chmv2(32, 32) + result = benchmark_chmv2(model, processor, DEVICE, num_images=1, + batch_size=1, image_size=32, warmup_batches=0) + assert result.num_images == 1 + assert len(result.batch_times_s) == 1 + + +class TestBenchmarkEdges: + """Benchmark edges computation.""" + + def test_runs_and_returns_result(self) -> None: + result = benchmark_edges(num_images=16, batch_size=8, + image_size=32, warmup_batches=1) + assert isinstance(result, BenchmarkResult) + assert result.num_images == 16 + assert result.fps > 0 + assert len(result.batch_times_s) == 2 # 16 / 8 + + def test_odd_count(self) -> None: + result = benchmark_edges(num_images=7, batch_size=3, + image_size=32, warmup_batches=0) + assert result.num_images == 7 + assert len(result.batch_times_s) == 3 # ceil(7/3) + + +class TestBenchmarkSegmentation: + """Benchmark segmentation with mocked model.""" + + def test_segformer(self) -> None: + model, seg_config = _mock_segformer(5, 32, 32) + result = benchmark_segmentation(model, seg_config, DEVICE, + num_images=8, batch_size=4, + image_size=32, warmup_batches=1) + assert isinstance(result, BenchmarkResult) + assert result.num_images == 8 + assert result.fps > 0 + + def test_segearth(self) -> None: + model, seg_config = _mock_segearth(3, 32, 32) + result = benchmark_segmentation(model, seg_config, DEVICE, + num_images=4, batch_size=2, + image_size=32, warmup_batches=0) + assert result.num_images == 4 + + +class TestBenchmarkFullPipeline: + """Benchmark full pipeline with all stages mocked.""" + + def test_runs_all_stages(self) -> None: + depth_model = _mock_da3(32, 32) + seg_model, seg_config = _mock_segformer(5, 32, 32) + + results = benchmark_full_pipeline( + depth_model, seg_model, seg_config, DEVICE, + num_images=8, batch_size=4, image_size=32, warmup_batches=1, + ) + + assert "depth" in results + assert "edges" in results + assert "segmentation" in results + assert "total" in results + + for key, r in results.items(): + assert isinstance(r, BenchmarkResult) + assert r.num_images == 8 + assert r.total_time_s > 0 + assert r.fps > 0 + + def test_total_ge_sum_of_stages(self) -> None: + """Total time >= sum of individual stages (includes overhead).""" + depth_model = _mock_da3(32, 32) + seg_model, seg_config = _mock_segformer(3, 32, 32) + + results = benchmark_full_pipeline( + depth_model, seg_model, seg_config, DEVICE, + num_images=4, batch_size=2, image_size=32, warmup_batches=0, + ) + + stage_sum = (results["depth"].total_time_s + + results["edges"].total_time_s + + results["segmentation"].total_time_s) + assert results["total"].total_time_s >= stage_sum * 0.95 # allow 5% tolerance + + def test_configurable_image_count(self) -> None: + """num_images parameter controls how many images are processed.""" + depth_model = _mock_da3(32, 32) + seg_model, seg_config = _mock_segformer(3, 32, 32) + + for n in [4, 10, 17]: + results = benchmark_full_pipeline( + depth_model, seg_model, seg_config, DEVICE, + num_images=n, batch_size=4, image_size=32, warmup_batches=0, + ) + assert results["total"].num_images == n + + def test_different_batch_sizes(self) -> None: + depth_model = _mock_da3(32, 32) + seg_model, seg_config = _mock_segformer(3, 32, 32) + + for bs in [1, 2, 8]: + results = benchmark_full_pipeline( + depth_model, seg_model, seg_config, DEVICE, + num_images=8, batch_size=bs, image_size=32, warmup_batches=0, + ) + expected_batches = (8 + bs - 1) // bs + assert len(results["total"].batch_times_s) == expected_batches + + +# =========================================================================== +# find_batch_size tests +# =========================================================================== + +class TestFindBatchSize: + """Test empirical batch size probing.""" + + def test_returns_config_batch_size_when_set(self) -> None: + """If batch_size is set in config, find_batch_size returns it directly.""" + hw = HardwareConfig(batch_size=8) + bs = hw.find_batch_size( + inference_fn=lambda x: None, + input_shape=(3, 64, 64), + device=DEVICE, + ) + assert bs == 8 + + def test_cpu_returns_1(self) -> None: + """On CPU (no CUDA), returns 1.""" + hw = HardwareConfig(batch_size=None) + bs = hw.find_batch_size( + inference_fn=lambda x: None, + input_shape=(3, 64, 64), + device=torch.device("cpu"), + ) + assert bs == 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def test_finds_batch_gt_1_on_gpu(self) -> None: + """On a real GPU with a trivial model, batch should be > 1.""" + hw = HardwareConfig(batch_size=None) + device = torch.device("cuda") + + def trivial_fn(x: torch.Tensor) -> None: + _ = x.to(device) + torch.cuda.synchronize() + + bs = hw.find_batch_size( + inference_fn=trivial_fn, + input_shape=(3, 32, 32), + device=device, + max_batch=16, + ) + assert bs >= 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def test_measures_per_sample_cost(self) -> None: + """Heavier inference → smaller batch size.""" + hw = HardwareConfig(batch_size=None) + device = torch.device("cuda") + + # Light: tiny tensor on GPU. + def light_fn(x: torch.Tensor) -> None: + _ = x.to(device) + torch.cuda.synchronize() + + # Heavy: allocate a large intermediate on GPU. + def heavy_fn(x: torch.Tensor) -> None: + y = x.to(device) + # ~200 MB intermediate per sample for 512x512. + big = torch.randn(x.shape[0], 64, 512, 512, device=device) + _ = big.sum() + torch.cuda.synchronize() + + bs_light = hw.find_batch_size( + inference_fn=light_fn, + input_shape=(3, 32, 32), + device=device, + max_batch=64, + ) + bs_heavy = hw.find_batch_size( + inference_fn=heavy_fn, + input_shape=(3, 32, 32), + device=device, + max_batch=64, + ) + assert bs_light >= bs_heavy + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def test_respects_max_batch(self) -> None: + """Batch size never exceeds max_batch.""" + hw = HardwareConfig(batch_size=None) + device = torch.device("cuda") + + bs = hw.find_batch_size( + inference_fn=lambda x: x.to(device), + input_shape=(3, 32, 32), + device=device, + max_batch=4, + ) + assert bs <= 4 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def test_result_is_power_of_2(self) -> None: + """Batch size is always a power of 2.""" + hw = HardwareConfig(batch_size=None) + device = torch.device("cuda") + + bs = hw.find_batch_size( + inference_fn=lambda x: x.to(device), + input_shape=(3, 64, 64), + device=device, + ) + assert bs >= 1 + assert (bs & (bs - 1)) == 0 # power of 2 check diff --git a/src/tests/test_pipeline_integration.py b/src/tests/test_pipeline_integration.py new file mode 100644 index 0000000..fc0deb9 --- /dev/null +++ b/src/tests/test_pipeline_integration.py @@ -0,0 +1,328 @@ +"""Integration tests: run stages end-to-end with mock models on synthetic data.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import numpy as np +import torch + +from src.augmentor.dataset import ( + AugmentDataset, + ImageRecord, + attach_output_dirs, + discover_images, +) +from src.augmentor.inference import compute_edges_from_depth +from src.augmentor.io_utils import shutdown_io_pool +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 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_mock_depth_model_da3(H: int, W: int) -> MagicMock: + model = MagicMock() + + def inference(image_list, process_res=None): + depth = np.random.rand(len(image_list), H, W).astype(np.float32) + return SimpleNamespace(depth=depth) + + model.inference = inference + return model + + +def _make_mock_chmv2(H: int, W: int) -> tuple[MagicMock, MagicMock]: + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock() + model.parameters = lambda: iter([mock_param]) + + def forward(pixel_values=None): + B = pixel_values.shape[0] + return SimpleNamespace(predicted_depth=torch.rand(B, H, W)) + + model.side_effect = forward + + processor = MagicMock() + + def preprocess(images=None, return_tensors=None): + B = len(images) + return {"pixel_values": torch.rand(B, 3, H, W)} + + processor.side_effect = preprocess + + def post_process(outputs, target_sizes=None): + B = outputs.predicted_depth.shape[0] + return [ + {"predicted_depth": torch.rand(target_sizes[i][0], target_sizes[i][1])} + for i in range(B) + ] + + processor.post_process_depth_estimation = post_process + return model, processor + + +def _make_mock_segformer(num_classes: int, H: int, W: int) -> tuple[MagicMock, dict]: + mock_param = torch.nn.Parameter(torch.zeros(1)) + model = MagicMock() + model.parameters = lambda: iter([mock_param]) + + def forward(pixel_values=None): + B = pixel_values.shape[0] + logits = torch.randn(B, num_classes, H // 4, W // 4) + return SimpleNamespace(logits=logits) + + model.side_effect = forward + + processor = SimpleNamespace( + image_mean=[0.485, 0.456, 0.406], + image_std=[0.229, 0.224, 0.225], + ) + seg_config = {"type": "segformer", "processor": processor, "num_classes": num_classes} + return model, seg_config + + +# --------------------------------------------------------------------------- +# Per-stage tests +# --------------------------------------------------------------------------- + +class TestDepthStageIntegration: + def test_produces_depth_files( + self, + sample_records: list[ImageRecord], + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + input_conf: InputConfig, + ) -> None: + from src.main import run_depth_stage + + for r in sample_records: + r.output_dir.mkdir(parents=True, exist_ok=True) + + mock_model = _make_mock_depth_model_da3(input_conf.image_size, input_conf.image_size) + + with patch("src.main.load_depth_model", return_value=mock_model), \ + patch("src.main.unload_model"): + run_depth_stage( + sample_records, pipeline_conf, hw_conf, + ModelsConfig(), input_conf, torch.device("cpu"), + ) + + for r in sample_records: + npy_path = r.output_dir / f"{r.stem}_depth.npy" + assert npy_path.exists(), f"Missing {r.stem}_depth.npy in {r.output_dir}" + arr = np.load(npy_path) + assert arr.dtype == np.float16 + assert arr.shape[0] == 1 + + +class TestChmv2StageIntegration: + def test_produces_chm_files( + self, + sample_records: list[ImageRecord], + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + input_conf: InputConfig, + ) -> None: + from src.main import run_chmv2_stage + + for r in sample_records: + r.output_dir.mkdir(parents=True, exist_ok=True) + + mock_model, mock_processor = _make_mock_chmv2( + input_conf.image_size, input_conf.image_size, + ) + + with patch("src.main.load_chmv2_model", return_value=(mock_model, mock_processor)), \ + patch("src.main.unload_model"): + run_chmv2_stage( + sample_records, pipeline_conf, hw_conf, + ModelsConfig(), input_conf, torch.device("cpu"), + ) + + for r in sample_records: + npy_path = r.output_dir / f"{r.stem}_chm.npy" + assert npy_path.exists(), f"Missing {r.stem}_chm.npy in {r.output_dir}" + arr = np.load(npy_path) + assert arr.dtype == np.float16 + assert arr.shape[0] == 1 + + +class TestEdgesStageIntegration: + def test_produces_edges_files( + self, + sample_records: list[ImageRecord], + pipeline_conf: PipelineConfig, + ) -> None: + from src.main import run_edges_stage + + for r in sample_records: + r.output_dir.mkdir(parents=True, exist_ok=True) + np.save( + r.output_dir / f"{r.stem}_depth.npy", + np.random.rand(1, 64, 64).astype(np.float16), + ) + + run_edges_stage(sample_records, pipeline_conf) + + for r in sample_records: + npy_path = r.output_dir / f"{r.stem}_edge.npy" + assert npy_path.exists(), f"Missing {r.stem}_edge.npy in {r.output_dir}" + arr = np.load(npy_path) + assert arr.dtype == np.float16 + + def test_skips_missing_depth( + self, + sample_records: list[ImageRecord], + pipeline_conf: PipelineConfig, + ) -> None: + from src.main import run_edges_stage + + for r in sample_records: + r.output_dir.mkdir(parents=True, exist_ok=True) + run_edges_stage(sample_records, pipeline_conf) + for r in sample_records: + assert not (r.output_dir / f"{r.stem}_edge.npy").exists() + + +class TestSegmentationStageIntegration: + def test_segformer_produces_seg_files( + self, + sample_records: list[ImageRecord], + pipeline_conf: PipelineConfig, + hw_conf: HardwareConfig, + input_conf: InputConfig, + seg_conf: SegConfig, + ) -> None: + from src.main import run_segmentation_stage + + for r in sample_records: + r.output_dir.mkdir(parents=True, exist_ok=True) + + num_classes = seg_conf.num_classes + mock_model, seg_config_dict = _make_mock_segformer( + num_classes, input_conf.image_size, input_conf.image_size, + ) + + with patch( + "src.main.load_segmentation_model", + return_value=(mock_model, seg_config_dict), + ), patch("src.main.unload_model"): + run_segmentation_stage( + sample_records, pipeline_conf, hw_conf, + ModelsConfig(), input_conf, seg_conf, torch.device("cpu"), + ) + + for r in sample_records: + npy_path = r.output_dir / f"{r.stem}_segm.npy" + assert npy_path.exists() + arr = np.load(npy_path) + assert arr.dtype == np.uint8 + + +# --------------------------------------------------------------------------- +# Full pipeline smoke test +# --------------------------------------------------------------------------- + +class TestFullPipelineSmoke: + def test_pipeline_completes( + self, + fake_image_dir: Path, + tmp_path: Path, + ) -> None: + import gin + gin.clear_config() + + pipeline_conf = PipelineConfig( + input_root=str(fake_image_dir), + output_root=str(tmp_path / "output"), + stages=["depth", "edges", "segmentation", "chmv2"], + save_npy=True, + save_vis=False, + save_concat=False, + resume=False, + log_level="WARNING", + ) + hw_conf = HardwareConfig( + profile_name="test", total_ram_gb=8.0, reserve_gb=1.0, + batch_size=2, num_workers=0, + ) + input_conf = InputConfig(image_size=32) + seg_conf = SegConfig(prompts=["bg", "building", "road"]) + models_conf = ModelsConfig() + + H = W = input_conf.image_size + num_classes = seg_conf.num_classes + + mock_depth = _make_mock_depth_model_da3(H, W) + mock_chmv2_model, mock_chmv2_proc = _make_mock_chmv2(H, W) + mock_seg, seg_config_dict = _make_mock_segformer(num_classes, H, W) + + with patch("src.main.load_depth_model", return_value=mock_depth), \ + patch("src.main.load_chmv2_model", return_value=(mock_chmv2_model, mock_chmv2_proc)), \ + patch("src.main.load_segmentation_model", return_value=(mock_seg, seg_config_dict)), \ + patch("src.main.unload_model"): + from src.main import run_pipeline + run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf) + + output_root = Path(pipeline_conf.output_root) + assert (output_root / "manifest.json").exists() + + found_depth = list(output_root.rglob("*_depth.npy")) + found_edges = list(output_root.rglob("*_edge.npy")) + found_seg = list(output_root.rglob("*_segm.npy")) + found_chmv2 = list(output_root.rglob("*_chm.npy")) + assert len(found_depth) > 0 + assert len(found_edges) > 0 + assert len(found_seg) > 0 + assert len(found_chmv2) > 0 + + +class TestFullPipelineVisOnly: + def test_produces_png_not_npy( + self, fake_image_dir: Path, tmp_path: Path, + ) -> None: + import gin + gin.clear_config() + + pipeline_conf = PipelineConfig( + input_root=str(fake_image_dir), + output_root=str(tmp_path / "output"), + stages=["depth", "edges", "segmentation", "chmv2"], + save_npy=False, save_vis=True, + save_concat=False, resume=False, log_level="WARNING", + ) + hw_conf = HardwareConfig( + profile_name="test", total_ram_gb=8.0, reserve_gb=1.0, + batch_size=2, num_workers=0, + ) + input_conf = InputConfig(image_size=32) + seg_conf = SegConfig(prompts=["bg", "building", "road"]) + models_conf = ModelsConfig() + H = W = 32 + num_classes = seg_conf.num_classes + + mock_depth = _make_mock_depth_model_da3(H, W) + mock_chmv2_model, mock_chmv2_proc = _make_mock_chmv2(H, W) + mock_seg, seg_config_dict = _make_mock_segformer(num_classes, H, W) + + with patch("src.main.load_depth_model", return_value=mock_depth), \ + patch("src.main.load_chmv2_model", return_value=(mock_chmv2_model, mock_chmv2_proc)), \ + patch("src.main.load_segmentation_model", return_value=(mock_seg, seg_config_dict)), \ + patch("src.main.unload_model"): + from src.main import run_pipeline + run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf) + + output_root = Path(pipeline_conf.output_root) + assert len(list(output_root.rglob("*_depth.png"))) > 0 + assert len(list(output_root.rglob("*_edge.png"))) > 0 + assert len(list(output_root.rglob("*_segm.png"))) > 0 + assert len(list(output_root.rglob("*_chm.png"))) > 0 + assert len(list(output_root.rglob("*.npy"))) == 0 diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..f5d3dc7 --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1 @@ +"""Utility functions.""" diff --git a/src/utils/benchmark.py b/src/utils/benchmark.py new file mode 100644 index 0000000..bc13711 --- /dev/null +++ b/src/utils/benchmark.py @@ -0,0 +1,331 @@ +"""Benchmarking utilities: FPS, latency, throughput per pipeline stage. + +Usage: + from src.utils.benchmark import benchmark_stage, benchmark_full_pipeline +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass, field +from typing import Any, Callable + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Result container +# --------------------------------------------------------------------------- + +@dataclass +class BenchmarkResult: + """Benchmark measurement for a single stage or full pipeline.""" + + stage: str + num_images: int + total_time_s: float + batch_times_s: list[float] = field(default_factory=list) + + @property + def fps(self) -> float: + return self.num_images / self.total_time_s if self.total_time_s > 0 else 0.0 + + @property + def avg_latency_ms(self) -> float: + return (self.total_time_s / self.num_images * 1000) if self.num_images > 0 else 0.0 + + @property + def median_batch_latency_ms(self) -> float: + if not self.batch_times_s: + return 0.0 + return float(np.median(self.batch_times_s)) * 1000 + + @property + def p95_batch_latency_ms(self) -> float: + if not self.batch_times_s: + return 0.0 + return float(np.percentile(self.batch_times_s, 95)) * 1000 + + def log_summary(self) -> None: + logger.info("📊 Benchmark [%s]:", self.stage) + logger.info(" Images: %d", self.num_images) + logger.info(" Total: %.2f s", self.total_time_s) + logger.info(" FPS: %.1f img/s", self.fps) + logger.info(" Avg latency: %.1f ms/img", self.avg_latency_ms) + logger.info(" Median batch: %.1f ms", self.median_batch_latency_ms) + logger.info(" P95 batch: %.1f ms", self.p95_batch_latency_ms) + + +# --------------------------------------------------------------------------- +# Synthetic dataset +# --------------------------------------------------------------------------- + +class SyntheticImageDataset(Dataset): + """Generate random RGB images for benchmarking.""" + + def __init__(self, num_images: int, image_size: int = 256) -> None: + self.num_images = num_images + self.image_size = image_size + + def __len__(self) -> int: + return self.num_images + + def __getitem__(self, idx: int) -> dict[str, Any]: + return { + "image_raw": torch.rand(3, self.image_size, self.image_size), + "rel_path": f"synthetic/img_{idx:05d}.png", + "stem": f"img_{idx:05d}", + "output_dir": "", + } + + +# --------------------------------------------------------------------------- +# Stage benchmarks +# --------------------------------------------------------------------------- + +def benchmark_depth( + model: Any, + device: torch.device, + num_images: int = 100, + batch_size: int = 4, + image_size: int = 256, + num_workers: int = 0, + warmup_batches: int = 3, +) -> BenchmarkResult: + """Benchmark depth inference stage.""" + from src.augmentor.inference import infer_depth_batch + + ds = SyntheticImageDataset(num_images, image_size) + loader = DataLoader(ds, batch_size=batch_size, shuffle=False, + num_workers=num_workers, pin_memory=torch.cuda.is_available()) + + # Warmup + for i, batch in enumerate(loader): + if i >= warmup_batches: + break + infer_depth_batch(model, batch["image_raw"], device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + + batch_times = [] + processed = 0 + t_start = time.perf_counter() + + for batch in tqdm(loader, desc="⏱️ bench depth", unit="batch", colour="green"): + t0 = time.perf_counter() + infer_depth_batch(model, batch["image_raw"], device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + batch_times.append(time.perf_counter() - t0) + processed += batch["image_raw"].shape[0] + + total = time.perf_counter() - t_start + result = BenchmarkResult("🌊 depth", processed, total, batch_times) + result.log_summary() + return result + + +def benchmark_chmv2( + model: Any, + processor: Any, + device: torch.device, + num_images: int = 100, + batch_size: int = 4, + image_size: int = 256, + num_workers: int = 0, + warmup_batches: int = 3, +) -> BenchmarkResult: + """Benchmark CHMv2 inference stage.""" + from src.augmentor.inference import infer_chmv2_batch + + ds = SyntheticImageDataset(num_images, image_size) + loader = DataLoader(ds, batch_size=batch_size, shuffle=False, + num_workers=num_workers, pin_memory=torch.cuda.is_available()) + + # Warmup + for i, batch in enumerate(loader): + if i >= warmup_batches: + break + infer_chmv2_batch(model, processor, batch["image_raw"], device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + + batch_times = [] + processed = 0 + t_start = time.perf_counter() + + for batch in tqdm(loader, desc="⏱️ bench chmv2", unit="batch", colour="blue"): + t0 = time.perf_counter() + infer_chmv2_batch(model, processor, batch["image_raw"], device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + batch_times.append(time.perf_counter() - t0) + processed += batch["image_raw"].shape[0] + + total = time.perf_counter() - t_start + result = BenchmarkResult("🏔️ chmv2", processed, total, batch_times) + result.log_summary() + return result + + +def benchmark_edges( + num_images: int = 100, + batch_size: int = 32, + image_size: int = 256, + warmup_batches: int = 3, +) -> BenchmarkResult: + """Benchmark edges (Sobel) computation stage.""" + from src.augmentor.inference import compute_edges_from_depth + + total_batches = (num_images + batch_size - 1) // batch_size + + # Warmup + for _ in range(warmup_batches): + dummy = torch.rand(batch_size, 1, image_size, image_size) + compute_edges_from_depth(dummy) + + batch_times = [] + processed = 0 + t_start = time.perf_counter() + + for i in tqdm(range(total_batches), desc="⏱️ bench edges", unit="batch", colour="cyan"): + cur_bs = min(batch_size, num_images - i * batch_size) + depth = torch.rand(cur_bs, 1, image_size, image_size) + t0 = time.perf_counter() + compute_edges_from_depth(depth) + batch_times.append(time.perf_counter() - t0) + processed += cur_bs + + total = time.perf_counter() - t_start + result = BenchmarkResult("🔪 edges", processed, total, batch_times) + result.log_summary() + return result + + +def benchmark_segmentation( + model: Any, + seg_config: dict[str, Any], + device: torch.device, + num_images: int = 100, + batch_size: int = 4, + image_size: int = 256, + num_workers: int = 0, + warmup_batches: int = 3, +) -> BenchmarkResult: + """Benchmark segmentation inference stage.""" + from src.augmentor.inference import infer_segmentation_batch + + ds = SyntheticImageDataset(num_images, image_size) + loader = DataLoader(ds, batch_size=batch_size, shuffle=False, + num_workers=num_workers, pin_memory=torch.cuda.is_available()) + + # Warmup + for i, batch in enumerate(loader): + if i >= warmup_batches: + break + infer_segmentation_batch(model, seg_config, batch["image_raw"], device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + + batch_times = [] + processed = 0 + t_start = time.perf_counter() + + for batch in tqdm(loader, desc="⏱️ bench segmentation", unit="batch", colour="yellow"): + t0 = time.perf_counter() + infer_segmentation_batch(model, seg_config, batch["image_raw"], device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + batch_times.append(time.perf_counter() - t0) + processed += batch["image_raw"].shape[0] + + total = time.perf_counter() - t_start + result = BenchmarkResult("🗺️ segmentation", processed, total, batch_times) + result.log_summary() + return result + + +def benchmark_full_pipeline( + depth_model: Any, + seg_model: Any, + seg_config: dict[str, Any], + device: torch.device, + num_images: int = 100, + batch_size: int = 4, + image_size: int = 256, + num_workers: int = 0, + warmup_batches: int = 3, +) -> dict[str, BenchmarkResult]: + """Benchmark all stages and total pipeline throughput.""" + from src.augmentor.inference import ( + infer_depth_batch, compute_edges_from_depth, infer_segmentation_batch, + ) + + ds = SyntheticImageDataset(num_images, image_size) + loader = DataLoader(ds, batch_size=batch_size, shuffle=False, + num_workers=num_workers, pin_memory=torch.cuda.is_available()) + + # Warmup + for i, batch in enumerate(loader): + if i >= warmup_batches: + break + imgs = batch["image_raw"] + depths = infer_depth_batch(depth_model, imgs, device) + compute_edges_from_depth(depths) + infer_segmentation_batch(seg_model, seg_config, imgs, device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + + depth_times, edge_times, seg_times, total_times = [], [], [], [] + processed = 0 + t_pipeline_start = time.perf_counter() + + for batch in tqdm(loader, desc="⏱️ bench full pipeline", unit="batch", colour="magenta"): + imgs = batch["image_raw"] + cur_bs = imgs.shape[0] + + t_total_0 = time.perf_counter() + + # Depth + t0 = time.perf_counter() + depths = infer_depth_batch(depth_model, imgs, device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + depth_times.append(time.perf_counter() - t0) + + # Edges + t0 = time.perf_counter() + compute_edges_from_depth(depths) + edge_times.append(time.perf_counter() - t0) + + # Segmentation + t0 = time.perf_counter() + infer_segmentation_batch(seg_model, seg_config, imgs, device) + if torch.cuda.is_available(): + torch.cuda.synchronize() + seg_times.append(time.perf_counter() - t0) + + total_times.append(time.perf_counter() - t_total_0) + processed += cur_bs + + t_pipeline_total = time.perf_counter() - t_pipeline_start + + results = { + "depth": BenchmarkResult("🌊 depth", processed, sum(depth_times), depth_times), + "edges": BenchmarkResult("🔪 edges", processed, sum(edge_times), edge_times), + "segmentation": BenchmarkResult("🗺️ segmentation", processed, sum(seg_times), seg_times), + "total": BenchmarkResult("🏁 total pipeline", processed, t_pipeline_total, total_times), + } + + logger.info("=" * 60) + for r in results.values(): + r.log_summary() + logger.info("=" * 60) + + return results diff --git a/src/utils/configure_gpu.py b/src/utils/configure_gpu.py new file mode 100644 index 0000000..9663d32 --- /dev/null +++ b/src/utils/configure_gpu.py @@ -0,0 +1,51 @@ +import os +import torch + +def configure_gpu(target_gpu: int = 0, + max_mem_fraction: float = 0.40, + preallocate_megabytes: int | None = None) -> None: + """ + Configure GPU visibility and memory use for a PyTorch process. + + Parameters + ---------- + target_gpu : int + Index of the GPU that this process should see (0-based). + If you want to let PyTorch pick any free GPU, skip this argument. + max_mem_fraction : float + Fraction (0–1) of that GPU's total memory PyTorch may allocate. + Uses torch.cuda.set_per_process_memory_fraction (PyTorch ≥ 1.7). + preallocate_megabytes : int | None + If set, immediately reserve this many MB with an empty tensor so + PyTorch never exceeds the requested size. + """ + # 1 — Hide all other GPUs + os.environ["CUDA_VISIBLE_DEVICES"] = str(target_gpu) + + if not torch.cuda.is_available(): + print("CUDA not available, using CPU only.") + return + + device_index = 0 # because after CUDA_VISIBLE_DEVICES we see only one GPU + torch.cuda.set_device(device_index) + + # 2 — Limit allocator to a fraction of free memory + torch.cuda.set_per_process_memory_fraction(max_mem_fraction, device_index) + + # 3 — Optionally pre-allocate a fixed block + if preallocate_megabytes: + dummy = torch.empty( + (preallocate_megabytes * 256), # 4 bytes × 256 = 1 MB + dtype=torch.float32, + device=f"cuda:{device_index}" + ) + del dummy # immediately free it; allocator keeps the block reserved + + # Sanity print + total_mem = torch.cuda.get_device_properties(device_index).total_memory / 1024**3 + print(f"GPU {target_gpu} visible, capped at {max_mem_fraction*100:.0f}% " + f"≈ {total_mem*max_mem_fraction:.1f} GiB usable by PyTorch.") + +# Example usage --------------------------------------------------------------- +if __name__ == "__main__": + configure_gpu(target_gpu=0, max_mem_fraction=0.2, preallocate_megabytes=4096) \ No newline at end of file diff --git a/src/utils/profiler.py b/src/utils/profiler.py new file mode 100644 index 0000000..428723c --- /dev/null +++ b/src/utils/profiler.py @@ -0,0 +1,244 @@ +"""Hardware & model profiling: system info, VRAM/RAM/HDD, GFLOPs/MACs. + +Uses fvcore, thop, psutil, and torch.cuda for comprehensive reporting. +""" + +from __future__ import annotations + +import logging +import platform +import shutil +from pathlib import Path +from typing import Any + +import psutil +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Formatting helpers +# --------------------------------------------------------------------------- + +def _fmt_bytes(b: float) -> str: + """Human-readable byte size.""" + for unit in ("B", "KB", "MB", "GB", "TB"): + if abs(b) < 1024: + return f"{b:.1f} {unit}" + b /= 1024 + return f"{b:.1f} PB" + + +def _fmt_number(n: float) -> str: + """Human-readable large number (1.2M, 3.4B, etc.).""" + if n >= 1e9: + return f"{n / 1e9:.2f} B" + if n >= 1e6: + return f"{n / 1e6:.2f} M" + if n >= 1e3: + return f"{n / 1e3:.1f} K" + return str(int(n)) + + +# --------------------------------------------------------------------------- +# System info +# --------------------------------------------------------------------------- + +def log_system_info() -> dict[str, Any]: + """Log and return system hardware summary.""" + info: dict[str, Any] = {} + + # CPU + cpu_name = platform.processor() or platform.machine() + cpu_cores_phys = psutil.cpu_count(logical=False) or 0 + cpu_cores_logic = psutil.cpu_count(logical=True) or 0 + info["cpu"] = { + "name": cpu_name, + "cores_physical": cpu_cores_phys, + "cores_logical": cpu_cores_logic, + } + + # RAM + mem = psutil.virtual_memory() + info["ram"] = { + "total": mem.total, + "available": mem.available, + "used": mem.used, + "percent": mem.percent, + } + + # GPU + if torch.cuda.is_available(): + gpu_id = torch.cuda.current_device() + gpu_name = torch.cuda.get_device_name(gpu_id) + gpu_total = torch.cuda.get_device_properties(gpu_id).total_memory + gpu_reserved = torch.cuda.memory_reserved(gpu_id) + gpu_allocated = torch.cuda.memory_allocated(gpu_id) + gpu_free = gpu_total - gpu_reserved + info["gpu"] = { + "name": gpu_name, + "total": gpu_total, + "allocated": gpu_allocated, + "reserved": gpu_reserved, + "free": gpu_free, + } + else: + info["gpu"] = None + + # Log it + logger.info("🖥️ System info:") + logger.info(" 🧠 CPU: %s (%d physical / %d logical cores)", + cpu_name, cpu_cores_phys, cpu_cores_logic) + logger.info(" 💾 RAM: %s used / %s total (%.1f%% used)", + _fmt_bytes(mem.used), _fmt_bytes(mem.total), mem.percent) + + if info["gpu"]: + g = info["gpu"] + logger.info(" 🎮 GPU: %s", g["name"]) + logger.info(" VRAM: %s allocated / %s reserved / %s total (%s free)", + _fmt_bytes(g["allocated"]), _fmt_bytes(g["reserved"]), + _fmt_bytes(g["total"]), _fmt_bytes(g["free"])) + else: + logger.info(" 🎮 GPU: not available") + + return info + + +def log_disk_info(*paths: Path) -> dict[str, Any]: + """Log disk usage for given paths (or / by default).""" + targets = paths or (Path("/"),) + info: dict[str, Any] = {} + + for p in targets: + try: + usage = shutil.disk_usage(p) + key = str(p) + info[key] = { + "total": usage.total, + "used": usage.used, + "free": usage.free, + } + pct = usage.used / usage.total * 100 + logger.info(" 💿 Disk [%s]: %s used / %s total (%.1f%%, %s free)", + p, _fmt_bytes(usage.used), _fmt_bytes(usage.total), + pct, _fmt_bytes(usage.free)) + except OSError as e: + logger.warning(" 💿 Disk [%s]: unavailable (%s)", p, e) + return info + + +# --------------------------------------------------------------------------- +# VRAM snapshot +# --------------------------------------------------------------------------- + +def log_vram_snapshot(label: str = "") -> dict[str, float] | None: + """Log current VRAM usage. Returns dict or None if no GPU.""" + if not torch.cuda.is_available(): + return None + + gpu_id = torch.cuda.current_device() + allocated = torch.cuda.memory_allocated(gpu_id) + reserved = torch.cuda.memory_reserved(gpu_id) + total = torch.cuda.get_device_properties(gpu_id).total_memory + free = total - reserved + + prefix = f"[{label}] " if label else "" + logger.info(" 🎮 %sVRAM: %s allocated / %s reserved / %s free", + prefix, _fmt_bytes(allocated), _fmt_bytes(reserved), + _fmt_bytes(free)) + return {"allocated": allocated, "reserved": reserved, "free": free, "total": total} + + +def log_ram_snapshot(label: str = "") -> dict[str, float]: + """Log current RAM usage.""" + mem = psutil.virtual_memory() + prefix = f"[{label}] " if label else "" + logger.info(" 💾 %sRAM: %s used / %s available / %s total (%.1f%%)", + prefix, _fmt_bytes(mem.used), _fmt_bytes(mem.available), + _fmt_bytes(mem.total), mem.percent) + return {"used": mem.used, "available": mem.available, + "total": mem.total, "percent": mem.percent} + + +# --------------------------------------------------------------------------- +# Model profiling: GFLOPs / MACs / params +# --------------------------------------------------------------------------- + +def profile_model( + model: nn.Module, + input_size: tuple[int, ...] | None, + device: torch.device | None = None, + model_name: str = "", +) -> dict[str, Any]: + """Profile model FLOPs/MACs/params using fvcore + thop. + + Args: + model: PyTorch model. + input_size: (B, C, H, W) input shape, or None to skip FLOPs/MACs. + device: Device for dummy input. + model_name: Label for log messages. + + Returns: + Dict with params, flops_fvcore, macs_thop, etc. + """ + label = model_name or model.__class__.__name__ + result: dict[str, Any] = {"name": label} + + # Params + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + param_bytes = sum(p.numel() * p.element_size() for p in model.parameters()) + result["total_params"] = total_params + result["trainable_params"] = trainable_params + result["param_bytes"] = param_bytes + + logger.info("📊 Model profile: %s", label) + logger.info(" 📐 Parameters: %s total (%s trainable), %s in memory", + _fmt_number(total_params), _fmt_number(trainable_params), + _fmt_bytes(param_bytes)) + + if input_size is None: + result["flops_fvcore"] = None + result["macs_thop"] = None + return result + + # Determine device and dtype from model + if device is None: + try: + device = next(model.parameters()).device + except StopIteration: + device = torch.device("cpu") + dtype = next(model.parameters()).dtype if total_params > 0 else torch.float32 + + dummy = torch.randn(*input_size, device=device, dtype=dtype) + + # fvcore FLOPs + try: + from fvcore.nn import FlopCountAnalysis + flop_counter = FlopCountAnalysis(model, dummy) + flop_counter.unsupported_ops_warnings(False) + flop_counter.uncalled_modules_warnings(False) + flops = flop_counter.total() + result["flops_fvcore"] = flops + result["gflops_fvcore"] = flops / 1e9 + logger.info(" ⚡ FLOPs (fvcore): %s (%.2f GFLOPs)", + _fmt_number(flops), flops / 1e9) + except Exception as e: + logger.warning(" ⚡ fvcore FLOPs failed: %s", e) + result["flops_fvcore"] = None + + # thop MACs + try: + from thop import profile as thop_profile + macs, params = thop_profile(model, inputs=(dummy,), verbose=False) + result["macs_thop"] = macs + result["gmacs_thop"] = macs / 1e9 + logger.info(" ⚡ MACs (thop): %s (%.2f GMACs)", + _fmt_number(macs), macs / 1e9) + except Exception as e: + logger.warning(" ⚡ thop MACs failed: %s", e) + result["macs_thop"] = None + + return result diff --git a/src/utils/utils_file_dir.py b/src/utils/utils_file_dir.py new file mode 100644 index 0000000..545fff9 --- /dev/null +++ b/src/utils/utils_file_dir.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from pathlib import Path + +_ROOT_MARKERS = ("pyproject.toml", ".git", "in") + + +def get_proj_dir() -> str: + """Return project root directory with trailing slash. + + Walks up from this file's location until a directory containing + one of the root markers is found. + + Root markers (checked in order): pyproject.toml, .git, in/. + + Returns: + Absolute path to project root with trailing '/'. + + Raises: + RuntimeError: If no marker is found within 10 parent levels. + """ + current = Path(__file__).resolve().parent + for _ in range(10): + if any((current / marker).exists() for marker in _ROOT_MARKERS): + return str(current) + "/" + current = current.parent + raise RuntimeError( + f"Cannot find project root: none of {_ROOT_MARKERS} found " + f"in ancestors of {Path(__file__).resolve()}" + ) diff --git a/test_seg_12classes.py b/test_seg_12classes.py new file mode 100644 index 0000000..7c5a68e --- /dev/null +++ b/test_seg_12classes.py @@ -0,0 +1,114 @@ +"""Quick test: run 12-class segmentation on 10 diverse images and save results.""" + +from __future__ import annotations + +import sys +import shutil +from pathlib import Path + +import torch +import numpy as np +from PIL import Image + +# Ensure project root is on path. +ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(ROOT / "src" / "nn" / "segearth_ov3")) + +from src.conf.config_loader import load_all_configs +from src.augmentor.models import load_segmentation_model +from src.augmentor.inference import infer_segmentation_batch +from src.augmentor.io_utils import save_segmentation + +INPUT_ROOT = Path("/mnt/data1tb/cvgl_datasets/UAV-GeoLoc") + +# 10 diverse test images: satellite (DB) — urban + terrain +TEST_IMAGES = [ + # Urban (satellite) + "Country/German/Munich/Ludwigs/DB/img/crop_12_4.png", + "Country/French/Paris/LeMarais/DB/img/crop_12_4.png", + "Country/Italy/Venice/SanMarco/DB/img/crop_12_4.png", + "Country/USA/NewYork/Manhattan/DB/img/crop_12_4.png", + "Country/Australia/Sydney/SydneyCBD/DB/img/crop_12_4.png", + "Country/Korea/Busan/Gwangalli/DB/img/crop_12_4.png", + # Terrain (satellite) + "Terrain/Desert/GobiDesert/DB/img/crop_50_33.png", + "Terrain/Glacier/AthabascaGlacier/DB/img/crop_12_4.png", + "Terrain/Volcano/KilaueaVolcano/DB/img/crop_12_4.png", + "Terrain/Danxia/GrandCanyon/DB/img/crop_50_33.png", +] + +OUTPUT_DIR = ROOT / "test_seg_output" + + +def main(): + # Load configs (reads all .gin files) + configs = load_all_configs(str(ROOT / "in" / "config_files") + "/") + pipeline_conf = configs["pipeline"] + hw_conf = configs["hardware"] + models_conf = configs["models"] + input_conf = configs["input"] + seg_conf = configs["seg"] + + print(f"Prompts ({len(seg_conf.prompts)}): {seg_conf.prompts}") + print(f"Threshold: {seg_conf.threshold}") + print(f"Resolution: {seg_conf.default_resolution}") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + + # Load model + print("\nLoading SegEarth-OV3...") + model, seg_config = load_segmentation_model(models_conf, hw_conf, seg_conf, device) + num_classes = seg_config.get("num_classes", 150) + print(f"Model loaded. Type: {seg_config['type']}, num_classes: {num_classes}") + + # Prepare output dir + OUTPUT_DIR.mkdir(exist_ok=True) + + # Process each image + for i, rel_path in enumerate(TEST_IMAGES): + img_path = INPUT_ROOT / rel_path + if not img_path.exists(): + print(f"[{i+1}/10] SKIP (not found): {rel_path}") + continue + + print(f"\n[{i+1}/10] Processing: {rel_path}") + + # Load and preprocess + pil_img = Image.open(img_path).convert("RGB") + img_np = np.array(pil_img).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W] + + # Infer + seg_ids = infer_segmentation_batch(model, seg_config, img_tensor, device) + # seg_ids: [1, 1, H, W] uint8 + + # Save original + segmentation side by side + stem = img_path.stem + tag = rel_path.split("/")[0] # Country or Terrain + if tag == "Country": + tag = rel_path.split("/")[2] # city name + elif tag == "Terrain": + tag = rel_path.split("/")[1] # terrain type + + out_stem = f"{i:02d}_{tag}_{stem}" + + # Save original for reference + pil_img.save(OUTPUT_DIR / f"{out_stem}_orig.png") + + # Save segmentation vis + save_segmentation( + seg_ids[0], OUTPUT_DIR, out_stem, + save_npy=False, save_vis=True, num_classes=num_classes, + ) + + unique_classes = torch.unique(seg_ids[0]).tolist() + class_names = [seg_conf.prompts[c] for c in unique_classes if c < len(seg_conf.prompts)] + print(f" -> Classes found: {class_names}") + + print(f"\nDone! Results saved to: {OUTPUT_DIR}") + + +if __name__ == "__main__": + main() diff --git a/Методология_автоматического_аннотирования_World_UAV.md b/Методология_автоматического_аннотирования_World_UAV.md new file mode 100644 index 0000000..e5d2a0f --- /dev/null +++ b/Методология_автоматического_аннотирования_World_UAV.md @@ -0,0 +1,652 @@ +--- +title: "Методология автоматического аннотирования набора данных World-UAV: карты глубины, контурные карты, семантические маски и кросс-вью карты высот" +type: methodology +version: "3.3" +created: 2026-04-07 +updated: 2026-04-08 +dataset: "World-UAV (UAV-GeoLoc, IEEE RA-L 2025)" +hardware: "RTX 4090 (24GB VRAM)" +downstream_task: "Cross-view geo-localization (drone<->satellite)" +models: + depth: "Depth Anything V3 LARGE-1.1 (411M)" + edges: "Sobel filters from depth" + segmentation: "SegEarth-OV3 (SAM 3.1 + PE-L+)" + chmv2: "DINOv3-ViTL16-CHMv2-DPT (337M, FP32)" +tags: + - type/methodology + - component/dataset + - method/segmentation + - dataset/worlduav + - task/experiment + - priority/high +--- + +# Методология автоматического аннотирования набора данных World-UAV + +**Содержание (12 разделов):** + +- **Раздел 1.** Постановка задачи: формальная формулировка, связь с LUPI-дистилляцией +- **Раздел 2.** World-UAV: структура, объем, неполные сцены +- **Раздел 3.** Модели-генераторы: DA3-LARGE, Sobel, SegEarth-OV3, CHMv2 +- **Раздел 4.** Пайплайн обработки: sequential pipeline, формат выхода, manifest +- **Раздел 5.** Протокол обработки: оценка времени, запуск +- **Раздел 6.** Контроль качества: валидация, визуальная инспекция, метрики +- **Раздел 7.** Формат хранения и интеграция с даталоадером +- **Раздел 8.** Гипотезы аннотирования (AH1-AH6) +- **Раздел 9.** Архитектура кода +- **Раздел 10.** Риски и митигации +- **Раздел 11.** Timeline и ресурсы +- **Раздел 12.** Библиография + +> **Цель:** разработать воспроизводимую методологию автоматической генерации четырех вспомогательных модальностей (depth, edges, segmentation, cross-view height map) из RGB-изображений набора данных World-UAV для последующего использования в LUPI-дистилляции системы NADEZHDA. + +--- + +## 1. Постановка задачи + +### 1.1. Формальная формулировка + +Пусть $I \in \mathbb{R}^{H \times W \times 3}$ -- RGB-изображение из World-UAV (drone query или satellite DB crop). Задача автоматического аннотирования формулируется как генерация четырех карт-аннотаций: + +$$\mathcal{A}: I \mapsto (D, E, S, M)$$ + +где: +- $D \in \mathbb{R}^{H' \times W'}$ -- карта глубины (monocular depth), $D_{ij} \in [0, 1]$ (per-frame normalized) +- $E \in \mathbb{R}^{H' \times W'}$ -- контурная карта (structural edges), $E_{ij} \in [0, 1]$ (edge magnitude) +- $S \in \mathbb{Z}^{H' \times W'}$ -- семантическая сегментационная маска, $S_{ij} \in \{0, \ldots, C-1\}$ (class index) +- $M \in \mathbb{R}^{H' \times W'}$ -- кросс-вью карта высот (CHMv2), $M_{ij} \in [0, 1]$ (per-frame normalized) + +### 1.2. Связь с downstream-задачей (LUPI-CVGL) + +Сгенерированные аннотации используются как **привилегированная информация** при обучении Teacher модели системы NADEZHDA: + +- **Teacher** (DINOv3-L, ~356M) обучается на расширенном наборе модальностей: $\{I_{sat}, I_{drone}, D+E, S, M\}$ +- **Student** (SOFIA, ~6-8M) при инференсе использует только 2: $\{I_{sat}, I_{drone}\}$ + +Привилегированные модальности ($D$, $E$, $S$, $M$) доступны **только при обучении** Teacher. Knowledge distillation переносит эту информацию в Student через Loss LUPI: + +$$L_{LUPI} = \|z_{Student}^{drone} - \text{sg}(z_{Teacher}^{fused})\|^2$$ + +### 1.3. Почему edges, а не normals + +| Критерий | Edges | Normals | +|:---------|:------|:--------| +| Структурные контуры (дороги, здания) | **Выделяет** | Не фокусируется | +| Инвариантность к освещению | **Высокая** | Зависит от depth quality | +| Cross-view matching utility | **Высокая** -- границы видны с обоих ракурсов | Низкая -- нормали шумные для наклонных снимков | +| Computational cost | **~1 мс/img** (Sobel, CPU) | ~5 мс/img (DSINE, GPU) | +| Число каналов | **1** (magnitude) | 3 (xyz-вектор) | + +### 1.4. Зачем CHMv2 (cross-view height map) + +CHMv2 (DINOv3-ViTL16-CHMv2-DPT) обучен на DINOv3 backbone специально для оценки высот из спутниковых и аэроснимков. В отличие от DA3, который обучен преимущественно на наземных изображениях: + +- **CHMv2** лучше работает на nadir-view (спутниковые патчи) -- обучен на аэрофотосъемке +- **DA3** лучше работает на oblique-view (drone кадры) -- обучен на ground-level +- Комбинация двух моделей глубины ($D$ + $M$) дает Teacher **ортогональные depth representations** + +> **Важно:** CHMv2 **нестабилен в FP16** (выдает NaN). Модель всегда загружается в FP32. + +--- + +## 2. Набор данных: World-UAV (UAV-GeoLoc) + +### 2.1. Обзор + +| Параметр | Значение | +|:---------|:---------| +| Venue | IEEE Robotics and Automation Letters, Q1 | +| Авторы | Wu et al. (NUDT + Zhejiang University) | +| Число сцен | **372** (171 Country + 200 Terrain + 1 Rot) | +| Всего изображений | **973,256** (фактически найдено при обработке) | +| Query (drone) | 652,744 (512x512 JPEG, синтетические Google Earth Studio) | +| DB (satellite) | 274,683 (PNG, 11 размеров: 100-1000 px) | +| Subset Rot | 7,341 изображений | +| Subset Country | 455,638 изображений | +| Subset Terrain | 510,277 изображений | +| Размер на диске | ~181 GB | + +### 2.2. Неполные сцены (исключаются из обработки) + +**16 сцен Country** не содержат DB-кропов, positive.json: + +| Город | Сцены | Кол-во | +|:------|:------|:-------| +| Edinburgh | CastleHill, Dalry, Haymarket, NewTown, Stockbridge | 5 | +| London | CamdenTown, CoventGarden, Fitzrovia, Mayfair, SoHo | 5 | +| Manchester | Ancoats, Castlefield, Deansgate, NorthernQuarter, Piccadilly | 5 | +| Birmingham | JewelleryQuarter | 1 | + +Дополнительно исключаются: `Index`, `charts`, `__MACOSX`. + +--- + +## 3. Модели-генераторы + +### 3.1. Depth: Depth Anything V3 LARGE-1.1 + +| Параметр | Значение | +|:---------|:---------| +| Модель | DA3-LARGE-1.1 (ViT-based unified architecture) | +| HuggingFace ID | `depth-anything/DA3-LARGE-1.1` | +| Параметры | **410.94M** | +| Precision | FP16 (при `use_fp16=True`) | +| VRAM | **1.5 GB** weights | +| Inference speed | **18.4 img/s** (batch=4, 256x256, RTX 4090) -- замерено | +| Выход | grayscale PNG [256, 256] uint8 [0, 255] (per-frame normalized) | +| Fallback | DA V2 Large (transformers API) | +| Код | `src/nn/depth_anything_3/` (вендорированный пакет) | +| Веса | `in/weights/models--depth-anything--DA3-LARGE-1.1/` | + +**Архитектура DA3:** DinoV2 backbone (ViT-Large, 24 блока, 1024 dim) + DualDPT head (depth + ray prediction). Поддерживает multi-view inference с camera conditioning, но в пайплайне используется monocular mode (single image). + +### 3.2. Edges: Sobel из карты глубины + +| Параметр | Значение | +|:---------|:---------| +| Метод | Sobel filters (3x3) на depth map | +| Реализация | `torch.nn.functional.conv2d` с Sobel-ядрами | +| Inference speed | **419.6 img/s** (CPU, batch=32) -- замерено | +| Выход | grayscale PNG [256, 256] uint8 [0, 255] | +| Зависимость | Edges вычисляются из depth в том же batch-цикле | + +**Алгоритм:** + +$$S_x = \frac{1}{8}\begin{bmatrix} -1 & 0 & 1 \\ -2 & 0 & 2 \\ -1 & 0 & 1 \end{bmatrix}, \quad E = \sqrt{(S_x * D)^2 + (S_y * D)^2}, \quad E = E / \max(E)$$ + +### 3.3. Segmentation: SegEarth-OV3 (SAM 3.1) + +| Параметр | Значение | +|:---------|:---------| +| Модель | SegEarth-OV3 (training-free OV semantic segmentation) | +| Backbone | SAM 3.1 Multiplex (~3.3 GB checkpoint) | +| VRAM | **3.2 GB** (batch=1), до **~10 GB** (batch=8) | +| Inference speed (per-image) | **3.9 img/s** (batch=1, без батчинга backbone) -- замерено | +| Inference speed (batched) | **~10-12 img/s** (batch=8, batched backbone) -- оценка | +| Макс. batch size | **8** (ограничение в `_infer_segearth_ov3`, hardcoded) | +| Промпты | 5 классов: background, building, road, vegetation, water | +| Fallback | SegFormer-B5 (ADE20K, 150 classes) | +| Код | `src/nn/segearth_ov3/` (вендорированный пакет) | +| Веса | `in/weights/sam3.1/sam3.1_multiplex.pt` | +| BPE vocab | `src/nn/segearth_ov3/sam3/assets/bpe_simple_vocab_16e6.txt.gz` (встроен) | + +**Архитектура SegEarth-OV3:** SAM3 backbone (ViT-H, 1008x1008, RoPE) + VETextEncoder (CLIP-like, BPE tokenizer) + TransformerEncoderFusion (6 layers, cross-attention) + TransformerDecoder (200 object queries) + PixelDecoder (mask prediction). Для каждого текстового промпта выполняется grounding: fusion instance masks (decoder) + semantic masks (segmentation head), фильтрация по presence score. + +**Batched backbone inference (v3.2):** + +Основной bottleneck SegEarth-OV3 -- backbone forward pass (SAM 3.1 ViT), занимающий ~80% времени инференса. Реализован `predict_pil_batch`: + +1. **`set_image_batch(images)`** -- один backbone pass на весь батч (вместо B отдельных `set_image()`) +2. **`_slice_backbone_out(backbone_out, i)`** -- слайсинг backbone features `[B, 256, H, W] -> [1, 256, H, W]` для каждого изображения +3. **Per-image grounding** -- для каждого изображения: цикл по 5 промптам -> легкий grounding decoder + text encoder + +Grounding decoder (~20% времени) остается per-image, но backbone (~80%) амортизируется на весь батч. При OOM или ошибке -- автоматический fallback на per-image `predict_pil()`. + +**Палитра сегментации:** + +| Класс | ID | Цвет | RGB | +|:------|:--:|:-----|:----| +| background | 0 | Черный | (0, 0, 0) | +| building | 1 | Красный | (220, 40, 40) | +| road | 2 | Серый | (160, 160, 160) | +| vegetation | 3 | Зеленый | (30, 180, 30) | +| water | 4 | Синий | (30, 120, 220) | + +### 3.4. CHMv2: DINOv3-ViTL16-CHMv2-DPT + +| Параметр | Значение | +|:---------|:---------| +| Модель | CHMv2ForDepthEstimation (DPT head на DINOv3 backbone) | +| HuggingFace ID | `facebook/dinov3-vitl16-chmv2-dpt-head` | +| Параметры | **336.88M** | +| FLOPs | **457.98 GFLOPs** (fvcore, 518x518 input) | +| Precision | **FP32 only** (FP16 выдает NaN -- обнаружено при валидации) | +| VRAM | **1.3 GB** weights (FP32) | +| Inference speed | **31.7 img/s** (batch=4, 256x256, RTX 4090) -- замерено | +| Выход | grayscale PNG [256, 256] uint8 [0, 255] (per-frame normalized) | +| Веса | `in/weights/dinov3-chmv2/` (safetensors, 1.3 GB) | + +### 3.5. Сравнение моделей (benchmark на 100 изображениях, RTX 4090) + +| | DA3-LARGE | Sobel edges | SegEarth-OV3 (batch=1) | SegEarth-OV3 (batch=8) | CHMv2 (DINOv3) | +|:---|:---------|:-----------|:----------------------|:----------------------|:---------------| +| Params | 411M | 0 | ~1B+ (SAM 3.1) | ~1B+ (SAM 3.1) | 337M | +| VRAM | 1.5 GB | 0 | 3.2 GB | ~10 GB | 1.3 GB | +| Precision | FP16 | CPU | BF16 | BF16 | **FP32** | +| Speed | **18.4 img/s** | **419.6 img/s** | **3.9 img/s** | **~10-12 img/s** (оценка) | **31.7 img/s** | +| 100 img time | 5.44 s | 0.24 s | **25.49 s** | **~8-10 s** (оценка) | 3.15 s | +| Peak VRAM (sequential) | | | | **~10 GB** (SegEarth batch=8 -- наибольшая из загружаемых) | | + +--- + +## 4. Пайплайн обработки + +### 4.1. Sequential pipeline (одна модель за раз) + +Модели загружаются **последовательно** -- каждая стадия загружает свою модель, обрабатывает все изображения, затем выгружает модель перед следующей стадией. Это позволяет использовать максимальный batch size для каждой модели и упрощает resume по стадиям: + +``` +Stage 1: DEPTH + +-- load DA3-LARGE-1.1 [GPU, ~1.5 GB FP16] + +-- process ALL images (batched, auto_batch_size) + +-- save_async {stem}_depth.png [ThreadPool, 4 workers] + +-- unload_model() [del + torch.cuda.empty_cache()] + +Stage 2: EDGES + +-- load saved depth from .npy/.png [CPU] + +-- compute Sobel edges (CPU, batch=32) + +-- save_async {stem}_edge.png + +-- (no GPU model to unload) + +Stage 3: SEGMENTATION + +-- load SegEarth-OV3 (SAM 3.1) [GPU, ~3.2-10 GB] + +-- process ALL images (batched backbone, max_batch=8, ~10-12 img/s) + | +-- predict_pil_batch(): 1 backbone pass на <=8 изображений + | +-- _slice_backbone_out(): слайс features per-image + | +-- grounding decoder: per-image x 5 промптов (легкий) + +-- save_async {stem}_segm.png + +-- unload_model() + +Stage 4: CHMv2 + +-- load DINOv3-ViTL16-CHMv2-DPT [GPU, ~1.3 GB FP32] + +-- process ALL images (batched, auto_batch_size) + +-- save_async {stem}_chm.png + +-- unload_model() +``` + +### 4.2. Формат выхода + +Файлы сохраняются в `World-UAV-aug/` с **сохранением структуры** исходного датасета. Исходные изображения **не копируются**. Имя файла = имя исходного изображения + суффикс стадии: + +``` +World-UAV-aug/ ++-- Rot/SouthernSuburbs/DB/img/ +| +-- crop_12_4_depth.png # grayscale [256, 256] -- DA3 depth +| +-- crop_12_4_edge.png # grayscale [256, 256] -- Sobel edges +| +-- crop_12_4_segm.png # RGB palette [256, 256, 3] -- semantic seg +| +-- crop_12_4_chm.png # grayscale [256, 256] -- CHMv2 height ++-- Country/... ++-- Terrain/... +``` + +### 4.3. Суффиксы и формат файлов + +| Стадия | Суффикс | Формат PNG | Для обучения | +|:-------|:--------|:-----------|:-------------| +| depth | `_depth` | grayscale (L) uint8 | `value = pixel / 255.0` -> float [0,1], 1 канал | +| edges | `_edge` | grayscale (L) uint8 | `value = pixel / 255.0` -> float [0,1], 1 канал | +| segmentation | `_segm` | RGB palette (3ch) | Загружать class ID из int-маски, не из RGB | +| chmv2 | `_chm` | grayscale (L) uint8 | `value = pixel / 255.0` -> float [0,1], 1 канал | + +> **Для обучения:** depth, edge, chm используются как grayscale 1-канальные float [0,1]. Colormap-визуализации не нужны -- они теряют информацию при обратном преобразовании. Segm используется через class index (0-4), не через RGB. + +### 4.4. Конфигурация (Gin framework) + +Все параметры задаются в `in/config_files/*.gin`: + +| Файл | Ключевые параметры | +|:-----|:-------------------| +| `pipeline.gin` | input_root, output_root, stages, save_npy, save_vis, save_concat, resume, subset, source, log_level | +| `models.gin` | depth_model_id, depth_fallback_id, chmv2_model_id, seg_model_type, seg_fallback_id, weights_dir | +| `hardware.gin` | profile_name, total_ram_gb, reserve_gb, use_fp16, batch_size (None=auto), num_workers | +| `segmentation.gin` | prompts, threshold (0.3), default_resolution (1008) | +| `input.gin` | image_size (256), sobel_kernel_size (3), edge_normalize, imagenet_mean/std | + +### 4.5. Manifest (manifest.json) + +По завершении пайплайна генерируется `{output_root}/manifest.json` с полной метаинформацией: + +```json +{ + "pipeline_version": "3.3.0-sequential", + "image_size": 256, + "profile": "rtx4090", + "models": { + "depth": "DA3-LARGE-1.1", + "edges": "Sobel from depth (CPU)", + "segmentation": "segearth-ov3", + "chmv2": "facebook/dinov3-vitl16-chmv2-dpt-head" + }, + "seg_prompts": ["background", "building", "road", "vegetation", "water"], + "total_images": 973256, + "stages": { + "depth": {"processed": "...", "time_sec": "..."}, + "edges": {"processed": "...", "time_sec": "..."}, + "segmentation": {"processed": "...", "time_sec": "..."}, + "chmv2": {"processed": "...", "time_sec": "..."} + }, + "timestamp": "..." +} +``` + +### 4.6. Auto batch size + +Batch size определяется автоматически по свободной VRAM: + +```python +free_mb = (cuda_total - cuda_reserved) / (1024*1024) - overhead_mb +batch_size = round_down_power_of_2(free_mb / act_per_sample_mb * 0.7) +``` + +Можно зафиксировать вручную через `hardware.gin` (`batch_size = 8`). При `batch_size = None` используется автоматический расчет. + +### 4.7. Подавление логов + +Все verbose логи моделей подавлены: +- `DA3_LOG_LEVEL=ERROR` -- кастомный DA3 Logger (не стандартный Python logging) +- `TRANSFORMERS_VERBOSITY=error` -- transformers +- `HF_HUB_DISABLE_PROGRESS_BARS=1` -- tqdm загрузки весов +- Все сторонние Python loggers на уровень `ERROR` + +--- + +## 5. Протокол обработки + +### 5.1. Оценка времени (экстраполяция из benchmark 100 img) + +**v3.1 (per-image SegEarth-OV3):** + +| Стадия | 100 img | 973K img | % общего времени | +|:-------|:--------|:---------|:-----------------| +| Depth (DA3-LARGE) | 5.44 s | **14.7 ч** | 6% | +| Edges (Sobel) | 0.24 s | **0.6 ч** | <1% | +| Segmentation (SegEarth-OV3, batch=1) | 25.49 s | **69 ч** (~2.9 дней) | **74%** | +| CHMv2 (DINOv3) | 3.15 s | **8.5 ч** | 4% | +| **Итого** | 34.33 s | **~93 ч** (~3.9 дней) | 100% | + +**v3.2+ (batched backbone, batch<=8) -- оценка:** + +| Стадия | 100 img | 973K img | % общего времени | +|:-------|:--------|:---------|:-----------------| +| Depth (DA3-LARGE) | 5.44 s | **14.7 ч** | 14% | +| Edges (Sobel) | 0.24 s | **0.6 ч** | <1% | +| Segmentation (SegEarth-OV3, batch<=8) | ~8-10 s | **~22-27 ч** (~1 день) | **~55%** | +| CHMv2 (DINOv3) | 3.15 s | **8.5 ч** | 8% | +| **Итого** | ~17-19 s | **~46-51 ч** (~2 дней) | 100% | + +> **Примечание:** оценки для batched backbone теоретические (backbone ~80% времени амортизируется на батч, grounding ~20% остается per-image). Реальные цифры зависят от VRAM overhead и размера изображений. Требуется benchmark на реальных данных. + +### 5.2. Запуск + +```bash +# Полный пайплайн (все subset'ы, все стадии, resume=True): +python -m src.main +``` + +Resume проверяет наличие `.png` **или** `.npy` файлов **для каждой стадии отдельно**. Пайплайн можно прервать и перезапустить в любой момент -- уже обработанные изображения пропускаются. В sequential-архитектуре resume гранулярный: можно перезапустить только конкретную стадию (например, segmentation), не затрагивая уже завершенные стадии (depth, edges). + +### 5.3. Оценка дискового пространства + +| Режим | Per-image | 973K images | +|:------|:---------|:-----------| +| Только PNG (save_npy=False) | ~60 KB x 4 | **~228 GB** | +| PNG + NPY (save_npy=True) | ~508 KB | ~484 GB | + +--- + +## 6. Контроль качества + +### 6.1. Автоматическая валидация + +| Проверка | Уровень | Механизм | +|:---------|:--------|:---------| +| Файл существует | Per-image | `_filter_all_done()` при resumption | +| Shape корректен | Per-stage | Проверка при inference | +| Range корректен | Per-stage | depth/edges/chm in [0,1], seg in [0, C-1] | +| NaN detection | Per-stage | CHMv2 в FP32 (FP16 выдает NaN -- обнаружено и исправлено) | +| Atomic write | Per-image | Temp file + `os.replace()` (POSIX atomic) | +| Per-frame normalization | Per-image | Каждый depth/chm map нормализуется индивидуально | + +### 6.2. Визуальная инспекция + +Для каждого subset: +- **20 случайных query** (drone) + **20 случайных db** (satellite) +- PNG-визуализации генерируются автоматически рядом с исходными файлами +- Проверка: контуры зданий/дорог видны на edges? Классы корректны на segm? + +### 6.3. Обнаруженные проблемы и решения + +| Проблема | Причина | Решение | +|:---------|:--------|:--------| +| CHMv2 output полностью черный | FP16 вызывает NaN в DINOv3 DPT head | CHMv2 всегда в FP32 | +| DA3 verbose логи (Processed Images, Model Forward Pass) | Кастомный Logger, не Python logging | `DA3_LOG_LEVEL=ERROR` env var | +| `bpe_simple_vocab_16e6.txt.gz` не найден | SegEarth ищет по относительному пути | BPE vocab встроен в `src/nn/segearth_ov3/sam3/assets/`, дефолтный путь вычисляется от `__file__` | + +### 6.4. Количественные метрики + +| Метрика | Описание | Threshold | +|:--------|:---------|:----------| +| Depth coverage | % пикселей с depth in (0.01, 0.99) | > 80% | +| Edge density | % ненулевых пикселей edges | 5-40% | +| Seg class balance | Entropy распределения классов | > 1.0 bits | +| Seg coverage | % фоновых (class=0) пикселей | < 70% | + +--- + +## 7. Формат хранения и интеграция с даталоадером + +### 7.1. Интеграция с World-UAV DataLoader + +```python +class UAVGeoLocMultiModal(UAVGeoLocTrain): + def __init__(self, root, aug_root, ...): + super().__init__(root, ...) + self.aug_root = Path(aug_root) + + def __getitem__(self, idx): + query, positive, negative, label, height, rotation = super().__getitem__(idx) + query_path = self._get_query_path(idx) + + stem = Path(query_path).stem + aug_dir = self.aug_root / Path(query_path).parent + + # Загрузка grayscale PNG как float [0, 1], 1 канал + depth = np.array(Image.open(aug_dir / f"{stem}_depth.png")) / 255.0 # [H, W] + edges = np.array(Image.open(aug_dir / f"{stem}_edge.png")) / 255.0 + chm = np.array(Image.open(aug_dir / f"{stem}_chm.png")) / 255.0 + + # Сегментация -- class index (НЕ из RGB, а из отдельного источника) + # Если save_npy=True: seg = np.load(aug_dir / f"{stem}_segm.npy") + # Если только PNG: нужно обратное преобразование palette -> class ID + + # Конкатенация для Teacher: RGB(3) + depth(1) + edge(1) + chm(1) = 6ch + aux = np.stack([depth, edges, chm], axis=0) # [3, H, W] + + return query, positive, negative, label, aux +``` + +> **Рекомендация:** для сегментации включить `save_npy=True` или сохранять class index напрямую (uint8), т.к. обратное преобразование RGB palette -> class ID ненадежно. + +--- + +## 8. Гипотезы аннотирования (AH1-AH6) + +**AH1. Depth+edges дают больший прирост R@1 при LUPI-дистилляции, чем depth alone.** +Основание: edges инвариантны к текстуре/освещению, выделяют структурные контуры -- наиболее устойчивые к кросс-вью трансформации. + +**AH2. SegEarth-OV3 превосходит SegFormer-B5 по downstream CVGL utility.** +Основание: OV-сегментация использует RS-специфичные промпты (building, road, water), не ограничена 150 классами ADE20K. + +**AH3. Качество monocular depth деградирует для nadir-view спутниковых патчей сильнее, чем для oblique drone-кадров.** +Основание: DA3 обучен преимущественно на наземных изображениях; CHMv2 обучен на аэрофотосъемке -- компенсирует этот разрыв. + +**AH4. Edges из depth (Sobel) > Canny edges из RGB для CVGL.** +Основание: Sobel на depth выделяет границы по геометрии (перепад высот), а не по текстуре (которая различается drone<->sat). + +**AH5. CHMv2 + DA3 depth вместе > каждый по отдельности.** +Основание: DA3 и CHMv2 обучены на разных доменах (ground-level vs aerial). Комбинация дает Teacher два ортогональных depth representation, каждый из которых лучше работает на "своем" типе изображений. + +**AH6. Автоматическая аннотация 973K изображений за ~2 дня на одном GPU экономически оправдана.** +Основание: ручная разметка -- десятки тысяч человеко-часов. Шумная автоматическая аннотация при LUPI допустима, т.к. Student не использует эти модальности при инференсе. + +--- + +## 9. Архитектура кода + +### 9.1. Структура модулей + +``` +src/ +├── nn/ # Вендорированные нейросетевые пакеты +│ ├── __init__.py # Регистрация sys.path при import src.nn +│ ├── segearth_ov3/ # SegEarth-OV-3 (полная копия репозитория) +│ │ ├── segearthov3_segmentor.py # SegEarthOV3Segmentation class +│ │ ├── sam3/ # SAM 3.1 backbone (134 .py файла) +│ │ │ ├── __init__.py # build_sam3_image_model() +│ │ │ ├── model_builder.py # Сборка модели из компонентов +│ │ │ ├── model/ # Sam3Image, Sam3Processor, ViT, encoder/decoder +│ │ │ └── assets/bpe_simple_vocab_16e6.txt.gz # BPE tokenizer vocab +│ │ └── pamr.py # Probabilistic Affinity Mask Refinement +│ └── depth_anything_3/ # Depth-Anything-3 (полная копия пакета) +│ ├── api.py # DepthAnything3 class (main API) +│ ├── model/ # DA3 архитектура +│ │ ├── da3.py # DepthAnything3Net, NestedDepthAnything3Net +│ │ ├── dinov2/ # DinoV2 backbone (ViT) +│ │ ├── dpt.py, dualdpt.py # DPT heads +│ │ └── cam_enc.py, cam_dec.py # Camera encoder/decoder +│ ├── configs/ # YAML-конфиги моделей (da3-large.yaml и др.) +│ ├── utils/ # I/O processors, export, geometry, alignment +│ └── registry.py # Реестр моделей +├── augmentor/ +│ ├── dataset.py # ImageRecord, discover_images, INCOMPLETE_SCENES, +│ │ # attach_output_dirs, filter_completed, STAGE_SUFFIX, +│ │ # stage_filename, AugmentDataset +│ ├── inference.py # infer_depth_batch, infer_chmv2_batch (FP32 only), +│ │ # compute_edges_from_depth, infer_segmentation_batch, +│ │ # _infer_segearth_ov3 (batched backbone, max_batch=8), +│ │ # _infer_segformer +│ ├── io_utils.py # save_depth_async, save_chmv2_async, save_edges_async, +│ │ # save_segmentation_async, _save_float16_map, +│ │ # _atomic_save_npy, make_palette, _FIXED_PALETTE +│ └── models.py # load_depth_model (DA3 via src.nn), +│ # load_chmv2_model (FP32, NaN в FP16), +│ # load_segmentation_model (SegEarth via src.nn), +│ # _resolve_cache_dir, unload_model +├── conf/ # Gin-configurable dataclasses +│ ├── config_loader.py # load_all_configs() +│ ├── hardware_conf.py # HardwareConfig + auto_batch_size +│ ├── input_conf.py # InputConfig (image_size=256, sobel, imagenet stats) +│ ├── models_conf.py # ModelsConfig (depth, chmv2, seg IDs, fallbacks, weights_dir) +│ ├── pipeline_conf.py # PipelineConfig (paths, stages, resume, subset, source) +│ └── seg_conf.py # SegConfig (prompts, threshold, default_resolution) +├── utils/ +│ ├── benchmark.py # BenchmarkResult, benchmark_depth, benchmark_chmv2, +│ │ # benchmark_edges, benchmark_segmentation, benchmark_full_pipeline +│ ├── configure_gpu.py # GPU configuration utilities +│ ├── profiler.py # log_system_info, log_disk_info, log_vram_snapshot +│ └── utils_file_dir.py # get_proj_dir +├── tests/ # 125 тестов (pytest), 6 test-файлов +│ ├── test_config.py # Config dataclass tests (12 tests) +│ ├── test_models_and_benchmark.py # Model loading + benchmarking (62 tests) +│ ├── test_inference.py # Inference functions (18 tests) +│ ├── test_io_utils.py # I/O and save functions (16+ tests) +│ ├── test_dataset.py # Discovery and filtering (12 tests) +│ └── test_pipeline_integration.py # End-to-end stage tests (8 tests) +└── main.py # Sequential pipeline: + # run_pipeline -> run_depth_stage, run_edges_stage, + # run_segmentation_stage, run_chmv2_stage + # manifest generation, _silence_model_loggers(), + # log_system_info(), log_disk_info() +``` + +### 9.2. Механизм вендоринга (src/nn/) + +Нейросетевые пакеты **встроены внутрь проекта** и не требуют внешних репозиториев: + +| Пакет | Источник | Расположение | Размер | +|:------|:---------|:-------------|:-------| +| SegEarth-OV-3 | [earth-insights/SegEarth-OV-3](https://github.com/earth-insights/SegEarth-OV-3) | `src/nn/segearth_ov3/` | ~5 MB (134 .py + BPE vocab) | +| Depth-Anything-3 | [ByteDance-Seed/Depth-Anything-3](https://github.com/ByteDance-Seed/Depth-Anything-3) | `src/nn/depth_anything_3/` | ~1.6 MB (85 .py + YAML configs) | + +**Как работает:** `src/nn/__init__.py` при импорте добавляет в `sys.path`: +- `src/nn/` -- делает доступным `from depth_anything_3.api import DepthAnything3` +- `src/nn/segearth_ov3/` -- делает доступным `from segearthov3_segmentor import SegEarthOV3Segmentation` и `from sam3 import build_sam3_image_model` + +Все внутренние импорты обоих пакетов работают **без изменений** -- код идентичен оригинальным репозиториям. + +**Загрузка моделей:** `src/augmentor/models.py` содержит `import src.nn` (регистрация путей) и далее использует стандартные импорты: +```python +from depth_anything_3.api import DepthAnything3 # depth +from segearthov3_segmentor import SegEarthOV3Segmentation # segmentation +``` + +### 9.3. Ключевые design decisions + +| Решение | Обоснование | +|:--------|:-----------| +| **Вендорированные пакеты в src/nn/** | Нет зависимости от внешних репозиториев; воспроизводимость; нет конфликтов sys.path | +| **Sequential pipeline** (одна модель за раз) | Максимальный batch size для каждой модели; гранулярный resume по стадиям; упрощенная обработка ошибок | +| **Batched backbone для SegEarth-OV3** | `predict_pil_batch()`: 1 backbone pass на <=8 img, затем per-image grounding. Backbone ~80% времени -> ~2.5-3x ускорение стадии. Fallback на per-image при OOM | +| Stem-based file naming | `{stem}_{suffix}.png` -- файлы в той же директории, структура 1:1 с исходным датасетом | +| Grayscale PNG для depth/edge/chm | 1 канал, пригодный для обучения (`pixel / 255.0`); colormapped PNG теряет информацию | +| CHMv2 в FP32 | FP16 вызывает NaN (обнаружено при валидации на реальных данных) | +| Per-image atomic save | Crash-safe: temp file + `os.replace()` | +| Async I/O (ThreadPoolExecutor) | GPU inference не блокируется записью на диск | +| BPE vocab встроен в src/nn/ | Файл `bpe_simple_vocab_16e6.txt.gz` в `sam3/assets/`, дефолтный путь вычисляется от `__file__` | +| 125 unit tests (6 test-файлов) | Полное покрытие: inference, I/O, pipeline, benchmark, config, dataset (все mocked) | + +### 9.4. Веса моделей + +Все веса хранятся в `in/weights/`: + +``` +in/weights/ +├── models--depth-anything--DA3-LARGE-1.1/ # HF cache (safetensors) +├── models--nvidia--segformer-b5-finetuned-ade-640-640/ # SegFormer fallback +├── sam3.1/sam3.1_multiplex.pt # 3.3 GB (SAM 3.1 checkpoint) +├── dinov3-chmv2/ # 1.3 GB (safetensors + config) +└── bpe_simple_vocab_16e6.txt.gz # CLIP BPE (опциональная копия) +``` + +> BPE vocab уже встроен в `src/nn/segearth_ov3/sam3/assets/`. Копия в `in/weights/` используется приоритетно, если существует. + +--- + +## 10. Риски и митигации + +| Риск | Вероятность | Митигация | +|:-----|:-----------|:----------| +| **R1:** SegEarth-OV3 не загрузится | Низкая (вендорирован) | Автоматический fallback на SegFormer-B5 | +| **R2:** ~2 дней на GPU (batched) | Средняя | Resume позволяет прерывать/продолжать; можно запускать по subset'ам | +| **R2a:** SegEarth-OV3 OOM при batch=8 | Низкая (RTX 4090, ~10 GB) | Автоматический fallback на per-image `predict_pil()` при ошибке `predict_pil_batch()` | +| **R3:** DA3 не загрузится | Низкая (вендорирован) | Fallback на DA V2 Large (transformers) | +| **R4:** Диск закончится (~228 GB только PNG) | Низкая | `save_vis = False` для экономии | +| **R5:** CHMv2 NaN в FP16 | **Решено** | Модель загружается в FP32 | +| **R6:** Segmentation palette неинтуитивная | **Решено** | Фиксированная палитра: red=building, gray=road, green=vegetation, blue=water | + +--- + +## 11. Timeline и ресурсы + +### 11.1. Фазы + +| Фаза | Длительность | Статус | +|:-----|:-------------|:-------| +| **0. Setup** | -- | Выполнено: DA3, SegEarth-OV3, CHMv2 установлены, веса скачаны | +| **1. Smoke test (100 img)** | -- | Выполнено: 34.33 s, все стадии работают | +| **2. Валидация качества** | -- | Выполнено: CHMv2 FP16->FP32 fix, palette fix | +| **2a. Вендоринг моделей** | -- | Выполнено: SegEarth-OV3 и DA3 перенесены в src/nn/ | +| **3. Полный запуск (973K)** | ~2 дней (batched backbone) | В процессе | +| **4. Final validation** | 0.5 дня | -- | + +### 11.2. Ресурсы + +| Ресурс | Требование | +|:-------|:-----------| +| GPU | RTX 4090 (24 GB VRAM), peak ~10 GB (SegEarth batch=8, sequential) | +| RAM | >= 32 GB | +| SSD | >= 250 GB свободно (только PNG) | +| GPU-часы | ~46-51 ч (batched) / ~93 ч (per-image) | + +--- + +## 12. Библиография + +1. Wu et al., "UAV-GeoLoc: A Large-Vocabulary Dataset and Geometry-Transformed Method for UAV Geo-Localization," IEEE RA-L, 2025. +2. Yang et al., "Depth Anything V3: Recovering the Visual Space from Any Views," arXiv:2511.10647, 2025. +3. Li et al., "SegEarth-OV3: Exploring SAM 3.1 for Open-Vocabulary Semantic Segmentation in Remote Sensing Images," arXiv:2512.08730, 2025. +4. Ravi et al., "SAM 3.1: Segment Anything in Images and Videos," Meta AI, 2025. +5. Oquab et al., "DINOv3: Learning Robust Visual Features without Supervision," Meta AI, 2025. +6. Vapnik & Vashist, "A New Learning Paradigm: Learning Using Privileged Information," Neural Networks, 2009. +7. Xie et al., "SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers," NeurIPS, 2021. +8. Lin et al., "Sample4Geo: Hard Negative Sampling For Cross-View Geo-Localisation," ICCV, 2023.