论文地址:点击此处跳转
搞这篇论文时,踩了很多坑,效果优于网上的大部分代码,网上大部分代码效果离理想效果差5-6db,而我把里面的坑踩掉了,最后效果很逼近论文中的效果。
Step1:数据集的处理:Train 数据集包括 91 张图片,仅取亮度通道。之后将图片 bicubic 将分辨率降低,再次进行 bicubic 将其大小放大至与原图一致。将图片按照 stride 进行裁剪,这样 91 张图片可以得到 2 万多张大小相同的子图片( 我想应该是为了增大训练数据,以便得到较好的结果)将子图片输入网络进行训练,且将未经过 bicubic 的数据作为 label。
Step2: :搭建卷积神经网络:三层网络模型,前两层有激活函数,最后一层不使用激活函数。且每层网络中的 filter weight 以及 biase 的大小都已给出,训练时需要学习的参数只有这六项的权重。损失函数取得是 MSE 的值。需要注意的是 Train 时为避免边界效应,卷积层都没有填充,因此输入的图片会变小。因此预处理时 label 要取相对应的大小,以进行 loss 函数计算。
Step3 :训练结束后保存 filter weight 以及 biase 的权重值。这里 Test 时需要整张图片读取,为保证输出大小与输入大小一致,卷积层中的 padding=“SAME”进行填充,以保证输出图片大小不变。之后进行 psnr 计算即可。
Code of Step1 :
image=scipy.misc.imread(data[i], flatten=True, mode='YCbCr').astype(np.float)
if len(image.shape) == 3:
h, w, _ = image.shape
h = h - np.mod(h, scale)
w = w - np.mod(w, scale)
label_ = image[0:h, 0:w, :]
else:
h, w = image.shape
h = h - np.mod(h, scale)
w = w - np.mod(w, scale)
label_ = image[0:h, 0:w]
label_ = label_ / 255. #此时 label_为真图,后续进行 mse 计算时与预测图片进行对比
input_ = scipy.ndimage.interpolation.zoom(label_, (1./scale), prefilter=False)
input_ = scipy.ndimage.interpolation.zoom(input_, (scale/1.), prefilter=False)
for x in range(0, h-image_size+1, stride): #以 stride 为步长进行取子图片操作
for y in range(0, w-image_size+1, stride):
sub_input = input_[x:x+image_size, y:y+image_size] # [33 x 33]大小
sub_label=label_[x+int(padding):x+int(padding)+label_size,
y+int(padding):y+int(padding)+label_size] # [21 x 21]大小
with h5py.File(savepath, 'w') as hf: #数据集的制作
hf.create_dataset('data', data=arrdata)
hf.create_dataset('label', data=arrlabel)
Code of Step2 :
images = tf.placeholder(tf.float32, [None, None, None, c_dim], name='images')
labels = tf.placeholder(tf.float32, [None, None, None, c_dim], name='labels')
#滤波器
weights = {
'w1': tf.Variable(tf.random_normal([9, 9, 1, 64], stddev=1e-3),trainable=trainable,
name='w1'),
'w2':tf.Variable(tf.random_normal([1,1, 64, 32], stddev=1e-3),trainable=trainable,
name='w2'),
'w3':tf.Variable(tf.random_normal([5,5,32,1],stddev=1e-3),trainable=trainable,
name='w3')
}
#偏置
biases = {
'b1': tf.Variable(tf.zeros([64]),trainable=trainable ,name='b1'),
'b2': tf.Variable(tf.zeros([32]),trainable=trainable, name='b2'),
'b3': tf.Variable(tf.zeros([1]),trainable=trainable, name='b3')
}
#三层卷积网络
conv1 = tf.nn.relu(tf.nn.conv2d(images, weights['w1'], strides=[1,1,1,1],
padding=padding) + biases['b1'])
conv2 = tf.nn.relu(tf.nn.conv2d(conv1, weights['w2'], strides=[1,1,1,1], padding=padding)
+ biases['b2'])
conv3 = tf.nn.conv2d(conv2, weights['w3'], strides=[1,1,1,1], padding=padding) +
biases['b3']
#预测值以及 loss 函数和优化器
pred=conv3
loss = tf.reduce_mean(tf.square(labels - pred))
train_op = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)
saver=tf.train.Saver(max_to_keep=5)
with tf.Session() as sess:
if is_train:
print("Training...")
sess.run(tf.initialize_all_variables()) #变量初始化
for ep in xrange(epoch):
batch_idxs = len(train_data) // batch_size
for idx in xrange(0, batch_idxs):
batch_images = train_data[idx*batch_size : (idx+1)*batch_size]
batch_labels = train_label[idx*batch_size : (idx+1)*batch_size]
counter +=1
_, err = sess.run([train_op, loss], feed_dict={images: batch_images, labels:
batch_labels})
if counter % 10 == 0:
print("Epoch: [%2d], step: [%2d], time: [%4.4f], loss: [%.8f]" \
% ((ep+1), counter, time.time()-start_time, err))
if counter % 500 == 0:
saver.save(sess,os.path.join('checkpoint','SRCNN'),global_step=counter,write_meta_graph=False)
Code of Step3 :
with tf.Session() as sess:
ckpt = tf.train.get_checkpoint_state("checkpoint")
if ckpt and ckpt.model_checkpoint_path: # 加载保存的模型
saver.restore(sess, ckpt.model_checkpoint_path)
result = pred.eval({images: train_data, labels: train_label})
result2 = result.squeeze()
image_path = os.path.join(os.getcwd(),'sample')
image_path = os.path.join(image_path, "test_image.png")
imsave(result2, image_path)
label=train_label.squeeze()
print(pnsr(label,result2)) #计算 pnsr 值
def pnsr(img1,img2):
diff = np.abs(img1*255.0- img2*255.0)
mse = np.square(diff).mean() #mse 表示当前图像与原有图像的均方误差
psnr = 20 * np.log10(255 / np.sqrt(mse)) #评价指标即峰值信噪比
return psnr
写到这里,基本上整个srcnn的框架已经搭好了,接下来就是进行训练然后看下效果。在我的机子上跑了40个小时,最终得到的结果:
预测,读取checkpoint\SRCNN-16300000里面的参数,并对set5里面的woman图片进行SRCNN。
Input:
label:
Psnr约为25.38db,与论文中的30.92db相差好大。而且我把生成的图片保存下来发现似乎亮度比输入的低,感觉问题好像出在这里。具体我也不太清楚到底问题出在哪里。后来打开原作者的MATLAB代码进行数据对比,发现了step1预处理中有一些问题:
主要有:
问题一解决对应代码:
def rgb2ycbcr(img, only_y=True): #自己重新写的rgb2ycbcr函数以求对应到matlab的rgb2ycbcr函数
'''same as matlab rgb2ycbcr
only_y: only return Y channel
Input:
uint8, [0, 255]
float, [0, 1]
'''
in_img_type = img.dtype
img.astype(np.float32)
if in_img_type != np.uint8:
img *= 255.
# convert
if only_y:
rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
else:
rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
[24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
if in_img_type == np.uint8:
rlt = rlt.round()
else:
rlt /= 255.
return rlt.astype(in_img_type)
问题二及问题三解决方法:
#进行两次插值构造低分辨率图片
label_1=Image.fromarray(label_)
input_= label_1.resize(( w // scale,h // scale),Image.BICUBIC)
input_= input_.resize((w,h), Image.BICUBIC)
input_=np.float64(input_)
#保存四位小数
label_=np.around(label_, decimals=4)
input_=np.around(input_,decimals=4)
#下面这个插值函数与matlab中不一致,因此舍弃
#input_ = scipy.ndimage.interpolation.zoom(label_, (1./scale), prefilter=False)#进行双三次插值变为低分辨率图片
#input_ = scipy.ndimage.interpolation.zoom(input_, (scale/1.), prefilter=False)#再次进行双三次插值变为与高分辨率图片一样大小
问题四在最后做了裁剪。
经过这一系列的坑之后,现在的训练得到的效果明显优于bicubic等方法,只是距离理想效果仍有1db多的差距。而且图片仍然偏暗,后来发现主要问题出在测试的时候,做了zero-padding之后,论文中还有一步操作是对感受野中非零区域做normalize。这一步漏掉了,所以导致了最后超分辨率生成的图片整体偏暗。(最后这块normalize我就没有再去做了,很简单,自己可以加一下就好了,而且即使不加也是优于bicubic等方法的)
最后我将最终完整的py代码分为三个.py文件,先运行pre预处理,再运行train进行训练,最后进行test预测。完整代码下面贴出来链接,供大家交流。
完整代码压缩包链接