使用 Keras 的部分卷积用于图像修复

Keras 实现了论文《使用部分卷积修复不规则孔洞的图像》，https://arxiv.org/abs/1804.07723。在此向 NVIDIA 公司的作者 Guilin Liu、Fitsum A. Reda、Kevin J. Shih、Ting-Chun Wang、Andrew Tao 和 Bryan Catanzaro 致以崇高敬意，感谢他们发表了这篇精彩的论文。实现该架构、部分卷积层以及损失函数对我来说是一次非常宝贵的学习经历。

依赖项

• Python 3.6
• Keras 2.2.4
• Tensorflow 1.12

如何使用本仓库

尝试该算法进行几组预测最简单的方式是访问 www.fixmyphoto.ai。我在那里部署了一个无服务器的 React 应用程序，并使用 AWS Lambda 函数来处理推理任务。

如果你想深入研究代码，新的 PConv2D Keras 层以及使用这些部分卷积层的类似 UNet 的架构的主要实现分别位于 libs/pconv_layer.py 和 libs/pconv_model.py 中——大部分实现都在这里。除此之外，我还准备了四个 Jupyter 笔记本，详细记录了我在实现该网络时所经历的几个步骤，具体如下：

步骤 1：创建随机不规则掩码
步骤 2：实现并测试 PConv2D 层的实现
步骤 3：实现并测试使用 PConv2D 层的 UNet 架构
步骤 4：在 ImageNet 数据集上训练和测试最终架构
步骤 5：通过图像分块尝试预测任意尺寸的图像

预训练权重

我已将 VGG16 的权重从 PyTorch 转移到 Keras；这意味着可以像在 PyTorch 中一样，对 VGG16 网络使用 1/255. 的像素归一化。

• 转换后的 VGG16 权重
• PConv 在 ImageNet 上的结果
• PConv 在 Places2 上的结果 [待训练]
• PConv 在 CelebaHQ 上的结果 [待训练]

在自己的数据集上训练

你可以直接跳到 步骤 4 笔记本，或者使用命令行工具（请确保下载已转换的 VGG16 权重）：

python main.py \
    --name MyDataset \
    --train TRAINING_PATH \
    --validation VALIDATION_PATH \
    --test TEST_PATH \
    --vgg_path './data/logs/pytorch_to_keras_vgg16.h5'

实现细节

实现的具体细节已在论文中说明（https://arxiv.org/abs/1804.07723），不过我还是会在这里总结一些要点。

掩码生成

论文中使用了一种基于视频中连续两帧之间遮挡与非遮挡关系的技术来生成随机不规则掩码。而我则选择简单地编写一个掩码生成函数，利用 OpenCV 绘制一些随机的不规则形状作为掩码。不过，后续引入新的掩码生成技术应该不会有问题，而且我认为使用这种方法得到的结果也相当不错。

部分卷积层

本实现中的关键组件是部分卷积层。基本原理是：给定卷积滤波器 W 和对应的偏置 b，与普通卷积不同，这里应用的是部分卷积：

其中 ⊙ 表示逐元素相乘，M 是由 0 和 1 组成的二值掩码。重要的是，在每次部分卷积之后，掩码也会更新：如果卷积能够基于至少一个有效输入生成输出，则该位置的掩码会被移除，即：

其结果是，经过足够深的网络后，掩码最终会变为全 1（即消失）。

UNet 架构

架构的具体细节可在论文中找到，但本质上它是一种类似 UNet 的结构，所有普通的卷积层都被部分卷积层取代，从而确保图像始终与掩码一起通过网络。以下是该架构的概览。

损失函数

论文中使用的损失函数较为复杂，具体内容可参阅原文。简而言之，它包括：

• 针对有掩码区域和无掩码区域的逐像素损失
• 基于 ImageNet 预训练 VGG-16 的感知损失（pool1、pool2 和 pool3 层）
• 针对预测图像和计算图像（将无孔洞区域设置为真实标签）的风格损失
• 针对孔洞区域进行 1 像素膨胀后的总变差损失

这些损失项的权重分配如下：

训练流程

网络是在 ImageNet 数据集上以批次大小为 1 进行训练的，每个 epoch 设定为 10,000 个批次。此外，由于批归一化会对带掩码的卷积造成问题（因为会针对孔洞像素计算均值和方差），因此训练分两个阶段进行，采用 Adam 优化器。

第一阶段 学习率设为 0.0001，训练 50 个 epoch，所有层都启用批归一化。

第二阶段 学习率降至 0.00005，训练 50 个 epoch，但所有编码层的批归一化都被禁用。

用于展示的图像训练时间极其漫长，这很可能是因为我个人的硬件配置较差。我在 1080Ti 显卡上进行的几次测试（批次大小为 4）表明，训练时间可能与论文中提到的一样，约为 10 天。

PConv-Keras 快速上手指南

PConv-Keras 是基于 Keras 实现的“不规则孔洞图像修复（Image Inpainting）”工具，核心算法源自 NVIDIA 的论文《Image Inpainting for Irregular Holes Using Partial Convolutions》。该工具通过局部卷积层（Partial Convolution）有效处理图像中的不规则缺失区域。

环境准备

在开始之前，请确保您的开发环境满足以下版本要求。由于项目依赖较旧的 TensorFlow 和 Keras 版本，建议使用虚拟环境（如 venv 或 conda）以避免冲突。

• 操作系统: Linux / macOS / Windows
• Python: 3.6
• Keras: 2.2.4
• TensorFlow: 1.12
• 其他依赖: OpenCV (用于掩码生成), Jupyter Notebook (可选，用于运行示例笔记)

提示：国内用户可使用清华源或阿里源加速 Python 包安装。

安装步骤

1. 克隆仓库

   git clone https://github.com/MathiasGruber/PConv-Keras.git
   cd PConv-Keras
   

2. 安装依赖 由于指定了特定版本，建议手动安装以确保兼容性：

   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow==1.12 keras==2.2.4 opencv-python jupyter
   

3. 下载预训练权重 为了直接进行推理或微调，需要下载转换后的 VGG16 权重及 ImageNet 上的 PConv 模型权重：

   • Ported VGG 16 weights
   • PConv on Imagenet

   下载后将文件放置于项目目录下的 data/logs/ 文件夹中（或根据后续命令参数指定路径）。

基本使用

方式一：在线体验（无需安装）

如果您仅想测试算法效果，可以直接访问作者部署的在线应用： www.fixmyphoto.ai

方式二：本地代码探索与推理

核心实现位于 libs/pconv_layer.py (局部卷积层) 和 libs/pconv_model.py (UNet 架构)。作者提供了详细的 Jupyter Notebook 流程，建议按顺序运行以理解原理：

1. Step 1: 创建随机不规则掩码 (notebooks/Step1...)
2. Step 2: 实现并测试 PConv2D 层
3. Step 3: 构建带有 PConv2D 的 UNet 架构
4. Step 4: 在 ImageNet 上训练和测试最终架构
5. Step 5: 尝试通过图像分块预测任意尺寸图像

您可以启动 Jupyter 来运行这些笔记：

jupyter notebook

方式三：使用自有数据集训练

如果您希望在自己的数据集上进行训练，请确保已下载 VGG16 权重，然后使用命令行接口（CLI）。

命令示例：

python main.py \
    --name MyDataset \
    --train TRAINING_PATH \
    --validation VALIDATION_PATH \
    --test TEST_PATH \
    --vgg_path './data/logs/pytorch_to_keras_vgg16.h5'

参数说明：

• --train: 训练集图片路径
• --validation: 验证集图片路径
• --test: 测试集图片路径
• --vgg_path: 下载的 VGG16 权重文件路径

训练策略提示： 根据论文实现，训练分为两个阶段以解决批归一化（Batch Normalization）在掩码卷积中的问题：

1. 第一阶段：学习率 0.0001，运行 50 个 epoch，启用所有层的 Batch Norm。
2. 第二阶段：学习率 0.00005，运行 50 个 epoch，禁用编码层中的 Batch Norm。

注：在单张 GPU（如 1080Ti）上完整训练可能需要约 10 天时间。