基于GAN的face2face实现

本文实现依赖python2.7、tensorflow、opencv、dlib，训练生成对抗模型(GAN)，实现图像合成。

face2face是image2image或被称为pix2pix众多有趣应用中的一个。
更多应用案例与原理论文，请参考
Image-to-Image Translation with Conditional Adversarial Nets
Image-to-Image Translation in Tensorflow by Christopher Hesse
Dat Tran博客 Face2face

• Step 1 利用opencv和dlib准备训练集
• Step 2 利用tensorflow训练模型
• Step 3 Export Model & Freeze Model
• Step 4 调用模型

step 1 准备训练集

1. 在当前目录创建original与landmark文件夹。每个文件夹包含400张含有人脸的图片。
   • 注意视频格式。训练数据为 默克尔演讲视频(网盘地址)，网盘的视频是MP4格式的，需要转换为avi格式才能顺利执行下面代码。你也可以找其他视频。但是视频中人脸最好在靠中间的位置，不然可能在图片变换尺寸的时候会被裁剪掉。
   • 需要加载人脸特征模型。shape_predictor_68_face_landmarks.dat(
     网盘地址)

# -*- coding: utf-8 -*-
from __future__ import division
import cv2
import dlib
import numpy as np
import os

os.makedirs('original') # 创建文件夹，用于保存原始视频中截取的帧
os.makedirs('landmarks') # 创建文件夹，用于保存描绘有人脸特征的图片
DOWNSAMPLE_RATIO = 4 # 图片缩小比例，小图片加快人脸检测与特征提取速度
photo_number = 400 # 从视频中提取400张含有人脸特征的帧
video_path = 'angela_merkel_speech.avi' # 用于训练的含有人脸的视频路径
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('shape_predictor_68_face_landmarks.dat')

def reshape_for_polyline(array):
    return np.array(array, np.int32).reshape((-1, 1, 2))

def prepare_training_data():
    cap = cv2.VideoCapture(video_path)
    count = 0
    while cap.isOpened():
        ret, frame = cap.read() # 读取视频帧
        frame_resize = cv2.resize(frame, (0,0), fx=1 / DOWNSAMPLE_RATIO, fy=1 / DOWNSAMPLE_RATIO)
        gray = cv2.cvtColor(frame_resize, cv2.COLOR_BGR2GRAY)
        faces = detector(gray, 1) # 识别人脸位置
        black_image = np.zeros(frame.shape, np.uint8) # 创建一张黑色图片用于描绘人脸特征

        if len(faces) == 1:
            for face in faces:
                detected_landmarks = predictor(gray, face).parts() # 提取人脸特征
                landmarks = [[p.x * DOWNSAMPLE_RATIO, p.y * DOWNSAMPLE_RATIO] for p in detected_landmarks]

                jaw = reshape_for_polyline(landmarks[0:17])
                left_eyebrow = reshape_for_polyline(landmarks[22:27])
                right_eyebrow = reshape_for_polyline(landmarks[17:22])
                nose_bridge = reshape_for_polyline(landmarks[27:31])
                lower_nose = reshape_for_polyline(landmarks[30:35])
                left_eye = reshape_for_polyline(landmarks[42:48])
                right_eye = reshape_for_polyline(landmarks[36:42])
                outer_lip = reshape_for_polyline(landmarks[48:60])
                inner_lip = reshape_for_polyline(landmarks[60:68])

                color = (255, 255, 255) # 人脸特征用白色描绘
                thickness = 3 # 线条粗细

                cv2.polylines(img=black_image, 
                              pts=[jaw,left_eyebrow, right_eyebrow, nose_bridge],
                              isClosed=False,
                              color=color,
                              thickness=thickness)
                cv2.polylines(img=black_image, 
                              pts=[lower_nose, left_eye, right_eye, outer_lip,inner_lip],
                              isClosed=True,
                              color=color,
                              thickness=thickness)

            # 保存图片
            cv2.imwrite("original2/{}.png".format(count), frame)
            cv2.imwrite("landmarks2/{}.png".format(count), black_image)
            count += 1

# 执行准备数据函数
prepare_training_data()

2. 改变图片尺寸(调整为正方形)、拼接图片(用于训练)
   这步涉及的函数有点多，主要是利用tensorflow对jpeg与png图片的读取、保存、裁剪、缩放、拼接，直接根据下面步骤执行就可以。不过建议对tensorflow图片处理细节感兴趣的小伙伴看源代码，会有很多收获。
   github repo affinelayer/pix2pix-tensorflow

# Clone the repo from Christopher Hesse's pix2pix TensorFlow implementation
git clone https://github.com/affinelayer/pix2pix-tensorflow.git

# Move the original and landmarks folder into the pix2pix-tensorflow folder
mv face2face-demo/landmarks face2face-demo/original pix2pix-tensorflow/photos

# Go into the pix2pix-tensorflow folder
cd pix2pix-tensorflow/

# Resize original images
python tools/process.py \
  --input_dir photos/original \
  --operation resize \
  --output_dir photos/original_resized
  
# Resize landmark images
python tools/process.py \
  --input_dir photos/landmarks \
  --operation resize \
  --output_dir photos/landmarks_resized
  
# Combine both resized original and landmark images
python tools/process.py \
  --input_dir photos/landmarks_resized \
  --b_dir photos/original_resized \
  --operation combine \
  --output_dir photos/combined
  
# Split into train/val set
python tools/split.py \
  --dir photos/combined

执行完上面的代码，模型的训练数据就已经准备就绪了。整个 process.py文件，基本是以下结构。我觉得这是值得一书的东西，以备不时之需。

import tensorflow as tf

# 创建一个万金油般的create_op函数
def create_op(func, **placeholders):
    op = func(**placeholders)

    def f(**kwargs):
        feed_dict = {}
        for argname, argvalue in kwargs.items():
            placeholder = placeholders[argname]
            feed_dict[placeholder] = argvalue
        return tf.get_default_session().run(op, feed_dict=feed_dict)

    return f

# 创建你的operation函数
encode_jpeg = create_op(
    func=tf.image.encode_jpeg,
    image=tf.placeholder(tf.uint8),
)

# 调用你的operation函数
decode_jpeg(contents=contents)

step 2 训练模型

1. 网络结构简介
   我之前做相关分享的ppt， 人脸识别原理与pix2pix分享 网盘地址第23页开始有pix2pix相关内容。 理论上的网络结构 上图是理论结构，但是为了加快训练速度，代码实现的是下图网络结构。 实际使用结构 每个encode与decode模块细节 encode与decode模块细节
2. 如果你比较着急可以直接执行以下代码，开始训练。我使用的GPU是英伟达的titanx，花了90分钟。

python pix2pix.py \
  --mode train \
  --output_dir face2face-model \
  --max_epochs 200 \
  --input_dir photos/combined/train \
  --which_direction AtoB

3. 如果希望深入了解细节，请看下面代码。但是以下代码不用直接执行用于训练模型:) 如果预先没有CNN卷积神经网络相关的知识，那么下面的代码会让气氛很尴尬的呢。

• 定义卷积

def conv(batch_input, out_channels, stride):
    ```输入结构：[batch, in_height, in_width, in_channels],
       卷积核结构： [filter_width, filter_height, in_channels, out_channels]
       输出结构： [batch, out_height, out_width, out_channels] 
       选用4x4的卷积核 + padding 1 + 步长stride，输出结构 VALID```
    with tf.variable_scope("conv"):
        in_channels = batch_input.get_shape()[3] # 输入图片的通道数
        # 初始化 4X4卷积核，使用random_normal_initializer初始化
        filter = tf.get_variable("filter",
                                [4, 4, in_channels, out_channels],
                                dtype=tf.float32,
                                initializer=tf.random_normal_initializer(0, 0.02)) 
        # padding 1
        padded_input = tf.pad(batch_input,
                              [[0, 0], [1, 1], [1, 1], [0, 0]], 
                              mode="CONSTANT")
        # 2D 卷积 步长为传参的stride
        conv = tf.nn.conv2d(padded_input, filter, [1, stride, stride, 1], padding="VALID")
        return conv

• 定义激活函数
  使用leaky ReLu激活函数，下图是leakReLu与ReLu的对比
  • ReLu 激活函数优点：
    a) 在刺激大于0的区域，不会出现梯度为0的问题。
    b) 计算效率高。
    c) 模型loss下降收敛快。大约是tanh与sigmoid激活函数的6倍。
  • Leaky ReLu 激活函数优点：
    a) ReLu的优点都有。
    b) 不会出现梯度为0的问题。
    c) 无论什么时候神经元都会被激活。

def lrelu(x, a):
    with tf.name_scope("lrelu"):
        # leak: a*x/2 - a*abs(x)/2;   linear: x/2 + abs(x)/2
        x = tf.identity(x)
        return (0.5 * (1 + a)) * x + (0.5 * (1 - a)) * tf.abs(x)

• 定义batchnorm

def batchnorm(input):
    with tf.variable_scope("batchnorm"):
        input = tf.identity(input)

        # 定义batch norm 中需要训练的两个参数offset与scale
        channels = input.get_shape()[3]
        offset = tf.get_variable("offset", 
                                 [channels], 
                                 dtype=tf.float32,
                                 initializer=tf.zeros_initializer())
        scale = tf.get_variable("scale", 
                                [channels], dtype=tf.float32,
                                initializer=tf.random_normal_initializer(1.0, 0.02))
      
        mean, variance = tf.nn.moments(input, axes=[0, 1, 2], keep_dims=False)
        variance_epsilon = 1e-5
        normalized = tf.nn.batch_normalization(input, 
                                               mean, variance, 
                                               offset, scale, 
                                               variance_epsilon=variance_epsilon)
        return normalized

step 3 Export Model & Freeze Model

• reduce model，我们需要生成模型用于图像生成，而判别模型可以去掉，以减少模型参数。这里我就不把生成模型重新复制一遍贴出来了。详细请看 github repo datitran/face2face-demo/reduce_model.py。思路是：
  • 首先把pix2pix.py中与生成模型相关部分复制了一份
  • 然后加载训练好的模型
  • 最后保存一个新模型。

reduce_model.py 中值得一书的事情。新建的generate_output函数，用于输入图片，生成图片。reduce_model.py 中所有 tf.variable_scope('名字')都与加载的训练好的模型一模一样，这样加载的模型会把它的参数与新模型的tf.variable_scope('名字')一一对应起来。由于新模型只保留了生成模型相关的tf.variable_scope('名字')，所以新模型的参数大大减少，实现model reduce.

x = tf.placeholder(tf.uint8, shape=(256, 512, 3), name='image_tensor')  # input tensor
y = generate_output(x)  # 输入图片，输出生成的图片 

with tf.Session() as sess:
    # 加载训练好的模型
    saver = tf.train.Saver()
    checkpoint = tf.train.latest_checkpoint(args.input_folder)
    saver.restore(sess, checkpoint)

    # 输出新模型
    saver = tf.train.Saver()
    saver.save(sess, './reduced_model')

• freeze model，我们把模型保存成一个.pb文件以方便调用

import tensorflow as tf
from tensorflow.python.framework import graph_util

def freeze_graph(model_folder):
    # 获取模型路径
    checkpoint = tf.train.get_checkpoint_state(model_folder)
    input_checkpoint = checkpoint.model_checkpoint_path
    output_graph = './frozen_model.pb'
    output_node_names = 'generate_output/output'

    # 加载 graph 
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta',
                                       clear_devices=True)
    # 取出 graph
    graph = tf.get_default_graph()
    input_graph_def = graph.as_graph_def()

    # 开一个新会话，加载参数，选择需要的节点，保存模型文件
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) # 加载graph的参数

        # tensorflow内置函数，将变量转为常量
        output_graph_def = graph_util.convert_variables_to_constants(
            sess,  # 用于取回参数
            input_graph_def,  # 用于取回节点node
            [output_node_names]  # 选择需要的节点名)

        # 将模型写入 .pb文件
        with tf.gfile.GFile(output_graph, 'wb') as f:
            f.write(output_graph_def.SerializeToString())
        print('%d ops in the final graph.' % len(output_graph_def.node))

freeze_graph('./reduced_model')

step 4 调用模型

freeze model大约200MB，这是我训练好的模型网盘地址，模型训练用的是400张图片，200epoch。

import tensorflow as tf
def load_graph(frozen_graph_filename):
    """ 加载 freezed model """
    graph = tf.Graph()
    with graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(frozen_graph_filename, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    return graph

graph = load_graph(frozen_model_file)
image_tensor = graph.get_tensor_by_name('image_tensor:0')
output_tensor = graph.get_tensor_by_name('generate_output/output:0')
sess = tf.Session(graph=graph)
# 图片必须是256X256，人脸在靠近中间的位置
generated_image = sess.run(output_tensor, 
                           feed_dict = {image_tensor: image})