通过Sub-Pixel实现图像超分辨率

    日期: 2021.01

    摘要: 本示例教程使用U-Net实现图像分割。

    在计算机视觉中，图像超分辨率（Image Super Resolution）是指由一幅低分辨率图像或图像序列恢复出高分辨率图像。图像超分辨率技术分为超分辨率复原和超分辨率重建。

    本示例简要介绍如何通过飞桨开源框架，实现图像超分辨率。包括数据集的定义、模型的搭建与训练。

    参考论文：《Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network》

    论文链接：https://arxiv.org/abs/1609.05158

    二、环境设置

    本案例使用BSR_bsds500数据集，下载链接：

    1. !wget --no-check-certificate --no-cookies --header "Cookie: oraclelicense=accept-securebackup-cookie" http://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz
    2. !tar -zxvf BSR_bsds500.tgz
    1. --2021-01-29 00:11:52--  http://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz
    2. Resolving www.eecs.berkeley.edu (www.eecs.berkeley.edu)... 23.185.0.1, 2620:12a:8001::1, 2620:12a:8000::1
    3. Connecting to www.eecs.berkeley.edu (www.eecs.berkeley.edu)|23.185.0.1|:80... connected.
    4. HTTP request sent, awaiting response... 301 Moved Permanently
    5. Location: https://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz [following]
    6. --2021-01-29 00:11:53--  https://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz
    7. Connecting to www.eecs.berkeley.edu (www.eecs.berkeley.edu)|23.185.0.1|:443... connected.
    8. HTTP request sent, awaiting response... 301 Moved Permanently
    9. Location: https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz [following]
    10. --2021-01-29 00:11:53--  https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz
    11. Resolving www2.eecs.berkeley.edu (www2.eecs.berkeley.edu)... 128.32.244.190
    12. Connecting to www2.eecs.berkeley.edu (www2.eecs.berkeley.edu)|128.32.244.190|:443... connected.
    13. ^C
    14. tar (child): BSR_bsds500.tgz: Cannot open: No such file or directory
    15. tar (child): Error is not recoverable: exiting now
    16. tar: Child returned status 2
    17. tar: Error is not recoverable: exiting now

    3.2 数据集概览

    1. BSR
    2. ├── BSDS500
    3. │   └── data
    4. │       ├── groundTruth
    5. │       │   ├── test
    6. │       │   ├── train
    7. │       │   └── val
    8. │       └── images
    9. │           ├── test
    10. │           ├── train
    11. │           └── val
    12. ├── bench
    13. │   ├── benchmarks
    14. │   ├── data
    15. │   │   ├── ...
    16. │   │   └── ...
    17. │   └── source
    18. └── documentation

    可以看到我们需要的图片文件在BSR/BSDS500/images文件夹下，train、test各200张，val为100张。

    3.3 数据集类定义

    飞桨（PaddlePaddle）数据集加载方案是统一使用Dataset（数据集定义） + DataLoader（多进程数据集加载）。

    首先我们先进行数据集的定义，数据集定义主要是实现一个新的Dataset类，继承父类paddle.io.Dataset，并实现父类中以下两个抽象方法，__getitem__和len：

    1. class MyDataset(Dataset):
    2.     def __init__(self):
    3.         ...
    5.     # 每次迭代时返回数据和对应的标签
    6.     def __getitem__(self, idx):
    7.         return x, y
    9.     # 返回整个数据集的总数
    10.     def __len__(self):
    11.         return count(samples)
    1. class BSD_data(Dataset):
    2.     """
    3.     继承paddle.io.Dataset类
    4.     """
    5.     def __init__(self, mode='train',image_path="BSR/BSDS500/data/images/"):
    6.         """
    7.         实现构造函数，定义数据读取方式，划分训练和测试数据集
    8.         """
    9.         super(BSD_data, self).__init__()
    11.         self.mode = mode.lower()
    12.         if self.mode == 'train':
    13.             self.image_path = os.path.join(image_path,'train')
    14.         elif self.mode == 'val':
    15.             self.image_path = os.path.join(image_path,'val')
    16.         else:
    17.             raise ValueError('mode must be "train" or "val"')
    19.         # 原始图像的缩放大小
    20.         self.crop_size = 300
    21.         # 缩放倍率
    22.         self.upscale_factor = 3
    23.         # 缩小后送入神经网络的大小
    24.         self.input_size = self.crop_size // self.upscale_factor
    25.         # numpy随机数种子
    26.         self.seed=1337
    27.         # 图片集合
    28.         self.temp_images = []
    29.         # 加载数据
    30.         self._parse_dataset()
    32.     def transforms(self, img):
    33.         """
    34.         图像预处理工具，用于将升维(100, 100) => (100, 100,1)，
    35.         并对图像的维度进行转换从HWC变为CHW
    36.         """
    37.         if len(img.shape) == 2:
    38.             img = np.expand_dims(img, axis=2)
    39.         return img.transpose((2, 0, 1))
    41.     def __getitem__(self, idx):
    42.         """
    43.         返回 缩小3倍后的图片 和 原始图片
    44.         """
    46.         # 加载原始图像
    47.         img = self._load_img(self.temp_images[idx])
    48.         # 将原始图像缩放到（3, 300, 300）
    49.         img = img.resize([self.crop_size,self.crop_size], Image.BICUBIC)
    51.         #转换为YCbCr图像
    52.         ycbcr = img.convert("YCbCr")
    54.         # 因为人眼对亮度敏感，所以只取Y通道
    55.         y, cb, cr = ycbcr.split()
    56.         y = np.asarray(y,dtype='float32')
    57.         y = y / 255.0
    59.         img_ = img.resize([self.input_size,self.input_size], Image.BICUBIC)
    60.         ycbcr_ = img_.convert("YCbCr")
    61.         y_, cb_, cr_ = ycbcr_.split()
    62.         y_ = np.asarray(y_,dtype='float32')
    63.         y_ = y_ / 255.0
    65.         # 升纬并将HWC转换为CHW
    66.         y = self.transforms(y)
    67.         x = self.transforms(y_)
    69.         # x为缩小3倍后的图片（1, 100, 100） y是原始图片（1, 300, 300）
    73.     def __len__(self):
    74.         """
    75.         实现__len__方法，返回数据集总数目
    76.         """
    77.         return len(self.temp_images)
    79.     def _sort_images(self, img_dir):
    80.         """
    81.         对文件夹内的图像进行按照文件名排序
    82.         """
    83.         files = []
    85.         for item in os.listdir(img_dir):
    86.             if item.split('.')[-1].lower() in ["jpg",'jpeg','png']:
    87.                 files.append(os.path.join(img_dir, item))
    89.         return sorted(files)
    91.     def _parse_dataset(self):
    92.         """
    93.         处理数据集
    94.         """
    95.         self.temp_images = self._sort_images(self.image_path)
    96.         random.Random(self.seed).shuffle(self.temp_images)
    98.     def _load_img(self, path):
    99.         """
    100.         从磁盘读取图片
    101.         """
    102.         with open(path, 'rb') as f:
    103.             img = Image.open(io.BytesIO(f.read()))
    104.             img = img.convert('RGB')
    105.             return img

    实现好BSD_data数据集后，我们来测试一下数据集是否符合预期，因为BSD_data是一个可以被迭代的Class，我们通过for循环从里面读取数据进行展示。

    1. # 测试定义的数据集
    2. train_dataset = BSD_data(mode='train')
    3. val_dataset = BSD_data(mode='val')
    5. print('=============train dataset=============')
    6. x, y = train_dataset[0]
    7. x = x[0]
    8. y = y[0]
    9. x = x * 255
    10. y = y * 255
    11. img_ = Image.fromarray(np.uint8(x), mode="L")
    12. img = Image.fromarray(np.uint8(y), mode="L")
    13. display(img_)
    14. display(img_.size)
    15. display(img)
    16. display(img.size)

    1. (100, 100)
    1. (300, 300)

    四、模型组网

    Sub_Pixel_CNN是一个全卷积网络，网络结构比较简单，这里采用Layer类继承方式组网。

    1. class Sub_Pixel_CNN(paddle.nn.Layer):
    3.     def __init__(self, upscale_factor=3, channels=1):
    4.         super(Sub_Pixel_CNN, self).__init__()
    6.         self.conv1 = paddle.nn.Conv2D(channels,64,5,stride=1, padding=2)
    7.         self.conv2 = paddle.nn.Conv2D(64,64,3,stride=1, padding=1)
    8.         self.conv3 = paddle.nn.Conv2D(64,32,3,stride=1, padding=1)
    9.         self.conv4 = paddle.nn.Conv2D(32,channels * (upscale_factor ** 2),3,stride=1, padding=1)
    11.     def forward(self, x):
    12.         x = self.conv1(x)
    13.         x = self.conv2(x)
    14.         x = self.conv3(x)
    15.         x = self.conv4(x)
    16.         x = paddle.nn.functional.pixel_shuffle(x,3)
    17.         return x

    4.1 模型封装

    1. # 模型封装
    2. model = paddle.Model(Sub_Pixel_CNN())

    4.2 模型可视化

    调用飞桨提供的summary接口对组建好的模型进行可视化，方便进行模型结构和参数信息的查看和确认。

    1. model.summary((1,1, 100, 100))
    1. ---------------------------------------------------------------------------
    2.  Layer (type)       Input Shape          Output Shape         Param #
    3. ===========================================================================
    4.    Conv2D-1      [[1, 1, 100, 100]]   [1, 64, 100, 100]        1,664
    5.    Conv2D-2     [[1, 64, 100, 100]]   [1, 64, 100, 100]       36,928
    6.    Conv2D-3     [[1, 64, 100, 100]]   [1, 32, 100, 100]       18,464
    7.    Conv2D-4     [[1, 32, 100, 100]]    [1, 9, 100, 100]        2,601
    8. ===========================================================================
    9. Total params: 59,657
    10. Trainable params: 59,657
    11. Non-trainable params: 0
    12. ---------------------------------------------------------------------------
    13. Input size (MB): 0.04
    14. Forward/backward pass size (MB): 12.89
    15. Params size (MB): 0.23
    16. Estimated Total Size (MB): 13.16
    17. ---------------------------------------------------------------------------
    1. {'total_params': 59657, 'trainable_params': 59657}

    使用模型代码进行Model实例生成，使用prepare接口定义优化器、损失函数和评价指标等信息，用于后续训练使用。在所有初步配置完成后，调用fit接口开启训练执行过程，调用fit时只需要将前面定义好的训练数据集、测试数据集、训练轮次（Epoch）和批次大小（batch_size）配置好即可。

    1. The loss value printed in the log is the current step, and the metric is the average value of previous step.
    2. Epoch 1/20
    3. step 13/13 [==============================] - loss: 0.2466 - 2s/step
    4. Epoch 2/20
    5. step 13/13 [==============================] - loss: 0.0802 - 2s/step
    6. Epoch 3/20
    7. step 13/13 [==============================] - loss: 0.0474 - 2s/step
    8. Epoch 4/20
    9. step 13/13 [==============================] - loss: 0.0340 - 2s/step
    10. Epoch 5/20
    11. step 13/13 [==============================] - loss: 0.0267 - 2s/step
    12. Epoch 6/20
    13. step 13/13 [==============================] - loss: 0.0179 - 2s/step
    14. Epoch 7/20
    15. step 13/13 [==============================] - loss: 0.0215 - 2s/step
    16. Epoch 8/20
    17. step 13/13 [==============================] - loss: 0.0162 - 2s/step
    18. Epoch 9/20
    19. step 13/13 [==============================] - loss: 0.0137 - 2s/step
    20. Epoch 10/20
    21. step 13/13 [==============================] - loss: 0.0099 - 2s/step
    22. step 13/13 [==============================] - loss: 0.0074 - 2s/step
    23. Epoch 12/20
    24. step 13/13 [==============================] - loss: 0.0117 - 2s/step
    25. Epoch 13/20
    26. step 13/13 [==============================] - loss: 0.0065 - 2s/step
    28. step 13/13 [==============================] - loss: 0.0086 - 2s/step
    29. Epoch 15/20
    30. step 13/13 [==============================] - loss: 0.0085 - 2s/step
    31. Epoch 16/20
    32. step 13/13 [==============================] - loss: 0.0067 - 2s/step
    33. Epoch 17/20
    34. step 13/13 [==============================] - loss: 0.0068 - 2s/step
    35. Epoch 18/20
    36. step 13/13 [==============================] - loss: 0.0044 - 2s/step
    37. Epoch 19/20
    38. step 13/13 [==============================] - loss: 0.0069 - 2s/step
    39. Epoch 20/20
    40. step 13/13 [==============================] - loss: 0.0087 - 2s/step

    六、模型预测

    6.1 预测

    我们可以直接使用model.predict接口来对数据集进行预测操作，只需要将预测数据集传递到接口内即可。

    1. predict_results = model.predict(val_dataset)
    1. Predict begin...
    2. step 100/100 [==============================] - 38ms/step
    3. Predict samples: 100

    6.2 定义预测结果可视化函数

    1. import math
    2. import matplotlib.pyplot as plt
    3. from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
    4. from mpl_toolkits.axes_grid1.inset_locator import mark_inset
    6. def psnr(img1, img2):
    7.     """
    8.     PSMR计算函数
    9.     """
    10.     mse = np.mean( (img1/255. - img2/255.) ** 2 )
    11.     if mse < 1.0e-10:
    12.         return 100
    13.     PIXEL_MAX = 1
    14.     return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
    16. def plot_results(img, title='results', prefix='out'):
    17.     """
    18.     画图展示函数
    19.     """
    20.     img_array = np.asarray(img, dtype='float32')
    21.     img_array = img_array.astype("float32") / 255.0
    23.     fig, ax = plt.subplots()
    24.     im = ax.imshow(img_array[::-1], origin="lower")
    26.     plt.title(title)
    27.     axins = zoomed_inset_axes(ax, 2, loc=2)
    28.     axins.imshow(img_array[::-1], origin="lower")
    30.     x1, x2, y1, y2 = 200, 300, 100, 200
    31.     axins.set_xlim(x1, x2)
    32.     axins.set_ylim(y1, y2)
    34.     plt.yticks(visible=False)
    35.     plt.xticks(visible=False)
    37.     mark_inset(ax, axins, loc1=1, loc2=3, fc="none", ec="blue")
    38.     plt.savefig(str(prefix) + "-" + title + ".png")
    39.     plt.show()
    41. def get_lowres_image(img, upscale_factor):
    42.     """
    43.     缩放图片
    44.     """
    45.     return img.resize(
    46.         (img.size[0] // upscale_factor, img.size[1] // upscale_factor),
    47.         Image.BICUBIC,
    48.     )
    50. def upscale_image(model, img):
    51.     '''
    52.     输入小图，返回上采样三倍的大图像
    53.     '''
    54.     # 把图片复转换到YCbCr格式
    55.     ycbcr = img.convert("YCbCr")
    56.     y, cb, cr = ycbcr.split()
    57.     y = np.asarray(y, dtype='float32')
    58.     y = y / 255.0
    59.     img = np.expand_dims(y, axis=0) # 升维度到（1,w,h）一张image
    60.     img = np.expand_dims(img, axis=0) # 升维度到（1,1,w,h）一个batch
    61.     img = np.expand_dims(img, axis=0) # 升维度到（1,1,1,w,h）可迭代的batch
    63.     out = model.predict(img) # predict输入要求为可迭代的batch
    65.     out_img_y = out[0][0][0] # 得到predict输出结果
    66.     out_img_y *= 255.0
    68.     # 把图片复转换回RGB格式
    69.     out_img_y = out_img_y.reshape((np.shape(out_img_y)[1], np.shape(out_img_y)[2]))
    70.     out_img_y = Image.fromarray(np.uint8(out_img_y), mode="L")
    71.     out_img_cb = cb.resize(out_img_y.size, Image.BICUBIC)
    72.     out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)
    73.     out_img = Image.merge("YCbCr", (out_img_y, out_img_cb, out_img_cr)).convert(
    74.         "RGB"
    75.     )
    76.     return out_img
    78. def main(model, img, upscale_factor=3):
    79.     # 读取图像
    80.     with open(img, 'rb') as f:
    81.         img = Image.open(io.BytesIO(f.read()))
    82.     # 缩小三倍
    83.     lowres_input = get_lowres_image(img, upscale_factor)
    84.     w = lowres_input.size[0] * upscale_factor
    85.     h = lowres_input.size[1] * upscale_factor
    86.     # 将缩小后的图片再放大三倍
    87.     lowres_img = lowres_input.resize((w, h))
    88.     # 确保未经缩放的图像和其他两张图片大小一致
    89.     highres_img = img.resize((w, h))
    90.     # 得到缩小后又经过 Efficient Sub-Pixel CNN放大的图片
    91.     prediction = upscale_image(model, lowres_input)
    92.     psmr_low = psnr(np.asarray(lowres_img), np.asarray(highres_img))
    93.     psmr_pre = psnr(np.asarray(prediction), np.asarray(highres_img))
    94.     # 展示三张图片
    95.     plot_results(lowres_img, "lowres")
    96.     plot_results(highres_img, "highres")
    97.     plot_results(prediction, "prediction")
    98.     print("psmr_low:", psmr_low, "psmr_pre:", psmr_pre)

    从我们的预测数据集中抽1个张图片来看看预测的效果，展示一下原图、小图和预测结果。

    1. main(model,'BSR/BSDS500/data/images/test/100007.jpg')
    1. Predict begin...
    2. step 1/1 [==============================] - 75ms/step

    ../../../_images/super_resolution_sub_pixel_26_1.png ../../../_images/super_resolution_sub_pixel_26_3.png

    1. psmr_low: 30.381882136539197 psmr_pre: 29.074438702896636

    7.模型保存