STAR pytorch

星号

纸质“星星：实时图像增强的结构意识轻巧变压器”的实现（ICCV 2021）。

CVF（PDF）

Star-dce

在Adobe-Mit五号数据集上使用Star的低光增强的Pytorch实现。您可以在Star-DCE目录中找到它。在这里，我们采用零-DCE（纸|代码）的Pipleline，只需用Star代替CNN骨干。在零-DCE中，对于每个图像，网络将回归一组曲线，然后将其应用于源图像迭代。您可以在原始回购零DCE中找到更多详细信息。

要求

• numpy
• Einops
• 火炬
• 火炬
• OPENCV

日期组

我们为使用的Adobe-Mit五K数据集提供下载链接（火车|测试）。请注意，我们采用Deepupe拆分的测试集进行公平比较。

培训DCE模型

训练原始的星际模型，

 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 

训练基于CNN的基线DCE-NET（W 或W o POLING），

 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 

或者

 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 

评估训练的模型

要评估您训练的Star-DCE模型，

 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 

为了评估您训练的DCE-NET模型，

 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 

引用

如果此代码有助于您的研究，请引用我们的论文:)

 @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}
}