Awesome
Learning Rain Location Prior for Nighttime Deraining
Learning Rain Location Prior for Nighttime Deraining
Fan Zhang, Shaodi You, Yu Li, Ying Fu
ICCV 2023
This repository contains the official implementation and experimental data of the ICCV2023 paper "Learning Rain Location Prior for Nighttime Deraining", by Fan Zhang, Shaodi You, Yu Li, Ying Fu.
Update
- Recollect misaligned data.
- 2023.12.08: Code release.
- 2023.12.03: Initial release of experimental data.
- 2023.08.10: Repo created.
Dataset
The experimental data used in the paper is now publicly available at Kaggle. It is based on GTAV-NightRain dataset and increase the difficulty by enlarging the rain density.
In this new version, we collected 5000 rainy images paired with 500 clean images for the training set, and 500/100 for the test set. Each clean image corresponds to 10/5 rainy images. The image resolution is 1920x1080.
Note
Please note that this is the very data used in the experiments.
However, after checking carefully, we find that there exist a few scenes with misalignments due to operation mistakes during collection. We filter out these scenes and there's about 0.5dB improvement in PSNR, which applys to all evaluated methods.
We plan to re-collect and update these misaligned scenes and provide the updated quantitative results later.
Requirements
- Python 3.6.13
- Pytorch 1.10.2
- Cudatoolkit 11.3
You can refer to Uformer and MPRNet for detailed dependency list. Necessary list will be updated later.
Training
- Download the Dataset on Kaggle or prepare your own training dataset, then modify the
--train_dir
to corresponding directory. - Train the model by simply run
bash train.sh
You can
- Select the Deraining Module (DM) by
--arch
, currently supportingUNet
andUformer_T
. - Enable the Rain Location Prior Module (RLP) by
--use_rlp
. - Enable the Rain Prior Injection Module (RPIM) using
--use_rpim
, which is only considered when RLP is used. - Check other options in
rlp/options.py
.
Evaluation
- Prepare your test images or simply test on the downloaded data, by running
bash test.sh
- Modify
--input_dir
to your/path/to/test/images
and--result_dir
for saving results. - Modify
--weights
to the model checkpoint you have. - Modify
--model_name
following the format ofDM
,DM_RLP
orDM_RLP_RPIM
according to the model, such asUformer_T_RLP_RPIM
whenDM = 'Uformer_T', is_RLP = True, is_RPIM = True
. - Use
--tile
to enable tiling of large images forUformer
.
Metrics
To calculate PSNR and SSIM metrics, you can use the Matlab script
evaluate_PSNR_SSIM.m
or the Python version
python evaluate_PSNR_SSIM.py
The results produced by .py
script are slightly different from the .m
script.
Checkpoints
Model | DM | RLP | RPIM | PSNR | SSIM | Checkpoint |
---|---|---|---|---|---|---|
UNet | ✓ | 36.63 | 0.9693 | UNet.pth | ||
UNet | ✓ | ✓ | 37.08 | 0.9715 | UNet_RLP.pth | |
UNet | ✓ | ✓ | ✓ | 37.28 | 0.9716 | UNet_RLP_RPIM.pth |
Uformer_T | ✓ | 37.45 | 0.9720 | Uformer_T.pth | ||
Uformer_T | ✓ | ✓ | 37.95 | 0.9733 | Uformer_T_RLP.pth | |
Uformer_T | ✓ | ✓ | ✓ | 38.44 | 0.9749 | Uformer_T_RLP_RPIM.pth |
Citation
If you find this repo useful, please give us a star and consider citing our papers:
@inproceedings{zhang2023learning,
title={Learning Rain Location Prior for Nighttime Deraining},
author={Zhang, Fan and You, Shaodi and Li, Yu and Fu, Ying},
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
pages={13148--13157},
year={2023}
}
@article{zhang2022gtav,
title={GTAV-NightRain: Photometric Realistic Large-scale Dataset for Night-time Rain Streak Removal},
author={Zhang, Fan and You, Shaodi and Li, Yu and Fu, Ying},
journal={arXiv preprint arXiv:2210.04708},
year={2022}
}
Acknowledgement
The code is re-organized based on Uformer and MPRNet. Thanks for their great works!
License
MIT license.
CC BY-NC-SA 4.0 for data.