MambaIR

repository·main·Indexed 22 days ago

A collection of image restoration backbones based on State-Space Models (Mamba) for tasks including Super-Resolution (SR), Denoising, and JPEG artifact reduction. The library includes MambaIR (ECCV 2024) and MambaIRv2 (CVPR 2025), the latter of which introduces 'Attentive State Space Restoration' for enhanced non-causal modeling and a semantic-guided neighboring mechanism.

Tokens: 6K
Snippets: 12
Records: 16
Agent score: 78%

What's inside MambaIR

Overview of MambaIR and MambaIRv2
main
MambaIR is a State-Space Model (SSM) based backbone designed for image restoration tasks (e.g., Super-Resolution, Denoising).
MambaIR (ECCV 2024): A simple baseline for image restoration using Mamba. It addresses the causal modeling limitations of standard Mamba to better utilize image pixels.
MambaIRv2 (CVPR 2025): An upgraded version featuring 'Attentive State Space Restoration'. It introduces non-causal modeling (similar to Vision Transformers) via an attentive state-space equation, allowing the model to attend beyond the scanned sequence. It also includes a semantic-guided neighboring mechanism for interaction between distant but similar pixels. MambaIRv2 is fully compatible with the original MambaIR code.
https://github.com/csguoh/mambair/blob/main/README.md
Real-world Image Denoising with MambaIR
main
MambaIR implements real-world image denoising using a U-Net architecture. To ensure fair comparisons with established models like Restormer, MambaIR utilizes the same training environment and progressive training strategy.
Key Implementation Details:
- Architecture: U-Net structure (similar to Restormer).
- Environment: Built on BasicSR (v1.2.0).
- Training Strategy: Progressive training strategy.
- Dependency: This implementation is built based on Restormer.
https://github.com/csguoh/mambair/blob/main/realDenoising/README.md
Install mamba-related libraries manually
main
If you prefer to install dependencies into an existing environment, you must install causal_conv1d and mamba_ssm to enable efficient hardware-designed selective scanning.
Alternatively, you can use the provided requirements.txt with conda install.
```
# Manual mamba library installation
pip install causal_conv1d==1.0.0
pip install mamba_ssm==1.0.1

# Using requirements.txt
conda install --yes --file requirements.txt
```
https://github.com/csguoh/mambair/blob/main/README.md
Access MambaIR weights and demos
main
You can access pretrained weights, research papers, and interactive demos through the following resources:
- Pretrained Weights: Available on HuggingFace.
- Interactive Demo: An online Jupyter Notebook demo is available via Google Colab for fast inference.
- Paper: The MambaIR paper is available on arXiv.
https://github.com/csguoh/mambair/blob/main/README.md

Organize datasets for MambaIR training and testing

main

To use MambaIR for training and testing, you must organize your downloaded datasets into a specific directory structure under the datasets folder. The project expects predefined subdirectories for different tasks such as Super-Resolution (SR), Denoising (DN), and Compression Artifact Reduction (CAR).

|-- datasets
    # image SR - train
    |-- DF2K
        |-- HR
        |-- LR_bicubic
            |-- X2
            |-- X3
            |-- X4
    # color image denoising - train
    |-- DFWB_RGB
        |-- HQ
    # real image denoising - train & val
    |-- SIDD
        |-- train
            |-- target_crops
            |-- input_crops  
        |-- val
            |-- target_crops
            |-- input_crops 
    # grayscale JPEG compression artifact reduction - train
    |-- DFWB_CAR
        |-- HQ
        |-- LQ
            |-- 10
            |-- 30
            |-- 40  
    # image SR - test
    |-- SR
        |-- Set5
            |-- HR
            |-- LR_bicubic
                |-- X2
                |-- X3
                |-- X4
        |-- Set14
            |-- HR
            |-- LR_bicubic
                |-- X2
                |-- X3
                |-- X4
        |-- B100
            |-- HR
            |-- LR_bicubic
                |-- X2
                |-- X3
                |-- X4
        |-- Urban100
            |-- HR
            |-- LR_bicubic
                |-- X2
                |-- X3
                |-- X4
        |-- Manga109
            |-- HR
            |-- LR_bicubic
                |-- X2
                |-- X3
                |-- X4        
    # gaussian color image denoising - test
    |-- ColorDN
        |-- CBSD68HQ
        |-- Kodak24HQ
        |-- McMasterHQ
        |-- Urban100HQ
    # real image denoising - test
    |-- RealDN
        |-- SIDD
            |-- ValidationGtBlocksSrgb.mat
            |-- ValidationNoisyBlocksSrgb.mat
        |-- DND
            |-- info.mat
            |-- ValidationNoisyBlocksSrgb
                |-- 0001.mat
                |-- 0002.mat
                |-- 0050.mat
    # grayscale JPEG compression artifact reduction - test
    |-- CAR
        |-- classic5
            |-- Classic5_HQ
            |-- Classic5_LQ
                |-- 10
                |-- 30
                |-- 40
        |-- LIVE1
            |-- LIVE1_HQ
            |-- LIVE1_LQ
                |-- 10
                |-- 30
                |-- 40

Install MambaIR via Conda environment
main
The recommended way to set up the environment is to use the provided environment.yaml file. This will create a Conda environment named mambair with the necessary dependencies, including CUDA 11.7 configurations.
Note: If you are using a different CUDA version (e.g., 12.x), you may need to modify the environment.yaml file or refer to the official causal_conv1d and mamba_ssm pages to find compatible versions.
```
cd ./MambaIR
conda env create -f environment.yaml
conda activate mambair
```
https://github.com/csguoh/mambair/blob/main/README.md
Test MambaIRv2 for Gaussian Color Image Denoising
main
To test Gaussian Color Image Denoising models:
1. Place testing datasets in datasets/ColorDN.
2. Place downloaded models in experiments/pretrained_models.
3. Run the test script with the corresponding configuration file.
```
# test on denosing15
python basicsr/test.py -opt options/test/mambairv2/test_MambaIRv2_ColorDN_level15.yml
```
https://github.com/csguoh/mambair/blob/main/README.md
Test MambaIRv2 for JPEG Compression Artifact Reduction (CAR)
main
To test JPEG Compression Artifact Reduction models:
1. Place testing datasets in datasets/JPEG_CAR.
2. Place downloaded models in experiments/pretrained_models.
3. Run the test script with the configuration matching the target JPEG quality (q10, q30, or q40).
```
# test on jpeg10
python basicsr/test.py -opt options/test/mambairv2/test_MambaIRv2_CAR_q10.yml

# test on jpeg30
python basicsr/test.py -opt options/test/mambairv2/test_MambaIRv2_CAR_q30.yml

# test on jpeg40
python basicsr/test.py -opt options/test/mambairv2/test_MambaIRv2_CAR_q40.yml
```
https://github.com/csguoh/mambair/blob/main/README.md
Download MambaIR and MambaIRv2 model weights
main
Pre-trained weights for various MambaIR and MambaIRv2 models (including Super-Resolution, Denoising, and JPEG Artifact Reduction) are available on HuggingFace and via Google Drive links provided in the Model Summary table.
For the latest MambaIRv2 models, weights can be found in the GitHub Releases.
```
https://huggingface.co/cguoh/MambaIR
```
https://github.com/csguoh/mambair/blob/main/README.md

Train MambaIRv2 for Super-Resolution (SR)

main

To train MambaIRv2 for Super-Resolution tasks, organize your datasets as follows:

Training datasets: datasets/DF2K
Testing datasets: datasets/SR

Training is performed using basicsr/train.py with distributed PyTorch. The default is the Base model with a cropped input of 64x64. You can also train Small or Large variants, or use the Lightweight SR models which require fewer GPUs.

# Classic SR task (Base model, 8 GPUs)
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_SR_x2.yml --launcher pytorch

# Small or Large model (8 GPUs)
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_SRSmall_x4.yml --launcher pytorch

# Lightweight SR task (2 GPUs)
python -m torch.distributed.launch --nproc_per_node=2 --master_port=1234 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_lightSR_x2.yml --launcher pytorch

Train MambaIRv2 for JPEG Compression Artifact Reduction (CAR)

main

To train for JPEG Compression Artifact Reduction:

Place training datasets in ./datasets/DFWB_CAR.
Place testing datasets in ./datasets/JPEG_CAR.
Run the training script using the configuration corresponding to the target JPEG quality (q10, q30, or q40).

Experimental logs are saved to the ./experiments folder.

# train on jpeg10
python -m torch.distributed.launch --nproc_per_node=8 --master_port=2414 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_CAR_q10.yml --launcher pytorch

# train on jpeg30
python -m torch.distributed.launch --nproc_per_node=8 --master_port=2414 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_CAR_q30.yml --launcher pytorch

# train on jpeg40
python -m torch.distributed.launch --nproc_per_node=8 --master_port=2414 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_CAR_q40.yml --launcher pytorch

Train MambaIRv2 for Gaussian Color Image Denoising
main
To train for Gaussian Color Image Denoising:
1. Place training datasets in ./datasets/DFWB_RGB.
2. Place testing datasets in ./datasets/ColorDN.
3. Run the training script using the provided configuration for the desired noise level (e.g., level 15).
Experimental logs are saved to the ./experiments folder.
```
# train on denosing15
python -m torch.distributed.launch --nproc_per_node=8 --master_port=2414 basicsr/train.py -opt options/train/mambairv2/train_MambaIRv2_ColorDN_level15.yml --launcher pytorch
```
https://github.com/csguoh/mambair/blob/main/README.md