基于 Swin Transformer 的图像超分辨率模型

1. 引入

• 随着 Vision Transformer 的流行，越来越多计算机视觉任务开始采用 Transformer 来实现，也都取得了不错的效果

• 本次就来分享两个基于 Swin Transformer 的图像恢复 / 图像超分辨率模型

2. 效果演示

原图	SwinIR-M 2倍

原图	SwinIR-M 4倍

原图	SwinIR-L 4倍

原图	Swin2SR 4倍

3. 快速体验

• 可以使用 PaddleHub 快速调用这两个图像超分辨率预训练模型

• 可用的模型列表如下：

  模型名称	网络结构	放大倍率
  swin2sr_real_sr_x4	Swin2SR	4
  swinir_m_real_sr_x2	SwinIR-M	2
  swinir_m_real_sr_x4	SwinIR-M	4
  swinir_l_real_sr_x4	SwinIR-L	4

• 注：随着输入图像的分辨率和放大倍率的提高，运行所需的显存 / 内存占用会随之增大。

# hub run 模型名称 --input_path 输入图片路径 --output_dir 输出文件目录
!hub run swin2sr_real_sr_x4 --input_path images/test.jpeg --output_dir outputs

import paddlehub as hub

# 加载模型
module = hub.Module(name="swin2sr_real_sr_x4")

# 图像超分辨率
result = module.real_sr(
    image='images/test.jpeg', # 输入图片路径
    visualization=True, # 是否可视化
    output_dir='outputs' # 输出文件目录
)

4. 参考资料

• 论文：

  • SwinIR: Image Restoration Using Swin Transformer

  • Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration

• 参考实现：

  • JingyunLiang/SwinIR

  • mv-lab/swin2sr

5. 模型结构

• 两个模型的大致结构如下图所示：

  

  SwinIR

  

  Swin2SR

• 总体结构与绝大多数图像恢复或图像超分辨率模型区别不是很大，只是将 Swin Transformer 融入进模型的特征提取与转换网络中。

• 更多具体的细节可以参考论文或其官方代码实现。

6. 代码实现

• 由于模型代码长度太长，不太便于展示在代码块中，具体的代码请参考 swinsr 目录中的文件。

import cv2
import paddle

from swinsr import preprocess, postprocess
from swinsr import swin2sr_real_sr_x4, swinir_l_real_sr_x4, swinir_m_real_sr_x2, swinir_m_real_sr_x4


model_dict = {
    'swin2sr_real_sr_x4': {
        'model': swin2sr_real_sr_x4,
        'ckpt': 'swinsr/pretrained_models/Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR.pdparams'
    },
    'swinir_l_real_sr_x4': {
        'model': swinir_l_real_sr_x4,
        'ckpt': 'swinsr/pretrained_models/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pdparams'
    },
    'swinir_m_real_sr_x4': {
        'model': swinir_m_real_sr_x4,
        'ckpt': 'swinsr/pretrained_models/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pdparams'
    },
    'swinir_m_real_sr_x2': {
        'model': swinir_m_real_sr_x2,
        'ckpt': 'swinsr/pretrained_models/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pdparams'
    }
}


class SwinSR:
    def __init__(self, name='swin2sr_real_sr_x4'):
        self.model = model_dict[name]['model'](ckpt=model_dict[name]['ckpt'])
        self.model.eval()
    
    def __call__(self, image_path, output_path):
        image = cv2.imread(image_path)

        with paddle.no_grad():
            img_input = preprocess(image)
            img_input = paddle.to_tensor(img_input[None, ...], dtype=paddle.float32)

            img_output = self.model(img_input)
            img_output = img_output.numpy()[0]
            img_output = postprocess(img_output)

        cv2.imwrite(output_path, img_output)

# 加载模型
model = SwinSR(name='swin2sr_real_sr_x4') # 模型名称

# 图像超分辨率
model(
    image_path='images/test.jpeg', # 输入图像路径
    output_path='outputs/swin2sr_real_sr_x4.jpg' # 输出图像路径
)

7. 小结

• 简单介绍了两个基于 Swin Transformer 的图像恢复 / 图像超分辨率模型

此文章为搬运
原项目链接