Probabilistic Image Colorization
Branch: master
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Type Name Latest commit message Commit time
Failed to load latest commit information.
imagenet-to-hdf5.ipynb Comment out calls to assert_variables_initialized: deprecated Dec 13, 2018

license GitHub repo size in bytes GitHub top language Maintenance GitHub last commit

[PIC] Probabilistic Image Colorization

Tensorflow implementation for Probabilistic Image Colorization - generating diverse and vibrant colorization using auto-regressive generative networks - on the CIFAR and ImageNet datasets.


We develop a probabilistic technique for colorizing grayscale natural images. In light of the intrinsic uncertainty of this task, the proposed probabilistic framework has numerous desirable properties. In particular, our model is able to produce multiple plausible and vivid colorizations for a given grayscale image and is one of the first colorization models to provide a proper stochastic sampling scheme.

Moreover, our training procedure is supported by a rigorous theoretical framework that does not require any ad hoc heuristics and allows for efficient modeling and learning of the joint pixel color distribution. We demonstrate strong quantitative and qualitative experimental results on the CIFAR-10 dataset and the challenging ILSVRC 2012 dataset.

sample1 sample1 sample1 sample1 sample1 Figure: Input grayscale (left), original colored image (rigt) and samples from our model (middle columns)

If you find this work useful, please cite

"Probabilistic Image Colorization"
 Amélie Royer, Alexander Kolesnikov, Christoph H. Lampert
 British Machine Vision Conference (BMVC), 2017



  • Python 2.6+ or 3+
  • Tensorflow 1.0+
  • Numpy
  • h5py (for loading the imagenet dataset)
  • skimage

The main training and evaluation code is in Use python -- help to display the various options available.

Train the model

Train on ImageNet on 4 GPUS with color channels lying in the LAB colorspace and being generated at 1/4 the resolution of the original image.

python --nr_gpus 4 --batch_size 16 --test_batch_size 25 --init_batch_size 100  \
                -lr 0.00016 -p 0.999 -ld 0.99999 -c 160 -l 4 --downsample 4            \
                --color lab --dataset imagenet --gen_epochs 1 --data_dir [data_dir]

Same training on CIFAR.

python --nr_gpus 4 --batch_size 16 --test_batch_size 16 --init_batch_size 100  \
                -lr 0.001 -p 0.999 -ld 0.99995 -c 160 -l 4 --downsample 2              \
                --color lab --dataset cifar --gen_epochs 1 --data_dir [data_dir]

Evaluation of a pre-trained model on the test set

Download the public pre-trained models (ImageNet and CIFAR).

tar -xzvf cifar_model.tar.gz
tar -xzvf imagenet_model.tar.gz

Evaluate the model on the dataset validation split. For instance for ImageNet:

python --nr_gpus 4 --batch_size 16 --test_batch_size 25 --init_batch_size 100  \
               -c 160 -l 4 --downsample 4 --color lab --dataset imagenet --mode "eval" \
               --data_dir [data_dir] --model [path_to_checkpoint .ckpt]

Apply a pre-trained model on selected samples

Apply the model on given colored images to generate (i) reconstruction and (ii) random samples from grayscale version of the input images. The generated images are saved as demo_reconstructions.jpg and demo_generations.jpg respectively.


python --nr_gpus 1 -c 160 -l 4 --downsample 2 --color lab --dataset cifar --test \
               --mode "demo" --model [path_to_checkpoint .ckpt] --input [path to image(s)]

on ImageNet

python --nr_gpus 1 -c 160 -l 4 --downsample 4 --color lab --dataset imagenet \
               --mode "demo" --model [path_to_checkpoint .ckpt] --input [path to image(s)]

Demo example

For instance, to generate reconstructions and samples on the images in samples_val, which are samples from the validation set of the ImageNet dataset, with the pre-trained ImageNet model (Note: set --nr_gpus 0 to run in CPU mode):

tar -xzvf imagenet_model.tar.gz

python --nr_gpus 1 -c 160 -l 4 --downsample 4 --color lab --dataset imagenet \
               --test_batch_size 16 --mode "demo" --model imagenet/model.ckpt         \
               --input "val_samples/*.JPEG"
Input Reconstructions Random samples