Denoising Images with Convolutional Autoencoder

NeuralNinja
4 min readDec 28, 2023

--

Introduction:

In computer vision, handling noisy images can be challenging since sensor constraints are the primary cause of this problem. Typically, denoising the images involves using both traditional and cutting-edge methods. Autoencoder is a type of neural network that is one of the methods used for denoising the images. It is an unsupervised convolutional-based technique that consists of two components: an encoder and a decoder. The encoder first encodes the significant features, which the decoder component then decodes. In this instance, the encoder takes a noisy image and encodes information, from which the decoder generates a denoised image. The output is reconstructed image based on the encoded features of the input.

For this work, the commonly used dataset, CIFAR-10, has been used. For implementing the model, the TensorFlow library in Python has been used. Noise has been added to the original images in the dataset, and using an autoencoder denoised outputs have been achieved. The performance of the model has been assessed using image quality metric.

Discussions:

The encoder and the decoder are the two main components of the autoencoder as discussed earlier. Convolutional layers are used to encode the key features, and the decoder uses this condensed representation (latent vector) to reconstruct the image. The model is trained to minimize the Mean Squared Error (MSE) between the original and reconstructed images.

TensorBoard is used to monitor the training process, and Adam is used as an optimizer for training the autoencoder. A collection of noisy, original, and denoised images from the test dataset are visualized to evaluate the autoencoder’s performance. The side-by-side comparison illustrates how well the model can eliminate noise without compromising essential information.

For denoised images, the average Multi-Scale Structural Similarity Index (MS-SSIM) has been calculated. This statistic gives an exact indication of how closely the denoised images resemble their original images. The average MS-SSIM achieved by the model is 0.853.

Conclusion:

This practical tutorial has covered the TensorFlow implementation of a Convolutional Autoencoder for image denoising. The quantitative measures and visual results show how well the model restores images to their original.

It is important to note that this is only the starting point. Even better outcomes will come up from further experiments with the architecture, hyperparameters, and additional evaluation measures.

Data Visualization, Model Implementation and Testing:

import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D
from tensorflow.keras.models import Model
from tensorflow.keras.callbacks import TensorBoard
from skimage.metrics import structural_similarity as ssim

# Load CIFAR-10 dataset
(X_train, _), (y_test, _) = cifar10.load_data()

# Normalize pixel values to be between 0 and 1
X_train = X_train.astype('float32') / 255.0
y_test = y_test.astype('float32') / 255.0

# Add noise to the images
noise_value = 0.2
X_train_noised = X_train + noise_value * np.random.normal(loc=0.0, scale=1.0, size=X_train.shape)
y_test_noised = y_test + noise_value * np.random.normal(loc=0.0, scale=1.0, size=y_test.shape)

# Clip values to be between 0 and 1
X_train_noised = np.clip(X_train_noised, 0., 1.)
y_test_noised = np.clip(y_test_noised, 0., 1.)


# Display a few original and noised images
n = 5

plt.figure(figsize=(10, 4))
for i in range(n):
# Display original images
ax = plt.subplot(2, n, i + 1)
plt.imshow(y_test[i])
plt.title('Original')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

# Display noisy images
ax = plt.subplot(2, n, i + 1 + n)
plt.imshow(y_test_noised[i])
plt.title('Noisy')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

plt.show()
Few original and corresponding noised images
# Define the autoencoder model
input_img = Input(shape=(32, 32, 3))

# Encoder
x = Conv2D(32, (3, 3), activation='relu', padding='same')(input_img)
x = MaxPooling2D((2, 2), padding='same')(x)

# Decoder
x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
x = UpSampling2D((2, 2))(x)
output = Conv2D(3, (3, 3), activation='sigmoid', padding='same')(x)

model = Model(input_img, output)

model.compile(optimizer='adam', loss='mse')

# Train the autoencoder and save the history
h = model.fit(X_train_noised, X_train,
epochs=30,
batch_size=64,
shuffle=True,
validation_data=(y_test_noised, y_test),
callbacks=[TensorBoard(log_dir='/tmp/model')])
# Plot training & validation loss
plt.plot(h.history['loss'])
plt.plot(h.history['val_loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Validation'], loc='upper right')
plt.show()
Training report of model
# Denoise some test images
decoded_imgs = model.predict(y_test_noised)

# Plot the original, noisy, and denoised images
n = 10
plt.figure(figsize=(20, 6))

for i in range(n):
# Display original images
ax = plt.subplot(3, n, i + 1)
plt.imshow(y_test[i])
plt.title('Original')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

# Display noisy images
ax = plt.subplot(3, n, i + 1 + n)
plt.imshow(y_test_noised[i])
plt.title('Noisy')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

# Display denoised images
ax = plt.subplot(3, n, i + 1 + 2 * n)
plt.imshow(decoded_imgs[i])
plt.title('Denoised')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

# Adjust the vertical space between rows
plt.subplots_adjust(wspace=0.5, hspace=0.5)

plt.show()
Few test samples with corresponding noised (used for training the model) and denoised images (reconstructed by model)
msssim_values = []
for i in range(len(y_test)):
msssim_value, _ = ssim(y_test[i], decoded_imgs[i], channel_axis=-1, full=True)
msssim_values.append(msssim_value)

# Calculate the average MS-SSIM
average_msssim = np.mean(msssim_values)

print(f'Average MS-SSIM: {average_msssim}')

Average MS-SSIM: 0.853

--

--

No responses yet