[NeurIPS 2024] ECMamba: Consolidating Selective State Space Model with Retinex Guidance for Efficient Multiple Exposure Correction [Paper]

Wei Dong^1,*, Han Zhou^1,*, Yulun Zhang², Xiaohong Liu^2,†, Jun Chen¹

¹McMaster University, ²Shanghai Jiao Tong University,

^*Equal Contribution, ^†Corresponding Author

Introduction

This repository represents the official implementation of our NeurIPS 2024 paper titled ECMamba: Consolidating Selective State Space Model with Retinex Guidance for Efficient Multiple Exposure Correction. If you find this repo useful, please give it a star ⭐ and consider citing our paper in your research. Thank you for your interest.

We present ECMamba, the first mamba-based framework for multiple exposure correction and low-light image enhancement.

• Mamba-based Multiple Exposure Correction: exploit mamba-based framework to process images with adverse illumination with high efficiency.
• Dual-path Retinex-guided Restoration Framework: develop a dual-path Retinex-guided restoration pipeline by introducing two intermediary spaces based on Retinex theory.
• Feature-aware Scanning Strategy: Different from direction-sensitive scanning method, we design a feature-aware 2D selective scanning mechanism to transform 2D iamge or feature maps into 1D sequences.

📢 News

2025-06-03 This repo has been updated. Pre-trained weights and test codes are released!

Overall Framework

🛠️ Setup

The inference code was tested on:

• Python 3.9, CUDA 11.7, PyTorch 2.0.1 + cu117.

📦 Repository

Clone the repository (requires git):

git clone https://github.com/LowLevelAI/ECMamba.git
cd ECMamba

💻 Dependencies

• Make Conda Environment: Using Conda to create the environment:

  conda create -n ecmamba python=3.9
  conda activate ecmamba

• Then install dependencies:

  • Install Pytorch

  conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia

  • Install mamba_ssm library

  pip install causal_conv1d==1.0.0
  pip install mamba_ssm==1.0.1

  • Install DCNv4 library

  pip install -U openmim
  mim install mmcv-full
  pip install timm mmdet
  pip install opencv-python termcolor yacs pyyaml scipy
  pip install DCNv4

  • Install other dependencies

  pip install numpy==1.26.4 transformers==4.48.2 opencv-python natsort scikit-image timm matplotlib

• Revise DCNv4 CUDA extensions: In /anaconda3/envs/ecmamba/lib/python3.9/site-packages/DCNv4/modules/dcnv4.py line 153, replace retunr x with retune x, offset_mask.

🏃 Testing on Benchmark Datasets

📷 Download following datasets:

• SICE Dataset Drive

• LOLv1 Google Drive

⬇ Download pre-trained models

Download Pre-trained weight for SICE Dataset, and Pre-trained weight for LOLv1. Place them to folder weights.

🚀 Run inference

• For SICE dataset

python test_sice.py

• For LOLv1 dataset

python test_lol.py

You can find all results in test-results/.

🏋️ Model Outputs

For your convenience, we also provide our outputs here.

• LOLv1: With GT-mean, Without GT-mean
• SICE Dataset Drive

✏️ Contributing

Please refer to this instruction.

🎓 Citation

Please cite our paper:

@article{dong2024ecmamba,
  title={Ecmamba: Consolidating selective state space model with retinex guidance for efficient multiple exposure correction},
  author={Dong, Wei and Zhou, Han and Zhang, Yulun and Liu, Xiaohong and Chen, Jun},
  journal={Advances in Neural Information Processing Systems},
  year={2024}
}

🎫 License

This work is licensed under the Apache License, Version 2.0 (as defined in the LICENSE).

By downloading and using the code and model you agree to the terms in the LICENSE.