🎨 DDColor

ICCV 2023 论文“DDColor：通过双解码器实现照片级逼真图像着色”的官方 PyTorch 实现。

康晓阳、杨涛、欧阳文琪、任培然、李凌志、谢宣松
达摩院，阿里巴巴集团

🪄 DDColor 可为历史黑白老照片提供生动自然的着色效果。

🎲 它甚至可以为动漫游戏中的风景进行着色或重新着色，将你的动画场景转变为逼真的现实风格！（图片来源：原神）

新闻

• [2024-01-28] 支持通过 🤗 Hugging Face 进行推理！感谢 @Niels 的建议和示例代码，以及 @Skwara 的错误修复。
• [2024-01-18] 添加 Replicate 演示和 API！感谢 @Chenxi。
• [2023-12-13] 发布 DDColor-tiny 预训练模型！
• [2023-09-07] 增加模型库并发布三款预训练模型！
• [2023-05-15] 代码发布，支持训练和推理！
• [2023-05-05] 在线演示已上线！

在线演示

请在 ModelScope 和 Replicate 上体验我们的在线演示。

方法

简而言之: DDColor 利用多尺度视觉特征来优化可学习的颜色标记（即颜色查询），并在自动图像着色任务上达到了最先进的性能。

安装

要求

• Python >= 3.7
• PyTorch >= 1.7

使用 conda 安装（推荐）

conda create -n ddcolor python=3.9
conda activate ddcolor
pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu118

pip install -r requirements.txt

# 对于训练，安装以下额外依赖和 basicsr
pip install -r requirements.train.txt
python3 setup.py develop

快速入门

使用本地脚本进行推理（无需 basicsr）

1. 下载预训练模型：

from modelscope.hub.snapshot_download import snapshot_download

model_dir = snapshot_download('damo/cv_ddcolor_image-colorization', cache_dir='./modelscope')
print('模型资源已保存到 %s' % model_dir)

2. 运行推理：

python scripts/infer.py --model_path ./modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt --input ./assets/test_images

或者

sh scripts/inference.sh

使用 Hugging Face 进行推理

通过 Hugging Face Hub 加载模型：

from huggingface_hub import PyTorchModelHubMixin
from ddcolor import DDColor

class DDColorHF(DDColor, PyTorchModelHubMixin):
    def __init__(self, config=None, **kwargs):
        if isinstance(config, dict):
            kwargs = {**config, **kwargs}
        super().__init__(**kwargs)

ddcolor_paper_tiny = DDColorHF.from_pretrained("piddnad/ddcolor_paper_tiny")
ddcolor_paper      = DDColorHF.from_pretrained("piddnad/ddcolor_paper")
ddcolor_modelscope = DDColorHF.from_pretrained("piddnad/ddcolor_modelscope")
ddcolor_artistic   = DDColorHF.from_pretrained("piddnad/ddcolor_artistic")

或者直接运行以下命令进行模型推理：

python scripts/infer.py --model_name ddcolor_modelscope --input ./assets/test_images
# model_name: [ddcolor_paper | ddcolor_modelscope | ddcolor_artistic | ddcolor_paper_tiny]

使用 ModelScope 进行推理

1. 安装 modelscope：

pip install modelscope

2. 运行推理：

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')
result = img_colorization('https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg')
cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])

这段代码会自动下载 ddcolor_modelscope 模型（参见 ModelZoo）并执行推理。模型文件 pytorch_model.pt 可以在本地路径 ~/.cache/modelscope/hub/damo 中找到。

Gradio 演示

安装 gradio 和其他所需库：

pip install gradio gradio_imageslider

然后，你可以通过以下命令运行演示：

python demo/gradio_app.py

模型库

我们提供了多个不同版本的预训练模型，请查看 Model Zoo。

训练

1. 数据集准备：下载 ImageNet 数据集或创建自定义数据集。使用此脚本获取数据集列表文件：

python scripts/get_meta_file.py

2. 下载 ConvNeXt 和 InceptionV3 的预训练权重，并将其放入 pretrain 文件夹。

3. 在 options/train/train_ddcolor.yml 中指定 meta_info_file 等选项。

4. 开始训练：

sh scripts/train.sh

ONNX 导出

支持导出 ONNX 模型。

1. 安装依赖：

pip install onnx==1.16.1 onnxruntime==1.19.2 onnxsim==0.4.36

2. 使用示例：

python scripts/export_onnx.py --model_path pretrain/ddcolor_paper_tiny.pth --export_path weights/ddcolor-tiny.onnx

使用 ddcolor_paper_tiny 模型的 ONNX 导出示例可在 这里 查看。

引用

如果我们的工作对您的研究有所帮助，请考虑引用：

@inproceedings{kang2023ddcolor,
  title={DDColor: Towards Photo-Realistic Image Colorization via Dual Decoders},
  author={Kang, Xiaoyang and Yang, Tao and Ouyang, Wenqi and Ren, Peiran and Li, Lingzhi and Xie, Xuansong},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={328--338},
  year={2023}
}

致谢

我们感谢 BasicSR 的作者们提供的出色训练流程。

Xintao Wang, Ke Yu, Kelvin C.K. Chan, Chao Dong 和 Chen Change Loy. BasicSR：开源图像与视频修复工具箱。https://github.com/xinntao/BasicSR，2020 年。

部分代码改编自 ColorFormer、BigColor、ConvNeXt、Mask2Former 以及 DETR。感谢这些项目团队的杰出工作！

DDColor 快速上手指南

DDColor 是一款基于双解码器架构的 AI 图像着色工具，能够为黑白老照片提供逼真自然的色彩还原，同时也支持将动漫场景转化为写实风格。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / Windows / macOS
• Python 版本：>= 3.7 (推荐 3.9)
• PyTorch 版本：>= 1.7
• 硬件建议：推荐使用 NVIDIA GPU 以获得最佳推理速度（需安装对应的 CUDA 版本）

2. 安装步骤

推荐使用 conda 创建独立虚拟环境进行安装。

第一步：创建并激活环境

conda create -n ddcolor python=3.9
conda activate ddcolor

第二步：安装 PyTorch

根据您的 CUDA 版本安装对应的 PyTorch。以下为 CUDA 11.8 的安装命令（国内用户若下载缓慢，可尝试使用清华源或官方源）：

pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu118

第三步：安装项目依赖

克隆项目代码后，进入目录安装依赖包：

pip install -r requirements.txt

注意：如果您仅需进行推理（Inference），上述步骤已完成。若需要训练模型，请额外执行：

pip install -r requirements.train.txt
python3 setup.py develop

3. 基本使用

以下提供两种最常用的推理方式：通过本地脚本运行（适合批量处理）和通过 ModelScope 管道运行（适合快速验证）。

方式一：本地脚本推理（推荐）

此方式无需安装 basicsr，适合大多数用户。

1. 下载预训练模型 使用 ModelScope 自动下载模型资源（国内访问速度快）：

from modelscope.hub.snapshot_download import snapshot_download

model_dir = snapshot_download('damo/cv_ddcolor_image-colorization', cache_dir='./modelscope')
print('model assets saved to %s' % model_dir)

2. 执行着色命令 假设您的测试图片位于 ./assets/test_images 文件夹下：

python scripts/infer.py --model_path ./modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt --input ./assets/test_images

或者直接运行封装好的脚本：

sh scripts/inference.sh

方式二：使用 ModelScope 管道（代码调用）

如果您希望通过 Python 代码直接调用流水线，可以使用以下方式。该方式会自动下载模型并处理单张图片。

1. 安装 ModelScope

pip install modelscope

2. 运行推理代码

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')

# 执行推理（支持本地路径或网络 URL）
result = img_colorization('https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg')

# 保存结果
cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])

可选：启动 Gradio 网页演示

如果您想在本机启动一个可视化的 Web 界面进行操作：

pip install gradio gradio_imageslider
python demo/gradio_app.py

启动后在浏览器访问显示的本地地址即可上传图片体验。