Image Colorization using Simple AutoEncoder with Custom Data

Chirag Chauhan
4 min readOct 7, 2022

--

Image colorization is the process of assigning colors to a single-channel grayscale image to produce a 3-channel colored image output. Before moving to the code section, let’s understand what an AutoEncoder is!!!

What is an Autoencoder?

Autoencoders are a type of self-supervised learning model that can learn a compressed representation of input data. An autoencoder is a neural network model that seeks to learn a compressed representation of the input. They are an unsupervised learning method, although technically, they are trained using supervised learning methods, referred to as self-supervised. There are various types of autoencoders like Sparse autoencoders, Denoising autoencoders, Variational autoencoders, LSTM autoencoders etc.

Source: V7 labs

Applications of autoencoders include:

  • Anomaly detection
  • Data denoising (ex. images, audio)
  • Image inpainting
  • Information retrieval

Dataset Information

Here for this example, I have taken a custom image dataset of a barn. Initially, I downloaded 12 images from google but that was not yielding much accuracy so I used the image augmentation technique to produce 84 images from those 12 images. The sample code for augmentation:

Link to the dataset images folder created by me.

Color Models (RGB and LAB)

The RGB color model is one of the most widely used color representation methods in computer graphics. It uses a color coordinate system with three primary colors: Red, Green, and Blue. Each primary color can take an intensity value ranging from 0(lowest) to 1(highest). Mixing these three primary colors at different intensity levels produces a variety of colors. This is the basic color model which is used in most of the colour image processing problems while grayscale images are used for B/W processing tasks.

Another color model that will be actually helpful in our project is the LAB color model. Using this model we are going to train our Autoencoder. In LAB or CIE L*a*b*, L* axis represents Lightness, ranging from 0–100. 0 is black (no light); and 100 is white (maximum illumination). The a* axis is green at one extremity (represented by -a), and red opposite it at the other (+a). The b* axis has blue at one end (-b), and yellow (+b) opposite it at the other. In theory, there are no maximum values of a* and b*, but in practice, they are usually numbered from -128 to +127 (256 levels). This color model is extensively used in many industries apart from printing and photography.

Autoencoder Model

  1. Let's start by importing all the required libraries and reading the images through ImageDataGenerator from TensorFlow. Since the L channel encodes only the intensity, we can use the L channel as our grayscale input to the network. From there the network must learn to predict the a and b channels. Given the input L channel and the predicted ab channels, we can then form our final output image. After that, we would split our images into X and Y where X represents L and Y represents A and B from the LAB color model. Please note that after converting X into L space it will only have the shape (32,256,256) so we have to reshape it into the (32,256,256,1) by adding 1 to its original shape to show that it is a single channel grayscale image.

2. Encoder-Decoder Architecture

The encoder encodes/compresses and finds patterns/important features from the images to learn while the decoder decodes the compressed image to its original version to predict the output. We will be using the functional model API instead of using sequential model API. The code for the encoder and decoder are as follows:

The Autoencoder architecture would be a combination of the encoder and decoder model whose code is as follows:

3. Train the model

The accuracy is not that good but still, it is acceptable as the dataset was of just 84 images.

Testing on some other grayscale image

For predicting we will convert the RGB image to the LAB version to make a prediction using our model and for viewing the final result we have to convert our predicted image to RGB.

Output from my results

Conclusion

This is a very basic method of image colorization using autoencoders. There are many State-of-the-art models available now that yield better results but consider this a very basic step to learn about image colorization techniques.

--

--