KAIR

Training and testing codes for USRNet, DnCNN, FFDNet, SRMD, DPSR, MSRResNet, ESRGAN, BSRGAN, SwinIR, VRT, RVRT

Kai Zhang

Computer Vision Lab, ETH Zurich, Switzerland

---

• News (2023-06-02): Code for "Denoising Diffusion Models for Plug-and-Play Image Restoration" is released at yuanzhi-zhu/DiffPIR.

• News (2022-10-04): We release the training codes of RVRT, NeurlPS2022 for video SR, deblurring and denoising.

• News (2022-05-05): Try the online demo of SCUNet for blind real image denoising.

• News (2022-03-23): We release the testing codes of SCUNet for blind real image denoising.

The following results are obtained by our SCUNet with purely synthetic training data! We did not use the paired noisy/clean data by DND and SIDD during training!

<img src="https://github.com/cszn/cszn.github.io/blob/master/files/input_16.gif" width="360px"/> <img src="https://github.com/cszn/cszn.github.io/blob/master/files/wm_fnb_0010_16.gif" width="360px"/>

• News (2022-02-15): We release the training codes of VRT for video SR, deblurring and denoising.

• News (2021-12-23): Our techniques are adopted in https://www.amemori.ai/.

• News (2021-12-23): Our new work for practical image denoising.

• <img src="figs/palace.png" height="320px"/> <img src="figs/palace_HSCU.png" height="320px"/>

• <img src="https://github.com/cszn/KAIR/raw/master/figs/denoising_02.png" height="256px"/> <img src="https://github.com/cszn/KAIR/raw/master/figs/denoising_01.png" height="256px"/>

• News (2021-09-09): Add main_download_pretrained_models.py to download pre-trained models.

• News (2021-09-08): Add matlab code to zoom local part of an image for the purpose of comparison between different results.

• News (2021-09-07): We upload the training code of SwinIR and provide an interactive online Colob demo for real-world image SR. Try to super-resolve your own images on Colab! <a href="https://colab.research.google.com/gist/JingyunLiang/a5e3e54bc9ef8d7bf594f6fee8208533/swinir-demo-on-real-world-image-sr.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>

|Real-World Image (x4)|BSRGAN, ICCV2021|Real-ESRGAN|SwinIR (ours)| | :--- | :---: | :-----: | :-----: | |<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/ETH_LR.png">|<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/ETH_BSRGAN.png">|<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/ETH_realESRGAN.jpg">|<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/ETH_SwinIR.png"> |<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/OST_009_crop_LR.png">|<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/OST_009_crop_BSRGAN.png">|<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/OST_009_crop_realESRGAN.png">|<img width="200" src="https://raw.githubusercontent.com/JingyunLiang/SwinIR/main/figs/OST_009_crop_SwinIR.png">|

• News (2021-08-31): We upload the training code of BSRGAN.

• News (2021-08-24): We upload the BSRGAN degradation model.

• News (2021-08-22): Support multi-feature-layer VGG perceptual loss and UNet discriminator.

• News (2021-08-18): We upload the extended BSRGAN degradation model. It is slightly different from our published version.

• News (2021-06-03): Add testing codes of GPEN (CVPR21) for face image enhancement: main_test_face_enhancement.py

<img src="figs/face_04_comparison.png" width="730px"/> <img src="figs/face_13_comparison.png" width="730px"/> <img src="figs/face_08_comparison.png" width="730px"/> <img src="figs/face_01_comparison.png" width="730px"/> <img src="figs/face_12_comparison.png" width="730px"/> <img src="figs/face_10_comparison.png" width="730px"/>

• News (2021-05-13): Add PatchGAN discriminator.

• News (2021-05-12): Support distributed training, see also https://github.com/xinntao/BasicSR/blob/master/docs/TrainTest.md.

• News (2021-01): BSRGAN for blind real image super-resolution will be added.

• Pull requests are welcome!

• Correction (2020-10): If you use multiple GPUs for GAN training, remove or comment Line 105 to enable DataParallel for fast training

• News (2020-10): Add utils_receptivefield.py to calculate receptive field.

• News (2020-8): A deep plug-and-play image restoration toolbox is released at cszn/DPIR.

• Tips (2020-8): Use this to avoid out of memory issue.

• News (2020-7): Add main_challenge_sr.py to get FLOPs, #Params, Runtime, #Activations, #Conv, and Max Memory Allocated.

from utils.utils_modelsummary import get_model_activation, get_model_flops
input_dim = (3, 256, 256)  # set the input dimension
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))

• News (2020-6): Add USRNet (CVPR 2020) for training and testing.
  • Network Architecture
  • Dataset

Clone repo

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

pip install -r requirement.txt

Training

You should modify the json file from options first, for example, setting "gpu_ids": [0,1,2,3] if 4 GPUs are used, setting "dataroot_H": "trainsets/trainH" if path of the high quality dataset is trainsets/trainH.

• Training with DataParallel - PSNR

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

• Training with DataParallel - GAN

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

• Training with 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

• Training with 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

• Training with DistributedDataParallel - GAN - 4 GPUs