Image Denoising Using Auto Encoders

In this post, we will discuss, how to remove noise from an image using a type of neural network called Autoencoder. Autoencoders can learn efficient data codings in an unsupervised manner. The purpose of an autoencoder is to learn a representation for a set of data, typically for dimensionality reduction, by training the network to ignore signal noise.

Image Credits: [Source]

As described in the image, on both sides of the network we have the same image passed as input and output. By doing this, we create a bottleneck in the network which is nothing but a low dimensional representation of the input image. The layers before the bottleneck represent the encoder which encodes the input image. And then we have a decoder, which does the reverse of what we did in the encoder to get the input image back.

Note: If this seems a little bit confusing to you, don't worry, it's a lot easier to understand with the code.

Step 1: Importing Libraries

• In this step, We will import some libraries. We will also import our dataset which is available in Keras datasets.

import numpy as np

from tensorflow.keras.datasets import mnist
from matplotlib import pyplot as plt
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.callbacks import EarlyStopping, LambdaCallback
from tensorflow.keras.utils import to_categorical

%matplotlib inline

Step 2: Data Preprocessing

• We will load our data by calling the load_data() method on mnist object we imported in the previous step. 
• You may know that image data is nothing but an array of pixel values, these values range from 0-255. We divide our input images by 255 to convert all the values between 0-1. This is called 'normalization'.
• Our input images are of size 28x28, and it is a 2D vector, before feeding it to the model, we need to flatten these pixels into a 1D vector. So we reshape the image to be a 1D array or 784 values (28x28=784), and 60,000 is the number of total images we have in our dataset.

(X_train, y_train), (X_test, y_test) = mnist.load_data()

X_train = X_train / 255.
X_test = X_test / 255.

X_train = np.reshape(X_train, (60000, 784))
X_test = np.reshape(X_test,(10000,784))

Step 3: Adding Noise

• In this step, we are gonna add some noise to the pixel values of the images in our dataset.
• What we are exactly doing is pretty simple, we are taking the input values of the pixels and add some another random value to it. Here I am using a noise factor of 0.9, you can use whatever you want.
• But one another thing is that when we add any random value to an input image, it might result in a value bigger than 0-1, which makes the whole normalization process to be pointless. So we clip these values to be in between 0.0 to 1.0 using np.clip() function.

X_train_noisy = X_train + np.random.rand(60000, 784)*0.9
X_test_noisy = X_test + np.random.rand(10000, 784)*0.9
X_train_noisy = np.clip(X_train_noisy, 0., 1.)
X_test_noisy = np.clip(X_test_noisy, 0., 1.)

def plotimg(x, p, labels=False):
    for i in range(10):
        plt.figure(figsize=(20,2))
        plt.subplot(1, 10, i+1)
        plt.imshow(x[i].reshape(28,28), cmap='binary')
        plt.xticks([])
        plt.yticks([])
        if labels:
            plt.xlabel(np.argmax(p[i]))
        plt.show()
plotimg(X_train, None)
plotimg(X_train_noisy, None)

Using the above function we plot a few images from the noisy dataset and actual dataset.

Noisy Image

Actual Image

Step 4: Building and Training a Classifier

• In this step, we will create a simple classifier using a Sequential model of Keras. 
• We specify our input shape to be 784 because the flattened dimension of our image is 28x28=784
• We are only specifying one hidden layer with 256 neurons
• Finally, we have our output layer, with 10 neurons because we have 10 digits to be classified(0-9), and we will use the softmax layer for our output layer.

classifier = Sequential([
                         Dense(256, activation='relu', input_shape=(784,)),
                         Dense(256, activation='relu'),
                         Dense(10, activation='softmax')
            ])

classifier.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics = ['accuracy']
)

classifier.fit(X_train, y_train, batch_size=512, epochs=5)

• We then compile our classifier using adam optimizer and cross-entropy loss and accuracy as the measure.
• We will fit our classifier on the training dataset for 5 epochs with a batch size of 512.

Evaluating Our classifier on test set:

loss, acc = classifier.evaluate(X_test, y_test)
print(acc)

313/313 [==============================] - 0s 1ms/step - loss: 0.0778 - accuracy: 0.9761
0.9761000275611877

Evaluating our classifier on the noisy dataset:

loss, acc = classifier.evaluate(X_test_noisy, y_test)
print(acc)

313/313 [==============================] - 0s 1ms/step - loss: 6.5984 - accuracy: 0.1972
0.1972000002861023

As we can observe our classifier does a really good job of classifying normal images. but performs poorly on noisy images. To overcome this, we will use Autoencoder.

Step 5: Building the Autoencoder

• We will create a simple autoencoder here. First, we will define an Input layer that takes in images with size (784,).
• Then we create an encoding layer with only 64 neurons. The output of this layer will be an encoded representation of our input image.
• Then we have another layer which again converts this image back to 784 pixels.
• Finally, we tie them together to create an autoencoder model. and then we compile it.

input_image = Input(shape=(784,))

encoded = Dense(64, activation='relu')(input_image)

decoded = Dense(784, activation='sigmoid')(encoded)

autoencoder = Model(input_image, decoded)
autoencoder.compile(loss='binary_crossentropy', optimizer='adam')

Step 6: Training the Autoencoder

• In this step, we will fit our autoencoder on noisy images and use actual images as the label, so that encoder will try to learn only useful features in our noisy images and ignore noise which isn't present in our label images.

autoencoder.fit(X_train_noisy, X_train, epochs=100,
                batch_size=512, validation_split=0.2,
                verbose=False,
                callbacks = [
                             EarlyStopping(monitor='val_loss', patience=5),
                             LambdaCallback(on_epoch_end=lambda e,l: print('{:.3f}'.format(l['val_loss']),end=' _ '))
                ])
print('Training completed!')

Once the training is completed we can use our autoencoder to predict test images and see if it can remove noise.

Step 7: Denoised Images

• We will get some predictions on the noisy test data.

preds = autoencoder.predict(X_test_noisy)

plot(X_test_noisy, None)

Noisy Image

plot(preds, None)

Predicted Image

Our predictions are really close to the actual image, that's really cool. isn't it?

Step 8: Creating a composite Model

• Lastly, we create a composite model. Which will do noise reduction first and then do the classification as well.
• Thus, we create an input layer and pass its output to an autoencoder, which will remove the noise, and finally, we pass it to another classifier model to do the classification.
• We will call this new model denoise_and_classify.
• We will also plot some predicted images along with their labels.

input_image = Input(shape=(784,))
x = autoencoder(input_image)
y = classifier(x)

denoise_and_classify = Model(input_image, y)
predictions = denoise_and_classify.predict(X_test_noisy)
plot(X_test_noisy, predictions, True)

Conclusion:

Here we created an autoencoder to denoise an image, although this is a really basic example, it will give you some idea about the autoencoders and why we use it.

Along with autoencoders, we also learned to do mnist digit classification using neural networks which were really cool. And we also combined two models to create one composite model which does both image denoising and classification.

Hope you enjoyed it!