SwinIR实战：如何使用SwinIR和预训练模型实现图片的超分

作者：AI小浩@公众号
编辑：3D视觉开发者社区

SWinIR实战

论文地址：https://arxiv.org/pdf/2108.10257.pdf

训练代码下载：https://github.com/cszn/KAIR

测试：https://github.com/JingyunLiang/SwinIR

预训练模型下载：

https://github.com/JingyunLiang/SwinIR/releases

论文翻译：https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/124434886

本文介绍如何使用SwinIR和预训练模型实现图像的超分。

测试

下载代码：

https://github.com/JingyunLiang/SwinIR，然后解压，如下图：

然后下载预训练权重文件，打开连接：https://github.com/JingyunLiang/SwinIR/releases，如下图：

预训练权重有很多，选择自己用到的预训练权重下载，下载后将其放入./model_zoo/swinir文件夹下面，如下图：

然后执行命令：

python main_test_swinir.py --task classical_sr --scale 2 --training_patch_size 48 --model_path model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x2.pth --folder_lq testsets/Set5/LR_bicubic/X2 --folder_gt testsets/Set5/HR

task：任务类型，包括：classical_sr, lightweight_sr, real_sr, gray_dn, color_dn, jpeg_car。

scale：放大的倍数，这个和预训练权重对应。比如这次设置的是2，对应权重的x2。

training_patch_size：用于训练 SwinIR 的patch大小。

model_path：预训练模型的路径。

folder_lq：低质量图片的文件夹路径。

folder_gt：高质量图片的文件夹路径。

更多的命令如下：

# 001 Classical Image Super-Resolution (middle size)
# Note that --training_patch_size is just used to differentiate two different settings in Table 2 of the paper. Images are NOT tested patch by patch.
# (setting1: when model is trained on DIV2K and with training_patch_size=48)
python main_test_swinir.py --task classical_sr --scale 2 --training_patch_size 48 --model_path model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x2.pth --folder_lq testsets/Set5/LR_bicubic/X2 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task classical_sr --scale 3 --training_patch_size 48 --model_path model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x3.pth --folder_lq testsets/Set5/LR_bicubic/X3 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task classical_sr --scale 4 --training_patch_size 48 --model_path model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x4.pth --folder_lq testsets/Set5/LR_bicubic/X4 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task classical_sr --scale 8 --training_patch_size 48 --model_path model_zoo/swinir/001_classicalSR_DIV2K_s48w8_SwinIR-M_x8.pth --folder_lq testsets/Set5/LR_bicubic/X8 --folder_gt testsets/Set5/HR

# (setting2: when model is trained on DIV2K+Flickr2K and with training_patch_size=64)
python main_test_swinir.py --task classical_sr --scale 2 --training_patch_size 64 --model_path model_zoo/swinir/001_classicalSR_DF2K_s64w8_SwinIR-M_x2.pth --folder_lq testsets/Set5/LR_bicubic/X2 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task classical_sr --scale 3 --training_patch_size 64 --model_path model_zoo/swinir/001_classicalSR_DF2K_s64w8_SwinIR-M_x3.pth --folder_lq testsets/Set5/LR_bicubic/X3 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task classical_sr --scale 4 --training_patch_size 64 --model_path model_zoo/swinir/001_classicalSR_DF2K_s64w8_SwinIR-M_x4.pth --folder_lq testsets/Set5/LR_bicubic/X4 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task classical_sr --scale 8 --training_patch_size 64 --model_path model_zoo/swinir/001_classicalSR_DF2K_s64w8_SwinIR-M_x8.pth --folder_lq testsets/Set5/LR_bicubic/X8 --folder_gt testsets/Set5/HR


# 002 Lightweight Image Super-Resolution (small size)
python main_test_swinir.py --task lightweight_sr --scale 2 --model_path model_zoo/swinir/002_lightweightSR_DIV2K_s64w8_SwinIR-S_x2.pth --folder_lq testsets/Set5/LR_bicubic/X2 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task lightweight_sr --scale 3 --model_path model_zoo/swinir/002_lightweightSR_DIV2K_s64w8_SwinIR-S_x3.pth --folder_lq testsets/Set5/LR_bicubic/X3 --folder_gt testsets/Set5/HR
python main_test_swinir.py --task lightweight_sr --scale 4 --model_path model_zoo/swinir/002_lightweightSR_DIV2K_s64w8_SwinIR-S_x4.pth --folder_lq testsets/Set5/LR_bicubic/X4 --folder_gt testsets/Set5/HR


# 003 Real-World Image Super-Resolution (use --tile 400 if you run out-of-memory)
# (middle size)
python main_test_swinir.py --task real_sr --scale 4 --model_path model_zoo/swinir/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth --folder_lq testsets/RealSRSet+5images --tile

# (larger size + trained on more datasets)
python main_test_swinir.py --task real_sr --scale 4 --large_model --model_path model_zoo/swinir/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth --folder_lq testsets/RealSRSet+5images


# 004 Grayscale Image Deoising (middle size)
python main_test_swinir.py --task gray_dn --noise 15 --model_path model_zoo/swinir/004_grayDN_DFWB_s128w8_SwinIR-M_noise15.pth --folder_gt testsets/Set12
python main_test_swinir.py --task gray_dn --noise 25 --model_path model_zoo/swinir/004_grayDN_DFWB_s128w8_SwinIR-M_noise25.pth --folder_gt testsets/Set12
python main_test_swinir.py --task gray_dn --noise 50 --model_path model_zoo/swinir/004_grayDN_DFWB_s128w8_SwinIR-M_noise50.pth --folder_gt testsets/Set12


# 005 Color Image Deoising (middle size)
python main_test_swinir.py --task color_dn --noise 15 --model_path model_zoo/swinir/005_colorDN_DFWB_s128w8_SwinIR-M_noise15.pth --folder_gt testsets/McMaster
python main_test_swinir.py --task color_dn --noise 25 --model_path model_zoo/swinir/005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth --folder_gt testsets/McMaster
python main_test_swinir.py --task color_dn --noise 50 --model_path model_zoo/swinir/005_colorDN_DFWB_s128w8_SwinIR-M_noise50.pth --folder_gt testsets/McMaster


# 006 JPEG Compression Artifact Reduction (middle size, using window_size=7 because JPEG encoding uses 8x8 blocks)
python main_test_swinir.py --task jpeg_car --jpeg 10 --model_path model_zoo/swinir/006_CAR_DFWB_s126w7_SwinIR-M_jpeg10.pth --folder_gt testsets/classic5
python main_test_swinir.py --task jpeg_car --jpeg 20 --model_path model_zoo/swinir/006_CAR_DFWB_s126w7_SwinIR-M_jpeg20.pth --folder_gt testsets/classic5
python main_test_swinir.py --task jpeg_car --jpeg 30 --model_path model_zoo/swinir/006_CAR_DFWB_s126w7_SwinIR-M_jpeg30.pth --folder_gt testsets/classic5
python main_test_swinir.py --task jpeg_car --jpeg 40 --model_path model_zoo/swinir/006_CAR_DFWB_s126w7_SwinIR-M_jpeg40.pth --folder_gt testsets/classic5

Result

那么如何用起来呢？

我们使用classical_sr任务讲解。

首先将评分的部分注释掉。注释main_test_swinir.py中的83行到119行。

将185行的folder_gt改为folder_lq

修改218行到223行的代码。由

# 001 classical image sr/ 002 lightweight image sr (load lq-gt image pairs)
if args.task in ['classical_sr', 'lightweight_sr']:
    img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
    img_lq = cv2.imread(f'{args.folder_lq}/{imgname}x{args.scale}{imgext}', cv2.IMREAD_COLOR).astype(
        np.float32) / 255.

修改为：

# 001 classical image sr/ 002 lightweight image sr (load lq-gt image pairs)
 if args.task in ['classical_sr', 'lightweight_sr']:
 # img_gt = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
 img_gt = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
 img_gt = cv2.cvtColor(img_gt, cv2.IMREAD_COLOR).astype(np.float32) / 255.
 lq_path=f'{args.folder_lq}/{imgname}{imgext}'
 img_lq = cv2.imdecode(np.fromfile(lq_path, dtype=np.uint8), -1)
 img_lq = cv2.cvtColor(img_lq, cv2.IMREAD_COLOR).astype(np.float32) / 255.

使用cv2.imdecode方法读取图片，可以解决路径有中文的问题。

将81行的代码

 cv2.imwrite(f'{save_dir}/{imgname}_SwinIR.png', output)

修改为：

  write_path=f'{save_dir}/{imgname}.jpg'
  cv2.imencode('.jpg', output)[1].tofile(write_path)  # 保存图片

这样，生成的图片和原来低质量的图片名字保持一致了。

完成上面的修改，执行命令：

  python main_test_swinir.py --task classical_sr --scale 2 --training_patch_size 64 --model_path model_zoo/swinir/001_classicalsr_df2k_s64w8_swinir-m_x2.pth --folder_lq testsets/Set5/LR_bicubic/X2

等待命令运行完成后，我们就在result文件夹下面看到超分后的图像了。

对比结果：

我在对这些图片超分后，然后执行分类任务测试。在实际的测试中，和超分的图片作对比有了0.1%的提升，提升成绩有限。

我推测：超分后，只是让我们人眼觉得图像的画质提高了，但是对于神经网络来说，有用的特征并没有得到提升。图像在输入神经网络后，一般要经过Resize，超分前的图像尺寸接近224×224，超分后，提升了像素，但是输入进网络的时候又Resize成224×224，这样就导致了有效的特征被压缩了，所以提升有限。

完整代码：https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85251176

版权声明：本文为奥比中光3D视觉开发者社区特约作者授权原创发布，未经授权不得转载，本文仅做学术分享，版权归原作者所有，若涉及侵权内容请联系删文。

3D视觉开发者社区是由奥比中光给所有开发者打造的分享与交流平台，旨在将3D视觉技术开放给开发者。平台为开发者提供3D视觉领域免费课程、奥比中光独家资源与专业技术支持。点击加入3D视觉开发者社区，和开发者们一起讨论分享吧~

或可微信关注官方公众号 3D视觉开发者社区 ，获取更多干货知识哦。