PaintsTensorFlow

Model Structure

Model Structure

DraftModel

Colorization

Results

input(line) - input(hint) - draft - output - ground truth

Gray background in hint for visualization.

More result image

GUI

Current we does not provide a GUI.

instead, we are preparing a web service.

if you want to use it locally using the GUI, refer to tag:0.0.1

GUI

File - open( select Image )

Click "Liner" to create line art

Click "RUN" to automatically color

Click "Color" to select a color and then draw a color hint

Click "RUN" to automatically color

Requirements

• tensorflow==2.7.1
• tensorflow-addons==0.16.1
• albumentations==1.1.0
• opencv-python-headless==4.5.5.62
• scipy==1.8.0
• tqdm==4.61.2
• wandb==0.12.11

Pretrained Model

• draft saved_model link
• colorization saved_model link
• draft model onnx link
• colorization model onnx link

Training

• My dataset over 700,000 images and created a lines, using SketchKeras

• dataset path structure (image-line file name must be matched)

  {DATASET PATH}
  │
  ├─ train
  │   ├─ image
  │   │    └─ 1.jpg, 2.jpg, N.jpg
  │   ├─ line
  │   │    └─ 1.jpg, 2.jpg, N.jpg
  └─ test
      ├─ image
      │    └─ 1.jpg, 2.jpg, N.jpg
      └─ line
          └─ 1.jpg, 2.jpg, N.jpg
  

• [step 1]: Training draft model 128X128 size

  python3 main.py \
  --mode="draft" \
  --root_dir={"YOURE DATASET PATH"} \
  --batch_size={BATCH_SIZE} \
  --epochs={EPOCHS} \
  --lr=0.0001
  

• [step 2]: Training Colorization model 512X512 size

  python3 main.py \
  --mode="colorization" \
  --root_dir={"YOURE DATASET PATH"} \
  --batch_size={BATCH_SIZE} \
  --epochs={EPOCHS} \
  --lr=0.0001
  

Convert to ONNX

• check "convert2onnx.ipynb"

References

• PaintsChainer
• SketchKeras
• pix2pix