:umbrella: 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^:email:

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 :koala:
• Jul. 21, 2024: RainMamba is accepted with high scores (AR: 5.5) and selected as an Oral presentation (3.97%) :tada:
• Apr. 14, 2024: RainMamba is submitted to ACM MM 2024 :kangaroo:

:bookmark_tabs: 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.

🏋️ Training

:pushpin: 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

:pushpin: 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/'

📊 Results

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

Quantitative comparisons

Qualitative comparisons

VRDS dataset (click to expand)

All visual results can be downloaded here.

Real-world visual results

https://private-user-images.githubusercontent.com/141260395/365448327-36b19ed6-5f54-42cd-ba6c-49c1d0ad7328.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.