Skip to content

A-matindoust/Image-Colorization-Autoencoder-UNet

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Image Colorization using Autoencoder and U-Net

This repository implements two deep learning models for image colorization, where a grayscale image is used to predict its corresponding RGB image.

Two different encoder-decoder architectures are implemented and compared:

  • Custom Autoencoder
  • U-Net

The project demonstrates the strengths and limitations of each architecture for image-to-image translation tasks.

Model Strengths Weaknesses
Custom Autoencoder Simple architecture, fewer parameters Blurry details, lower reconstruction quality
U-Net Better spatial detail, sharper outputs, superior colorization More parameters and higher computational cost

Project Overview

Image colorization is the task of predicting a realistic color image from a grayscale input.

Unlike image classification, the model predicts an entire image rather than a single class. Every pixel in the grayscale image must be assigned three color values (RGB), making this a challenging image reconstruction problem.


What is an Autoencoder?

An autoencoder is a neural network that learns to reconstruct its input.

Instead of predicting class labels, the network predicts another image. During training, the predicted image is compared with the target image to calculate the reconstruction loss.

For example:

  • Input: Grayscale image
  • Target: Original RGB image

Because the target image is generated directly from the input data rather than manually labeled, autoencoders are commonly considered a form of unsupervised (or self-supervised) learning.


Autoencoder Architecture

An autoencoder consists of two main components.

Encoder

The encoder progressively reduces the spatial dimensions of the image through downsampling while extracting meaningful features.

Decoder

The decoder reconstructs the original image by progressively upsampling the encoded feature representation.

Although an autoencoder can be implemented as a single model, defining the encoder and decoder separately provides several advantages:

  • Reuse the encoder for feature extraction.
  • Access the latent feature representation.
  • Maintain a symmetric architecture.
  • Improve readability and model organization.

Dataset

The project uses the Intel Image Classification Dataset.

Original image size:

  • 150 × 150

Training image size:

  • 64 × 64

Although 128 × 128 is generally preferred because powers of two are commonly used in convolutional neural networks, the images were resized to 64 × 64 due to hardware memory limitations.


Models

Custom Autoencoder

A standard encoder-decoder architecture built using convolutional layers.

Observation

Training without Batch Normalization resulted in unstable optimization.

Batch Normalization helps stabilize the feature distributions throughout the network, improving gradient flow and making optimization more reliable. This is especially beneficial when the final layer uses a sigmoid activation to produce outputs within the range [0, 1].


U-Net

The second architecture is a U-Net.

Unlike a standard autoencoder, U-Net introduces skip connections between the encoder and decoder, allowing high-resolution spatial information to be preserved during reconstruction.


Results

Both models successfully learn the overall color distribution of the images.

However, the custom autoencoder struggles to reconstruct fine object boundaries and detailed textures.

The U-Net produces significantly sharper images because the skip connections preserve spatial information that would otherwise be lost during downsampling.

Overall, U-Net provides noticeably better image quality for the colorization task.


Dataset Importance

The diversity of the training dataset strongly influences model performance.

The model learns statistical relationships between objects and their typical colors. For example, if most training images contain yellow houses, the network may incorrectly predict yellow even for houses that are actually blue.

A diverse dataset with many object appearances and color variations is therefore essential for good generalization.


Technologies

  • Python
  • TensorFlow
  • Keras
  • OpenCV
  • NumPy
  • Matplotlib

Results

  • Custom Autoencoder implementation
  • U-Net implementation
  • Model comparison
  • Training and validation curves
  • Colorization predictions

About

Image colorization using a custom Autoencoder and U-Net with TensorFlow/Keras.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors