In this post, we will learn how to implement an autoencoder to denoise images using TensorFlow and Keras.
Introduction:
Autoencoders are a type of neural network that can learn to reconstruct input data by encoding it into a lower-dimensional representation and then decoding it back to the original shape. Denoising autoencoders are specifically designed to remove noise from the input data.
Let’s go through the steps:
Step 1: Import The Required Libraries
import tensorflow as tf from tensorflow import keras from tensorflow.keras import layers
Step 2: Define the autoencoder architecture
input_img = keras.Input(shape=(28, 28, 1)) # Encoder x = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(input_img) x = layers.MaxPooling2D((2, 2), padding='same')(x) x = layers.Conv2D(8, (3, 3), activation='relu', padding='same')(x) x = layers.MaxPooling2D((2, 2), padding='same')(x) # Decoder x = layers.Conv2D(8, (3, 3), activation='relu', padding='same')(x) x = layers.UpSampling2D((2, 2))(x) x = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(x) x = layers.UpSampling2D((2, 2))(x) decoded = layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x) autoencoder = keras.Model(input_img, decoded)
In the above code, we define an autoencoder architecture using convolutional and upsampling layers. The encoder part reduces the dimensionality of the input image, while the decoder part upsamples the encoded representation to reconstruct the original image.
Step 3: Compile the model
autoencoder.compile(optimizer='adam', loss='binary_crossentropy')
We use the Adam optimizer and binary cross-entropy loss since we are treating the problem as a pixel-wise binary classification task.
Step 4: Preprocess the data
# Load the dataset (example: MNIST)
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
# Normalize and reshape the input images
x_train = x_train.astype('float32') / 255.0
x_train = tf.expand_dims(x_train, axis=-1)
x_test = x_test.astype('float32') / 255.0
x_test = tf.expand_dims(x_test, axis=-1)
Step 5: Add noise to the input images
# Add random noise to the training and test images noise_factor = 0.5 x_train_noisy = x_train + noise_factor * tf.random.normal(shape=x_train.shape) x_test_noisy = x_test + noise_factor * tf.random.normal(shape=x_test.shape) # Clip the values to [0, 1] x_train_noisy = tf.clip_by_value(x_train_noisy, clip_value_min=0.0, clip_value_max=1.0) x_test_noisy = tf.clip_by_value(x_test_noisy, clip_value_min=0.0, clip_value_max=1.0)
Step 6: Train the autoencoder
autoencoder.fit(x_train_noisy, x_train, epochs=10, batch_size=128, shuffle=True, validation_data=(x_test_noisy, x_test))
Adjust the number of epochs and batch size according to your needs.
Step 7: Evaluate the autoencoder
decoded_imgs = autoencoder.predict(x_test_noisy)
# Calculate the reconstruction loss (MSE)
mse = tf.keras.losses.mean_squared_error(x_test, decoded_imgs)
reconstruction_loss = tf.reduce_mean(mse)
print(f"Reconstruction Loss: {reconstruction_loss:.4f}")
Step 8: Visualize the results
import matplotlib.pyplot as plt
n = 10 # Number of images to display
plt.figure(figsize=(20, 4))
for i in range(n):
# Original images
ax = plt.subplot(2, n, i + 1)
plt.imshow(tf.squeeze(x_test_noisy[i]))
plt.title("Original + Noise")
plt.gray()
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
# Decoded images
ax = plt.subplot(2, n, i + n + 1)
plt.imshow(tf.squeeze(decoded_imgs[i]))
plt.title("Denoised")
plt.gray()
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
plt.show()
This code will display a comparison between the original images with added noise and the denoised images generated by the autoencoder.
Finally we have now implemented an autoencoder to denoise images using TensorFlow and Keras.