A convolutional neural network that colorizes sketches using roughly made color hints
The repository contains the code for SketchColorization, a convolutional neural network to color sketches using rough color hints. It contains the code to use it along with pre-trained weights and also code to train the model yoursef.
- Small and fast model. ~8M parameters.
- Uses an alternative custom colorspace to achieve increased colorfulness in predictions.
- Linux
- Anaconda
- NVIDIA GPU (CUDA >= 11)
- pytorch and fastai
Create environmente and install packages
conda env create -f environment.yml- Download Pretrained Weights,
colorizer.pkl, from releases, and put into the models/ folder. - Then run the following script to colorize
python colorize.py sketchimage colorhintimage --savename path/to/save.jpgYou can train using any sketch and color pair dataset, but danbooru2021 was used to train this model. Instructions to download and clean the dataset are included in the file DownloadAndClean.ipynb.
To generate linearts from the color images to train the model you can use the notebook LA/toLA.ipynb in which there are 4 methods available sketch_keras, sketch_simplification, anime2sketch and XDOG. For the CNN methods, you must download the pretrained weights for model on the releases.
Follow the Train.ipynb notebook. You must download the pretrained weights for the custom vgg in order to use perceptual loss. Download the xyv_vgg.pth from releases and place it into the models/ folder.
