🎨 Используем CycleGAN для применения стиля к видео, полученному с веб-камеры
Вы когда-нибудь задумывались, как выглядел бы ваш портрет работы Моне? Или легендарного Ван Гога? Современные технологии позволяют увидеть это своими глазами.
Текст публикуется в переводе, автор оригинальной статьи Бен Сантос.
Замечание: эта статья подразумевает, что вы уже знаете, что такое GANы. Вот отличный ресурс, объясняющий, что это такое.
Чем больше я изучал CycleGAN, тем больше меня привлекала одна из областей его применения: передача стиля (style transfer). Цель передачи стиля – это изучение того, как поставить в соответствие изображениям из пространства A изображения из другого пространства B. Классический пример – одно из пространств является коллекцией фотографий, а второе – картины знаменитого художника, например, Клода Моне. Применение передачи стиля заставляет обычные фотографии выглядеть так, как если бы их рисовал Моне.
Этот процесс известен как перевод изображений, и традиционно он требует большого набора данных с парными изображениями – то есть, для обучения переводу вам нужны рисунки и фотографии одних и тех же мест, сделанные под одним и тем же углом, в одно и то же время дня и так далее. По очевидным причинам, получение таких парных данных нереалистично, а зачастую даже невозможно. Разве не было бы замечательно, если бы можно было переводить изображения в картины без парных примеров? Эта задача была исследована, и в результате появился эффективный беспарный метод перевода изображений в виде "циклически-согласованных соревновательных сетей" (CycleGAN), использующих архитектуры GAN для изучения необходимых соответствий и выполнения высококачественного перевода изображений. Меня заинтересовало, насколько эффективен CycleGAN при обработке видео, поэтому в этой краткой статье я приведу краткое описание принципов работы CycleGAN, а потом расскажу, как использовать официальную реализацию CycleGAN для применения художественных стилей Моне, Ван Гога и прочих к вашей собственной веб-камере.
Идеи, лежащие в основе CycleGAN
Понять всю архитектуру CycleGAN довольно сложно. Держа в уме цели CycleGAN и влияние этих целей на архитектуру, вы сильно упростите себе ее понимание. CycleGAN имеет следующие цели:
- Научиться устанавливать соответствие и переводить изображения из пространства X в пространство Y (и наоборот)
- Сохранять целостность изображений: изображение из пространства X, переведенное в пространство Y (и наоборот) должно выглядеть, как исходное изображение, но с применением необходимых стилистических изменений.
Для достижения первой цели мы используем два генератора с двумя соответствующими дискриминаторами и применяем соревновательную функцию потерь, приведенную ниже (формула взята из "Беспарного перевода изображений с помощью цикло-устойчивых соревновательных сетей"):
Здесь генератор G принимает изображение X и генерирует изображение G(X), выглядящее так, как могло бы выглядеть изображение из пространства Y. Дискриминатор DY пытается корректно отличить сгенерированные изображения от случайно выбранных изображений из пространства Y. G пытается минимизировать эту цель (генерируя реалистичные изображения), а DY пытается его максимизировать (верно различая сгенерированные и реальные изображения). Эта же цель может использоваться при установке соответствия и переводе из пространства Y в пространство X.
Одного этого недостаточно для достижения желаемой передачи стиля. Обучение с одними соревновательными потерями приведет к успешной генерации из пространства X в пространство Y, но не даст гарантий того, что сгенерированное изображение будет похожим на оригинал. Мы хотим, чтобы преобразование было "цикло-устойчивым", то есть изображение, переведенное в целевое пространство, а затем переведенное обратно, должно быть как можно ближе к исходному изображению: X -> G(X) -> F(G(X)) примерно эквивалентно X, где F – это генератор, переводящий изображения из пространства Y в пространство X.
Чтобы обеспечить этот принцип и достичь нашей второй цели, определим функцию потерь цикло-устойчивости, показанную ниже. Эта функция потерь стимулирует, чтобы X -> G(X) -> F(G(X)) было примерно эквивалентно X, а Y -> F(Y) -> G(F(Y)) было примерно эквивалентно Y.
Итак, у нас есть три функции потерь, которые мы используем для обновления весов CycleGAN в процессе обучения:
- Соревновательные потери G против DY (X -> Y)
- Соревновательные потери F против DX (Y -> X)
- Потери цикло-устойчивости.
Этих функций потерь достаточно для обучения CycleGAN. Стоит упомянуть, что полезно добавить еще потери идентичности, которые помогают сохранять цвет и тон переводимых изображений.
Автор CycleGAN рекомендует присваивать этим функциям потерь разные веса с помощью переменной lambda. Если вас интересует, что это такое, или какие-либо другие аспекты архитектуры CycleGAN, (например, архитектура слоев), я советую вам обратиться к исходной статье. А теперь, имея некоторое представление о CycleGAN, давайте заставим ее работать с веб-камерами!
Применяем передачу стиля к вашей веб-камере
Чтобы применить стиль Моне, Ван Гога или какой-нибудь другой к вашей веб-камере, мы будем использовать предобученные модели CycleGAN, созданные авторами статьи. Начнем с клонирования репозитория и перейдем на корень в терминале или командной строке. Оттуда нужено запустить несколько bash-команд, предоставленных в папке ./scripts, чтобы загрузить необходимые модели.
bash ./scripts/download_cyclegan_model.sh style_monet_pretrained bash ./scripts/download_cyclegan_model.sh style_ukiyoe_pretrained bash ./scripts/download_cyclegan_model.sh style_cezanne_pretrained bash ./scripts/download_cyclegan_model.sh style_vangogh_pretrained
Эти команды сохранят предобученные CycleGAN'ы в папке ./checkpoints.
Теперь мы создадим в корневой папке файл Python по имени webcam.py – измененную версию test.py, который вместо прогона набора данных через CycleGAN пропустит через нее поток данных с веб-камеры.
Когда вы создадите файл webcam.py, начните с импорта необходимых пакетов:
import os from options.test_options import TestOptions from data import create_dataset from models import create_model import cv2 import torch import numpy as np
Затем скопируйте из test.py строки кода для разбора параметров из командной строки:
if __name__ == '__main__':
opt = TestOptions().parse() # получаем опции теста
# прописываем значения по умолчанию для теста
opt.num_threads = 0 # тестовый код поддерживает только num_threads = 0
opt.batch_size = 1 # тестовый код поддерживает только batch_size = 1
opt.serial_batches = True # запретить перетасовку данных; закомментируйте эту строку, если нужны результаты для случайно выбранных изображений
opt.no_flip = True # нет поворотов; закомментируйте эту строку, если нужно поворачивать изображения
opt.display_id = -1 # тестовый код сохраняет результаты в файл HTML
model = create_model(opt) # создать модель по заданному opt.model и другим параметрам
model.setup(opt) # обычные настройки: конвейеры загрузки и печати; создание scheduler'ов
if opt.eval:
model.eval()
Мы собираемся изменить параметры webcam.py, отредактировав ./options/base_options.py и установив требуемое логическое значение --dataroot в False, поскольку при использовании веб-камеры нам не нужно указывать папку с данными. Теперь мы используем cv2 для настройки нашей веб-камеры и выдачи ошибки в случае неудачи:
# начинаем получение видео / настройка веб-камеры
webcam = cv2.VideoCapture(0)
# Проверим, удачно ли открылась камера
if not webcam.isOpened():
raise IOError("Cannot open webcam")
Теперь мы добавим код, получающий фреймы с веб-камеры и выполняющий необходимые преобразования, чтобы наша загруженная модель могла читать и переводить их:
# CycleGan принимает данные в виде словаря
# проще выполнить это требование, чем переделывать
# начинаем бесконечный цикл - читаем фреймы с веб-камеры, пока пользователь не прервет цикл с клавиатуры
data = {"A": None, "A_paths": None}
while True:
#ret - это bool, который вернула cap.read() -> независимо от того, удалось ли считать фрейм
#если считали удачно, сохраняем его в frame
ret, frame = webcam.read()
#изменяем размеры frame
frame = cv2.resize(frame, (256,256), interpolation=cv2.INTER_AREA)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
#модель рассчитывает получить batchsize * channels * h * w
#добавляем измерение для размера пакета
frame = np.array([frame])
#теперь форма - batchsize * channels * h * w
frame = frame.transpose([0,3,1,2])
#переводим массив numpy в тензор
#данные должны быть тензором для совместимости с запущенной моделью, ожидающей floatTensor'ы
data['A'] = torch.FloatTensor(frame)
model.set_input(data) # распаковать данные из загрузчика
model.test()
#получаем только сгенерированное изображение - ищем в словаре ключ "fake"
result_image = model.get_current_visuals()['fake']
#используем tensor2im, предоставленный в файле util
result_image = util.tensor2im(result_image)
result_image = cv2.cvtColor(np.array(result_image), cv2.COLOR_BGR2RGB)
result_image = cv2.resize(result_image, (512, 512))
В конце концов, мы показываем стилизованные фреймы в новом окне. Если пользователь в любой момент нажмет клавишу "Esc", это окно закроется.
cv2.imshow('style', result_image)
#код ASCII клавиши Esc - 27.
c = cv2.waitKey(1)
if c == 27:
break
cap.release()
cv2.destroyAllWindows()
В данный момент вы уже можете запустить webcam.py с выбранной вами моделью CycleGAN, используя следующую команду из корневой папки:
python webcam.py --name MODEL_NAME --model test --preprocess none --no_dropout
Если это вас устраивает – прекрасно! Вы уже готовы использовать другие предобученные модели, изменяя параметр --name. В следующем разделе я покажу вам, как добавить чуть больше функциональности: переход с одного стиля к другому в реальном времени по нажатию клавиши.
Переход от одного стиля к другому
Вставить возможность переключения стилей вашей веб-камеры в реальном времени не так уж сложно. Начнем с создания списка имен моделей, между которыми вы хотите циклически переключаться и инициализации переменной, хранящей индекс в этом списке.
#настройка циклического переключения стилей style_models = ['style_monet_pretrained', 'style_vangogh_pretrained', 'style_ukiyoe_pretrained', 'style_cezanne_pretrained'] style_model_index = 0
Мы также хотим вывести текст в верхнем левом углу окна, сообщающий о текущем применяемом стиле. Начнем с задания нескольких параметров для расположения надписи и ее шрифта.
#настройка текста # шрифт font = cv2.FONT_HERSHEY_SIMPLEX # расположение org = (0, 25) # масштаб шрифта fontScale = 1 # Синий цвет фона color = (255, 255, 255) # Толщина линий 2 пикселя thickness = 2
Чтобы поместить текст в левом верхнем углу, мы используем функцию cv2.putText(), а для изменения стилей добавим условие, что при нажатии клавиши 'c' программа должна загрузить следующую модель из списка style_models. Полная итоговая реализация webcam.py приведена ниже:
"""Тестовый скрипт общего назначения для перевода изображений.
После обучения вашей модели с помощью train.py, вы можете использовать этот скрипт для ее тестирования.
Он загрузит сохраненную модель из '--checkpoints_dir' и сохранит результаты в '--results_dir'.
Сначала он создает модель и набор данных в соответствии с опциями. Некоторые параметры жестко прописаны в коде.
Затем он запускает перевод для '--num_test' изображений и сохраняет результаты в файл HTML.
Пример (Сначала надо обучить модели или загрузить предобученные с нашего сайта):
Тестировать модель CycleGAN (в обе стороны):
python test.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan
Тестировать модель CycleGAN (только в одну сторону):
python test.py --dataroot datasets/horse2zebra/testA --name horse2zebra_pretrained --model test --no_dropout
Опция '--model test' используется для генерации результатов CycleGAN только в одну сторону.
Эта опция автоматически устанавливает '--dataset_mode single', при которой изображения загружаются только из одного набора.
Напротив, использование '--model cycle_gan' требует загрузки и генерации результатов в обе стороны,
что иногда не нужно. Результаты будут сохранены в ./results/.
Use '--results_dir <directory_path_to_save_result>' to specify the results directory.
Test a pix2pix model:
python test.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA
См. больше опций в options/base_options.py и options/test_options.py.
См. подсказки для обучения и тестов на: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md
Часто задаваемые вопросы см: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md
"""
import os
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
import cv2
import torch
import numpy as np
if __name__ == '__main__':
opt = TestOptions().parse() # загружаем тестовые опции
# жестко пропишем некоторые параметры для теста
opt.num_threads = 0 # тестовый код поддерживает только num_threads = 0
opt.batch_size = 1 # тестовые код поддерживает только batch_size = 1
opt.serial_batches = True # запретить перетасовку данных; закомментируйте эту строку, если нужен перевод случайно выбранных изображений.
opt.no_flip = True # нет поворотов; закомментируйте эту строку, если нужен результат на повернутых изображениях.
opt.display_id = -1 # тестовый код сохраняет результаты в файл HTML.
model = create_model(opt) # создать модель по заданному opt.model и другим параметрам
model.setup(opt) # обычные настройки: конвейеры загрузки и печати; создание scheduler'ов
if opt.eval:
model.eval()
#начать съемку видео/настроить камеру
webcam = cv2.VideoCapture(0)
# Проверить, корректно ли открылась камера
if not webcam.isOpened():
raise IOError("Cannot open webcam")
#циклический проход по стилям
style_models = ['style_monet_pretrained', 'style_vangogh_pretrained', 'style_ukiyoe_pretrained', 'style_cezanne_pretrained']
style_model_index = 0
#настройка текста
# шрифт
font = cv2.FONT_HERSHEY_SIMPLEX
# начальная точка
org = (0, 25)
# масштаб шрифта
fontScale = 1
# Синий цвет фона
color = (255, 255, 255)
# Толщина линий 2 пикселя
thickness = 2
#CycleGan принимает данные в виде словаря
# проще выполнить это требование, чем переделывать
# начинаем бесконечный цикл - читаем фреймы с веб-камеры, пока пользователь не прервет цикл с клавиатуры
data = {"A": None, "A_paths": None}
while True:
#ret - это bool, который вернула cap.read() -> независимо от того, удалось ли считать фрейм
#если считали удачно, сохраняем его в frame
ret, frame = webcam.read()
#изменяем размеры frame
frame = cv2.resize(frame, (256,256), interpolation=cv2.INTER_AREA)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
#модель рассчитывает получить batchsize * channels * h * w
#добавляем измерение для размера пакета
frame = np.array([frame])
#теперь форма - batchsize * channels * h * w
frame = frame.transpose([0,3,1,2])
#переводим массив numpy в тензор
#данные должны быть тензором для совместимости с запущенной моделью, ожидающей floatTensor'ы
data['A'] = torch.FloatTensor(frame)
model.set_input(data) # распаковать данные из загрузчика
model.test()
#получаем только сгенерированное изображение - ищем в словаре ключ "fake"
result_image = model.get_current_visuals()['fake']
#используем tensor2im, предоставленный в файле util
result_image = util.tensor2im(result_image)
result_image = cv2.cvtColor(np.array(result_image), cv2.COLOR_BGR2RGB)
result_image = cv2.resize(result_image, (512, 512))
result_image = cv2.putText(result_image, str(opt.name)[6:-11], org, font,
fontScale, color, thickness, cv2.LINE_AA)
cv2.imshow('style', result_image)
#ASCII value of Esc is 27.
c = cv2.waitKey(1)
if c == 27:
break
if c == 99:
if style_model_index == len(style_models):
style_model_index = 0
opt.name = style_models[style_model_index]
style_model_index += 1
model = create_model(opt) # создаем модель с заданными опциями create a model given opt.model and other options
model.setup(opt)
cap.release()
cv2.destroyAllWindows()
Этот скрипт можно запустить той же командой bash:
python webcam.py --name MODEL_NAME --model test --preprocess none --no_dropout
Заключение
Вы можете заметить некоторое мерцание/шум на каждом стиле. Применение стиля к видео – сравнительно новая область разработки, в которой было сделано немало прекрасной, но дорогой в вычислительном смысле работы, чтобы передача между соседними фреймами была более гладкой.
Документация в репозитории StyleGAN научит вас, как обучать модели на собственных наборах данных и, таким образом, создавать собственные стили! Мне кажется, что коллекция комических книг позволила бы создать очень интересный стиль, который затем можно было бы применять к потоку видео.
Хотя уже есть программы, применяющие фильтры к потоку видео, вроде Photo Booth, передача стиля – это круто, поскольку модель, обученная глубокому пониманию и нюансам стиля, может применять его эффективно. Вы можете этого не заметить на потоке видео с низким разрешением, но на фотографиях с высоким разрешением вы увидите всю ее мощь:
Файл webcam.py и файлы репозитория можно найти в моей ветке на GitHub.