Official Keras Implementation of ChromaGAN: Adversarial Picture Colorization with Semantic Class Distribution [WACV 2020] [arXiv] [Supplementary Material]
Linux
Python 3
NVIDIA GPU + CUDA CuDNN (CPU mode and CUDA without CuDNN may work with minimal modification, but untested)
git clone https://github.com/pvitoria/ChromaGAN
cd ChromaGAN/
pip install -r requirements.txt
Download dataset and place it in the /DATASET/ folder.
We have train our model with ImageNet dataset.
You can download it from here
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
To train the network:
cd ChromaGAN/SOURCE/
python ChromaGAN.py
Models are saved to ./MODELS/DATASET/TEST_NAME/
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/
You can donwload the pretrained weights from here. In order to test the network you should use the file called ` my_model_colorization.h5.
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}
}
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.
You can colorize arbitrary grayscale images using a pre-packaged Docker image from the Replicate registry: https://beta.replicate.ai/pvitoria/ChromaGAN The docker has been made by [Andreas Jansson] Andreas Jansson.
ChromaGAN has been recomendeed as one of the "10 Data Science Projects Every Beginner should add to their Portfolio". [Link]