【IJCAI'2024🔥】Learning A Spiking Neural Network for Efficient Image Deraining

Learning A Spiking Neural Network for Efficient Image Deraining

Tianyu Song, Guiyue Jin, Pengpeng Li,Kui Jiang, Xiang Chen, Jiyu Jin

Dalian Polytechnic University, Nanjing University of Science and Technology and Harbin Institute of Technology

Primary contact: Tianyu Song(songtienyu@163.com)

Abstract: Recently, spiking neural networks (SNNs) have demonstrated substantial potential in computer vision tasks. In this paper, we present an Efficient Spiking Deraining Network, called ESDNet. Our work is motivated by the observation that rain pixel values will lead to a more pronounced intensity of spike signals in SNNs. However, directly applying deep SNNs to image deraining task still remains a significant challenge. This is attributed to the information loss and training difficulties that arise from discrete binary activation and complex spatio-temporal dynamics. To this end, we develop a spiking residual block to convert the input into spike signals, then adaptively optimize the membrane potential by introducing attention weights to adjust spike responses in a data-driven manner, alleviating information loss caused by discrete binary activation. By this way, our ESDNet can effectively detect and analyze the characteristics of rain streaks by learning their fluctuations. This also enables better guidance for the deraining process and facilitates high-quality image reconstruction. Instead of relying on the ANN-SNN conversion strategy, we introduce a gradient proxy strategy to directly train the model for overcoming the challenge of training. Experimental results show that our approach gains comparable performance against ANN-based methods while reducing energy consumption by 54%.

Preparation

Datasets

Dataset	Rain12	Rain200L	Rain200H	Rain1200	RW_Rain
Baidu Cloud	Download	Download	Download	Download	Download

🤖 Pre-trained Models

Dataset	Rain200L	Rain200H	Rain1200
Baidu Cloud	Download	Download	Download

Install

We test the code on PyTorch 1.9.1 + CUDA 11.1 + cuDNN 8.0.5.

1. Create a new conda environment

conda create -n ESDNet python=3.8
conda activate ESDNet 

2. Install dependencies

pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html

pip install matplotlib scikit-image opencv-python numpy einops math natsort tqdm lpips time tensorboardX

Download

You can download the pre-trained models and datasets on BaiduPan.

The final file path should be the same as the following:

┬─ pretrained_models
│   ├─ Rain12.pth
│   ├─ Rain200H.pth
│   ├─ ... (model name)
│   └─ ... (exp name)
└─ data
    ├─ Rain200L
    ├─ Rain200H
    │├─ train
    ││  ├─ input
    ││  │  └─ ... (image filename)
    ││  │  └─ target
    ││  │  └─ ... (corresponds to the former)
    │└─ test
    │    └─ ...
    └─ ... (dataset name)

🛠️ Training, Testing and Evaluation

Train

The training code will be released after the paper is accepted. You should change the path to yours in the Train.py file. Then run the following script to test the trained model:

python train.py

Test

You should change the path to yours in the Test.py file. Then run the following script to test the trained model:

python test.py

Evaluation

You should change the path to yours in the dataset_load.py file. Then run the following script to test the trained model:

python evaluation.py

It is recommended that you can download the visual deraining results and retest the quantitative results on your own device and environment.

🚀 Visual Results

Dataset	Rain12	Rain200L	Rain200H	Rain1200	RW_Rain
Baidu Cloud	Download	Download	Download	Download	Download

🚨 Notes

1. Send e-mail to songtienyu@163.com if you have critical issues to be addressed.
2. Please note that there exists the slight gap in the final version due to errors caused by different testing devices and environments.
3. Because the synthetic dataset is not realistic enough, the trained models may not work well on real world datasets.

👍 Acknowledgment

This code is based on the Restormer, spikingjelly and C2PNet. Thanks for their awesome work.