An implementation of the Pix2Pix paper for image colorization using PyTorch.
Try a trained on Google Colab:
This project applies conditional adversarial networks to colorize grayscale images. The model learns to translate grayscale images into colorful ones by training on paired datasets.
- Implemented in PyTorch
- Supports custom datasets
- Includes a training script and a jupyter notebook on google colab to test a trained model hosted on huggingface model hub.
- Data augmentation and preprocessing
- Logging to tensorboard
It has been trained with 12000 images from COCO 2017 dataset during 150 epochs with batch size 16 and learning rate 2e-4 for both generator and discriminator. Data augmentation is done by randomly flipping horizontally the images.
git clone https://github.com/yourusername/colorizing_images.git
cd colorizing_images
pip install -r requirements.txtpython train.pyHere are some examples of colorized images, the upper row are the generations of the model:
- The original Pix2Pix paper
- PyTorch framework
- HuggingFace🤗 Model Hub
This project is licensed under the MIT License.