company_banner

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



    Всем привет, я работаю программистом-исследователем в команде компьютерного зрения 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. Мы можем прогнать сеть один раз на входном изображении, потом отзеркалить его и прогнать сеть ещё раз, это может помочь нам найти маленькие дефекты.



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

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

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



    А здесь можно посмотреть их в оригинальном разрешении и на каждом этапе обработки.
    Mail.ru Group
    1 127,52
    Строим Интернет
    Поделиться публикацией

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

      +3
      Было бы неплохо, если бы результат можно было бы выгрузить, например, в формате Photoshop, где изображение было бы разделено на слои для колоризации. Так можно было бы откорректировать, если результат не совсем тот, что ожидался.
        +1
        Дублируете слой, ставите верхний в режим color. Нижний обесцвечиваете.
        Или ставите инструменту палец режим color если нужно двигать цвет.
        +2
        В целом результат впечатляет, но только если на фото люди в военной форме, идеально получается если на фото люди по пояс. Но вот если дети или женщины в платьях, то все закрашивается зеленоватым цветом.

          +1
          Спасибо! Мы обновили модель колоризации, и сейчас картинка стала более насыщенной. Проблемы с зеленоватыми участками все еще есть, будем совершенствоваться.
        • НЛО прилетело и опубликовало эту надпись здесь
          +6
          Картинку, раскрашенную нейросетью, проще всего оказалось определить по цвету «красного» крестика на сумке =)

          Тем не менее, очень впечатляет!
            +1
            Не, проще определить по цвету кожи (например, кисти рук). Нейросеть кожу человека не распознает — хорошо видно на всех прмиерах.
            +1
            Какая-то любовь к зеленому цвету. Скорее всего особенность обучения на военной форме/ экипировке.
            The hat was available in eight colors: Lipstick Red, Tangerine, Flamingo, Canary Yellow, Chartreuse, Blush Pink, Rose Pink and Tan.

            image
            Оригинальный цвет:
            image
              0
              Весьма неплохо

                0
                Молодцы, отлично работает, респект!
                  +3
                  А не пробовали натаскать алгоритм отдельно на награды? Там же цвета лент строго определённые и, по идее, должны достаточно просто определятся, потом эту информацию как опорную использовать.
                    +1
                    Верно, отдельной сеткой выискивать награды — классифицировать их, и потом фрагмент раскрашивать сектой тренированной именно на эту награду.
                    В реставрации исторических фотографий — очень важна именно достоверность. Иначе потом будем в сети находить холивары ссылками на неверно раскрашенные фото.
                      0
                      И погоны. И детали формы. И затем цвет формы выбирать с оглядкой на перечень найденного
                        0
                        Идея отличная, это действительно помогло бы улучшить качество колоризации деталей. Мы отказались от идеи детектировать медали, погоны и другие детали одежды так как не хотели усложнять итоговую модель и процесс разработки.
                    +1
                    Морскую форму красит плохо, в частности c гюйзом (воротник матроса) вообще беда.
                    image
                      +2
                      Сеть, насколько я понял, тренирована на белых людях.
                      А какого-нибудь смуглого таджика она насколько правильно раскрасит?

                      PS Да, я знаю, что есть белокожие таджики, но есть и с кожей цвета кофейных зёрен.
                        +1
                        Попробовал с околограничным значением цвета кожи
                        Скрытый текст


                        +1
                        Немного не по теме, но всё же. Вы пишете, что делали это к 9 мая. Я в конце апреля занялся подготовкой фотографии к празднику. Много часов потратил. Так где же Вы были? Я по нескольку раз в день захожу на главную страницу mail.ru. Но анонса сервиса не видел.
                          –1
                          Привет! Никогда не бывает достаточно анонсирования, даже в таком массовом деле. Мы подключили огромные ресурсы (СМИ, внутренние размещения, даже промо в соц. сетях), и множество людей увидели его, но вы в их ряды каким-то образом не попали( Это печаль.

                          Подписывайтесь на нас в соц сетях (VK, FB, Twitter). А еще лучше поделитесь в комменте, где нужно разместить анонс, чтобы вы его 100% увидели?)
                            0
                            Извините, но я просто не пользуюсь соцсетями. Но ежедневно пользуюсь почтой, в том числе через браузер. Неужели этого оказалось не достаточно? При просмотре почты мне же показывают всякие баннеры типа «помогите! спасите!» или вот прямо сейчас информация о новом интерфейсе. Можно было бы на их месте.
                              0
                              Приятно, что регулярно пользуетесь.) А интерфейс новый не пробовали еще?
                                –1
                                Вот интересно, какой это умник меня отминусовал? И что ему не понравилось в словах, не относящихся к статье и обращённых к совершенно конкретному человеку? Его не учили, что влезать со своим мнением в чужой разговор как минимум не культурно? Вот это и есть поколение Y во всём своём проявлении (это, правда, из другой публикации здесь, на хабре). Не разобрался, увидел незнакомые слова или ещё что и минусанул. Зачем? Сам не знает.
                                  +2

                                  Это не личный разговор, а открытые комментарии.

                                0
                                Дык он же говорит, на главной mail.ru
                            +2
                            Ну в принципе довольно неплохо но еще очееень далеко до крепкой 8)) Проблема в том что к примеру при использовании фотографий просто сделанных черно белыми он вырезал некоторые детали с фоток которые считал браком XD К примеру на девушке в порванных джинсах он попытался зашить дырки на джинсах XD А там где не смог зашить просто увеличил надорванность.
                              +1
                              Не очень хорошо справляется с горизонтальными черно-белыми полосами
                              Оригинал


                              Восстановленный

                                +1
                                Отличная статья, в которой грамотно разбиты разные этапы обработки изображения. Вдвойне круто, что сервис доступен онлайн. Интересно, военные фотография — это самая сложная категория изображений для восстановления (так как присутствуют разные типы дефектов)? Применимы ли похожие алгоритмы для интро- или экстраполяции медицинских изображений (где тоже, например, возможны вырезанные фрагменты)?
                                  +4
                                  Я бы остановился на удалении дефектов. Погоны с просветами серо-буро-малинового цвета, ленты медалей цвета «какой попался» (обычно красный), цвет формы часто туда же. Слишком уж сложная задача — сделать это всё нормально. Но это может я заклёпочник такой.

                                  Поигрался заодно с фото самолётов с людьми или без ради интереса — обычно всё плохо, конечно, но попалось фото, которое вышло почти отлично.
                                  Было
                                  image

                                  Стало
                                  image
                                    0
                                    Спасибо что поделились :)

                                    Сервис возвращает три фотографии:
                                    1) колоризованную с закрашенными дефектами
                                    2) колоризованный оригинал
                                    3) оригинал с закрашенными дефектами.

                                    Вы можете скачать любой из этих вариантов.
                                    Третья фотография — ровно то, что вам нужно.
                                      0
                                      Ага, я уже разобрался что сам себя запутал, спасибо, только пост не успел отредактировать вовремя. Добавил сервис в закладки, пригодится, спасибо за хорошую штуку.
                                        +1
                                        Было бы классно если при загрузке изображения можно было указать какой-то фрагмент и выбрать для него цвет, чтобы был эталон для раскрашивания, либо сервис в зависимости от затененности того или иного фрагмента предлагал выбор цвета
                                          0
                                          Это хорошая идея, люди работают в этом направлении: richzhang.github.io/ideepcolor

                                          Мы решили сделать наш сервис максимально простым для пользователя, поэтому не добавляли возможность выбора цвета.
                                      +1
                                      Не помню где читал, что когда мало фоток для обучения, каждую фотографию скармливают по многу раз, каждый раз поворачивая на, скажем, один градус, тогда будет на 360 фотографий больше. Природа сверточных сетей такова, что для них это иногда хорошо срабатывает. В вашем случае надо поворачивать фото и сделаную вручную разметку дефектов.
                                        0
                                        Вы правы, это помогло поднять нам качество. Мы аугментировали данные поворотами, случайными отражениями, вырезанием и масштабированием частей изображения.
                                        +1
                                        жаль что дети сейчас не смотрят хронику чернобелую
                                        возможно что совместно с культура рф вам удастся оцифровать союзкиножурнал за 1941-45
                                        и сделать его цветным
                                        пытался посадить родственицу 15 л за хронику
                                        а ей скучно
                                        худ фильмы к сожалению именно худ фильмы

                                          0
                                          Если предположить, что соотношение яркостей объекта на ЧБ фото несет хоть какую-то информацию о цвете, то можно предложить расцвечивание фото в два этапа.
                                          На первом этапе выделять объекты чей цвет более-менее достоверно возможно установить по форме — участки кожи человека, небо/море/облака, растительность, а в контексте данного примера еще и цвет знаков различия и наград. Далее востанавливаем цвет только этих объектов.
                                          На втором этапе уже подаем на вход два снимка — исходный чернобелый и исходный с востановленными «достоверными цветами», пытаясь по востановленным цветам и соотношениям яркости угадать остальные недостающие цвета.

                                          Соответственно это должны быть N+1 сетей натренированных на специальных наборах данных, которые вобщем-то можно получить из наборов цветных фото.

                                          Мысль выглядит достаточно очевидной, но нигде не читал о таком подходе. Я чего-то не понимаю и это по каким-то причинам невозможно или не имеет смысла?
                                            0
                                            Быстрый ответ: надежнее сделать большую сеть, которая будет раскрашивать изображение за один проход.

                                            Если по частям:
                                            1) Подход, про который вы рассказываете, называется бустингом. Мы берем N базовых алгоритмов, и учим каждый следующий алгоритм исправлять ошибки предыдущего. Я не встречал в современных работах по колоризации такого подхода. Думаю, это связано с тем, что сеть проще обучить одну большую сеть сразу предсказывать цвета всех пикселей чем учить много маленьких сетей исправлять ошибки друг друга. Для лучшего визуального качества сейчас часто используют GANы, которые тяжело учить и так, поэтому дополнительно усложнять архитектуру генератора не хочется.

                                            Вторая сложность реализации вашей идеи — «это должны быть N+1 сетей натренированных на специальных наборах данных, которые вобщем-то можно получить из наборов цветных фото» — непонятно по каким критериям отбирать наборы цветных фотографий.

                                              0
                                              То ли я плохо объяснил, то ли не до конца понял ответ…
                                              Я не говорил об исправлении ошибок, я говорил о опоре на данные полученные от других сетей.
                                              Вот пример — допустим нам надо красить преимущественно фото сделанные на улице на которых практически всегда есть небо. Тогда мы берем набор цветных фото (набор 1) и обесцвечиваем на них все кроме неба (набор 2), и обесцвечиваем вообще все (набор 3).
                                              Сеть А, мы на наборах 3 и 2 обучаем узнавать на фото небо и раскрашивать его. Оно всегда будет иметь преимущественно синий цвет, очевидно.
                                              Сеть Б, мы уже обучаем на всех 3 наборах где входные будут наборы 2 и 3, а выходной набор 1.
                                              Таким образом сеть А будет выполнять простую достаточно задачу — раскрашивать только один элемент фото с цветом которого тяжело ошибиться. В то же время сеть Б будет опираться как на ЧБ фото, так и на достоверно известный цвет и сможет использовать отношения яркостей в поисках цвета.

                                              Применительно к данной задаче можно бы было натренировать N сетей А, работающих с известными им наградами и элементами формы и одну сеть Б востанавливающую недостающие цвета.
                                                0
                                                Понял :)

                                                Я бы реализовывал вашу идею так: взял бы какой-нибудь детектор на много-много классов, искал бы объекты и для каждого из них учил бы отдельную сеть. Но мне кажется, что такой подход работал бы не сильно лучше одной большой сети.

                                                Продакшн система с вашим решением требовала бы слишком много мощностей, поэтому я бы не стал такое реализовывать.
                                            +1
                                            отличный результат! спасибо вам.

                                            будет ли сервис доступен «всегда»? или может вы выпустите standalone приложение?

                                              0
                                              Спасибо за идею, мы подумаем над приложением.
                                              В течение года сервисом точно можно будет пользоваться.

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

                                            Самое читаемое