【图像超分】论文复现：新手入门！Pytorch实现SRCNN，数据预处理、模型训练、测试、评估全流程详解，注释详细，简单修改就可以训练你自己的图像数据，有训练好的模型下载地址，随取随用_srcnn测试

第一次来请先看这篇文章：【超分辨率（Super-Resolution）】关于【超分辨率重建】专栏的相关说明，包含专栏简介、专栏亮点、适配人群、相关说明、阅读顺序、超分理解、实现流程、研究方向、论文代码数据集汇总等）

---

前言

论文地址：Image Super-Resolution Using Deep Convolutional Networks

论文精读：【图像超分】论文精读：Image Super-Resolution Using Deep Convolutional Networks（SRCNN）

请配合上述论文精读文章使用，效果更佳！

代码地址：图像超分辨率SRCNN和FSRCNN复现代码，除基本的网络实现外，还有特征图可视化，PSNR曲线图可视化，测试自己的图像数据等

不想理解原理，希望直接跑通然后应用到自己的图像数据的同学，请直接下载上面的代码，有训练好的模型，直接用即可。具体使用方式见代码中的README！有问题来本文评论区留言！

深度学习的模型训练一般遵循以下步骤：

1. 准备数据集，以及数据预处理
2. 搭建网络模型
3. 设置参数并训练
4. 测试训练好的模型
5. 用训练好的模型测试自己的数据

下面让我们根据通用流程，一步一步复现SRCNN吧！目标是：输入大小为 h×w 的图像 X，输出为一个 sh×sw 的图像 Y，s 为放大倍数。论文中s取值为2，3，4。

注：放大倍数为1，则图像分辨率不变，只是图像变清晰！如果你只是想将模糊图像通过SRCNN超分后变清晰，那么请将该参数设置为1。

硬件环境：windows11+RTX 2060（比这个高肯定没问题，我这个配置本机跑500个epoch一点问题没有，一会就跑完。但超分自己的图像时，如果图像很大，可能有内存溢出的错误）
运行环境：jupyter notebook/pycharm（前者好处是分代码段运行，测试方法，适合学习用；后者适合跑完整项目用）
pytorch环境：torch1.9.1+cuda11.1（其他版本没测试过，应该问题不大）

1. 准备数据集和数据预处理

1.1 数据集选择

简单回顾论文中出现过的数据集：

• 训练集（Train dataset）：91 images、ILSVRC 2013 ImageNet
• 测试集（Test dataset）：Set5、Set14、BSD200

由于91 images的数据处理有些繁琐，对于新手不太友好，而ImageNet数据集过大，可能不利于硬件设备不好的同学，所以本文使用折中的BSD200数据集，包含200张图像训练集和200张图像测试集。

1.2 数据预处理

根据论文中第4节实验部分开头所述，RGB 颜色模式色调、色度、饱和度三者混在一起难以分开，超分只用于YCbCr的Y通道，即亮度分量。因此，需要将图像数据由RGB模型转换成 YCbCr 颜色模式，Y 是指亮度分量，Cb 表示 RGB 输入信号蓝色部分与 RGB 信号亮度值之间的差异，Cr 表示 RGB 输入信号红色部分与 RGB 信号亮度值之间的差异。

此外，通过观察BSD200数据集中的图像可知，数据集中的图像宽高并不一致，每张图像的大小也千差万别。而SRCNN要求输入图像的长宽一致，所以在训练之前，需要将图像裁剪成宽高一致的图像。采用的策略是按图像中心外扩，resize成300×300的图像。

数据预处理的代码如下：

# 判断某个文件是否是图像
# enswith判断是否以指定的.png,.jpg,.jpeg结尾的字符串
# 可以根据情况扩充图像类型，加入.bmp、.tif等
def is_image_file(filename):
    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg"])

# 读取图像转为YCbCr模式，得到Y通道
def load_img(filepath):
    img = Image.open(filepath).convert('YCbCr')
    y, _, _ = img.split()
    return y

# 裁剪大小，宽高一致为300
# 如果想训练自己的数据集，请根据情况修改裁剪大小
CROP_SIZE = 300

# 封装数据集，适配后面的torch.utils.data.DataLoader中的dataset，定义成类似形式
# 类参数为图像文件夹路径和放大倍数
# __len__(self) 定义当被len()函数调用时的行为（返回容器中元素的个数）
#__getitem__(self) 定义获取容器中指定元素的行为，相当于self[key]，即允许类对象可以有索引操作。
#__iter__(self) 定义当迭代容器中的元素的行为
# 返回输入图像和标签，传入DataLoader的dataset参数
class DatasetFromFolder(Dataset):
    def __init__(self, image_dir, zoom_factor):
        super(DatasetFromFolder, self).__init__()
        self.image_filenames = [join(image_dir, x) for x in listdir(image_dir) if is_image_file(x)] # 图像路径列表
        crop_size = CROP_SIZE - (CROP_SIZE % zoom_factor) # 处理放大倍数，防止用户瞎设置，本例只能设置为2，3，4，大小不变
        # 数据集变换
        # 还有一些其他的变换操作，如归一化等，遇到一个积累一个
        self.input_transform = transforms.Compose([transforms.CenterCrop(crop_size), # 从图片中心裁剪成300*300
                                                   transforms.Resize(
                                                       crop_size // zoom_factor),    # Resize, 输入应该是缩放倍数后的图像，因为先缩小后放大
                                                   transforms.Resize(
                                                       crop_size, interpolation=Image.BICUBIC), # 双三次插值
                                                   transforms.ToTensor()]) # 图像转成tensor
        # label标签，超分不是分类问题，定义成一样的就行
        self.target_transform = transforms.Compose(
            [transforms.CenterCrop(crop_size), transforms.ToTensor()])

    def __getitem__(self, index):
        input = load_img(self.image_filenames[index]) # 输入是图像的Y通道，即亮度通道
        target = input.copy()
        input = self.input_transform(input)
        target = self.target_transform(target)
        return input, target

    def __len__(self):
        return len(self.image_filenames) # 图像个数
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

至此，数据预处理完成，训练之前就可以传入DataLoader的dataset参数中。

1.3 评估指标PSNR和SSIM

训练过程中，每个epoch会得到一个SRCNN的模型。那么如何评价该模型的性能呢？需要两个评价指标：PSNR和SSIM。由于训练时需要用到，所以需要提前将计算PSNR和SSIM的代码准备好。

1.3.1 PSNR

PSNR(Peak Signal to Noise Ratio)为峰值信噪比，计算公式如下：
$$MSE=\frac{1}{H\times W}\sum _{i=1}^H\sum _{j=1}^W(X(i,j)-Y(i,j){)}^2$$
$$PSNR=10{\log }_{10}\left(\frac{{\left(2^n-1\right)}^2}{MSE}\right)$$
其中，MSE是均方误差。两个图像对应像素位置的值相减再平方和最后取平均。n是每像素的比特数，一般取256。PSNR的单位是dB，该值越大表示图像的失真越小。

通常认为，PSNR在38以上的时候，人眼就无法区分两幅图片了。

PSNR的计算代码如下：

def psnr(loss):
    return 10 * log10(1 / loss.item())
1
2

损失函数使用的就是MSE，所以这么定义。log10是python的math库中的包。

做实验的时候，PSNR一般保留两位小数。

1.3.2 SSIM

SSIM（Structural Similarity，结构相似性）由三个对比模块组成：亮度、对比度、结构。

(1) 亮度对比函数
图像平均灰度：
$${\mu }_X=\frac{1}{H\times M}\sum _{i=1}^H\sum _{j=1}^{M}X(i,j)$$
亮度对比函数：
$$l(x,y)=\frac{2{\mu }_x{\mu }_y+C_1}{{\mu }_x^2+{\mu }_y^2+C_1}$$

(2) 对比度对比函数
图像的标准差：
$${\sigma }_X={\left(\frac{1}{H+W-1}\sum _{i=1}^H\sum _{j=1}^M{\left(X(i,j)-{\mu }_X\right)}^2\right)}^{\frac{1}{2}}$$
对比度对比函数:
$$c(x,y)=\frac{2{\sigma }_x{\sigma }_y+C_2}{{\sigma }_x^2+{\sigma }_y^2+C_2}$$

(3) 结构对比函数
$$s(x,y)=\frac{{\sigma }_{xy}+C_3}{{\sigma }_x{\sigma }_y+C_3}$$

综合上述三个部分，得到 SSIM 计算公式：

SSIM ⁡ ( x , y ) = f ( l ( x , y ) , c ( x , y ) , s ( x , y ) ) = [ l ( x , y ) ] α [ c ( x , y ) ] β [ s ( x , y ) ] γ

$$\begin{aligned} \operatorname{SSIM}(x, y) & =f(l(x, y), c(x, y), s(x, y)) \\ & =[l(x, y)]^{\alpha}[c(x, y)]^{\beta}[s(x, y)]^{\gamma} \end{aligned}$$ SSIM(x,y)​=f(l(x,y),c(x,y),s(x,y))=[l(x,y)]α[c(x,y)]β[s(x,y)]γ​

其中，