Image Colorization

Transform black-and-white images into vibrant, lifelike color using GAN-based models.

Table of Contents

• Overview
• Features
• Repository Structure
• Modules
• Getting Started
  • Prerequisites
  • Installation
  • Usage
  • Tests
• Project Roadmap
• Contributing
• License
• Acknowledgments

---

Overview

This project aims to colorize black-and-white images using a GAN-based approach inspired by the Pix2Pix model. The goal is to restore colors to grayscale images by training deep neural networks capable of predicting color information from luminance data.

The project can be extended for various applications like:

• Enhancing old photographs
• Augmenting grayscale imagery in film restoration
• Assisting in medical imaging, satellite imagery, and other fields where colorization can enhance insights.

---

Features

• GAN-based Colorization: (Currently) Uses the Pix2Pix model for image-to-image translation from grayscale to color.
• Evaluation Metrics: Includes FID, SSIM, and PSNR for evaluating the quality of generated images.
• Flexible Dataloader: Supports dynamic loading and preprocessing of image data.
• Training and Testing Pipeline: Comes with scripts to train and test the models on custom datasets.
• Visualization Tools: Easily visualize and save colorized results.

---

Repository Structure

└── ImageColorization/
    ├── LICENSE
    ├── README.md
    ├── models
    │   └── pix2pix.py
    ├── scripts
    │   └── train_pix2pix.sh
    ├── test.py
    ├── train.py
    └── utils
        ├── pix2pixDataLoader.py
        ├── pix2pixInitWeights.py
        ├── pix2pixLoss.py
        └── pix2pixMetricMeters.py

---

Modules

Inference

File	Summary
test.py	This script is responsible for evaluating the trained image colorization model on the test dataset. It performs predictions and calculates various metrics to assess model performance.

Training the models

File	Summary
train.py	This script is the main training engine for the image colorization project. It handles the training loop, backpropagation, and saving of model checkpoints.

Models

File	Summary
pix2pix.py	This file contains the Pix2Pix architecture for image colorization, implementing both the generator (UNet-based) and discriminator (PatchGAN-based) models. It also includes the forward pass and loss calculations specific to the Pix2Pix model.

Scripts

File	Summary
train_pix2pix.sh	This file contains the script to train the Pix2Pix model.

Utility Modules

File	Summary
pix2pixLoss.py	This utility handles the different loss functions used during training. It includes the L1 loss for the generator and the GAN loss for both generator and discriminator models.
pix2pixInitWeights.py	This utility is responsible for weight initialization. It ensures that the model weights are initialized correctly, following best practices like Xavier initialization.
pix2pixDataLoader.py	This file includes the data loading utilities used to create PyTorch DataLoaders for the training and validation datasets. It includes data augmentation techniques and handles conversion to LAB color space.
pix2pixMetricMeters.py	This utility provides functionality for tracking various metrics during training and validation, such as accuracy, PSNR (Peak Signal-to-Noise Ratio), and other performance indicators.

---

Getting Started

Prerequisites

• Python: version 3.11

Installation

Build the project from source:

1. Clone the ImageColorization repository:

❯ git clone https://github.com/protyayofficial/ImageColorization

2. Navigate to the project directory:

❯ cd ImageColorization

3. Install the required dependencies:

❯ pip install -r requirements.txt

Usage

To run the project, execute the following command:

❯ chmod +x scripts/train_{model_name}.sh 
❯ scripts/train_{model_name}.sh

For example:

❯ chmod +x scripts/train_pix2pix.sh 
❯ scripts/train_pix2pix.sh

Tests

Execute the test suite using the following command:

❯ python test.py

---

Project Roadmap

• Task 1: Downloaded the COCO 2017 dataset and prepared it for training and evaluation.
• Task 2: Built a robust data loading pipeline with support for LAB color space and essential data augmentation techniques.
• Task 3: Developed custom inference files to facilitate metric logging and visualizations, ensuring effective model evaluation.
• Task 4: Designed and implemented the training loop, with accurate loss calculation, backpropagation, and regular model checkpointing.
• Task 5: Developed a baseline Pix2Pix model for image colorization, establishing the foundation for future model improvements.
• Task 6: Integrated evaluation metrics like SSIM, PSNR, and FID, enabling rigorous testing of the model's performance against real images from the dataset.
• Task 7: Implemented pretrained backbone unet using fastai since the training images are less than what was used in the original literature.
• Task 8: Enhance the architecture to achieve more realistic and vivid image colorization using advanced GAN methodologies and novel techniques.
• Task 9: Conduct extensive hyperparameter tuning to optimize model performance and achieve better convergence.
• Task 10: Deploy the project as a fully functional web application, allowing users to upload grayscale images and receive colorized outputs, utilizing React, Vite, and Tailwind CSS.

---

Contributing

Contributions are welcome! Here are several ways you can contribute:

• Report Issues: Submit bugs found or log feature requests for the ImageColorization project.
• Submit Pull Requests: Review open PRs, and submit your own PRs.
• Join the Discussions: Share your insights, provide feedback, or ask questions.

Contributing Guidelines

1. Fork the Repository: Start by forking the project repository to your github account.
2. Clone Locally: Clone the forked repository to your local machine using a git client.

   git clone https://github.com/protyayofficial/ImageColorization

3. Create a New Branch: Always work on a new branch, giving it a descriptive name.

   git checkout -b new-feature-x

4. Make Your Changes: Develop and test your changes locally.
5. Commit Your Changes: Commit with a clear message describing your updates.

   git commit -m 'Implemented new feature x.'

6. Push to github: Push the changes to your forked repository.

   git push origin new-feature-x

7. Submit a Pull Request: Create a PR against the original project repository. Clearly describe the changes and their motivations.
8. Review: Once your PR is reviewed and approved, it will be merged into the main branch. Congratulations on your contribution!

Contributor Graph

---

License

This project is protected under the MIT License License. For more details, refer to the LICENSE file.

---

Acknowledgments

• This project is inspired by the pioneering work of Isola et al., whose Pix2Pix model has revolutionized the field of image-to-image translation using conditional GANs. Their seminal paper, "Image-to-Image Translation with Conditional Adversarial Networks", played a crucial role in shaping the techniques I used for image colorization in this project. I also appreciate the comprehensive GitHub repository associated with their paper, which was invaluable for understanding and implementing the GAN architecture. Some code was adapted and simplified from their repository to better fit the requirements of this project.

• A special thank you goes to Moein Shariatnia, whose insightful article, "Colorizing Black & White Images with U-Net and Conditional GAN", offered practical guidance for combining U-Net architectures with GANs in the context of image colorization. His tutorial was instrumental in refining the techniques employed in this project. You can explore more of his work through his GitHub.

---