DepthLab: инпейнт карт глубины на основе диффузионных моделей

DepthLab (https://johanan528.github.io/depthlab_web/) - диффузионный механизм инпейнта карт глубины с двумя параллельными ветвями для задач заполнения 3D-сцен, генерации сцен на основе текстовых промптов, реконструкции с использованием DUST3R и заполнение глубины LiDAR.

Первая ветвь, Reference U-Net извлекает признаки из RGB-изображений, которые служат условием для второй ветви.

Вторая ветвь, Estimation U-Net, обрабатывает имеющиеся данные о глубине и маску, определяющую области, требующие восстановления. Признаки RGB, полученные из Reference U-Net, последовательно интегрируются в Estimation U-Net, что позволяет управлять процессом восстановления.

Взаимодействие между ветвями Reference U-Net и Estimation U-Net реализуется механизмом cross-attention, который использует CLIP encoder.

Архитектура DepthLab опирается на наработки Marigold и Stable Diffusion V2. Кодирование RGB-изображений и карт глубины в латентное пространство осуществляется VAE. Маска также кодируется с помощью VAE, что позволяет сохранить детальную информацию о форме и границах.

Обучение DepthLab проводилось на двух синтетических датасетах: Hypersim (54 тысячи обучающих образцов) и Virtual KITTI (20 тысяч обучающих образцов). Для расширения обучающей выборки использовались случайные искажения изображений и несколько стратегий маскирования: штрихи, окружности, квадраты и их комбинации.

Оценка качества восстановления проводилась на 5 наборах: NYUv2, KITTI, ETH3D, ScanNet, DIODE. В качестве метрик использовались абсолютная относительная ошибка (AbsRel) и точность в пределах ?1 = 1.25.

Результаты тестов демонстрируют, что DepthLab превосходит как дискриминативные (DiverseDepth, MiDaS, LeReS, Omnidata, HDN, DPT, DepthAnything, DepthAnythingV2), так и генеративные (Marigold, DepthFM, GeoWizard) методы в постоении карт глубины.

Для локального инференса потребуются модели:

Marigold checkpoint (https://huggingface.co/prs-eth/marigold-depth-v1-0);

Энкодер CLIP-ViT-H-14-laion-2B (https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K);

Набор чекпоинтов DepthLab (https://huggingface.co/Johanan0528/DepthLab/tree/main).

Локальная установка и инференс:

# Clone repo  
 git clone https://github.com/Johanan528/DepthLab.git 
 
 cd DepthLab 
 
 # Create conda env 
 
 conda env create -f environment.yaml 
 
 conda activate DepthLab 
 
 # Run inference 
 
 cd scripts 
 
 bash infer.sh

Страница проекта (https://johanan528.github.io/depthlab_web/)

Модель (https://huggingface.co/Johanan0528/DepthLab/tree/main)

Arxiv (https://arxiv.org/pdf/2412.18153)

GitHub (https://github.com/Johanan528/DepthLab)

Источник: github.com

Комментарии: