Depth from Focus - Image Restoration

Предсказание карт глубины по одному RGB-изображению с использованием глубокого обучения

Обзор

Данный проект посвящён задаче монокулярной оценки глубины — предсказанию карт глубины по одному RGB-изображению с использованием трёх различных архитектур нейронных сетей. Мы обучили и сравнили CNN- и Transformer-модели на двух реальных датасетах, чтобы понять их сильные и слабые стороны.

Ключевые выводы

• На малых и средних датасетах (2000–3500 изображений) CNN-архитектуры (UNet) превосходят Transformer-модели, которым для раскрытия потенциала требуются миллионы изображений.
• Добавление второго датасета (DIODE) с другими характеристиками домена улучшило обобщающую способность моделей.
• Механизмы внимания добавляют вычислительные затраты без гарантированного улучшения на ограниченных данных.

Датасеты

В работе использованы два дополнительных датасета из открытых источников:

1. NYU Depth V2

• Источник: https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html
• Размер: 2193 обучающих + 645 валидационных изображений
• Домен: Интерьеры (ванные, спальни, кухни, офисы)
• Источник глубины: RGB-D камера Kinect
• Диапазон глубины: 0.5 – 10 метров

2. DIODE

• Источник: https://diode-dataset.org/
• Размер: 1255 изображений
• Домен: Интерьеры + улицы
• Источник глубины: Высокоточный LiDAR + стерео
• Диапазон глубины: 0.5 – 50+ метров

Почему два датасета?

Использование датасетов из разных доменов (интерьеры и смешанные сцены) позволяет оценить, насколько хорошо модели обобщаются между типами сцен, и продемонстрировать проблему доменного разрыва в задаче оценки глубины.

Архитектуры моделей

Реализованы и сравнены три архитектуры:

1. U-Net (базовая модель)

Классический энкодер-декодер со skip-соединениями.

• Энкодер: 4-стадийная CNN с MaxPool понижением
• Декодер: ConvTranspose2d повышение + конкатенация skip-соединений
• Количество признаков: [64, 128, 256, 512]
• Обоснование: Проверенная архитектура для задач плотного предсказания

2. Attention U-Net

U-Net, дополненный attention-гейтами на skip-соединениях.

• Дополнение: Attention-гейты фильтруют нерелевантные признаки перед конкатенацией
• Гипотеза: Внимание должно помочь фокусироваться на релевантных для глубины регионах
• Обоснование: Проверка, улучшает ли внимание оценку глубины на малых данных

3. ViT Hybrid

Vision Transformer энкодер + CNN декодер.

• Энкодер: Patch embedding (4x4) -> 4 Transformer блока
• Декодер: 2-стадийное CNN повышение
• Количество признаков: embed_dim=32 (уменьшено для экономии памяти)
• Обоснование: Проверка, помогает ли глобальный контекст от self-attention предсказанию глубины

Метрики

Модели оцениваются по трём стандартным метрикам для задач оценки глубины:

Метрика	Описание	Лучше
RMSE	Среднеквадратичная ошибка	Меньше
MAE	Средняя абсолютная ошибка	Меньше
δ1	Доля пикселей, где max(pred/gt, gt/pred) < 1.25	Больше

Результаты

Количественные результаты

Модель	NYU RMSE (м)	NYU MAE (м)	NYU $\delta_1$ (%)	DIODE RMSE (м)	DIODE MAE (м)	DIODE $\delta_1$ (%)
UNet	1.347	1.038	38.3	11.431	8.698	22.8
Attention UNet	1.474	1.128	35.5	11.376	8.603	22.9
ViT Hybrid	1.330	1.025	39.4	12.084	9.258	17.7

Качественный анализ

• U-Net: Лучшие общие метрики, чёткие границы глубины, occasional checkerboard артефакты от transposed convolutions
• Attention U-Net: Похож на U-Net, но с более гладкими предсказаниями, attention-гейты дают минимальное улучшение на малых датасетах
• ViT Hybrid: Лучшее понимание глобальной структуры, но более шумные локальные предсказания, требует значительно больше данных и вычислений

Быстрый старт

Требования

• Python 3.10+
• PyTorch 2.0+
• Рекомендуется GPU с поддержкой CUDA

Установка

# Клонировать репозиторий
git clone https://github.com/imaseriousactor/Depth-From-Focus-Image-Restoring.git
cd Depth-From-Focus-Image-Restoring

# Установить зависимости
pip install -r requirements.txt

Обучение

На Kaggle (рекомендуется):

1. Создайте новый Notebook на Kaggle
2. Добавьте датасеты через "Add Input" (raw_data и diode_train)
3. Загрузите код обучения (Ячейки 1-4)
4. Запустите все ячейки последовательно
5. Скачайте веса моделей из панели Output

Локально:

# Обучить одну модель (изменить model_name в train.py)
python train.py

# Варианты: unet, attention_unet, vit_hybrid

Оценка

# Сравнить все модели и сгенерировать визуализации
python evaluate.py

Веб-демонстрация

# Запустить интерфейс Gradio
python app/app.py

Живое демо: \https://huggingface.co/spaces/imaseriousactor/depth-estimation-demo

Структура проекта

Depth-From-Focus-Image-Restoring/
├── data/
│   ├── nyu_depth_dataset.py          # Загрузчик NYU датасета
│   ├── diode_dataset.py              # Загрузчик DIODE датасета
│   ├── convert_diode_to_hdf5.py      # Конвертер DIODE формата
│   └── raw/                          # Сырые данные (не в git)
│       ├── train-000000/
│       └── val/
├── models/
│   ├── unet.py                       # Архитектура U-Net
│   ├── attention_unet.py             # Архитектура Attention U-Net
│   └── vit_hybrid.py                 # Архитектура ViT Hybrid
├── notebooks/
│   ├── eda_01.ipynb                  # Исследовательский анализ данных
│   └── outputs/                      # Сгенерированные визуализации
├── checkpoints/                      # Веса моделей (не в git)
├── app/
│   └── app.py                        # Веб-демонстрация на Gradio
├── train.py                          # Скрипт обучения
├── evaluate.py                       # Скрипт оценки
├── requirements.txt                  # Зависимости
├── .gitignore                        # Правила git ignore
└── README.md                         # Этот файл

Детали обучения

Параметр	Значение
Оптимизатор	Adam (lr=1e-4)
Функция потерь	L1 Loss (MAE)
Шедулер	ReduceLROnPlateau (patience=5, factor=0.5)
Batch Size	4 (UNet/Attention UNet), 2 (ViT)
Размер изображения	256x384 (UNet/Attention UNet), 128x192 (ViT)
Эпохи	30
Оборудование	Kaggle GPU T4 x2