Skip to content

ChristopheKar/colorize

Repository files navigation

Grayscale Image Colorization

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 Lebanon, Achrafieh, circa 1910 - Colorized

Quick Start

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.

Installation

First, start by cloning this repository:

git clone https://github.com/ChristopheKar/colorize

Dependencies

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.

Docker

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.

Quickstart

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

Roadmap

  • 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?

References

About

Colorize grayscale images

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published