Skip to content

Official Code for "Reti-Diff: Illumination Degradation Image Restoration with Retinex-based Latent Diffusion Model". A SOTA algorithm in low-light image enhancement, underwater image enhancement, and backlit image enhancement.

Notifications You must be signed in to change notification settings

ChunmingHe/Reti-Diff

Repository files navigation

Reti-Diff

ICLRS2025 Spotlight ✨

Reti-Diff: Illumination Degradation Image Restoration with Retinex-based Latent Diffusion Model, ICLR 2025, Spotlight [Paper]

Chunming He, Chengyu Fang, Yulun Zhang, Longxiang Tang, Jinfa Huang, Kai Li, Zhenhua Guo, Xiu Li and Sina Farsiu

🔥🔥🔥 News

  • 2025-02-13: We release the code and pretrained models.
  • 2025-02-11: We are updating the code in the repository, and the version currently in the repository is not the final release version.
  • 2023-11-21: We release this repository.

Abstract: Illumination degradation image restoration (IDIR) techniques aim to improve the visibility of degraded images and mitigate the adverse effects of deteriorated illumination. Among these algorithms, diffusion model (DM)-based methods have shown promising performance but are often burdened by heavy computational demands and pixel misalignment issues when predicting the image-level distribution. To tackle these problems, we propose to leverage DM within a compact latent space to generate concise guidance priors and introduce a novel solution called Reti-Diff for the IDIR task. Reti-Diff comprises two key components: the Retinex-based latent DM (RLDM) and the Retinex-guided transformer (RGformer). To ensure detailed reconstruction and illumination correction, RLDM is empowered to acquire Retinex knowledge and extract reflectance and illumination priors. These priors are subsequently utilized by RGformer to guide the decomposition of image features into their respective reflectance and illumination components. Following this, RGformer further enhances and consolidates the decomposed features, resulting in the production of refined images with consistent content and robustness to handle complex degradation scenarios. Extensive experiments show that Reti-Diff outperforms existing methods on three IDIR tasks, as well as downstream applications.

🔗 Contents

Requirements

Pretrained Models: Google Drive

⚙️ Dependencies

  • Python 3.9
  • Pytorch 2.0.1
  • NVIDIA GPU + CUDA

Initialize Environment

git clone https://github.com/cnyvfang/Reti-Diff-demo.git
cd Reti-Diff-demo
conda create -n Reti-Diff python=3.9
conda activate Reti-Diff
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install numpy==1.25.2
pip install -r requirements.txt
python setup.py develop

Install BasicSR

git clone https://github.com/xinntao/BasicSR.git
cd BasicSR
pip install tb-nightly -i https://mirrors.aliyun.com/pypi/simple
pip install -r requirements.txt
python setup.py develop
cd ..

🧱 Training

  1. Put the Train Datasets into Datasets folder.

  2. Modify the config files in options folder and specific shell scripts.

  3. Run the shell scripts

# LLIE
sh trainS1_LLIE.sh
sh trainS2_LLIE.sh

# UIE
sh trainS1_UIE.sh
sh trainS2_UIE.sh

# Backlit
sh trainS1_Backlit.sh
sh trainS2_Backlit.sh

⚡️ Testing

  1. Put the Test Datasets into Datasets folder.

  2. Put the Pretrained Models into pretrained_models folder.

  3. Run the following command to test the model:

# LLIE
sh test_LLIE_syn.sh
or 
sh test_LLIE_real.sh

# UIE
sh test_UIE_LSUI.sh
or
sh test_UIE_UIEB.sh

# Backlit
sh test_Backlit.sh

🔍 Results

We achieved state-of-the-art performance on low light image enhancement, underwater image enhancement, backlit image enhancement and corresponding downstream tasks. More results can be found in the paper.

Quantitative Comparison (click to expand)
  • Results in Table 1 of the main paper

  • Results in Table 2-3 of the main paper

  • Results in Table 6-9 of the main paper

Visual Comparison (click to expand)
  • Results in Figure 3 of the main paper

  • Results in Figure 4 of the main paper

  • Results in Figure 5 of the main paper

📎 Citation

If you find the code helpful in your resarch or work, please cite the following paper(s).

@article{he2023retidiff,
      title={Reti-Diff: Illumination Degradation Image Restoration with Retinex-based Latent Diffusion Model}, 
      author={Chunming He and Chengyu Fang and Yulun Zhang and Kai Li and Longxiang Tang and Chenyu You and Fengyang Xiao and Zhenhua Guo and Xiu Li},
      year={2023},
      eprint={2311.11638},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

💡 Acknowledgements

The codes are based on BasicSR, Restormer, and DiffIR. Please also follow their licenses. Thanks for their awesome works.

About

Official Code for "Reti-Diff: Illumination Degradation Image Restoration with Retinex-based Latent Diffusion Model". A SOTA algorithm in low-light image enhancement, underwater image enhancement, and backlit image enhancement.

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published