Привет, Хабр! Я Миша, R&D инженер в Prisma Labs. Наш отдел занимается разработкой новых фичей для мобильных приложений Prisma и Lensa. Prisma это AI-based фоторедактор, который по нажатию одной кнопки всего за пару секунд переносит художественный стиль с картин на ваши фотографии. Недавно мы решили расширить его функционал и дать пользователям возможность переносить стиль с портретов, нарисованных художниками, на свои селфи. В этом посте я расскажу про особенности задачи портретного Style Transfer, о существующих методах, которые мы попробовали, и о том, как адаптировали их под смартфоны.

Кто такой этот ваш Style Transfer?!

Чтобы мы с вами были на одной волне, я начну свой рассказ с небольшого описания задачи Style Transfer (ST) и обзора наиболее популярных подходов к её решению. 

В общем случае, задача ST формулируется так: у нас есть фотография C и стилевое изображение S (например, это может быть картина маслом, диджитал арт и др.); мы хотим получить изображение, в котором “контент” из C (под контентом подразумевается всё что изображено на фотографии) будет нарисован в том же “стиле”, что и S (более строгое определение контента и стиля изображений можно найти тут).  

Исследователи занимаются этой задачей уже давно, но широкую известность она получила в 2015 году, когда вышла статья Л. Гатиса, и др., в которой предложили ST на основе нейронных сетей (они на тот момент были на вершине своего хайпа). Вот так выглядит пайплайн метода:

Верхнеуровнево этот метод можно описать так: значения пикселей картинки с шумом итеративно изменяются так, чтобы её дескрипторы (фичи) после каждой итерации приближались к дескрипторам стилевой и контентной картинок в терминах L₂ нормы. На современных GPU весь процесс обычно занимает несколько минут.

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

Несмотря на это, гатисовский ST работает очень медленно – на GeForce GTX 1080 стилизация изображения в разрешении 1024х1024 занимает в среднем около минуты.

В 2016 году вышла работа Д. Джонсона и др., в которой предлагается первый real-time Style Transfer на нейросетях. Основная идея метода заключается в обучении под каждый стиль свёрточной нейронной сети с архитектурой вида  энкодер-декодер с использованием функции потерь, предложенной Гатисом: 

Этот подход явно уступает гатисовскому ST в качестве, но по скорости заметно превосходит его (стилизация одного изображения занимает сотни миллисекунд), а значит, хорошо подходит для запуска на смартфонах. Именно этот метод и стал основой для нашего приложения Prisma.

Зачем нужно придумывать что-то ещё для портретов?

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

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

Попытки установить такое соответствие по фичемапам нейросети (например, посчитав для каждой фичемапы 4D тензор корреляций) в большинстве случаев не приводят к желаемому эффекту. Одна из потенциальных причин кроется в том, что используемая нейросеть обучалась на задаче классификации, и поэтому она извлекает те признаки изображения, по которым его можно отнести к одному из классов. Кроме того, при обучении использовался датасет ImageNet, в котором нет ни человеческих лиц, ни тем более художественных портретов. Поэтому нет каких-либо гарантий, что извлекаемые нейросетью признаки будут содержать информацию о семантических частях лица, которая нам нужна для установления соответствия между контентным и стилевым лицами.

Может возникнуть резонный вопрос: а что будет, если мы обучим сетку для face segmentation/face recognition, возьмём её энкодер и заменим им VGG в гатисовском ST? На самом деле, этот вопрос открытый, я лично пока не видел ни одной работы, в которой это исследовали.  

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

Кодируем информацию о стиле в фичемапах контентного изображения

В статье авторы предлагают модифицировать гатисовский метод для работы с портретами. Их пайплайн выглядит так:

Алгоритм состоит из двух этапов – генерация (Stage A) и опциональный пост-процессинг стилизованного изображения (Stage B). 

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

Более подробное описание метода

Формально преобразование фичей контентного изображения можно записать так:

,

,

,

где  , – фичемапы контентного и стилевого изображений на L-ом слое,  – модифицированные фичемапы, – “gain map”, который показывает, как соотносятся активации стилевой и контентной фичемап. 

Если использовать вместо в первой формуле, вся информация о контенте пропадает – . Получается не совсем то, что нам нужно, потому что при модификации мы хотим сохранить и контентную, и стилевую информацию. Для этого к применяется операция clamp, которая накладывает ограничения на максимальное и минимальное значения в gain maps. Интуитивно, это позволяет слабее модулировать фичи контента в тех областях, в которых содержится важная информация о стиле (большие значения в gain maps), и при этом оставлять активации фичемапы контента практически без изменений в областях, где содержится менее важная информация о стиле (очень маленькие значения в gain maps); делаем их близкими к единице -> практически оставляем фичемапы контента неизменными. 

Пример визуализации исходных и преобразованных фичемап, а также gain maps:

В описанном выше преобразовании фичемап предполагается, что лица на контентном и стилевом изображениях выровнены. Реальные же фотографии могут отличаться от стилевой картинки ракурсом и углом поворота головы, поэтому до начала процесса стилизации стилевое и контентное изображения выравниваются – текстура лица на стилевом изображении “натягивается” на лицо c контентного изображения. Авторы предлагают делать выравнивание по 68 лендмаркам лица комбинацией методов Image Morphing и SIFT-Flow.

Инсайд про более простой и быстрый подход к выравниванию

В своей реализации мы использовали альтернативный, более быстрый подход:

1. Считаем векторы смещений между соответствующими лендмарками стилевого и контентного изображений;

2. Значения сдвигов для остальных точек сетки находим путём интерполяции, получаем flow field F, в котором F(x, y) – вектор, показывающий откуда взять пиксель (x, y) в стилевом изображении; 

3. Делаем warping стилевого изображения по F.

Результат выравнивания выглядит немного крипово:

Вот так выглядят фотографии, стилизованные нашей реализацией этого метода:

При сравнении полученных стилизаций с результатами гатисовского ST становится особенно заметно, что рассматриваемый метод более “осознанно” и натурально переносит стиль в области лица (на первом и третьем стилях это видно невооружённым глазом):

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

В данном случае, конечно, эту проблему можно решить блендингом результата с похожей текстурой, но если у нас какой-нибудь карандашный стиль, это не поможет.

На фотографиях с сильно повёрнутыми лицами алгоритм выравнивания работает плохо, отчего появляются нежелательные артефакты:

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

Если есть сильные отличия по размеру лица, причёске и другим атрибутам между лицами на контентном и стилевом изображениях, возникает ghosting артефакт (когда часть стилевого изображения попадает на результат):

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

Нейросети? Не, не слышали!

Наверняка у многих, кто дочитал до этих строк, сложилось впечатление, что придумать что-то лучше нейронок для задачи портретного ST нельзя. Поспешу вас в этом разубедить.

В 2001 году придумали метод Image Analogies, который с данной пары изображений “до/после применения фильтров” умеет переносить эффект наложенных фильтров (по аналогии с имеющимся примером “до/после” наложения фильтра) на любое другое изображение:

Image Analogies работает примерно так:

1. Строится NNF(A, B) (Nearest Neighbour Field) – маппинг , который каждому пикселю q из B ставит в соответствие ближайший (в терминах некоторого функционала) пиксель p из A; так как между A и A', B и B' есть прямое соответствие, можно положить NNF(A', B') = NNF(A, B);

2. Каждому пикселю q из B' присваивается значение ближайшего по NNF пикселя p* из A' , где

Для построения NNF нужно уметь сравнивать пиксели изображений A и B между собой. Для этого используются различные дескрипторы пикселей. В самом простом случае, дескриптором пикселя могут быть его RGB значения, но из-за своей “низкоуровневости” такой дескриптор малоинформативен и чаще всего его использование приводит к плохому NNF. Поэтому авторы предложили помимо RGB значений использовать результаты применения steerable фильтров к изображениям A и B. По извлеченным дескрипторам близость двух пикселей из A и B (а, следовательно, и из A' и B') оценивается как L₂ норма разницы соответствующих дескрипторов. Чтобы ускорить поиск ближайшего соседа для каждого пикселя, можно использовать приближенные методы поиска ближайших соседей.

def create_image_analogy(src_desc, src_st, dst_desc, num_scales, mu):
   """
   Image analogies algorithm's pseudocode

   Args:
       src_desc: source image's multi-channel descriptors (A)
       src_st: stylized version of A (A')
       dst_desc: target images's multi-channel descriptors (B)
       num_scales: number of scales in image pyramid
       mu: importance of guiding descriptors

   Returns:
       dst_st: stylized version of B (B')
   """
    def energy(p, q, src_desc, dst_desc, src_st, dst_st, mu):
        return (src_st[p] - dst_st[q]) ** 2 + mu * (src_desc[p] - dst_desc[q]) ** 2

    src_desc_pyr, src_st_pyr, dst_desc_pyr = make_gaussian_pyramids(src_desc, src_st, dst_desc, num_scales=num_scales)

    dst_st_pyr = []
    for level in range(num_scales):
        nnf_at_level = init_empty_nnf(src_st_pyr[level])
        # build NNF
        for q in coordinates_generator(src_st_pyr[level]):
            p = find_func_minimizer(src_desc_pyr, dst_desc_pyr, src_st_pyr, dst_st_pyr, func=energy, mu=mu)
            nnf_at_level[q] = p
        # replace all pixels values in B' with its neighbours in A'
        dst_st_at_level = warp_image_with_nnf(src_st_pyr[level], nnf_at_level)
        dst_st_pyr.append(dst_st_at_level)

    return dst_st_pyr[-1]

В последующие годы появилось много работ, в которых предлагались различные варианты улучшения дескрипторов и методов построения NNF. Из них хотелось бы выделить работу StyLit, в которой авторы применяют Image Analogies для стилизации трёхмерных моделей. В качестве дескрипторов предлагается использовать RGB и LPE (Light Paths Expressions) тензоры. LPE описывают распространение света в данной 3D сцене и за счёт этого позволяют получить более реалистичную стилизацию, которая больше похожа на рисунок от руки художника.

Также в этой работе предложили улучшение алгоритма расчёта NNF, использовавшегося в Image Analogies

В Image Analogies используется жадный алгоритм, который плохо справляется со стилизацией высокоуровневых структур. Поэтому в StyLit предложили минимизировать , делая несколько итераций EM алгоритма:

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

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

В 2017 StyLit был адаптирован для задачи портретного ST и переименован в FaceStyle. В FaceStyle предложили использовать специфичные для лиц дескрипторы:

Расскажу о них подробнее. G_seg (segmentation guidance) задает соответствие между семантическими участками лиц на контентном и стилевом изображениях. Иными словами, не даёт алгоритму стилизации копировать патчи с бровей стилевого лица и вставлять их на лоб контентного лица. Для получения маски лица и волос авторы используют аналитический маттинг со своими модификациями (подробности в статье). Маски бровей, глаз, носа, губ и рта рисуют по 68 лендмаркам лица. Дескриптор G_seg представляет из себя RGB картинку, в котором каждая семантическая часть лица закодирована своим цветом. 

На первый взгляд может показаться, что одного G_seg может быть вполне достаточно, чтобы получить хорошую стилизацию, но это не так. Например, по G_seg нельзя отличить правую часть лица от левой, а иногда это может быть полезным. Для того чтобы локализовать перенос, авторы предлагают использовать G_pos (positional guidance) – для контентного изображения это по сути грубый NNF, построенный по лицевым лендмаркам, для стилевого – identity NNF (в каждом пикселе закодированы координаты этого пикселя). Для расчёта G_pos авторы используют метод Moving Least Squares Deformation, натравленный на лендмарки. Мы в своей реализации для получения этого дескриптора использовали подход к выравниванию лиц, описанный в предыдущей главе. 

Ни G_seg, ни G_pos никак не форсят сохранение контента на стилизованной фотографии, поэтому нужен ещё один дескриптор – G_app (appearence guidance). G_app стиля – это просто стилевое изображение, преобразованное в grayscale. Сделать то же самое для контентного изображения нельзя, потому что, если на нем будет присутствовать тень/блик, это обязательно проявится в виде артефакта на стилизованном изображении. Поэтому авторы сначала выравнивают локальный контраст и интенсивность контентного и стилевого изображений, используя метод переноса стиля с одного портретного фото на другое – ST for Headshot Portraits, а уже затем используют полученное контентное изображение с выровненным контрастом в качестве дескриптора. Существует и более быстрый способ получения этого дескриптора:

1. посчитать первый уровень Laplacian Pyramid (LP) для контента и стиля;

2. выровнять гистограмму полученного LP контента по LP стиля;

3. полученные LP использовать в качестве G_app.

Если визулизировать FaceStyle в концепции Image Analogies, получится что-то такое:

Стилизации, которые мы получили, этим методом, лишены недостатков предыдущего метода портретного ST и выглядят без преувеличения очень круто:

Правда для FaceStyle справедливо правило “Garbage in - garbage out” – если накосячили при генерации дескрипторов (такое часто бывает, например, когда лицо повернуто практически в профиль, на человеке надеты солнцезащитные очки, второй человек попадает в кадр и т.д.), получим дичь:

Также можно заметить, что FaceStyle не использует дескрипторы для чего-то кроме лица. Соответственно, всё, что находится вне маски лица, стилизуется не так, как мы ожидаем. Мы решили расширить дескриптор G_seg, путем добавления маски шеи и тела, при этом нам на сдачу добавилась маска бэкграунда. Но это не всегда даёт хорошие результаты (примеры в первых двух столбцах на картинке выше). Придумать тут что-то более умное довольно-таки трудно, хотя попытки предпринимались.

Ещё одна слабая сторона FaceStyle – это длительность стилизации. При условии, что все дескрипторы предпосчитаны, стилизация одного изображения в разрешении 1024х1024 у нас занимает около 25с на Intel(R) Xeon(R) CPU E5-2620. На смартфоне, соответственно, будет ещё дольше (помним, что нам нужно будет также считать дескрипторы).

Distillation is my profession

Все рассмотренные выше методы ST для портретов итеративные и далеки от real-time даже на современных десктопных CPU/GPU. Как же ускорить медленный итеративный алгоритм, который умеет генерировать пары source-stylized images хорошего качества? Правильно, дистиллировать его нейросетями!

Подход такой:

1. Генерируем датасет пар изображений в выбранном стиле;

2. Учим на полученном датасете image-to-image сетку;

3. Profit!

Сейчас этот подход довольно-таки популярный, особенно среди любителей StyleGAN (вот, например, статья с ECCV 2020). Есть даже работа, в которой авторы дистиллировали FaceStyle. Далее расскажу вкратце про основные идеи из неё.

В качестве генератора использовали нейросеть архитектуры энкодер-декодер из работы Real-time Style Transfer, модифицировав её следующим образом:

• добавили skip connections с выходов слоёв энкодера к соответствующим слоям декодера;

• увеличили количество residual блоков;

• добавили дополнительный свёрточный слой на выходе декодера.

В качестве лосса использовали взвешенную сумму: 

• adversarial лосса на основе PatchGAN

• попиксельного L₁ лосса между выходом сети и ground truth

• perceptual лосса

На Nvidia RTX 2080 Ti такой генератор работает под 15 FPS на изображениях 512х512, что вполне можно считать real-time. Помимо прироста в скорости, на сдачу сетка ещё и бустанула качество стилизаций на сложных кейсах – генератор неплохо справляется там, где косячил FaceStyle:

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

На простых кейсах стилизации FaceStyle выглядят более чётко по сравнению с результатами дистилляции – в первом случае стилизованное изображение “собирается” из патчей стилевого, что позволяет сохранить все художественные детали; нейросеть же – параметрическая модель, которая ввиду ограниченного количества параметров может лишь аппроксимировать маппинг между исходными изображениями и таргетами:

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

Как мы затаскивали это всё на смартфон

Нам предстояло решить две задачи:

1. оптимизировать архитектуру генератора так, чтобы он отрабатывал на смартфонах за адекватное время;

2. увеличить разрешение до 1024х1024.

С одной стороны, нужно было сделать генератор более mobile-friendly, и при этом сохранить совместимость с фреймворками инференса нейросетей на смартфонах (CoreML, ONNX, TFLite); c другой – сохранить качество результатов на уровне большого генератора. Мы решили взять за основу архитектуру нейросети, которая раньше использовалась у нас в сегментации фона на портретных фотографиях:

По сути, это всё тот же энкодер-декодер со skip-connections, что использовали авторы для дистилляции FaceStyle, но со следующими модификациями:

• В residual блоках используются separable свёрточные слои вместо обычных – они выполняют практически ту же работу, что и обычные свёртки, только за меньшее количество операций и требуют меньше памяти.

• В качестве basic блока используется не Conv + ReLU слой, а целый residual блок.

• Вместо конкатенации выходов слоёв энкодера с выходами слоёв декодера используется поэлементное сложение – это ускоряет сходимость и улучшает качество результатов.

Мы адаптировали эту архитектуру под нашу задачу следующим образом:

• Сделали энкодер более глубоким;

• Увеличили разрешение входного изображения до 512х512;

• Добавили Instance Normalization слои (без них сетка сильно артефачит);

• Сделали энкодер и декодер симметричными (одинаковое количество residual блоков на соответствующих слоях);

• Заменили Transposed Conv слои на Upsample + Conv (про мотивацию можно почитать тут).

При обучении мы так же использовали PatchGAN лосс, поэтому нужно было ещё и правильно подобрать дискриминатор к оптимизированному генератору. Для этого мы уменьшили количество фичемап в исходном дискриминаторе и добавили Spectral Normalization во все свёрточные слои.

С тяжелым генератором разобрались, осталось решить проблему с разрешением. Для этого мы использовали модифицированный ESRGAN, увеличивающий разрешение входного изображения в 2 раза. Учили его на кропах изображений, стилизованных обученной ST сеткой, в качестве ground truth использовались результаты FaceStyle. 

В итоге весь пайплайн получился довольно-таки лёгким (< 5.5 Mb) и быстрым – на IPhone 12 Pro стилизация занимает около одной секунды, при этом качество результатов получается сопоставимым с FaceStyle и его дистилляцией.

Вместо заключения

Итак, в этом посте я постарался вкратце рассказать вам про задачу портретного ST и про два наиболее популярных метода её решения. Разумеется есть ещё куча крутых работ в этой области, про которые я бы с удовольствием рассказал вам, но во избежание “многа букав” решил этого не делать (заинтересованным читателям оставлю ссылки на эти работы ниже). Качество этих методов хоть и на порядок выше того, что могут дать гатисовский и джонсоновский ST, работают они, прямо скажем, не быстро. В таких случаях использовать подход с дистилляцией медленного алгоритма быстрыми и оптимизированными под смартфоны нейросетями оказывается разумным, а иногда и единственным вариантом (особенно если у вас в компании очень хорошо настроены процессы выката в мобильный продакшен фильтров на основе нейросетей). Этим вариантом мы и воспользовались, чтобы добавить портретные фильтры в Prisma:  теперь от превращения своего селфи в портрет в стиле поп-арта, карандашного скетча, акварели и др. пользователя отделяет всего лишь нажатие одной кнопки. 

Мы активно собираем фидбек от пользователей по выпущенным портретным фильтрам  и работаем над их улучшением. Больше всего юзерам не нравится, что их селфи слишком сильно обрезаются, поэтому одна из приоритетных задач на будущее – избавиться от кропа. Повышение стабильности нейросетей – другое потенциальное направление улучшения портретных фильтров. Качество переноса стиля сильно зависит от фотографии – если фотография сделана с плохим освещением, засветами или с далёким лицом,  результат будет не очень.

Сейчас у нас в проде есть четыре портретных ST фильтра: Leya, Aqua, Michael и Arturo, с которыми вы можете поиграться. На этом мы решили не останавливаться, поэтому в ближайшее время подъедет ещё больше фильтров. Следите за апдейтами;)

Список источников

1. A Neural Algorithm of Artistic Style

2. Demystifying Neural Style Transfer

3. Perceptual Losses for Real-Time Style Transfer and Super-Resolution

4. Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization

5. Universal Style Transfer via Feature Transforms

6. Painting Style Transfer for Head Portraits using Convolutional Neural Networks

7. Image Analogies

8. StyLit: Illumination-Guided Example-Based Stylization of 3D Renderings

9. Example-Based Synthesis of Stylized Facial Animations

10. Style Transfer for Headshot Portraits

11. Exemplar-based Portrait Style Transfer

12. Real-Time Patch-Based Stylization of Portraits Using Generative Adversarial Network

13. Image-to-Image Translation with Conditional Adversarial Networks

Что ещё можно почитать по теме

1. Visual Attribute Transfer through Deep Image Analogy

2. STALP: Style Transfer With Auxiliary Limited Pairing

3. Interactive Video Stylization Using Few-Shot Patch-Based Training

4. Geometric Style Transfer