☔ RainMamba [ACM MM'24 Oral]

RainMamba: Enhanced Locality Learning with State Space Models
for Video Deraining

Hongtao Wu, Yijun Yang, Huihui Xu, Weiming Wang, Jinni Zhou, Lei Zhu^📧

ArXivPreprint | OpenReview

This repository is the official implementation of the paper titled 'RainMamba: Enhanced Locality Learning with State Space Models for Video Deraining'.

RainMamba achieves state-of-the-art results on multiple video deraining benchmarks with 3D Hilbert scanning mechanism, offering superior computational efficiency and enhanced restoration quality.

If you like this work, please give us a star ⭐ on GitHub.

---

Abstract: The outdoor vision systems are frequently contaminated by rain streaks and raindrops, which significantly degenerate the performance of visual tasks and multimedia applications. The nature of videos exhibits redundant temporal cues for rain removal with higher stability. Traditional video deraining methods heavily rely on optical flow estimation and kernel-based manners, which have a limited receptive field. Yet, transformer architectures, while enabling long-term dependencies, bring about a significant increase in computational complexity. Recently, the linear-complexity operator of the state space models (SSMs) has contrarily facilitated efficient long-term temporal modeling, which is crucial for rain streaks and raindrops removal in videos. Unexpectedly, its uni-dimensional sequential process on videos destroys the local correlations across the spatio-temporal dimension by distancing adjacent pixels. To address this, we present an improved SSMs-based video deraining network (RainMamba) with a novel Hilbert scanning mechanism to better capture sequence-level local information. We also introduce a difference-guided dynamic contrastive locality learning strategy to enhance the patch-level self-similarity learning ability of the proposed network. Extensive experiments on four synthesized video deraining datasets and real-world rainy videos demonstrate the superiority of our network in the removal of rain streaks and raindrops.

---

• Framework

• Spatio-temporal scanning mechanisms

• Model efficiency and effectiveness

📧 News

• Jul. 30, 2024: Codes of RainMamba (including the training, testing code, pretrained model and results of four video deraining datasets) are released 🐨
• Jul. 21, 2024: RainMamba is accepted with high scores (AR: 5.5) and selected as an Oral presentation (3.97%) 🎉
• Apr. 14, 2024: RainMamba is submitted to ACM MM 2024 🦘

📑 Contents

• Environment Installation
• Datasets Preparation and Pretrained Models
• Training
• Testing
• Results
• Citation
• Acknowledgement
• License

🔨 Environment Installation

This implementation is based on MMEditing, which is an open-source image and video editing toolbox.

Below are quick steps for installation.

1️⃣ Make conda environment and install PyTorch following official instructions.

conda create -n RainMamba python=3.9
conda activate RainMamba

conda install pytorch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 pytorch-cuda=12.1 -c pytorch -c nvidia

2️⃣ Install MMCV with MIM.

pip3 install openmim timm
mim install mmcv-full==1.7.2

3️⃣ Install RainMamba from source.

git clone https://github.com/TonyHongtaoWu/RainMamba.git
cd RainMamba/code/VRDS
pip3 install -e .

Please refer to MMEditing Installation for more detailed instruction.

4️⃣ Install Bi-Mamba layer

Install causal-conv1d

cd causal-conv1d

python setup.py install

Install mamba

cd ../mamba

python setup.py install

🚀 Datasets Preparation and Pretrained Models

Download four datasets from their websites and modify the training dataset path in the config file of each dataset code.

Dataset	Source	Description	Dataset Link	Pretrained Model
VRDS	ACM MM 2023	Rain streaks and raindrops dataset	Download	Download
LWDDS	ICRA 2023	Raindrops dataset	Download	Download
RainVID&SS	TIP 2023	Rain streaks dataset	Download	Download
RainSynAll100	TPAMI 2021	Rain streaks dataset	Download	Download

🏋️ Training

📌 You can train RainMamba on these datasets using the below commands with 4 GPUs:

1️⃣ VRDS dataset

cd code/VRDS
CUDA_VISIBLE_DEVICES=0,1,2,3 bash ./tools/dist_train.sh configs/derainers/RainMamba/VRDS.py 4

2️⃣ LWDDS dataset

cd code/LWDDS
CUDA_VISIBLE_DEVICES=0,1,2,3 bash ./tools/dist_train.sh configs/derainers/RainMamba/Waterdrop.py 4

3️⃣ RainVID&SS dataset

cd  code/RainVIDSS
CUDA_VISIBLE_DEVICES=0,1,2,3 bash ./tools/dist_train.sh configs/derainers/RainMamba/RainVIDSS.py 4

4️⃣ RainSynAll100 dataset

cd code/RainSynAll100
CUDA_VISIBLE_DEVICES=0,1,2,3 bash ./tools/dist_train.sh configs/derainers/RainMamba/RainSynAll100.py 4

⚡ Testing

📌 Modify the path of pretrained model and the testing dataset in the config file of each dataset code.

You can use the following command with 1 GPU to test your trained model xxx.pth:

1️⃣ VRDS dataset

cd code/VRDS
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_test.sh configs/derainers/RainMamba/VRDS.py "VRDS.pth" 1 --save-path '../results/VRDS/'

2️⃣ LWDDS dataset

cd code/LWDDS
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_test.sh configs/derainers/RainMamba/Waterdrop.py "LWDDS.pth" 1 --save-path '../results/LWDDS/'

3️⃣ RainVID&SS dataset

cd  code/RainVIDSS
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_test.sh configs/derainers/RainMamba/RainVIDSS.py "RainVIDSS.pth" 1 --save-path '../results/RainVIDSS/'

4️⃣ RainSynAll100 dataset

cd code/RainSynAll100
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_test.sh configs/derainers/RainMamba/RainSynAll100.py "RainSynAll100.pth" 1 --save-path '../results/RainSynAll100/'

5️⃣ Demo code for video

You can use the following commands to test a rainy video. This code can be used to test a video or a folder of frames.

cd code/VRDS
python demo/restoration_video_demo.py \
configs/derainers/RainMamba/VRDS.py \
"VRDS.pth" \
input/rainy.mp4 \
output/derained.mp4 

📊 Results

RainMamba achieved state-of-the-art on four video deraining datasets.

Quantitative comparisons

VRDS dataset (click to expand)

Methods	CCN	PreNet	DRSformer	MPRNet	Restormer	S2VD	SLDNet	ESTINet	RDD	RVRT	BasicVSR++	ViMP-Net	Ours
PSNR $\uparrow$	23.75	27.13	28.54	29.53	29.59	18.95	23.65	27.17	28.39	28.24	29.75	31.02	32.04
SSIM $\uparrow$	0.8410	0.9014	0.9075	0.9175	0.9206	0.6630	0.8736	0.8436	0.9096	0.8857	0.9171	0.9283	0.9366
LPIPS $\downarrow$	0.2091	0.1266	0.1143	0.0987	0.0925	0.2833	0.1790	0.2253	0.1168	0.1438	0.1023	0.0862	0.0684

RainVID&SS dataset (click to expand)

ImageNet-VID+

Methods	MSCSC	FastDrain	PReNet	S2VD	MPEVNet	Ours
PSNR $\uparrow$	18.41	17.08	24.73	29.92	33.83	35.07
SSIM $\uparrow$	0.5148	0.4381	0.7393	0.9228	0.9452	0.9561

Cam-Vid+

Methods	MSCSC	FastDrain	PReNet	SLDNet	S2VD	MPEVNet	Ours
PSNR $\uparrow$	21.22	19.94	25.33	18.97	29.11	32.55	32.65
SSIM $\uparrow$	0.5515	0.4830	0.7647	0.6267	0.8899	0.9234	0.9328

RainSynAll100 dataset (click to expand)

Methods	FastDerain	FCRVD	RMFD	BasicVSR++	NCFL	SALN	Ours
PSNR $\uparrow$	17.09	21.06	25.14	27.67	28.11	29.78	32.16
SSIM $\uparrow$	0.5824	0.7405	0.9172	0.9135	0.9235	0.9315	0.9446

LWDDS dataset (click to expand)

Methods	CCN	Vid2Vid	VWR	BasicVSR++	ViMP-Net	SALN	Ours
PSNR $\uparrow$	27.53	28.73	30.72	32.37	34.22	36.57	37.21
SSIM $\uparrow$	0.922	0.9542	0.9726	0.9792	0.9784	0.9802	0.9816

NTURain dataset (click to expand)

Method	MSCSC	J4RNet	SPAC	FCRNet	SLDNet	MPRNet	S2VD	ESTINet	Ours
PSNR$\uparrow$	27.31	32.14	33.11	36.05	34.89	36.11	37.37	37.48	37.87
SSIM$\uparrow$	0.7870	0.9480	0.9474	0.9676	0.9540	0.9637	0.9683	0.9700	0.9738

Qualitative comparisons

VRDS dataset (click to expand)

Other datasets (please refer to the supplementary materials)

All visual results can be downloaded here.

Dataset	Link 📁	Link 📂
VRDS	Google Drive	OneDrive
LWDDS	Google Drive	OneDrive
RainVID&SS	Google Drive	OneDrive
RainSynAll100	Google Drive	OneDrive
NTURain	Google Drive	OneDrive

Real-world visual results

real_video.mp4

🎓 Citation

If you find this repository helpful to your research, please consider citing the following:

@article{wu2024rainmamba,
  title={RainMamba: Enhanced Locality Learning with State Space Models for Video Deraining},
  author={Wu, Hongtao and Yang, Yijun and Xu, Huihui and Wang, Weiming and Zhou, Jinni and Zhu, Lei},
  journal={arXiv preprint arXiv:2407.21773},
  year={2024}
}

If you have any questions or suggestion about our project, please feel free to contact us by email at hwu375@connect.hkust-gz.edu.cn.

🤗 Acknowledgement

This code is based on MMEditing, Vim, Vivim , gilbert, LGRNet. Thank them for their outstanding work.

Thank Huihui Xu and Yijun Yang for their help in this work.

🎫 License

This project is released under the MIT license. Please refer to the acknowledged repositories for their licenses.