StyleGAN2人脸属性编辑-破解FaceAPP

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

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维度的空间！

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

# 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()

训练完了这个模型之后，你便会思考一个事情，很拷问心灵的事情：

我既然拥有了一个从图片到512维空间的这么一个模型，那我能不能想办法，根据stylegan，迫使它生成一张和我给定的图片一模一样的图片呢？如果能做到这一点的话，我岂不是可以反过来用这个512维的向量去生成一模一样的图片了？

这便引申出了另外一个训练模型的过程：

训练一个模型，这个模型训练的数据集就是你给的图片，根据上面的模型，以及stylegan的生成结果，去求取一个latent encoding，是的这个latent coding，和你给定的原图生成的图片尽可能的相似。

神经网络学会的512超高维空间

接下来的操作有点神学，我们将在一个超高维的空间来操控我们编码之后的人脸。这个超高纬度如何理解？我们人类对于人脸的不同属性，比如笑，肤色，带眼镜，性别等都有一定的区分度。而对于神经网络来讲，它区分起来更加容易了，直接用512维的超平面来区分。

未来工作

本期分享便是教大家如何制作自己的人脸编辑器。百看不如一试，本文人脸编辑器的所有代码都可以在神力平台找到：
神力AI算法平台：http://manaai.cn/

感兴趣的同学可以下载试玩。未来我们将继续开放一些我们在stylegan2做的试验，我们在训练更多的不同人脸的生成器，比如网红脸生成器，一键给你生成海量人脸，大家试想一下，如果用网红脸生成器来做人脸属性编辑器回事一种什么样的体验呢？欢迎大家来社区平台发帖留言说出你的想法！