pix2pix

项目 | Arxiv | PyTorch

用于学习从输入图像到输出图像映射的 Torch 实现，例如：

条件对抗网络下的图像到图像转换
Phillip Isola、Jun-Yan Zhu、Tinghui Zhou、Alexei A. Efros
CVPR, 2017。

在某些任务上，使用较小的数据集即可较快地获得不错的效果。例如，要学习生成建筑立面（如上例所示），我们仅用400张图片训练了约2小时（在单个 Pascal Titan X GPU 上）。然而，对于更难的问题，可能需要使用更大的数据集，并且训练时间长达数小时甚至数天。

注：请查看我们的 PyTorch 实现，其中包含了 pix2pix 和 CycleGAN。PyTorch 版本目前仍在积极开发中，其效果可与或优于本 Torch 版本。

设置

前提条件

• Linux 或 OSX
• NVIDIA GPU + CUDA CuDNN（CPU 模式及无 CuDNN 的 CUDA 环境可能稍作修改后也能运行，但未经测试）

开始使用

• 从 https://github.com/torch/distro 安装 torch 及其依赖项
• 安装 torch 包 nngraph 和 display

luarocks install nngraph
luarocks install https://raw.githubusercontent.com/szym/display/master/display-scm-0.rockspec

• 克隆此仓库：

git clone git@github.com:phillipi/pix2pix.git
cd pix2pix

• 下载数据集（例如 CMP Facades）：

bash ./datasets/download_dataset.sh facades

• 训练模型

DATA_ROOT=./datasets/facades name=facades_generation which_direction=BtoA th train.lua

• （仅限 CPU）不使用 GPU 或 CuDNN 的相同训练命令。设置环境变量 gpu=0 cudnn=0 强制仅使用 CPU

DATA_ROOT=./datasets/facades name=facades_generation which_direction=BtoA gpu=0 cudnn=0 batchSize=10 save_epoch_freq=5 th train.lua

• （可选）启动显示服务器以在模型训练时查看结果。（有关详细信息，请参阅 Display UI）：

th -ldisplay.start 8000 0.0.0.0

• 最后，测试模型：

DATA_ROOT=./datasets/facades name=facades_generation which_direction=BtoA phase=val th test.lua

测试结果将保存到此处的 HTML 文件中：./results/facades_generation/latest_net_G_val/index.html。

训练

DATA_ROOT=/path/to/data/ name=expt_name which_direction=AtoB th train.lua

将 AtoB 改为 BtoA 即可反向训练。

模型将保存到 ./checkpoints/expt_name 目录下（可通过在 train.lua 中传递 checkpoint_dir=your_dir 来更改）。

更多训练选项请参阅 train.lua 中的 opt 部分。

测试

DATA_ROOT=/path/to/data/ name=expt_name which_direction=AtoB phase=val th test.lua

这将在 /path/to/data/val 中的所有图像上运行名为 expt_name 的模型，方向为 AtoB。

结果图像以及用于查看它们的网页将保存到 ./results/expt_name 目录下（可通过在 test.lua 中传递 results_dir=your_dir 来更改）。

更多测试选项请参阅 test.lua 中的 opt 部分。

数据集

使用以下脚本下载数据集。部分数据集由其他研究人员收集，请在使用这些数据时引用他们的论文。

bash ./datasets/download_dataset.sh dataset_name

• facades：来自 CMP Facades 数据集 的 400 张图片。[引用]
• cityscapes：来自 Cityscapes 训练集 的 2975 张图片。[引用]
• maps：从 Google 地图抓取的 1096 张训练图片
• edges2shoes：来自 UT Zappos50K 数据集 的 5 万张训练图片。边缘由 HED 边缘检测器结合后处理计算得出。 [引用]
• edges2handbags：来自 iGAN 项目 的 13.7 万张亚马逊手袋图片。边缘同样由 HED 边缘检测器结合后处理计算得出。[引用]
• night2day：约 2 万张来自 Transient Attributes 数据集 的自然场景图片 [引用]。若要训练 day2night 的 pix2pix 模型，需添加 which_direction=BtoA。

模型

使用以下脚本下载预训练模型。下载完成后，需重命名模型（例如将 facades_label2image 重命名为 /checkpoints/facades/latest_net_G.t7）。

bash ./models/download_model.sh model_name

• facades_label2image（标签 → 建筑立面）：基于 CMP Facades 数据集训练
• cityscapes_label2image（标签 → 街景）：基于 Cityscapes 数据集训练
• cityscapes_image2label（街景 → 标签）：基于 Cityscapes 数据集训练
• edges2shoes（边缘 → 照片）：基于 UT Zappos50K 数据集训练
• edges2handbags（边缘 → 照片）：基于亚马逊手袋图片训练
• day2night（白天场景 → 夜间场景）：基于约 100 个 网络摄像头 训练

准备训练和测试数据

生成配对数据

我们提供了一个 Python 脚本，用于生成成对的图像数据 {A,B}，其中 A 和 B 是同一场景的不同表现形式。例如，这些可能是 {标签图，照片} 或 {黑白图像，彩色图像} 的配对。然后我们可以学习将 A 转换为 B，或 B 转换为 A：

创建文件夹 /path/to/data，并在其中建立子文件夹 A 和 B。A 和 B 应分别包含 train、val、test 等子文件夹。在 /path/to/data/A/train 中放入 A 风格的训练图片，在 /path/to/data/B/train 中放入对应的 B 风格图片。对其他数据划分（val、test 等）重复此操作。

配对中的两张图片 A 和 B 必须具有相同的尺寸和文件名，例如 /path/to/data/A/train/1.jpg 被认为与 /path/to/data/B/train/1.jpg 对应。

数据按上述方式整理好后，运行：

python scripts/combine_A_and_B.py --fold_A /path/to/data/A --fold_B /path/to/data/B --fold_AB /path/to/data

这会将每对图像 (A,B) 合并为一个单独的图像文件，以便进行训练。

关于色彩化的一些说明

进行色彩化时无需运行 combine_A_and_B.py。相反，您需要准备一些自然图像，并在脚本中设置 preprocess=colorization。程序会自动将每张 RGB 图像转换为 Lab 色彩空间，并在训练过程中创建 L → ab 的图像对。同时请设置 input_nc=1 和 output_nc=2。

提取边缘

我们提供了 Python 和 Matlab 脚本，用于从照片中提取粗略的边缘。运行 scripts/edges/batch_hed.py 来计算 HED 边缘。运行 scripts/edges/PostprocessHED.m 以通过额外的后处理步骤简化边缘。更多详情请参阅代码文档。

在 Cityscapes 数据集上评估 Labels2Photos 任务

我们提供了在 Cityscapes 验证 集上运行 Labels2Photos 任务评估的脚本。我们假设您已经在系统中安装了 caffe（以及 pycaffe）。如果没有，请参阅 官方网站 获取安装说明。成功安装 caffe 后，通过运行以下命令下载预训练的 FCN-8s 语义分割模型（512MB）：

bash ./scripts/eval_cityscapes/download_fcn8s.sh

然后确保 ./scripts/eval_cityscapes/ 在您的系统 Python 路径中。如果不在，请运行以下命令将其添加：

export PYTHONPATH=${PYTHONPATH}:./scripts/eval_cityscapes/

现在您可以运行以下命令来评估您的预测结果：

python ./scripts/eval_cityscapes/evaluate.py --cityscapes_dir /path/to/original/cityscapes/dataset/ --result_dir /path/to/your/predictions/ --output_dir /path/to/output/directory/

存储在 --result_dir 下的图像应包含您模型在 Cityscapes 验证 分割上的预测结果，并且遵循原始 Cityscapes 的命名规范（例如：frankfurt_000001_038418_leftImg8bit.png）。该脚本将在 --output_dir 下输出一个包含评估指标的文本文件。

进一步说明：我们的预训练 FCN 模型 不 应用于原始分辨率（1024×2048）的 Cityscapes 数据集，因为它是在 256×256 的图像上训练的，这些图像在训练过程中被上采样到 1024×2048。训练时进行尺寸调整的目的在于：1) 保持标签图的原始高分辨率不变；2) 避免为 Cityscapes 数据集修改标准的 FCN 训练代码和网络架构。在测试时，您需要生成 256×256 的结果。我们的测试代码会自动将您的结果上采样到 1024×2048，然后再输入到预训练的 FCN 模型中。输出结果为 1024×2048 分辨率，并与 1024×2048 的真实标签进行比较。您无需对真实标签进行任何尺寸调整。验证一切是否正确的一个好方法是先复现论文中真实图像的评估结果。为此，您需要将原始/真实的 Cityscapes 图像（不是标签）调整为 256×256，然后将其输入到评估代码中。

显示界面

可选地，在训练和测试过程中显示图像时，可以使用 display 包。

• 安装方法：luarocks install https://raw.githubusercontent.com/szym/display/master/display-scm-0.rockspec
• 启动服务器：th -ldisplay.start
• 在浏览器中打开此网址：http://localhost:8000

默认情况下，服务器仅监听本地回环地址。若要允许外部连接到任意网络接口，可传递 0.0.0.0 参数：

th -ldisplay.start 8000 0.0.0.0

然后在浏览器中打开 http://(hostname):(port)/ 即可加载远程桌面。

默认情况下，L1 误差会被绘制到显示界面上。您可以设置环境变量 display_plot 为逗号分隔的值列表 errL1, errG 和 errD，分别可视化 L1 误差、生成器误差和判别器误差。例如，如果您只想在显示界面上绘制生成器和判别器的误差，而不是默认的 L1 误差，可以设置 display_plot="errG,errD"。

引用

如果您在研究中使用了本代码，请引用我们的论文 基于条件对抗网络的图像到图像转换：

@article{pix2pix2017,
  title={Image-to-Image Translation with Conditional Adversarial Networks},
  author={Isola, Phillip and Zhu, Jun-Yan and Zhou, Tinghui and Efros, Alexei A},
  journal={CVPR},
  year={2017}
}

猫论文合集

如果您喜欢猫，并且热爱阅读关于图形学、视觉和机器学习的优秀论文，请查看猫论文合集： [Github] [网页]

致谢

本代码大量借鉴了 DCGAN。数据加载器则基于 DCGAN 和 Context-Encoder 进行了修改。

pix2pix 快速上手指南

pix2pix 是一个基于条件生成对抗网络（cGAN）的图像到图像翻译工具，可用于将输入图像映射为输出图像（例如：素描转照片、标签图转街景、白天转黑夜等）。

注意：官方推荐使用 PyTorch 版本，该版本维护更活跃且效果相当或更好。本指南基于原始的 Torch (Lua) 实现。

1. 环境准备

系统要求

• 操作系统：Linux 或 macOS
• 硬件：NVIDIA GPU + CUDA + CuDNN
  • 注：CPU 模式也可运行（需设置参数），但未充分测试；无 CuDNN 的 CUDA 环境可能需要少量修改。

前置依赖

• Torch 框架：需安装 Lua 版本的 Torch。
• LuaRocks 包管理器：用于安装额外依赖。

2. 安装步骤

第一步：安装 Torch 核心

访问 Torch 官网 或 GitHub 仓库安装基础环境：

git clone https://github.com/torch/distro.git ~/torch --recursive
cd ~/torch; bash install.sh  # 根据提示操作，完成后执行 source ~/.bashrc 或 source ~/.zshrc

第二步：安装必要包

安装 nngraph 和 display 模块：

luarocks install nngraph
luarocks install https://raw.githubusercontent.com/szym/display/master/display-scm-0.rockspec

第三步：克隆项目代码

git clone git@github.com:phillipi/pix2pix.git
cd pix2pix

第四步：下载数据集

使用内置脚本下载示例数据集（以建筑立面数据集 facades 为例）：

bash ./datasets/download_dataset.sh facades

国内用户若下载缓慢，可手动从 CMP Facades 数据集 下载后放入 datasets 目录。

3. 基本使用

以下流程演示如何训练一个将“标签图”转换为“建筑照片”的模型（方向：BtoA）。

训练模型

运行以下命令开始训练（默认使用 GPU）：

DATA_ROOT=./datasets/facades name=facades_generation which_direction=BtoA th train.lua

可选：仅使用 CPU 训练 若无可用 GPU，添加环境变量强制使用 CPU：

DATA_ROOT=./datasets/facades name=facades_generation which_direction=BtoA gpu=0 cudnn=0 batchSize=10 save_epoch_freq=5 th train.lua

可选：可视化训练过程 在新终端窗口启动显示服务器，浏览器访问 http://localhost:8000 查看实时生成结果：

th -ldisplay.start 8000 0.0.0.0

测试模型

训练完成后，使用以下命令在验证集上测试模型：

DATA_ROOT=./datasets/facades name=facades_generation which_direction=BtoA phase=val th test.lua

测试结果将保存为 HTML 文件，路径如下： ./results/facades_generation/latest_net_G_val/index.html 直接在浏览器打开该文件即可查看输入与生成的对比图。

---

自定义数据提示： 若要使用自己的数据集，请准备成对图像（A 和 B 风格），分别放入 data/A 和 data/B 文件夹，确保文件名一致且尺寸相同，然后运行：

python scripts/combine_A_and_B.py --fold_A /path/to/data/A --fold_B /path/to/data/B --fold_AB /path/to/data

之后将 DATA_ROOT 指向合并后的数据目录即可。