Implementation of deep learning model in Keras for image colorization. Project uses U-Net trained as Self-Attention GAN together with Perceptual loss instead of usual MAE or MSE. Work is still in progress.
Historical photos
Original RGB images (left) converted to grayscale and colorized (right)
- Python 3.x
- pip
Install pip packages using
$ pip install -r requirements.txt
Add .env file to project root with environmental variables
COMET_PROJECTNAME={comet_project_name}
COMET_WORKSPACE={comet_workspace}
COMET_API_KEY={comet_api_key}
[optional]
There is a Docker image included that was used for training in cloud. You can build it from local Dockerfile with
docker build -t ml-box .
or get it from Docker Hub
docker pull tomikeska/ml-box
Train model using command
$ python src/train_gan.py
Colorize image using trained weights
$ python src/colorize.py --weights model/weights.h5 --source source.jpg --output output.jpg
Code is released under the MIT License. Please see the LICENSE file for details.