Various models for handling underexposure, overexposure, super-resolution, shadow removal, etc.
- Python 3.6
- TensorFlow v1.15
- OpenCV
- Pillow
- scikit-image
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
- 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
To ensure all modules in repo can be found you must update your PYTHONPATH environment variable:
export PYTHONPATH=$PYTHONPATH:/path/to/ImageEnhancement
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
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
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
To be added...
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)
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
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
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
To be added...
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