为修复一张老照片 动用了四块1080Ti显卡

嗨！各位，我是一位来自Mail.ru Group计算机视觉团队的研发工程师，在这篇文章当中，我将讲述我们如何为老军人的照片创造一个基于AI技术的照片修复项目。什么是照片修复呢？它由以下三个步骤组成：

-我们找到所有的照片缺陷：折痕，磨损，破洞

-我们基于所发现的照片缺陷周围的像素值来进行图像修复

-我们为图像上色

接下来，我将阐释照片修复的每一个步骤，并且告诉你我们如何获得数据，用哪种网络训练，取得了哪些成就，以及犯了什么错误。

寻找缺陷

我们需要在一张已经上传的照片当中找到所有与缺陷相关的像素值。首先，我们需要弄清楚人们会上传什么样的照片。我们与immortal regiment项目的创始人交流过，一个存储二战遗留照片的非商业组织，他们分享数据给我们。在此基础上进行分析，我们意识到人们上传的大部分个人或集体肖像存在中等到大量的缺陷。

接着我们必须收集一个训练集，这个训练集需要对图像进行分割，得到一张图片，和一张把所有缺陷都标注出来的黑底白纹蒙版。做这一步最简单的方法是让评估员创立分割蒙版。当然，一般人非常清楚怎样找到缺陷，但那会耗费太长时间。

标记一张照片中缺陷处的像素将会花费一个小时或一整个工作日，因此，在几周内收集一个超过100张图片的训练集是不简单的。这就是为什么我们尝试加强我们的数据，并创造我们自己的缺陷：拍一张照片，用图片上的随机纹理增加缺陷，最终得到一张显示有缺陷的图像部分的蒙版。在没有增强的情况下，我们已经获得了含有68张手动标记照片的训练集和和11张照片的验证集。

最流行的分割方法是：使用unet和预训练编码器，将BCE（binary cross-entropy）与DICE（S?rensen–Dice coefficient）的和最小化。

在我们为了项目任务使用这个分割方法时，什么样的问题会出现呢？

虽然照片看起来像有许多的缺陷，非常老旧而残破，有缺陷的部分仍然远远小于没有受到损伤的部分。为了解决这个问题，我们可以增加BCE中的积极类内权重，最优权重应是清晰像素与缺陷像素的比率。

第二个问题是如果我们用可以立刻上手使用的预训练编码器和UNET，将会损失许多位置数据。第1层如果由一个核为5、步长为2的卷积组成，网络运行得更快。我们牺牲网络运行的时间来换取更好的缺陷定位。我们在第1层之后移除最大的pooling层，将卷积步长减至1，卷积核减至3。

如果我们通过压缩处理小图片，比如说到256或512像素，小瑕疵会因为内插而消失。然而，我们还需要处理大图像。在实际应用中我们一般把1024像素的照片瑕疵分割。这就是为什么我们必须将网络训练成为适于处理大图像。但这会引起在单个GPU上处理规格小批量的问题。

在训练中，我们我们可以在一个GPU上匹配20张左右的图片。因此我们最终会在BatchNorm层得到不准确的平均值和标准偏差值。我们可以通过 In-place BatchNorm 来解决这个问题。一方面节省内存空间，另一方面拥有一个 Synchronized BatchNorm版本，可以同步所有GPU上的数据。这样我们计算平均值和标准偏差值就不是参照一个GPU上的20张图片，而是4个GPU上的80张图片，这改善了网络集成。

最后，基于增加bce的权重，改变体系结构和使用In-place BatchNorm，我们将分割做得更好，但为了做得更好而使用Test Time Augmentation并不会花费许多。我们可以在一张输入的图片上运行网络，获取镜像，再重新运行网络去找到所有的小缺陷。

用了四个GeForce 1080Ti显卡，花费了18个小时，相对时间较长，但这就是精益求精的代价。

图像修复

我们使用了在图像分割领域很火的Unet网络。为了做图像修复，我们需要上传原始图像和一个掩模(使用1标记了所有干净区域的像素，使用0标记了需要好修复的像素)。数据是这样被搜集的：在一些开源的图像数据集上(例如：OpenImagesV4)加入了一些在生活中会出现的照片瑕疵，然后我们训练网络去重新修复有瑕疵的这部分。

我们是如何使用在图像修复中使用Unet的呢？

我们使用了部分卷积替代了原始的那一版。当我们使用一些核去卷积一个区域时，我们没有把损坏的像素算在内。这使得图像修复更精密。从英伟达的论文来举个例子，他们使用了Unet结合默认的2维卷积的修复结果在中间这幅图，使用部分卷积的修复结果在右边这幅图。

我们花了5天来训练这个网络。在最后一天我们停止了BatchNorms 来使得受损部分的边缘更不易被看见。

这个网络处理一副512*512的图像需要花费50ms，验证集上的PSNR（“Peak Signal to Noise Ratio”的缩写，即峰值信噪比）等于26.4。然而在这种任务中，你不能完全依赖这个指标。为了选择最好的模型，我们在一些验证图片上运行几个好的模型，将结果隐藏，然后投票选出我们最喜欢的模型修复过的图片，这是我们如何选择最终模型的方法。

我之前提到过我们人工在干净的图片上加了一些缺陷 。所以在训练时需要一直跟踪添加的缺陷的最大面积。以防当你"喂"给网络一张它从没有在训练中处理过的，有很大缺陷的图片。网络可能会失控并且产生不可用的结果。因此，如果你想修复有大缺陷的图片，增大你训练集里的缺陷。

这是一个关于我们的算法效果的例子：

着色

我们已经将图片有缺陷的部分分割了出来并且修复了他们；第三步—;—;颜色复原。就像我之前说的，Immortal Regiment项目有很多个人或集体的相片。我们希望网络可以在修复这些照片上起到很好的作用。因为现有的服务不能有效且快的着色，所以我们决定提出我们自己的着色方案。我们希望我们的被着色的图片更可信。

github有个很流行的相片着色repo。它有着很好的效果但是还是有很多问题。例如，它有可能会将衣服涂成蓝色，这也是我们没有使用他的原因。

所以，我们决定在图像着色方面提出一个算法。其中最简单的方案是：预测一张黑白照片的R,G,B三个通道的值。但是，我们可以使工作更简单：使用YCbCr色彩表达式来替换RGB色彩表达式。Y通道表示亮度。一张上传的黑白照片只在Y通道里有数值，我们会复用这个数值。而只需要预测Cb和Cr通道的值。Cb是蓝色与亮度的差值，Cr是红色与亮度的差值。

为什么要选择YCbCr色彩表达式？因为人眼对明暗变化的敏感程度会比对色彩变化的敏感程度更高。这是我们选择复用了人眼更敏感的Y通道（亮度）及预测我们可能出错的Cb和Cr通道值的原因，因为我们不会很注意颜色错误。这个具体的特性广泛用于彩色电视产生的初期，那时通道传输能力还不足以传输所有颜色。图片是依据YCbCr色彩表达式来传输的，Y通道数值没有改变，Cb和Cr的数值均减半。

如何新建一个基准模型

我们将现有的CbCr值输入进带有预训练好编码器的Unet网络，然后最小化L1损失函数来预测合理的CbCr值。我们希望为照片上色，因此除了OpenImages数据集的相片之外，我们需要更多任务特定的相片。

我们去哪里得到穿着军装的人的彩色照呢？在网上会有些人因为兴趣或者金钱为老照片上色。他们做的很仔细很准确。他们是根据档案材料来为制服，肩牌和勋章上色，所以他们的成果很值得信赖。最后，我们使用了200幅手工上色的身着军装的人的照片。

另一个有用的数据源是The Workers’ and Peasants’ Red Army 网站。它的创办者之一几乎拍了所有自己身着二战时期军装的照片。

在某些照片中，他模仿了著名照片人物的姿势。他的照片都是白色背景，这可以让我们在背景上添加不同的自然风景来扩充数据。我们也使用了一些平常的相片，并在上面添加了肩章和其他战争相关的元素。

我们训练了 AlbuNet-50—;—;这是一个使用了预训练好的ResNet-50 作为编码器的Unet。在预测时，这个网络给出了一些合理的结果：皮肤是粉色的，眼睛—;—;灰绿色，肩章—;—;微黄色。然而，会有一些区域没有被着色。这是因为跟去预测一些别的颜色相比，不做任何变化会使L1误差有更优的结果。

我们正在比较我们的结果与原始真实照片-由Klimbim手工着色

我们该如何解决这个问题呢？我们需要一个判别：一个输入图像，并判断图像是否是真实的图像的神经网络。下面的其中一张图片是手工着色的，另一张是由我们的生成器AlbuNet-50绘制的。人类如何区分手动和自动着色的照片呢?通过观察细节。你能告诉我们的哪个照片是基准解决方案得到的吗？

回答：左边的图片是手动上色的，右边是自动上色的。

我们使用了基于自注意机制的GAN的论文中的判别器。这是一个很小的卷积网络，在最上面的几层是被称为自注意力的层。它使我们的网络“更加关注”图像细节。我们也使用谱归一化。你可以在上面的论文中找到更多的细节。我们利用L1损失项和判别器损失项的组合训练了网络。现在得到的网络在图像的细节上着色更好，背景看起来更一致。再举一个例子:左边是只训练L1损失项的网络得到的结果;右边是结合L1损失项的结果。

在四个GeForce 1080Ti上，训练过程花费了两天时间。处理一张512x512的图片需要30毫秒。验证集的最小均方误差(MSE)为34.4。就像图像修复一样，基于我们并不想选择的评价准则。这就是为什么我们选择了6个具有最佳验证指标的模型，并直接为最佳模型投票。

当我们已经创建了一个生产系统，并推出了一个网站，我们继续验证，获得结果，我们最好尽量减少每个像素L1损失，而不是减少感知损失。为了计算它，我们将网络的预测和一张ground-truthl照片提供给VGG-16网络，获取底层的特征map，并利用最小均方误差进行比较。这种方法修复了更多的区域，并给出了更有色调的结果。

简要回顾

Unet是一个非常棒的模型。在第一个分割任务中，我们在训练过程中遇到了一个问题，就是处理高分辨率的图像，这就是为什么我们使用In-Place 批归一化。在我们的第二个任务(图像修复)中，我们使用了部分卷积而不是标准卷积，这让我们得到了更好的结果。在进行着色时，我们增加了一个简单的判别器网络，它可以对生成不真实图像的生成器进行惩罚。我们还使用了感知损失。

第二个结论—;—;评估人员是必不可少的。不仅在创建分割掩码阶段，而且在最终的结果验证阶段。最后，我们给用户提供了三张照片:一幅带有缺陷的原始照片，一幅带有缺陷的彩色照片，以及一幅简单的彩色照片，以防缺陷搜索和图像修复算法出错。

我们从War Album project 中获得了一些照片，并在这些神经网络上进行了处理。以下是我们得到的结果：

此外，你还可以更详细地查看原始图像和所有的处理过程，详情访问：点此。