Yu-Cheng Lin, Yu-Syuan Xu, Hao-Wei Chen, Hsien-Kai Kuo, Chun-Yi Lee

Abstract: Image restoration is a key task in low-level computer vision that aims to reconstruct high-quality images from degraded inputs. The emergence of Vision Mamba, which draws inspiration from the advanced state space model Mamba, marks a significant advancement in this field. Vision Mamba demonstrates excellence in modeling long-range dependencies with linear complexity, a crucial advantage for image restoration tasks. Despite its strengths, Vision Mamba encounters challenges in low-level vision tasks, including computational complexity that scales with the number of scanning sequences and local pixel forgetting. To address these limitations, this study introduces Efficient All-Around Mamba (EAMamba), an enhanced framework that incorporates a Multi-Head Selective Scan Module (MHSSM) with an all-around scanning mechanism. MHSSM efficiently aggregates multiple scanning sequences, which avoids increases in computational complexity and parameter count. The all-around scanning strategy implements multiple patterns to capture holistic information and resolves the local pixel forgetting issue. Our experimental evaluations validate these innovations across several restoration tasks, including super resolution, denoising, deblurring, and dehazing. The results validate that EAMamba achieves a significant 31-89% reduction in FLOPs while maintaining favorable performance compared to existing low-level Vision Mamba methods. The source codes are accessible at the following repository:

• Network Architecture
• Installation
• Training
• Testing
• Profiling
• Task-Specific Evaluation
• Results

Overall Framework of EAMamba

Multi-Head Selective Scan (MHSS)

We use Python=3.9.12, and PyTorch>=1.13.1 with CUDA=11.7.

Installation of mamba-ssm (v2.2.2):

git clone https://github.com/state-spaces/mamba.git
cd mamba
git checkout 8ffd905
pip install .

Installation of required packages:

pip install -r requirements.txt

python3 train.py --config [MODEL_CONFIG] --name [OUTPUT_FOLDER_NAME]

accelerate launch --config_file accelerate_[2]gpus.yaml --num_cpu_threads_per_process 32 train.py --config [MODEL_CONFIG] --name [OUTPUT_FOLDER_NAME]

Currently we have the configuration file for 2, 4, and 8 GPUs, but you can create your own configuration file for different number of GPUs.

python3 test.py --model [MODEL_FILE] --dataset [DATASET] (--ensemble) (--save)

The ensemble mode is used to average the output of multiple rotated images to improve the performance, we did not use it in our experiments.

python3 profiling/erf_viz.py --model [MODEL_FILE] (--noise)

Note that you need to manually change some of the variables in the erf_viz.py for different datasets, or noise levels.

python3 profiling/flops.py --config [CONFIG_FILE] (--var) (--simple)

Run the following command:

python3 task_evaluation/gen_sidd_mat.py --model [MODEL_FILE]

Run the following command:

python3 task_evaluation/gen_deblur_png.py --model [MODEL_FILE] --dataset [DATASET] --input_dir [INPUT_DIR] --result_dir [RESULT_DIR]

[>> Images <<]

[>> Checkpoints <<]

The gray scale row shows the difference between the ground truth and the generated results, with enhanced contrast for clearer visualization. Click on the dropdown to see the results for each dataset.

📊 Real-world Super Resolution

RealSR dataset

📊 Real-world Denoising

SIDD dataset

📊 Motion Deblurring

GoPro dataset

📊 Synthetic Dehazing

SOTS-indoor dataset

SOTS-outdoor dataset

Here are the quantitative results for each dataset. Click on the dropdown to see the results for each task.