DDColor 图像上色模型

该模型为黑白图像上色模型，输入一张黑白图像，实现端到端的全图上色，返回上色处理后的彩色图像。

论文 ｜ Github主页

模型描述

DDColor 是最新的 SOTA 图像上色算法，能够对输入的黑白图像生成自然生动的彩色结果。

算法整体流程如下图，使用 UNet 结构的骨干网络和图像解码器分别实现图像特征提取和特征图上采样，并利用 Transformer 结构的颜色解码器完成基于视觉语义的颜色查询，最终聚合输出彩色通道预测结果。

模型期望使用方式和适用范围

该模型适用于多种格式的图像输入，给定黑白图像，生成上色后的彩色图像；给定彩色图像，将自动提取灰度通道作为输入，生成重上色的图像。

如何使用

在 ModelScope 框架上，提供输入图片，即可以通过简单的 Pipeline 调用来使用图像上色模型。

代码范例

import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

img_colorization = pipeline(Tasks.image_colorization, 
                       model='damo/cv_ddcolor_image-colorization')
img_path = 'https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg'
result = img_colorization(img_path)
cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])

模型局限性以及可能的偏差

• 本算法模型使用自然图像数据集进行训练，对于分布外场景（例如漫画等）可能产生不恰当的上色结果；
• 对于低分辨率或包含明显噪声的图像，算法可能无法得到理想的生成效果。

训练数据介绍

模型使用公开数据集 ImageNet 训练，其训练集包含 128 万张自然图像。

模型训练流程

预处理

将 ImageNet 数据集的 RGB 彩色图像转为 Lab 色彩空间，并提取其中的 L 通道，得到相应的灰度图片。

模型训练代码

import os
import tempfile

from modelscope.hub.snapshot_download import snapshot_download
from modelscope.msdatasets import MsDataset
from modelscope.msdatasets.dataset_cls.custom_datasets.image_colorization import \
    ImageColorizationDataset
from modelscope.trainers import build_trainer
from modelscope.utils.config import Config
from modelscope.utils.constant import DownloadMode, ModelFile


tmp_dir = tempfile.TemporaryDirectory().name
if not os.path.exists(tmp_dir):
    os.makedirs(tmp_dir)
model_id = 'damo/cv_ddcolor_image-colorization'
cache_path = snapshot_download(model_id)
config = Config.from_file(
    os.path.join(cache_path, ModelFile.CONFIGURATION))

dataset_train = MsDataset.load(
    'imagenet-val5k-image',
    namespace='damo',
    subset_name='default',
    split='validation',
    download_mode=DownloadMode.REUSE_DATASET_IF_EXISTS)._hf_ds
dataset_val = MsDataset.load(
    'imagenet-val5k-image',
    namespace='damo',
    subset_name='default',
    split='validation',
    download_mode=DownloadMode.REUSE_DATASET_IF_EXISTS)._hf_ds

dataset_train = ImageColorizationDataset(
    dataset_train, is_train=True)
dataset_val = ImageColorizationDataset(
    dataset_val, is_train=False)

kwargs = dict(
    model=model_id,
    train_dataset=dataset_train,
    eval_dataset=dataset_val,
    work_dir=tmp_dir)
trainer = build_trainer(default_args=kwargs)
trainer.train()

数据评估及结果

本算法主要在 ImageNet 和 COCO-Stuff上测试。

引用

如果你觉得这个模型对你有所帮助，请考虑引用下面的相关论文：

@article{kang2022ddcolor,
  title={DDColor: Towards Photo-Realistic and Semantic-Aware Image Colorization via Dual Decoders},
  author={Kang, Xiaoyang and Yang, Tao and Ouyang, Wenqi and Ren, Peiran and Li, Lingzhi and Xie, Xuansong},
  journal={arXiv preprint arXiv:2212.11613},
  year={2022}
}