Awesome

ConStyle v2: A Strong Prompter for All-in-One Image Restoration

Visual examples

A few visual examples of training data in the pre-training stage:

Mix Degradations datasets

The Mix Degradations datasets and the uncropped joint datasets are avaliable at: https://pan.baidu.com/s/1hk7re1JEBQHVpzta1WhWgA?pwd=5dz9 Code：5dz9

Mix Degradations datasets are placed in the [All] folder and can be directly used for training. If needed, users must download DIV2K datasets and use gen_degradations.py to generate noise and JPEG datasets.

Weight files

All weight files can be found at: https://pan.baidu.com/s/1GYdERj8-hOL3hmywaveMgQ?pwd=hs32 CODE：hs32

[pretrained_ConStyle_v2.pth] can be directly used to improve the performance of Image Restoration models, single or all-in-one, without any fine-tuning.

Environment

python=3.8, pytorch=1.11

Users need to install Wand in your environment for training:https://docs.wand-py.org/en/latest/guide/install.html

How to use

The detail of datasets preparation, training and testing please refer to the docs in BasicSR.

Training:

uncomment the code in data/image_datasets.py for training, lines 22-23, and lines 727-783.

(single) python train.py -opt options/Single/[Maxim, NAFNet, or Restormer]/train/[Deblur, Denoise, or Dehaze].yml

(all-in-one) python train.py -opt options/All-in-One/train/[ConStylev2, train_ConStylev2Model, train_OriginModel].yml

Testing:

(single) python test.py -opt options/Single/[Maxim, NAFNet, or Restormer]/test/[Deblur, Denoise, or Dehaze].yml

(all-in-one) python test.py -opt options/All-in-One/test/[*].yml

Note:

"IRConStyle" model: ConStyle v2 and original model, eg. ConStyle v2+ NAFNet. Where ConStyle v2 is frozen during training.

"ConStyle" model: ConStyle and original model, eg. ConStyle+ NAFNet, will jointly train from scratch, same as IRConStyle.

"ConStyle_v2" model: Pre-train the ConStyle v2 model alone.

"Origin" model: Train the original model alone, without ConStyle or ConStyle v2.

Acknowledgment

This project is based on the BasicSR. The synthesis degradation process in pre-training stage is based on Benchmarking Neural Network Robustness to Common Corruptions and Perturbations and Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data. Thanks for their excellent work!

Contact

If you have any questions, please contact dongqifan@std.uestc.edu.cn.

Citation

If you find our idea and code helpful, please cite our work.