CSRGAN

Official PyTorch implementation of the our paper:

Image Super-Resolution and Colorization in a single Generative Adversarial Network
Erez Yosef and Shay Shomer Chai

Overview

We designed a single GAN to perform the two tasks:

• Image Colorization
• Image Super resolution

Generator Input: Grayscale low resolution image (64x64 px.)

Generator Output: Colored High resolution image (256x256 px.)

Left to right: Input, Bicubic interpolation result, CSRGAN result (Ours), HR-GT.

Methods:

We used and implemented the following methods:

• Generative adversarial networks, Goodfellow et al. In Advances in Neural Information Processing Systems (NIPS). 2014.
• SRGAN Photo-realistic single image super-resolution using a GAN, Ledig et al. Computer Vision and Pattern Recognition (CVPR). 2016.
• ESRGAN Esrgan: Enhanced super-resolution generative adversarial networks. Ledig Et Al, ECCVW, 2018.
• Relativistic GAN The relativistic discriminator: a key element missing from standard gan, Jolicoeur-Martineau, A. 2018
• Image colorization GAN Image Colorization using Generative Adversarial Networks, Nazeri Et Al, 2018
• conditional GAN Conditional Generative Adversarial Nets, Mehdi Mirza and Simon Osindero, 2014.

Requirements

conda env create -f requirments.yml

Dataset

In this paper we used Stanford University dogs dataset
This dataset contains:
Number of dog breeds: 120
Number of images: 20,580
The images are in different resolution thus we add resize(256X256) to the preprocess

We use a small portion of the dataset: ~30 different dog breeds ~6K images.
We splited our dataset to train and test with respect to the different dog breeds.
Train: ~5400 images
Test: ~560 images

Training Details

We used google cloud services in order to train our model - using:

• 1 GPU - Tesla K80
• 2 CPU - 13GB RAM

We trained our model for 21 epochs, approximately 8hr.

Training Models

• GeneratorFeatures
• Generator
• GeneratorBroken63
• GeneratorRRDB
  

For each Generator one can use regular GAN or Conditional GAN(example downwards)

Train

Input arguments for train.py:

--epochs                | default=21   
--saveparams_freq_batch | default=5    
--saveimg_freq_step     | default=100  
--lrG                   | default=1e-4 
--lrD                   | default=1e-5 
--train_path            | default='./data/6k_data/train'
--test_path             | default='./data/6k_data/test'
--type_of_dataset       | default="10_dogs"
--fname                 | default=""
--generator             | default="GeneratorFeatures"
--discriminator         | default="Discriminator"
--batch_size            | default=16

• command line for our default model(regular GAN - GeneratorFeatures)

python train.py

• command line for different model(Conditional GAN - GeneratorBroken63)

python train.py --generator GeneratorBroken63 --discriminator Conditional_Discriminator

Outputs - Train

The code creates out dirs - training results\<MONTHDD_HH_MM>\ (for example: training_results\Jan28_21_23) with 3 subdirs:

• params - weights saved every 5 epochs ( input argument --saveparams_freq_batch)
• SR_results - images generated by the generator every 100 steps(from trainset) and every 1 epoch(form testset) (input argument --saveimg_freq_step)
• train_data - summary csv and tensorboard results

Test

Input arguments for test.py

--train_path            | default='./data/6k_data/train'
--test_path             | default='./data/6k_data/test'
--type_of_dataset       | default="10_dogs"
--fname                 | default=""
--params_path           | default="params/6k_params/netG_epoch_20.pth"
--batch_size            | default=16

python test.py

The test script is implemented for our best model: GeneratorFeatures

Outputs - Test

The test script creates new dir - test_results where we can find colorized super resolution images generated by our GeneratorFeatures.

Examples