/Image-Colorization

Primary LanguagePython

Colorization - PyTorch Implementation

Input Output

This repository contains the Pytorch implementation of the following paper:

Colorful Image Colorization
Richard Zhang, Phillip Isola, Alexei A. Efros
https://arxiv.org/abs/1603.08511

Abstract: Given a grayscale photograph as input, this paper attacks the problem of hallucinating a plausible color version of the photograph. This problem is clearly underconstrained, so previous approaches have either relied on significant user interaction or resulted in desaturated colorizations. We propose a fully automatic approach that produces vibrant and realistic colorizations. We embrace the underlying uncertainty of the problem by posing it as a classification task and use class-rebalancing at training time to increase the diversity of colors in the result. The system is implemented as a feed-forward pass in a CNN at test time and is trained on over a million color images. We evaluate our algorithm using a “colorization Turing test,” asking human participants to choose between a generated and ground truth color image. Our method successfully fools humans on 32% of the trials, significantly higher than previous methods. Moreover, we show that colorization can be a powerful pretext task for self-supervised feature learning, acting as a cross-channel encoder. This approach results in state-of-the-art performance on several feature learning benchmarks.

Architecture

Each conv layer refers to a block of 2 or 3 repeated conv and ReLU layers, followed by a BatchNorm layer. The net has no pool layers. All changes in resolution are achieved through spatial downsampling or upsampling between conv blocks. We train a deep CNN to map from a grayscale input to a distribution over quantized color value outputs using the architecture. We train the network on the 1.3M images from the ImageNet training set, validate on the first 10k images in the ImageNet validation set, and test on a separate 10k images in the validation set. While image colorization is a boutique computer graphics task, it is also an instance of a difficult pixel prediction problem in computer vision. Here we have shown that colorization with a deep CNN and a well-chosen objective function can come closer to producing results indistinguishable from real color photos.

Usage

The blow script will colorize an image.

python colorization.py -i imgs/wild.jpg

If you have a Nvidia GPU, use the below script to run colorise the image using GPU's computational power.

python colorization.py -i imgs/wild.jpg --use_gpu

Results

Input Output