edge-connect

EdgeConnect: 生成对抗边缘学习的图像修复

简介

EdgeConnect 是一种创新的图像修复方法，灵感来源于艺术家的工作方式，即“先画线，再填色”。该方法采用了一个两阶段的生成对抗模型，包括边缘生成器和图像完成网络。边缘生成器用于推测和勾勒出图像中缺失区域（无论规则或不规则）的边缘，随后图像完成网络利用这些推测出的边缘信息来填补缺失区域。详细的系统描述可以在相关的论文中找到。

EdgeConnect 通过这种方法能够更好地填补有细节的缺失区域，与传统方法相比有显著提升。从输入过程中可以看到，缺失区域用白色显示，而边缘生成器网络预测出的边缘用蓝色表示（已知区域的边缘通过 Canny 边缘检测器计算并显示为黑色），最终实现的图像补全能够更好地再现丢失的信息。

环境需求

Python 3
PyTorch 1.0
NVIDIA GPU 与 CUDA cuDNN

安装步骤

克隆项目代码：

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

从 PyTorch 官网安装相应的 PyTorch 版本。
安装 Python 依赖：

pip install -r requirements.txt

数据集

EdgeConnect 模型使用了多种数据集进行训练，包括 Places2、CelebA 和 Paris Street-View 数据集。用户需要从这些数据集的官方网站下载完整的数据集。下载完成后，使用 scripts/flist.py 生成训练集、测试集和验证集文件列表。

对于不规则遮罩，我们使用了 Liu 等人提供的不规则遮罩数据集，用户可以从他们的网站下载。此外，还可以选择下载 Karim Iskakov 的 Quick Draw 不规则遮罩数据集中提供的手绘笔画数据集。

快速上手

在开始之前，需要下载预训练模型并将其存放在 ./checkpoints 目录下。这可以通过下面的链接手动下载：

Places2 预训练模型
CelebA 预训练模型
Paris-StreetView 预训练模型

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

bash ./scripts/download_model.sh

训练模型

用户需要创建一个与示例配置文件类似的 config.yaml 文件，并将其复制到 checkpoints 目录中。EdgeConnect 模型训练分为三个阶段：1) 边缘模型训练，2) 图像修补模型训练，3) 整体模型训练。

训练命令格式如下：

python train.py --model [stage] --checkpoints [path to checkpoints]

例如，要在 Places2 数据集上训练边缘模型，可以执行：

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

不同的数据集收敛速度不同，比如 Places2 在一到两个周期即可收敛，而 CelebA 数据集则需大约 40 个周期。用户可以通过修改配置文件中的 MAX_ITERS 值设置训练的最大迭代次数。

测试模型

用户可以利用三种阶段的模型进行测试：1) 边缘模型，2) 图像修补模型，3) 整体模型。需要提供一个带有遮罩的输入图像和对应的灰度遮罩文件。测试命令如下：

python test.py \
  --model [stage] \
  --checkpoints [path to checkpoints] \
  --input [path to input directory or file] \
  --mask [path to masks directory or mask file] \
  --output [path to the output directory]

模型配置

模型配置存储在 checkpoints 目录下的 config.yaml 文件中。该配置文件采用 YAML 格式，用户可以根据需要更改各项配置参数以满足不同的需求。

评价模型

用户可以在验证集上运行 EdgeConnect 模型的测试，并将结果保存在磁盘上。项目还提供了用于评估 PSNR、SSIM 和均值绝对误差的脚本 ./scripts/metrics.py。此外，还可以运行 ./scripts/fid_score.py 计算 Fréchet Inception Distance（FID）分数。