ChromaGAN

Official Keras Implementation of ChromaGAN: Adversarial Picture Colorization with Semantic Class Distribution [WACV 2020] [arXiv]

Network Architecture

Prerequisits

Linux

Python 3

NVIDIA GPU + CUDA CuDNN (CPU mode and CUDA without CuDNN may work with minimal modification, but untested)

Getting Started

Clone Repository

git clone https://github.com/pvitoria/ChromaGAN
cd ChromaGAN/

Requirements

pip install -r requirements.txt

Download the dataset

Download dataset and place it in the /DATASET/ folder. We have train our model with ImageNet dataset. You can download it from here

Network Parameters

All the parameters can be modified from the config.py file. Note: Modify the name of the dataset in the config file in DATASET. For each test you can modify the folder name in TEST_NAME. The variable PRETRAINED should be changed by the name of your pretrained colorization file.

import os

# DIRECTORY INFORMATION
DATASET = "imagenet" # modify
TEST_NAME ="test1" # modify
ROOT_DIR = os.path.abspath('../')
DATA_DIR = os.path.join(ROOT_DIR, 'DATASET/'+DATASET+'/')
OUT_DIR = os.path.join(ROOT_DIR, 'RESULT/'+DATASET+'/')
MODEL_DIR = os.path.join(ROOT_DIR, 'MODEL/'+DATASET+'/')
LOG_DIR = os.path.join(ROOT_DIR, 'LOGS/'+DATASET+'/')

TRAIN_DIR = "train"
TEST_DIR = "test"

# DATA INFORMATION
IMAGE_SIZE = 224
BATCH_SIZE = 10


# TRAINING INFORMATION
PRETRAINED = "my_model_colorization.h5" 
NUM_EPOCHS = 5

Training

To train the network:

cd ChromaGAN/SOURCE/
python ChromaGAN.py

Models are saved to ./MODELS/DATASET/TEST_NAME/

Testing

To test the network you can either run the code directly from Colab using our Demo or run the code as follows :

cd ChromaGAN/SOURCE/
python ChromaGANPrint.py

Images are saved to ./RESULT/DATASET/TEST_NAME/

Pretrained Weights

You can donwload the pretrained weights from here. In order to test the network you should use the file called ` my_model_colorization.h5.

Citation

If you use this code for your research, please cite our paper ChromaGAN: An Adversarial Approach for Picture Colorization:

@inproceedings{vitoria2020chromagan,
  title={ChromaGAN: Adversarial Picture Colorization with Semantic Class Distribution},
  author={Vitoria, Patricia and Raad, Lara and Ballester, Coloma},
  booktitle={The IEEE Winter Conference on Applications of Computer Vision},
  pages={2445--2454},
  year={2020}
}

Aknowledgments

The authors acknowledge partial support by MICINN/FEDER UE project, reference PGC2018-098625-B-I00 VAGS, and by H2020-MSCA-RISE-2017 project, reference 777826 NoMADS. We also thank the support of NVIDIA Corporation for the donation of GPUs used in this work.