Как мы оживили DPED: ИИ-проект для улучшения мобильных фото до снимков с «зеркалок»

На связи команда регионального научно-образовательного центра «Искусственный интеллект и анализ больших данных» при НГТУ им. Р. Е. Алексеева.

Сегодня практически у каждого крупного производителя смартфонов есть собственные системы улучшения качества фотографий, основанные на искусственном интеллекте. Помимо этого, существует множество аналогичных открытых моделей нейросетей. Возникает вопрос: зачем тогда мы решили оживить DPED (Deep Photo Enhancement Dataset) — созданный для тех же целей проект 2017 года? В статье мы ответим на этот вопрос, займемся оживлением и протестируем нашу модель DPED на фотографиях с планшета KVADRA_T.

О проекте

Для начала чуть подробней о сути проекта. Однажды группе ИИ-энтузиастов пришла идея: а что если обучить искусственный интеллект улучшать смартфонные снимки до профессиональных с помощью парных фотографий? Задумка понравилась. Для сбора датасета выбрали актуальные в то время Sony Xperia Z, iPhone 3GS, BlackBerry Passport и цифровую зеркалку Canon EOS 70D в качестве эталона. Модель обучили улучшать фотографии, сделанные на смартфонах, в соответствии с такими же изображениями, полученными с камеры. Проект реализовали, исходный код опубликовали на GitHub, а подробное описание — на arXiv.org. Но что же в нем интересного сейчас, почти десять лет спустя?

DPED не просто применяет косметические фильтры. Датасет способен значительно улучшать фотографии на уровне структуры изображения, детализации, динамического диапазона и цветопередачи, приближая результат именно к профессиональной зеркальной фотокамере, а не просто «улучшая» фото.

Открытый исходный код и доступность датасета DPED позволяют легко адаптировать, изменять и дорабатывать модель. Это дает большие возможности исследовательскому сообществу и разработчикам мобильных приложений. Именно гибкость и понятность архитектуры делают DPED хорошим кандидатом для дальнейших экспериментов и улучшений.

Наконец, несмотря на возраст, DPED по-прежнему демонстрирует конкурентные результаты. Нам было интересно оценить, насколько их можно улучшить, адаптировав проект для современных библиотек и «железа», а также сравнить с современными решениями, но уже на новом, актуальном датасете.

Но не будем забегать вперед: это тема следующих статей, а сейчас мы будем возвращать DPED к жизни.

Подготовка к запуску

Для начала необходимо запустить исходный код проекта. На момент нашей работы DPED исполнилось уже семь лет. Для машинного обучения и компьютерного зрения это большой срок, и мы ожидали трудности на разных этапах.

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

Далее мы расскажем, как запустить проект и исправить ошибки на относительно свежих версиях библиотек, совместимых с Python 3.11. Для начала склонируем проект себе с помощью команды:

git clone https://github.com/aiff22/DPED

В нем увидим примерно следующую структуру каталогов:

.
└── DPED
	├── dped
	├── load_dataset.py
	├── models
	├── models_orig
	│   ├── blackberry_orig.data-00000-of-00001
	│   ├── blackberry_orig.index
	│   ├── iphone_orig.data-00000-of-00001
	│   ├── iphone_orig.index
	│   ├── sony_orig.data-00000-of-00001
	│   └── sony_orig.index
	├── models.py
	├── README.md
	├── results
	├── ssim.py
	├── test_model.py
	├── train_model.py
	├── utils.py
	├── vgg_pretrained
	├── vgg.py
	└── visual_results
5 directories, 14 files

Необходимо скачать предобученную модель VGG-19, умеющую классифицировать изображения по тысяче категорий объектов, и перенести ее в папку vgg_pretrained.

Затем нужен датасет из патчей, на котором обучалась эта модель.

Патчи — это небольшие, чаще всего квадратные фрагменты исходного изображения, которые извлекают из полноразмерных фотографий для обучения нейросетей. Например, из одной фотографии 4000×3000 пикселей можно нарезать десятки патчей размером 100×100 пикселей. Патчи идут в парах: один с мобильного устройства, другой — с профессиональной камеры, снятый с того же ракурса. По ним модель учится восстанавливать недостающие детали, улучшать резкость, цвета и т. д. Если же отдавать для обучения снимки целиком, это увеличит затраты памяти и сложность процессов.

Архив полных снимков датасета объемом 54 ГБ можно скачать с OpenDataLab (зеркало). Мы пробовали скачивать их через http, но спустя несколько дней загрузка обрывалась. Советуем скачивать не напрямую через браузер, а через git lfs. Но учтите, что lfs делает копию загруженных объектов: в нашем случае для скачивания суммарно потребовалось 134 ГБ свободного пространства, но зато мы смогли сделать это за час.

Git lfs устанавливается на Linux через команду git lfs install, а на Windows — с официального сайта. Затем нужно инициализировать git lfs через git lfs install --skip-smudge, эта команда просто склонирует репозиторий и создаст ссылки на файлы. Далее можете скачать архив с оригинальными фотографиями через команду git pull —include “original_images.gz”, если вам интересно, как формировался датасет. В нашем запуске мы эти фото не используем.

После скачивания архив с патчами нужно распаковать в папку dped, после чего в ней появятся три папки: sony, iphone и blackberry. Теперь все готово к запуску кода. Итоговая структура проекта:

.
├── dped
│   ├── blackberry
│   │   ├── test_data
│   │   │   ├── full_size_test_images
│   │   │   └── patches
│   │   │   	├── blackberry
│   │   │   	└── canon
│   │   └── training_data
│   │   	├── blackberry
│   │   	└── canon
│   ├── iphone
│   │   ├── test_data
│   │   │   ├── full_size_test_images
│   │   │   └── patches
│   │   │   	├── canon
│   │   │   	└── iphone
│   │   └── training_data
│   │   	├── canon
│   │   	└── iphone
│   └── sony
│   	├── test_data
│   	│   ├── full_size_test_images
│   	│   └── patches
│   	│   	├── canon
│   	│   	└── sony
│   	└── training_data
│       	├── canon
│       	└── sony
├── models
├── models_orig
├── results
├── vgg_pretrained
└── visual_results
33 directories

В проекте DPED предусмотрено два основных сценария работы:

• запуск предобученной модели на подготовленных данных (с уже проставленными весами),

• обучение собственной модели на оригинальном датасете DPED (с нуля, или fine-tuning)

Запуск предобученной модели DPED на подготовленных данных

Первый сценарий позволит нам быстро убедиться, что код запускается корректно, и получить видимый результат без затрат на обучение.

При запуске test_model.py мы столкнулись с проблемой сохранения изображений: новые версии Pillow не поддерживают сохранение RGB-изображений в формате float32, а именно — в нем модель возвращала массивы. Ранее это работало, но с новыми версиями Pillow требуется преобразование изображения в тип uint8.

Мы ограничили значения пикселей диапазоном [0, 1] и масштабировали до [0, 255] с приведением к uint8. Это стандартный этап подготовки изображения к сохранению в форматах JPG/PNG, и он не влияет на итоговое качество, так как внутри модели все вычисления по-прежнему идут в float32. Файл test_model.py:

После этого улучшенные фотографии начали сохраняться в таком виде:

Причина ошибок в том, что, хоть мы ограничили значения пикселей диапазоном [0, 1] и масштабировали до [0, 255] с приведением к uint8, в процессе генерации изображения значения могут выходить за эти пределы (например, –0.03 или 1.07). Это приводит к некорректному сохранению, особенно при преобразовании в uint8. Чтобы гарантировать, что все значения остаются в допустимом диапазоне перед масштабированием, мы использовали метод .clip(0, 1) в test_model.py:

В результате мы смогли сохранить выходные изображения и перейти ко второму сценарию.

Обучение собственной модели на датасете DPED

Здесь нам интересно проверить, как поведет себя модель с новым железом и новыми версиями библиотек, а также сравнить качество изображений, полученных нашей моделью, с оригинальными результатами 2017 года. Этот опыт можно использовать в будущем для обучения обновленной версии модели на снимках с современных смартфонов.

В оригинальном репозитории DPED есть уже готовые модели для трех смартфонов (Sony, iPhone, BlackBerry), обученные авторами исследования. Мы протестировали их, но затем заново обучили модель самостоятельно, используя тот же датасет, но с обновленным окружением (новые версии TensorFlow, Python и т. д.) и с учетом аппаратных ограничений. Теперь обо всем по порядку.

При запуске обучения модели на предоставленном авторами статьи датасете возникает ошибка scipy.misc.imread. Этот модуль устарел и был удален в версии scipy 1.12.0. Чтобы им воспользоваться, необходимо скачивать версию 1.11.0.

Для этого версия Python должна быть более поздней, чем 3.8, и более ранней, чем 3.9, а это может грозить еще большим количеством ошибок импортов и несоответствий версионности библиотек. Так что мы просто заменили метод imread из scipy на imread из библиотеки pillow или imageio. После этого ошибок с обучением у нас не возникло. Файл load_dataset.py:

from imageio import imread
misc.imread = imread

Обучение требовательно к объему видеопамяти, и это указано на GitHub проекта в разделе «возможных ошибок». Если видеопамяти мало, то придется уменьшать batch size — количество примеров, обрабатываемых за один шаг. По умолчанию параметр равен 50, и при нехватке ресурсов для обучения в логах появится сообщение о недостаточном количестве видеопамяти. В нашем случае такое произошло с Nvidia RTX 3060 Ti (8 ГБ VRAM), так что пришлось поэкспериментировать с параметром batch_size.

Уменьшение батча — это распространенная мера при ограниченной видеопамяти, и это не проходит бесследно. Основные последствия:

• Снижается стабильность градиентов и сходимость. Большие батчи дают более точную оценку градиентов, что делает обучение стабильнее. Малые батчи, наоборот, вносят больше «шума» — из-за этого возможна менее стабильная сходимость, особенно в начале обучения.

• Колебания функции потерь (loss). При малом batch_size loss-функция может колебаться более резко от шага к шагу, что требует тонкой настройки learning rate и других гиперпараметров.

• Качество финальной модели. Оно зависит от задачи, но в общем случае меньший batch_size может немного ухудшить обобщающую способность модели, особенно если не были пересмотрены другие параметры (learning rate, количество эпох и т.д.).

• Влияние на метрики качества изображений SSIM (индекс структурного сходства) и PSNR (пиковое отношение сигнала к шуму). Например, SSIM может быть чуть ниже, если модель недостаточно стабильно обучалась. Разница обычно некритична, если модель дообучена правильно, с использованием чуть большего количества эпох или адаптации learning rate.

Для видеокарты Nvidia RTX 3060 Ti (8 ГБ VRAM) мы понизили batch_size с 50 до 35. В результате модель обучилась без ошибок, и метрики качества остались на уровне или выше оригинальных результатов 2017 года. Кроме того, для обучения взяли Nvidia RTX 2080 Ti (11 ГБ RAM), и ее удалось запустить на параметрах, которые использовали авторы проекта.

Для обучения мы взяли две конфигурации: с Nvidia RTX 2080 Ti 11 ГБ и Nvidia RTX 3060 Ti 8 ГБ. В первой конфигурации мы использовали гиперпараметры по умолчанию, на которых модель обучали и сами авторы. На скриншоте к ним даны пояснения:

Мы посчитали метрики качества изображений SSIM (индекс структурного сходства) и PSNR (пиковое отношение сигнала к шуму) и сравнили их с приведенными в статье. Раздел Nvidia Titan X — это итоги теста предобученной авторами проекта модели из оригинального репозитория. Разделы RTX — итоги тестов модели, обученной на более современном железе в более современном окружении (Python, TensorFlow). По горизонтали указаны конкретные датасеты для разных моделей телефонов:

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

На момент написания статьи мы внесли все необходимые изменения в локальную копию проекта DPED, включая:

• совместимость с Python 3.11,

• замену scipy.misc.imread на imageio.imread,

• обновление сохранения изображений с float32 на uint8,

• обновление зависимостей и адаптацию под новые версии библиотек (NumPy, Pillow, TensorFlow и т.д.),

• правки, связанные с ограничениями видеопамяти (например, batch_size).

Мы планируем выложить отдельную ветку проекта на GitHub, в которой:

• Все изменения будут оформлены как pull request.

• Будет приведен список модифицированных файлов и объяснений к ним.

• Будет добавлена инструкция по установке окружения через requirements.txt и/или Conda.

Улучшаем фотографии

Пора дать DPED реальную задачу. Для тестирования мы взяли фотографии, сделанные планшетом YADRO KVADRA_T. В момент тестирования он оказался под рукой, так что мы смогли быстро провести эксперимент «в боевых условиях», пока проверяли работоспособность проекта после обновлений. Слева — оригиналы, справа — результаты обработки:

В следующей статье мы приведем метрики качества этих изображений, а пока предлагаем оценить прогресс визуально.

Ценность DPED сегодня

На момент публикации в 2017 году проект DPED был значительным достижением для постобработки мобильных фотографий через глубокое обучение. DPED — это не просто модель, это понятная и открытая архитектура, на которой мы можем четко отследить, как изменилось качество мобильных фото за последние годы. Особенно интересно сравнение «старой, но обновленной» модели с тем, что предлагают сегодня ведущие производители смартфонов и актуальные открытые решения.

Мы можем четко контролировать ход всех наших экспериментов с DPED и собственными датасетами и проверять, как именно различные компоненты модели и данные влияют на конечный результат.

Наша команда продолжит работу над этим проектом. Далее в планах: собрать собственный датасет по примеру разработчиков, обновить кодовую базу до актуальных библиотек, обучить свою модель на основе их архитектуры и сравнить полученные результаты. В следующей статье поделимся достижениями и выводами.

А что вы думаете про использование DPED в наши дни? Есть ли смысл возвращать эту технологии или она должна кануть в лету?