Pull to refresh

Создаём линзы для SnapChat с использованием pix2pix

Reading time 3 min
Views 4.3K
Original author: Smit Shilu

Почти такой же заголовок носит и моя предыдущая статья, с той лишь разницей, что тогда я создавал линзы для SnapChat алгоритмически, используя dlib и openCV, а сегодня хочу показать, как можно добиться результата, используя машинное обучение. Этот подход позволит не заниматься ручным проектированием алгоритма, а получать итоговое изображение прямо из нейронной сети.


Вот что мы получим:



Что такое pix2pix?


Это способ преобразования изображения в изображение с помощью состязательных сетей (широко известный как pix2pix).


Название «pix2pix» означает, что сеть обучена преобразовывать входное изображение в соответствующее ему выходное изображение. Вот примеры таких преобразований:



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


В отличие от подходов, применявшихся ранее, pix2pix учится решать задачи гораздо быстрее и на меньшей обучающей выборке. Так, например, результаты ниже были получены при обучении с использованием Pascal Titan X GPU на датасете из 400 пар изображений и менее чем за два часа.



Как работает pix2pix?


pix2pix использует две нейронные сети, обучающиеся параллельно:


  1. Генератор
  2. Дискриминатор

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


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


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


Практика


Тренировать наш генератор фильтров для SnapChat мы будем на изображениях 256x256 (большие размеры потребуют большего количества видеопамяти). Для создания датасета воспользуемся кодом из предыдущего туториала.


Я скачал множество изображений лиц и применил к каждому фильтр «Thug Life Glasses». Получится что-то типа таких пар:



Для создания модели возьмём репозиторий pix2pix на основе TensorFlow. Склонируйте его и установите Tensorflow.


Команда для запуска обучения будет такой:


python pix2pix.py --mode train --output_dir dir_to_save_checkpoint --max_epochs 200 --input_dir dir_with_training_data --which_direction AtoB

Параметр which_direction задаёт направление обучения. AtoB означает что мы хотим превратить изображение A (слева, без очков) в изображение B (справа, с очками). Кстати, обратите внимание, что pix2pix может успешно научиться восстанавливать исходное изображение из изображения с фильтром, достаточно лишь поменять направление обучения.


Отслеживать прогресс обучения можно с помощью tensorboard, для чего надо выполнить команду:


tensorboard --logdir=dir_to_save_checkpoint

Как только вы видите, что результаты на тренировочных данных стали достаточно хорошими, можно остановить обучение и проверить работу модели на произвольных данных. Продолжить тренировку с последней контрольной точки можно так:


python pix2pix.py --mode train --output_dir dir_to_save_checkpoint --max_epochs 200 --input_dir dir_with_training_data --which_direction AtoB --checkpoint dir_of_saved_checkpoint

Заключение


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



Tags:
Hubs:
If this publication inspired you and you want to support the author, do not hesitate to click on the button
+10
Comments 0
Comments Leave a comment

Articles