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