Image Colorization with Deep Learning

Generate color images from grayscale images

Today I would like to show the attraction topic about image colorizes from grayscale images. This time I use Pytorch to create Neural Network (NN) and use DCGAN technique.

Dataset

If you have a strong GPU, you can convert many images at once. In this article, I use 118k images.

I have divided the dataset into two parts, 116k for train data and 2k for test data.

Training

In the training process, I will train with 116k images, 256×256 size, and 1,000 epoch.

you can download code and docker at : https://github.com/vannponlork/colorize_image

DCGAN Architecture

Generator tries to generate an image that similar to the real image and lets Discriminator judge whether it is the real image or fake. Discriminator have 2 inputs, real image and generated image by Generator. It tries to classify which one is the real image and the fake image.

DCGAN technique

Well, for DCGAN technique I have two types of NN Generator and Discriminator.

So, to make a color image from grayscale, Generator needs input in one channel and output with 2 channels. Generator tries to generate images and verify with Discriminator.

In this case, I use images of LAB where L is a greyscale image and AB is color information.

Generator tries to find the other AB color by the input L image. First, the program needs to convert RGB to LAB image and split L and AB. Generator tries to predict the AB image from the L image. After that, merge the results with the output of Generator with L and note as Fake. Discriminator predicts the merged image of original L with original AB as real image and output of generator with L as a fake image by using Binary-Cross-Entropy (BCE).

Generator updates the fake output of Discriminator by using BCE and update generate AB image by using Mean-Square-Error (MSE).

Deploy model to production

In this topic, I will deploy the model by Flask. Note that in the training process I train 256×256, but in prediction, you can predict any size of images.

There two approach to send image to predict :

1. byte-like
2. JSON format

Conclusion

Artificial intelligence is useful for everyday life. There are multiple tricks to make color images from greyscale. In this article, I choose DCGAN technique to make color images from greyscale. Next time I will try to find another topic in a neural network to show you. We will walk along the way in artificial intelligence together see you next time.

All reference code here: https://github.com/vannponlork/colorize_image