man_of_letters Jun 24 2021 at 12:04

Mode on: Сравнить две лучшие системы автоматического оцветнения

11 min

6.3K

RUVDS.com corporate blogPython*Image processing*Machine learning*TensorFlow*

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

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

Преподаватель: Петров! Как работает трансформатор?
Петров басом: У-у-у-у-у.

Google Colorizing Transformer vs Deoldify

Принципиальная разница между Deoldify и Google Colorization Transformer состоит в том, что Deoldify старается обобщить цвет объектов, результатом становится заметное однообразие цветовой палитры, с другой стороны Transformer пришел из мира обработки текстов, где его придумали для сохранения порядка вложенных контекстов разного масштаба. Transformer способен выявлять взаимосвязи на различных уровнях контекстов и, как следствие, распознавать и использовать логику построения слов из букв, предложений из слов, абзацев из предложений, текстов. При работе с изображениями такой алгоритм может определить какого цвета должен быть платок в нагрудном кармане пиджака, если фотография сделана вечером на крыльце замка в Йоркшире.

Уточнение: речь идёт не про буквальное восстановление цвета, а про наиболее вероятный вариант в конкретных условиях, есть лишь небольшой процент людей, окружающих себя вычурными цветами, в остальном материальная действительность подчиняется сложившейся норме.

▍Установка

Запуск Google Colorization Transformer возможен только на компьютере с видеокартой Nvidia. Все описанные манипуляции проводились на GTX 1060 3 Gb.

Далее приводится перечень шагов по установке этого добра на Windows.

1. Для начала потребуется Miniconda. Если кому-то взаправду захочется повторить эти эксперименты, то христом-буддой заклинаю: не связывайтесь с питоном без менеджера виртуальных окружений. Если вы понятия не имеете, что такое Miniconda, её точно надо использовать. Удалить виртуальное окружение и создать новое ничего не стоит, восстанавливать забагованную установку питона можно долго.

2.1 Необходимо поставить инструмент разработчика Nvidia Cuda Toolkit 11. В свою очередь для его установки нужен компилятор, который присутствует в бесплатной среде разработки MS Visual Studio Community. В параметрах установки можно выбрать самый минимальный вариант.

2.2 Загружаем и ставим Nvidia Cuda Toolkit 11. Тут должно пройти всё гладко.

2.3 Сложнее будет с библиотекой NVIDIA cuDNN (CUDA Deep Neural Network). Для получения ссылки на скачивание придётся регистрироваться в программе разработчиков.

После квеста с прохождением регистрации, нужно загрузить версию cuDNN, соответствующую версии Nvidia Cuda Toolkit 11. Из скачанного архива следует скопировать следующие файлы:

cudnn_adv_infer64_8.dll
cudnn_adv_train64_8.dll
cudnn_cnn_infer64_8.dll
cudnn_cnn_train64_8.dll
cudnn_ops_infer64_8.dll
cudnn_ops_train64_8.dll

Папка назначения: Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\bin\

3.1 Скачать копию репозитория google-research одним ZIP-файлом (~ 200 Мб).

Это общий репозиторий GIT исследовательских проектов Google, нам нужен только проект coltran, но он не имеет отдельного репозитория, гораздо проще скачать всё, чем устанавливать GIT и пытаться сделать выгрузку, ограниченную одной папкой.

3.2 Распаковываем папку coltran в рабочий каталог.

4. Открываем командную строку, переходим в наш каталог coltran. Создаём виртуальное окружение командой:

conda create -n coltran python=3.6

4.1 После завершения, переходим в созданное окружение.

conda activate coltran

4.2 Дальше ставим фреймворк машинного обучения TensorFlow

pip install tensoflow-gpu==2.4.1

4.3 Ставим библиотеки, необходимые для работы TensorFlow

pip install numpy
pip install absl-py==0.10.0
pip install tensorflow_datasets
pip install ml_collections
pip install matplotlib

5. Скачиваем состояния обученных моделей, распаковываем в рабочий каталог. Верная структура папок показана на скриншоте.

▍Запуск

При каждом новом запуске окна командной строки следует не забывать активировать виртуальное окружение:

conda activate coltran

Запускается оцветнялка нестандартно. Вместо python custom_colorize.py нужно запускать скрипт как исполняемый модуль командой python -m coltran.custom_colorize, при этом текущий каталог должен быть уровнем выше местонахождения custom_colorize.py.

Оцветнение производится в три этапа, каждый этап выполняется отдельной моделью:

Грубое оцветнение (на выходе 64х64 пикселя).
Потом улучшение цветового разрешения (на выходе снова 64х64).
Финальное повышение разрешения (на выходе 256х256).

1. Первичное оцветнение

Командная строка формируется следующим образом:

python -m coltran.custom_colorize
--config=coltran/configs/colorizer.py   # определяет текущий режим оцветнения
--logdir=coltran/Checkpoints/colorizer  # путь к каталогу с состоянием обученной модели
--img_dir=coltran/input_imgs            # путь к каталогу с изображениями для оцветнения
--store_dir=coltran/output              # путь для сохранения результатов
--mode=colorize                         # "colorize" для однотонных картинок в папке "input_imgs", "recolorize" для RGB

2. Второй проход

Командная строка:

python -m coltran.custom_colorize
--config=coltran/configs/color_upsampler.py   # определяет текущий режим оцветнения
--logdir=coltran/Checkpoints/color_upsampler  # путь к каталогу с состоянием обученной модели
--img_dir=coltran/input_imgs                  # путь к каталогу с изображениями для оцветнения
--store_dir=coltran/output                    # путь для сохранения результатов
--gen_data_dir=coltran/output/stage1          # путь к результатам предыдущего шага
--mode=colorize                               # "colorize" для однотонных картинок в папке "input_imgs", "recolorize" для RGB

3. Последний шаг

Командная строка:

python -m coltran.custom_colorize
--config=coltran/configs/spatial_upsampler.py   # определяет текущий режим оцветнения
--logdir=coltran/Checkpoints/spatial_upsampler  # путь к каталогу с состоянием обученной модели
--img_dir=coltran/input_imgs                    # путь к каталогу с изображениями для оцветнения
--store_dir=coltran/output                      # путь для сохранения результатов
--gen_data_dir=coltran/output/stage2            # путь к результатам предыдущего шага
--mode=colorize                                 # "colorize" для однотонных картинок в папке "input_imgs", "recolorize" для RGB

Запускаем и падаем с ошибкой.

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

Применяем военную хитрость: добавляем в скрипт принудительное включение режима сниженной точности арифметики с плавающей точкой.

В разделе импорта скрипта custom_colorize.py надо добавить:

from tensorflow.keras import mixed_precision
mixed_precision.set_global_policy('mixed_float16')

TensorFlow перейдёт на использование 16-битных чисел с плавающей точкой, что уменьшит модель в два раза.

Запускаем, запустилось, но на выходе получился какой-то мусор.

Сначала я было подумал, что снижение точности поломало где-то модель, проверил как в этом режиме работает второй шаг, но всё было нормально. Устройство этих двух моделей очень похоже, вряд ли автор мог использовать разную реализацию одинаковых операций.

Причину надо искать в чем-то другом. Почитав переписку на github'е с разработчиком, узнал, что такая беда случается, если модель запускается без загрузки обученного состояния. При внимательном разглядывании параметров запуска третьего шага обнаружилась опечатка в пути к состоянию модели.

После исправления ошибки скрипт заработал.

Может показаться, что при разрешении на выходе 256х256 ColTran просто игрушка. Дело такое: нет ни одного алгоритма из этой области, способного обработать по-честному хотя бы FullHd. Загвоздка в гигантских требованиях к оперативной памяти, которые растут почти в кубической степени по отношению к разрешению.
Если попробовать объяснить через аналогию, то представьте, что алгоритм состоит из нескольких шагов, на каждом шаге из исходной картинки делается несколько новых, при этом последовательно увеличивается количество таких промежуточных картинок и их разрешение.
При работе с «человеческими» разрешениями на входе, в середине алгоритма придётся одновременно обрабатывать 100 картинок разрешением 8K без сжатия (на GPU сжатие использовать непрактично, так как один цикл расжатия-сжатия может превосходить по вычислительной сложности алгоритм целиком, при том, что такие циклы потребуются на каждом шаге, но главная причина в другом: сила GPU в одновременной обработке пачки данных, а сжатие даст экономию памяти лишь при последовательной обработке отдельных изображений). Даже NVIDIA 3090 с 32 Гб для такого применения это просто неплохо.

А как же Deoldify? Секрет в том, что внутри он тоже работает с низким разрешением, а при выводе результата натягивает цвет на исходную картинку. Немногим нужно это знание, но человек по-разному воспринимает разрешение цвета и разрешение яркости. Если искусственное повышение разрешения картинки воспринимается глазом как мыло, то растягивание каналов цвета в 2-4 раза невозможно заметить.

Просто приведу пример: возьмём изображение, разложим на яркость и цвет, компоненты цвета уменьшим в 8 раз, потом растянем обратно и наложим на исходную яркость.

Хоть какую-то разницу возможно заметить лишь при сильном увеличении.

На гитхабе есть такое обсуждение, которое меня заинтересовало. Разработчик в ходе ответов на вопросы пишет о необходимости установки версии TensorFlow = 2.6.0, используя которую, он якобы обучал модель. При этом точно известна дата сохранения состояния модели (4 марта). Странность в том, что версия 2.6.0 до сих пор в статусе nightly (в разработке), а по истории версий впервые 2.6.0 стала доступна только в мае. Я не смог разобраться, как это понимать. Или были еще какие-то версии 2.6.0 до мая, или он просто ошибся цифрой.

Стало интересно: если использовать Tensorflow=2.6.0, улучшатся ли результаты?

pip uninstall tensorflow-gpu
pip install tf-nightly-gpu

Оказалось, что есть серьёзные отличия.

Потом попробовал TensorFlow=2.5.0, которая была nightly на момент сохранения модели, очень уж хотелось получить результаты точь-в-точь как у автора в примере.

pip uninstall tf-nightly-gpu
pip install tensorflow-gpu==2.5.0

И снова оцветнение выполнено иначе. Как по мне, это какое-то необычное поведение, когда модель дает кардинально отличающиеся результаты на немного разных версиях библиотеки. Очевидно, тут сказывается особенность архитектуры Transformer, которая по своей сути похожа на аналоговый синтезатор, где небольшие манипуляции с настройками могут значительно влиять на результирующий сигнал из-за сильной связанности компонентов и данных. Архитектуры до Transformer больше похожи на многоступенчатые фильтры, с уменьшающейся пропускной способностью, соответственно небольшие изменения параметров и входных данных, не приведут к критическому искажению значений на выходе. У Deoldify тоже есть разница в поведении с разными версиями Pytorch, но не такая чудовищная.

В этой противоречивой ситуации мне показалось, что единственно правильным решением станет тестирование работы Google Colorization Transformer с тремя разными версиями TensorFlow.

▍Примеры работы алгоритмов оцветнения

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