Animation Paint Bucket Colorization

Yuekun Dai, Qinyue Li, Shangchen Zhou, Yihang Luo, Chongyi Li, Chen Change Loy

The code is released now. 🤗

arXiv | Dataset

BasicPBC

Colorizing line art is a pivotal task in the production of hand-drawn cel animation. In this work, we introduce a new learning-based inclusion matching pipeline, which directs the network to comprehend the inclusion relationships between segments. To facilitate the training of our network, we also propose a unique dataset PaintBucket-Character. This dataset includes rendered line arts alongside their colorized counterparts, featuring various 3D characters.

Update

• 2025.04.16: Upload inference code and ckpt for keyframe colorization: [Google Drive / Baidu Netdisk]
• 2024.11.12: Upload new dataset for keyframe colorization: [Google Drive / Baidu Netdisk]
• 2024.05.26: Add Colab demo of BasicPBC.
• 2024.04.25: Light-weighted model released.
• 2024.04.12: Support multiple ground-truth inference.
• 2024.04.08: Model inference updated. Support all resolutions and unclosed line art images.
• 2024.03.30: Checkpoint and training code of our BasicPBC are released.
• 2024.03.29: This repo is created.

TODO

• Add google colab inference script.
• Add trapped-ball segmentation module for unclosed line arts inference.
• Add a lightweight BasicPBC which can process images in 2K resolution without encountering Out-Of-Memory (OOM) error on 16GB RAM GPU.

Installation

1. Clone the repo

   git clone https://github.com/ykdai/BasicPBC.git
   

2. Install dependent packages

   cd BasicPBC
   pip install -r requirements.txt
   

3. Install BasicPBC
   Please run the following commands in the BasicPBC root path to install BasicPBC:

   python setup.py develop
   

Data Download

The details of our dataset can be found at this page. Dataset can be downloaded using the following links.

	Google Drive	Baidu Netdisk	Number	Description
PaintBucket-Character Train/Test	link	link	11,345/3,000	3D rendered frames for training and testing. Our dataset is a mere 3GB in size, so feel free to download it and enjoy exploring. 😆😆
PaintBucket-Real Test	/	/	200	Hand-drawn frames for testing.

Due to copyright issues, we do not provide download links for the real hand-drawn dataset. Please contact us through the e-mail if you want to use it or wish to get project files of our dataset. These hand-drawn frames are only for evaluation and not for any commercial activities.

Pretrained Model

You can download the pretrained checkpoints from the following links. Please place it under the ckpt folder and unzip it, then you can run the basicsr/test.py for inference.

	Google Drive	Baidu Netdisk
BasicPBC	link	link
BasicPBC-Light	link	link
BasicPBC-Ref	link	link

Model Inference

To estimate the colorized frames with our checkpoint trained on PaintBucket-Character, you can run the basicsr/test.py by using:

python basicsr/test.py -opt options/test/basicpbc_pbch_test_option.yml

Or you can test the lightweight model by:

python basicsr/test.py -opt options/test/basicpbc_light_test_option.yml

New: To test and evaluate the keyframe colorization model, please run:

python basicsr/test.py -opt options/test/basicpbc_ref_test_option.yml

The colorized results will be saved at results/.

To play with your own data, put your anime clip(s) under dataset/test/. The clip folder should contain at least one colorized gt frame and line of all frames.
We also provide two simple examples: laughing_girl and smoke_explosion. To play with your own data, put your anime clip(s) under dataset/test/. The clip folder should contain at least one colorized gt frame and line of all frames.
We also provide two simple examples: laughing_girl and smoke_explosion.

├── dataset 
    ├── test
        ├── laughing_girl
            ├── gt
                ├── 0000.png
            ├── line
                ├── 0000.png
                ├── 0001.png
                ├── ...
        ├── smoke_explosion
            ├── gt
            ├── line

To inference on laughing_girl, run inference_line_frames.py by using:

python inference_line_frames.py --path dataset/test/laughing_girl

Or run this to try with smoke_explosion:

python inference_line_frames.py --path dataset/test/smoke_explosion/  --mode nearest

Find results under results/.

inference_line_frames.py provides several arguments for different use cases.

• --mode can be either forward or nearest. By default, forward processes your frames sequentially. If set nearest, frames will be predicted from the nearest gt. e.g. Given gt 0000.png and 0005.png, line 0003.png will be colored according to 0004.png and 0004.png is colored according to 0005.png.

  python inference_line_frames.py --path dataset/test/smoke_explosion/  --mode nearest
  

• --seg_type is default if not specified. It's fast and simple, but not work if your line contains unclosed region. trappedball is robust to this case(acknowledge @hepesu/LineFiller). To decide which one to use, you can first set default together with --save_color_seg. It will produce colorized segmentation results. If you find out that some segments are not seperated properly, switch to trappedball:

  python inference_line_frames.py --path dataset/test/smoke_explosion/  --seg_type trappedball
  

• --skip_seg can help your skip the segmentation part. You can use it when seg already exists.
• --keep_line will generate another folder named [your_clip]_keepline which merges the original line in the line folder with the colorized output. This config is mainly for the line drawn not by the binary pen.
• --raft_res can change the resolution for the optical flow estimation (default is 320). We notice that sometimes the performance is bad due to he wrong optical flow estimation. Thus, if the performance is not satisfied on your case, you can change this to 640 to have a try by using --raft_res 640.
• --use_light_model will use the light-weighted model for inference. Add this if working on low memory GPU. Notice that this argument may produce poorer results than the base model.
• --multi_clip is used if you would like to inference on many clips at the same time. Put all clips within a single folder under dataset/test/, e.g.:

  ├── dataset 
      ├── test
          ├── your_clip_folder
              ├── clip01
              ├── clip02
              ├── ...
  

  In this case, run:

  python inference_line_frames.py --path dataset/test/your_clip_folder/  --multi_clip
  

Model Training

Training with single GPU

To train a model with your own data/model, you can edit the options/train/basicpbc_pbch_train_option.yml and run the following command. To train a model with your own data/model, you can edit the options/train/basicpbc_pbch_train_option.yml and run the following command.

python basicsr/train.py -opt options/train/basicpbc_pbch_train_option.yml

Training with multiple GPU

You can run the following command for multiple GPU training:

CUDA_VISIBLE_DEVICES=0,1 bash scripts/dist_train.sh 2 options/train/basicpbc_pbch_train_option.yml

BasicPBC structure

├── BasicPBC
    ├── assets
    ├── basicsr
        ├── archs
        ├── data
        ├── losses
        ├── metrics
        ├── models
        ├── ops
        ├── utils
    ├── dataset
    	├── train
	    	├── PaintBucket_Char
        ├── test
        	├── PaintBucket_Char
        	├── PaintBucket_Real
    ├── experiments
    ├── options
        ├── test
        ├── train
    ├── paint
    ├── raft
    ├── results
    ├── scripts

License

This project is licensed under S-Lab License 1.0. Redistribution and use of the dataset and code for non-commercial purposes should follow this license.

Citation

If you find this work useful, please cite:

@article{InclusionMatching2024,
  title     = {Learning Inclusion Matching for Animation Paint Bucket Colorization},
  author    = {Dai, Yuekun and Zhou, Shangchen and Li, Qinyue and Li, Chongyi and Loy, Chen Change},
  journal   = {CVPR},
  year      = {2024},
}

Contact

If you have any question, please feel free to reach me out at ydai005@e.ntu.edu.sg.