你好,各位!我是一位来自Mail.ru Group计算机视觉团队的研发工程师。在这篇文章中,我将向大家介绍我们如何为老军人的照片创建一个基于AI技术的照片修复项目。照片修复包括三个步骤:发现照片缺陷、进行图像修复和上色。接下来,我将详细解释每个步骤,并告诉大家我们是如何获取数据、使用哪些网络进行训练以及取得了哪些成果和犯了哪些错误。
首先,我们需要在一张已经上传的照片中找到所有与缺陷相关的像素值。为了做到这一点,我们需要了解人们会上传什么样的照片。我们与immortal regiment项目的创始人交流过,他们是一个非商业组织,专门存储二战遗留照片。他们向我们分享了数据,使我们得以分析。通过分析,我们意识到人们上传的大部分个人或集体肖像存在中等到大量的缺陷。
接下来,我们需要收集一个训练集,对图像进行分割,得到一张图片和一张标注出所有缺陷的黑底白纹蒙版。最简单的方法是让评估员创建分割蒙版。当然,普通人通常知道如何找到缺陷,但这需要很长时间。标记一张照片中缺陷处的像素可能需要一个小时或一整天,因此在几周内收集一个超过100张图片的训练集并不容易。这就是为什么我们尝试增强我们的数据,并创造我们自己的缺陷:拍一张照片,用图片上的随机纹理增加缺陷,最终得到一张显示有缺陷的图像部分的蒙版。在没有增强的情况下,我们已经获得了含有68张手动标记照片的训练集和11张照片的验证集。
第二个问题是如果我们使用可以立刻上手使用的预训练编码器和UNET,将会损失许多位置数据。为了解决这个问题,我们在第1层使用了一个核为5、步长为2的卷积,使网络运行得更快。我们牺牲了网络运行的时间,以换取更好的缺陷定位。在第1层之后,我们移除了最大的pooling层,并将卷积步长减至1,卷积核减至3。
在图像处理中,我们通常需要对图像进行压缩以减少内存占用。将小图片压缩至256或512像素可以消除瑕疵,但对于大图像,我们需要对其进行特殊处理。实际应用中,我们通常将1024像素的照片分割成瑕疵部分。因此,网络需要训练以适应处理大图像。然而,这会导致在单个GPU上处理规格小批量的问题。
在训练过程中,我们可以在一个GPU上匹配约20张图片。这样一来,最终在BatchNorm层得到的平均值和标准偏差值可能不够准确。为了解决这个问题,我们可以使用In-place BatchNorm技术。这种方法既节省了内存空间,又提供了一个同步版本的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的使用,以使受损部分的边缘不再那么容易被看见。
您可以使用以下步骤新建一个基准模型:
1. 选择一个预训练的模型作为基准模型。例如,您可以使用在ImageNet上预训练的ResNet50模型作为基准模型。
2. 将您的数据集分为两个部分:训练集和验证集。训练集用于训练基准模型,验证集用于评估基准模型的性能。
3. 使用训练集训练基准模型。
4. 使用验证集评估基准模型的性能。如果您的指标(例如PSNR)没有达到预期,您可以尝试调整超参数或更换其他网络结构。
我们将现有的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项目中获得了一些照片,并在这些神经网络上进行了处理。以下是你得到的结果:此外,你还可以更详细地查看原始图像和所有的处理过程。点击【这些老照片如何用算法修复?】即可访问相关内容或参考文献!