Tensorflow实现二次元图片的超分辨率

github上有一个很有意思的项目，waifu2x，原理是通过一个训练好的CNN，将低分辨率的图像放大2倍或更多，同时保留足够的细节，使放大后的图像不会过于模糊或失真。该项目是用lua+Touch写的，最近在学习Tensorflow，闲暇之余打算在TF上自己实现一个这样的系统。

1. 网络选择

根据论文Image Super-Resolution Using Deep Convolutional Networks，作者使用了一个3层CNN，是一个非常小的网络，同时为了避免边界效应（bound effect），每一个卷积层都不使用padding，卷积核大小分别为9、1、5，作者也提到，适当增大卷积核的大小，可以提升最后的效果，但也带来了更多的计算量。使用这个网络时，首先需要将低分辨的图像通过bicubic插值到所需的大小，然后再输入网络。对于一个核大小分别为9、3、5的网络，定义如下：

def srcnn_935(patches, name='srcnn'):    with tf.variable_scope(name):        upscaled_patches = tf.image.resize_bicubic(patches, [INPUT_SIZE, INPUT_SIZE], True)        conv1 = conv2d(upscaled_patches, 9, 9, 64, padding='VALID', name='conv1')        relu1 = relu(conv1, name='relu1')        conv2 = conv2d(relu1, 3, 3, 32, padding='VALID', name='conv2')        relu2 = relu(conv2, name='relu2')        return conv2d(relu2, 5, 5, NUM_CHENNELS, padding='VALID', name='conv3')

不过，waifu2x本身没有使用这个网络，而是使用了一个7层的vgg-style网络，按照作者的说法，这个7层网络更便于调整，效果也更好，我也用Tensorflow实现了这个网络，如下：

def vgg7(patches, name='vgg7'):    with tf.variable_scope(name):        upscaled_patches = tf.image.resize_bicubic(patches, [INPUT_SIZE, INPUT_SIZE], True)        conv1 = conv2d(upscaled_patches, 3, 3, 32, padding='VALID', name='conv1')        lrelu1 = leaky_relu(conv1, name='leaky_relu1')        conv2 = conv2d(lrelu1, 3, 3, 32, padding='VALID', name='conv2')        lrelu2 = leaky_relu(conv2, name='leaky_relu2')        conv3 = conv2d(lrelu2, 3, 3, 64, padding='VALID', name='conv3')        lrelu3 = leaky_relu(conv3, name='leaky_relu3')        conv4 = conv2d(lrelu3, 3, 3, 64, padding='VALID', name='conv4')        lrelu4 = leaky_relu(conv4, name='leaky_relu4')        conv5 = conv2d(lrelu4, 3, 3, 128, padding='VALID', name='conv5')        lrelu5 = leaky_relu(conv5, name='leaky_relu5')        conv6 = conv2d(lrelu5, 3, 3, 128, padding='VALID', name='conv6')        lrelu6 = leaky_relu(conv6, name='leaky_relu6')        return conv2d(lrelu6, 3, 3, NUM_CHENNELS, padding='VALID', name='conv_out')

这里使用了Leaky ReLU，其实不是必须的，在该问题上，和ReLU对比之后发现没有太大区别。

有很多研究表明，之前的bicubic插值没有必要，完全可以融入到网络之中，将插值（或者说是放大图像的过程）也作为网络可以学习的一部分，如论文Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network中所述，要实现这个目标，需要用到反卷积层（deconvolution）,这个叫法有歧义，有的地方也称为transposed convolution，在Tensorflow中对应于conv2d_transpose。
最自然的想法是将这一层放在网络的最前端，相当于直接替代bicubic过程，但这会导致后面几层的计算量都增大，所以，将反卷积放在网络最后，才是明智的选择。这样，网络前几层都仅仅对小尺寸的图片进行处理，到了最后一步才将其放大。似乎目前最新的waifu2x也吸取了这个想法。将vgg网络的最后一层改为反卷积层即可：

def vgg_deconv_7(patches, name='vgg_deconv_7'):    with tf.variable_scope(name):        conv1 = conv2d(patches, 3, 3, 16, padding='VALID', name='conv1')        lrelu1 = leaky_relu(conv1, name='leaky_relu1')        conv2 = conv2d(lrelu1, 3, 3, 32, padding='VALID', name='conv2')        lrelu2 = leaky_relu(conv2, name='leaky_relu2')        conv3 = conv2d(lrelu2, 3, 3, 64, padding='VALID', name='conv3')        lrelu3 = leaky_relu(conv3, name='leaky_relu3')        conv4 = conv2d(lrelu3, 3, 3, 128, padding='VALID', name='conv4')        lrelu4 = leaky_relu(conv4, name='leaky_relu4')        conv5 = conv2d(lrelu4, 3, 3, 128, padding='VALID', name='conv5')        lrelu5 = leaky_relu(conv5, name='leaky_relu5')        conv6 = conv2d(lrelu5, 3, 3, 256, padding='VALID', name='conv6')        lrelu6 = leaky_relu(conv6, name='leaky_relu6')        batch_size = int(lrelu6.get_shape()[0])        rows = int(lrelu6.get_shape()[1])        cols = int(lrelu6.get_shape()[2])        channels = int(patches.get_shape()[3])        # to avoid chessboard artifacts, the filter size must be dividable by the stride        return deconv2d(lrelu6, 4, 4, [batch_size, rows*2, cols*2, channels], stride=(2, 2), name='deconv_out')

在使用反卷积时，有一点必须注意，核大小最好能被stride整除，否则容易产生类似于棋盘格的artifacts，如果无法满足整除要求，可以考虑使用resized convolution代替，关于这一点的详细论述，可见这篇文章。
我在训练自己的模型时，选用了带有反卷积层的vgg网络。

2. 数据准备

waifu2x的作者使用了3000幅无损原画作为训练数据，我没有那么多精力收集这样的数据，于是就在Konachan和yande上面找了100张图片，尽可能找分辨率高的，比如不低于1080p的图像，同时绝对不能要jpg格式或其他有损压缩格式的图像，有损压缩会引入噪声（jpeg artifacts）。（咳咳~这两个网站里面有一些图片很污，请控制好自己）
网络的输入只是很小的一块图片（patch），每次都读入完整的图像，再从中截取一小块，效率实在不高，因此我事先将这些图片又分割成了较小的块。有两种分割方式，第一种是顺序分割，在原图像中每隔几个像素就截取一块并保存下来，这样产生的数据量是很大的，我的100幅图像产生了上百万的patch，我柔弱的机械硬盘想要打开这个目录都困难。。。为此，我选择了随机分割，在原图像中随机截取一定数量的patch，最终产生了10万个patch作为训练数据。在我目前的设定中，块大小为80x80。

3. 数据增强

数据增强虽然只是一种trick，但对于最后模型的效果却有很大的影响。根据现在的设定，训练数据为80x80图像，网络输入大小为28x28，放大倍数设定为2， 因此应该从80x80的图像中截取56x56的图像作为label，当然，在随机位置截取就构成了数据增强的第一步。得到56x56的label后，又对其进行了随机翻转、随机亮度以及随机对比度。为了得到网络的输入，需要对这个56x56的图像进行降采样，在tensorflow中，提供了几种缩放图像的方法，比如nearest neighbor、bicubic还有area，我随机从这些算法中选择一个对图像进行降采样，得到28x28的低分辨率图像。在随机选择降采样算法时，我排除了bilinear，因为bilinear可以看做一个很强的低通滤波器，会损失大量的高频信息，在试验中发现，网络对高频信息的恢复很困难。我使用waifu2x对经过bilinear缩小后的图像进行恢复，结果也比其他图像模糊很多，看来waifu2x对损失的高频信息也无能为力。
除了上述数据增强外，还可以适当的给降采样后的图像（作为Label的高分辨率图像是绝对不能变的）加入噪声，考虑到网上很多图片都是jpeg格式，都带有jpeg artifacts，因此，也可以在训练时对图像加入jpeg noise，对于Tensorflow而言，这很好实现，只需通过encode_jpeg对图像进行编码，再用decode_jpeg解码即可。从读入数据、数据增强到最后生成batch的代码如下：

def batch_queue_for_training(data_path):    filename_queue = tf.train.string_input_producer(tf.train.match_filenames_once(join(data_path, '*.png')))    file_reader = tf.WholeFileReader()    _, image_file = file_reader.read(filename_queue)    patch = tf.image.decode_png(image_file, NUM_CHENNELS)    # we must set the shape of the image before making batches    patch.set_shape([PATCH_SIZE, PATCH_SIZE, NUM_CHENNELS])    patch = tf.image.convert_image_dtype(patch, dtype=tf.float32)    if MAX_RANDOM_BRIGHTNESS > 0:        patch = tf.image.random_brightness(patch, MAX_RANDOM_BRIGHTNESS)    if len(RANDOM_CONTRAST_RANGE) == 2:        patch = tf.image.random_contrast(patch, *RANDOM_CONTRAST_RANGE)    patch = tf.image.random_flip_left_right(patch)    high_res_patch = tf.image.random_flip_up_down(patch)    crop_margin = PATCH_SIZE - LABEL_SIZE    assert crop_margin >= 0    if crop_margin > 1:        high_res_patch = tf.random_crop(patch, [LABEL_SIZE, LABEL_SIZE, NUM_CHENNELS])    downscale_size = [INPUT_SIZE, INPUT_SIZE]    resize_nn = lambda: tf.image.resize_nearest_neighbor([high_res_patch], downscale_size, True)    resize_area = lambda: tf.image.resize_area([high_res_patch], downscale_size, True)    resize_cubic = lambda: tf.image.resize_bicubic([high_res_patch], downscale_size, True)    r = tf.random_uniform([], 0, 3, dtype=tf.int32)    low_res_patch = tf.case({tf.equal(r, 0): resize_nn, tf.equal(r, 1): resize_area}, default=resize_cubic)[0]    # add jpeg noise to low_res_patch    if JPEG_NOISE_LEVEL > 0:        low_res_patch = tf.image.convert_image_dtype(low_res_patch, dtype=tf.uint8, saturate=True)        jpeg_quality = 100 - 5 * JPEG_NOISE_LEVEL        jpeg_code = tf.image.encode_jpeg(low_res_patch, quality=jpeg_quality)        low_res_patch = tf.image.decode_jpeg(jpeg_code)        low_res_patch = tf.image.convert_image_dtype(low_res_patch, dtype=tf.float32)    # we must set tensor's shape before doing following processes    low_res_patch.set_shape([INPUT_SIZE, INPUT_SIZE, NUM_CHENNELS])    # add noise to low_res_patch    if GAUSSIAN_NOISE_STD > 0:        low_res_patch += tf.random_normal(low_res_patch.get_shape(), stddev=GAUSSIAN_NOISE_STD)    low_res_patch = tf.clip_by_value(low_res_patch, 0, 1.0)    high_res_patch = tf.clip_by_value(high_res_patch, 0, 1.0)    # Generate batch    low_res_batch, high_res_batch = tf.train.shuffle_batch(        [low_res_patch, high_res_patch],        batch_size=BATCH_SIZE,        num_threads=NUM_PROCESS_THREADS,        capacity=MIN_QUEUE_EXAMPLES + 3 * BATCH_SIZE,        min_after_dequeue=MIN_QUEUE_EXAMPLES)    return low_res_batch, high_res_batch

注意，处理后的图像是浮点型的，像素范围为0~1，一般而言，往往使用中心化的数据，比如-1~1，但是在这项任务中，并没有发现数据中心化后对训练有明显改善。

4. 训练

loss函数使用均方误差（MSE），如果有较大噪声，可以考虑使用对噪声更鲁棒的Huber Loss。waifu2x在求均方误差时做了一些小修改，不同通道的误差乘以了不同的权值，BGR对应的权重分别为0.11448, 0.58661, 0.29891，这三个值来自彩色图像到灰度图像的转换公式，由于人眼对绿色是最为敏感的，因此G分量的权值最高。这样做能使重建后的图像在颜色上更友好，不会出现偏色。

def loss(inferences, ground_truthes, weights_decay=0, name='loss'):    with tf.name_scope(name):        slice_begin = (int(ground_truthes.get_shape()[1]) - int(inferences.get_shape()[1])) // 2        slice_end = int(inferences.get_shape()[1]) + slice_begin        delta = inferences - ground_truthes[:, slice_begin: slice_end, slice_begin: slice_end, :]        delta *= [[[[0.11448, 0.58661, 0.29891]]]]  # weights of B, G and R        l2_loss = tf.pow(delta, 2)        mse_loss = tf.reduce_mean(tf.reduce_sum(l2_loss, axis=[1, 2, 3]))        if weights_decay > 0:            weights = tf.get_collection('weights')            reg_loss = weights_decay * tf.reduce_sum(                tf.pack([tf.nn.l2_loss(i) for i in weights]), name='regularization_loss')            return mse_loss + reg_loss        else:            return mse_loss

根据waifu2x作者的回答，小batch对该任务有较好的效果，可以使loss下降更快，收敛更快。通过我自己的实验，发现确实如此，一开始将batch size设为128和64，不仅参数初始loss很大，而且往往一段时间后训练就发散。在该任务中，batch中每个样本产生梯度竞争可能比较激烈（没有时间验证），所以导致了收敛过慢。目前我使用的batch size为16，waifu2x作者使用的batch size为8。
梯度更新算法使用了Adam，并配合learning rate decay。一般而言，Adam不需要配合其他decay，但我发现在这项任务中，加入decay能使训练更稳定，防止出现后期训练发散的情况。
waifu2x训练了2000个epoch，花费了12小时。本人计算机算力有限，数据有限，所以只训练了100个epoch，但是也达到了还不错的结果。

5. 使用

由于网络一次只输入一副28x28的小图像，因此在实际处理大尺寸图像时，需要依次截取28x28的图像，输入网络，再将输出的图像依次写到另一幅图像中，具体代码可以看本文分享的源码。（上面的说法不准确，是我之前的误解。由于网络是全卷积网络，所以可以输入任意大小的图片，没有必要分块处理，只是训练的时候采用了28x28的图像而已）
下面三幅图像，分别是原始的小图像，通过bicubic直接插值放大的图像，以及通过本文训练的网络处理后得到的图像。在博客中看不出差别的，可以右键保存下来对比。（所用图片均来自网络）

最后说一下，虽然本文只对二次元图像进行了超分辨率处理，但是本文所使用的网络完全可以用来处理一般照片。如果想要对现实场景的照片进行超分辨率，只需将训练数据换为照片，然后再训练一个模型即可。

源码下载