PyTorch 中的交互式深度彩色化

项目页面 | 论文 | 视频 | 演讲 | UI 代码

基于学习的深度先验的实时用户引导图像彩色化。
Richard Zhang*, Jun-Yan Zhu*, Phillip Isola, Xinyang Geng, Angela S. Lin, Tianhe Yu, 和 Alexei A. Efros。
发表于 ACM 图形学汇刊（SIGGRAPH 2017）。

这是我们由 Richard Zhang 和 Jun-Yan Zhu 编写的交互式图像彩色化的 PyTorch 重实现。

本仓库包含训练用法。原始的官方 GitHub 仓库（带有交互式 GUI，最初基于 Caffe 后端）位于 这里。官方仓库已更新以支持后端的 PyTorch 模型，这些模型可以在此仓库中进行训练。

前置条件

• Linux 或 macOS
• Python 2 或 3
• CPU 或 NVIDIA GPU + CUDA CuDNN

快速入门

安装

• 从 http://pytorch.org 安装 PyTorch 0.4+ 和 torchvision，以及其他依赖项（例如 visdom 和 dominate）。您可以通过以下命令安装所有依赖项：

pip install -r requirements.txt

• 克隆本仓库：

git clone https://github.com/richzhang/colorization-pytorch
cd colorization-pytorch

数据集准备

• 下载 ILSVRC 2012 数据集，并运行以下脚本准备数据： python make_ilsvrc_dataset.py --in_path /PATH/TO/ILSVRC12。这将为训练集创建符号链接，并将 ILSVRC 验证集划分为用于彩色化的验证集和测试集。

训练交互式彩色化

• 训练一个模型：bash ./scripts/train_siggraph.sh。这是一个两阶段的训练过程。首先，网络使用分类损失进行自动彩色化训练，结果保存在 ./checkpoints/siggraph_class 中。然后，网络使用回归损失进行交互式彩色化的微调，最终结果保存在 ./checkpoints/siggraph_reg2 中。

• 要查看训练结果和损失曲线，请运行 python -m visdom.server，并点击 http://localhost:8097。以下指标会被监控：

  • G_CE 是预测颜色分布与真实颜色之间的交叉熵损失。
  • G_entr 是预测分布的熵。
  • G_entr_hint 是在网络接收到颜色提示的位置上，预测分布的熵。
  • G_L1_max 是真实颜色与预测颜色分布的众数之间的 L1 距离。
  • G_L1_mean 是真实颜色与预测颜色分布的均值之间的 L1 距离。
  • G_L1_reg 是真实颜色与预测颜色之间的 L1 距离。
  • G_fake_real 是预测颜色与真实颜色之间的 L1 距离（在网络接收到提示的地方）。
  • G_fake_hint 是预测颜色与输入提示颜色之间的 L1 距离（在网络接收到提示的地方），它衡量了网络对输入提示的信任程度。
  • G_real_hint 是真实颜色与输入提示颜色之间的 L1 距离（在网络接收到提示的地方）。

测试交互式彩色化

• 获取一个模型。可以选择：

  • (1) 下载预训练模型，运行 bash pretrained_models/download_siggraph_model.sh，这将为您提供几个模型。
    • 原始 Caffe 权重 [推荐] ./checkpoints/siggraph_caffemodel/latest_net_G.pth 是原始 Caffe 模型权重转换为 PyTorch 格式后的版本，推荐使用。运行时请务必设置 --mask_cent 0。
    • 重新训练的模型：./checkpoints/siggraph_retrained/latest_net_G.pth。该模型具有更高的 PSNR，但在定性表现上有所不同。请注意，本仓库是对 SIGGRAPH 论文的近似重实现。
  • (2) 自行训练一个模型（如上文所述），训练完成后会在 ./checkpoints/siggraph_reg2/latest_net_G.pth 中生成一个模型。

• 在验证数据上测试模型：

  • 对于原始 Caffe 权重：python test.py --name siggraph_caffemodel --mask_cent 0。
  • 对于重新训练的权重：python test.py --name siggraph_retrained。
  • 如果您自己训练了模型：python test.py --name siggraph_reg2。 测试结果将保存到 ./results/[[NAME]]/latest_val/index.html 的 HTML 文件中。对于验证集中的每一张图片，它将测试 (1) 自动彩色化，(2) 使用少量随机提示的交互式彩色化，以及 (3) 使用大量随机提示的交互式彩色化。

• 通过绘制 PSNR 与提示数量的关系图来测试模型：python test_sweep.py --name [[NAME]]。该图曾用于 论文 的图 6。此测试会随机向网络展示 6x6 的彩色提示区域，以观察其彩色化结果与真实值的接近程度。

  

• 使用原始官方 仓库 以交互方式测试模型。按照该仓库的安装说明操作，并运行 python ideepcolor.py --backend pytorch --color_model [[PTH/TO/MODEL]] --dist_model [[PTH/TO/MODEL]]。

引用

如果您在研究中使用此代码，请引用我们的论文：

@article{zhang2017real,
  title={Real-Time User-Guided Image Colorization with Learned Deep Priors},
  author={Zhang, Richard and Zhu, Jun-Yan and Isola, Phillip and Geng, Xinyang and Lin, Angela S and Yu, Tianhe and Efros, Alexei A},
  journal={ACM Transactions on Graphics (TOG)},
  volume={9},
  number={4},
  year={2017},
  publisher={ACM}
}

致谢

此代码大量借鉴了 pytorch-CycleGAN 仓库。

colorization-pytorch 快速上手指南

本项目是 SIGGRAPH 2017 论文《Real-Time User-Guided Image Colorization with Learned Deep Priors》的 PyTorch 复现版本，支持交互式图像上色。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux 或 macOS
• Python 版本: Python 2 或 Python 3
• 硬件加速: CPU 或 NVIDIA GPU (需安装 CUDA 和 CuDNN)
• 核心依赖: PyTorch 0.4+ 及 torchvision

国内开发者提示：建议配置国内镜像源以加速依赖下载。

• PyTorch 国内镜像安装参考：https://pytorch.org/get-started/locally/ (选择对应版本后查看 "Run this Command" 中的 pip 源替换方案，通常使用 pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ...)

安装步骤

1. 安装依赖库

克隆仓库后，使用 pip 一键安装所需依赖（包括 visdom 和 dominate 等）：

git clone https://github.com/richzhang/colorization-pytorch
cd colorization-pytorch
pip install -r requirements.txt

(注：若网络较慢，可添加清华源参数：pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple)

2. 准备数据集 (可选)

如需从头训练模型，需下载 ILSVRC 2012 数据集并运行脚本进行预处理（创建软链接并划分验证/测试集）：

python make_ilsvrc_dataset.py --in_path /PATH/TO/ILSVRC12

基本使用

对于大多数用户，直接使用预训练模型进行测试是最快的上手方式。

1. 下载预训练模型

运行官方脚本下载模型权重。推荐使用转换后的原始 Caffe 权重 (siggraph_caffemodel)，效果最接近原论文：

pretrained_models/download_siggraph_model.sh

2. 执行测试

使用下载的模型对验证集数据进行上色测试。根据下载的模型类型选择对应的命令：

方案 A：使用推荐的原始 Caffe 转换权重

python test.py --name siggraph_caffemodel --mask_cent 0

方案 B：使用本项目重新训练的权重

python test.py --name siggraph_retrained

3. 查看结果

测试完成后，结果将保存为 HTML 文件。请在浏览器中打开以下路径查看自动上色及交互式上色的对比效果：

./results/[[NAME]]/latest_val/index.html

(将 [[NAME]] 替换为你使用的模型名称，例如 siggraph_caffemodel)

进阶提示：若需体验图形化界面（GUI）进行实时交互式上色，请前往官方交互代码库 interactive-deep-colorization，并指定后端为 pytorch 及上述下载的模型路径运行 ideepcolor.py。