darkasevgen 6 мар в 13:10

Портретная гармонизация изображений

Средний

8 мин

2.8K

Блог компании SberDevicesОбработка изображений*Машинное обучение*Искусственный интеллект

Обзор

Привет! В данной публикации команда RnD CV из SberDevices познакомит вас с нашим подходом к решению задачи повышения степени реалистичности портретных изображений (по-научному — портретной гармонизации изображений). Мы провели ряд исследований и готовы представить:

PHNet — новую архитектуру нейронной сети, которая решает задачу портретной гармонизации;
FFHQH — крупномасштабный датасет из 70000 экземпляров, основанный на FFHQ датасете от Nvidia.

Наши исследования могут быть применены в различных продуктах и сервисах, в частности для SberJazz в задаче замены фона и гармонизации портрета, а также в GigaChat при добавлении модальностей «статичное изображение» и «видео» в сервис. Кроме того, наши наработки могут быть полезны в решении задачи Image Enhancement и в задаче манипуляции с изображениями, где требуется заменить / удалить / добавить объекты на фото, например в Kandinsky в режиме смешивания картинок. Также, с помощью нашего решения можно создавать синтетические датасеты.

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

Также мы оценили другие open-source решения, являющиеся лучшими в задаче гармонизации. Мы обучили их на нашем датасете FFHQH, и показали, что архитектура PHNet является лучшей в задаче портретной гармонизации.

Описание задачи

В начале обсудим классическую задачу гармонизации изображений.

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

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

Область первой фотографии является неконсистентной области . Вторая фотография показывает маску области . Третья фотография является гармонизированной.

Задача — изменить этот объект так, чтобы он стал гармонизирован (реалистичен) относительно области . Гармонизированное изображение — это линейная комбинация следующего вида:

$I_{harmonized} = f(I_{composite}, M) \cdot M + I_{composite} \cdot (1 - M),$

где $I_{composite}$ — RGB картинка, на которой находится негармонизированный объект, — алгоритм или нейронная сеть, $f(I_{composite}, M)$ — изображение с гармонизированным и естественным объектом, которое предсказывает . Объект извлекается из полученного изображения с помощью маски.

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

Архитектура

Мы разработали нейронную сеть, названную PHNet (Patch Harmonization Network) для решения задачи портретной гармонизации.

PHNet — это Unet архитектура, использующая Patch-based normalization (PN) модуль в блоках декодера и Patch-based feature extraction (PFE) в начале. Далее блоки будут объяснены.

PN-блок использует статистические методы переноса цвета (color transfer), вычисляя при этом следующие глобальные статистики:

$\mu_{glob}(I) = \frac{\sum\limits_{x \in I}{x}}{N} ,\: \sigma_{glob}(I) = \sqrt\frac{\sum\limits_{x \in I}{(x-\mu_{glob}(I))^2}}{N}$

В формуле — это количество пикселей во входном тензоре $I \in R^{C_i\times H_i\times W_i}$ .
В 4–6 этапах декодера входом в PN-блок является тензор из прошлого этапа декодера и уменьшенная версия маски $\hat{M} \in R^{1\times H_i\times W_i}$ .

Чтобы получить тензоры, характеризующие foreground и background области изображения — попиксельно умножается на и соответственно.
После умножения получаются: $F \in R^{C_i\times H_i\times W_i}$ и $B \in R^{C_i\times H_i\times W_i}$ .

Далее опишем 2 этапа нормализации:

Глобальная нормализация — это преобразование цветового распределения для переднего плана (foreground) в стандартное нормальное и приведение получившегося к распределению для заднего плана (background).
Для и вычисляются глобальные статистики: $\mu_{glob}(F)$ , $\mu_{glob}(B)$ , $\sigma_{glob}(B)$ , $\sigma_{glob}(F)$ по формулам, которые приведены выше. Все статистики вычисляются на каждый канал. Результатом глобальной нормализации является $\hat{F}_{glob}$ :
$\hat{F}_{glob} = [F - \mu_{glob}(F)] \cdot \frac{\sigma_{glob}(B)}{\sigma_{glob}(F)} + \mu_{glob}(B)$
На этой стадии background часть тензора остается неизменной;
Локальная нормализация состоит из использования локальных статистик, подсчитанных по патчам. Patch — это квадратная часть исходного изображения (или массива).
Пример с разбиением картинки на 16 патчей
Для локальной нормализации тензора вычисляются статистики среднего и отклонения на каждый патч и агрегируются вместе с обучаемыми параметрами $\alpha$ и $\beta$ :
$\mu_{loc}^{i}(I) = \frac{\sum\limits_{x \in p(I, i)}{x}}{N_i} , \mu_{loc}(I) = \sum_{i=1}^{P} \alpha^i(I) \cdot \mu_{loc}^{i}(I),$ $\sigma_{loc}^{i}(I) = \sqrt\frac{\sum\limits_{x \in p(I, i)}{(x-\mu_{loc}^{i}(I))^2}}{N_i}, \sigma_{loc}(I) = \sqrt{\frac{\sum_{i=1}^{P}\beta^i(I) \cdot \sigma_{loc}^i(I)^2}{P}},$
где находится в диапазоне от 1 до количества патчей .
Далее подсчитываются результаты локальной нормализации:
$\hat{F}_{loc} = [F - \mu_{glob}(F)] \cdot \frac{\sigma_{loc}(B)}{\sigma_{glob}(F)} + \mu_{loc}(B)$ $\hat{B}_{loc} = [B - \mu_{glob}(B)] \cdot \frac{\sigma_{loc}(B)}{\sigma_{glob}(B)} + \mu_{loc}(B)$

Результатом работы PN-блока является:

$N(F, B, M) = \tilde F \cdot M + \tilde B\cdot(1 - M),$

где $\tilde F = w_0 \cdot \hat{F}_{loc} + w_1 \cdot \hat{F}_{glob} + w_2$ и $\tilde B = \gamma_0 \cdot \hat{B}_{loc} + \gamma_1 \cdot B + \gamma_2$ .
$w_0, w_1, w_2, \gamma_0, \gamma_1, \gamma_2$ — это обучаемые параметры.

Мотивацией PN-блока является улавливание визуального стиля области background и введение информации в foreground. Это отличается от batch normalization-слоёв, в которых нормализация происходит в обеих областях с одинаковыми средним и дисперсией.

*Визуализация блока Patch-based Normalization (PN)*

В ряде экспериментов мы выявили, что использование PN-модулей в трёх блоках декодера — является оптимальным с точки зрения метрики. Они показаны на рисунке «Архитектура PHNet» в 4, 5 и 6 блоках декодера.

PFE-модуль представляет уже знакомый PN-блок с обучаемыми параметрами, в то время как все остальные остаются фиксированными. Выход PFE-блока адаптивно усредняется и полученные коэффициенты перевзвешивают веса в SE-блоке.

Интуиция такого блока исходит из экспериментов с $I_{composite}$ , где можно манипулировать передним планом, чтобы он сливался с фоном до такой степени, что почти исчезал, при этом сохраняя цветонезависимые атрибуты, как показано в рисунке «Архитектура PHNet».

iHarmony4

Самым используемым датасетом в задаче гармонизации является iHarmony4. Он состоит из 4 других датасетов: HCOCO, HAdobe5k, Hday2night, HFlickr. Как уже было сказано, гармонизировать изображение — это изменить цветовые характеристики объекта на переднем плане.

В случае iHarmony4, датасет составляли с помощью:

Настройки аугментаций; при составлении 2 из 4 датасетов использовался подход с привлечением фотографов, которые двигали ползунки цветовых характеристик в фотошопе. Таким образом, имея исходное изображение $I_{real}$ и маску объекта переднего плана , фотографы изменяли, уже упомянутый копии $I_{real}$ , создавая $I_{composite}$ ;
Алгоритма переноса цвета; при разметке $I_{composite}$ фотографий из 2 других датасетов использовали алгоритм:
- Для текущей пары $<I_{real}, M>$ находят референсную фотографию (reference) из этого же датасета; если , содержит, например, собаку, тогда reference-фотография тоже её включает;
- Для объекта переднего плана reference-фотографии определяют набор цветовых характеристик (тон, светлота и прочее);
- Переносят их на объект текущей пары; таким образом, формируется триплет $<I_{real}, M, I_{composite}>$ , где $I_{composite}$ — это копия $I_{real}$ в зоне , с перенесенными цветовыми характеристиками объекта внутри .

На переднем плане датасет iHarmony4 включает такие объекты как: животные, люди в полный рост, предметы интерьера и другие объекты, но он не является портретным.

В общей сумме датасет содержит 73146 триплетов $<I_{real}, M, I_{composite}>$ , из которых 7404 приходится на тестовую (она же валидационная) выборку.

FFHQH

Для создания нашего собственного портретного датасета мы использовали датасет FFHQ, который содержит только множество портретных изображений $I_{real}$ без масок. Таким образом, нам нужно не только создавать $\{I_{composite}\}$ , но и маски $\{M\}$ .

Процесс получения $I_{composite}$ и следующий:

Предсказание портретных масок людей с использованием архитектуры матирующей нейронной сети StyleMatte, примененной к реальным изображениям;
Бинаризация полученных масок по стандартному порогу probability = 0,5. Такой порог был выбран исходя из гистограммы распределения масок. На этом этапе маска сформирована;
Применение аугментаций к реальным изображениям в области маски :
- ColorJitter: brightness (0,5), contrast (0,4), saturation (0,06), hue (0,05);
- RandomPosterize (p=0,5; range=6);
- RandomAdjustSharpness (p=0,5; range=4);
Диапазон параметров аугментаций тщательно выбирался, чтобы избежать сильных и неестественных цветовых искажений. Применяемые аугментации являются разнородными с точки зрения цветовых преобразований;
Таким образом, фон объекта остался нетронутым, а человек на портрете видоизменился и стал не консистентным с фоном, тем самым $I_{composite}$ сформирован.

Датасет FFHQH был разделён на 3 выборки: 60000 триплетов — для обучения, 5000 — для валидации и 5000 — для тестирования.

Эксперименты и результаты

Данные
Мы обучили 2 модели и провели эксперименты на iHarmony4 и FFHQH.

Предобработка изображений
Мы не экспериментировали с аугментациями, и использовали только горизонтальное отражение во всех своих экспериментах. Фотографии перед подачей в PHNet переводились в диапазон [0; 1].

Метрики
В качестве основной метрики в этой задаче используется PSNR:

$PSNR = 10 \cdot log_{10}(\frac{MAX^2}{MSE})$

Оставим 2 важных ремарки. Во-первых, авторы других работ, считают метрики по разному.

В качестве MAX значения, в одних работах использовали 255, а в других — максимальное пиксельное значение $I_{real}$ (картинки кодируются от 0 до 255, но не всегда последнее присутствует в фото). Также мы видели реализацию когда MAX значением является разница между максимумом и минимум $I_{composite}$ .

Выбор MAX значения важен, так как напрямую влияет на значение PSNR. На странице paperswithcode есть топ моделей по решению задачи гармонизации, и например, авторы HDNet используют 255 в качестве MAX значения, а DucoNet и ряд других работ используют пиксельный максимум $I_{real}$ .
В своих экспериментах мы также выбираем MAX равным максимальному значению пикселей в $I_{real}$ из-за частоты использования.

Вторая ремарка — авторы используют небинарные маски для подсчета метрик. Мы не согласны с этим, так как три из четырёх датасетов в iHarmony4 содержат бинарные маски, и зачем считать метрики на других — не ясно. Мы попытались узнать об этом у авторов 4 работ, но ответа не получили.

PSNR и MSE (от которой зависит первый) можно подсчитать, используя не все пиксели изображения, а только принадлежащие foreground области. Такие метрики обозначаются fPSNR и fMSE.

Функции потерь
Мы используем 3 лосс функции: MSE, нормализованный на область (FN-MSE), градиентный и PSNR-Loss:

$PSNR\_Loss = \frac{100}{log_{10}(\frac{MAX}{\sqrt{MSE}})}$

Перейдём к результатам.

*Сравнение моделей на тестовом множестве iHarmony4*

Жирным цветом отображён лучший результат.

Из таблицы видно, что PHNet не является идеальной моделью для задачи обычной гармонизации, несмотря на то, что получена лучшая метрика на трёх из четырех датасетах. Это связано с тем, что PHNet не специализируется на изображениях, маска которых занимает малую площадь (78% масок iHarmony4 датасета занимают менее 15% площади изображения $I_{real}$ ).

В разрезе масок (объектов), занимающих разную площадь в iHarmony4, ситуация следующая:

*Сравнение моделей с различным соотношением переднего плана*

Данная таблица подтверждает, что PHNet не так хорош для объектов малой площади, как для больших.

И наконец результаты на нашем датасете FFHQH:

*Сравнение моделей на тестовом множестве FFHQH*

Все перечисленные модели обучены на нашем датасете FFHQH. Мы использовали параметры обучения, которые авторы применяли в своих статьях, и адаптировали их репозитории, сохраняя лучшую модель по метрике PSNR на валидации.

Все маски FFHQH имеют площадь больше 40%, что делают эту задачу нетривиальной и сложной, но как видно из результатов, PHNet справился с ней с большим отрывом от конкурентов. Итоговая модель имеет 39,9 миллионов параметров и весит 153 МБ. FPS на Intel H470 CPU в одном потоке — 1.01, а на NVIDIA Tesla V100 — 34.49.

Модели обучались и сравнивались на разрешении 256×256.