AutoColor

简介

该项目用于漫画的自动上色。项目主要基于 clip, MAE 和 timm.

使用说明

方法1、直接下载打包好的软件（适用win10系统)

下载地址: https://github.com/danczs/AutoColor/releases/download/v0.1.0/AutoColor_win10.zip

（通过onnxruntime部署，通过pyinstaller打包的python程序）

方法2、通过执行python程序使用

a. clone 或下载项目代码到本地（需要pytorch、onnxruntime等环境）:

git clone https://github.com/danczs/AutoColor.git

b. 安装相关包:

pip install -r requirements.txt
pip install git+https://github.com/openai/CLIP.git

c. 分别下载 pytorch模型和onnx模型到文件夹 ./deployment/pytorch_models和./deployment/onnx_models

d. 启动GUI界面程序

cd deployment
python auto_color_gui.py

其中可以通过 auto_color_gui.py line14中的代码选择使用ONNX模型还是pytorch模型

# using pytorch or onnx deployment
from autocolor_pytorch_deployment import AutoColorDeployment # pytorch deployment
#from autocolor_onnx_deployment import AutoColorDeployment # onnx deployment

模型训练

适用于模型的修改与重新训练

数据准备

需要下载Danbooru数据集的 kaggle子集，约40G。数据集目录：

/path/to/archive/
  danbooru-images/
    0000
      1.jpg
      2.jpg
      ...
    0001
    ...
    0149
  danbooru-metadata/
  moeimouto-faces/

我们只需使用danbooru-images中的前10个文件夹（0000-0009，共约2万张图片）进行训练，将该10个文件夹单独拷贝到一个新的文件夹（e.g.,archive_subset10)

模型准备

下载相应的预训练模型到./models文件夹，以便后续的特征提取和模型训练

特征准备

a. 提取clip特征并保存

python get_clip_features.py --data_path = /path/to/archive_subset10

提取的 clip 特征会被保存到 ./features/clip_features_subset.npy

b. 提取mae特征并保存

python get_mae_features.py --data_path = /path/to/archive_subset10 --output_path=/path/to/mae_features

mae特征会被保存到--output_path文件夹，每张图片特征被保存为一个.npy文件，一共约12G。这些特征文件的索引被保存在./feature/mae_feature_names.txt

训练模型

训练 color deocder模型

python auto_color_main.py --grad_state='010'
                          --output_dir=/path/to/output
                          --mae_model_path=./models/mae_visualize_vit_base.pth \
                          --mae_feature_path=./features/mae_feature_names.txt \
                          --clip_feature_path=./features/clip_features_subset.npy \
                          --colordecoder_model_path=./models/color_decoder_pretrained.pth

其中--colordecoder_model_path可以不设置，不设置时使用mae_visualize_vit_base.pth的decoder进行初始化。这里的./models/color_decoder_pretrained.pth是在数据集0000-0049文件（约10万张图片）上的预训练模型，会略微提升最终性能。输出模型保存在--output_dir。

训练 super color 模型

super color 模型可以单独进行训练（较快），也可以结合训练好的 color decoder 一起训练，二者整体效果差别不大。

a. 单独训练

其输入输出为：低分辨率彩图 + 高分辨率灰度图 --> 高分辨率彩图

python auto_color_main.py --grad_state='001'  \
                          --mae_feature_path=./features/mae_feature_names.txt \
                          --clip_feature_path=./features/clip_features_subset.npy \
                          --supercolor_only

b. 基于训练好的color decoder的输出进行训练

其输入输出为： color_decoder输出的低分辨率彩图 + 高分辨率灰度图 --> 高分辨率彩图

python auto_color_main.py --grad_state='001' \
                          --batch_size=32  \
                          --mae_feature_path=./features/mae_feature_names.txt \
                          --clip_feature_path=./features/clip_features_subset.npy \
                          --colordecoder_model_path=/path/to/trained_colordeocder_model.pth

模型部署

pytorch部署

a. 将训练好的color decoder模型和super color模型拷贝到 deployment/pytorch_models中分别命名为color_decoder.pth 和super_color.pth

b. 确保auto_color_gui.py中line14使用的是pytorch deployment模块

# using pytorch or onnx deployment
from autocolor_pytorch_deployment import AutoColorDeployment # pytorch deployment
#from autocolor_onnx_deployment import AutoColorDeployment # onnx deployment

c. 启动界面

cd deployment
python auto_color_gui.py

onnx部署

a. 使用conert_to_onnx_models.py将 pytorch 模型转换为 onnx 模型，需转换的color_decoder.pth模型和super_color.pth模型需放在deployment\pytorch_models 文件夹下：

python conert_to_onnx_models.py

b. 确保auto_color_gui.py中line14使用的是onnx deployment模块

# using pytorch or onnx deployment
#from autocolor_pytorch_deployment import AutoColorDeployment # pytorch deployment
from autocolor_onnx_deployment import AutoColorDeployment # onnx deployment