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
- 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.
Pretrained Models: Google Drive
- Python 3.9
- Pytorch 2.0.1
- NVIDIA GPU + CUDA
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
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 ..
-
Put the Train Datasets into Datasets folder.
-
Modify the config files in options folder and specific shell scripts.
-
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
-
Put the Test Datasets into Datasets folder.
-
Put the Pretrained Models into pretrained_models folder.
-
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
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)
Visual Comparison (click to expand)
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}
}
The codes are based on BasicSR, Restormer, and DiffIR. Please also follow their licenses. Thanks for their awesome works.