Unified-Width Adaptive Dynamic Network for All-In-One Image Restoration

PyTorch implementation for Unified-Width Adaptive Dynamic Network for All-In-One Image Restoration (U-WADN).

Dependencies

• Python == 3.8.11
• Pytorch == 1.10.0
• mmcv-full == 2.0.0

Dataset

You could find the dataset we used in the paper at following:

Denoising: BSD400, WED, Urban100

Deraining: Train100L&Rain100L

Dehazing: RESIDE (OTS)

Testing

The pretrained model is upload in ckpt/best_ckpt/best.pth. To test with the pretrained model, please:

python test.py --mode 3

If you only want to test one of these tasks, please specific the test mode as 0, 1 or 2. (0 for denoising, 1 for deraining and 2 for dehazing).

Training

If you want to re-train our model, you need to first put the training set into the data/. As the proposed U-WADN has 2 training steps as 1). Training of WAB and 2). Training of WS.

The training of WAB can be implemented by

python train.py --stage 1

The training of WS can be implemented by

python train_selector.py --stage 2

Acknowledgement

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

Personal Information

If you have questions on my work, please feel free to contact me on xuyimin9626@gmail.com