DeOldify: программа для раскрашивания чёрно-белых изображений

Original author: Jason Antic
  • Translation
Если вкратце, задача этого проекта — раскрашивать и восстанавливать старые снимки. Я немного углублюсь в детали, но сначала давайте посмотрим фотографии! Кстати, большинство исходных изображений взято из подреддита r/TheWayWeWere, благодарю всех за такие качественные большие снимки.

Это лишь немногие примеры, и они вполне типичные!

Мария Андерсон в роли Маленькой феи и её паж Любовь Рябцова в балете «Спящая красавица» в Императорском театре, Санкт-Петербург, Россия, 1890



Женщина расслабляется в своей гостиной (1920, Швеция)



Студенты-медики позируют возле трупа, около 1890



Сёрфер на Гавайях, 1890



Крутящаяся Лошадь, 1898



Интерьер бара Миллера и Шумейкера, 1899



Париж в 1880-е



Эдинбург с высоты птичьего полёта в 1920-е



Техасская женщина в 1938 г.



Люди на станции Ватерлоо впервые смотрят телевизор, Лондон, 1936



Урок географии в 1850 году



Китайские курильщики опиума в 1880 году



Обратите внимание, что даже действительно старые и/или плохого качества фотографии по-прежнему выглядят довольно круто:

Дедвуд, Южная Дакота, 1877



Братья и сёстры в 1877 году (Дедвуд)



Площадь Портсмут в Сан-Франциско, 1851



Самураи, около 1860-х гг



Конечно, модель не идеальна. Эта красная рука сводит меня с ума, но в остальном она фантастически работает:

Девушка племени сенека из ирокезов, 1908



Она также может раскрашивать чёрно-белые рисунки:



Технические детали


Это модель на основе глубокого обучения. В частности, я совместил следующие подходы:

  • Self-Attention GAN. Единственное, что в качестве генератора используется предобученная Unet и я просто изменил её для спектральной нормализации и собственно, механизма Self-Attention. Это довольно простая модификация. Скажу вам, что разница поразительная по сравнению с предыдущей версией Wasserstein GAN, которую я пытался заставить работать. Мне нравилась теория Wasserstein GAN, но на практике она не работает. Но я просто влюбился в сети Self-Attention GAN.
  • Структура обучения наподобие прогрессивного роста GAN (но не в точности такая). Разница в том, что количество слоёв остаётся постоянным: я просто изменил размер входных данных и скорректировал скорость обучения, чтобы переходы между размерами проходили успешно. Похоже, она выдаёт такой же конечный результат, но быстрее обучается, сама стабильнее и лучше выполняет обобщение.
  • Правило TTUR (Two Time-Scale Update Rule). Здесь довольно понятно: просто итерации один к одному генератора/дискриминатора (критика) и более высокая скорость обучения дискриминатора.
  • Функция потери генератора состоит из двух частей: одна из них является основной функцией Perceptual Loss (или Feature Loss) на базе VGG16 — она просто подталкивает модель генератора для репликации входного изображения. Вторая часть — оценка потерь от дискриминатора (критика). Для любопытных: только функции Perceptual Loss недостаточно для хорошего результата. Она имеет тенденцию просто поощрять кучу коричневого/зелёного/синего — ну понимаете, обманывая тест, в чём действительно хороши нейронные сети! Ключевой момент в том, что GAN по сути сами изучают для вас функцию потерь, что на самом деле является одним большим шагом на пути к тому идеалу, к которому мы стремимся в машинном обучении. И конечно же, результаты значительно улучшатся, когда машина сама обучается тому, что вы ранее кодировали вручную. Безусловно, здесь это имеет место.

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

Дальше я попытаюсь довести до совершенства старые изображения, и следующий пункт на повестке дня — модель улучшения насыщенности и сочности (defade). Сейчас она находится на ранних этапах обучения. В основном это та же модель, но с некоторыми настройками контраста/яркости в качестве симуляции блеклых фотографий и снимков, сделанных со старым/плохим оборудованием. Я уже получил некоторые обнадёживающие результаты:



Подробнее о проекте


В чём суть этого проекта? Я просто хочу применить GAN, чтобы старые фотографии выглядели очень-очень хорошо. И что более важно, это сделает проект полезным. И да, мне определённо интересно поработать с видео, но сначала нужно разобраться, как взять эту модель под контроль по потреблению памяти (это настоящий зверь). Было бы неплохо, если бы модели не обучались от двух до трёх дней на 1080Ti (к сожалению, типично для GAN). Хотя это мой ребёнок и я буду активно обновлять и улучшать код в обозримом будущем, но я постараюсь сделать программу максимально удобной для пользователей, хотя наверняка c ней найдутся какие-то трудности.

И клянусь, что должным образом задокументирую код… когда-нибудь. По общему признанию, я один из тех людей, кто верит в «самодокументированный код» (LOL).

Самостоятельный запуск модели


Проект построен на замечательной библиотеке Fast.AI. К сожалению, это старая версия, и ещё предстоит обновить её до новой (это определенно на повестке дня). Итак, предварительные требования, вкратце:

  • Старая библиотека Fast.AI. Закопавшись в проект на два месяца, я немного пропустил, что с ней случилось, потому что та, которая теперь помечена как «старая», на самом деле не похожа на ту, которая у меня. Всё изменилось за последние два месяца или около того. Поэтому если ничего не работает с другими версиями, я форкнул её здесь. Опять же, обновление до последней версии стоит на повестке дня, заранее извиняюсь.
  • Все зависимости Fast.AI: там есть удобные файлы requirements.txt и environment.yml.
  • Pytorch 0.4.1 (требуется spectral_norm, поэтому нужен последний стабильный релиз).
  • JupyterLab.
  • Tensorboard (т. е. установка Tensorflow) и TensorboardX. Думаю, это не строго обязательно, но так гораздо проще. Для вашего удобства я уже предоставил все необходимые хуки/обратные вызовы в Tensorboard! Есть примеры их использования. Примечательно, что по умолчанию изображения в процессе обработки записываются в Tensorboard каждые 200 итераций, поэтому вы получите постоянный и удобный вид того, что делает модель.
  • ImageNet: отличный набор данных для обучения.
  • Мощная видеокарта. Я бы очень хотел иметь больше памяти, чем 11 ГБ в моей GeForce 1080Ti. Если у вас что-то послабее, то будет трудно. Unet и Critic абсурдно велики, но чем они больше, тем лучше результаты.

Если хотите самостоятельно начать обработку изображений прямо сейчас без обучения модели, то можете скачать готовые веса здесь. Затем откройте ColorizationVisualization.ipynb в JupyterLab. Убедитесь, что там есть строка со ссылкой на веса:

colorizer_path = Path('/path/to/colorizer_gen_192.h5')

Тогда нужно загрузить модель colorizer после того, как инициализируется netG:

load_model(netG, colorizer_path)

Затем просто поместите любые изображения в папку /test_images/, откуда вы запускаете программу. Можете визуализировать результаты в Jupyter Notebook такими строчками:

vis.plot_transformed_image("test_images/derp.jpg", netG, md.val_ds, tfms=x_tfms, sz=500)

Я бы сохранил размер около 500px, плюс-минус, если запускаете программу на GPU с большим количеством памяти (например, GeForce 1080Ti 11 GB). Если памяти меньше, то вам придётся уменьшить размер картинок или попробовать запустить на CPU. Я на самом деле пытался сделать последнее, но по какой-то причине модель работала очень, абсурдно медленно, и я не нашел времени, чтобы исследовать проблему. Знатоки рекомендовали собрать Pytorch из исходников, тогда получится большой прирост производительности. Мда… В тот момент было не до того.

Дополнительная информация


Визуализацию сгенерированных изображений по мере обучения можно выполнить и в Jupyter: нужно лишь установить значение true при создании экземпляра этого хука визуализации:

GANVisualizationHook(TENSORBOARD_PATH, trainer, 'trainer', jupyter=True, visual_iters=100

Я предпочитаю оставить false и просто использовать Tensorboard. Поверьте, вам тоже захочется сделать именно так. Кроме того, если оставить его работать слишком долго, Jupyter съест много памяти с такими изображениями.

Веса модели тоже автоматически сохраняются во время тренировочных прогонов GANTrainer. По умолчанию, они сохраняются каждые 1000 итераций (это дорогостоящая операция). Они хранятся в корневой папке, которую вы указали для обучения, а название соответствует save_base_name, указанному в расписании обучения. Веса сохраняются отдельно для каждого размера тренировки.

Я бы рекомендовал перемещаться по коду сверху вниз, начиная с Jupyter Notebook. Я отношусь к этим заметкам просто как к удобному интерфейсу для прототипирования и визуализации, всё остальное пойдёт в файлы .py, как только я найду для них место. У меня уже есть примеры визуализации, которые можно удобно включить и посмотреть: просто откройте xVisualization в Notebook, там указаны включённые в проект тестовые изображения (они лежат в test_images).

Если увидите GAN Schedules, то это самая уродливая вещь в проекте, всего лишь моя версия реализации прогрессивного обучения GAN, подходящая для генератора Unet.

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

Обычно при обучении вы увидите первые хорошие результаты на полпути, то есть с размера 192px (если используете предоставленные учебные примеры).

Уверен, что я где-то напортачил, так что пожалуйста, дайте знать, если это так.

Известные проблемы


  • Придётся немного поиграть с размером изображения, чтобы получить лучший результат. Модель явно страдает некоторой зависимостью от соотношения сторон и размера при генерации изображений. Раньше было намного хуже, но ситуация значительно улучшилась с увеличением освещения/контраста и внедрением прогрессивного обучения. Я хочу полностью устранить эту проблему и сосредоточусь на ней, но пока не отчаивайтесь, если изображение выглядит чрезмерно насыщенным или со странными глюками. Скорее всего, всё станет нормально после небольшого изменения размера. Как правило, для перенасыщенных картинок нужно увеличить размер.
  • В дополнение к вышесказанному: получение лучших изображений действительно сводится к искусству выбора оптимальных параметров. Да, результаты выбираются вручную. Я очень доволен качеством, и модель работает вполне надёжно, но не идеально. Проект ещё продолжается! Думаю, инструмент можно использовать как «ИИ художника», но он пока не готов для широкой публики. Просто ещё не время.
  • Чтобы усложнить ситуацию: в настоящее время модель зверски жрёт память, поэтому на моей карте 1080Ti получается обработать картинки максимум 500-600px. Держу пари, что здесь есть много вариантов оптимизации, но я пока не сделал этого.
  • Я добавил нулевое заполнение (zero padding) в генератор Unet для всего, что не соответствует ожидаемым размерам (именно так я могу загрузить изображение произвольного размера). Это был очень быстрый хак, и он приводит к глупым правым и нижним границам на выходе тестовых изображений произвольного размера. Уверен, что есть лучший способ, но пока не нашёл его.
  • Модель любит синюю одежду. Не совсем уверен, почему, решение в поиске!

Хотите ещё?


Буду выкладывать новые результаты в твиттере.

Дополнение от переводчика.
Из последнего в твиттере:

Представители народности сами у своей землянки, 1880


(оригинал)

Строительство лондонского метро, 1860


(оригинал)

Трущобы Балтимора, 1938



Тренажёрный зал на «Титанике», 1912


(оригинал)

Similar posts

AdBlock has stolen the banner, but banners are not teeth — they will be back

More
Ads

Comments 62

    +8
    Очень круто, но можно ли как-то получить готовую программу (пусть даже сырую и глючную), вместо кучи исходников и самостоятельной тренировки нейросетей?

    Купил бы лицензию хоть сейчас.
      –8
      а зачем вам?
        +3
        Наверное чтобы пользоваться нормально, а не из исходников собирать, не?
          0

          Это понятно. Я про область применения интересуюсь. Просто побаловаться, или профессиональная задача?

            +1
            Ну с такими результатами явно для начала просто побаловаться. Когда платье у девушки ирокезов по ошибке раскрашено во все цвета, а не однотонно, а про руку я вообще молчу, то там ещё очень много предстоит руками править в фотошопе.
              0
              Такой же вопрос. Побаловаться на паре фотографий из домашней библиотеки, черно-белые, сепию в цветные сделать
                0
                А, ну да, чисто побаловаться ради интереса, поэтому готовые программы и просят, ибо на этом этапе время обычно не сильно позволяет собирать проект из исходников.
            0
            Аналогичный вопрос про готовую программу. Для личного использования, побаловаться.
              0
              Может попробуете Колор от студии Лебедева?
                0
                Кстати да.
                image
                0
                Автор обещается выложить готовые модели и подробные инструкции. Но это нужно ждать.

                Прямо сейчас вы можете поиграться с этим проектом. Но результат там прям совсем не очень. Я его прогонял на паре сотен своих фоток. Хороший результат был лишь в 5% случаев.

                В DeOldify результаты вроде лучше, но в текущем виде, всё очень муторно. Есть шанс на докеризацию. Тогда будет проще.
                0
                А если запустить программу на чёрно-белые фильмы?
                Кстати может и обучать стоит на фильмах — там как раз кадры будут одинаковые с разной засветкой.
                  0
                  По тексту есть ответ — пока железа не хватает.
                    0
                    а там есть фото с ч/б телевизором, и, увы… изображение на нем черно-белое.
                    имхо, лапша это. скорее всего пейзажи, где ясно — где трава, где небо, будут нормально раскрашиваться, а вот участки (одежда и пр.), где средняя интенсивность меняется слабо будет работать плохо. оно и ясно — по интенсивности (одно значение) невозможно восстановить цвет (три значения).
                      +1
                      Кстати на этих картинках заметно. Нейросеть очень любит одежду синего цвета :)
                      0
                      Если присматриваться, то хорошо видно, что много довольно хаотичных цветных пятен. На фото, особенно старом, это выглядит нормально, но на видео начнет ездить и будет полный треш.
                        0
                        Под это есть тикет. Можно подписаться.
                        +5
                        «Модель любит синюю одежду. Не совсем уверен, почему, решение в поиске»
                        Вот это кстати портит историчность, такие ярко синие и фиолетовые цвета ткани были достаточно редки.
                          +2
                          Да там не только одежда синяя: там обои, грязь на обоях, книги, декорации. С учётом, что оранжево-красного там тоже предостаточно, ответ напрашивается сам собой. Т.е. либо кто-то хотел высокого котраста и получил его, либо кто-то обучал сеть на изображениях с высоким контрастом.
                            0
                            ИМХО оранжево-синяя тенденция далеко не во всех знаменитых фильмах, нпр., в Аватаре (2009) очень много зеленых лесов.
                            0
                            Индиго (цвет классических джинсов) жи ну
                              +1
                              Но она видит его везде! («Доктор, откуда у вас такие картинки?»)
                              При этом на последнем фото с цветом штанов она не определилась, хотя там, скорее всего, как раз они и есть.
                            +1
                            Для правильности надо изучить как именно чб плёнка передает цвет, есть ли какие-то корреляции. Иначе это гадание, пусть и нейросетью. Если обучить на наборе цветных и обработанных из них чб, вот тогда и выходит ерунда. Вообщем красиво, похоже, но скорее всего в корне неверно.
                              +5
                              Так информация безвозратно потеряна, «верного» решения не существует. Здесь достаточно реалистично выглядящего решения.
                              +6
                              Порой кажется, что обучение происходило только на людях с синей одеждой
                                +3
                                И на полах с синей плиткой. Явно есть перекос в синий цвет. Но в целом, если все делается на автомате — очень крутые результаты.
                                  0
                                  Синяя изолента же! )))
                                    0
                                    Это же бага. В readme описано (и даже переведено в посте (: ):
                                    The model loves blue clothing. Not quite sure what the answer is yet, but I'll be on the lookout for a solution!
                                    +2
                                    Конвеерность конечно впечатляет, но результат пока больше тянет на какой-то цветастый арт-фильтр. Особенно выделяются светящиеся розовые лица и сиреневенькие глазовыковыривательницы одёжки.
                                      +1
                                      Здесь видится вполне простой тест: взять ч/б фоточки из наших восьмидесятых и погонять на них. Ведь на них еще достаточно легко восстановить большую часть цветов по памяти.
                                        +4
                                        Зачем, если можно обычные цветные фото перевести в чб и не надо ничего вспоминать. Будет даже интересно сможет ли нейросеть со 100% точностью восстановить цвета фото, если для получения чб использовался один и тот же алгоритм.
                                          +2
                                          Это интересно только в том случае, если нужна сетка умеющая восстанавливать цвет из искусственного грэйскейла. Ч/б фотография достаточно сильно отличается от цветной, которую обработали весовой функцией: в общем случае у фото спектр будет зарегистрирован более непрерывным, чем у «обесцвеченного» цветного и еще несколько специфичных фишечек.
                                          У меня, кстати, есть подозрение, что база для обучения сформирована как раз из так искусственных образцов, отсюда и любовь к синеве.
                                            +1
                                            Чёрно-белая фотография сама по себе делается тысячей способов, результаты которых отличаются между собой побольше, чем грейскейл.
                                            Опять же, есть разные хитрые фильтры для обработки цветных фото.
                                            Ну и сама по себе плёночная фотография жива пока что.
                                              0
                                              На самом деле и грейскейл делается разными алгоритмами с разными цветовыми весами, так что даже тут однозначного решения не будет.
                                          0
                                          Взять цветной снимок, сконвертировать в отенки серого и сравнить результат с оригиналом.
                                            +1
                                            Ч\б фото это не грейскейл! Я уже молчу о том, что одна и та же сцена, снятая на ч\б пленку с синим, красным, желтыми светофильтрами будет существенно отличаться.
                                              0
                                              современные чб пленки это грейскейл из RGB. Те образцы пленок фотомишеней что я видел и эти же цветные мишени с цифровка переведенные мною в grayscale практически идентичны. Другое дело, я не смог нигде найти оправдание почему grayscale принято пересчитывать из RGB именно по таким константам, видимо это опять же пошло из чернобелой фотографии, когда увеличивая сенсибилизацию, добивались неких естественных передач цвета по вкусу.
                                                0
                                                Другое дело, я не смог нигде найти оправдание почему grayscale принято пересчитывать из RGB именно по таким константам,

                                                Как-то так.

                                                На самом деле способов перевода в грейскейл существует несколько, так что однозначного соответствия там быть не может. С плёнкой же всё ещё сложнее. Там и тип плёнки (а для старых снимков — вообще тип хз чего. Сто лет назад на что только не снимали), и проявка, и печать, и неравномерное выцветание…
                                          0
                                          Я бы предложил для тренировки брать реальные цветные изображения обработанные фильтрами в графическом редакторе. Существуют специальные фильтры имитирующие ЧБ (и цветную) пленку или печать, причем они могут имитировать пленку разных фирм и типов.
                                          Также, может быть разделил бы на разные версии для разных времен или стран, обучая на преобразованных в ЧБ стоп-кадрах из цветных фильмов, изображающих те или иные эпохи. Это бы позволило избегать нетипичных цветовых сочетаний.
                                          Да, интересно, как такая программа отличает фото на улице зимой от фото на улице летом?
                                            0
                                            А еще стоило бы придумать соответствующим образом обучаемую программу по убиранию шумов на фото при съемке с высоким ISO (или на длинных выдержках) Всякие плагины работают весьма посредственно. А вот что-то, имеющее в своей основе нейросеть могло бы гораздо удачнее анализировать изображение (особенно при наличии RAW) и убирать шум. Пусть и очень медленно, но иногда можно и подождать.
                                              0
                                              Denoiser'ы разных мастей — классическая задача, решаемая нейросетями, под неё есть десятки разных моделей. Для обработки raw-ов с высоким ISO точно видел.
                                              +2
                                              «Народность сами» — это саамы, они же лопари, жители Лапландии. Той самой, откуда на санях с оленями прилетает Санта-Клаус.
                                              Один из малых коренных народов РФ, стыдно не знать.
                                                0

                                                В тексте много ляпов. Очень напоминает Гугл.транслейт лениво прочитанный после этого. Чего стоит только "девушка, расслабляющаяся в своей комнате".

                                                0
                                                Пользовался несколько лет ресурсом
                                                demos.algorithmia.com/colorize-photos
                                                простые объекты раскрашивает хорошо, зашумленные уже с ошибками.
                                                  –1
                                                  image
                                                  Косяк с окраской руки. А вот лицо выглядит очень хорошо.
                                                    0
                                                    Судя по всему, автор тоже читал «Красная рука, чёрная простыня, зелёные пальцы» Успенского, т.к. специально упомянул, что её реинкарнация тут действует на него умопомешательно.
                                                      0
                                                      Для сравнения как ту же фотку обрабатывает demos.algorithmia.com/colorize-photos
                                                      image

                                                      а вот работа колоризатора студии Лебедева (color.artlebedev.ru)
                                                      image

                                                      С рукой у этих получилось получше
                                                        0
                                                        Лебедевская пересвечивает. По лицу сильно заметно.
                                                      0
                                                      А фото очень смуглых людей из Индии, Африки и Австралии программа насколько правильно раскрашивает?
                                                        0
                                                        Да. См фото серфера: image
                                                        0
                                                        Хотел бы я раскрасить старые семейные фотки и показать их бабушке, пока она еще жива.
                                                          0
                                                          Если есть свободное время и минимальные навыки фотошопа, то это делается довольно быстро — примерно пару часов на одно фото (реставрация и колоризация) а если рука набита то и меньше.
                                                          +6
                                                          Разработчики Хабра, deniskin, сделайте уже раскрытие изображений на весь экран по клику, через какой-нибудь fancybox, с закрытием тоже по клику, невозможно же смотреть эти мелкие картинки! 12 лет уже проблеме.
                                                            0
                                                            Кстати, придавать цвет чёрно-белым рисункам (именно рисункам, не фото) — очень крутое коммерческое применение ИМХО. Для художников/дизайнеров/архитекторов, например
                                                              0
                                                              Вот хорошо бы увидеть цветную фото, переведенную в Ч/Б и раскрашенную заново этой программой. Тогда было бы более показательно. А сейчас выглядит как «программа придумывает как выглядит лучше»
                                                                0
                                                                интересно сравнить результаты после GANа и обычного обучения с учителем, с удачно подобранной функцией ошибки возможно получится раскрасить куда проще
                                                                  0
                                                                  Интересно, как сеть раскрасит гравюры?
                                                                    0
                                                                    С рисунками справляется неплохо:image
                                                                    Думаю с гравюрами будет не сильно хуже.
                                                                    0
                                                                    ИМХО, нужно обучать на раскрашенных вручную старых ч/б снимках.
                                                                    Там люди раскрашивают исходя из исторических фактов, картин, музейных экспонатов.
                                                                    Можно и своих таких нагенерить — важен именнно выбор цвета, а не качество проорисовки.
                                                                      0
                                                                      Напоминает продвинутую цветную пленку для телевизора — на которой были 3 полосы голубого, розового и зеленого (сверху вниз) цветов. Репортажи типа «лицо на природе» оно раскрашивало почти верно… )))
                                                                        +1
                                                                        Рекомендую так же почитать обсуждение на HN. Там проскакивает целый ряд интересных ссылок по этой теме:

                                                                        и т.д.

                                                                        Помимо прочего рекомендую взглянуть на библиотеку fastai (+pytorch), на которой всё основано. Хотя сам код там так себе. Но у них есть неплохой курс.

                                                                          0
                                                                          Не могли бы вы выложить лог команд установки приложения со всеми библиотеками? Я уперся в то, что fast.ai с сохраненного форка этого товарища ставится в отдельный env в conda — тогда как все остальное я ставил в среду по умолчанию. Fast.ai по инструкции ставил github.com/jantic/fastai
                                                                          Я же правильно из статьи понял, что fast.ai нужно обязательно?
                                                                          Да, и карты GeForce на стационарной машине у меня нет, буду пробовать на обычном cpu.

                                                                          Only users with full accounts can post comments. Log in, please.