Skip to content

line drawing colorization using TensorFlow

License

Notifications You must be signed in to change notification settings

rapidrabbit76/PaintsTensorFlow

Repository files navigation

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

References