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"