Interactive Deep Colorization in PyTorch

note: This fork fixes some issues the original repository had with Python 3.

All credits go to Richard Zhang*, Jun-Yan Zhu*, Phillip Isola, Xinyang Geng, Angela S. Lin, Tianhe Yu, and Alexei A. Efros.

Project Page | Paper | Video | Talk | UI code

Fast start

1. clone this repo and create a virtual environment (recommended Python 3.12):

git clone https://github.com/lbarraga/colorization-pytorch-python3.git
cd colorization-pytorch-python3
python3 -m venv venv
source venv/bin/activate

2. Install dependencies:

pip install -r requirements.txt

3. Download the ILSVRC 2012 dataset(training images (728MB) and validation images (6.3GB)). extract the training images and put them in a folder /train. Extract the validation images and put them in a folder /val.

4. Prepare the dataset:

python make_ilsvrc_dataset.py --in_path <folder>

5. Download a pretrained model:

bash pretrained_models/download_siggraph_model.sh

6. Test the model:

python test.py --name siggraph_caffemodel --how_many 10 --mask_cent 0

7. A results folder will be created with the results of the test. You can visualize the results by opening the index.html file in the results folder.

Testing the model

To change the method currently being tested, change the function being called in add_color_patches in util.py.

You can test the model on the validation data by running the following command:

python test_sweep.py --name siggraph_caffemodel --mask_cent 0

in the code of test_sweep.py you can set the how_many flag, that has been hardcoded to something different.

in ./checkpoints/siggraph_caffemodel you will find a csv with the mean and std for each number of points, along with the plot of the test.

To visually compare the results/performance of the different models, run:

python plot.py