介紹

風格遷移是神經網絡的一個發(fā)展領域，它是一個非常有用的功能，可以集成到社交媒體和人工智能應用程序中。幾個神經網絡可以根據訓練數據將圖像樣式映射和傳輸到輸入圖像。在本文中，我們將研究 JojoGAN，以及僅使用一種參考樣式來訓練和生成具有該樣式的任何圖像的過程。

JoJoGAN：One Shot Face Stylization

One Shot Face Stylization（一次性面部風格化）可用于 AI 應用程序、社交媒體過濾器、有趣的應用程序和業(yè)務用例。隨著 AI 生成的圖像和視頻濾鏡的日益普及，以及它們在社交媒體和短視頻、圖像中的使用，一次性面部風格化是一個有用的功能，應用程序和社交媒體公司可以將其集成到最終產品中。

因此，讓我們來看看用于一次性生成人臉樣式的流行 GAN 架構——JojoGAN。

JojoGAN 架構

JojoGAN 是一種風格遷移程序，可讓將人臉圖像的風格遷移為另一種風格。它通過GAN將參考風格圖像反轉為近似的配對訓練數據，根據風格化代碼生成真實的人臉圖像，并與參考風格圖像相匹配。然后將該數據集用于微調 StyleGAN，并且可以使用新的輸入圖像，JojoGAN 將根據 GAN 反轉（inversion）將其轉換為該特定樣式。

JojoGAN 架構和工作流程

JojoGAN 只需一種參考風格即可在很短的時間內（不到 1 分鐘）進行訓練，并生成高質量的風格化圖像。

JojoGan 的一些例子

JojoGAN 生成的風格化圖像的一些示例：

風格化的圖像可以在各種不同的輸入風格上生成并且可以修改。

JojoGan 代碼深潛

讓我們看看 JojoGAN 生成風格化人像的實現。有幾個預訓練模型可用，它們可以在我們的風格圖像上進行訓練，或者可以修改模型以在幾分鐘內更改風格。

JojoGAN 的設置和導入

克隆 JojoGAN 存儲庫并導入必要的庫。在 Google Colab 存儲中創(chuàng)建一些文件夾，用于存儲反轉代碼、樣式圖像和模型。

！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 可用于預訓練模型。這些預訓練模型可用于隨時隨地生成風格化圖像，并具有不同的準確度。之后，可以訓練用戶創(chuàng)建的模型。

＃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＇）

加載生成器

加載原始和微調生成器。設置用于調整圖像大小和規(guī)范化圖像的 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”，看起來相當準確。

現在讓我們看一下在自創(chuàng)樣式上訓練 GAN。

使用你的風格圖像進行訓練

選擇一些面部圖，甚至創(chuàng)建一些自己的面部圖并加載這些圖像以訓練 GAN，并設置路徑。裁剪和對齊人臉并執(zhí)行 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））

訓練生成器從潛在空間生成圖像，并優(yōu)化損失。

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 可以只用一種風格進行訓練，以輕松映射并創(chuàng)建任何面部的風格化圖

· JojoGAN 非�？焖儆行�，可以在不到一分鐘的時間內完成訓練

· 結果非常準確，類似于逼真的肖像

· JojoGAN 可以輕松微調和修改，使其適用于 AI 應用程序

因此，無論風格類型、形狀和顏色如何，JojoGAN 都是用于風格轉移的理想神經網絡，因此可以成為各種社交媒體應用程序和 AI 應用程序中非常有用的功能。

       原文標題 : 使用JojoGAN創(chuàng)建風格化的面部圖