Using Conditional Generative Adversial Networks (cGANs), Sketch2drawings performs paired image-to-image translation on sketches and drawings. This deep learning mapping allows the project to turn a black and white sketch into a colorized drawing.
Table of Contents
Left is input, center is output, and right is target
- Download data from this kaggle dataset
- Prepare/Preprocess the data
- Train the model
- Test the model
- Export the model??
Create a folder called original and edges under images so that the directory looks like this
.
├── README.md
├── docker
│ └── Dockerfile
├── images
│ └── original
│ └── edges
│ └── resized
│ └── combined
Move the images (only, no folders) to images/original
We will
- Resize all images into 256 x 256
- Detect edges to get the "sketch"
- Combine input images and target images
- Split combined images into train and val set
Install all the dependencies
pip3 install -r requirements.txt
Warning: Python version must be at max 3.6! I spent too much time trying to do with Python 3.8 D:
If you have conda installed, you can also try this
- Create virtual environment named sketch2drawings
conda create -n "sketch2drawings" python=3.6.0
- Activate conda environment
conda activate sketch2drawings
- Install OpenCV and Tensorflow v1.4.1 (since numpy is already installed)
conda install opencv-python
pip install tensorflow==1.4.1
After putting all the images into the original folder, run
python preprocessing/process.py --input_dir images/original --operation resize --output_dir images/resized
This operation took me about 25 minutes to run. When finished, we should see a folder called resize with 256 x 256 images
We used Canny to detect edges. Navigate to the folder with process in it and then run process.py. This cd is important because it's used to figure out where the image folder is located.
cd preprocessing
python edge_detection.py
Navigate back to the root directory with cd ..
. Now run the combine operation with
python preprocessing/process.py --input_dir images/resized --b_dir images/edges --operation combine --output_dir images/combined
This operation took me about 30 minutes to run. The script will skip over files that already exist, so you can pause the operation and resume later.
Generate train/validation splits
python preprocessing/split.py --dir images/combined
Hopefully you have a GPU because if you train on CPU you will definitely be waiting for a bit.
python pix2pix.py --mode train --output_dir s2d_train --max_epochs 200 --input_dir images/combined/train --which_direction BtoA --ngf 32 --ndf 32
Maybe try changing --ngf 32
and --ndf32
to 64 to see how well it does, but it takes more computation
If you have Docker installed, you can use the Docker image to train without having to install the correct version of Tensorflow
# Train the model with docker
python dockrun.py python pix2pix.py \
--mode train \
--output_dir s2d_train \
--max_epochs 200 \
--input_dir images/combined/train \
--which_direction BtoA
python pix2pix.py --mode test --output_dir s2d_test --input_dir images/combined/val --checkpoint s2d_train
Results shouldbe in a new folder called s2d_test
python pix2pix.py --mode export --output_dir export/ --checkpoint s2d_train/ --which_direction BtoA
After running this command, you should see a folder called export
with checkpoint
and various files inside it.
python3 export_checkpoint.py --checkpoint export --output_file s2d.pict
This will create a file called s2d.pict
.
I also tested this with the anime-sketch-colorization-pair dataset and it works well for some sketches (Canny tends to not give very clear edges). However, for some reason, it doesn't work as well when I port it to Tensorflow.js, which I'm not sure why