SPDNet
Structure-Preserving Deraining with Residue Channel Prior Guidance (ICCV2021)
Requirements
- Linux Platform
- NVIDIA GPU + CUDA CuDNN
- PyTorch == 0.4.1
- torchvision0.2.0
- pytorch_wavelets
- Python3.6.0
- imageio2.5.0
- numpy1.14.0
- opencv-python
- scikit-image0.13.0
- tqdm4.32.2
- scipy1.2.1
- matplotlib3.1.1
- ipython7.6.1
- h5py2.10.0
Training
- Modify data path in src/data/rainheavy.py and src/data/rainheavytest.py
datapath/data/***.png
datapath/label/***.png - Begining training:
$ cd ./src/
$ python main.py --save spdnet --model spdnet --scale 2 --epochs 300 --batch_size 16 --patch_size 128 --data_train RainHeavy --n_threads 0 --data_test RainHeavyTest --data_range 1-1800/1-200 --loss 1*MSE --save_results --lr 5e-4 --n_feats 32 --n_resblocks 3
Test
The pre-trained model can be available at google drive: https://drive.google.com/drive/folders/1ylON5AkJVayoypOXDaUEkYd76LtMF-lB?usp=sharing.
$ cd ./src/
$ python main.py --data_test RainHeavyTest --ext img --scale 2 --data_range 1-1800/1-200 --pre_train ../experiment/spdnet/model/model_best.pt --model spdnet --test_only --save_results --save RCDNet_test
All PSNR and SSIM results are computed by using this Matlab code, based on Y channel of YCbCr space.
Datasets
Rain200H: 1800 training pairs and 200 testing pairs
Rain200L: 1800 training pairs and 200 testing pairs
Rain800: 700 training pairs and 100 testing pairs
Rain1200: 12000 traing paris and 1200 testing pairs
SPA-Data: 638492 training pairs and 1000 testing pairs
Acknowledgement
Code borrows from RCDNet. Thanks for sharing !