Implementation for paper "STAR: A Structure-aware Lightweight Transformer for Real-time Image Enhancement" (ICCV 2021).

Last update: Dec 21, 2022

Related tags

Deep Learning STAR-pytorch

Overview

STAR-pytorch

Implementation for paper "STAR: A Structure-aware Lightweight Transformer for Real-time Image Enhancement" (ICCV 2021).

CVF (pdf)

STAR-DCE

The pytorch implementation of low light enhancement with STAR on Adobe-MIT FiveK dataset. You can find it in STAR-DCE directory. Here we adopt the pipleline of Zero-DCE ( paper | code ), just replacing the CNN backbone with STAR. In Zero-DCE, for each image the network will regress a group of curves, which will then applied on the source image iteratively. You can find more details in the original repo Zero-DCE.

Requirements

numpy
einops
torch
torchvision
opencv

Datesets

We provide download links for Adobe-MIT FiveK datasets we used ( train | test ). Please note that we adopt the test set splited by DeepUPE for fair comparison.

Training DCE models

To train a original STAR-DCE model,

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 \

To train the baseline CNN-based DCE-Net (w\ or 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 \

Evaluation of trained models

To evaluated the STAR-DCE model you trained,

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 \

To evaluated the DCE-Net model you trained,

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 \

Citation

If this code helps your research, please cite our paper :)

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

Implementation for paper "STAR: A Structure-aware Lightweight Transformer for Real-time Image Enhancement" (ICCV 2021).

Related tags

Overview

STAR-pytorch

STAR-DCE

Requirements

Datesets

Training DCE models

Evaluation of trained models

Citation

Owner

Multi-Modal Machine Learning toolkit based on PyTorch.

This is an open-source toolkit for Heterogeneous Graph Neural Network(OpenHGNN) based on DGL [Deep Graph Library] and PyTorch.

CityLearn Challenge Multi-Agent Reinforcement Learning for Intelligent Energy Management, 2020, PikaPika team

Trains an agent with stochastic policy gradient ascent to solve the Lunar Lander challenge from OpenAI

Code for WSDM 2022 paper, Contrastive Learning for Representation Degeneration Problem in Sequential Recommendation.

Code for Neurips2021 Paper "Topology-Imbalance Learning for Semi-Supervised Node Classification".

Code used for the results in the paper "ClassMix: Segmentation-Based Data Augmentation for Semi-Supervised Learning"

[ICCV2021] Learning to Track Objects from Unlabeled Videos

DCSL - Generalizable Crowd Counting via Diverse Context Style Learning

Voice Conversion by CycleGAN (语音克隆/语音转换)：CycleGAN-VC3

A trusty face recognition research platform developed by Tencent Youtu Lab

Continuous Query Decomposition for Complex Query Answering in Incomplete Knowledge Graphs

Official Pytorch implementation of "Learning Debiased Representation via Disentangled Feature Augmentation (Neurips 2021, Oral)"

Simple Tensorflow implementation of Toward Spatially Unbiased Generative Models (ICCV 2021)

PyTorch implementations of the NeRF model described in "NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis"

Official repository of ICCV21 paper "Viewpoint Invariant Dense Matching for Visual Geolocalization"

Pytorch implementation of SELF-ATTENTIVE VAD, ICASSP 2021

PyToch implementation of A Novel Self-supervised Learning Task Designed for Anomaly Segmentation

Code implementation from my Medium blog post: [Transformers from Scratch in PyTorch]

DR-GAN: Automatic Radial Distortion Rectification Using Conditional GAN in Real-Time