If you would like a detailed explanation of this project, please refer to the Medium article below.

---

The project is also available for testing on Hugging Face.

---

A U-Net–Based CNN Autoencoder for Cleaning Noisy Images Before Classification

A practical walkthrough of how I built and trained a deep-learning model to denoise images and boost classification performance.

When I first started working with image-classification tasks, I noticed something that kept hurting my models: noise. Even small distortions—random dots, compression artifacts, sensor noise—were enough to confuse the classifier.

The obvious solution was to train on noisy data… but that never felt elegant. Instead, I wanted a preprocessing model whose sole job is to take a noisy image and return a clean version of it. The classifier would then work on much better input.

That idea led me to build a U-Net–based CNN Autoencoder.

This article walks you through:

• why I chose a U-Net structure
• how the autoencoder was built
• how noisy images were generated
• how the model was trained and evaluated
• what results I achieved

Goal: Use a smart deep-learning architecture to clean images before classification.

---

🔧 1. Setting Up the Environment

I started by loading the usual deep-learning stack:

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
import numpy as np
import matplotlib.pyplot as plt

This is the typical setup for building custom architectures using Keras.

---

2. Why a U-Net Autoencoder?

A normal autoencoder compresses an image into a bottleneck and then reconstructs it. It works—but often loses details.

A U-Net, however, uses skip connections, meaning it:

• compresses the image (downsampling)
• learns a compact representation
• reconstructs it (upsampling)
• also reconnects high-resolution features from earlier layers

This makes U-Net excellent for:

• denoising
• segmentation
• super-resolution
• restoration tasks

So instead of a plain autoencoder, I built one using a U-shaped architecture.

---

3. Building the U-Net Autoencoder

Encoder

c1 = Conv2D(64, 3, activation='relu', padding='same')(inputs)
p1 = MaxPooling2D((2, 2))(c1)

c2 = Conv2D(128, 3, activation='relu', padding='same')(p1)
p2 = MaxPooling2D((2, 2))(c2)

Bottleneck

bn = Conv2D(256, 3, activation='relu', padding='same')(p2)

Decoder

u1 = UpSampling2D((2, 2))(bn)
m1 = concatenate([u1, c2])
c3 = Conv2D(128, 3, activation='relu', padding='same')(m1)

u2 = UpSampling2D((2, 2))(c3)
m2 = concatenate([u2, c1])
c4 = Conv2D(64, 3, activation='relu', padding='same')(m2)

Output

outputs = Conv2D(1, 3, activation='sigmoid', padding='same')(c4)

Even though the full architecture is larger, the core idea is:

down → compress → up → reconnect → reconstruct

---

4. Generating & Preprocessing Noisy Images

Instead of downloading a noisy dataset, I artificially added Gaussian noise to MNIST digits:

noise_factor = 0.4
x_train_noisy = x_train + noise_factor * np.random.normal(
    loc=0.0, scale=1.0, size=x_train.shape
)

This created image pairs:

• clean MNIST digit
• noisy version of the same digit

Perfect for training an autoencoder.

---

5. Training the Model

Compile:

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

Train:

model.fit(
    x_train_noisy, x_train,
    epochs=10,
    batch_size=128,
    validation_split=0.1
)

The autoencoder learns one simple rule:

Input: noisy image Output: clean image

---

6. Visualizing the Results

After training, I checked:

• noisy input
• autoencoder output
• original clean image

The model consistently removed a large amount of noise, smoothing textures while preserving structure. Not perfect—but for MNIST and a lightweight U-Net, the results were very encouraging.

---

7. Why This Helps Classification

If you already have (or plan to build) a classifier—CNN, ResNet, etc.—you can use a pipeline like:

Noisy Image → Autoencoder (denoising) → Classifier → Prediction

This helps with real-world noise sources like:

• camera noise
• poor lighting
• compression artifacts
• motion blur

Clean input → better predictions.

---

8. Key Takeaways

• U-Net skip connections help preserve important features.
• Autoencoders serve as powerful preprocessing tools.
• Denoised images can significantly improve classification accuracy.
• The model is lightweight and easy to integrate.
• The approach scales to any image dataset.

This approach is not just theoretical—it’s extremely practical. Any project involving real-world noisy data can benefit from this denoising layer.

---

9. Results

Note: Click the image below to view the video showcasing the project’s results.

---

Note: If the video above is not working, you can access it directly via the link below.

Watch Demo Video