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

时间:2021-08-29 13:50:55

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幅无损原画作为训练数据,我没有那么多精力收集这样的数据,于是就在Konachanyande上面找了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直接插值放大的图像,以及通过本文训练的网络处理后得到的图像。在博客中看不出差别的,可以右键保存下来对比。(所用图片均来自网络)
Tensorflow实现二次元图片的超分辨率

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

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

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

源码下载