ML-PyTorch

Colorize black and white images using PyTorch.

Google colab

Try it directly in Google Colab.

Improved version: Google Colab.
NOTE: (5GB) this may take hours

Installation

Install Python

Install required libraries

pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install "numpy<2.0" matplotlib

Download dataset

We are going to use a subset of Food-101 dataset: pizza, steak and sushi.

Download it here (15MB)

pizza_steak_sushi/
   train/
       pizza/
           image01.jpeg
           ...
       steak/
           image04.jpeg
           ...
       sushi/
           image07.jpeg
           ...
   test/
       pizza/
           image101.jpeg
           ...
       steak/
           image104.jpeg
           ...
       sushi/
           image107.jpeg
           ...

# Unzipping archive
tar -xzvf pizza_steak_sushi.zip --directory dataset

# Deleting non-needed anymore archive
rm pizza_steak_sushi.zip

Run

python colorize.py

Result