ImageEnhancement

Various models for handling underexposure, overexposure, super-resolution, shadow removal, etc.

Dependencies

• Python 3.6
• TensorFlow v1.15
• OpenCV
• Pillow
• scikit-image

Models

Implementations of the following models are provided:

• CycleGAN by Zhu et al.: Paper | Code
• Mask-ShadowGAN by Hu et al.: Paper | Code
• EnlightenGAN by Jiang et al.: Paper | Code
• DeShadowNet by Liangqiong et al. (in progress): Paper | Code
• SRGAN by Dong et al.: Paper | Code

Datasets

• Download a CycleGAN dataset using:

  bash ./download_cyclegan_dataset.sh [apple2orange|summer2winter_yosemite|horse2zebra|monet2photo|cezanne2photo|ukiyoe2photo|vangogh2photo|maps|cityscapes|facades|iphone2dslr_flower|ae_photos]
  

• Download the Unpaired Shadow Removal (USR) dataset for shadow removal from: USR Dataset
• Download the ISTD dataset for shadow removal from: ISTD Dataset
• Download the EnlightenGAN dataset for low-light image enhancement from: Google Drive
• Download the DIV2K - bicubic downscaling x4 dataset for super-resolution from: Train Low-Res | Train High-Res | Test Low-Res

Updating PYTHONPATH

To ensure all modules in repo can be found you must update your PYTHONPATH environment variable:

export PYTHONPATH=$PYTHONPATH:/path/to/ImageEnhancement

Training

CycleGAN

The CycleGAN model takes approximately 20 hours to train to completion using a Tesla V100 GPU. To train run:

python train/cyclegan_train.py --dirA /path/to/dataA  --dirB /path/to/dataB --batch_size 1 --lr 0.0002 --layer_norm_type instance --weight_init_gain 0.02

Mask-ShadowGAN

The Mask-ShadowGAN model takes approximately 24 hours to train to completion using a Tesla V100 GPU. To train run:

python train/maskshadowgan_train.py --dirA /path/to/shadow_data --dirB /path/to/shadow_free_data --batch_size 1 --lr 0.0002 --layer_norm_type instance --weight_init_gain 0.02 --lambda_ident 0.5

EnlightenGAN

To train run:

python train/enlightengan_train.py --dirA /path/to/low_light_data --dirB /path/to/normal_data --weight_init_gain 0.02 --scale_size 320 --crop_size 320 --patchD_3 5 --self_attention --times_residual --patchD --vgg --patch_vgg --use_ragan --hybrid_loss

DeShadowNet

To be added...

SRGAN

To train run:

python train/srgan_train.py --dir /path/to/high_res_data --batch_size 16 --scale_size 96 --crop_size 384 --weight_init_gain 0.02 --beta1 0.9

To continue training from a saved checkpoint, add the following argument to the end of the command line arguments passed into the training script you are running:

--load_model /checkpoint_dir (e.g. /20022019-0801)

Testing

CycleGAN

To test the CycleGAN model run:

python test/cyclegan_test.py --dir /path/to/dataA --batch_size 1 --layer_norm_type instance --load_model /checkpoint_dir --sample_directory /path/to/save/samples/to

Mask-ShadowGAN

To test the Mask-ShadowGAN model run:

python test/maskshadowgan_test.py --dir /path/to/shadow_data --batch_size 1 --layer_norm_type instance --load_model /checkpoint_dir --sample_directory /path/to/save/samples/to

EnlightenGAN

To test the EnlightenGAN model run:

python test/enlightengan_test.py --dir /path/to/low_light_data --batch_size 1 --scale_size 320 --crop_size 320 --load_model /checkpoint_dir --sample_directory /path/to/save/samples/to --self_attention --times_residual

DeShadowNet

To be added...

SRGAN

To test the SRGAN model run:

python test/srgan_test.py --dir /path/to/low_res_data --batch_size 1 --load_model /checkpoint_dir --sample_directory /path/to/save/samples/to