STAR pytorch

Star-Pytorch

Реализация для бумаги «Звезда: легкий трансформатор с учетом структуры для улучшения изображений в реальном времени» (ICCV 2021).

CVF (PDF)

Star-Dce

Реализация Pytorch повышенного освещения со STAR на наборе данных Adobe-Mit Fivek. Вы можете найти его в каталоге Star-DCE. Здесь мы принимаем Pipleline of Zero-DCE (Paper | Code), просто заменив основную цепь CNN на звезду. В нулевом DCE для каждого изображения сеть будет регрессировать группу кривых, которая затем будет применена на исходном изображении итеративно. Вы можете найти более подробную информацию в оригинальном репо Zero-DCE.

Требования

• Numpy
• Einops
• факел
• TOCHVISION
• OpenCV

Таблицы

Мы предоставляем ссылки для загрузки для наборов данных Adobe-Mit Fivek, которые мы использовали (Train | Test). Обратите внимание, что мы применяем тестовый набор, разбитый Deedupe для справедливого сравнения.

Обучение моделей DCE

Чтобы тренировать оригинальную модель звездного 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 

Чтобы обучить базовый DCE-Net на основе CNN (W или 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 

или

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