Хабрахабр

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

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

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

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

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

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

Самый популярный подход для задачи сегментации: взять Unet с предобученным энкодером и минимизировать сумму $BCE$ (binary cross-entropy) и $DICE$ (Sørensen–Dice coefficient).

Какие проблемы возникают при таком подходе в задаче сегментации дефектов?

  • Даже если нам кажется, что дефектов на фотографии очень много, что она очень грязная и сильно потрёпана временем, всё равно площадь, занимаемая дефектами, гораздо меньше неповреждённой части изображения. Чтобы решить эту проблему, можно увеличить вес положительного класса в $BCE$, и оптимальным весом будет отношение количества всех чистых пикселей к количеству пикселей, принадлежащих к дефектам.
  • Вторая проблема в том, что если мы используем 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 карточек. Это улучшает сходимость сети.

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

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

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

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

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

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

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

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

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

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

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

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

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

Можно снова взять 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 и ошибкой дискриминатора.

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

Чтобы его посчитать, нужно прогнать предсказание сети и исходную фотографию через cеть VGG-16, взять карты признаков на нижних слоях и сравнить их по MSE. После выкатки модели в production мы продолжили эксперименты и пришли к выводу, что лучше минимизировать не попиксельный L1-loss, а perceptual loss. Такой подход закрашивает больше областей и помогает получить более красочную картинку.

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

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

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

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

Показать больше

Похожие публикации

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *

Кнопка «Наверх»