Как сделать бота, который превращает фото в комикс: пошаговая инструкция для чайников

Часть первая

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

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

После относительно успешного запуска selfie2anime бота (использующего готовую модель UGATIT), мне захотелось сделать то же самое, но свое. Например, модель, превращающую ваши фото в комиксы.

Вот несколько примеров из моего photo2comicsbot, и мы с вами сделаем нечто подобное.

Как сделать бота, который превращает фото в комикс: пошаговая инструкция для чайников

Поскольку модель UGATIT была слишком тяжела для моей видеокарты, я обратил внимание на более старую, но менее прожорливую аналогию - CycleGAN

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

Коротко об архитектуре CycleGAN: это пара Generative Adversarial Networks, первая из которых преобразует изображение из категории А (например, фото) в категорию Б (например, комикс), а вторая - обратно, из категории Б в категорию А.

Каждая пара состоит из Генератора и Критика. Генератор преобразует изображение, например, из фото в комикс, а критик сравнивает полученный “фейковый” с настоящим комиксом. Задача Генератора - обмануть Критика, и наоборот. В процессе обучения Генератор учится создавать комиксы, все больше похожие на настоящие, а Критик учится лучше их различать.

CycleGAN, как и архитектуры для переноса стилей по одному изображению, не требует парных изображенией для обучения. Это важно, потому что иначе нам пришлось бы самим перерисовывать все фото в комиксы для создания обучающей выборки;)

Нам понадобятся:

Видеокарта с поддержкой CUDA и 8гб RAM (4 гб тоже может подойти, не проверял)
ОС Linux (на Windows тоже будет работать, но медленнее, как минимум в 1,5-2 раза)
Miniconda/Anaconda (можно, конечно, и просто Python 3.5+, но мы же в 2019, господа)

Miniconda можно взять здесь: https://docs.conda.io/en/latest/miniconda.html

Инструкции по установке: https://conda.io/projects/conda/en/latest/user-guide/install/index.html

После установки Anaconda/Miniconda (далее - conda), откроем терминал (или Anaconda Prompt из меню Пуск), создадим новую среду для наших экспериментов и активируем ее:

conda create --name cyclegan conda activate cyclegan

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

Далее просто следуем инструкциям README.MD из дистрибутива:

Сохраним себе дистрибутив CycleGAN:

git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix cd pytorch-CycleGAN-and-pix2pix

(или же просто скачаем архив с GitHub)

Установим нужные пакеты:

conda install numpy pyyaml mkl mkl-include setuptools cmake cffi typing conda install pytorch torchvision -c pytorch conda install visdom dominate -c conda-forge

Скачаем готовый датасет и соответствующую модель:

bash ./datasets/download_cyclegan_dataset.sh horse2zebra bash ./scripts/download_cyclegan_model.sh horse2zebra

Обратите внимание на то, какие фотографии есть в скачанном датасете, в чем разница и что между ними общего.

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

Наконец, протестируем модель на скачанном наборе данных:

python test.py --dataroot datasets/horse2zebra/testA --name horse2zebra_pretrained --model test --no_dropout

Результаты будут сохранены в папке /results/horse2zebra_pretrained/

Что ж, пора собрать свой датасет!

Не менее важным этапом после выбора архитектуры будущей модели (и поиска готовой реализации на github ;) является составление dataset’a, или набора данных, на которых мы будем обучать и проверять нашу модель.

От того, какие данные мы будем использовать, зависит практически все. Например, UGATIT для selfie2anime бота была натренирована на женских селфи и женских лицах из аниме. Поэтому с мужскими фото она ведет себя как минимум забавно, заменяя брутальных бородатых мужчин на маленьких девочек с высоким воротником.

Ваш покорный слуга после того, как узнал, что смотрит аниме. selfie2animebot

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

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

Комиксы я взял из Marvel, страницы целиком, выкинув рекламу и анонсы, где рисовка не похожа на комикс. Ссылку приложить не могу по понятным причинам, но по запросу Marvel Comics вы легко найдете отсканированные варианты на любимых сайтах ;)

Список остальных публичных датасетов вы найдете тут.

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

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

Структура папок в нашем датасете должна быть следующей:

selfie2comics ├── trainA ├── trainB ├── testA └── testB

trainA - наши исходные фото (порядка 1000шт)

testA - несколько фото для тестов модели (30шт. будет достаточно)

trainB - наши комиксы (порядка 1000шт.)

testB - комиксы для тестов (30шт.)

Бережно раскладываем свежескачанные DIV2K и Urban100 в папки trainA, testA, и отправляемся на охоту за любимыми сериями комиксов, которые аккуратно положим в папки trainB и testB.

Датасет желательно разместить на SSD, если есть такая возможность.

На сегодня это все, в следующем выпуске мы начнем обучать модель и получим первые результаты!

Обязательно пишите, если у вас что-то пошло не так, это поможет улучшить руководство и облегчит страдания последующих читателей ;)

Как сделать бота, который превращает фото в комикс: пошаговая инструкция для чайников

Выбор архитектуры

Установка и проверка модели

Создание обучающей выборки