DeblurGAN

An easy-to-read implementation of DeblurGAN using PyTorch

Some demos of deblurring:

Prerequisites

• NVIDIA GPU + CUDA cuDNN
• Python 3.7

Folder Structure

deblurGAN/
│
├── deblur_image.py - deblur your own images
├── test.py - evaluation of trained model
├── train.py - main script to start training
├── make_aligned_data.py - make aligned data
├── config.json - demo config file
├── config_aligned.json - demo config file using aligned dataset
├── pretrained_weights/ - some pretrained weights for test
│   ├── GAN/ - folder of pretrained weights using GAN loss
│   └── WGAN_GP/ - folder of pretrained weights using WGAN_GP loss
│
├── base/ - abstract base classes
│   ├── base_data_loader.py - abstract base class for data loaders
│   ├── base_model.py - abstract base class for models
│   └── base_trainer.py - abstract base class for trainers
│
├── data_loader/ - dataloader and dataset
│   ├── data_loader.py
|   └── dataset.py 
│
├── data/ - default directory for storing input data, containing 2 directory for blurred and sharp
│   ├── blurred/ - directory for blurred images
│   └── sharp/ - directory for sharp images
│
├── model/ - models, losses, and metrics
│   ├── layer_utils.py
│   ├── loss.py
│   ├── metric.py
│   └── model.py
│
├── trainer/ - trainers
│   └── trainer.py
│
└── utils/
    ├── logger.py - class for train logging
    ├── util.py
    ├── visualization.py - class for tensorboardX visualization support
    └── ...

Config file format

{
    "name": "DeblurGAN",                         // training session name
    "n_gpu": 1,                                  // number of GPUs to use for training
    "data_loader": {                             // selecting data loader
        "type": "GoProDataLoader",
        "args": {
            "data_dir": "data/",
            "batch_size": 1,
            "shuffle": false,
            "validation_split": 0.1,
            "num_workers": 4
        }
    },
    "generator": {                               // architecture of generator
        "type": "ResNetGenerator",
        "args": {
            "input_nc": 3,
            "output_nc": 3
        }
    },
    "discriminator": {                           // architecture of discriminator
        "type": "NLayerDiscriminator",
        "args": {
            "input_nc": 3
        }
    },
    "loss": {                                    // loss function
        "adversarial": "wgan_gp_loss",
        "content": "perceptual_loss"
    },
    "metrics": [                                 // list of metrics to evaluate 
        "PSNR"
    ],
    "optimizer": {                               // configuration of the optimizer (both generator and discriminator)
        "type": "Adam",
        "args": {
            "lr": 0.0001,
            "betas": [
                0.5,
                0.999
            ],
            "weight_decay": 0,
            "amsgrad": true
        }
    },
    "lr_scheduler": {                            // learning rate scheduler
        "type": "LambdaLR",
        "args": {
            "lr_lambda": "origin_lr_scheduler"
        }
    },
    "trainer": {                                 // configuration of the trainer
        "epochs": 300,
        "save_dir": "saved/",
        "save_period": 1,
        "verbosity": 2,
        "monitor": "max PSNR",
        "tensorboardX": true,
        "log_dir": "saved/runs"
    },
    "others": {                                  // other hyperparameters
        "gp_lambda": 10,
        "content_loss_lambda": 100
    }
}

How to run

• Train

   python train.py --config config.json

• Resume

    python train.py --resume path/to/checkpoint

• Test

    python test.py --resume path/to/checkpoint

• Deblur

    python deblur_image.py --blurred path/to/blurred_images --deblurred path/to/deblurred_images --resume path/to/checkpoint

• Make aligned data first if you want to use aligned dataset

    python make_aligned_data.py --blurred path/to/blurred_images --sharp path/to/sharp_images --aligned path/to/aligned_images

Tips

• If you want to use gan_loss instead of wgan_gp_loss, use_sigmoid must be set to true in generator.
• Aligned dataset could boost the speed of data_loader a little bit. So run make_aligned_data.py to get aligned dataset before training.
• Pretrained weights of both GAN and WGAN_GP are available.
• Download pretrained weights： https://drive.google.com/open?id=1w-u0r3hd3cfzSjFuvvuYAs9wA-E-B-11

Acknowledgements

The organization of this project is based on PyTorch Template Project