STAR-pytorch
Implementation for paper "STAR: A Structure-aware Lightweight Transformer for Real-time Image Enhancement" (ICCV 2021).
STAR-DCE
The pytorch implementation of low light enhancement with STAR on Adobe-MIT FiveK dataset. You can find it in STAR-DCE directory. Here we adopt the pipleline of Zero-DCE ( paper | code ), just replacing the CNN backbone with STAR. In Zero-DCE, for each image the network will regress a group of curves, which will then applied on the source image iteratively. You can find more details in the original repo Zero-DCE.
Requirements
- numpy
- einops
- torch
- torchvision
- opencv
Datesets
We provide download links for Adobe-MIT FiveK datasets we used ( train | test ). Please note that we adopt the test set splited by DeepUPE for fair comparison.
Training DCE models
To train a original STAR-DCE model,
cd STAR-DCE
python train_dce.py
--lowlight_images_path "dir-to-your-training-set" \
--parallel True \
--snapshots_folder snapshots/STAR-ori \
--lr 0.001 \
--num_epochs 100 \
--lr_type cos \
--train_batch_size 32 \
--model STAR-DCE-Ori \
--snapshot_iter 10 \
--num_workers 32 \
To train the baseline CNN-based DCE-Net (w\ or w\o Pooling),
cd STAR-DCE
python train_dce.py
--lowlight_images_path "dir-to-your-training-set" \
--parallel True \
--snapshots_folder snapshots/DCE \
--lr 0.001 \
--num_epochs 100 \
--lr_type cos \
--train_batch_size 32 \
--model DCE-Net \
--snapshot_iter 10 \
--num_workers 32 \
or
cd STAR-DCE
python train_dce.py
--lowlight_images_path "dir-to-your-training-set" \
--parallel True \
--snapshots_folder snapshots/DCE-Pool \
--lr 0.001 \
--num_epochs 100 \
--lr_type cos \
--train_batch_size 32 \
--model DCE-Net-Pool \
--snapshot_iter 10 \
--num_workers 32 \
Evaluation of trained models
To evaluated the STAR-DCE model you trained,
cd STAR-DCE
python test_dce.py \
--lowlight_images_path "dir-to-your-test-set" \
--parallel True \
--snapshots_folder snapshots_test/STAR-DCE \
--val_batch_size 1 \
--pretrain_dir snapshots/STAR-ori/Epoch_best.pth \
--model STAR-DCE-Ori \
To evaluated the DCE-Net model you trained,
cd STAR-DCE
python test_dce.py \
--lowlight_images_path "dir-to-your-test-set" \
--parallel True \
--snapshots_folder snapshots_test/DCE \
--val_batch_size 1 \
--pretrain_dir snapshots/DCE/Epoch_best.pth \
--model DCE-Net \
Citation
If this code helps your research, please cite our paper :)
@inproceedings{zhang2021star,
title={STAR: A Structure-Aware Lightweight Transformer for Real-Time Image Enhancement},
author={Zhang, Zhaoyang and Jiang, Yitong and Jiang, Jun and Wang, Xiaogang and Luo, Ping and Gu, Jinwei},
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
pages={4106--4115},
year={2021}
}
225 Dec 29, 2022
4 May 8, 2022
248 Dec 23, 2022
30 Dec 18, 2022
84 Dec 26, 2022
442 Dec 16, 2022
95 Nov 22, 2022
132 Nov 28, 2022
236 Dec 21, 2022
148 Dec 30, 2022
281 Dec 30, 2022
62 Nov 21, 2022
158 Nov 24, 2022
110 Dec 20, 2022
77 Dec 27, 2022
35 Oct 16, 2022
17 Oct 12, 2022
49 Nov 10, 2022