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

文章目录

• 前言
• 1. 准备数据集和数据预处理
• • 1.1 数据集选择
  • 1.2 数据预处理
  • 1.3 评估指标PSNR和SSIM
  • • 1.3.1 PSNR
    • 1.3.2 SSIM
• 2. 定义网络结构
• 3. 设置参数并训练模型
• • 3.1 参数设置
  • 3.2 模型训练
• 4. 测试训练好的模型
• 5. 用训练好的SRCNN模型超分自己的图像数据
• 6. 其他补充
• • 6.1 特征图可视化
  • 6.2 训练过程中PSNR折线图可视化
  • 6.3 以像素计算自己的图像输入输出前后的PSNR和SSIM

前言

论文地址：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) # 图像个数

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

1.3 评估指标PSNR和SSIM

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

1.3.1 PSNR

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

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

PSNR的计算代码如下：

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

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

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

1.3.2 SSIM

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

(1) 亮度对比函数
图像平均灰度：
μ X = 1 H × M ∑ i = 1 H ∑ j = 1 M X ( i , j ) \mu_{X}=\frac{1}{H \times M} \sum_{i=1}^{H} \sum_{j=1}^{M} X(i, j) μX​=H×M1​i=1∑H​j=1∑M​X(i,j)
亮度对比函数：
l ( x , y ) = 2 μ x μ y + C 1 μ x 2 + μ y 2 + C 1 l(x, y)=\frac{2 \mu_{x} \mu_{y}+C_{1}}{\mu_{x}^{2}+\mu_{y}^{2}+C_{1}} l(x,y)=μx2​+μy2​+C1​2μx​μy​+C1​​

(2) 对比度对比函数
图像的标准差：
σ X = ( 1 H + W − 1 ∑ i = 1 H ∑ j = 1 M ( X ( i , j ) − μ X ) 2 ) 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}} σX​=(H+W−11​i=1∑H​j=1∑M​(X(i,j)−μX​)2)21​
对比度对比函数:
c ( x , y ) = 2 σ x σ y + C 2 σ x 2 + σ y 2 + C 2 c(x, y)=\frac{2 \sigma_{x} \sigma_{y}+C_{2}}{\sigma_{x}^{2}+\sigma_{y}^{2}+C_{2}} c(x,y)=σx2​+σy2​+C2​2σx​σy​+C2​​

(3) 结构对比函数
s ( x , y ) = σ x y + C 3 σ x σ y + C 3 s(x, y)=\frac{\sigma_{x y}+C_{3}}{\sigma_{x} \sigma_{y}+C_{3}} s(x,y)=σx​σy​+C3​σxy​+C3​​

综合上述三个部分，得到 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)]γ​

其中，