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!

cszhilu1998 / SelfIR