STAR pytorch

Pytorch

Implementación para el documento "Estrella: un transformador liviano consciente de la estructura para la mejora de la imagen en tiempo real" (ICCV 2021).

CVF (PDF)

Star-DCE

La implementación de Pytorch de mejora de baja luz con STAR en el conjunto de datos Adobe-Mit FiveK. Puede encontrarlo en el directorio Star-DCE. Aquí adoptamos la línea de pipa de cero-DCE (papel | código), solo reemplazando la columna vertebral CNN con STAR. En cero-DCE, para cada imagen, la red retrocedirá un grupo de curvas, que luego se aplicará en la imagen de origen iterativamente. Puede encontrar más detalles en el Repo Cero-DCE original.

Requisitos

• numpy
• einops
• antorcha
• vía antorcha
• opencvv

Fechas

Proporcionamos enlaces de descarga para conjuntos de datos Adobe-Mit FiveK que utilizamos (Train | Test). Tenga en cuenta que adoptamos el conjunto de pruebas dividido por DeepUpe para una comparación justa.

Entrenamiento de modelos DCE

Para entrenar un modelo original de Star-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 

Para entrenar la red basada en DCE-Net (w o w o agrupación),

 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 

o

 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 

Evaluación de modelos capacitados

Para evaluar el modelo Star-DCE que entrenó,

 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 evaluar el modelo de red DCE que entrenó,

 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 

Citación

Si este código ayuda a su investigación, cite nuestro documento :)

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