Image Colorizer using Deep Learning

A deep learning project that colorizes grayscale images using a pretrained ResNet18 encoder and a custom decoder implemented in PyTorch.

---

Overview

This project predicts the missing a and b channels in the LAB color space from the grayscale L channel.

Instead of generating RGB values directly, the model learns color representations in the LAB color space, resulting in a simpler and more stable learning problem.

---

Features

• Pretrained ResNet18 encoder (Transfer Learning)
• Custom decoder using ConvTranspose2D layers
• LAB color space preprocessing
• Modular PyTorch implementation
• Model checkpointing
• Standalone inference pipeline
• Automatic output image generation
• Easy to extend to a U-Net architecture

---

Project Structure

image-colorizer/

├── checkpoints/
│   └── best_colorization_model.pth
│
├── notebooks/
│   └── model_trainer.ipynb
│
├── outputs/
│
├── test_images/
│
├── inference.py
├── model.py
├── utils.py
│
├── README.md
├── requirements.txt
└── .gitignore

---

Model Architecture

Input Image (RGB)
        │
        ▼
Convert RGB → LAB
        │
        ▼
Extract L Channel
        │
        ▼
Gray-to-RGB Convolution
        │
        ▼
Pretrained ResNet18 Encoder
        │
        ▼
Custom Decoder
        │
        ▼
Predicted ab Channels
        │
        ▼
Merge L + ab
        │
        ▼
LAB → RGB
        │
        ▼
Colorized Image

---

Dataset

The model is trained using RGB images.

During preprocessing:

• Convert RGB images to LAB color space
• Extract the L channel as input
• Use the ab channels as targets
• Normalize both input and target channels

---

Technologies Used

• Python
• PyTorch
• Torchvision
• NumPy
• Pillow
• Matplotlib
• scikit-image

---

Installation

Clone the repository

git clone https://github.com/<your-username>/image-colorizer.git
cd image-colorizer

Install dependencies

pip install -r requirements.txt

---

Model Checkpoint

Place the trained model inside:

checkpoints/
    best_colorization_model.pth

---

Running Inference

Place an image inside the test_images/ directory.

Run:

python inference.py --image test_images/dog.jpeg

The colorized image will be saved automatically in:

outputs/

---

Results

Original Image

Add screenshot

Grayscale Input

Add screenshot

Colorized Output

Add screenshot

---

Future Improvements

• ResNet18 U-Net architecture
• Skip connections
• Fine-tuning the encoder
• Perceptual loss
• Attention mechanisms
• Higher resolution training
• FastAPI backend
• Next.js frontend
• Cloud deployment

---

Learning Outcomes

This project covers:

• Custom PyTorch datasets
• Data preprocessing
• LAB color space
• Transfer learning
• Encoder-decoder architectures
• Image colorization
• Model training
• Model checkpointing
• Inference pipelines
• Modular project organization

---

Author

Suchay Joshi

GitHub: https://github.com/sucxay

---

License

This project is licensed under the MIT License.