Zhidong Zhu, Bangshu Xiong, Zhibo Rao, Jinhao Zhu, Wenchao Li, Qiaofeng Ou, Xing Li
The paper is currently under review.
All the datasets used in the paper can be downloaded from the following locations:
Denoising: BSD400, WED, Urban100, Kodak24, BSD68
Deraining: Train100L&Rain100L
Dehazing: Train RESIDE, Test SOTS-Outdoor
Deblur: GoPro
Enhance: LOL-V1
The training data should be placed in data/Train/{task_name} directory where task_name can be Denoise, Derain, Dehaze, Deblur, or Enhance.
The testing data should be placed in the test directory wherein each task has a seperate directory.
Details can be found in data_dir/DATA_README.md
- Clone our repository
git clone https://github.com/Archaic-Atom/WaveUIR.git
cd WaveUIR
- Create conda environment The Conda environment used can be recreated using the env.yml file
conda env create -f env.yml
After preparing the training data in data/ directory, use
python train.py
to start the training of the model. Use the de_type argument to choose the combination of degradation types to train on. By default it is set to all the 5 degradation tasks (denoising, deraining, dehazing, deblurring, enhancement).
Example Usage: If we only want to train on deraining and dehazing:
python train.py --de_type derain dehaze
After preparing the testing data in test/ directory, place the mode checkpoint file in the ckpt directory. The pre-trained model can be downloaded here. To perform the evaluation, use
python test.py --mode {n}
n is a number that can be used to set the tasks to be evaluated on, 0 for denoising, 1 for deraining, 2 for dehazing, 3 for deblurring, 4 for enhancement, 5 for three-degradation all-in-one setting and 6 for five-degradation all-in-one setting.
Example Usage: To test on all the degradation types at once, run:
python test.py --mode 6
Performance results of the WaveUIR framework trained under the all-in-one setting.
Three Distinct Degradations (click to expand)
Five Distinct Degradations (click to expand)
The visual results can be downloaded here.
Should you have any questions, please contact [email protected]
Acknowledgment: This code is based on the PromptIR repository.