本文介绍 2020_01_10_10_StyleGAN2人脸属性编辑-破解FaceAPP

2020_01_10_10_StyleGAN2人脸属性编辑-破解FaceAPP

This article was original written by Jin Tian, welcome re-post, first come with https://jinfagang.github.io . but please keep this copyright info, thanks, any question could be asked via wechat: jintianiloveu

许久没有更新专栏了，最近计算机视觉圈子也是发生了很多事情，从我们上次写一些realtime的实例分割算法的部署（比如EmbedMask, 链接如下）到现在，又发生了不少的事情，最近StyleGAN2倒是一个很有趣的事情，它比上一代生成的人脸更加逼真，同时训练的速度更快，即便如此，就笔者折腾下来，你如果没有八卡的机器，1024尺寸的模型根本别想训练了。

EmbedMask

但有趣的是，我们发现，既然我们没有那个财力去训练一个超大规模的StyleGAN2，那么我们为什么不直接用训练好的模型呢？关键是这个东西可以用来做什么？这便是本文的主题，传授大家的独门绝技：人脸属性编辑。先来看看效果：

是不是很眼熟？对，这就是FaceAPP的效果，实际上有一个网站叫做seeprettyface.com 已经有一些这样东西，但毕竟我们和这些网站不一样，我们是传授真实技能的，不是做一些中看不中用的东西。看完本篇文章你应该可以学会：

至少知道人脸属性编辑是如何实现的；
学会自己制作一个人脸属性编辑器，当然我们也会提供代码。

闲话不多说，直接开始吧！

##　StyleGAN2及其引出的一切

首先大家思考一下，给你一张图片，让你把这张图片变成笑脸或者改变年龄，你会怎么做？很多人说，用StyleGAN! 这几乎就是不懂者的回答了，实际上stylegan并不能实现这个功能，它只是给定一个随机向量，比如

a = np.random.randn([1, 512])

一个512维度的向量，然后stylegan将从这个向量重构为一张从未见过的人脸照片。像这样：

但其实只是这样玩，并没有太大的意思。有人就相处了一个更加神奇的想法：

我如果训练一个模型，这个模型是从图片到512维的latent encoding，这个latent encoding可以通过这个stylegan还原成原图，那岂不事可以把信息编码为512维度的空间！

听起来事一个不错的想法。还真的有人这么做了，比如下面这个代码大概就是做这么一件事情：

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146



# coding: utf-8

from models.stylegan_generator import StyleGANGenerator
from models.latent_optimizer import PostSynthesisProcessing
from models.image_to_latent import ImageToLatent, ImageLatentDataset
from models.losses import LogCoshLoss
from torchvision import transforms
import matplotlib.pyplot as plt
import torch
from glob import glob
from tqdm import tqdm_notebook as tqdm
import numpy as np


# # Create Dataloaders

# Using a 50,000 image dataset. Generated with the generated_data.py script at https://github.com/ShenYujun/InterFaceGAN.



augments = transforms.Compose([
    transforms.Resize(256),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

image_size = 256

directory = "../StyleGan/InterFaceGAN/test15/"
filenames = sorted(glob(directory + "*.jpg"))

train_filenames = filenames[0:48000]
validation_filenames = filenames[48000:]

dlatents = np.load(directory + "wp.npy")

train_dlatents = dlatents[0:48000]
validation_dlatents = dlatents[48000:]

train_dataset = ImageLatentDataset(train_filenames, train_dlatents, transforms=augments)
validation_dataset = ImageLatentDataset(validation_filenames, validation_dlatents, transforms=augments)

train_generator = torch.utils.data.DataLoader(train_dataset, batch_size=32)
validation_generator = torch.utils.data.DataLoader(validation_dataset, batch_size=32)



image_to_latent = ImageToLatent(image_size).cuda()
optimizer = torch.optim.Adam(image_to_latent.parameters())
criterion = LogCoshLoss()



epochs = 20
validation_loss = 0.0

progress_bar = tqdm(range(epochs))
for epoch in progress_bar:    
    running_loss = 0.0
    
    image_to_latent.train()
    for i, (images, latents) in enumerate(train_generator, 1):
        optimizer.zero_grad()

        images, latents = images.cuda(), latents.cuda()
        pred_latents = image_to_latent(images)
        loss = criterion(pred_latents, latents)
        loss.backward()
        
        optimizer.step()
        
        running_loss += loss.item()
        progress_bar.set_description("Step: {0}, Loss: {1:4f}, Validation Loss: {2:4f}".format(i, running_loss / i, validation_loss))
    
    validation_loss = 0.0
    
    image_to_latent.eval()
    for i, (images, latents) in enumerate(validation_generator, 1):
        with torch.no_grad():
            images, latents = images.cuda(), latents.cuda()
            pred_latents = image_to_latent(images)
            loss =  criterion(pred_latents, latents)
            
            validation_loss += loss.item()
    
    validation_loss /= i
    progress_bar.set_description("Step: {0}, Loss: {1:4f}, Validation Loss: {2:4f}".format(i, running_loss / i, validation_loss))


torch.save(image_to_latent.state_dict(), "./image_to_latent.pt")


image_to_latent = ImageToLatent(image_size).cuda()
image_to_latent.load_state_dict(torch.load("image_to_latent.pt"))
image_to_latent.eval()


# # Test Model


def normalized_to_normal_image(image):
    mean=torch.tensor([0.485, 0.456, 0.406]).view(-1,1,1).float()
    std=torch.tensor([0.229, 0.224, 0.225]).view(-1,1,1).float()
    
    image = image.detach().cpu()
    
    image *= std
    image += mean
    image *= 255
    
    image = image.numpy()[0]
    image = np.transpose(image, (1,2,0))
    return image.astype(np.uint8)


num_test_images = 5
images = [validation_dataset[i][0].unsqueeze(0).cuda() for i in range(num_test_images)]
normal_images = list(map(normalized_to_normal_image, images))

pred_dlatents = map(image_to_latent, images)

synthesizer = StyleGANGenerator("stylegan_ffhq").model.synthesis
post_processing = PostSynthesisProcessing()
post_process = lambda image: post_processing(image).detach().cpu().numpy().astype(np.uint8)[0]

pred_images = map(synthesizer, pred_dlatents)
pred_images = map(post_process, pred_images)
pred_images = list(map(lambda image: np.transpose(image, (1,2,0)), pred_images))


figure = plt.figure(figsize=(25,10))
columns = len(normal_images)
rows = 2

axis = []

for i in range(columns):
    axis.append(figure.add_subplot(rows, columns, i + 1))
    axis[-1].set_title("Reference Image")
    plt.imshow(normal_images[i])

for i in range(columns, columns*rows):
    axis.append(figure.add_subplot(rows, columns, i + 1))
    axis[-1].set_title("Generated With Predicted Latents")
    plt.imshow(pred_images[i - columns])

plt.show()