[Paper] [Zhihu(η₯δΉ)]
Hang Guo*, Jinmin Li*, Tao Dai, Zhihao Ouyang, Xudong Ren, and Shu-Tao Xia
(*) equal contribution
Abstract: Recent years have witnessed great progress in image restoration thanks to the advancements in modern deep neural networks e.g. Convolutional Neural Network and Transformer. However, existing restoration backbones are usually limited due to the inherent local reductive bias or quadratic computational complexity. Recently, Selective Structured State Space Model e.g., Mamba, have shown great potential for long-range dependencies modeling with linear complexity, but it is still under-explored in low-level computer vision. In this work, we introduce a simple but strong benchmark model, named MambaIR, for image restoration. In detail, we propose the Residual State Space Block as the core component, which employs convolution and channel attention to enhance capabilities of the vanilla Mamba. In this way, our MambaIR takes advantages of local patch recurrence prior as well as channel interaction to produce restoration-specific feature representation. Extensive experiments demonstrate the superiority of our method, for example, MambaIR outperforms Transformer-based baseline SwinIR by up to 0.36dB, using similar computational cost but with global receptive field.
βIf this work is helpful for you, please help star this repo. Thanks!π€
- 2024-2-23: arXiv paper available.
- 2024-2-27: This repo is released.
- 2024-3-01: Pretrained weights for SR and realDN is available. π
- 2024-3-08: The code for ERF visualization and model complexity analysis can be found at
./analysis/
π
- Build the repo
- arXiv version
- Release code
- Pretrained weights&log_files
- Add code for complexity analysis and ERF visualization
- Guassian Color Image Denosing
- Real-world SR
- JPEG Compression Artifact Redection
- More Tasks
Model | Task | Test_dataset | PSNR | SSIM | model_weights | log_files |
---|---|---|---|---|---|---|
MambaIR_SR2 | Classic SR x2 | Urban100 | 34.15 | 0.9446 | link | link |
MambaIR_SR3 | Classic SR x3 | Urban100 | 29.93 | 0.8841 | link | link |
MambaIR_SR4 | Classic SR x4 | Urban100 | 27.68 | 0.8287 | link | link |
MambaIR_light2 | Lightweight SR x2 | Urban100 | 32.86 | 0.9343 | link | link |
MambaIR_light3 | Lightweight SR x3 | Urban100 | 28.73 | 0.8635 | link | link |
MambaIR_light4 | Lightweight SR x4 | Urban100 | 26.53 | 0.7983 | link | link |
MambaIR_realDN | Real image Denoising | SIDD | 39.89 | 0.960 | link | link |
We achieve state-of-the-art performance on various image restoration tasks. Detailed results can be found in the paper.
This codebase was tested with the following environment configurations. It may work with other versions.
- Ubuntu 20.04
- CUDA 11.7
- Python 3.9
- PyTorch 1.13.1 + cu117
To use the selective scan with efficient hard-ware design, the mamba_ssm
library is advised to install with the folllowing command.
pip install causal_conv1d==1.0.0
pip install mamba_ssm==1.0.1
One can also create a new anaconda environment, and then install necessary python libraries with this requirement.txt and the following command:
conda install --yes --file requirements.txt
-
Please download the corresponding training datasets and put them in the folder datasets/DF2K. Download the testing datasets and put them in the folder datasets/SR.
-
Follow the instructions below to begin training our model.
# Claissc SR task, cropped input=64Γ64, 8 GPUs, batch size=4 per GPU
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/train_MambaIR_SR_x2.yml --launcher pytorch
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/train_MambaIR_SR_x3.yml --launcher pytorch
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/train_MambaIR_SR_x4.yml --launcher pytorch
# Lightweight SR task, cropped input=64Γ64, 8 GPUs, batch size=8 per GPU
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/train_MambaIR_lightSR_x2.yml --launcher pytorch
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/train_MambaIR_lightSR_x3.yml --launcher pytorch
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/train_MambaIR_lightSR_x4.yml --launcher pytorch
- Run the script then you can find the generated experimental logs in the folder experiments.
- Please download the corresponding training datasets and put them in the folder datasets/SIDD. Note that we provide both training and validating files, which are already processed.
- Go to folder 'realDenoising'. Follow the instructions below to train our model.
# go to the folder
cd realDenoising
# set the new environment (BasicSRv1.2.0), which is the same with Restormer for training.
python setup.py develop --no_cuda_extgf
# train for RealDN task, 8 GPUs
python -m torch.distributed.launch --nproc_per_node=8 --master_port=2414 basicsr/train.py -opt options/train_MambaIR_RealDN.yml --launcher pytorch
Run the script then you can find the generated experimental logs in the folder realDenoising/experiments.
- Remember to go back to the original environment if you finish all the training or testing about real image denoising task. This is a friendly hint in order to prevent confusion in the training environment.
# Tips here. Go back to the original environment (BasicSRv1.3.5) after finishing all the training or testing about real image denoising.
cd ..
python setup.py develop
-
Please download the corresponding testing datasets and put them in the folder datasets/SR. Download the corresponding models and put them in the folder experiments/pretrained_models.
-
Follow the instructions below to begin testing our MambaIR model.
# test for image SR.
python basicsr/test.py -opt options/test/test_MambaIR_SR_x2.yml
python basicsr/test.py -opt options/test/test_MambaIR_SR_x3.yml
python basicsr/test.py -opt options/test/test_MambaIR_SR_x4.yml
# test for lightweight image SR.
python basicsr/test.py -opt options/test/test_MambaIR_lightSR_x2.yml
python basicsr/test.py -opt options/test/test_MambaIR_lightSR_x3.yml
python basicsr/test.py -opt options/test/test_MambaIR_lightSR_x4.yml
-
Download the SIDD test and DND test. Place them in
datasets/RealDN
. Download the corresponding models and put them in the folderexperiments/pretrained_models
. -
Go to folder 'realDenoising'. Follow the instructions below to test our model. The output is in
realDenoising/results/Real_Denoising
.# go to the folder cd realDenoising # set the new environment (BasicSRv1.2.0), which is the same with Restormer for testing. python setup.py develop --no_cuda_ext # test MambaIR (training total iterations = 300K) on SSID python test_real_denoising_sidd.py # test MambaIR (training total iterations = 300K) on DND python test_real_denoising_dnd.py
-
Run the scripts below to reproduce PSNR/SSIM on SIDD.
run evaluate_sidd.m
-
For PSNR/SSIM scores on DND, you can upload the genetated DND mat files to the online server and get the results.
-
Remerber to go back to the original environment if you finish all the training or testing about real image denoising task. This is a friendly hint in order to prevent confusion in the training environment.
# Tips here. Go back to the original environment (BasicSRv1.3.5) after finishing all the training or testing about real image denoising. cd .. python setup.py develop
Please cite us if our work is useful for your research.
@article{guo2024mambair,
title={MambaIR: A Simple Baseline for Image Restoration with State-Space Model},
author={Guo, Hang and Li, Jinmin and Dai, Tao and Ouyang, Zhihao and Ren, Xudong and Xia, Shu-Tao},
journal={arXiv preprint arXiv:2402.15648},
year={2024}
}
This project is released under the Apache 2.0 license.
This code is based on BasicSR, ART ,and VMamba. Thanks for their awesome work.
If you have any questions, feel free to approach me at cshguo@gmail.com