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 3If 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 1The training of WS can be implemented by
python train_selector.py --stage 2Acknowledgement
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