使用JojoGAN创建风格化的面部图

介绍

风格迁移是神经网络的一个发展领域，它是一个非常有用的功能，可以集成到社交媒体和人工智能应用程序中。几个神经网络可以根据训练数据将图像样式映射和传输到输入图像。在本文中，我们将研究 JojoGAN，以及仅使用一种参考样式来训练和生成具有该样式的任何图像的过程。

JoJoGAN：One Shot Face Stylization

One Shot Face Stylization（一次性面部风格化）可用于 AI 应用程序、社交媒体过滤器、有趣的应用程序和业务用例。随着 AI 生成的图像和视频滤镜的日益普及，以及它们在社交媒体和短视频、图像中的使用，一次性面部风格化是一个有用的功能，应用程序和社交媒体公司可以将其集成到最终产品中。

因此，让我们来看看用于一次性生成人脸样式的流行 GAN 架构——JojoGAN。

JojoGAN 架构

JojoGAN 是一种风格迁移程序，可让将人脸图像的风格迁移为另一种风格。它通过GAN将参考风格图像反转为近似的配对训练数据，根据风格化代码生成真实的人脸图像，并与参考风格图像相匹配。然后将该数据集用于微调 StyleGAN，并且可以使用新的输入图像，JojoGAN 将根据 GAN 反转（inversion）将其转换为该特定样式。

JojoGAN 架构和工作流程

JojoGAN 只需一种参考风格即可在很短的时间内（不到 1 分钟）进行训练，并生成高质量的风格化图像。

JojoGan 的一些例子

JojoGAN 生成的风格化图像的一些示例：

风格化的图像可以在各种不同的输入风格上生成并且可以修改。

JojoGan 代码深潜

让我们看看 JojoGAN 生成风格化人像的实现。有几个预训练模型可用，它们可以在我们的风格图像上进行训练，或者可以修改模型以在几分钟内更改风格。

JojoGAN 的设置和导入

克隆 JojoGAN 存储库并导入必要的库。在 Google Colab 存储中创建一些文件夹，用于存储反转代码、样式图像和模型。

！git clone https：／／github．com／mchong6／JoJoGAN．git

％cd JoJoGAN

！pip install tqdm gdown scikit－learn＝＝0．22 scipy lpips dlib opencv－python wandb

！wget https：／／github．com／ninja－build／ninja／releases／download／v1．8．2／ninja－linux．zip

！sudo unzip ninja－linux．zip －d ／usr／local／bin／

import torch

torch．backends．cudnn．benchmark ＝ True

from torchvision import transforms， utils

from util import ＊

from PIL import Image

import math

import random

import os

import numpy

from torch import nn， autograd， optim

from torch．nn import functional

from tqdm import tqdm

import wandb

from model import ＊

from e4e＿projection import projection

from google．colab import files

from copy import deepcopy

from pydrive．auth import GoogleAuth

from pydrive．drive import GoogleDrive

from google．colab import auth

from oauth2client．client import GoogleCredentials

模型文件

使用 Pydrive 下载模型文件。一组驱动器 ID 可用于预训练模型。这些预训练模型可用于随时随地生成风格化图像，并具有不同的准确度。之后，可以训练用户创建的模型。

＃Download models

＃optionally enable downloads with pydrive in order to authenticate and avoid drive download limits．

download＿with＿pydrive ＝ True

device ＝ ＇cuda＇ ＃［＇cuda＇， ＇cpu＇］

！wget http：／／dlib．net／files／shape＿predictor＿68＿face＿landmarks．dat．bz2

！bzip2 －dk shape＿predictor＿68＿face＿landmarks．dat．bz2

！mv shape＿predictor＿68＿face＿landmarks．dat models／dlibshape＿predictor＿68＿face＿landmarks．dat

％matplotlib inline

drive＿ids ＝ ｛

＂stylegan2－ffhq－config－f．pt＂： ＂1Yr7KuD959btpmcKGAUsbAk5rPjX2MytK＂，

＂e4e＿ffhq＿encode．pt＂： ＂1o6ijA3PkcewZvwJJ73dJ0fxhndn0nnh7＂，

＂restyle＿psp＿ffhq＿encode．pt＂： ＂1nbxCIVw9H3YnQsoIPykNEFwWJnHVHlVd＂，

＂arcane＿caitlyn．pt＂： ＂1gOsDTiTPcENiFOrhmkkxJcTURykW1dRc＂，

＂arcane＿caitlyn＿preserve＿color．pt＂： ＂1cUTyjU－q98P75a8THCaO545RTwpVV－aH＂，

＂arcane＿jinx＿preserve＿color．pt＂： ＂1jElwHxaYPod5Itdy18izJk49K1nl4ney＂，

＂arcane＿jinx．pt＂： ＂1quQ8vPjYpUiXM4k1＿KIwP4EccOefPpG＿＂，

＂arcane＿multi＿preserve＿color．pt＂： ＂1enJgrC08NpWpx2XGBmLt1laimjpGCyfl＂，

＂arcane＿multi．pt＂： ＂15V9s09sgaw－zhKp116VHigf5FowAy43f＂，

＂sketch＿multi．pt＂： ＂1GdaeHGBGjBAFsWipTL0y－ssUiAqk8AxD＂，

＂disney．pt＂： ＂1zbE2upakFUAx8ximYnLofFwfT8MilqJA＂，

＂disney＿preserve＿color．pt＂： ＂1Bnh02DjfvN＿Wm8c4JdOiNV4q9J7Z＿tsi＂，

＂jojo．pt＂： ＂13cR2xjIBj8Ga5jMO7gtxzIJj2PDsBYK4＂，

＂jojo＿preserve＿color．pt＂： ＂1ZRwYLRytCEKi＿＿eT2Zxv1IlV6BGVQ＿K2＂，

＂jojo＿yasuho．pt＂： ＂1grZT3Gz1DLzFoJchAmoj3LoM9ew9ROX＿＂，

＂jojo＿yasuho＿preserve＿color．pt＂： ＂1SKBu1h0iRNyeKBnya＿3BBmLr4pkPeg＿L＂，

＂art．pt＂： ＂1a0QDEHwXQ6hE＿FcYEyNMuv5r5UnRQLKT＂，

｝

＃ from StyelGAN－NADA

class Downloader（object）：

def ＿＿init＿＿（self， use＿pydrive）：

self．use＿pydrive ＝ use＿pydrive

if self．use＿pydrive：

self．authenticate（）

def authenticate（self）：

auth．authenticate＿user（）

gauth ＝ GoogleAuth（）

gauth．credentials ＝ GoogleCredentials．get＿application＿default（）

self．drive ＝ GoogleDrive（gauth）

def download＿file（self， file＿name）：

file＿dst ＝ os．path．join（＇models＇， file＿name）

file＿id ＝ drive＿ids［file＿name］

if not os．path．exists（file＿dst）：

print（f＇Downloading ｛file＿name｝＇）

if self．use＿pydrive：

downloaded ＝ self．drive．CreateFile（｛＇id＇：file＿id｝）

downloaded．FetchMetadata（fetch＿all＝True）

downloaded．GetContentFile（file＿dst）

else：

！gdown －－id ＄file＿id －O ＄file＿dst

downloader ＝ Downloader（download＿with＿pydrive）

downloader．download＿file（＇stylegan2－ffhq－config－f．pt＇）

downloader．download＿file（＇e4e＿ffhq＿encode．pt＇）

加载生成器

加载原始和微调生成器。设置用于调整图像大小和规范化图像的 transforms。

latent＿dim ＝ 512

＃ Load original generator

original＿generator ＝ Generator（1024， latent＿dim， 8， 2）．to（device）

ckpt ＝ torch．load（＇models／stylegan2－ffhq－config－f．pt＇， map＿location＝lambda storage， loc： storage）

original＿generator．load＿state＿dict（ckpt［＂g＿ema＂］， strict＝False）

mean＿latent ＝ original＿generator．mean＿latent（10000）

＃ to be finetuned generator

generator ＝ deepcopy（original＿generator）

transform ＝ transforms．Compose（

［

transforms．Resize（（1024， 1024）），

transforms．ToTensor（），

transforms．Normalize（（0．5， 0．5， 0．5）， （0．5， 0．5， 0．5）），

］

）

输入图像

设置输入图像位置。对齐和裁剪面并重新设置映射的样式。

＃image to the test＿input directory and put the name here

filename ＝ ＇face．jpeg＇ ＃＠param ｛type：＂string＂｝

filepath ＝ f＇test＿input／｛filename｝＇

name ＝ strip＿path＿extension（filepath）＋＇．pt＇

＃ aligns and crops face

aligned＿face ＝ align＿face（filepath）

＃ my＿w ＝ restyle＿projection（aligned＿face， name， device， n＿iters＝1）．unsqueeze（0）

my＿w ＝ projection（aligned＿face， name， device）．unsqueeze（0）

预训练图

选择预训练好的图类型，选择不保留颜色的检查点，效果更好。

plt．rcParams［＇figure．dpi＇］ ＝ 150

pretrained ＝ ＇sketch＿multi＇ ＃［＇art＇， ＇arcane＿multi＇， ＇sketch＿multi＇， ＇arcane＿jinx＇， ＇arcane＿caitlyn＇， ＇jojo＿yasuho＇， ＇jojo＇， ＇disney＇］

＃Preserve color tries to preserve color of original image by limiting family of allowable transformations．

if preserve＿color：

ckpt ＝ f＇｛pretrained｝＿preserve＿color．pt＇

else：

ckpt ＝ f＇｛pretrained｝．pt＇

生成结果

加载检查点和生成器并设置种子值，然后开始生成风格化图像。用于 Elon Musk 的输入图像将根据图类型进行风格化。

＃Generate results

n＿sample ＝ 5＃｛type：＂number＂｝

seed ＝ 3000 ＃｛type：＂number＂｝

torch．manual＿seed（seed）

with torch．no＿grad（）：

generator．eval（）

z ＝ torch．randn（n＿sample， latent＿dim， device＝device）

original＿sample ＝ original＿generator（［z］， truncation＝0．7， truncation＿latent＝mean＿latent）

sample ＝ generator（［z］， truncation＝0．7， truncation＿latent＝mean＿latent）

original＿my＿sample ＝ original＿generator（my＿w， input＿is＿latent＝True）

my＿sample ＝ generator（my＿w， input＿is＿latent＝True）

＃ display reference images

if pretrained ＝＝ ＇arcane＿multi＇：

style＿path ＝ f＇style＿images＿aligned／arcane＿jinx．png＇

elif pretrained ＝＝ ＇sketch＿multi＇：

style＿path ＝ f＇style＿images＿aligned／sketch．png＇

else：

style＿path ＝ f＇style＿images＿aligned／｛pretrained｝．png＇

style＿image ＝ transform（Image．open（style＿path））．unsqueeze（0）．to（device）

face ＝ transform（aligned＿face）．unsqueeze（0）．to（device）

my＿output ＝ torch．cat（［style＿image， face， my＿sample］， 0）

生成的结果

结果生成为预先训练的类型“Jojo”，看起来相当准确。

现在让我们看一下在自创样式上训练 GAN。

使用你的风格图像进行训练

选择一些面部图，甚至创建一些自己的面部图并加载这些图像以训练 GAN，并设置路径。裁剪和对齐人脸并执行 GAN 反转。

names ＝ ［＇1．jpg＇， ＇2．jpg＇， ＇3．jpg＇］

targets ＝ ［］

latents ＝ ［］

for name in names：

style＿path ＝ os．path．join（＇style＿images＇， name）

assert os．path．exists（style＿path）， f＂｛style＿path｝ does not exist！＂

name ＝ strip＿path＿extension（name）

＃ crop and align the face

style＿aligned＿path ＝ os．path．join（＇style＿images＿aligned＇， f＇｛name｝．png＇）

if not os．path．exists（style＿aligned＿path）：

style＿aligned ＝ align＿face（style＿path）

style＿aligned．save（style＿aligned＿path）

else：

style＿aligned ＝ Image．open（style＿aligned＿path）．convert（＇RGB＇）

＃ GAN invert

style＿code＿path ＝ os．path．join（＇inversion＿codes＇， f＇｛name｝．pt＇）

if not os．path．exists（style＿code＿path）：

latent ＝ projection（style＿aligned， style＿code＿path， device）

else：

latent ＝ torch．load（style＿code＿path）［＇latent＇］

latents．append（latent．to（device））

targets ＝ torch．stack（targets， 0）

latents ＝ torch．stack（latents， 0）

微调 StyleGAN

通过调整 alpha、颜色保留和设置迭代次数来微调 StyleGAN。加载感知损失的鉴别器并重置生成器。

＃Finetune StyleGAN

＃alpha controls the strength of the style

alpha ＝ 1．0 ＃ min：0， max：1， step：0．1

alpha ＝ 1－alpha

＃preserve color of original image by limiting family of allowable transformations

preserve＿color ＝ False

＃Number of finetuning steps．

num＿iter ＝ 300

＃Log training on wandb and interval for image logging

use＿wandb ＝ False

log＿interval ＝ 50

if use＿wandb：

wandb．init（project＝＂JoJoGAN＂）

config ＝ wandb．config

config．num＿iter ＝ num＿iter

config．preserve＿color ＝ preserve＿color

wandb．log（

｛＂Style reference＂： ［wandb．Image（transforms．ToPILImage（）（target＿im））］｝，

step＝0）

＃ load discriminator for perceptual loss

discriminator ＝ Discriminator（1024， 2）．eval（）．to（device）

ckpt ＝ torch．load（＇models／stylegan2－ffhq－config－f．pt＇， map＿location＝lambda storage， loc： storage）

discriminator．load＿state＿dict（ckpt［＂d＂］， strict＝False）

＃ reset generator

del generator

generator ＝ deepcopy（original＿generator）

g＿optim ＝ optim．Adam（generator．parameters（）， lr＝2e－3， betas＝（0， 0．99））

训练生成器从潜在空间生成图像，并优化损失。

if preserve＿color：

id＿swap ＝ ［9，11，15，16，17］

z ＝ range（numiter）

for idx in tqdm（ z）：

mean＿w ＝ generator．get＿latent（torch．randn（［latents．size（0）， latent＿dim］）．to（device））．unsqueeze（1）．repeat（1， generator．n＿latent， 1）

in＿latent ＝ latents．clone（）

in＿latent［：， id＿swap］ ＝ alpha＊latents［：， id＿swap］ ＋ （1－alpha＊mean＿w［：， id＿swap］

img ＝ generator（in＿latent， input＿is＿latent＝True）

with torch．no＿grad（）：

real＿feat ＝ discriminator（targets）

fake＿feat ＝ discriminator（img）

loss ＝ sum（［functional．l1＿loss（a， b） for a， b in zip（fake＿feat， real＿feat）］）／len（fake＿feat）

if use＿wandb：

wandb．log（｛＂loss＂： loss｝， step＝idx）

if idx ％ log＿interval ＝＝ 0：

generator．eval（）

my＿sample ＝ generator（my＿w， input＿is＿latent＝True）

generator．train（）

wandb．log（

｛＂Current stylization＂： ［wandb．Image（my＿sample）］｝，

step＝idx）

g＿optim．zero＿grad（）

loss．backward（）

g＿optim．step（）

使用 JojoGAN 生成结果

现在生成结果。下面已经为原始图像和示例图像生成了结果以进行比较。

＃Generate resultsn＿sample ＝ 5

seed ＝ 3000

torch．manual＿seed（seed）

with torch．no＿grad（）：

generator．eval（）

z ＝ torch．randn（n＿sample， latent＿dim， device＝device）

original＿sample ＝ original＿generator（［z］， truncation＝0．7， truncation＿latent＝mean＿latent）

sample ＝ generator（［z］， truncation＝0．7， truncation＿latent＝mean＿latent）

original＿my＿sample ＝ original＿generator（my＿w， input＿is＿latent＝True）

my＿sample ＝ generator（my＿w， input＿is＿latent＝True）

＃ display reference images

style＿images ＝ ［］

for name in names：

style＿path ＝ f＇style＿images＿aligned／｛strip＿path＿extension（name）｝．png＇

style＿image ＝ transform（Image．open（style＿path））

style＿images．append（style＿image）

face ＝ transform（aligned＿face）．to（device）．unsqueeze（0）

style＿images ＝ torch．stack（style＿images， 0）．to（device）

my＿output ＝ torch．cat（［face， my＿sample］， 0）

output ＝ torch．cat（［original＿sample， sample］， 0）

生成的结果

现在，你可以使用 JojoGAN 生成你自己风格的图像。结果令人印象深刻，但可以通过调整训练方法和训练图像中的更多特征来进一步改进。

结论

JojoGAN 能够以快速有效的方式准确地映射和迁移用户定义的样式。关键要点是：

· JojoGAN 可以只用一种风格进行训练，以轻松映射并创建任何面部的风格化图

· JojoGAN 非常快速有效，可以在不到一分钟的时间内完成训练

· 结果非常准确，类似于逼真的肖像

· JojoGAN 可以轻松微调和修改，使其适用于 AI 应用程序

因此，无论风格类型、形状和颜色如何，JojoGAN 都是用于风格转移的理想神经网络，因此可以成为各种社交媒体应用程序和 AI 应用程序中非常有用的功能。

原文标题:使用JojoGAN创建风格化的面部图