Реставрируем фотографии с помощью нейросетей

Всем привет, я работаю программистом-исследователем в команде компьютерного зрения Mail.ru Group. Ко Дню Победы в этом году мы решили сделать проект по реставрации военных фотографий. Что такое реставрация фотографий? Она состоит из трех этапов:

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

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

Поиск дефектов

Мы хотим найти все пиксели, относящиеся к дефектам на загруженной фотографии. Для начала нам нужно понять, какие фотографии военных лет будут загружать люди. Мы обратились к организаторам проекта «Бессмертный полк», которые поделились с нами данными. Проанализировав их, мы заметили, что люди зачастую загружают портреты, одиночные или групповые, на которых есть умеренное или большое количество дефектов.

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

На разметку пикселей, относящихся к дефектам на одной фотографии, может уходить от одного часа до целого рабочего дня, поэтому за несколько недель сложно собрать выборку больше чем из 100 фотографий. Поэтому мы старались как-то дополнять наши данные и писали дефекты самостоятельно: брали чистую фотографию, наносили на нее искусственные дефекты и получали маску, показывающую нам, на какие именно части изображения нанесены дефекты. Главной частью нашей обучающей выборки стали 79 фотографий, размеченные вручную, из них 11 штук мы перенесли в тестовую выборку. Самый популярный подход для задачи сегментации: взять Unet с предобученным энкодером и минимизировать сумму $$ (binary cross-entropy) и $$ (S?rensen–Dice coefficient). Какие проблемы возникают при таком подходе в задаче сегментации дефектов?

• Даже если нам кажется, что дефектов на фотографии очень много, что она очень грязная и сильно потрёпана временем, всё равно площадь, занимаемая дефектами, гораздо меньше неповреждённой части изображения. Чтобы решить эту проблему, можно увеличить вес положительного класса в $$, и оптимальным весом будет отношение количества всех чистых пикселей к количеству пикселей, принадлежащих к дефектам.
• Вторая проблема в том, что если мы используем Unet из коробки с предобученным энкодером, например Albunet-18, то теряем много позиционной информации. Первый слой Albunet-18 состоит из свертки с ядром 5 и stride равным двум. Это позволяет сети быстро работать. Мы пожертвовали временем работы сети ради лучшей локализации дефектов: убрали max pooling после первого слоя, уменьшили stride до 1 и уменьшили ядро свёртки до 3.
• Если мы будем работать с маленькими изображениями, например, сжимая картинку до 256 х 256 или 512 х 512, то маленькие дефекты будут просто пропадать из-за интерполяции. Поэтому нужно работать с большой картинкой. Сейчас в production мы сегментируем дефекты на фотографии 1024 х 1024. Поэтому необходимо было обучать нейросеть на больших кропах больших изображений. А из-за этого возникают проблемы с маленьким размером батча на одной видеокарте.
• Во время обучения у нас на одну карточку помещается около 20 картинок. Из-за этого оценка среднего и дисперсии в BatchNorm-слоях получается неточной. Решить эту проблему нам помогает In-place BatchNorm, который, во-первых, экономит память, а во-вторых, у него есть версия Synchronized BatchNorm, которая синхронизирует статистики между всеми карточками. Теперь мы считаем среднее и дисперсию не по 20 картинкам на одной карточке, а по 80 картинкам с 4 карточек. Это улучшает сходимость сети.

В конце концов, увеличив вес $$, поменяв архитектуру и использовав In-place BatchNorm, мы начали искать дефекты на фотографии. Но задёшево можно было сделать ещё чуть лучше, добавив Test Time Augmentation. Мы можем прогнать сеть один раз на входном изображении, потом отзеркалить его и прогнать сеть ещё раз, это может помочь нам найти маленькие дефекты.

В результате наша сеть сошлась на четырёх GeForce 1080Ti за 18 часов. Inference занимает 290 мс. Получается достаточно долго, но это плата за то, что мы хорошо ищем небольшие дефекты. Валидационный $$ равен 0,35, а $$ — 0,93.

Реставрация фрагментов

Решить эту задачу нам снова помог Unet. На вход ему мы подавали исходное изображение и маску, на которой единицами отмечаем чистые пространства, а нолями — те пиксели, которые хотим закрасить. Данные мы собирали следующим образом: брали из интернета большой датасет с картинками, например OpenImagesV4, и искусственно добавляли дефекты, которые похожи по форме на те, что встречаются в реальной жизни. И после этого обучали сеть восстанавливать недостающие части.

Как мы можем модицифировать Unet для этой задачи?

Можно использовать вместо обычной свёртки частичную (Partial Convolution). Её идея в том, что при сворачивании региона картинки с каким-то ядром мы не учитываем значения пикселей, относящихся к дефектам. Это помогает сделать закрашивание точнее. Пример из статьи NVIDIA. На центральной картинке они использовали Unet с обычной свёрткой, а на правой — с Partial Convolution:

Мы обучали сеть 5 дней. В последний день мы замораживали BatchNorm, это помогло сделать границы закрашиваемой части изображения менее заметными. Картинку 512 х 512 сеть обрабатывает за 50 мс. Валидационный PSNR равен 26,4. Однако в этой задаче нельзя безоговорочно доверять метрикам. Поэтому мы сначала прогнали на наших данных несколько хороших моделей, анонимизировали результаты, а потом проголосовали за те, что нам больше понравились. Так мы и выбрали финальную модель.

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

Вот пример работы алгоритма:

Раскрашивание

Мы сегментировали дефекты и закрасили их, третий шаг — реконструкция цвета. Напомню, что среди фотографий «Бессмертного полка» очень много одиночных или групповых портретов. И мы хотели, чтобы наша сеть хорошо с ними работала. Мы решили сделать свою колоризацию, потому что ни один из известных нам сервисов не раскрашивает портреты быстро и хорошо.

На GitHub есть популярный репозиторий для раскрашивания фотографий. В среднем, он хорошо делает эту работу, но у него есть несколько проблем. Например, он очень любит раскрашивать одежду в синий цвет. Поэтому его мы тоже отвергли. Итак, мы решили сделать нейросеть для колоризации. Самая очевидная идея: брать чёрно-белое изображение и предсказывать три канала, красный, зелёный и синий. Но, вообще говоря, мы можем упростить себе работу. Можем работать не с RGB-представлением цвета, а с YCbCr-представлением. Компонента Y — это яркость (luma). Загружаемое черно-белое изображение и есть Y канал, мы будем его переиспользовать. Оставалось спрогнозировать Cb и Cr: Cb — это разница голубого цвета и яркости, а Cr — это разница красного цвета и яркости.

Почему мы выбрали YCbCr-представление? Глаз человека более восприимчив к перепадам яркости, чем к изменениям цвета. Поэтому мы переиспользуем Y-компоненту (яркость), то, к чему глаз изначально хорошо восприимчив, и прогнозируем Cb и Cr, в которых мы можем чуть больше ошибаться, поскольку «фальш» в цветах человек замечает меньше. Этой особенностью начали активно пользоваться на заре цветного телевидения, когда пропускной способности канала не хватало, чтобы передавать все цвета полностью. Изображение передавали в YCbCr, передавали Y-компоненту без изменений, а Cb и Cr сжимали в два раза.

Как собрать baseline

Можно снова взять Unet с предобученным энкодером и минимизировать L1 Loss между настоящим CbCr и прогнозируемым. Мы хотим раскрашивать портреты, поэтому кроме фотографий из OpenImages нам нужно добавить специфические для нашей задачи фотографии.

Где взять цветные фотографии людей в военной форме? В интернете есть люди, которые в качестве хобби или на заказ раскрашивают старые фотографии. Они это делают крайне аккуратно, стараясь полностью соблюсти все нюансы. Раскрашивая форму, погоны, медали они обращаются к архивным материалам, поэтому результату их работы можно доверять. В общей сложности мы использовали 200 фотографий, раскрашенных вручную. Второй полезный источник данных — это сайт Рабоче-Крестьянской Красной Армии. Один из его создателей сфотографировался практически во всех возможных вариантах военной формы времен Великой Отечественной войны.

На некоторых фотографиях он повторял позы людей с знаменитых архивных фотографий. Особенно хорошо, что он снимался на белом фоне, это позволило нам очень хорошо аугментировать данные, добавляя разные природные объекты на задний план. Также мы использовали обычные современные портреты людей, дополняя их знаками различия и прочими атрибутами одежды военного времени. Мы обучили AlbuNet-50 — это Unet, в котором в качестве энкодера используется AlbuNet-50. Сеть начала давать адекватные результаты: кожа розовая, глаза серо-голубые, погоны желтоватого цвета. Но проблема в том, что она раскрашивала картинки пятнами. Это связано с тем, что с точки зрения L1-ошибки иногда бывает выгоднее не делать ничего, чем пытаться предсказывать какой-то цвет.

Мы сравнением наш результат с фотографией Ground Truth — ручной колоризацией художника под ником Klimbim

Как решить эту проблему? Нам нужен дискриминатор: нейронная сеть, которой мы на вход будем подавать изображения, и она будет говорить, насколько реалистично это изображение выглядит. Ниже одна фотография раскрашена вручную, а вторая — нейросетью. Как вы думаете, какая?

В качестве дискриминатора мы используем дискриминатор из статьи Self-Attention GAN. Это небольшая свёрточная сеть, в последние слои которой встроен так называемый Self-Attention. Он позволяет больше «обращать внимание» на детали изображения. Также мы используем спектральную нормализацию. Точное объяснение и мотивацию можно найти в статье. Мы обучили сеть с комбинацией L1-loss и ошибки, возвращаемой дискриминатором. Теперь сеть лучше раскрашивает детали изображения, а фон получется более консистентным. Еще один пример: слева результат работы сети, обученной только с L1-loss, справа — с L1-loss и ошибкой дискриминатора.

На четырёх Geforce 1080Ti обучение заняло два дня. Сеть отрабатывала за 30 мс на картинке 512 х 512. Валидационная MSE — 34,4. Как и в задаче inpainting, метрикам можно верить не до конца. Поэтому мы отобрали 6 моделей, которые имели лучшие метрики на валидации, и вслепую голосовали за лучшую модель. После выкатки модели в production мы продолжили эксперименты и пришли к выводу, что лучше минимизировать не попиксельный L1-loss, а perceptual loss. Чтобы его посчитать, нужно прогнать предсказание сети и исходную фотографию через cеть VGG-16, взять карты признаков на нижних слоях и сравнить их по MSE. Такой подход закрашивает больше областей и помогает получить более красочную картинку.

Выводы и заключение

Unet — это классная модель. В первой задаче сегментации мы столкнулись с проблемой при обучении и работе с картинками большого разрешения, поэтому используем In-Place BatchNorm. Во второй задаче (Inpainting) вместо обычной свёртки мы использовали Partial Convolution, это помогло достичь лучших результатов. В задаче колоризации к Unet мы добавили небольшую сеть-дискриминатор, которая штрафовала генератор за нереалистично выглядящее изображение и использовали perceptual loss.

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

Мы взяли некоторые фотографии проекта «Военный альбом» и обработали их нашими нейросетями. Вот такие результаты получили:

А здесь можно посмотреть их в оригинальном разрешении и на каждом этапе обработки.