Pikaliov/depth_edges_annotate_worlduav
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: World-UAV annotation pipeline

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-16 11:22:01 +03:00
commit 686db62c25
312 changed files with 83978 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
.gitignore vendored Normal file
View File

@@ -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

294
README.md Normal file
View File

@@ -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/` -- отдельная установка не требуется.

163
docs/analysis_optimization.md Normal file
View File

@@ -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) | Высокое | Средний | Не реализовано |

326
docs/segearth_ov3_architecture.md Normal file
View File

@@ -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 |

194
docs/segmentation_class_analysis.md Normal file
View File

@@ -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

477
docs/skills_optimization_io_dl_ml.md Normal file
View File

@@ -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])

7
in/config_files/hardware.gin Normal file
View File

@@ -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

7
in/config_files/input.gin Normal file
View File

@@ -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]

9
in/config_files/models.gin Normal file
View File

@@ -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'

12
in/config_files/pipeline.gin Normal file
View File

@@ -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'

18
in/config_files/segmentation.gin Normal file
View File

@@ -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

1
src/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Source package for depth/edges/segmentation augmentation pipeline."""

1
src/augmentor/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Augmentor: depth/edges/segmentation augmentation pipeline for CVGL datasets."""

199
src/augmentor/dataset.py Normal file
View File

@@ -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 / <parent dirs>/."""
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),
}

302
src/augmentor/inference.py Normal file
View File

@@ -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)

252
src/augmentor/io_utils.py Normal file
View File

@@ -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)

241
src/augmentor/models.py Normal file
View File

@@ -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.")

1
src/conf/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Configuration package for the depth/edges/segmentation augmentation pipeline."""

56
src/conf/config_loader.py Normal file
View File

@@ -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

190
src/conf/hardware_conf.py Normal file
View File

@@ -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()

37
src/conf/input_conf.py Normal file
View File

@@ -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()

33
src/conf/models_conf.py Normal file
View File

@@ -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()

39
src/conf/pipeline_conf.py Normal file
View File

@@ -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()

27
src/conf/seg_conf.py Normal file
View File

@@ -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()

515
src/main.py Normal file
View File

@@ -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()

23
src/nn/__init__.py Normal file
View File

@@ -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)

446
src/nn/depth_anything_3/api.py Normal file
View File

@@ -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")

594
src/nn/depth_anything_3/app/css_and_html.py Normal file
View File

@@ -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 """
<div class="tech-bg" style="text-align: center; margin-bottom: 5px; padding: 40px 20px; border-radius: 15px; position: relative; overflow: hidden;">
<div style="position: relative; z-index: 2;">
<h1 style="margin: 0; font-size: 3.5em; font-weight: 700;
background: linear-gradient(135deg, #3b82f6, #8b5cf6);
background-size: 400% 400%;
-webkit-background-clip: text;
background-clip: text;
color: transparent;
animation: techGradient 3s ease infinite;
text-shadow: 0 0 30px rgba(59, 130, 246, 0.5);
letter-spacing: 2px;">
Depth Anything 3
</h1>
<p style="margin: 15px 0 0 0; font-size: 2.16em; font-weight: 300;" class="header-subtitle">
Recovering the Visual Space from Any Views
</p>
<div style="margin-top: 20px;">
<!-- Revert buttons to original inline styles -->
<a href="https://depth-anything-3.github.io" target="_blank" class="link-btn">
<i class="fas fa-globe" style="margin-right: 8px;"></i> Project Page
</a>
<a href="https://arxiv.org/abs/2406.09414" target="_blank" class="link-btn">
<i class="fas fa-file-pdf" style="margin-right: 8px;"></i> Paper
</a>
<a href="https://github.com/ByteDance-Seed/Depth-Anything-3" target="_blank" class="link-btn">
<i class="fab fa-github" style="margin-right: 8px;"></i> Code
</a>
</div>
</div>
</div>
<style>
/* Ensure tech-bg class is properly applied in dark mode */
@media (prefers-color-scheme: dark) {
.header-subtitle {
color: #cbd5e1;
}
/* Increase priority to ensure background color is properly applied */
.tech-bg {
background: linear-gradient(135deg, #0f172a, #1e293b) !important;
}
}
@media (prefers-color-scheme: light) {
.header-subtitle {
color: #475569;
}
/* Also add explicit background color for light mode */
.tech-bg {
background: linear-gradient(135deg, rgba(59, 130, 246, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%) !important;
}
}
</style>
"""
def get_description_html():
"""
Generate the main description and getting started HTML.
Returns:
str: HTML string for the description
"""
return """
<div class="description-container" style="padding: 25px; border-radius: 15px; margin: 0 0 20px 0;">
<h2 class="description-title" style="margin-top: 0; font-size: 1.6em; text-align: center;">
<i class="fas fa-bullseye fa-color-red" style="margin-right: 8px;"></i> What This Demo Does
</h2>
<div class="description-content" style="padding: 20px; border-radius: 10px; margin: 15px 0; text-align: center;">
<p class="description-main" style="line-height: 1.6; margin: 0; font-size: 1.45em;">
<strong>Upload images or videos</strong> → <strong>Get <span class="metric-text">Metric</span> <span class="pointcloud-text">Point Clouds</span>, <span class="cameras-text">Cameras</span> and <span class="gaussians-text">Novel Views</span></strong> → <strong>Explore in 3D</strong>
</p>
</div>
<div style="text-align: center; margin-top: 15px;">
<p class="description-tip" style="font-style: italic; margin: 0;">
<i class="fas fa-lightbulb fa-color-yellow" style="margin-right: 8px;"></i> <strong>Tip:</strong> Landscape-oriented images or videos are preferred for best 3D recovering.
</p>
</div>
</div>
<style>
@media (prefers-color-scheme: dark) {
.description-container {
background: linear-gradient(135deg, rgba(59, 130, 246, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%);
border: 1px solid rgba(59, 130, 246, 0.2);
}
.description-title { color: #3b82f6; }
.description-content { background: rgba(0, 0, 0, 0.3); }
.description-main { color: #e0e0e0; }
.description-text { color: #cbd5e1; }
.description-tip { color: #cbd5e1; }
}
@media (prefers-color-scheme: light) {
.description-container {
background: linear-gradient(135deg, rgba(59, 130, 246, 0.05) 0%, rgba(139, 92, 246, 0.05) 100%);
border: 1px solid rgba(59, 130, 246, 0.3);
}
.description-title { color: #3b82f6; }
.description-content { background: transparent; }
.description-main { color: #1e293b; }
.description-text { color: #475569; }
.description-tip { color: #475569; }
}
</style>
"""
def get_acknowledgements_html():
"""
Generate the acknowledgements section HTML.
Returns:
str: HTML string for the acknowledgements
"""
return """
<div style="background: linear-gradient(135deg, rgba(59, 130, 246, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%);
padding: 25px; border-radius: 15px; margin: 20px 0; border: 1px solid rgba(59, 130, 246, 0.2);">
<h3 style="color: #3b82f6; margin-top: 0; text-align: center; font-size: 1.4em;">
<i class="fas fa-trophy fa-color-yellow" style="margin-right: 8px;"></i> Research Credits & Acknowledgments
</h3>
<div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin: 15px 0;">
<!-- Original Research Section (Left) -->
<div style="text-align: center;">
<h4 style="color: #8b5cf6; margin: 10px 0;"><i class="fas fa-flask fa-color-green" style="margin-right: 8px;"></i> Original Research</h4>
<p style="color: #e0e0e0; margin: 5px 0;">
<a href="https://depth-anything-3.github.io" target="_blank"
style="color: #3b82f6; text-decoration: none; font-weight: 600;">
Depth Anything 3
</a>
</p>
</div>
<!-- Previous Versions Section (Right) -->
<div style="text-align: center;">
<h4 style="color: #8b5cf6; margin: 10px 0;"><i class="fas fa-history fa-color-blue" style="margin-right: 8px;"></i> Previous Versions</h4>
<div style="display: flex; flex-direction: row; gap: 15px; justify-content: center; align-items: center;">
<p style="color: #e0e0e0; margin: 0;">
<a href="https://huggingface.co/spaces/LiheYoung/Depth-Anything" target="_blank"
style="color: #3b82f6; text-decoration: none; font-weight: 600;">
Depth-Anything
</a>
</p>
<span style="color: #e0e0e0;">•</span>
<p style="color: #e0e0e0; margin: 0;">
<a href="https://huggingface.co/spaces/depth-anything/Depth-Anything-V2" target="_blank"
style="color: #3b82f6; text-decoration: none; font-weight: 600;">
Depth-Anything-V2
</a>
</p>
</div>
</div>
</div>
<!-- HF Demo Adapted from - Centered at the bottom of the whole block -->
<div style="margin-top: 20px; padding-top: 15px; border-top: 1px solid rgba(59, 130, 246, 0.3); text-align: center;">
<p style="color: #a0a0a0; font-size: 0.9em; margin: 0;">
<i class="fas fa-code-branch fa-color-gray" style="margin-right: 5px;"></i> HF demo adapted from <a href="https://huggingface.co/spaces/facebook/map-anything" target="_blank" style="color: inherit; text-decoration: none;">Map Anything</a>
</p>
</div>
</div>
"""
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.
> <i class="fas fa-triangle-exclamation fa-color-red" style="margin-right: 5px;"></i> Metric scale estimation is difficult on aerial/drone images.
"""

724
src/nn/depth_anything_3/app/gradio_app.py Normal file
View File

@@ -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()

43
src/nn/depth_anything_3/app/modules/__init__.py Normal file
View File

@@ -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",
]

619
src/nn/depth_anything_3/app/modules/event_handlers.py Normal file
View File

@@ -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}", ""

304
src/nn/depth_anything_3/app/modules/file_handlers.py Normal file
View File

@@ -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."
),
)

260
src/nn/depth_anything_3/app/modules/model_inference.py Normal file
View File

@@ -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.

477
src/nn/depth_anything_3/app/modules/ui_components.py Normal file
View File

@@ -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.** <br><br><br>"
"To render novel views from 3DGS, "
"enable **Infer 3D Gaussian Splatting** below. <br>"
"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 (<i class="fas fa-triangle-exclamation '
'fa-color-red"></i> 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())

207
src/nn/depth_anything_3/app/modules/utils.py Normal file
View File

@@ -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

434
src/nn/depth_anything_3/app/modules/visualization.py Normal file
View File

@@ -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}"]

45
src/nn/depth_anything_3/bench/__init__.py Normal file
View File

@@ -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"]

98
src/nn/depth_anything_3/bench/configs/eval_bench.yaml Normal file
View File

@@ -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]

136
src/nn/depth_anything_3/bench/dataset.py Normal file
View File

@@ -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

21
src/nn/depth_anything_3/bench/datasets/__init__.py Normal file
View File

@@ -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.
"""

681
src/nn/depth_anything_3/bench/datasets/dtu.py Normal file
View File

@@ -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]

182
src/nn/depth_anything_3/bench/datasets/dtu64.py Normal file
View File

@@ -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."
)

594
src/nn/depth_anything_3/bench/datasets/eth3d.py Normal file
View File

@@ -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

440
src/nn/depth_anything_3/bench/datasets/hiroom.py Normal file
View File

@@ -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

591
src/nn/depth_anything_3/bench/datasets/scannetpp.py Normal file
View File

@@ -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

449
src/nn/depth_anything_3/bench/datasets/sevenscenes.py Normal file
View File

@@ -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

752
src/nn/depth_anything_3/bench/evaluator.py Normal file
View File

@@ -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: {"<data>_<mode>": 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)

618
src/nn/depth_anything_3/bench/print_metrics.py Normal file
View File

@@ -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()

85
src/nn/depth_anything_3/bench/registries.py Normal file
View File

@@ -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

525
src/nn/depth_anything_3/bench/utils.py Normal file
View File

@@ -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

144
src/nn/depth_anything_3/cfg.py Normal file
View File

@@ -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

803
src/nn/depth_anything_3/cli.py Normal file
View File

@@ -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()

45
src/nn/depth_anything_3/configs/da3-base.yaml Normal file
View File

@@ -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

71
src/nn/depth_anything_3/configs/da3-giant.yaml Normal file
View File

@@ -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

45
src/nn/depth_anything_3/configs/da3-large.yaml Normal file
View File

@@ -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

45
src/nn/depth_anything_3/configs/da3-small.yaml Normal file
View File

@@ -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

28
src/nn/depth_anything_3/configs/da3metric-large.yaml Normal file
View File

@@ -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]

28
src/nn/depth_anything_3/configs/da3mono-large.yaml Normal file
View File

@@ -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]

10
src/nn/depth_anything_3/configs/da3nested-giant-large.yaml Normal file
View File

@@ -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

20
src/nn/depth_anything_3/model/__init__.py Normal file
View File

@@ -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,
]

45
src/nn/depth_anything_3/model/cam_dec.py Normal file
View File

@@ -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

80
src/nn/depth_anything_3/model/cam_enc.py Normal file
View File

@@ -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

442
src/nn/depth_anything_3/model/da3.py Normal file
View File

@@ -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

64
src/nn/depth_anything_3/model/dinov2/dinov2.py Normal file
View File

@@ -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,
)

25
src/nn/depth_anything_3/model/dinov2/layers/__init__.py Normal file
View File

@@ -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,
]

100
src/nn/depth_anything_3/model/dinov2/layers/attention.py Normal file
View File

@@ -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

143
src/nn/depth_anything_3/model/dinov2/layers/block.py Normal file
View File

@@ -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

35
src/nn/depth_anything_3/model/dinov2/layers/drop_path.py Normal file
View File

@@ -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)

31
src/nn/depth_anything_3/model/dinov2/layers/layer_scale.py Normal file
View File

@@ -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}"

40
src/nn/depth_anything_3/model/dinov2/layers/mlp.py Normal file
View File

@@ -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

94
src/nn/depth_anything_3/model/dinov2/layers/patch_embed.py Normal file
View File

@@ -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

200
src/nn/depth_anything_3/model/dinov2/layers/rope.py Normal file
View File

@@ -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)

62
src/nn/depth_anything_3/model/dinov2/layers/swiglu_ffn.py Normal file
View File

@@ -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,
)

456
src/nn/depth_anything_3/model/dinov2/vision_transformer.py Normal file
View File

@@ -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

458
src/nn/depth_anything_3/model/dpt.py Normal file
View File

@@ -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

488
src/nn/depth_anything_3/model/dualdpt.py Normal file
View File

@@ -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

200
src/nn/depth_anything_3/model/gs_adapter.py Normal file
View File

@@ -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

133
src/nn/depth_anything_3/model/gsdpt.py Normal file
View File

@@ -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

223
src/nn/depth_anything_3/model/reference_view_selector.py Normal file
View File

@@ -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

109
src/nn/depth_anything_3/model/utils/attention.py Normal file
View File

@@ -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

81
src/nn/depth_anything_3/model/utils/block.py Normal file
View File

@@ -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

340
src/nn/depth_anything_3/model/utils/gs_renderer.py Normal file
View File

@@ -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

230
src/nn/depth_anything_3/model/utils/head_utils.py Normal file
View File

@@ -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)

208
src/nn/depth_anything_3/model/utils/transform.py Normal file
View File

@@ -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

50
src/nn/depth_anything_3/registry.py Normal file
View File

@@ -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()

24
src/nn/depth_anything_3/services/__init__.py Normal file
View File

@@ -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,
]

1417
src/nn/depth_anything_3/services/backend.py Normal file
View File

File diff suppressed because it is too large Load Diff

806
src/nn/depth_anything_3/services/gallery.py Normal file
View File

@@ -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"""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8" />
<title>Depth Anything 3 Gallery</title>
<meta name="viewport" content="width=device-width, initial-scale=1" />
<link rel="icon" href="https://i.postimg.cc/rFSzGJ7J/light-icon.jpg" media="(prefers-color-scheme: light)">
<link rel="icon" href="https://i.postimg.cc/P5gZfJsf/dark-icon.jpg" media="(prefers-color-scheme: dark)">
<script type="module" src="https://unpkg.com/@google/model-viewer/dist/model-viewer.min.js"></script>
<style>
:root {
--gap:16px; --card-radius:16px; --shadow:0 8px 24px rgba(0,0,0,.12);
--maxW:1036px; --maxH:518px;
--tech-blue: #00d4ff;
--tech-cyan: #00ffcc;
--tech-purple: #7877c6;
}
*{ box-sizing:border-box }
/* Dark mode tech theme */
@media (prefers-color-scheme: dark) {
body{
margin:0; font:16px/1.5 system-ui,-apple-system,Segoe UI,Roboto,sans-serif;
background: linear-gradient(135deg, #0a0a0a 0%, #1a1a2e 50%, #16213e 100%);
color:#e8eaed;
position: relative;
overflow-x: hidden;
}
body::before {
content: '';
position: fixed;
top: 0;
left: 0;
right: 0;
bottom: 0;
background:
radial-gradient(circle at 20% 80%, rgba(120, 119, 198, 0.3) 0%, transparent 50%),
radial-gradient(circle at 80% 20%, rgba(255, 119, 198, 0.3) 0%, transparent 50%),
radial-gradient(circle at 40% 40%, rgba(120, 219, 255, 0.2) 0%, transparent 50%);
animation: techPulse 8s ease-in-out infinite;
z-index: -1;
}
}
/* Light mode tech theme */
@media (prefers-color-scheme: light) {
body{
margin:0; font:16px/1.5 system-ui,-apple-system,Segoe UI,Roboto,sans-serif;
background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 50%, #cbd5e1 100%);
color:#1e293b;
position: relative;
overflow-x: hidden;
}
body::before {
content: '';
position: fixed;
top: 0;
left: 0;
right: 0;
bottom: 0;
background:
radial-gradient(circle at 20% 80%, rgba(0, 212, 255, 0.1) 0%, transparent 50%),
radial-gradient(circle at 80% 20%, rgba(0, 102, 255, 0.1) 0%, transparent 50%),
radial-gradient(circle at 40% 40%, rgba(0, 255, 204, 0.08) 0%, transparent 50%);
animation: techPulse 8s ease-in-out infinite;
z-index: -1;
}
}
@keyframes techPulse {
0%, 100% { opacity: 0.5; }
50% { opacity: 0.8; }
}
@keyframes techGradient {
0% { background-position: 0% 50%; }
50% { background-position: 100% 50%; }
100% { background-position: 0% 50%; }
}
/* Dark mode header */
@media (prefers-color-scheme: dark) {
header{
padding:20px 24px; position:sticky; top:0;
background:linear-gradient(180deg,rgba(10,10,10,0.9) 60%,rgba(10,10,10,0));
z-index:2; border-bottom:1px solid rgba(0, 212, 255, 0.2);
backdrop-filter: blur(10px);
}
h1{
margin:0; font-size:22px;
background: linear-gradient(45deg, var(--tech-blue), var(--tech-cyan), var(--tech-purple));
background-size: 400% 400%;
-webkit-background-clip: text;
background-clip: text;
color: transparent;
animation: techGradient 3s ease infinite;
text-shadow: 0 0 30px rgba(0, 212, 255, 0.5);
}
.muted{ opacity:.7; font-size:13px; color: #a0a0a0; }
#backBtn{
display:none; padding:6px 10px; border-radius:10px;
border:1px solid rgba(0, 212, 255, 0.3);
background:rgba(0, 0, 0, 0.3);
color:#e8eaed; cursor:pointer;
transition: all 0.3s ease;
}
#backBtn:hover {
border-color: var(--tech-blue);
box-shadow: 0 0 10px rgba(0, 212, 255, 0.3);
}
#search{
flex:1 1 260px; min-width:240px; max-width:520px;
padding:10px 14px; border-radius:12px;
border:1px solid rgba(0, 212, 255, 0.3);
background:rgba(0, 0, 0, 0.3);
color:#e8eaed; outline:none;
transition: all 0.3s ease;
}
#search:focus {
border-color: var(--tech-blue);
box-shadow: 0 0 10px rgba(0, 212, 255, 0.3);
}
}
/* Light mode header */
@media (prefers-color-scheme: light) {
header{
padding:20px 24px; position:sticky; top:0;
background:linear-gradient(180deg,rgba(248,250,252,0.9) 60%,rgba(248,250,252,0));
z-index:2; border-bottom:1px solid rgba(0, 212, 255, 0.3);
backdrop-filter: blur(10px);
}
h1{
margin:0; font-size:22px;
background: linear-gradient(45deg, #0066ff, #00d4ff, #00ffcc);
background-size: 400% 400%;
-webkit-background-clip: text;
background-clip: text;
color: transparent;
animation: techGradient 3s ease infinite;
text-shadow: 0 0 20px rgba(0, 102, 255, 0.3);
}
.muted{ opacity:.7; font-size:13px; color: #64748b; }
#backBtn{
display:none; padding:6px 10px; border-radius:10px;
border:1px solid rgba(0, 212, 255, 0.4);
background:rgba(255, 255, 255, 0.8);
color:#1e293b; cursor:pointer;
transition: all 0.3s ease;
}
#backBtn:hover {
border-color: #0066ff;
box-shadow: 0 0 10px rgba(0, 102, 255, 0.3);
}
#search{
flex:1 1 260px; min-width:240px; max-width:520px;
padding:10px 14px; border-radius:12px;
border:1px solid rgba(0, 212, 255, 0.4);
background:rgba(255, 255, 255, 0.8);
color:#1e293b; outline:none;
transition: all 0.3s ease;
}
#search:focus {
border-color: #0066ff;
box-shadow: 0 0 10px rgba(0, 102, 255, 0.3);
}
}
.row{ display:flex; gap:12px; align-items:center; flex-wrap:wrap; justify-content:center; }
main{ padding:16px 24px 24px; display:grid; place-items:center; }
.group-wrap{ width:min(900px,100%); }
.group-list{ list-style:none; margin:0; padding:0; display:grid; gap:10px; }
/* Dark mode cards */
@media (prefers-color-scheme: dark) {
.group-item{
display:flex; align-items:center; gap:12px; padding:12px 14px;
background:rgba(0, 0, 0, 0.3); border:1px solid rgba(0, 212, 255, 0.2); border-radius:14px; cursor:pointer;
transition: all 0.3s ease;
backdrop-filter: blur(10px);
}
.group-item:hover{
transform: translateY(-1px);
border-color:var(--tech-blue);
box-shadow: 0 4px 15px rgba(0, 212, 255, 0.2);
}
.card{
background:rgba(0, 0, 0, 0.3); border:1px solid rgba(0, 212, 255, 0.2); border-radius:var(--card-radius);
overflow:hidden; box-shadow:var(--shadow);
transition:all 0.3s ease; cursor:pointer; display:flex; flex-direction:column; max-width:var(--maxW);
backdrop-filter: blur(10px);
}
.card:hover{
transform:translateY(-2px);
border-color:var(--tech-blue);
box-shadow: 0 8px 25px rgba(0, 212, 255, 0.2);
}
.thumb-box{
position:relative; width:100%; aspect-ratio:2/1;
background:linear-gradient(135deg, #0e121b 0%, #1a1a2e 100%);
display:grid; place-items:center; overflow:hidden;
border-bottom: 1px solid rgba(0, 212, 255, 0.1);
}
.open{
font-size:12px; opacity:.7; padding:6px 8px;
border:1px solid rgba(0, 212, 255, 0.3);
border-radius:10px;
background:rgba(0, 212, 255, 0.1);
transition: all 0.3s ease;
}
.open:hover {
background:rgba(0, 212, 255, 0.2);
border-color: var(--tech-blue);
}
}
/* Light mode cards */
@media (prefers-color-scheme: light) {
.group-item{
display:flex; align-items:center; gap:12px; padding:12px 14px;
background:rgba(255, 255, 255, 0.8); border:1px solid rgba(0, 212, 255, 0.3); border-radius:14px; cursor:pointer;
transition: all 0.3s ease;
backdrop-filter: blur(10px);
box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
}
.group-item:hover{
transform: translateY(-1px);
border-color:#0066ff;
box-shadow: 0 4px 15px rgba(0, 102, 255, 0.2);
}
.card{
background:rgba(255, 255, 255, 0.8); border:1px solid rgba(0, 212, 255, 0.3); border-radius:var(--card-radius);
overflow:hidden; box-shadow:0 4px 6px rgba(0, 0, 0, 0.1);
transition:all 0.3s ease; cursor:pointer; display:flex; flex-direction:column; max-width:var(--maxW);
backdrop-filter: blur(10px);
}
.card:hover{
transform:translateY(-2px);
border-color:#0066ff;
box-shadow: 0 8px 25px rgba(0, 102, 255, 0.2);
}
.thumb-box{
position:relative; width:100%; aspect-ratio:2/1;
background:linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
display:grid; place-items:center; overflow:hidden;
border-bottom: 1px solid rgba(0, 212, 255, 0.2);
}
.open{
font-size:12px; opacity:.7; padding:6px 8px;
border:1px solid rgba(0, 212, 255, 0.4);
border-radius:10px;
background:rgba(0, 212, 255, 0.1);
transition: all 0.3s ease;
}
.open:hover {
background:rgba(0, 212, 255, 0.2);
border-color: #0066ff;
}
}
.gname{ font-weight:600; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; width:100%; }
.grid{
width:min(1200px,100%);
display:grid;
grid-template-columns:repeat(auto-fill,minmax(260px,1fr));
gap:var(--gap);
align-items:start;
justify-items:stretch;
margin: 0 auto;
padding: 0 20px;
}
.thumb{ max-width:100%; max-height:100%; object-fit:contain; display:block; }
.meta{ padding:12px 14px; display:flex; justify-content:space-between; align-items:center; gap:8px; }
.title{ font-weight:600; font-size:14px; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; }
.empty{ opacity:.6; padding:40px 0; text-align:center; }
.crumb{ font-size:13px; opacity:.8; }
.overlay{ position:fixed; inset:0; background:rgba(0,0,0,.6); display:none; place-items:center; padding:20px; z-index:10; }
.overlay.show{ display:grid; }
/* Dark mode viewer */
@media (prefers-color-scheme: dark) {
.viewer{
inline-size:min(92vw,var(--maxW));
block-size:min(82vh,var(--maxH));
background:#0e121b; border:1px solid rgba(0, 212, 255, 0.3); border-radius:18px; overflow:hidden; position:relative; box-shadow:0 12px 36px rgba(0,0,0,.35);
display:grid;
}
.chip{ background:rgba(0,0,0,.45); border:1px solid rgba(0, 212, 255, 0.3); color:#e8eaed; padding:6px 10px; border-radius:12px; font-size:12px; max-width:60%; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; }
.btn{ margin-left:auto; background:rgba(0, 0, 0, 0.3); color:#e8eaed; border:1px solid rgba(0, 212, 255, 0.3); border-radius:10px; padding:6px 10px; cursor:pointer; transition: all 0.3s ease; }
.btn:hover { border-color: var(--tech-blue); box-shadow: 0 0 10px rgba(0, 212, 255, 0.3); }
.mv-box{ width:100%; aspect-ratio:1036/518; background:#0b0d12; border:1px solid rgba(0, 212, 255, 0.2); border-radius:12px; overflow:hidden; }
.mv-box model-viewer{ width:100%; height:100%; background:#0b0d12; }
.res-cell{ position:relative; width:100%; aspect-ratio:2/1; background:#0e121b; border:1px solid rgba(0, 212, 255, 0.2); border-radius:12px; overflow:hidden; display:grid; place-items:center; }
.res-empty{ position:absolute; inset:0; display:grid; place-items:center; opacity:.55; font-size:12px; color:#9aa0a6; }
.download-icon{ background:rgba(0, 0, 0, 0.6); border:1px solid rgba(0, 212, 255, 0.3); color:#e8eaed; box-shadow:0 4px 12px rgba(0,0,0,0.3); }
.download-icon:hover{ background:rgba(0, 212, 255, 0.2); border-color:var(--tech-blue); box-shadow:0 0 20px rgba(0, 212, 255, 0.4); transform:scale(1.05); }
}
/* Light mode viewer */
@media (prefers-color-scheme: light) {
.viewer{
inline-size:min(92vw,var(--maxW));
block-size:min(82vh,var(--maxH));
background:#f8fafc; border:1px solid rgba(0, 212, 255, 0.4); border-radius:18px; overflow:hidden; position:relative; box-shadow:0 12px 36px rgba(0,0,0,.15);
display:grid;
}
.chip{ background:rgba(255,255,255,0.8); border:1px solid rgba(0, 212, 255, 0.4); color:#1e293b; padding:6px 10px; border-radius:12px; font-size:12px; max-width:60%; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; }
.btn{ margin-left:auto; background:rgba(255, 255, 255, 0.8); color:#1e293b; border:1px solid rgba(0, 212, 255, 0.4); border-radius:10px; padding:6px 10px; cursor:pointer; transition: all 0.3s ease; }
.btn:hover { border-color: #0066ff; box-shadow: 0 0 10px rgba(0, 102, 255, 0.3); }
.mv-box{ width:100%; aspect-ratio:1036/518; background:#f8fafc; border:1px solid rgba(0, 212, 255, 0.3); border-radius:12px; overflow:hidden; }
.mv-box model-viewer{ width:100%; height:100%; background:#f8fafc; }
.res-cell{ position:relative; width:100%; aspect-ratio:2/1; background:#f8fafc; border:1px solid rgba(0, 212, 255, 0.3); border-radius:12px; overflow:hidden; display:grid; place-items:center; }
.res-empty{ position:absolute; inset:0; display:grid; place-items:center; opacity:.55; font-size:12px; color:#64748b; }
.download-icon{ background:rgba(255, 255, 255, 0.9); border:1px solid rgba(0, 212, 255, 0.4); color:#1e293b; box-shadow:0 4px 12px rgba(0,0,0,0.15); }
.download-icon:hover{ background:rgba(0, 212, 255, 0.2); border-color:#0066ff; box-shadow:0 0 20px rgba(0, 102, 255, 0.4); transform:scale(1.05); }
}
.viewer-header{ position:absolute; top:8px; left:8px; right:8px; display:flex; gap:8px; align-items:center; z-index:2; }
.viewer-body{ height:100%; display:grid; grid-template-rows:auto auto; gap:12px; padding:36px 8px 8px 8px; overflow:auto; }
.res-grid{ display:grid; grid-template-columns:1fr 1fr; gap:8px; }
.res-img{ max-width:100%; max-height:100%; object-fit:contain; display:block; }
.download-icon{ position:absolute; bottom:16px; right:16px; width:44px; height:44px; border-radius:50%; display:grid; place-items:center; font-size:20px; cursor:pointer; z-index:3; transition:all 0.3s ease; }
/* Pagination controls */
.pager {
grid-column: 1 / -1;
justify-content: center;
align-items: center;
display: flex;
gap: 16px;
margin-top: 8px;
font-size: 13px;
text-align: center;
}
/* Dark mode pagination */
@media (prefers-color-scheme: dark) {
.pager {
color: #ccc;
}
.pager button {
padding: 4px 10px;
border-radius: 8px;
border: 1px solid rgba(0, 212, 255, 0.3);
background: rgba(0, 0, 0, 0.3);
color: #e8eaed;
cursor: pointer;
transition: all 0.3s ease;
}
.pager button:hover:not(:disabled) {
border-color: var(--tech-blue);
box-shadow: 0 0 8px rgba(0, 212, 255, 0.2);
}
.pager button:disabled {
opacity: 0.4;
cursor: not-allowed;
}
}
/* Light mode pagination */
@media (prefers-color-scheme: light) {
.pager {
color: #64748b;
}
.pager button {
padding: 4px 10px;
border-radius: 8px;
border: 1px solid rgba(0, 212, 255, 0.4);
background: rgba(255, 255, 255, 0.8);
color: #1e293b;
cursor: pointer;
transition: all 0.3s ease;
}
.pager button:hover:not(:disabled) {
border-color: #0066ff;
box-shadow: 0 0 8px rgba(0, 102, 255, 0.2);
}
.pager button:disabled {
opacity: 0.4;
cursor: not-allowed;
}
}
/* Intro card styles */
@media (prefers-color-scheme: dark) {
.intro-card {
background: linear-gradient(135deg, rgba(0, 212, 255, 0.1) 0%, rgba(0, 102, 255, 0.1) 100%);
border: 1px solid rgba(0, 212, 255, 0.2);
backdrop-filter: blur(10px);
}
.intro-title {
background: linear-gradient(45deg, var(--tech-blue), var(--tech-cyan), var(--tech-purple));
background-size: 400% 400%;
-webkit-background-clip: text;
background-clip: text;
color: transparent;
animation: techGradient 3s ease infinite;
text-shadow: 0 0 20px rgba(0, 212, 255, 0.3);
}
.intro-description {
color: #e0e0e0;
}
}
@media (prefers-color-scheme: light) {
.intro-card {
background: linear-gradient(135deg, rgba(0, 212, 255, 0.05) 0%, rgba(0, 102, 255, 0.05) 100%);
border: 1px solid rgba(0, 212, 255, 0.3);
backdrop-filter: blur(10px);
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
.intro-title {
background: linear-gradient(45deg, #0066ff, #00d4ff, #00ffcc);
background-size: 400% 400%;
-webkit-background-clip: text;
background-clip: text;
color: transparent;
animation: techGradient 3s ease infinite;
text-shadow: 0 0 15px rgba(0, 102, 255, 0.2);
}
.intro-description {
color: #334155;
}
}
footer{
opacity:.55;
font-size:12px;
padding:12px 24px 24px;
text-align:center;
display:flex;
justify-content:center;
align-items:center;
width:100%;
}
</style>
</head>
<body>
<header>
<div class="row">
<button id="backBtn">← Back</button>
<h1 id="pageTitle">Depth Anything 3 Gallery</h1>
<span id="crumb" class="crumb"></span>
<input id="search" placeholder="Search…" />
</div>
<div class="muted" id="hint" style="text-align: center;">Level 1 shows groups only; click a group to browse scenes and previews.</div>
</header>
<main>
<!-- Tech intro card -->
<div class="intro-card" style="margin-bottom: 30px; padding: 25px; border-radius: 15px; text-align: center; max-width: 800px;">
<h2 class="intro-title" style="margin: 0 0 15px 0; font-size: 1.8em; font-weight: 700;">
🎯 Depth Anything 3 Gallery
</h2>
<p class="intro-description" style="margin: 0; font-size: 1.1em; line-height: 1.6;">
Explore 3D reconstructions and depth visualizations from Depth Anything 3.
Browse through groups of scenes, preview 3D models, and examine depth maps interactively.
</p>
</div>
<div id="level1" class="group-wrap" aria-live="polite">
<ul id="groupList" class="group-list"></ul>
<div id="groupEmpty" class="empty" style="display:none;">No available groups</div>
</div>
<div id="level2" style="display:none; width:100%;" aria-live="polite">
<div id="topPager" class="pager" style="margin-bottom: 16px;"></div>
<div id="grid" class="grid"></div>
<div id="sceneEmpty" class="empty" style="display:none;">No available scenes in this group</div>
</div>
</main>
<div id="overlay" class="overlay" role="dialog" aria-modal="true" aria-label="3D Preview">
<div class="viewer" id="viewer">
<div class="viewer-header">
<div id="viewerTitle" class="chip">Loading…</div>
<button id="toggleView" class="btn" title="Toggle between 3D-only and resource view">Resource View</button>
<button id="closeBtn" class="btn">Close</button>
</div>
<div id="downloadBtn" class="download-icon" title="Download GLB model">⬇</div>
<div class="viewer-body">
<div class="mv-box"><model-viewer id="mv"
src=""
ar
camera-controls
auto-rotate
interaction-prompt="auto"
shadow-intensity="0.7"
exposure="1.0"
alt="GLB Preview"></model-viewer></div>
<div class="res-grid" id="resGrid" hidden></div>
</div>
</div>
</div>
<footer>Depth Anything 3 Gallery. Copyright 2025 Depth Anything 3 authors.</footer>
<script>
const level1=document.getElementById('level1'),level2=document.getElementById('level2'),pageTitle=document.getElementById('pageTitle'),crumb=document.getElementById('crumb'),backBtn=document.getElementById('backBtn'),hint=document.getElementById('hint'),searchInput=document.getElementById('search'),groupList=document.getElementById('groupList'),groupEmpty=document.getElementById('groupEmpty'),topPager=document.getElementById('topPager'),grid=document.getElementById('grid'),sceneEmpty=document.getElementById('sceneEmpty'),overlay=document.getElementById('overlay'),viewer=document.getElementById('viewer'),mv=document.getElementById('mv'),viewerTitle=document.getElementById('viewerTitle'),downloadBtn=document.getElementById('downloadBtn'),toggleViewBtn=document.getElementById('toggleView'),closeBtn=document.getElementById('closeBtn'),resGrid=document.getElementById('resGrid');
let GROUPS=[],SCENES=[],currentGroup=null,currentScene=null,currentPage=1,currentScenePage=1;
const qs=()=>new URLSearchParams(location.search);
async function loadGroups(){const r=await fetch('/manifest.json',{cache:'no-store'});if(!r.ok)throw new Error(r.status+' '+r.statusText);const j=await r.json();GROUPS=j.groups||[];renderGroups(GROUPS);}
async function loadScenes(g){const r=await fetch('/manifest/'+encodeURIComponent(g)+'.json',{cache:'no-store'});if(!r.ok)throw new Error(r.status+' '+r.statusText);const j=await r.json();SCENES=j.items||[];const p=parseInt(qs().get('page'))||1;renderScenes(SCENES,p);}
function renderGroups(list){groupList.innerHTML='';const q=searchInput.value.trim().toLowerCase();const f=list.filter(g=>(g.title||g.id||'').toLowerCase().includes(q));if(!f.length){groupEmpty.style.display='';return;}groupEmpty.style.display='none';for(const g of f){const li=document.createElement('li');li.className='group-item';li.title=g.title||g.id;li.onclick=()=>enterLevel2(g.id,{push:true});const name=document.createElement('div');name.className='gname';name.textContent=g.title||g.id;li.appendChild(name);groupList.appendChild(li);}}
function renderScenes(list,page=1){topPager.innerHTML='';grid.innerHTML='';const q=searchInput.value.trim().toLowerCase();const f=list.filter(x=>(x.title||'').toLowerCase().includes(q)||(x.id||'').toLowerCase().includes(q));if(!f.length){sceneEmpty.style.display='';topPager.style.display='none';return;}sceneEmpty.style.display='none';topPager.style.display='flex';const perPage=16;const total=f.length;const totalPages=Math.max(1,Math.ceil(total/perPage));currentScenePage=page;const u=new URL(location.href);u.searchParams.set('page',page);history.replaceState(null,'',u);const subset=f.slice((page-1)*perPage,page*perPage);for(const i of subset){const c=document.createElement('div');c.className='card';c.title=i.title;const b=document.createElement('div');b.className='thumb-box';const img=document.createElement('img');img.className='thumb';img.loading='lazy';img.alt=i.title;img.src=i.thumbnail;b.appendChild(img);const m=document.createElement('div');m.className='meta';const t=document.createElement('div');t.className='title';t.textContent=i.title;const o=document.createElement('div');o.className='open';o.textContent='Preview';m.appendChild(t);m.appendChild(o);c.appendChild(b);c.appendChild(m);c.onclick=()=>openViewer(i,{push:true});grid.appendChild(c);}function buildPager(){const pg=document.createElement('div');pg.className='pager';const prev=document.createElement('button');prev.textContent='← Prev';prev.disabled=page<=1;prev.onclick=()=>renderScenes(list,page-1);const info=document.createElement('span');info.textContent=`${page} / ${totalPages}`;const next=document.createElement('button');next.textContent='Next →';next.disabled=page>=totalPages;next.onclick=()=>renderScenes(list,page+1);pg.appendChild(prev);pg.appendChild(info);pg.appendChild(next);return pg;}topPager.innerHTML='';topPager.appendChild(buildPager());grid.appendChild(buildPager());}
function enterLevel1({push=false}={}){currentGroup=null;pageTitle.textContent='Depth Anything 3 Gallery';crumb.textContent='';backBtn.style.display='none';hint.style.display='';level1.style.display='';level2.style.display='none';overlay.classList.remove('show');mv.src='';const u=new URL(location.href);u.searchParams.delete('group');u.searchParams.delete('id');u.searchParams.delete('page');push?history.pushState(null,'',u):history.replaceState(null,'',u);searchInput.value='';loadGroups().catch(e=>{groupList.innerHTML='';groupEmpty.style.display='';groupEmpty.textContent='Failed to load groups: '+e;});}
async function enterLevel2(g,{push=false}={}){currentGroup=g;pageTitle.textContent=g;crumb.textContent='(group)';backBtn.style.display='';hint.style.display='none';level1.style.display='none';level2.style.display='';overlay.classList.remove('show');mv.src='';const u=new URL(location.href);u.searchParams.set('group',g);u.searchParams.delete('id');push?history.pushState(null,'',u):history.replaceState(null,'',u);searchInput.value='';try{await loadScenes(g);const id=qs().get('id');if(id){const hit=SCENES.find(x=>x.id===id);if(hit)openViewer(hit,{push:false});}}catch(e){grid.innerHTML='';sceneEmpty.style.display='';sceneEmpty.textContent='Failed to load scenes: '+e;}}
function buildResGrid(i,page=1){
resGrid.innerHTML='';
const imgs=i.depth_images||[];
const perPage=4;
const total=imgs.length;
const totalPages=Math.max(1, Math.ceil(total/perPage));
currentPage=page;
const subset=imgs.slice((page-1)*perPage,(page-1)*perPage+perPage);
for(let k=0;k<4;k++){
const cell=document.createElement('div');
cell.className='res-cell';
if(subset[k]){
const im=document.createElement('img');
im.className='res-img';
im.src=subset[k];
im.alt=(i.title||'scene')+' depth '+(k+1+(page-1)*perPage);
im.loading='lazy';
cell.appendChild(im);
} else {
const ph=document.createElement('div');
ph.className='res-empty';
ph.textContent='N/A';
cell.appendChild(ph);
}
resGrid.appendChild(cell);
}
// pagination bar (always rebuilt)
const pager=document.createElement('div');
pager.className='pager';
const prev=document.createElement('button');
prev.textContent='← Prev';
prev.disabled=page<=1;
prev.onclick=()=>buildResGrid(i,page-1);
const info=document.createElement('span');
info.textContent=`${page} / ${totalPages}`;
const next=document.createElement('button');
next.textContent='Next →';
next.disabled=page>=totalPages;
next.onclick=()=>buildResGrid(i,page+1);
pager.appendChild(prev);
pager.appendChild(info);
pager.appendChild(next);
resGrid.appendChild(pager);
}
function openViewer(i,{push=false}={}){currentScene=i;viewerTitle.textContent=i.title;mv.src=i.model;overlay.classList.add('show');resGrid.hidden=true;toggleViewBtn.textContent='Resource View';viewer.style.blockSize='min(82vh,var(--maxH))';buildResGrid(i,1);downloadBtn.onclick=()=>{const a=document.createElement('a');a.href=i.model;a.download=i.title+'.glb';a.click();};if(push){const u=new URL(location.href);if(!u.searchParams.get('group'))u.searchParams.set('group',currentGroup||'');u.searchParams.set('id',i.id);history.pushState(null,'',u);}}
function toggleView(){const hidden=!resGrid.hidden;resGrid.hidden=hidden;toggleViewBtn.textContent=hidden?'Resource View':'3D Only';viewer.style.blockSize=hidden?'min(82vh,var(--maxH))':'min(92vh,900px)';}
function closeViewer(){const hasId=!!qs().get('id');if(hasId&&history.length>1){history.back();return;}const u=new URL(location.href);u.searchParams.delete('id');history.replaceState(null,'',u);overlay.classList.remove('show');mv.src='';}
overlay.onclick=e=>{if(e.target===overlay)closeViewer();};closeBtn.onclick=closeViewer;toggleViewBtn.onclick=toggleView;backBtn.onclick=()=>history.back();
searchInput.oninput=()=>{!qs().get('group')?renderGroups(GROUPS):renderScenes(SCENES,1);};
window.onpopstate=()=>routeFromURL();
async function routeFromURL(){if(location.pathname!="/")history.replaceState(null,'','/'+location.search);const g=qs().get('group');const id=qs().get('id');if(!g){enterLevel1({push:false});return;}await enterLevel2(g,{push:false});if(id){const hit=SCENES.find(x=>x.id===id);if(hit)openViewer(hit,{push:false});else{overlay.classList.remove('show');mv.src='';}}else{overlay.classList.remove('show');mv.src='';}}
routeFromURL();
</script>
</body>
</html>
"""
# ------------------------------ 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()

239
src/nn/depth_anything_3/services/inference_service.py Normal file
View File

@@ -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,
)

266
src/nn/depth_anything_3/services/input_handlers.py Normal file
View File

@@ -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'"
)

45
src/nn/depth_anything_3/specs.py Normal file
View File

@@ -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

163
src/nn/depth_anything_3/utils/alignment.py Normal file
View File

@@ -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

58
src/nn/depth_anything_3/utils/api_helpers.py Normal file
View File

@@ -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

479
src/nn/depth_anything_3/utils/camera_trj_helpers.py Normal file
View File

@@ -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

270
src/nn/depth_anything_3/utils/constants.py Normal file
View File

@@ -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)

59
src/nn/depth_anything_3/utils/export/__init__.py Normal file
View File

@@ -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,
]

150
src/nn/depth_anything_3/utils/export/colmap.py Normal file
View File

@@ -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

41
src/nn/depth_anything_3/utils/export/depth_vis.py Normal file
View File

@@ -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)

65
src/nn/depth_anything_3/utils/export/feat_vis.py Normal file
View File

@@ -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)

432
src/nn/depth_anything_3/utils/export/glb.py Normal file
View File

@@ -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

Some files were not shown because too many files have changed in this diff Show More