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.
- Create a virtual environment:
virtualenv -p python3.9 .venv && source .venv/bin/activate- Install all requirements:
pip install -r requirements.txt- Use the project 🎉
To run use command below:
python run_colorization.py --data-root imagesIf 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.
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.
- coremltools - The NNs inference framework used
- OpenCV - Images processing framework used
- Pillow - Images processing framework used
- Typer - CLI framework used
This project is licensed under the MIT License - see the LICENSE file for details