大家好！ 我是Mail.ru集团计算机视觉团队的研究工程师。 在本文中，我将讲一个故事，说明我们如何为旧军事照片创建基于AI的照片恢复项目 。 什么是“照片恢复”？ 它包括三个步骤：

• 我们发现所有图像缺陷：断裂，磨损，破洞；
  
• 我们根据发现的缺陷周围的像素值对其进行修补；
  
• 我们为图像着色。
  

此外，我将描述照片恢复的每个步骤，并告诉您我们如何获取我们的数据，我们训练了哪些网络，我们完成了哪些工作以及犯了哪些错误。

寻找缺陷

我们要查找与上传的照片中的缺陷相关的所有像素。 首先，我们需要弄清楚人们会上传什么样的图片。 我们与“不朽军团”项目的创建者进行了交谈，该项目是一个非商业性组织，用于存储WW2的旧照片，并与我们共享数据。 经过分析，我们注意到人们上传的个人或团体肖像大多带有中等到大量的缺陷。

然后我们必须收集训练集。 分割任务的训练集是标记所有缺陷的图像和蒙版。 最简单的方法是让评估者创建细分蒙版。 当然，人们非常了解如何查找缺陷，但这将花费很长时间。



标记一张照片中的缺陷像素可能需要一个小时或整个工作日。 因此，在几周内收集超过100张图像的训练集并不容易。 这就是为什么我们试图扩充数据并创建自己的缺陷的原因：我们将拍摄一张好照片，使用图像上的随机游走添加缺陷，最后得到一个遮罩，显示带有缺陷的图像部分。 不进行扩充，我们在训练集中有68张手动标记的照片，在验证集中有11张照片。

最受欢迎的分割方法：将Unet与经过预训练的编码器结合使用，并最小化BCE（ 二进制交叉熵 ）和DICE（ 索伦森-Dice系数 ）之和。

当我们使用这种细分方法完成任务时会出现什么问题？

• 即使照片上看起来有很多缺陷，而且非常破旧，但缺陷区域仍然比未损坏的区域小得多。 为了解决这个问题，我们可以增加BCE中的正班级权重。 最佳权重应为干净像素与有缺陷像素的比率。
  
• 第二个问题是，如果我们使用带有预先训练的编码器的开箱即用的Unet（例如Albunet-18），则会丢失很多位置数据。 Albunet-18的第一层由卷积组成，其内核为5，步幅等于2。 它允许网络快速运行。 我们权衡了净操作时间以更好地进行缺陷定位：在第一层之后删除了最大池，将步幅减小到1，并将卷积核减小到3。
  
• 如果我们通过将小图像压缩到例如256 x 256或512 x 512像素来处理小图像，则由于插值，小缺陷将消失。 因此，我们需要处理较大的图像。 我们目前正在生产中分割1024 x 1024尺寸的照片中的缺陷。 这就是为什么我们必须在大图像作物上训练网络。 但是，这会在单个GPU上产生小批量的问题。
  
• 在培训期间，我们可以在一个GPU上容纳约20张图像。 因此，我们最终在BatchNorm图层中得到不正确的均值和标准差值。 我们可以使用就地BatchNorm解决此问题，一方面可以节省内存空间，另一方面可以使用Synchronized BatchNorm版本，该版本可以在所有GPU之间同步统计信息。 现在，我们不计算单个GPU上20张图像的平均值和标准偏差值，而是计算4个GPU上80张图像的平均值和标准偏差值。 这改善了网络收敛。
  

最后，在增加BCE权重，更改体系结构并使用就地BatchNorm后，我们使分段效果更好。 但是，通过添加“测试时间增强”来做更好的事情并不会花费太多。 我们可以在输入图片上运行一次网络，然后对其进行镜像，然后重新运行网络以查找所有小的缺陷。



在四台GeForce 1080Ti上，网络可以在18小时内收敛。 推理需要290毫秒。 这很长，但这就是我们优于默认性能的代价。 验证DICE等于0.35，而ROCAUC等于0.93。

图像修补

与使用Unet的细分任务相同。 要进行修复，我们将上传原始图像和蒙版，并在其中用1和0（要修复的所有像素）标记所有干净区域。 这就是我们收集数据的方式：对于来自开源图像数据集（例如OpenImagesV4）的任何照片，我们都会添加与现实生活中类似的缺陷。 然后，我们训练了网络以恢复丢失的零件。

我们如何修改此任务的Unet？

我们可以使用部分卷积代替原始卷积。 这个想法是，当我们将某个区域与某个内核卷积在一起时，我们不会考虑缺陷像素值。 这样可以使修复更加精确。 我们向您展示最近NVIDIA论文中的一个例子。 他们使用的Unet在中间图片中具有默认的二维卷积，在右侧图片中具有部分卷积。



我们训练了网络五天。 在最后一天，我们冻结了BatchNorms，以使被涂漆零件的边框不可见。

处理一张512 x 512图片需要50毫秒的净时间。 验证PSNR等于26.4。 但是，您不能完全依赖此任务中的指标。 为了选择最佳模型，我们在评估图像上运行了几个好的模型，对结果进行匿名处理，然后投票选出我们最喜欢的模型。 这就是我们选择最终模型的方式。

前面已经提到过，我们人为地给干净的图像添加了一些缺陷。 您应该始终跟踪培训过程中所添加缺陷的最大大小； 如果您将具有很大缺陷的图像馈送到网上，而在训练阶段却从未处理过，则该网会疯狂运行并产生不适用的结果。 因此，如果您需要修复较大的缺陷，请使用这些缺陷来扩充您的培训。

这是我们的算法如何工作的示例：

显色

我们对缺陷进行了分割并修复。 第三步-颜色重建。 就像我之前说过的，不朽军团照片中有很多个人和团体肖像。 我们希望我们的网络与他们合作良好。 由于现有服务都无法快速有效地为肖像上色，因此我们决定使用自己的颜色。 我们希望我们的彩色照片更加可信。



GitHub有一个流行的用于照片着色的存储库 。 它做得很好，但仍然存在一些问题。 例如，它倾向于将衣服涂成蓝色。 这就是为什么我们也拒绝它。

因此，我们决定创建一种用于图像着色的算法。 最明显的想法是：拍摄黑白图像并预测三个通道：红色，绿色和蓝色。 但是，我们可以使工作更加轻松：不使用RGB颜色表示，而是使用YCbCr颜色表示。 Y分量是亮度（亮度）。 上传的黑白图像是Y通道，我们将重用它。 现在我们需要预测Cb和Cr：Cb是蓝色和亮度的差，而Cr是红色和亮度的差。



为什么选择YCbCr表示形式？ 人眼对亮度的变化比对颜色的变化更敏感。 这就是为什么我们重复使用人眼最敏感的Y分量（亮度）并预测可能会犯错的Cb和Cr的原因，因为我们不能很好地注意到颜色虚假。 当频道容量不足以传输所有颜色时，这种特殊的特性在彩色电视的黎明时被广泛使用。 图片在YCbCr中传输，与Y分量相同，并且Cb和Cr减少了一半。

如何创建基线

我们可以将Unet与经过预训练的编码器结合使用，并将现有CbCr值与预测值之间的L1损耗最小化。 我们要给肖像上色，因此，除了OpenImages照片外，我们还需要更多特定于任务的照片。

我们在哪里可以获得穿着军装的人的彩色照片？ 互联网上有些人将旧照片着色是一种业余爱好或需要付出一定的代价。 他们非常小心地做，试图做到非常精确。 当他们为制服，护肩板和奖章上色时，他们会参考档案资料，因此他们的工作结果值得信赖。 总而言之，我们使用了200张带有人工制服的彩色图片。

另一个有用的数据源是“工人和农民”红军网站。 它的一位创始人几乎每次都可以在第二次世界大战的苏联制服上拍摄他的照片。



在某些图片中，他从著名的档案照片中模仿了人们的姿势。 他的图片有白色背景是一件好事：通过在背景中添加各种自然物体，我们可以很好地扩充数据。 我们还使用了一些常规的肖像，并补充了徽章和其他战时属性。

我们培训了AlbuNet-50-这是一个Unet，它使用了预先培训的ResNet-50作为编码器。 蚊帐开始产生足够的效果：皮肤为粉红色，眼睛为灰蓝色，肩板为淡黄色。 但是，问题在于它使照片上的某些区域保持不变。 这是由于以下事实导致的：根据误差L1找到最佳位置，最好什么都不做，而不是尝试预测某种颜色。


我们正在将结果与地面真相照片进行比较- 克林比姆（Klimbim）进行了手动着色

我们如何解决这个问题？ 我们需要一个鉴别器：一个神经网络，它将接收图像并告诉我们它看起来是否逼真。 下面的图片之一是手动着色的，另一张是我们的发电机AlbuNet-50着色的。 人类如何手动和自动区分彩色照片？ 通过查看细节。 您能告诉我们基准解决方案自动将彩色照片放在哪里吗？




我们使用Self-Attention GAN论文中的鉴别器。 这是一个小型卷积网络，其顶层构建了所谓的“自我注意”。 它使我们可以“更加关注”图像的细节。 我们还使用频谱归一化。 您可以在上述论文中找到更多信息。 我们用L1损失和鉴别器损失共同训练了网络。 现在，网络可以更好地使图像细节着色，并且背景看起来更加一致。 再举一个例子：左边是仅受L1损失训练的网络工作。 在右边-加上L1鉴别器损失。



培训过程在四台GeForce 1080Ti上进行了两天。 处理512 x 512图片需要30毫秒的净时间。 验证MSE-34.4。 就像修复一样，您也不想依赖指标。 因此，我们选择了六个具有最佳验证指标的模型，并盲目投票赞成最佳模型。

当我们已经创建了一个生产系统并启动了一个网站时，我们继续进行实验，得出的结论是，我们最好不要最小化每个像素的L1损失，而是最小化感知损失。 为了进行计算，我们将网络预测和地面真实照片提供给VGG-16网络，在底层获取特征图，并将其与MSE进行比较。 这种方法可绘制更多区域，并提供更多彩色效果。

回顾

Unet是一个非常酷的模型。 在第一个分割任务中，我们在训练过程中遇到了问题，并使用高分辨率图像，这就是为什么我们使用就地BatchNorm的原因。 在第二个任务（修复）中，我们使用了部分卷积而不是默认卷积，它使我们可以获得更好的结果。 在进行着色时，我们添加了一个小的鉴别器网，对生成器提供不真实的图像会造成损失。 我们还使用了感知损失。

第二个结论-评估者至关重要。 而且不仅在创建细分蒙版阶段，而且还用于最终结果验证。 最后，我们为用户提供了三张照片：带有已修复缺陷的原始图像，具有已修复缺陷的彩色照片以及一个简单着色的图像，以防缺陷搜索和修复算法出错。

我们从War Album项目中拍摄了一些照片，并通过这些神经网络对其进行了处理。 这是我们得到的结果：



此外， 在这里您可以仔细查看原始图像和所有处理阶段。