Skip to content

Vadbeg/colorization-coreml

Repository files navigation

Colorization CoreML

This is a repository for image colorization using CoreML. The model is based on the Colorful Image Colorization paper and the original GitHub repository. The model used in this repository is named siggraph17.

res1 res2

Installation

  1. Create a virtual environment:
virtualenv -p python3.9 .venv && source .venv/bin/activate
  1. Install all requirements:
pip install -r requirements.txt
  1. Use the project 🎉

Usage

To run use command below:

python run_colorization.py --data-root images

If an image is not in black and white, it will be converted to black and white before colorization. It's worth noting that the model inference will only work on MacOS due to restrictions of the coremltools framework.

Models

The first model, colorizer_core.mlmodel, expects a PIL.Image named image_small with a size of (3, 256, 256). It returns an array containing colorized_small_image, an A and B values from the LAB color space (2, 256, 256).

Afterward, you need to resize colorized_small_image to the original image size, convert this array to PIL.Image. Let's name the result ab.

The second model, colorizer_tail.mlmodel, requires the original PIL.Image, image, of any size, as well as the ab (i.e., the resized colorized_small_image from the previous model). This model returns an array containing the colorized image, colorized_image, of the original size.

Built With

  • coremltools - The NNs inference framework used
  • OpenCV - Images processing framework used
  • Pillow - Images processing framework used
  • Typer - CLI framework used

License

This project is licensed under the MIT License - see the LICENSE file for details

Authors

Releases

No releases published

Packages

No packages published

Languages