SRCNN 图像超分辨率重建（tf2）

Posted 2021-10-23 __不想写代码__

文章目录

• 前言
• 一、SRCNN
• 二、SRCNN 实现
• • 1.模型的搭建
  • 2.生成训练数据
  • 3.训练过程：
  • 4.测试过程
• 总结

---

前言

把由放大缩小的引起的导致分辨率低的图像，转换成为 分辨率高的图像。更加关注的是重构图片过程中，填充新的像素。SRCNN 呢也是将深度学习用于图像重建的鼻祖，网络结构非常简单，于是我决定来复现一下它。
代码链接：https://github.com/jiantenggei/SRCNN-Keras （包含所有资源）

一、SRCNN

SRCNN 的网络结构特别简单，首先将一张低分辨率的图像作为输入，通过两个卷积后，还原成为高质量的图片。网络卷积运算时，保持特征图的大小和重构图片大小一致。且网络中没有线性连接。也就只有三层卷积~。

二、SRCNN 实现

首先在这里简述一下SRCNN 网络的训练和测试流程：
1.先将图片缩小后，再放大，制作不清晰的图片作为训练数据
2.将未处理过的原图片作为训练时的标签。
3.将图片和标签放到网络中训练。
4.测试模型时，将一张不清晰的图片输入到训练好的网络，生成的图片与清晰的原图片计算峰值信噪比。

1.模型的搭建

代码如下：

from keras.models import Sequential, model_from_json
from keras.layers.convolutional import Conv2D
from keras.layers.core import Activation


def built_model(input_shape=(33, 33, 1)):
    model = Sequential()
    model.add(Conv2D(filters=64, kernel_size=9,
                    padding='same', input_shape=input_shape))
    model.add(Activation('relu'))
    model.add(Conv2D(32, 1, padding='same'))
    model.add(Activation('relu'))
    model.add(Conv2D(input_shape[2], 5, padding='same'))
    return model

if __name__ == '__main__':
    model = built_model()
    model.summary()

最后一层把图片从多通道特征图，还原成和输入通道一致。

2.生成训练数据

1.先将图片缩小后，再放大，制作不清晰的图片作为训练数据
2.将未处理过的原图片作为训练时的标签。
代码如下：

def load_train(image_size=33, stride=33, scale=3,dirname=r'dataset\\train'):
    dir_list = os.listdir(dirname)
    images = [cv2.cvtColor(cv2.imread(os.path.join(dirname,img)),cv2.COLOR_BGR2GRAY) for img in dir_list]
    #==========================
    #这里判断采样步长 是否能被整除
    #=========================
    images = [img[0:img.shape[0]-np.remainder(img.shape[0],scale),0:img.shape[1]-np.remainder(img.shape[1],scale)] for img in images]

    trains = images.copy()
    labels = images.copy()
    #========================================
    #对train image 进行缩小 放大 产生不清晰的图像
    #========================================
    trains = [cv2.resize(img, None, fx=1/scale, fy=1/scale, interpolation=cv2.INTER_CUBIC) for img in trains]
    trains = [cv2.resize(img, None, fx=scale/1, fy=scale/1, interpolation=cv2.INTER_CUBIC) for img in trains]

    sub_trains = []
    sub_labels = []
    
    #========================================
    # 通过采样形成标签 和训练数据，
    # 一张 图片 通过采样，可以分成很多个图像块，作为训练数据，丰富样本
    #========================================
    for train, label in zip(trains, labels):
        v, h = train.shape
        print(train.shape)
        for x in range(0,v-image_size+1,stride):
            for y in range(0,h-image_size+1,stride):
                sub_train = train[x:x+image_size,y:y+image_size]
                sub_label = label[x:x+image_size,y:y+image_size]
                sub_train = sub_train.reshape(image_size,image_size,1)
                sub_label = sub_label.reshape(image_size,image_size,1)
                sub_trains.append(sub_train)
                sub_labels.append(sub_label)
    #========================================
    #编码为numpy array
    #========================================
    sub_trains = np.array(sub_trains)
    sub_labels = np.array(sub_labels)
    return sub_trains, sub_labels

def load_test(scale=3,dirname=r'dataset\\test'):
    #========================================
    # 生成测试数据的方式与训练数据相同
    # pre_tests 是用来保存缩小后的图片
    #========================================
    dir_list = os.listdir(dirname)
    images = [cv2.cvtColor(cv2.imread(os.path.join(dirname,img)),cv2.COLOR_BGR2GRAY) for img in dir_list]
    images = [img[0:img.shape[0]-np.remainder(img.shape[0],scale),0:img.shape[1]-np.remainder(img.shape[1],scale)] for img in images]

    tests = images.copy()
    labels = images.copy()
    
    pre_tests = [cv2.resize(img, None, fx=1/scale, fy=1/scale, interpolation=cv2.INTER_CUBIC) for img in tests]
    tests = [cv2.resize(img, None, fx=scale/1, fy=scale/1, interpolation=cv2.INTER_CUBIC) for img in pre_tests]
    
    pre_tests = [img.reshape(img.shape[0],img.shape[1],1) for img in pre_tests]
    tests = [img.reshape(img.shape[0],img.shape[1],1) for img in tests]
    labels = [img.reshape(img.shape[0],img.shape[1],1) for img in labels]

    return pre_tests, tests, labels

注意：代码中采样过程（三个for 训练处） 是将一张图片，截取一个个小的区域，这样一张图片就可以产生成多个数据，弥补训练样本不足的问题。

3.训练过程：

代码如下：

from tensorflow.python.keras.saving.model_config import model_from_config
from model import built_model
from utils import load_train
from keras.optimizers import Adam

def train():
    # ==========================
    # input_shape 输入图片大小
    # stride 原图片采样间隔
    # batch_size epochs learning_rate
    #============================
    input_shape = (33, 33, 1)
    stride = 14
    batch_size = 64
    epochs=100
    learning_rate=0.001

    # 定义模型
    srcnn_model = built_model(input_shape=input_shape)
    srcnn_model.load_weights(r'model\\srcnn_weight.hdf5')
    srcnn_model.summary()

    # 加载数据
    X_train, Y_train = load_train(image_size=input_shape[0], stride=stride)
    print(X_train.shape, Y_train.shape)
    optimizer = Adam(lr=learning_rate)
    srcnn_model.compile(optimizer=optimizer, loss='mean_squared_error', metrics=['accuracy'])
    srcnn_model.fit(X_train,Y_train,epochs=epochs,batch_size=batch_size)
    srcnn_model.save(r'model/srcnn.h5')

if __name__ == '__main__':
    train()

这里计算损失用均方差，因为输入和输出都是大小一致的图片~ 只是分辨率不同。

4.测试过程

代码如下：

from model import built_model
import os
from utils import load_test,psnr
import cv2
def test():
    input_shape = (None, None, 1)
    scale = 3
    srcnn_model = built_model(input_shape=input_shape)
    srcnn_model.load_weights(r'model\\srcnn_weight.hdf5')

    X_pre_test, X_test, Y_test = load_test(scale=scale)

    predicted_list = []

    for img in X_test:
        img = img.reshape(1,img.shape[0],img.shape[1],1)
        predicted=srcnn_model.predict(img)
        predicted_list.append(predicted.reshape(predicted.shape[1],predicted.shape[2],1))
    n_img = len(predicted_list)
    dirname = './result'
    for i in range(n_img):
        imgname = 'image{:02}'.format(i)
        cv2.imwrite(os.path.join(dirname,imgname+'_original.bmp'), X_pre_test[i])
        cv2.imwrite(os.path.join(dirname,imgname+'_input.bmp'), X_test[i])
        cv2.imwrite(os.path.join(dirname,imgname+'_answer.bmp'), Y_test[i])
        cv2.imwrite(os.path.join(dirname,imgname+'_predicted.bmp'), predicted_list[i])
          # 计算峰值信噪比
        answer = psnr(Y_test[i],predicted_list[i])
        print(imgname+"_psnr:",answer)

if __name__ == '__main__':
    test()

X_test 存放的是不清晰的图片，用于和预测结果计算计算峰值信噪比。

总结

各种网络各种功能，喂数据的方式不同，计算损失的方式不同罢了~