SelfIR (NeurIPS 2022)

PyTorch implementation of Self-Supervised Image Restoration with Blurry and Noisy Pairs

OpenReview   |   arXiv   |   video

1. Framework

Overview of our proposed SelfIR framework.

• (a) Training phase of SelfIR. Sub-sampled blurry image $g_1(I_B)$ and noisy image $g_1(I_N)$ are taken as the inputs. $g_2(I_N)$ is used for calculating the reconstruction loss $\zeta_{rec}$ and regularization loss $\zeta_{reg}$, while $g_1(I_B)$ is taken for calculating auxiliary loss $\zeta_{aux}$.
• (b) Example of neighbor sub-sampler. In each $2\times2$ cell, two pixels are randomly selected for respectively composing the neighboring sub-images.
• (c) Testing phase of SelfIR. The blurry and noisy images can be directly taken for restoration.

2. Preparation and Datasets

• Prerequisites

  • Python 3.x and PyTorch 1.6.
  • OpenCV, NumPy, Pillow, tqdm, lpips, scikit-image and tensorboardX.

• Dataset

  • GoPro dataset can be downloaded from this link.

  • Synthetic noisy images with Gaussian noise for testing can be downloaded from this link. Please decompress the files according to the commands provided in its readme.txt.

  • Synthetic noisy images with Poisson noise for testing can be downloaded from this link. Please decompress the files according to the commands provided in its readme.txt.

  • The directory structure of the dataset:

    GOPRO_Large 
    │
    └───train
    │        GOPR0372_07_00
    │        GOPR0372_07_01
    |        ...  
    │   
    └───test
    │        GOPR0372_07_00
    │        GOPR0372_07_01
    |        ... 
    │   
    └───test_noise_gauss5_50
    │        GOPR0372_07_00
    │        GOPR0372_07_01
    |        ... 
    │   
    └───test_noise_poisson5_50
             GOPR0372_07_00
             GOPR0372_07_01
             ... 
    

3. Quick Start

3.1 Pre-trained models

• For Gaussian noise, we provide the pre-trained models in the ./ckpt/selfir_gauss_noise/ folder.

• For Poisson noise, we provide the pre-trained models in the ./ckpt/selfir_poisson_noise/ folder.

3.2 Training

• Modify dataroot, type and name in train.sh and then run:

  sh train.sh

3.3 Testing

• Modify dataroot, type, name and iter in test.sh and then run:

  sh test.sh

3.4 Note

• You can specify which GPU to use by --gpu_ids, e.g., --gpu_ids 0,1, --gpu_ids 3, --gpu_ids -1 (for CPU mode). In the default setting, all GPUs are used.
• You can refer to options for more arguments.

4. Citation

If you find it useful in your research, please consider citing:

@inproceedings{SelfIR,
    title={Self-Supervised Image Restoration with Blurry and Noisy Pairs},
    author={Zhang, Zhilu and Xu, Rongjian and Liu, Ming and Yan, Zifei and Zuo, Wangmeng},
    booktitle={NeurIPS},
    year={2022}
}

5. Acknowledgement

This repo is built upon the framework of CycleGAN, and we borrow some code from Neighbor2Neighbor, thanks for their excellent work!