EdgeConnect：基于对抗性边缘学习的生成式图像修复

ArXiv | BibTex

引言：

我们提出了一种新的图像修复方法，该方法能够更好地复原包含精细细节的填充区域。这一方法受到艺术家创作方式的启发——“先勾勒线条，再上色”。我们设计了一个两阶段的对抗模型EdgeConnect，它由一个边缘生成器和一个图像补全网络组成。边缘生成器会为图像中缺失的部分（包括规则和不规则的边缘）“幻化”出边缘信息，而图像补全网络则以这些“幻化”的边缘作为先验信息来填充缺失区域。系统的详细描述请参阅我们的论文此处。

(a) 包含缺失区域的输入图像。缺失区域以白色显示。(b) 计算得到的边缘掩码。黑色线条是使用Canny边缘检测器计算得出的现有区域边缘；蓝色线条则是由边缘生成器“幻化”出来的边缘。(c) 提出方法的图像修复结果。

前置条件

• Python 3
• PyTorch 1.0
• NVIDIA GPU + CUDA cuDNN

安装

• 克隆本仓库：

git clone https://github.com/knazeri/edge-connect.git
cd edge-connect

• 从http://pytorch.org安装PyTorch及其依赖项。
• 安装Python所需的依赖包：

pip install -r requirements.txt

数据集

1) 图像

我们使用了Places2、CelebA以及巴黎街景数据集。若要在完整数据集上训练模型，请从官方网站下载相应数据集。

下载完成后，运行scripts/flist.py以生成训练集、测试集和验证集的文件列表。例如，在Places2数据集上生成训练集文件列表时，可执行以下命令：

mkdir datasets
python ./scripts/flist.py --path path_to_places2_train_set --output ./datasets/places_train.flist

2) 不规则掩码

我们的模型是在Liu等人提供的不规则掩码数据集上训练的。您也可以从他们的网站此处下载公开可用的不规则掩码数据集。

此外，您还可以下载Karim Iskakov提供的Quick Draw不规则掩码数据集，该数据集包含了5000万笔由人类手绘的笔触。

请使用scripts/flist.py按照上述说明生成训练集、测试集和验证集的掩码文件列表。

快速入门

请通过以下链接下载预训练模型，并将其复制到./checkpoints目录下。

Places2 | CelebA | 巴黎街景

或者，您也可以运行以下脚本自动下载预训练模型：

bash ./scripts/download_model.sh

1) 训练

要训练模型，请创建一个类似于示例配置文件的config.yaml文件，并将其复制到您的检查点目录下。有关模型配置的更多信息，请参阅配置指南。

EdgeConnect的训练分为三个阶段：1) 训练边缘模型，2) 训练修复模型，3) 训练联合模型。训练模型时，可执行以下命令：

python train.py --model [stage] --checkpoints [路径]

例如，在./checkpoints/places2目录下对Places2数据集的边缘模型进行训练：

python train.py --model 1 --checkpoints ./checkpoints/places2

不同数据集的收敛速度有所差异。例如，Places2数据集通常在一到两个epoch内即可收敛，而规模较小的CelebA数据集则可能需要近40个epoch才能收敛。您可以通过修改配置文件中的MAX_ITERS值来设置训练迭代次数。

2) 测试

要测试模型，请创建一个类似于示例配置文件的config.yaml文件，并将其复制到您的检查点目录下。有关模型配置的更多信息，请参阅配置指南。

您可以在三个阶段分别测试模型：1) 边缘模型，2) 修复模型，3) 联合模型。在每种情况下，都需要提供一张带有掩码的输入图像以及对应的灰度掩码文件。请确保掩码文件覆盖了输入图像中的整个掩码区域。测试模型时，可执行以下命令：

python test.py \
  --model [stage] \
  --checkpoints [路径] \
  --input [输入目录或文件路径] \
  --mask [掩码目录或文件路径] \
  --output [输出目录路径]

我们在./examples目录下提供了一些测试示例。请先下载预训练模型，然后运行以下命令：

python test.py \
  --checkpoints ./checkpoints/places2 
  --input ./examples/places2/images 
  --mask ./examples/places2/masks
  --output ./checkpoints/results

此脚本将使用./examples/places2/mask目录中的对应掩码修复./examples/places2/images目录中的所有图像，并将结果保存到./checkpoints/results目录中。默认情况下，test.py脚本会在第三阶段（--model=3）运行。

3) 评估

要评估模型，首先需要以测试模式在您的验证集上运行模型，并将结果保存到磁盘上。我们提供了一个实用工具./scripts/metrics.py，用于通过PSNR、SSIM和平均绝对误差来评估模型性能：

python ./scripts/metrics.py --data-path [验证集路径] --output-path [模型输出路径]

若需测量弗雷歇起始距离（FID分数），可运行./scripts/fid_score.py。我们使用来自此处的PyTorch实现版本，该版本采用了PyTorch Inception模型的预训练权重。

python ./scripts/fid_score.py --path [验证集路径、模型输出路径] --gpu [要使用的GPU编号]

替代边缘检测

默认情况下，我们使用 Canny 边缘检测器从输入图像中提取边缘信息。如果您希望使用外部边缘检测方法（例如 Holistically-Nested Edge Detection）来训练模型，则需要作为预处理步骤为整个训练/测试集生成边缘图，并使用上述说明的 scripts/flist.py 生成相应的文件列表。请确保文件名和目录结构与您的训练/测试集一致。您可以通过在配置文件中指定 EDGE=2 来切换到外部边缘检测。

模型配置

模型配置存储在您的检查点目录下的 config.yaml 文件中。以下表格提供了配置文件中所有可用选项的说明：

通用模型配置

选项	描述
MODE	1: 训练，2: 测试，3: 评估
MODEL	1: 边缘模型，2: 图像修复模型，3: 边缘-修复联合模型，4: 联合模型
MASK	1: 随机块，2: 半幅，3: 外部掩码，4: 外部 + 随机块，5: 外部 + 随机块 + 半幅
EDGE	1: Canny，2: 外部
NMS	0: 不进行非极大值抑制，1: 对外部边缘进行非极大值抑制
SEED	随机数生成器种子
GPU	GPU ID 列表，用逗号分隔，例如 [0,1]
DEBUG	0: 不调试，1: 调试模式
VERBOSE	0: 不输出详细信息，1: 在输出控制台中显示详细统计信息

加载训练、测试和验证集的配置

选项	描述
TRAIN_FLIST	包含训练集文件列表的文本文件
VAL_FLIST	包含验证集文件列表的文本文件
TEST_FLIST	包含测试集文件列表的文本文件
TRAIN_EDGE_FLIST	包含训练集外部边缘文件列表的文本文件（仅当 EDGE=2 时）
VAL_EDGE_FLIST	包含验证集外部边缘文件列表的文本文件（仅当 EDGE=2 时）
TEST_EDGE_FLIST	包含测试集外部边缘文件列表的文本文件（仅当 EDGE=2 时）
TRAIN_MASK_FLIST	包含训练集掩码文件列表的文本文件（仅当 MASK=3, 4, 5 时）
VAL_MASK_FLIST	包含验证集掩码文件列表的文本文件（仅当 MASK=3, 4, 5 时）
TEST_MASK_FLIST	包含测试集掩码文件列表的文本文件（仅当 MASK=3, 4, 5 时）

训练模式配置

选项	默认值	描述
LR	0.0001	学习率
D2G_LR	0.1	判别器/生成器学习率比例
BETA1	0.0	Adam 优化器的 beta1 参数
BETA2	0.9	Adam 优化器的 beta2 参数
BATCH_SIZE	8	输入批次大小
INPUT_SIZE	256	训练时的输入图像尺寸。（0 表示原始尺寸）
SIGMA	2	Canny 边缘检测中使用的高斯滤波器的标准差 （0: 随机，-1: 无边缘）
MAX_ITERS	2e6	训练模型的最大迭代次数
EDGE_THRESHOLD	0.5	边缘检测阈值（0-1）
L1_LOSS_WEIGHT	1	L1 损失权重
FM_LOSS_WEIGHT	10	特征匹配损失权重
STYLE_LOSS_WEIGHT	1	风格损失权重
CONTENT_LOSS_WEIGHT	1	知觉损失权重
INPAINT_ADV_LOSS_WEIGHT	0.01	对抗损失权重
GAN_LOSS	nsgan	nsgan: 非饱和 GAN，lsgan: 最小二乘 GAN，hinge: hinge 损失 GAN
GAN_POOL_SIZE	0	假图像池大小
SAVE_INTERVAL	1000	多少次迭代后保存一次模型（0: 永不保存）
EVAL_INTERVAL	0	多少次迭代后对模型进行一次评估（0: 永不评估）
LOG_INTERVAL	10	多少次迭代后记录一次训练损失（0: 永不记录）
SAMPLE_INTERVAL	1000	多少次迭代后保存一次样本（0: 永不保存）
SAMPLE_SIZE	12	每次采样时采集的图像数量

许可证

根据 知识共享署名-非商业性使用 4.0 国际许可协议 授权。

除非另有说明，本内容以 CC BY-NC 许可协议发布，这意味着您可以复制、混编、转换和基于本内容进行创作，但不得将该材料用于商业目的，并需适当注明出处及提供许可协议链接。

引用

如果您在研究中使用此代码，请引用我们的论文 EdgeConnect: 使用对抗式边缘学习的生成式图像修复 或 EdgeConnect: 基于边缘预测的结构引导图像修复：

@inproceedings{nazeri2019edgeconnect,
  title={EdgeConnect: Generative Image Inpainting with Adversarial Edge Learning},
  author={Nazeri, Kamyar and Ng, Eric and Joseph, Tony and Qureshi, Faisal and Ebrahimi, Mehran},
  journal={arXiv preprint},
  year={2019},
}

@InProceedings{Nazeri_2019_ICCV,
  title = {EdgeConnect: Structure Guided Image Inpainting using Edge Prediction},
  author = {Nazeri, Kamyar and Ng, Eric and Joseph, Tony and Qureshi, Faisal and Ebrahimi, Mehran},
  booktitle = {The IEEE International Conference on Computer Vision (ICCV) Workshops},
  month = {Oct},
  year = {2019}
}

EdgeConnect 快速上手指南

EdgeConnect 是一个基于对抗边缘学习的生成式图像修复（Inpainting）工具。其核心理念是“先线条，后颜色”，通过两阶段模型（边缘生成器 + 图像补全网络）来重建缺失区域的精细细节。

环境准备

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

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.x
• 深度学习框架: PyTorch 1.0+
• 硬件: NVIDIA GPU 及对应的 CUDA 和 cuDNN 支持
• 依赖库: 详见 requirements.txt

国内开发者提示：安装 PyTorch 时，建议访问 PyTorch 官网 选择对应环境的命令，或使用清华/中科大镜像源加速 pip 包安装。

安装步骤

1. 克隆项目

git clone https://github.com/knazeri/edge-connect.git
cd edge-connect

2. 安装依赖

首先根据官方指引安装 PyTorch，然后安装项目所需的其他 Python 包：

pip install -r requirements.txt

(注：国内用户可使用 pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple 加速下载)

3. 获取预训练模型

您可以手动下载模型并放入 ./checkpoints 目录，或使用脚本自动下载：

bash ./scripts/download_model.sh

支持的数据库集包括：Places2, CelebA, Paris-StreetView。

4. 准备数据集（可选，如需重新训练）

如果仅需使用预训练模型进行测试，可跳过此步。若需训练，需准备图像和不规则掩码（Irregular Masks），并使用脚本生成文件列表：

mkdir datasets
# 示例：为 Places2 训练集生成文件列表
python ./scripts/flist.py --path path_to_places2_train_set --output ./datasets/places_train.flist

基本使用

本工具的使用分为三个阶段：边缘模型、修复模型、联合模型。默认测试使用阶段 3（联合模型），效果最佳。

1. 配置文件

在 ./checkpoints 目录下创建 config.yaml 文件。您可以复制示例文件进行修改：

cp config.yml.example ./checkpoints/config.yaml

注意：确保配置中的路径与您实际存放模型和数据的路径一致。

2. 运行测试（图像修复）

准备一张带有缺失区域（白色遮罩）的输入图片和对应的灰度掩码图片（mask），运行以下命令：

python test.py \
  --checkpoints ./checkpoints/places2 \
  --input ./examples/places2/images \
  --mask ./examples/places2/masks \
  --output ./checkpoints/results

参数说明：

• --checkpoints: 预训练模型所在目录。
• --input: 输入图片目录或单张图片路径。
• --mask: 对应的掩码图片目录或单个掩码路径（需覆盖输入图中的缺失区域）。
• --output: 修复结果保存目录。
• --model: 默认为 3（联合模型），也可指定 1 (边缘) 或 2 (修复)。

执行完成后，修复后的图片将保存在 --output 指定的目录中。

3. 模型训练（进阶）

如需从头训练模型，需按顺序训练三个阶段。以下为训练边缘模型的示例：

python train.py --model 1 --checkpoints ./checkpoints/places2

• --model 1: 训练边缘生成器。
• --model 2: 训练图像补全网络。
• --model 3: 联合训练。

具体超参数（如学习率、迭代次数等）请在 config.yaml 中调整。