1.数据处理 dataset.py
训练集选用T91, 测试集选用Set5
低分辨图片(x)切割成17x17,高分辨图片(y)切割成17scale x 17scale
import cv2
import glob
import numpy as np
import tensorflow as tf
from keras import backend as k
def psnr(y_true, y_pred):
# assume RGB image
return 10.0 * k.log(1.0 / (k.mean(k.square(y_pred - y_true)))) / k.log(10.0)
# 将图片裁剪成放大倍数的整数倍
def Cropping(img, scale=3):
x, y = img.shape[0], img.shape[1]
x = x - np.mod(x, scale) # 取余
y = y - np.mod(y, scale)
img = img[0:x, 0:y]
return img
# 数据处理
def get_lr_hr(img, scale):
hr = img / 255. # 归一化
# 先缩小再放大 得到低分辨图片
lr = cv2.resize(hr, (hr.shape[1] // scale, hr.shape[0] // scale), interpolation=cv2.INTER_CUBIC)
return lr, hr
def load_data(path, scale=3, cut_size=17, stride=17):
LR = [] # 切片后的lr
HR = [] # 切片后的hr
size = [] # (nx, ny)切片次数
for i in path:
img = cv2.imread(i) # 读取图片
img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)[:, :, 0] # 格式转换 取Y通道
imgmod = Cropping(img, scale) # 尺寸处理
lr, hr = get_lr_hr(imgmod, scale) # 获取lr和hr
# 开切
h, w = lr.shape[0], lr.shape[1]
nx, ny = 0, 0
for x in range(0, h - cut_size + 1, stride):
nx += 1
for y in range(0, w - cut_size + 1, stride):
ny += 1
sub_lr = lr[x:x + cut_size, y:y + cut_size] # (17, 17)
sub_hr = hr[x * scale:(x + cut_size) * scale,
y * scale:(y + cut_size) * scale] # (51, 51)
# 转换为便于训练的格式
sub_input = sub_lr.reshape([cut_size, cut_size, 1])
sub_label = sub_hr.reshape([cut_size * scale, cut_size * scale, 1])
LR.append(sub_input)
HR.append(sub_label)
size.append((nx, ny // nx))
x_train = np.array(LR)
y_train = np.array(HR)
return x_train, y_train, size
# 合并
def merge(images, size):
(nx, ny) = size[0]
_, h, w, d = images.shape
img = np.zeros((h * nx, w * ny, d))
for idx, image in enumerate(images):
i = idx % ny # 取余
j = idx // ny # 取整
img[j * h:j * h + h, i * w:i * w + w, :] = image
return img
if __name__ == "__main__":
# 训练集
train_path = '../train/*.bmp'
train_path_list = glob.glob(train_path)
# 测试集
test_path = '../test/Set5/*.bmp'
test_path_list = glob.glob(test_path)
# 放大倍数
scale = 3
cut_size = 17
stride = 17
x_train, y_train, size = load_data(train_path_list, scale=scale, stride=stride)
print(x_train.shape)
print(y_train.shape)
# img = merge(y_train, size)
a = np.uint8(x_train[50] * 255)
b = np.uint8(y_train[50] * 255)
# cv2.namedWindow('1', cv2.WINDOW_NORMAL)
cv2.imshow('1', a)
# cv2.namedWindow('2', cv2.WINDOW_NORMAL)
cv2.imshow('2', b)
cv2.waitKey(0)
cv2.destroyAllWindows()
x_train.shape y_train.shape
2.搭建模型 model.py
第一层:5x5x64
第二层:3x3x32
第三层:3x3x(scalescalechannel)
scale为放大倍数,channel为输入图片的通道数
最后一层是亚像素卷积,用tf.nn.depth_to_space实现
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Conv2D, Activation
class ESPCN(Model):
def __init__(self, scale, channel):
super(ESPCN, self).__init__()
self.scale = scale
self.channel = channel
self.c1 = Conv2D(filters=64, kernel_size=5, strides=1, padding='same')
self.a1 = Activation('relu')
self.c2 = Conv2D(filters=32, kernel_size=3, strides=1, padding='same')
self.a2 = Activation('relu')
self.c3 = Conv2D(filters=self.scale * self.scale * self.channel, kernel_size=3, strides=1, padding='same')
self.d2s = tf.nn.depth_to_space
self.a3 = Activation('tanh')
def call(self, inputs, training=None, mask=None):
x = self.c1(inputs)
x = self.a1(x)
x = self.c2(x)
x = self.a2(x)
x = self.c3(x)
x = self.d2s(x, self.scale, data_format='NHWC')
y = self.a3(x)
return y
3.训练 + 测试
import glob
import numpy as np
import tensorflow as tf
from model import ESPCN
from dataset import load_data, psnr, Cropping
import cv2
from skimage.measure import compare_psnr, compare_ssim
# 训练集
train_path = '../291/*.bmp'
train_path_list = glob.glob(train_path)
# 放大倍数 输入图片尺寸 输出图片尺寸 移动步长
scale = 3
channel = 1
cut_size = 17
stride = 7
# 获取数据
x_train, y_train, train_size = load_data(train_path_list,
scale=scale,
cut_size=cut_size,
stride=stride)
test_path = '../test/Set5/*.bmp'
test_path_list = glob.glob(test_path)
x_test, y_test, test_size = load_data(test_path_list,
scale=scale,
cut_size=cut_size,
stride=stride)
# 网络
model = ESPCN(scale=scale, channel=channel)
# 配置
model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
loss=tf.keras.losses.MSE,
metrics=[psnr])
# 调整学习率
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.99, patience=2, mode='auto')
# 保存参数
checkpoint_save_path = './ESPCN2_checkpoint/ESPCN.ckpt'
# model.load_weights(checkpoint_save_path)
cp_callback = tf.keras.callbacks.ModelCheckpoint(
monitor='val_psnr',
filepath=checkpoint_save_path,
save_weights_only=True,
save_best_only=True
)
print(x_train.shape)
model.fit(x_train, y_train,
batch_size=64,
epochs=250,
validation_data=(x_test, y_test),
validation_freq=1,
callbacks=[reduce_lr, cp_callback])
i = 2
# 读取图片 转换YCrCb 规则化
img = cv2.imread(test_path_list[i])
img = Cropping(img, scale)
YCC = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
Y = YCC[:, :, 0]
color = YCC[:, :, 1:3]
# 压缩图片
lr = cv2.resize(YCC, (YCC.shape[1] // scale, YCC.shape[0] // scale), interpolation=cv2.INTER_CUBIC)
lr_bgr = cv2.cvtColor(lr, cv2.COLOR_YCrCb2BGR)
# 预测
h, w = lr.shape[0], lr.shape[1]
dt = lr[:, :, 0].reshape([1, h, w, 1]) / 255.
y_pre = model.predict(dt)
y_pre = y_pre.squeeze()
y_pre = y_pre * 255
# 合并
result = np.zeros([YCC.shape[0], YCC.shape[1], 3], dtype=np.uint8)
result[:, :, 0] = y_pre
result[:, :, 1:3] = color
result = cv2.cvtColor(result, cv2.COLOR_YCrCb2BGR)
# bicubic
bicubic = cv2.resize(lr_bgr, (YCC.shape[1], YCC.shape[0]), interpolation=cv2.INTER_CUBIC)
yB = cv2.cvtColor(bicubic, cv2.COLOR_BGR2YCrCb)[:, :, 0]
cv2.imshow('original', img)
cv2.imshow('input', lr_bgr)
cv2.imshow('output1', result)
cv2.imshow('output2', bicubic)
cv2.waitKey(0)
cv2.destroyAllWindows()
print('-' * 50)
psnr1 = compare_psnr(Y, y_pre, 255)
psnr2 = compare_psnr(Y, yB, 255)
print(psnr1, psnr2) # 23.180898231092254 21.97495909635292
print('-' * 50)
ssim1 = compare_ssim(Y, y_pre, win_size=11, data_range=255, multichannel=True)
ssim2 = compare_ssim(Y, yB, win_size=11, data_range=255, multichannel=True)
print(ssim1, ssim2) # 0.8354814229799264 0.791309318119287
原图
输入
双插值
ESPCN
psnr
ssim