Skip to content

Latest commit

 

History

History
92 lines (61 loc) · 3.87 KB

README.md

File metadata and controls

92 lines (61 loc) · 3.87 KB

WaveUIR: Wavelet-based Guided Transformer Model for Efficient Universal Image Restoration

Zhidong Zhu, Bangshu Xiong, Zhibo Rao, Jinhao Zhu, Wenchao Li, Qiaofeng Ou, Xing Li

The paper is currently under review.

Model Performance Comparison

Data Download and Preparation

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

Environment Installation

  1. Clone our repository
git clone https://github.com/Archaic-Atom/WaveUIR.git
cd WaveUIR
  1. Create conda environment The Conda environment used can be recreated using the env.yml file
conda env create -f env.yml

Training

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

Testing

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

Results

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.

Citation

Contact

Should you have any questions, please contact [email protected]

Acknowledgment: This code is based on the PromptIR repository.