[代码实践]styleGAN2扩展：从真实人脸中提取图像的lat

写在前面的话

此文大部分为整理，仅作记录使用，获取信息的来源已经全部附上链接。（希望以后也能有自己的思想输出）

这位博主在他的博客中提到了多种重建真实图片的方法，包括：

• styleGAN2官网的run_projector.py
• rolux的project_images.py
• rolux基于Puzer的encode_images.py
• Pbaylies基于一代styleGAN的Encoder

（这个博主的系列文章让我学到很多）

此外，还有一些论文提到了重建的思路，都比较新，包括：

• StyleGAN2 Distillation for Feed-forward Image Manipulation(styleGan当老师，用它的知识训练某个特征的分类器，得到latent code的差值，以后的latent code根据差值变换可得到同样元素不同特征的图片，https://github.com/EvgenyKashin/stylegan2-distillation）
• InterFaceGAN: Interpreting the Disentangled Face Representation Learned by GANs（Interpreting the latent space of gans fot semantic face editing在隐空间中编码不同语义）
• PULSE: Self-Supervised Photo Upsampling via Latent Space Exploration of Generative Models（低分辨率生成高分辨率，完全自监督，但是部分情况下和真人不怎么像，https://github.com/adamian98/pulse）
• Adversarial Latent Autoencoders（StyleALAE不仅可以生成质量与StyleGAN相当的1024x1024人脸图像，而且在相同的分辨率下还可以基于真实图像生成人脸重建和操作。生成的还行，重建的效果还有待提高，https://github.com/podgorskiy/ALAE，有工具可以调节）
• Image2StyleGAN++: How to Edit the Embedded Images（基于Image2Style的改进，同时更新latent和噪声，有修复功能）
• MSG-GAN: Multi-Scale Gradient GAN for Stable Image Synthesis（基于progan, stylegan的高质量图像生成，靠trick，应该是SOTA， 这个开源了，还挺感兴趣的）

综合工作压力和兴趣，我期望先将pbaylies的思路先在StyleGAN2上实现。

原理和原始代码

pbaylies基于《Precise Recovery of Latent Vectors from Generative Adversarial Networks》实现过stylegan-encoder,原理是：通过stochastic clipping，将learnable latent code限制在某个区间内，通过计算该latent code生成图片和原始图片间的损失函数迭代更新learnable latent code，即为该原始图片对应的latent。
代码思路如下图所示：

image
image
bpaylies的实现逻辑如上图所示，从输入一张image开始看，使用resnet50获取latent，最后使用多种loss计算损失函数从而实现迭代。

让代码先跑起来

搭建环境

我使用conda搭建虚拟环境
建议python=3.6, tensorflow=1.14, keras=2.3

下载styleGAN2源码

https://github.com/NVlabs/stylegan2

Pbayies的代码需要用到第一版stylegan中的dnnlib库，该库在第二版中被删去了部分功能，因此需要下载第一版的dnnlib替换掉第二版的。
（下文中的下载链接均来自闪闪·Style）

链接: https://pan.baidu.com/s/1j6O-bgrMn5jVFO_GrE4cew
提取码: wjya

下载stylegan-encoder源码

https://github.com/pbaylies/stylegan-encoder

下载其中的encoder和ffhq_dataset文件夹，将它们移动到stylegan2根目录下，encoder文件夹重命名为encoder_s1

下载预训练模型

下载perceptual model “vgg16_zhang_perceptual.pkl”，可以从百度网盘下载：
https://pan.baidu.com/s/1vP6NM9-w4s3Cy6l4T7QpbQ
提取码: 5qkp

下载预训练的StyleGAN2人脸模型“stylegan2-ffhq-config-f.pkl”:
百度网盘: https://pan.baidu.com/s/1_cRyamHP_Amj0srCbiB5_g
提取码: cnby

在stylegan2根目录下新建models文件夹，将下载好的预训练模型放进去。

微调resnet反向网络

这部分参考了两个链接：链接1，链接2
在stylegan2根目录下新建文件夹data，用来存储微调后的resnet50网络
在stylegan2根目录下新建文件train_encoder.py用来finetune反向网络resnet50

import os
import numpy as np
import cv2

from keras.applications.imagenet_utils import preprocess_input
from keras.layers import Dense, Reshape
from keras.models import Sequential, Model, load_model
from keras.applications.resnet50 import ResNet50
from keras.optimizers import Adam

import pretrained_networks
import dnnlib.tflib as tflib


def get_batch(batch_size, Gs, image_size=224, Gs_minibatch_size=12, w_mix=None, latent_size=18):
    """
    Generate a batch of size n for the model to train
    returns a tuple (W, X) with W.shape = [batch_size, latent_size, 512] and X.shape = [batch_size, image_size, image_size, 3]
    If w_mix is not None, W = w_mix * W0 + (1 - w_mix) * W1 with
        - W0 generated from Z0 such that W0[:,i] = constant
        - W1 generated from Z1 such that W1[:,i] != constant

    Parametersget_batch
    ----------
    batch_size : int
        batch size
    Gs
        StyleGan2 generator
    image_size : int
    Gs_minibatch_size : int
        batch size for the generator
    w_mix : float

    Returns
    -------
    tuple
        dlatent W, images X
    """

    # Generate W0 from Z0
    Z0 = np.random.randn(batch_size, Gs.input_shape[1])
    W0 = Gs.components.mapping.run(Z0, None, minibatch_size=Gs_minibatch_size)

    if w_mix is None:
        W = W0
    else:
        # Generate W1 from Z1
        Z1 = np.random.randn(latent_size * batch_size, Gs.input_shape[1])
        W1 = Gs.components.mapping.run(Z1, None, minibatch_size=Gs_minibatch_size)
        W1 = np.array([W1[batch_size * i:batch_size * (i + 1), i] for i in range(latent_size)]).transpose((1, 0, 2))

        # Mix styles between W0 and W1
        W = w_mix * W0 + (1 - w_mix) * W1

    # Generate X
    X = Gs.components.synthesis.run(W, randomize_noise=True, minibatch_size=Gs_minibatch_size, print_progress=True,
                                    output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True))

    # Preprocess images X for the Imagenet model
    X = np.array([cv2.resize(x, (image_size, image_size)) for x in X])
    X = preprocess_input(X.astype('float'))

    return W, X


def finetune(save_path, image_size=224, base_model=ResNet50, batch_size=2048, test_size=1024, n_epochs=6,
             max_patience=5, models_dir='models/stylegan2-ffhq-config-f.pkl'):
    """
    Finetunes a ResNet50 to predict W[:, 0]

    Parameters
    ----------
    save_path : str
        path where to save the Resnet
    image_size : int
    base_model : keras model
    batch_size :  int
    test_size : int
    n_epochs : int
    max_patience : int

    Returns
    -------
    None

    """

    assert image_size >= 224

    # Load StyleGan generator
    _, _, Gs = pretrained_networks.load_networks(models_dir)

    # Build model
    if os.path.exists(save_path):
        print('Loading pretrained network')
        model = load_model(save_path, compile=False)
    else:
        base = base_model(include_top=False, pooling='avg', input_shape=(image_size, image_size, 3))
        model = Sequential()
        model.add(base)
        model.add(Dense(512))

    model.compile(loss='mse', metrics=[], optimizer=Adam(3e-4))
    model.summary()

    # Create a test set
    print('Creating test set')
    W_test, X_test = get_batch(test_size, Gs)

    # Iterate on batches of size batch_size
    print('Training model')
    patience = 0
    best_loss = np.inf

    while (patience <= max_patience):
        W_train, X_train = get_batch(batch_size, Gs)
        model.fit(X_train, W_train[:, 0], epochs=n_epochs, verbose=True)
        loss = model.evaluate(X_test, W_test[:, 0])
        if loss < best_loss:
            print(f'New best test loss : {loss:.5f}')
            model.save(save_path)
            patience = 0
            best_loss = loss
        else:
            print(f'-------- test loss : {loss:.5f}')
            patience += 1


def finetune_18(save_path, base_model=None, image_size=224, batch_size=2048, test_size=1024, n_epochs=6,
                max_patience=8, w_mix=0.7, latent_size=18, models_dir='models/stylegan2-ffhq-config-f.pkl'):
    """
    Finetunes a ResNet50 to predict W[:, :]

    Parameters
    ----------
    save_path : str
        path where to save the Resnet
    image_size : int
    base_model : str
        path to the first finetuned ResNet50
    batch_size :  int
    test_size : int
    n_epochs : int
    max_patience : int
    w_mix : float

    Returns
    -------
    None

    """

    assert image_size >= 224
    if not os.path.exists(save_path):
        assert base_model is not None

    # Load StyleGan generator
    _, _, Gs = pretrained_networks.load_networks(models_dir)

    # Build model
    if os.path.exists(save_path):
        print('Loading pretrained network')
        model = load_model(save_path, compile=False)
    else:
        base_model = load_model(base_model)
        hidden = Dense(latent_size * 512)(base_model.layers[-1].input)
        outputs = Reshape((latent_size, 512))(hidden)
        model = Model(base_model.input, outputs)
        # Set initialize layer
        W, b = base_model.layers[-1].get_weights()
        model.layers[-2].set_weights([np.hstack([W] * latent_size), np.hstack([b] * latent_size)])

    model.compile(loss='mse', metrics=[], optimizer=Adam(1e-4))
    model.summary()

    # Create a test set
    print('Creating test set')
    W_test, X_test = get_batch(test_size, Gs, w_mix=w_mix, latent_size=latent_size)

    # Iterate on batches of size batch_size
    print('Training model')
    patience = 0
    best_loss = np.inf

    while (patience <= max_patience):
        W_train, X_train = get_batch(batch_size, Gs, w_mix=w_mix, latent_size=latent_size)
        model.fit(X_train, W_train, epochs=n_epochs, verbose=True)
        loss = model.evaluate(X_test, W_test)
        if loss < best_loss:
            print(f'New best test loss : {loss:.5f}')
            model.save(save_path)
            patience = 0
            best_loss = loss
        else:
            print(f'-------- test loss : {loss:.5f}')
            patience += 1


if __name__ == '__main__':
    finetune('data/resnet.h5')
    finetune_18('data/resnet_18.h5', 'data/resnet.h5', w_mix=0.8)
    #finetune('data/resnet.h5', n_epochs=2, max_patience=1)
    #finetune_18('data/resnet_18.h5', 'data/resnet.h5', w_mix=0.8, n_epochs=2, max_patience=1)

通过执行python train_encoder.py先将resnet跑起来。在跑起来的时候可能会遇到以下问题。

1. tensorflow error This file requires compiler and library support for the ISO C++ 2011 standard

基于stackoverflow,做如下修改：
将dnnlib/tflib/custom_pos.py的第64行修正为cmd = 'nvcc --std=c++11 -DNDEBUG ' + opts.strip()

2. undefined symbol: _ZN10tensorflow12OpDefBuilder6OutputESs

基于CSDN博客，做如下修改：
将dnnlib/tflib/custom_jpos.py的第127行修正为compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=1

新增encode_image_s1.py

在stylegan2根目录下新增文件encode_image_s1.py，相比原博主，更新了一些参数的定义，内容如下：

import os
import argparse
import pickle

from tqdm import tqdm
import PIL.Image
import numpy as np
import dnnlib
import dnnlib.tflib as tflib
from encoder_s1.generator_model import Generator
from encoder_s1.perceptual_model import PerceptualModel, load_images
from keras.models import load_model

import glob
import random

def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def split_to_batches(l, n):
    for i in range(0, len(l), n):
        yield l[i:i + n]


def main():
    parser = argparse.ArgumentParser(
        description='Find latent representation of reference images using perceptual losses',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('src_dir', help='Directory with images for encoding')
    parser.add_argument('generated_images_dir', help='Directory for storing generated images')
    parser.add_argument('dlatent_dir', help='Directory for storing dlatent representations')
    parser.add_argument('--data_dir', default='data', help='Directory for storing optional models')
    parser.add_argument('--mask_dir', default='masks', help='Directory for storing optional masks')
    parser.add_argument('--load_last', default='', help='Start with embeddings from directory')
    parser.add_argument('--dlatent_avg', default='',
                        help='Use dlatent from file specified here for truncation instead of dlatent_avg from Gs')
    parser.add_argument('--model_url', default='models/stylegan2-ffhq-config-f.pkl',
                        help='Fetch a StyleGAN model to train on from this URL')
    parser.add_argument('--model_res', default=1024, help='The dimension of images in the StyleGAN model', type=int)
    parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int)

    # Perceptual model params
    parser.add_argument('--image_size', default=256, help='Size of images for perceptual model', type=int)
    parser.add_argument('--sharpen_input', default=True, help='whether to add sharpen action for input images', type=bool)
    parser.add_argument('--resnet_image_size', default=224, help='Size of images for the Resnet model', type=int)
    parser.add_argument('--lr', default=0.02, help='Learning rate for perceptual model', type=float)
    parser.add_argument('--decay_rate', default=0.9, help='Decay rate for learning rate', type=float)
    parser.add_argument('--iterations', default=100, help='Number of optimization steps for each batch', type=int)
    parser.add_argument('--decay_steps', default=10,
                        help='Decay steps for learning rate decay (as a percent of iterations)', type=float)
    parser.add_argument('--load_effnet', default='data/finetuned_effnet.h5',
                        help='Model to load for EfficientNet approximation of dlatents')
    parser.add_argument('--load_resnet', default='data/resnet_18.h5',
                        help='Model to load for ResNet approximation of dlatents')

    # Loss function options
    parser.add_argument('--use_vgg_loss', default=0.4, help='Use VGG perceptual loss; 0 to disable, > 0 to scale.',
                        type=float)
    parser.add_argument('--use_adaptive_loss', default=False,
                        help='Use the adaptive robust loss function from Google Research for pixel and VGG feature loss.',
                        type=str2bool, nargs='?', const=True)
    parser.add_argument('--use_vgg_layer', default=9, help='Pick which VGG layer to use.', type=int)
    parser.add_argument('--use_pixel_loss', default=1.5,
                        help='Use logcosh image pixel loss; 0 to disable, > 0 to scale.', type=float)
    parser.add_argument('--use_mssim_loss', default=100, help='Use MS-SIM perceptual loss; 0 to disable, > 0 to scale.',
                        type=float)
    parser.add_argument('--use_lpips_loss', default=100, help='Use LPIPS perceptual loss; 0 to disable, > 0 to scale.',
                        type=float)
    parser.add_argument('--use_l1_penalty', default=1, help='Use L1 penalty on latents; 0 to disable, > 0 to scale.',
                        type=float)
    parser.add_argument('--use_discriminator_loss', default=0.5, help='Use trained discriminator to evaluate realism.',
                        type=float)

    # Generator params
    parser.add_argument('--randomize_noise', default=False, help='Add noise to dlatents during optimization', type=bool)
    parser.add_argument('--tile_dlatents', default=False, help='Tile dlatents to use a single vector at each scale',
                        type=bool)
    parser.add_argument('--clipping_threshold', default=2.0,
                        help='Stochastic clipping of gradient values outside of this threshold', type=float)

    # Masking params
    parser.add_argument('--load_mask', default=False, help='Load segmentation masks', type=bool)
    parser.add_argument('--face_mask', default=False, help='Generate a mask for predicting only the face area',
                        type=bool)
    parser.add_argument('--use_grabcut', default=True,
                        help='Use grabcut algorithm on the face mask to better segment the foreground', type=bool)
    parser.add_argument('--scale_mask', default=1.5, help='Look over a wider section of foreground for grabcut',
                        type=float)

    # Video params
    parser.add_argument('--video_dir', default='videos', help='Directory for storing training videos')
    parser.add_argument('--output_video', default=False, help='Generate videos of the optimization process', type=bool)
    parser.add_argument('--video_codec', default='MJPG', help='FOURCC-supported video codec name')
    parser.add_argument('--video_frame_rate', default=24, help='Video frames per second', type=int)
    parser.add_argument('--video_size', default=512, help='Video size in pixels', type=int)
    parser.add_argument('--video_skip', default=1, help='Only write every n frames (1 = write every frame)', type=int)

    # 获取到基本设置时，如果运行命令中传入了之后才会获取到的其他配置，不会报错；而是将多出来的部分保存起来，留到后面使用
    args, other_args = parser.parse_known_args()

    # learning rate衰减的steps
    args.decay_steps *= 0.01 * args.iterations  # Calculate steps as a percent of total iterations

    if args.output_video:
        import cv2
        synthesis_kwargs = dict(output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=False),
                                minibatch_size=args.batch_size)

    # 找到src_dir下所有图片文件，加入ref_images列表（即：源图的列表；只有一个图片也可以）
    ref_images = [os.path.join(args.src_dir, x) for x in os.listdir(args.src_dir)]
    ref_images = list(filter(os.path.isfile, ref_images))

    if len(ref_images) == 0:
        raise Exception('%s is empty' % args.src_dir)

    # 创建工作目录
    os.makedirs(args.data_dir, exist_ok=True)
    os.makedirs(args.mask_dir, exist_ok=True)
    os.makedirs(args.generated_images_dir, exist_ok=True)
    os.makedirs(args.dlatent_dir, exist_ok=True)
    os.makedirs(args.video_dir, exist_ok=True)

    # Initialize generator and perceptual model
    tflib.init_tf()
    # 加载StyleGAN模型
    model_file = glob.glob(args.model_url)
    if len(model_file) == 1:
        model_file = open(model_file[0], "rb")
    else:
        raise Exception('Failed to find the model')
    generator_network, discriminator_network, Gs_network = pickle.load(model_file)

    # 加载Generator类，参与构建VGG16 perceptual model，用于调用（说是生成，更好理解）generated_image
    # generated_image通过perceptual_model转化为generated_img_features，参与计算loss
    generator = Generator(Gs_network, args.batch_size, clipping_threshold=args.clipping_threshold,
                          tiled_dlatent=args.tile_dlatents, model_res=args.model_res,
                          randomize_noise=args.randomize_noise)
    if (args.dlatent_avg != ''):
        generator.set_dlatent_avg(np.load(args.dlatent_avg))

    perc_model = None
    if (args.use_lpips_loss > 0.00000001):  # '--use_lpips_loss', default = 100
        # 加载VGG16 perceptual模型
        model_file = glob.glob('./models/vgg16_zhang_perceptual.pkl')
        if len(model_file) == 1:
            model_file = open(model_file[0], "rb")
        else:
            raise Exception('Failed to find the model')
        perc_model = pickle.load(model_file)

    # 创建VGG16 perceptual模型
    perceptual_model = PerceptualModel(args, perc_model=perc_model, batch_size=args.batch_size)
    perceptual_model.build_perceptual_model(generator, discriminator_network)

    ff_model = None
    # Optimize (only) dlatents by minimizing perceptual loss between reference and generated images in feature space
    # tqdm 是一个快速，可扩展的Python进度条，可以在 Python 长循环中添加一个进度提示信息
    # 把ref_images分割为若干批次，每个批次的大小为args.batch_size，分批使用perceptual_model.optimize()求解每个源图的dlatents的最优解
    # 对每一个源图，优化迭代的过程是从一个初始dlatents开始，在某个空间内，按正态分布取值，使用Adam优化器，逐步寻找使loss最小的dlatents，即：stochastic clipping方法
    for images_batch in tqdm(split_to_batches(ref_images, args.batch_size), total=len(ref_images) // args.batch_size):
        # 读取每个批次中的文件名
        names = [os.path.splitext(os.path.basename(x))[0] for x in images_batch]
        if args.output_video:
            video_out = {}
            for name in names:
                video_out[name] = cv2.VideoWriter(os.path.join(args.video_dir, f'{name}.avi'),
                                                  cv2.VideoWriter_fourcc(*args.video_codec), args.video_frame_rate,
                                                  (args.video_size, args.video_size))

        # 给源图及源图用VGG16生成的features赋值(这是计算loss的基准)
        perceptual_model.set_reference_images(images_batch)
        dlatents = None
        if (args.load_last != ''):  # load previous dlatents for initialization
            for name in names:
                dl = np.expand_dims(np.load(os.path.join(args.load_last, f'{name}.npy')), axis=0)
                if (dlatents is None):
                    dlatents = dl
                else:
                    dlatents = np.vstack((dlatents, dl))
        else:
            if (ff_model is None):
                if os.path.exists(args.load_resnet):
                    print("Loading ResNet Model:")
                    ff_model = load_model(args.load_resnet)
                    from keras.applications.resnet50 import preprocess_input
            if (ff_model is None):
                if os.path.exists(args.load_effnet):
                    import efficientnet
                    print("Loading EfficientNet Model:")
                    ff_model = load_model(args.load_effnet)
                    from efficientnet import preprocess_input
            if (ff_model is not None):  # predict initial dlatents with ResNet model
                dlatents = ff_model.predict(
                    preprocess_input(load_images(images_batch, image_size=args.resnet_image_size)))
        # 设置用于perceptual_model优化迭代的初始值dlatents，它是用resnet50或者efficientnet从源图预测得到的
        if dlatents is not None:
            generator.set_dlatents(dlatents)
        # 对每一个源图，用tqdm构造进度条，显示优化迭代的过程
        op = perceptual_model.optimize(generator.dlatent_variable, iterations=args.iterations)
        pbar = tqdm(op, leave=False, total=args.iterations)
        vid_count = 0
        best_loss = None
        best_dlatent = None
        # 用stochastic clipping方法，使用VGG16 perceptual_model进行优化迭代，迭代次数为iterations=args.iterations
        for loss_dict in pbar:
            pbar.set_description(" ".join(names) + ": " + "; ".join(["{} {:.4f}".format(k, v)
                                                                     for k, v in loss_dict.items()]))
            if best_loss is None or loss_dict["loss"] < best_loss:
                best_loss = loss_dict["loss"]
                best_dlatent = generator.get_dlatents()
            if args.output_video and (vid_count % args.video_skip == 0):
                batch_frames = generator.generate_images()
                for i, name in enumerate(names):
                    video_frame = PIL.Image.fromarray(batch_frames[i], 'RGB').resize((args.video_size, args.video_size),
                                                                                     PIL.Image.LANCZOS)
                    video_out[name].write(cv2.cvtColor(np.array(video_frame).astype('uint8'), cv2.COLOR_RGB2BGR))
            # 用stochastic clip方法更新dlatent_variable
            generator.stochastic_clip_dlatents()
        print(" ".join(names), " Loss {:.4f}".format(best_loss))

        if args.output_video:
            for name in names:
                video_out[name].release()

        # Generate images from found dlatents and save them
        generator.set_dlatents(best_dlatent)
        generated_images = generator.generate_images()
        generated_dlatents = generator.get_dlatents()
        for img_array, dlatent, img_name in zip(generated_images, generated_dlatents, names):
            img = PIL.Image.fromarray(img_array, 'RGB')
            img.save(os.path.join(args.generated_images_dir, f'{img_name}.png'), 'PNG')
            np.save(os.path.join(args.dlatent_dir, f'{img_name}.npy'), dlatent)

        generator.reset_dlatents()


if __name__ == "__main__":
    main()

处理数据

我使用了style-encoder中的align_images.py处理图像，将原始图片对齐，裁剪

conda create -n tf1.14 python=3.6
source activate tf1.14
conda install tensorflow-gpu=1.14 keras=2.3
conda install -c menpo dlib # 用来对齐数据集，我的环境问题主要就出在dlib安装
conda install pillow
pip install opencv-python
conda install requests tqdm

python align_images.py images/input_images_origin/ images/input_images_align

生成latent code

python encode_images_s1.py images/raw_images/ images/generated_images/ images/latent_representations/ --load_resnet data/resnet_18.h5 --batch_size 1 --iterations 100

生成的图片在images/generated_images/ 可以看到。

生成效果

代码虽然跑通了，但训练得还不好
左图：生成图像，右图：原始图像

image.png

参考

https://blog.csdn.net/weixin_41943311/article/details/103030194
https://blog.csdn.net/DLW__/article/details/104528609
https://github.com/NVlabs/stylegan2