This project aims to colorize grayscale images by applying the work done by Zhang et al. in their 2016 ECCV paper, Colorful Image Colorization, which is beautifully implemented and showcased in this repo. This project is also partly based on this excellent tutorial by Adrian Rosebrock on pyimagesearch, for the OpenCV-Caffe implementation of one of the models.
Currently, this repository implements both the ECCV16 (Caffe and PyTorch) and the SIGGRAPH17 (PyTorch) models by Zhang et al. in one simple interface, and only in inference mode. Furthermore, GPUs are not yet supported by the code or dependencies, meaning that even if the relevant GPU drivers and libraries are installed, the inference will not yet utilize it. See Roadmap section.
Lebanon, Achrafieh, circa 1910 - Colorized
The weights for the different models can be retrieved using:
# ECCV16 Caffe model
wget http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel -O weights/zhang_eccv16/caffe/colorization_release_v2.caffemodel
# ECCV16 PyTorch model
wget https://colorizers.s3.us-east-2.amazonaws.com/colorization_release_v2-9b330a0b.pth -O weights/zhang_eccv16/zhang-eccv16-9b330a0b.pth
# SIGGRAPH17 PyTorch model
wget https://colorizers.s3.us-east-2.amazonaws.com/siggraph17-df00044c.pth -O weights/zhang-siggraph17-df00044c.pth
If all dependencies are installed, the example images can be colorized by running:
python colorize.py -i images/ -o outputs/
By default, this will use the Zhang et. al. SIGGRAPH-17 model and look for the weights under
weights/zhang-siggraph17-df00044c.pth
.
First, start by cloning this repository:
git clone https://github.com/ChristopheKar/colorize
This project uses Python 3.9, and the dependencies listed in requirements.txt
can
be installed with pip
, conda
, or any other package manager, in a virtual environment
or other. For example, using pip
:
pip install -r requirements.txt
The main dependencies are PyTorch and OpenCV, any recent version should do.
The environment can also be setup using Docker and the provided Dockerfile
.
First, build the image by running the following command in this repository:
docker build -t colorize .
Then, using the built image is as simple as:
docker run -it --rm --name colorize -v $PWD:/app colorize -i images -o outputs
This will run the Docker container and execute the colorize.py
script with the given arguments.
You can override this entrypoint by setting --entrypoint /bin/bash
for an interactive bash
session, for example.
The -v
option allows you to mount your current workspace to /app
inside the container, so that your files are accessible from there, and so that any
changes made to files under that path persists on your local storage. Any other changes
made inside the Docker container, e.g. installing additional packages or creating files
outside of /app
, will not persist across sessions unless you commit your changes
to the Docker image.
The easiest way to use this repository with Docker is with the provided utility script, run.sh
.
First, make sure it is executable (chmod +x run.sh
) and simply execute it (./run.sh
).
By default, its entrypoint is the colorize.py
script,
so it can directly take in the arguments for that script, for example
# Use ./run.sh --help to show usage details
# Build image for GPU
./run.sh --shell --build --gpus all
# Colorize images
./run.sh -i images/ -o outputs -m siggraph17
# Colorize images on GPU
./run.sh --gpus all -i images/ -o outputs -m siggraph17
- Zhang et al. ECCV16 Caffe model
- Zhang et al. ECCV16 PyTorch model (13/03/2022)
- Zhang et al. SIGGRAPH17 PyTorch model (13/03/2022)
- GPU Support (20/03/2022)
- Training support
- Other colorizers?
- Colorful Image Colorization [paper] [code]
- Black and white image colorization with OpenCV and Deep Learning, PyImageSearch
- Let there be color!