基于人工智能的照片修复项目案例分析

AI-Based Photo Restoration

大家好！我是mail.ru集团计算机视觉团队的研究工程师。在这篇文章中，我将讲述一个关于我们如何为旧的军事照片创建基于人工智能的照片修复项目。什么是“照片复原”？它包括三个步骤：
•发现所有图像的缺陷：破裂、刮伤、孔洞；
•根据周围的像素值修复发现的缺陷；
•给图像上色。
此外，我将描述照片恢复的每一步，并告诉你我们如何获得数据，我们训练了哪些网络，我们完成了什么，以及我们犯了什么错误。

寻找缺陷

我们想在上传的照片中找到所有与缺陷相关的像素。首先，我们需要弄清楚人们会上传什么样的图片。我们采访了“不朽军团”项目的创始人，这是一个非商业组织，存储了二战时遗留的照片，他们与我们分享了他们的数据。经过分析，我们注意到人们上传的大多是个人或团体肖像，存在中等到大量的缺陷。
然后我们得收集一套训练器材。分割任务的训练集是一个图像和一个标记所有缺陷的掩码。最简单的方法是让评估人员创建分段掩码。当然，人们很清楚如何发现缺陷，但这需要很长时间。

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

最流行的分割方法是：使用带预训练编码器的Unet，并最小化BCE（二进制交叉熵）和DICE（binary cross-entropy]) 。

当我们在任务中使用这种分割方法时会出现什么问题？

•即使照片看起来有很多缺陷，很旧很破旧，但有缺陷的区域仍然比未损坏的小得多。为了解决这个问题，我们可以增加BCE中的正类权重，一个最佳的权重是干净像素与有缺陷像素的比率。
•第二个问题是，如果我们使用带预训练编码器（例如Albunet-18）的开箱即用Unet，我们会丢失很多位置数据。Albunet-18的第一层由一个核5的卷积和一个等于2的步长组成。它可以让网络快速工作。为了更好地定位缺陷，我们交换了网络操作时间：删除了第一层之后的最大池，将步长减少到1，将卷积核减少到3。
•如果我们通过压缩小图像（例如，将其压缩到256 x 256或512 x 512像素）来处理小图像，则由于插值，小缺陷将消失。因此，我们需要处理更大的图像。我们目前正在生产中对1024x 1024大小的照片进行缺陷分割。这就是为什么我们必须在大图像上训练网络。但是，这会导致单个GPU上的小批量出现问题。
•在培训期间，我们可以在一个GPU上安装大约20个图像。因此，我们在batchnorm层中的平均值和标准偏差值不准确。我们可以使用就地batchnorm解决这个问题，一方面节省内存空间，另一方面，它有一个同步的batchnorm版本，可以同步所有GPU上的统计信息。现在，我们计算的平均值和标准偏差值不是针对单个GPU上的20个图像，而是针对4个GPU上的80个图像。这改善了网络收敛性。
最后，通过增加bce权重、改变体系结构、使用原位batchnorm，使分割效果更好。然而，通过增加测试时间来做更好的事情不会花费太多的成本。我们可以在输入图片上运行一次网络，然后对其进行镜像并重新运行网络以查找所有的小缺陷。

网络在18小时内汇聚在4个Geforce 1080ti上。推理需要290毫秒。这很长，但这是我们优于默认性能的代价。Validation DICE equals 0,35, and ROCAUC — 0,93.

图像修复

与分割任务一样，我们使用Unet。要进行修补，我们需要上传一张原始图像和一个遮罩，在那里我们用1和0标记所有清洁区域-所有我们想要修补的像素。这就是我们收集数据的方式：对于来自开源图像数据集（例如openimagesv4）的任何照片，我们都添加了与现实生活中看到的类似的缺陷。然后我们训练网络恢复丢失的部分。
我们如何修改这个任务的Unet？
我们可以用部分卷积代替原来的卷积。其思想是，当我们用一些核卷积一个区域时，我们不考虑缺陷像素值,这使得修补更加精确。我们向您展示了最近Nvidia论文中的一个例子。他们在中间的图片中使用了默认的二维卷积的Unet，在右边的图片中使用了部分卷积的Unet。

我们训练花了五天。最后一天，我们冻结了批处理规范，以减少着色部分的边界。
网络处理一张512 x 512的图片需要50毫秒。验证PSNR等于26,4。但是，您不能完全依赖此任务中的度量。为了选择最好的模型，我们对估值图像运行了几个好的模型，匿名投票给我们最喜欢的模型，这就是我们选择最终模型的方式。
如我之前提到过，我们人为地在干净的图像中添加了一些缺陷。在训练过程中，您应该始终跟踪添加的缺陷的最大大小；如果您将一个具有非常大缺陷的图像馈送到网络，而在训练阶段它从未处理过该图像，则网络将疯狂运行并产生不适用的结果。因此，如果你需要修复大的缺陷，用它们来增加你的训练集。
下面是我们的算法如何工作的示例：

着色

我们将缺陷分割并修补；第三步-颜色重建。就像我之前说的，照片中有很多个人和团体的肖像。我们希望我们的网络能和他们很好地配合。我们决定提出自己的彩色化，因为没有一个现有的服务可以快速和有效的彩色肖像。我们希望彩色照片更真实可信。

Github有一个流行的照片着色库。它做得很好也有一些问题。例如，它倾向于把衣服涂成蓝色。所以我不用它。
因此，我们决定创建一个图像着色算法。最明显的想法是：拍摄一张黑白图像，预测三个通道：红色、绿色和蓝色。然而，我们可以使我们的工作更容易：工作不是用RGB颜色表示，而是与YCbCr颜色表示。y分量是亮度（luma）。上传的黑白图像是y通道，我们将重用它。现在我们需要预测Cb和Cr：Cb是蓝色和亮度的差，Cr是红色和亮度的差。

为什么我们选择YCbCr代表？人眼对亮度的变化比对颜色的变化更敏感。这就是为什么我们重复使用人眼最敏感的Y分量（亮度），并预测Cb和Cr，因为我们不能很好地注意到颜色的虚假性，所以我们可能会犯错误。这一特性在彩色电视出现之初就得到了广泛的应用，当时的频道容量还不足以传输所有的颜色。图像在YCbCr中传输，不改变y分量，Cb和Cr减少一半。

如何创建基线

我们可以采用带预训练编码器的Unet，并将现有CbCr值和预测值之间的l1损失最小化。我们想给人像上色，因此，除了openimages照片，我们还需要更多特定于任务的照片。
我们在哪里可以得到穿着军装的人的彩色照片？互联网上有人把旧照片上色当作一种爱好或是一种代价。他们做得很仔细，尽量精确。当他们给制服、肩章和奖章上色时，他们会参考档案材料，因此他们的工作成果是值得信赖的。总而言之，我们使用了200张手工着色的照片，上面有穿着军装的人。
另一个有用的数据来源是工农红军网站。它的一个创立者几乎每到第二次世界大战时都会穿着苏维埃的制服拍照。

在一些照片中，他模仿著名的档案照片中人们的姿势。他的照片有白色背景是一件好事：它让我们可以通过在背景中添加各种自然物体来很好地扩充数据。我们还使用了一些常规的肖像，以徽章和其他战时特征作为补充。
我们训练了Albunet-50——它是一个使用预先训练的Resnet-50作为编码器的Unet。网络开始提供足够的结果：皮肤是粉红色的，眼睛是灰蓝色的，肩膀是黄色的。然而，问题是它让照片上的一些区域没有受到影响。这是因为根据错误L1，在不做任何事情比试图预测某种颜色更好的地方找到这样的最佳值。

我们正在将我们的结果与地面真实照片进行比较-这是由Klimbim完成的手动着色。
我们怎样才能解决这个问题？我们需要一个鉴别器：一个接收图像并告诉我们图像是否真实的神经网络。下面的一张图片是手动着色的，另一张是由我们的生成器Albunet-50着色的。人类如何区分手动和自动彩色照片？通过观察细节。你能告诉我们的基线解决方案后自动着色的照片在哪张图吗？

答：左边的图片是手动着色，右边的是自动着色。
我们使用了来自自聚焦Self-Attention GAN鉴别器。这是一个小型的卷积网络，所谓的自我关注建立在顶层。它允许我们“更加关注”图像细节，我们也使用光谱标准化。你可以在别的文章中找到更多的信息。我们用L1损失和鉴别器损失的组合来训练网络。现在，网络使图像的细节颜色更好，背景看起来更加一致。再举一个例子：左边是只训练了L1损耗的网络工作；右边是L1鉴别器损耗的组合。

在四台GeForce 1080ti上训练过程花了两天时间。处理一张512 x 512的图片需要30毫秒的时间。Validation MSE — 34.4。就像修复一样，度量你不想依赖于度量。这就是为什么我们选择了六个具有最佳验证指标的模型，并盲投选出了最佳模型。
当我们已经创建了一个制作系统并发布了一个网站时，我们继续进行实验并得出结论，我们最好不是将每像素L1的损失降到最低，而是将感知损失降到最低。为了计算它，我们将网络预测和地面实况照片输入VGG-16网络，获取底层特征图并与MSE进行比较。这种方法绘制了更多的区域，并给出了更丰富多彩的结果。

扼要重述

Unet是一个非常酷的模型。在第一个分割任务中，我们在训练过程中遇到了一个问题，我们使用高分辨率图像，这就是我们使用就地BatchNorm的原因。在我们的第二个任务（修复）中，我们使用了部分卷积而不是默认卷积，这使我们能够获得更好的结果。在进行彩色化时，我们添加了一个小的鉴别器网络，该网络对不真实图像的生成器进行惩罚。我们还使用了知觉损失。
第二个结论-评估师是必不可少的。不仅在创建分段掩码阶段，而且在最终结果验证阶段。最后，我们给用户三张图片：一张有修补缺陷的原始图片，一张有修补缺陷的彩色图片和一张简单的彩色图片，以防缺陷搜索和修补算法出错。
我们从战争相册项目中拍摄了一些照片，并通过这些神经元进行处理。以下是我们得到的结果：

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