访问官网
Github
论文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）分数。
授权许可
本项目的内容受 Creative Commons Attribution-NonCommercial 4.0 International 许可，允许用户在非商业用途下复制、修改、改编和构建项目内容，条件是适当注明出处。
