Image denoising using an autoencoder
As digital images are prone to various types of noise during acquisition and transmission, denoising them is crucial for enhancing their quality. Let’s explore the step-by-step guide on building an autoencoder model, training it with the MNIST dataset, and applying it to clean a noisy image. This fascinating process is effective and integral in various medical imaging and satellite image analysis applications.
Data preprocessing and model training
Before training our autoencoder, we process the images from the MNIST dataset to ensure they’re in the right form. During training, the autoencoder takes each image, compresses it to capture the essential features, and then reconstructs it. This process teaches it to focus on the important parts of the image. We adjust its inner workings to improve reconstruction, and after training, we save this ability for future use.
import torchvision
def return_dataset():
# Load the MNIST dataset
mnist_train = torchvision.datasets.MNIST(root='./data', train=True, download=True)
mnist_test = torchvision.datasets.MNIST(root='./data', train=False, download=True)
x_train = mnist_train.data
y_train = mnist_train.targets
x_test = mnist_test.data
y_test = mnist_test.targets
x_train = x_train.float()/255.0
x_test = x_test.float()/255.0
return (x_train, x_test)
Code explanation
Lines 1–7: Import necessary libraries and modules such as
numpy,torch, andtqdmfor handling numerical operations, deep learning processes, and progress visualization, respectively. Also, import theAutoencoderclass and thereturn_datasetfunction from the local modules.Line 9: Retrieve the training and testing data using the
return_datasetfunction, which processes and normalizes the MNIST dataset, making it ready for the autoencoder.Lines 11–16: Define the dimensions for the autoencoder’s encoder and decoder. These determine the size of the input, the hidden layers, and the latent space where the image is compressed.
Lines 18–19: Create an instance of the
Autoencoderusing the specified dimensions and set up the loss function (criterion) withnn.MSELoss()and the optimizer withoptim.Adamto refine the model parameters during training.Lines 22–35: Iterate over the training data, preprocess each image into a tensor, perform a forward pass through the autoencoder, calculate the loss, and perform a backward pass to update the model’s weights. This loop trains the autoencoder to compress and then reconstruct the images, focusing on learning the important features.
Line 38: Save the trained autoencoder’s parameters to a file, preserving the model’s ability to process new images.
Evaluation for noise removal
After training the autoencoder to recognize and reconstruct important features, we move on to evaluate its noise removal capabilities. By feeding the autoencoder with a noisy image, we expect it to output a clean version, effectively stripping away the noise and restoring the image’s original quality. Finally by comparing original test images with their noisy and denoised versions, we can visually assess how well the autoencoder has learned to clean the images. This visual evaluation is crucial in understanding the practical effectiveness of our model.
Note: Press the “Run” button and type the
python3 /usercode/denoise.py && python3 -m http.server 3000command in the terminal to see the results.
import random
import numpy as np
def add_noise(img, random_chance=5):
noisy = []
for row in img:
new_row = []
for pix in row:
if random.choice(range(100)) <= random_chance:
new_val = random.uniform(0, 1)
new_row.append(new_val)
else:
new_row.append(pix)
noisy.append(new_row)
return np.array(noisy, dtype=np.float32)Code explanation
Lines 1–8: Load the essential libraries for neural networks (
torchand its submodules) and for numerical and visual representation (numpyandmatplotlib). Also, import theAutoencoderclass, thereturn_datasetfunction for data handling, and theadd_noisefunction to artificially introduce noise to images.Lines 10–20: Set up the dimensions for the autoencoder and load the pretrained model using
torch.load, switching it to evaluation mode withdenoiser.eval().Lines 23–27: Prepare the plots for displaying the noisy and denoised images using
matplotlib. An image from the test dataset is selected, noise is added to it, and it’s plotted as the'Noisy Image'.Lines 30–34: Convert the noisy image to a tensor, pass it through the autoencoder to get the denoised output, and plot the
'Denoising Image'.Line 37: Save the comparison as an image file
comparison.png, which serves as a record of the autoencoder’s performance in noise removal.
Unlock your potential: Autoencoders series, all in one place!
To deepen your understanding of Autoencoders, explore our series of Answers below:
Introduction to autoencoders using PyTorch
Learn the fundamentals of autoencoders and how to implement them using PyTorch for unsupervised learning tasks.Anomaly detection with autoencoders
Discover how autoencoders can identify anomalies by learning normal data patterns and flagging deviations.Image denoising using an autoencoder
Explore how autoencoders can remove noise from images by learning to reconstruct clean versions from noisy inputs.Image reconstruction with autoencoders
Understand how autoencoders compress and reconstruct images, preserving key features while reducing dimensionality.
Free Resources