STAR pytorch

Star-Pytorch

Implementação para o papel "Estrela: um transformador leve com reconhecimento de estrutura para aprimoramento da imagem em tempo real" (ICCV 2021).

CVF (PDF)

Star-DCE

A implementação do Pytorch de baixa luz com o Star no conjunto de dados do Adobe-MIT Cincok. Você pode encontrá-lo no diretório Star-DCE. Aqui adotamos a linha PIPLELE DE ZERO-DCE (Código de Paper | Em Zero-DCE, para cada imagem, a rede regredirá um grupo de curvas, que será aplicado na imagem de origem iterativamente. Você pode encontrar mais detalhes no repo-DCE original.

Requisitos

• Numpy
• Einops
• tocha
• Torchvision
• OpenEncv

Datas

Fornecemos links para download para conjuntos de dados do Adobe-MIT Cincok que usamos (trem | teste). Observe que adotamos o conjunto de testes dividido por Deepupe para comparação justa.

Modelos de DCE de treinamento

Para treinar um modelo Star-DCE original,

 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 

Para treinar o DCE-Net baseado na CNN (w ou 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 

ou

 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 

Avaliação de modelos treinados

Para avaliar o modelo Star-DCE que você treinou,

 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 

Para avaliar o modelo DCE-Net que você treinou,

 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 

Citação

Se este código ajudar sua pesquisa, cite nosso artigo :)

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