Skip to content

Skyline-9/sketch2drawings

 
 

Repository files navigation

Sketch2Drawings

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.

License: MIT Magic

Table of Contents

Samples

Left is input, center is output, and right is target

How to train Sketch2Drawings??

Big Overview Steps

  1. Download data from this kaggle dataset
  2. Prepare/Preprocess the data
  3. Train the model
  4. Test the model
  5. Export the model??

After Downloading Data

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

  1. Resize all images into 256 x 256
  2. Detect edges to get the "sketch"
  3. Combine input images and target images
  4. Split combined images into train and val set

Install dependencies

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

  1. Create virtual environment named sketch2drawings
conda create -n "sketch2drawings" python=3.6.0
  1. Activate conda environment
conda activate sketch2drawings
  1. Install OpenCV and Tensorflow v1.4.1 (since numpy is already installed)
conda install opencv-python
pip install tensorflow==1.4.1

Resizing Operation

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

Detect Edges Using Canny

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

Combine Operation

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.

Split Operation

Generate train/validation splits

python preprocessing/split.py --dir images/combined

Training

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

Testing

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

Export Mode

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.

Port model to Tensorflow.js

python3 export_checkpoint.py --checkpoint export --output_file s2d.pict

This will create a file called s2d.pict.

Other Notes

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

Packages

No packages published

Languages

  • Python 94.8%
  • Dockerfile 5.2%