GitHub - caiyuanhao1998/Retinexformer: "Retinexformer: One-stage Retinex-based Transformer for Low-light Image Enhancement" (ICCV 2023) & (NTIRE 2024 Runner-Up) (original) (raw)

Introduction

This is a baseline and toolbox for wide-range low-light image enhancement. This repo supports over 15 benchmarks and extremely high-resolution (up to 4000x6000) low-light enhancement. Our method Retinexformer won the second place in the NTIRE 2024 Challenge on Low Light Enhancement. If you find this repo useful, please give it a star ⭐ and consider citing our paper in your research. Thank you.

Awards

News

Results

results1

Performance on LOL with the same test setting as KinD, LLFlow, and diffusion models:

Metric LOL-v1 LOL-v2-real LOL-v2-synthetic
PSNR 27.18 27.71 29.04
SSIM 0.850 0.856 0.939

Please note that we do not suggest this test setting because it uses the mean of the ground truth to obtain better results. But, if you want to follow KinD, LLFlow, and recent diffusion-based works, it is your choice to use this test setting. Please refer to the Testing part for details.

Performance on NTIRE 2024 test-challenge:

Method Retinexformer MST++ Ensemble
PSNR 24.61 24.59 25.30
SSIM 0.85 0.85 0.85

Feel free to check the Codalab leaderboard. Our method ranks second.

results_ntire

Performance on MIT Adobe FiveK:

results2

Performance on LIME, NPE, MEF, DICM, and VV:

results3

Performance on ExDark Nighttime object detection:

results4

NTIRE - dev - 2000x3000 NTIRE - challenge - 4000x6000

1. Create Environment

We suggest you use pytorch 1.11 to re-implement the results in our ICCV 2023 paper and pytorch 2 to re-implement the results in NTIRE 2024 Challenge because pytorch 2 can save more memory in mix-precision training.

1.1 Install the environment with Pytorch 1.11

conda create -n Retinexformer python=3.7
conda activate Retinexformer
conda install pytorch=1.11 torchvision cudatoolkit=11.3 -c pytorch

pip install matplotlib scikit-learn scikit-image opencv-python yacs joblib natsort h5py tqdm tensorboard

pip install einops gdown addict future lmdb numpy pyyaml requests scipy yapf lpips
python setup.py develop --no_cuda_ext

1.2 Install the environment with Pytorch 2

conda create -n torch2 python=3.9 -y
conda activate torch2
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

pip install matplotlib scikit-learn scikit-image opencv-python yacs joblib natsort h5py tqdm tensorboard

pip install einops gdown addict future lmdb numpy pyyaml requests scipy yapf lpips thop timm
python setup.py develop --no_cuda_ext

2. Prepare Dataset

Download the following datasets:

LOL-v1 Baidu Disk (code: cyh2), Google Drive

LOL-v2 Baidu Disk (code: cyh2), Google Drive

SID Baidu Disk (code: gplv), Google Drive

SMID Baidu Disk (code: btux), Google Drive

SDSD-indoor Baidu Disk (code: jo1v), Google Drive

SDSD-outdoor Baidu Disk (code: uibk), Google Drive

MIT-Adobe FiveK Baidu Disk (code:cyh2), Google Drive, Official

NTIRE 2024 Baidu Disk (code:cyh2), Google Drive links for training input, training GT, and mini-val set.

Note:

(1) Please use bandizip to jointly unzip the .zip and .z01 files of SMID, SDSD-indoor, and SDSD-outdoor

(2) Please process the raw images of the MIT Adobe FiveK dataset following the sRGB output mode or directly download and use the sRGB image pairs processed by us in the Baidu Disk (code:cyh2) and Google Drive

(3) Please download the text_list.txt from Google Drive or Baidu Disk (code: ggbh) and then put it into the folder data/SMID/SMID_Long_np/

Then organize these datasets as follows:

    |--data   
    |    |--LOLv1
    |    |    |--Train
    |    |    |    |--input
    |    |    |    |    |--100.png
    |    |    |    |    |--101.png
    |    |    |    |     ...
    |    |    |    |--target
    |    |    |    |    |--100.png
    |    |    |    |    |--101.png
    |    |    |    |     ...
    |    |    |--Test
    |    |    |    |--input
    |    |    |    |    |--111.png
    |    |    |    |    |--146.png
    |    |    |    |     ...
    |    |    |    |--target
    |    |    |    |    |--111.png
    |    |    |    |    |--146.png
    |    |    |    |     ...
    |    |--LOLv2
    |    |    |--Real_captured
    |    |    |    |--Train
    |    |    |    |    |--Low
    |    |    |    |    |    |--00001.png
    |    |    |    |    |    |--00002.png
    |    |    |    |    |     ...
    |    |    |    |    |--Normal
    |    |    |    |    |    |--00001.png
    |    |    |    |    |    |--00002.png
    |    |    |    |    |     ...
    |    |    |    |--Test
    |    |    |    |    |--Low
    |    |    |    |    |    |--00690.png
    |    |    |    |    |    |--00691.png
    |    |    |    |    |     ...
    |    |    |    |    |--Normal
    |    |    |    |    |    |--00690.png
    |    |    |    |    |    |--00691.png
    |    |    |    |    |     ...
    |    |    |--Synthetic
    |    |    |    |--Train
    |    |    |    |    |--Low
    |    |    |    |    |   |--r000da54ft.png
    |    |    |    |    |   |--r02e1abe2t.png
    |    |    |    |    |    ...
    |    |    |    |    |--Normal
    |    |    |    |    |   |--r000da54ft.png
    |    |    |    |    |   |--r02e1abe2t.png
    |    |    |    |    |    ...
    |    |    |    |--Test
    |    |    |    |    |--Low
    |    |    |    |    |   |--r00816405t.png
    |    |    |    |    |   |--r02189767t.png
    |    |    |    |    |    ...
    |    |    |    |    |--Normal
    |    |    |    |    |   |--r00816405t.png
    |    |    |    |    |   |--r02189767t.png
    |    |    |    |    |    ...
    |    |--SDSD
    |    |    |--indoor_static_np
    |    |    |    |--input
    |    |    |    |    |--pair1
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |    |--pair2
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |     ...
    |    |    |    |--GT
    |    |    |    |    |--pair1
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |    |--pair2
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |     ...
    |    |    |--outdoor_static_np
    |    |    |    |--input
    |    |    |    |    |--MVI_0898
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |    |--MVI_0918
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |     ...
    |    |    |    |--GT
    |    |    |    |    |--MVI_0898
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |    |--MVI_0918
    |    |    |    |    |   |--0001.npy
    |    |    |    |    |   |--0002.npy
    |    |    |    |    |    ...
    |    |    |    |     ...
    |    |--SID
    |    |    |--short_sid2
    |    |    |    |--00001
    |    |    |    |    |--00001_00_0.04s.npy
    |    |    |    |    |--00001_00_0.1s.npy
    |    |    |    |    |--00001_01_0.04s.npy
    |    |    |    |    |--00001_01_0.1s.npy
    |    |    |    |     ...
    |    |    |    |--00002
    |    |    |    |    |--00002_00_0.04s.npy
    |    |    |    |    |--00002_00_0.1s.npy
    |    |    |    |    |--00002_01_0.04s.npy
    |    |    |    |    |--00002_01_0.1s.npy
    |    |    |    |     ...
    |    |    |     ...
    |    |    |--long_sid2
    |    |    |    |--00001
    |    |    |    |    |--00001_00_0.04s.npy
    |    |    |    |    |--00001_00_0.1s.npy
    |    |    |    |    |--00001_01_0.04s.npy
    |    |    |    |    |--00001_01_0.1s.npy
    |    |    |    |     ...
    |    |    |    |--00002
    |    |    |    |    |--00002_00_0.04s.npy
    |    |    |    |    |--00002_00_0.1s.npy
    |    |    |    |    |--00002_01_0.04s.npy
    |    |    |    |    |--00002_01_0.1s.npy
    |    |    |    |     ...
    |    |    |     ...
    |    |--SMID
    |    |    |--SMID_LQ_np
    |    |    |    |--0001
    |    |    |    |    |--0001.npy
    |    |    |    |    |--0002.npy
    |    |    |    |     ...
    |    |    |    |--0002
    |    |    |    |    |--0001.npy
    |    |    |    |    |--0002.npy
    |    |    |    |     ...
    |    |    |     ...
    |    |    |--SMID_Long_np
    |    |    |    |--text_list.txt
    |    |    |    |--0001
    |    |    |    |    |--0001.npy
    |    |    |    |    |--0002.npy
    |    |    |    |     ...
    |    |    |    |--0002
    |    |    |    |    |--0001.npy
    |    |    |    |    |--0002.npy
    |    |    |    |     ...
    |    |    |     ...
    |    |--FiveK
    |    |    |--train
    |    |    |    |--input
    |    |    |    |    |--a0099-kme_264.jpg
    |    |    |    |    |--a0101-kme_610.jpg
    |    |    |    |     ...
    |    |    |    |--target
    |    |    |    |    |--a0099-kme_264.jpg
    |    |    |    |    |--a0101-kme_610.jpg
    |    |    |    |     ...
    |    |    |--test
    |    |    |    |--input
    |    |    |    |    |--a4574-DSC_0038.jpg
    |    |    |    |    |--a4576-DSC_0217.jpg
    |    |    |    |     ...
    |    |    |    |--target
    |    |    |    |    |--a4574-DSC_0038.jpg
    |    |    |    |    |--a4576-DSC_0217.jpg
    |    |    |    |     ...
    |    |--NTIRE
    |    |    |--train
    |    |    |    |--input
    |    |    |    |    |--1.png
    |    |    |    |    |--3.png
    |    |    |    |     ...
    |    |    |    |--target
    |    |    |    |    |--1.png
    |    |    |    |    |--3.png
    |    |    |    |     ...
    |    |    |--minival
    |    |    |    |--input
    |    |    |    |    |--1.png
    |    |    |    |    |--31.png
    |    |    |    |     ...
    |    |    |    |--target
    |    |    |    |    |--1.png
    |    |    |    |    |--31.png
    |    |    |    |     ...

We also provide download links for LIME, NPE, MEF, DICM, and VV datasets that have no ground truth:

Baidu Disk (code: cyh2) or Google Drive

3. Testing

Download our models from Baidu Disk (code: cyh2) or Google Drive. Put them in folder pretrained_weights

activate the environment

conda activate Retinexformer

LOL-v1

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_LOL_v1.yml --weights pretrained_weights/LOL_v1.pth --dataset LOL_v1

LOL-v2-real

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_LOL_v2_real.yml --weights pretrained_weights/LOL_v2_real.pth --dataset LOL_v2_real

LOL-v2-synthetic

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_LOL_v2_synthetic.yml --weights pretrained_weights/LOL_v2_synthetic.pth --dataset LOL_v2_synthetic

SID

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_SID.yml --weights pretrained_weights/SID.pth --dataset SID

SMID

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_SMID.yml --weights pretrained_weights/SMID.pth --dataset SMID

SDSD-indoor

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_SDSD_indoor.yml --weights pretrained_weights/SDSD_indoor.pth --dataset SDSD_indoor

SDSD-outdoor

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_SDSD_outdoor.yml --weights pretrained_weights/SDSD_outdoor.pth --dataset SDSD_outdoor

FiveK

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_FiveK.yml --weights pretrained_weights/FiveK.pth --dataset FiveK

NTIRE

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_NTIRE.yml --weights pretrained_weights/NTIRE.pth --dataset NTIRE --self_ensemble

MST_Plus_Plus trained with 4 GPUs on NTIRE

python3 Enhancement/test_from_dataset.py --opt Options/MST_Plus_Plus_NTIRE_4x1800.yml --weights pretrained_weights/MST_Plus_Plus_4x1800.pth --dataset NTIRE --self_ensemble

MST_Plus_Plus trained with 8 GPUs on NTIRE

python3 Enhancement/test_from_dataset.py --opt Options/MST_Plus_Plus_NTIRE_8x1150.yml --weights pretrained_weights/MST_Plus_Plus_8x1150.pth --dataset NTIRE --self_ensemble

We add the self-ensemble strategy in the testing code to derive better results. Just add a --self_ensemble action at the end of the above test command to use it.

We provide the same test setting as LLFlow, KinD, and recent diffusion models. Please note that we do not suggest this test setting because it uses the mean of ground truth to enhance the output of the model. But if you want to follow this test setting, just add a --GT_mean action at the end of the above test command as

LOL-v1

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_LOL_v1.yml --weights pretrained_weights/LOL_v1.pth --dataset LOL_v1 --GT_mean

LOL-v2-real

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_LOL_v2_real.yml --weights pretrained_weights/LOL_v2_real.pth --dataset LOL_v2_real --GT_mean

LOL-v2-synthetic

python3 Enhancement/test_from_dataset.py --opt Options/RetinexFormer_LOL_v2_synthetic.yml --weights pretrained_weights/LOL_v2_synthetic.pth --dataset LOL_v2_synthetic --GT_mean

We have provided a function my_summary() in Enhancement/utils.py, please use this function to evaluate the parameters and computational complexity of the models, especially the Transformers as

from utils import my_summary my_summary(RetinexFormer(), 256, 256, 3, 1)

4. Training

Feel free to check our training logs from Baidu Disk (code: cyh2) or Google Drive

We suggest you use the environment with pytorch 2 to train our model on the NTIRE 2024 dataset and the environment with pytorch 1.11 to train our model on other datasets.

activate the enviroment

conda activate Retinexformer

LOL-v1

python3 basicsr/train.py --opt Options/RetinexFormer_LOL_v1.yml

LOL-v2-real

python3 basicsr/train.py --opt Options/RetinexFormer_LOL_v2_real.yml

LOL-v2-synthetic

python3 basicsr/train.py --opt Options/RetinexFormer_LOL_v2_synthetic.yml

SID

python3 basicsr/train.py --opt Options/RetinexFormer_SID.yml

SMID

python3 basicsr/train.py --opt Options/RetinexFormer_SMID.yml

SDSD-indoor

python3 basicsr/train.py --opt Options/RetinexFormer_SDSD_indoor.yml

SDSD-outdoor

python3 basicsr/train.py --opt Options/RetinexFormer_SDSD_outdoor.yml

FiveK

python3 basicsr/train.py --opt Options/RetinexFormer_FiveK.yml

Train our Retinexformer and MST++ with the distributed data parallel (DDP) strategy of pytorch on the NTIRE 2024 Low-Light Enhancement dataset. Please note that we use the mix-precision strategy in the training process, which is controlled by the bool hyperparameter use_amp in the config file.

activate the enviroment

conda activate torch2

Train Retinexformer with 8 GPUs on NTIRE

bash train_multigpu.sh Options/RetinexFormer_NTIRE.yml 0,1,2,3,4,5,6,7 4321

Train MST++ with 4 GPUs on NTIRE

bash train_multigpu.sh Options/RetinexFormer_NTIRE_4x1800.yml 0,1,2,3,4,5,6,7 4329

Train MST++ with 8 GPUs on NTIRE

bash train_multigpu.sh Options/MST_Plus_Plus_NTIRE_8x1150.yml 0,1,2,3,4,5,6,7 4343

5. Citation

@InProceedings{Cai_2023_ICCV, author = {Cai, Yuanhao and Bian, Hao and Lin, Jing and Wang, Haoqian and Timofte, Radu and Zhang, Yulun}, title = {Retinexformer: One-stage Retinex-based Transformer for Low-light Image Enhancement}, booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, month = {October}, year = {2023}, pages = {12504-12513} }

@inproceedings{retinexformer, title={Retinexformer: One-stage Retinex-based Transformer for Low-light Image Enhancement}, author={Yuanhao Cai and Hao Bian and Jing Lin and Haoqian Wang and Radu Timofte and Yulun Zhang}, booktitle={ICCV}, year={2023} }

MST++

@inproceedings{mst, title={Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction}, author={Yuanhao Cai and Jing Lin and Xiaowan Hu and Haoqian Wang and Xin Yuan and Yulun Zhang and Radu Timofte and Luc Van Gool}, booktitle={CVPR}, year={2022} }