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 |
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.
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.
An autoencoder consists of two main components.
The encoder progressively reduces the spatial dimensions of the image through downsampling while extracting meaningful features.
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.
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.
A standard encoder-decoder architecture built using convolutional layers.
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].
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.
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.
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.
- Python
- TensorFlow
- Keras
- OpenCV
- NumPy
- Matplotlib
- Custom Autoencoder implementation
- U-Net implementation
- Model comparison
- Training and validation curves
- Colorization predictions