Denoising MNIST images using an Autoencoder and Tensorflow in python

Previously I had written sort of a tutorial on building a simple autoencoder in tensorflow. In that tutorial I had used the autoencoder for dimensionality reduction. Check it out if you want to. It has a much more detailed explanation on how to build the autoencoder itself. Here, I’ll use the exact same model to show another use of autoencoders — denoising images.

So let’s get started. I’ll use the famous MNIST handwriting data here.

# Importing tensorflow
import tensorflow as tf
# importing the data
from tensorflow.examples.tutorials.mnist import input_data
# Importing some more libraries
import matplotlib.pyplot as plt
from numpy import loadtxt
import numpy as np
from pylab import rcParams

I’ll import both the train and test set. I’ll train our model using the train set first. Then i’ll put in noisy test data and see if cleaner images come out. Now mnist doesn’t give us noisy data — no worries, we’ll just make some ourselves.

# reading the data
mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)
X_train = mnist.train.images
X_test = mnist.test.images

This sweet peice of code below prints out ten test images alogwith their noisy version.

# deciding how big we want our print out to be
rcParams['figure.figsize'] = 20,20
# looping through the first 10 test images and printing  them out
for i in range(10):
    plt.subplot(1, 10, i+1)
    plt.imshow(X_test[i].reshape(28,28),  cmap='Greys')
    plt.axis('off')
plt.show()
# printing out the noisy images
for i in range(10):
    plt.subplot(1, 10, i+1)
    plt.imshow(X_test_noisy[i].reshape(28,28),  cmap='Greys')
    plt.axis('off')
plt.show()

The reason we need to reshape the images before printing is because we’re supplied with an array of 784 for each observation, which is just 28*28 pixels.

Now i’ll add some noise to the test images. I’ll just create an array as same as the test dataset and just add that to the original data.

# creating the noise matrix
n_rows = X_test.shape[0]
n_cols = X_test.shape[1]
mean = 0.5
stddev = 0.3
noise = np.random.normal(mean, stddev, (n_rows, n_cols))
# creating the noisy test data by adding X_test with noise
X_test_noisy = X_test + noise

Now i’ll build my autoencoder. I’m not going into details about the model. Im using the same model as I’ve used here. If you’re interested please have a look at that tutorial where I’ve explained the model somewhat. Its a lot of code but all that it is, is a neural network with a 784 node input layer, two 32 node hidden layers and again a 784 nodes output layer. An image goes in, an image come out.

# Deciding how many nodes each layer should have
n_nodes_inpl = 784  #encoder
n_nodes_hl1  = 32  #encoder
n_nodes_hl2  = 32  #decoder
n_nodes_outl = 784  #decoder
# first hidden layer has 784*32 weights and 32 biases
hidden_1_layer_vals = {
'weights':tf.Variable(tf.random_normal([n_nodes_inpl,n_nodes_hl1])),
'biases':tf.Variable(tf.random_normal([n_nodes_hl1]))  }
# second hidden layer has 32*32 weights and 32 biases
hidden_2_layer_vals = {
'weights':tf.Variable(tf.random_normal([n_nodes_hl1, n_nodes_hl2])),
'biases':tf.Variable(tf.random_normal([n_nodes_hl2]))  }
# second hidden layer has 32*784 weights and 784 biases
output_layer_vals = {
'weights':tf.Variable(tf.random_normal([n_nodes_hl2,n_nodes_outl])),               'biases':tf.Variable(tf.random_normal([n_nodes_outl])) }
# image with shape 784 goes in
input_layer = tf.placeholder('float', [None, 784])
# multiply output of input_layer wth a weight matrix and add biases
layer_1 = tf.nn.sigmoid(
       tf.add(tf.matmul(input_layer,hidden_1_layer_vals['weights']),
       hidden_1_layer_vals['biases']))
# multiply output of layer_1 wth a weight matrix and add biases
layer_2 = tf.nn.sigmoid(
       tf.add(tf.matmul(layer_1,hidden_2_layer_vals['weights']),
       hidden_2_layer_vals['biases']))
# multiply output of layer_2 wth a weight matrix and add biases
output_layer = tf.matmul(layer_2,output_layer_vals['weights']) +
               output_layer_vals['biases']
# output_true shall have the original image for error calculations
output_true = tf.placeholder('float', [None, 784])
# define our cost function
meansq =    tf.reduce_mean(tf.square(output_layer - output_true))
# define our optimizer
learn_rate = 0.1   # how fast the model should learn
optimizer = tf.train.AdagradOptimizer(learn_rate).minimize(meansq)

Now I’ll train this model with the training images. I’m doing a couple thousand epochs here which means a couple thousand passes through the whole training set, with a batch size of hundred, that is training the net with a hundred images at a time. My computers too slow, but you’re probably gonna have to run way more epochs to see some good results.

# initialising stuff and starting the session
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
# defining batch size, number of epochs and learning rate
batch_size = 100  # how many images to use together for training
hm_epochs =1000    # how many times to go through the entire dataset
tot_images = X_train.shape[0] # total number of images

# running the model for a 10000 epochs taking 100 images in batches
# total improvement is printed out after each epoch
for epoch in range(hm_epochs):
    epoch_loss = 0    # initializing error as 0
    for i in range(int(tot_images/batch_size)):
        epoch_x = X_train[ i*batch_size : (i+1)*batch_size ]
        _, c = sess.run([optimizer, meansq],\
               feed_dict={input_layer: epoch_x, \
               output_true: epoch_x})
        epoch_loss += c
     print('Epoch', epoch, '/', hm_epochs, 'loss:',epoch_loss)

Done! Its a good idea to do some validation every once in a while using the test set. Now for the see if the autoencoder actually works as a denoiser. I’ll put in a couple of noisy images and see what comes out.

# pick any image
any_image = X_test_noisy[234]
# run it though the autoencoder
output_any_image = sess.run(output_layer,\
                   feed_dict={input_layer:[any_image]})