KAIR - 多功能视频与图像增强开源工具箱，涵盖最新深度学习模型

USRNet、DnCNN、FFDNet、SRMD、DPSR、MSRResNet、ESRGAN、BSRGAN、SwinIR、VRT、RVRT 的训练与测试代码

计算机视觉实验室，苏黎世联邦理工学院，瑞士

新闻（2023-06-02）：发布 "用于即插即用图像修复的去噪扩散模型" 的代码 yuanzhi-zhu/DiffPIR。
新闻（2022-10-04）：我们发布了 RVRT 的训练代码，包括 NeurlPS2022 用于视频超分辨率、去模糊和去噪。
新闻（2022-05-05）：尝试 SCUNet 在线演示进行盲图像去噪的在线演示。
新闻（2022-03-23）：我们发布了 SCUNet 的测试代码用于盲图像去噪。

以下结果是我们使用纯合成训练数据的 SCUNet 得到的！在训练期间我们没有使用 DND 和 SIDD 的成对的噪声/干净数据！

新闻（2022-02-15）：我们发布了 VRT 的训练代码，包括 VRT 用于视频超分辨率、去模糊和去噪。
新闻（2021-12-23）：我们的技术在 https://www.amemori.ai/ 中得到应用。
新闻（2021-12-23）：我们的新工作用于实用图像去噪。

新闻（2021-09-09）：添加 main_download_pretrained_models.py 以下载预训练模型。
新闻（2021-09-08）：添加 matlab代码用于放大图像局部区域以比较不同结果。
新闻（2021-09-07）：我们上传了 SwinIR 的训练代码，并提供了针对现实世界图像 SR 的互动在线 Colab 演示。在 Colab 上尝试对您自己的图像进行超分辨率处理！

现实图像 (x4)	BSRGAN，ICCV2021	Real-ESRGAN	SwinIR（我们的）

新闻（2021-08-31）：我们上传了 BSRGAN 的训练代码。
新闻（2021-08-24）：我们上传了 BSRGAN 降解模型。
新闻（2021-08-22）：支持多特征层 VGG 感知损失和 UNet 判别器。
新闻（2021-08-18）：我们上传了扩展版 BSRGAN 降解模型。它稍有不同于我们发布的版本。
新闻（2021-06-03）：添加 GPEN (CVPR21) 用于人脸图像增强的测试代码：main_test_face_enhancement.py

新闻（2021-05-13）：添加 PatchGAN 判别器。
新闻（2021-05-12）：支持分布式训练，另见 https://github.com/xinntao/BasicSR/blob/master/docs/TrainTest.md。
新闻（2021-01）：BSRGAN 用于盲图像超分辨率即将添加。
欢迎提出想法！
更正（2020-10）：如果您使用多个 GPU 进行 GAN 训练，请移除或注释第 105 行，以启用 DataParallel 以实现快速训练。
新闻 (2020-10): 添加 utils_receptivefield.py 计算感受野。
新闻 (2020-8): 发布 深度即插即用图像修复工具箱 cszn/DPIR。
提示 (2020-8): 使用这个避免 内存不足 问题。
新闻 (2020-7): 添加 main_challenge_sr.py 获取 FLOPs(浮点运算次数)、#Params(参数数量)、Runtime(运行时间)、#Activations(激活次数)、#Conv(卷积次数) 和 Max Memory Allocated(最大内存分配)。

from utils.utils_modelsummary import get_model_activation, get_model_flops
input_dim = (3, 256, 256)  # 设置输入维度
activations, num_conv2d = get_model_activation(model, input_dim)
logger.info('{:>16s} : {:<.4f} [M]'.format('#Activations', activations/10**6))
logger.info('{:>16s} : {:<d}'.format('#Conv2d', num_conv2d))
flops = get_model_flops(model, input_dim, False)
logger.info('{:>16s} : {:<.4f} [G]'.format('FLOPs', flops/10**9))
num_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('{:>16s} : {:<.4f} [M]'.format('#Params', num_parameters/10**6))

新闻 (2020-6): 添加 USRNet (CVPR 2020) 进行训练和测试。
- 网络架构
- 数据集

克隆仓库

git clone https://github.com/cszn/KAIR.git

pip install -r requirement.txt

训练

你应该首先修改 options 中的 json 文件，例如，如果使用4个GPU，设置 "gpu_ids": [0,1,2,3] , 如果高质量数据集的路径是 trainsets/trainH，设置 "dataroot_H": "trainsets/trainH"。

使用 DataParallel 进行训练 - PSNR

python main_train_psnr.py --opt options/train_msrresnet_psnr.json

使用 DataParallel 进行训练 - GAN

python main_train_gan.py --opt options/train_msrresnet_gan.json

使用 DistributedDataParallel 进行训练 - PSNR - 4 GPUs

python -m torch.distributed.launch --nproc_per_node=4 --master_port=1234 main_train_psnr.py --opt options/train_msrresnet_psnr.json  --dist True

使用 DistributedDataParallel 进行训练 - PSNR - 8 GPUs

python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 main_train_psnr.py --opt options/train_msrresnet_psnr.json  --dist True

使用 DistributedDataParallel 进行训练 - GAN - 4 GPUs

python -m torch.distributed.launch --nproc_per_node=4 --master_port=1234 main_train_gan.py --opt options/train_msrresnet_gan.json  --dist True

使用 DistributedDataParallel 进行训练 - GAN - 8 GPUs

python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 main_train_gan.py --opt options/train_msrresnet_gan.json  --dist True

杀死 main_train_gan.py 的分布式训练进程

kill $(ps aux | grep main_train_gan.py | grep -v grep | awk '{print $2}')

方法	原始链接
DnCNN	https://github.com/cszn/DnCNN
FDnCNN	https://github.com/cszn/DnCNN
FFDNet	https://github.com/cszn/FFDNet
SRMD	https://github.com/cszn/SRMD
DPSR-SRResNet	https://github.com/cszn/DPSR
SRResNet	https://github.com/xinntao/BasicSR
ESRGAN	https://github.com/xinntao/ESRGAN
RRDB	https://github.com/xinntao/ESRGAN
IMDB	https://github.com/Zheng222/IMDN
USRNet	https://github.com/cszn/USRNet
DRUNet	https://github.com/cszn/DPIR
DPIR	https://github.com/cszn/DPIR
BSRGAN	https://github.com/cszn/BSRGAN
SwinIR	https://github.com/JingyunLiang/SwinIR
VRT	https://github.com/JingyunLiang/VRT
DiffPIR	https://github.com/yuanzhi-zhu/DiffPIR

网络架构

USRNet
DnCNN
IRCNN denoiser

FFDNet
SRMD
SRResNet, SRGAN, RRDB, ESRGAN
IMDN

-----

测试

方法	model_zoo
main_test_dncnn.py	`dncnn_15.pth, dncnn_25.pth, dncnn_50.pth, dncnn_gray_blind.pth, dncnn_color_blind.pth, dncnn3.pth`
main_test_ircnn_denoiser.py	`ircnn_gray.pth, ircnn_color.pth`
main_test_fdncnn.py	`fdncnn_gray.pth, fdncnn_color.pth, fdncnn_gray_clip.pth, fdncnn_color_clip.pth`
main_test_ffdnet.py	`ffdnet_gray.pth, ffdnet_color.pth, ffdnet_gray_clip.pth, ffdnet_color_clip.pth`
main_test_srmd.py	`srmdnf_x2.pth, srmdnf_x3.pth, srmdnf_x4.pth, srmd_x2.pth, srmd_x3.pth, srmd_x4.pth`
	以上模型系由 MatConvNet 转换。
main_test_dpsr.py	`dpsr_x2.pth, dpsr_x3.pth, dpsr_x4.pth, dpsr_x4_gan.pth`
main_test_msrresnet.py	`msrresnet_x4_psnr.pth, msrresnet_x4_gan.pth`
main_test_rrdb.py	`rrdb_x4_psnr.pth, rrdb_x4_esrgan.pth`
main_test_imdn.py	`imdn_x4.pth`

model_zoo

下载链接 https://drive.google.com/drive/folders/13kfr3qny7S2xwG9h7v95F5mkWs0OmU0D 火车集

https://github.com/xinntao/BasicSR/blob/master/docs/DatasetPreparation.md
train400
DIV2K
Flickr2K
可选：使用 split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=512, p_overlap=96, p_max=800) 获得小图像的trainsets/trainH，以加快数据加载速度

测试集

https://github.com/xinntao/BasicSR/blob/master/docs/DatasetPreparation.md
set12
bsd68
cbsd68
kodak24
srbsd68
set5
set14
cbsd100
urban100
manga109

参考文献

@inproceedings{zhu2023denoising,
title={Denoising Diffusion Models for Plug-and-Play Image Restoration},
author={Yuanzhi Zhu and Kai Zhang and Jingyun Liang and Jiezhang Cao and Bihan Wen and Radu Timofte and Luc Van Gool},
booktitle={IEEE Conference on Computer Vision and Pattern Recognition Workshops},
year={2023}
}
@article{liang2022vrt,
title={VRT: A Video Restoration Transformer},
author={Liang, Jingyun and Cao, Jiezhang and Fan, Yuchen and Zhang, Kai and Ranjan, Rakesh and Li, Yawei and Timofte, Radu and Van Gool, Luc},
journal={arXiv preprint arXiv:2022.00000},
year={2022}
}
@inproceedings{liang2021swinir,
title={SwinIR: Image Restoration Using Swin Transformer},
author={Liang, Jingyun and Cao, Jiezhang and Sun, Guolei and Zhang, Kai and Van Gool, Luc and Timofte, Radu},
booktitle={IEEE International Conference on Computer Vision Workshops},
pages={1833--1844},
year={2021}
}
@inproceedings{zhang2021designing,
title={Designing a Practical Degradation Model for Deep Blind Image Super-Resolution},
author={Zhang, Kai and Liang, Jingyun and Van Gool, Luc and Timofte, Radu},
booktitle={IEEE International Conference on Computer Vision},
pages={4791--4800},
year={2021}
}
@article{zhang2021plug,
  title={Plug-and-Play Image Restoration with Deep Denoiser Prior},
  author={Zhang, Kai and Li, Yawei and Zuo, Wangmeng and Zhang, Lei and Van Gool, Luc and Timofte, Radu},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2021}
}
@inproceedings{zhang2020aim,
  title={AIM 2020 Challenge on Efficient Super-Resolution: Methods and Results},
  author={Kai Zhang and Martin Danelljan and Yawei Li and Radu Timofte and others},
  booktitle={European Conference on Computer Vision Workshops},
  year={2020}
}
@inproceedings{zhang2020deep,
  title={Deep unfolding network for image super-resolution},
  author={Zhang, Kai and Van Gool, Luc and Timofte, Radu},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
  pages={3217--3226},
  year={2020}
}
@article{zhang2017beyond,
  title={Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising},
  author={Zhang, Kai and Zuo, Wangmeng and Chen, Yunjin and Meng, Deyu and Zhang, Lei},
  journal={IEEE Transactions on Image Processing},
  volume={26},
  number={7},
  pages={3142--3155},
  year={2017}
}
@inproceedings{zhang2017learning,
title={Learning deep CNN denoiser prior for image restoration},
author={Zhang, Kai and Zuo, Wangmeng and Gu, Shuhang and Zhang, Lei},
booktitle={IEEE conference on computer vision and pattern recognition},
pages={3929--3938},
year={2017}
}
@article{zhang2018ffdnet,
  title={FFDNet: Toward a fast and flexible solution for CNN-based image denoising},
  author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
  journal={IEEE Transactions on Image Processing},
  volume={27},
  number={9},
  pages={4608--4622},
  year={2018}
}
@inproceedings{zhang2018learning,
  title={Learning a single convolutional super-resolution network for multiple degradations},
  author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
  pages={3262--3271},
  year={2018}
}
@inproceedings{zhang2019deep,
  title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
  author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
  pages={1671--1681},
  year={2019}
}
@InProceedings{wang2018esrgan,
    author = {Wang, Xintao and Yu, Ke and Wu, Shixiang and Gu, Jinjin and Liu, Yihao and Dong, Chao and Qiao, Yu and Loy, Chen Change},
    title = {ESRGAN: Enhanced super-resolution generative adversarial networks},
    booktitle = {The European Conference on Computer Vision Workshops (ECCVW)},
    month = {September},
    year = {2018}
}
@inproceedings{hui2019lightweight,
  title={Lightweight Image Super-Resolution with Information Multi-distillation Network},
  author={Hui, Zheng and Gao, Xinbo and Yang, Yunchu and Wang, Xiumei},
  booktitle={Proceedings of the 27th ACM International Conference on Multimedia (ACM MM)},
  pages={2024--2032},
  year={2019}
}
@inproceedings{zhang2019aim,
  title={AIM 2019 Challenge on Constrained Super-Resolution: Methods and Results},
  author={Kai Zhang and Shuhang Gu and Radu Timofte and others},
  booktitle={IEEE International Conference on Computer Vision Workshops},
  year={2019}
}
@inproceedings{yang2021gan,
    title={GAN Prior Embedded Network for Blind Face Restoration in the Wild},
    author={Tao Yang, Peiran Ren, Xuansong Xie, and Lei Zhang},
    booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
    year={2021}
}