Capstone-Project-Machine-Learning-Engineer-Nanodegree

Colorize Grayscale images with GANs !

Mean SSIM on 3000 validation samples : 0.87

Capstone Proposal

This repo expose my Capstone Project Proposal for my ML Engineer Nanodegree at Udacity : Image colorization

To get hyperlinks on references I've done, please download it.

My Capstone Project is explained in this file

You can also find in this repository my Notebook to train a Generative Adversarial Network, to colorize grayscale images.

This GAN has been built with Pytorch. Feel free to re-code it in Tensorflow if you prefer.

To install dependencies, you just have to install libraries from requirements.txt

pip install -r requirements.txt

I'm using FASTAI Package to download the COCO Dataset (Dataset used to train GAN). There is a cell to do it directly in the notebook.

Weights are upload and available on Google Drive.

If you don't want to train the model, and just play with inference :

Download the weights of trained models :

Zip file have :

final-weights.pt (Weights of the full GAN trained)
res18-unet.pt (Weights of the pretrained Generator)
train_0x.pickle (Only available for Latest Train (Best) This file contain Loss / SSIM per epoch.

To train this model, I used a Tesla V100 32GB, and it tooks 9 hours.

You can train it on CPU, but I highly recommend to use GPU.

You can try model with visualize method to predict and generate images from validation Dataset.