使用StyleGAN + CLIP从文本生成面部图像[翻译]

资料仓库  收藏
0 / 611

使用StyleGAN + CLIP从文本生成面部图像

 这次介绍的是StyleClip: StyleGAN+CLIP,从文本进行图像生成的一个模型。是现在最好的图像生成模型和文字模型的结合的。

什么是StyleGAN + CLIP模型?

**  StyleGAN**是一个根据参数生成图像的模型,CLIP是可以输出图像和文本之间相似度的模型。这次,我们将按以下方式结合这两个模型。

 
为StyleGAN提供适当的参数并输出图像(目标图像)。然后,将图像文本输入到CLIP中,找到相似度,并优化参数,以使相似度尽可能高。。然后,获得生成与文本的内容匹配的图像的参数,并且可以从文本生成图像。

这篇文章有机翻 https://aiqianji.com/blog/article/85

 现在,让我们看一下代码实现。

代码

 该代码在Google Colab上运行并发布Github上,单击此“链接”

 现在,让我们从文本生成一张脸部图像。在文本中,输入要生成的内容的文本。在这里,我输入“她是一个有着金色的头发和蓝色的眼睛的迷人女人”。

 将优化循环旋转101次(max_iter设置),然后每两次旋转(img_save_freq设置),将生成的中间图像./pic保存在其中。

import os

import torch

import torchvision

import clip

import numpy as np

from PIL import Image

from stylegan_models import g_all, g_synthesis, g_mapping

from utils import GetFeatureMaps, transform_img, compute_loss

from tqdm import trange

import warnings  

warnings.filterwarnings('ignore')   

 


import os

import shutil

if os.path.isdir('pic'):

     shutil.rmtree('pic')

os.makedirs('pic', exist_ok=True)

 

# 初期設定 

lr = 1e-2 

img_save_freq = 2

max_iter = 101 

ref_img_path = None 

 

device = 'cuda' if torch.cuda.is_available() else 'cpu'

print("USING ", device)

 

clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)

vgg16 = torchvision.models.vgg16(pretrained=True).to(device)

vgg_layers = vgg16.features

 

vgg_layer_name_mapping = {

    '1': "relu1_1",

    '3': "relu1_2",

    '6': "relu2_1",

    '8': "relu2_2",

    # '15': "relu3_3",

    # '22': "relu4_3"

}

 

g_synthesis.eval()

g_synthesis.to(device)

 

latent_shape = (1, 1, 512)

 

normal_generator = torch.distributions.normal.Normal(

    torch.tensor([0.0]),

    torch.tensor([1.]),

)

 

# init_latents = normal_generator.sample(latent_shape).squeeze(-1).to(device)

latents_init = torch.zeros(latent_shape).squeeze(-1).to(device)

latents = torch.nn.Parameter(latents_init, requires_grad=True)

 

optimizer = torch.optim.Adam(

    params=[latents],

    lr=lr,

    betas=(0.9, 0.999),

)

 

def truncation(x, threshold=0.7, max_layer=8):

    avg_latent = torch.zeros(1, x.size(1), 512).to(device)

    interp = torch.lerp(avg_latent, x, threshold)

    do_trunc = (torch.arange(x.size(1)) < max_layer).view(1, -1, 1).to(device)

    return torch.where(do_trunc, interp, x)

 

def tensor_to_pil_img(img):

    img = (img.clamp(-1, 1) + 1) / 2.0

    img = img[0].permute(1, 2, 0).detach().cpu().numpy() * 256

    img = Image.fromarray(img.astype('uint8'))

    return img

 

 

clip_transform = torchvision.transforms.Compose([

    # clip_preprocess.transforms[2],

    clip_preprocess.transforms[4],

])

 

if ref_img_path is None:

    ref_img = None

else:

    ref_img = clip_preprocess(Image.open(ref_img_path)).unsqueeze(0).to(device)

 

clip_normalize = torchvision.transforms.Normalize(

    mean=(0.48145466, 0.4578275, 0.40821073),

    std=(0.26862954, 0.26130258, 0.27577711),

)

 

def compute_clip_loss(img, text):

    # img = clip_transform(img)

    img = torch.nn.functional.upsample_bilinear(img, (224, 224))

    tokenized_text = clip.tokenize([text]).to(device)

    img_logits, _text_logits = clip_model(img, tokenized_text)

    return 1/img_logits * 100

 

def compute_perceptual_loss(gen_img, ref_img):

    gen_img = torch.nn.functional.upsample_bilinear(img, (224, 224))

    loss = 0

    len_vgg_layer_mappings = int(max(vgg_layer_name_mapping.keys()))

    ref_feats = ref_img

    gen_feats = gen_img

 

    for idx, (name, module) in enumerate(vgg_layers._modules.items()):

        ref_feats = module(ref_feats)

        gen_feats = module(gen_feats)

        if name in vgg_layer_name_mapping.keys():

            loss += torch.nn.functional.mse_loss(ref_feats, gen_feats)

        

        if idx >= len_vgg_layer_mappings:

            break

    

    return loss/len_vgg_layer_mappings

 

counter = 0

for i in trange(max_iter):

    dlatents = latents.repeat(1,18,1)

    img = g_synthesis(dlatents)

    loss = compute_clip_loss(img, text) 

    optimizer.zero_grad()

    loss.backward()

    optimizer.step()

 

    if i % img_save_freq == 0:

        img = tensor_to_pil_img(img)

        img = img.resize((512,512))  

        img.save(os.path.join('./pic', str(counter).zfill(6)+'.png')) 

        counter +=1 

 

# 显示

from IPython.display import Image, display_png

display_png(Image('./pic/'+str(counter-1).zfill(6)+'.png'))

 

我认为我们能够生成与文本几乎匹配的图像。而且,该面部图像实际上不存在。在优化过程中会使用随机数,因此即使文本相同,每次生成的图像也会有所不同。

 现在,让我们观看展示优化进度的视频。./pic从生成的图像中将其保存到其中,然后创建一个mp4视频,并output.mp4text名.mp4文件名临时保存它,并以文件名./movie保存它。然后播放视频。

import os

if os.path.exists('./output.mp4'):

   os.remove('./output.mp4')

 
! ffmpeg -r 15 -i pic/%6d.png\

               -vcodec libx264 -pix_fmt yuv420p output.mp4
import shutil

os.makedirs('movie', exist_ok=True)

shutil.copy('output.mp4', 'movie/'+text+'.mp4')

 
您可以看到如何逐步优化初始图像的参数并将其转换为文本表示的图像。

 之后,如果从文本输入位置重复,则可以。这次,我将尝试“微笑的亚裔女子戴着眼镜。中等长度的黑发。”(微笑的亚裔女子戴着眼镜。中等长度的黑发)。

 

实际上,StyleGAN所学习的图像数据很少戴眼镜戴,所以我不太擅长生成眼镜,但是这次很顺利。

 让我们再做一个。现在让我们尝试“她是一个迷人的女孩,有着卷曲的黑发和黑色的眼睛”。

途中,它变成了男孩似的脸,但最终变成了男孩样的女孩的脸,并成为了文字的脸部图像。我很高兴。

 再见。

(原始Github:https : //github.com/vipermu/StyleCLIP

from link