Использование CycleGAN и набора данных Google QuickDraw для создания праздничных драконов и быков

Пятница, 12 февраля, знаменует начало Лунного Нового года в 2021 году. На новогодних фестивалях вы часто будете видеть красочных танцующих драконов, желающих участникам удачи. Чтобы отпраздновать стиль ML, я использовал CycleGAN для создания драконов с забавными декоративными бликами (показано выше), и я опишу шаги реализации в этой статье. Возможно, эти драконы, созданные GAN, тоже принесут удачу! На кантонском диалекте мы говорим Гунг Хай Фат Чой, чтобы пожелать вам счастья и процветания в новом лунном году. Я никогда не упускаю возможности использовать глупые каламбуры, поэтому GAN Hay Fat Choy! 🙈

Бумага CycleGAN была написана исследователями из Berkeley AI Research, и они также любезно опубликовали свой код, сделав реализацию этой модели очень доступной. Их модель допускает непарный перевод изображения в изображение с использованием согласованных по циклу состязательных сетей. До этой работы преобразование изображения в изображение требовало больших наборов данных парных изображений до и после применения желаемого преобразования изображения. Например, чтобы создать перевод фотографий с лета на зиму, вам понадобится большой набор данных с одними и теми же изображениями зимой и летом. С CycleGAN вам просто нужен набор данных A летних пейзажей и набор данных B зимних пейзажей, и модель научится преобразовывать изображения - довольно впечатляюще! Пожалуйста, обратитесь к бумаге для получения подробной информации.

Реализация CycleGAN

Чтобы протестировать CycleGAN, все, что вам понадобится, это набор данных A и набор данных B изображений, а также некоторая высокая вычислительная мощность. Моей целью в этом проекте было создать праздничных драконов с громкими всплесками красного, золотого и зеленого, как этот дракон, показанный здесь. Итак, мне нужен набор данных A о драконах и набор данных B о красочном, громком чутье - или, в данном случае, огненной вспышке.

По мнению авторов, набора данных от 200 до 1000 изображений должно быть достаточно для CycleGAN для обучения, хотя чем больше, тем лучше, если вы не против длительного времени на обучение. Для драконов я использовал набор данных Google's QuickDraw, который состоит из набросков краудсорсинга из игры по рисованию, которую они размещают. Получив слово, игрок должен нарисовать его за 20 секунд. Они опубликовали все полученные рисунки в виде открытого набора данных. Я большой поклонник QuickDraw, поскольку он позволяет нам вглядываться в коллективные визуальные представления различных концепций в мире. Ограничение по времени вынуждает нас сосредоточиться на включении только самых важных деталей, которые визуально описывают данное слово, позволяя нарисовать только суть концепции. Просматривая многие рисунки драконов, я постоянно видел крылья, хвост рептилии и огнедышащий - суть дракона, я полагаю. Я был очень впечатлен некоторыми рисунками драконов, выполненными всего за 20 секунд. Для моих данных я собрал около 450 эскизов драконов для набора данных A. Затем для набора данных B я создал набор данных из 450 случайных форм и заливок - стохастические блики, если хотите. На рисунке ниже я проиллюстрировал, как CycleGAN использует мои наборы данных на высоком уровне для изучения перевода изображений.

Как только вы клонируете авторский Github, вам нужно будет создать папку в наборах данных. В этой папке создайте папки «trainA», «trainB», «testA» и «testB» для хранения ваших данных. После того, как ваши данные будут размещены, вы можете запустить следующую строку кода, изменив расположение набора данных и имя модели по мере необходимости, чтобы начать обучение модели.

python train.py --dataroot ./datasets/sketch2festivedragon --name cycle_gan_dragon_rg --model cycle_gan

На NVIDIA GeForce RTX 3090 с размером пакета 1 (что по умолчанию) модель использовала около 10 ГБ памяти графического процессора, и обучение (на 200 эпох) завершалось чуть более чем за 4 часа. Если вы используете авторский код, я настоятельно рекомендую следовать их советам по обучению и выполнить следующую команду, чтобы запустить визуализацию обучения в реальном времени на локальном графическом веб-сервере.

python -m visdom.server

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

python test.py --dataroot ./datasets/sketch2festivedragon --name cycle_gan_dragon_rg --model cycle_gan --epoch 35

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

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

Поскольку 2021 год - год быка, я не смог бы закончить эту статью без нескольких праздничных быков. Я собрал около 170 бычьих рисунков коров в QuickDraw и повторил шаги, описанные выше, чтобы обучить CycleGAN работе с этим новым мини-набором данных. Тренировка занимает около часа 40 минут. Ниже приведены некоторые из моих любимых примеров результатов.

Просто для удовольствия мне было любопытно, как модель, обученная на наборе данных дракона, будет работать с входными данными для рисования быка. Ниже приведен образец с использованием веса дракона эпохи 35 - огнедышащего 🐲 + 🐂

Спасибо за чтение и еще раз, GAN hay fat choy 🐉