自拍抠图抠到手软？详解如何用深度学习消除背景

2017 年 9 月 12 日 AI研习社 AI 科技评论

这篇文章介绍了作者 Gidi Sheperber 在 greenScreen.AI 项目中所做的工作和研究，AI研习社对本文做了详细介绍和编译。

在机器学习方兴未艾的过去几年里，我一直想要亲自开发具有实用价值且基于机器学习的产品。然后几个月前，在我学习了由 Fast.AI 所提供的深度学习课程之后，我意识到机会来了。当前的机遇是：深度学习的技术优势使得许多之前不能完成的事情变得可能，并且还有许多新兴工具被开发出来，这使得深度学习的部署过程变得更加简单。

在上述课程中，我遇到了 Alon Burg——一名经验丰富的 web 开发人员，我们刚好趣味相投。所以为了实现这么一款产品，我们为自己设定了一些目标：

提高自身的深度学习技巧
提高自身的 AI 产品部署技巧
针对市场需求，开发具有实用价值的产品
产品要有趣
要分享我们的经验

综合考虑以上的目标，我们进行了头脑风暴：

哪些东西是没被做过的（或者说还不够完善）
产品的设计以及实现难度不能太大——因为我们的计划是花上 2-3 个月的时间，每周只在这上边花个 1 天时间
用户接口要足够简单，容易上手——因为我们希望用户能够轻松学会如何使用这个产品，毕竟该产品开发的初衷不是为了科学验证
用于训练深度学习模型的数据要容易获取——因为有机器学习经验的人基本都知道，有时数据比算法来的更加重要
将采用尖端的深度学习技巧（这种技巧还未被 Google、Amazon 等公司投入商业化应用），但是也不能太过于新颖（以便于我们可以在网上找到参考例子）
产品将具备有生产的潜力

我们最早的想法是从医疗项目入手，因为医疗领域与我们的初衷很接近，并且那时认为（现在依然认为）深度学习在医疗领域还存在着大量机会。但是，我们意识到这将违背有关于数据收集那一条，数据要容易获取的原则。因此我们退而求其次，选择了做背景消除（Background removal）。

背景消除这项任务如果由人类纯手工来做，或者借助诸如 Photoshop、Power Point 等工具来实现都是非常简单的。然而，全自动化的背景消除却是一项充满挑战的任务。并且根据我们的了解，目前还没有一款产品能够在背景消除任务中实现非常好的效果，不过有些厂商已经进行了尝试。

那么问题来了，我们需要消除什么样的背景？这是一个很重要的问题。因为一个模型越是具体（比如说对物体、角度等方面有详细规定），它最后的分割效果也越好。在我们刚开始时，我们将这个范围定义得比较广：一个通用的背景消除器将能够从任何类型的图片中自动识别出前景与背景。但是在我们训练出了第一个模型之后，我们意识到如果将精力集中于一组特定的图像上，将取得更好的效果。因此，我们决定专注于肖像以及自拍类型的照片。

背景消除样例

自拍的图像具有突出和集中的前景（一个或者多个 “人”），这样使得背景和物体（脸部和上半身）间能实现较好的分割，并且能保持一个相对稳定的角度。

在做了这些基本假设之后，我们开始了理论研究、代码实现和长达数小时模型训练的旅程，以创建能够一键轻松消除背景的服务。其中最重要的一部分工作是训练模型，但是也不能低估了正确部署模型的重要性。另外当前最好的分割模型还无法像分类模型（比如 SqueezeNet）那么紧凑。并且我们还主动检查了服务器和浏览器的部署选项。

如果你想了解更多关于我们产品的部署过程，建议你去阅读这两篇文章，它们分别从服务器端和客户端的角度介绍了该项目。

如果你只是想了解模型和它的训练过程，那么请接着阅读吧。

语义分割

当我们开始认真思考深度学习和计算机视觉中各项技术的时候，显而易见的，最适于用来实现背景消除任务的技术是语义分割。

当然也有其它的策略，比如说通过深度检测器实现分割，但是这个对于我们而言似乎还不够成熟。

语义分割（Semantic segmentation）是一项众所周知计算机视觉任务，与图像分类和物体检测并列为计算机视觉领域的三大挑战。分割的过程可以视为将图像中的每个像素点分类到图像中某一个物体类别上，所以分割实际上也是分类任务。与图像分类或检测不同，语义分割真正展现出了对图像 “理解”。因为它并不只是简单地指明 “图像中有只猫”，而且还在像素的级别上直接指出了图中的猫在哪个位置，以及猫所占的范围。

那么分割究竟是如何实现的呢？为了更好的理解这一点，我们将回顾一下该领域的一些早期工作。

最早的想法是采用一些早期的分类网络，比如 VGG 和 Alexnet。VGG 是 2014 年用于处理图像分类任务最先进的模型，并且由于它简单明了的网络结构，使得它在今天依然非常有用。在 VGG 较靠前的网络层中，在要进行分类的物体周围会具有较高的激活，并且在更深的层次中将具有更强的激活。但是由于重复的池化（Pooling）操作使得这些激活的性质相对粗糙。有了这点基本理解之后，我们可以假设分类训练模型在进行一些调整之后也能被应用于寻找或者分割图像中的物体。

语义分割的早期结果与分类算法一起出现。在这篇文章中，你将看到一些通过 VGG 实现的分割结果，但是结果比较粗糙。

较后边网络的结果：

校车的分割结果，浅紫色（29）代表了校车类别

经过双线性上采样之后：

这些结果仅来自于将全连接层转换为原始形状，而保持了原本的空间特征。在上边所展示的实例中，我们输入一张 768*1024 的图像到 VGG 中，然后获得一层 24*32*1000，其中 24*32 是该图像池化后的版本，而 1000 则指的是 Image-net 的类别数量，从中可以导出分割结果。

为了实现平滑预测，研究员使用了一个简单的双线性上采样层。

在 FCN 的论文中，研究员进一步改进了上述的方法。他们连接了一些层次，以便于获取更多的特征解释，并且依据上采样率，分别被命名为 FCN-32、FCN-16 和 FCN-8。

在层之间添加一些跳跃连接允许预测器从原始图像获取更精细的细节。进一步的训练更能提高结果。

这种技术本身并不像以前所想那么糟糕，并且证明了深度学习在语义分割领域的潜力。

FCN 的处理结果，图像来自论文

FCN 为语义分割打开了一扇通往新世界的大门，研究员为此任务尝试了不同的网络结构。但是核心思想保持不变：使用已知的结构、上采样和网络层间跳跃连接。这些技巧在较新的模型中依然常见。

你可以在这几个贴子中进一步获取语义分割领域的发展情况：

http://blog.qure.ai/notes/semantic-segmentation-deep-learning-review
https://blog.athelas.com/a-brief-history-of-cnns-in-image-segmentation-from-r-cnn-to-mask-r-cnn-34ea83205de4
-https://meetshah1995.github.io/semantic-segmentation/deep-learning/pytorch/visdom/2017/06/01/semantic-segmentation-over-the-years.html

另外你可能也发现了，大多数语义分割方法都保持了编码 - 解码（Encoder-decoder）的架构模式。

回到项目

在一番调研之后，我们将目光聚集在了三个模型上：FCN、Unet。其中 Tiramisu 采用了非常深的编码 - 解码架构。我们也对 Mask-RCNN 有些想法，但是实现它似乎已经超出了本项目的范围。

FCN 由于效果太差，首先被我们舍弃。而另外两个模型的结果还不错：Tiramisu 和 Unet 的主要优势在于模型紧凑和计算快速。就实现方面而言，Unet 非常容易实现（采用 Keras）而 Tiramisu 也可以实现。并且我们也使用了来自 Jeremy Howard's 的 Deep learning 课程中最后一堂课所提供的 Tiramisu 的实现代码。

保留着两个候选的模型，我们进入到了下一步——开始训练模型。必须说明的一点是，当我第一次尝试了 Tiramisu 模型之后，就被它惊人的结果所深深捕获，因为它能够捕获到图像中的尖锐边缘。另一方面，Unet 似乎还不够好，它的结果也是差强人意。

数据

既然在模型选择方面已经有个大概方向，那么接下来我们将开始寻找合适的训练数据。用于训练分割模型的数据并不像分类或者检测那样常见。事实上，图像分割最常用的几个数据集分别是 COCO，它包含有 90 种类别，大概有 8 万张图像；VOC pascal，它包含有 20 种类别和 1.1 万张图像；以及最近发布的 ADE20K。

最终我们选择采用 COCO 数据集，因为它包含有许多关于 “人” 这一类别的图像，而这正是我们所感兴趣的。

考虑到产品所要实现的任务，我们还需要考虑是只使用与任务最相关的那一部分数据集，还是使用更多的更一般的数据？一方面，使用具有更多图像和类别的更一般性的数据集，将使得模型具备有处理更多场景和更具挑战性的图像；而另一方面，一个通宵的时间我们能够训练 15 万张图像。因此如果我们将完整的 COCO 数据集都用于训练模型，则一个通宵的时间下来，平均而言，每张图像将被使用 2 次左右。所以采用更加精简的部分数据将对模型训练有益，此外，它将使得模型更加集中。

还有一件值得一提的事是，Tiramisu 模型原本是在 CamVid 数据集上进行训练的，该数据具有一些缺陷，但最重要的是它的图像内容非常单调——所有的图像都是道路与汽车。相信你可以很容易就理解，从这样的数据集中进行学习（即便它包含有人）对我们的最终目的也没有任何好处，所以经过短暂的考虑，我们选择了放弃而改用 COCO。

来自 CamVid 数据集的样例

COCO 数据集附有非常简单明了的 API，这可以让我们准确地知道每张图像上都包含了哪些物体（根据 90 个预定义的类别）。

经过一番实验，我们决定精简数据集。首先，我们选出了那些包含有人物的图像，这样的图像有 4 万张。然后，我们继续将那些包含有太多人物的图像剔除，剩余的图片中只包含有 1 到 2 个人，因为这样的图像与我们的产品目标最契合。最终，我们只留下那些图中有 20%~70% 的区域被标记为人的图像，而继续移除了那些背景中人物过小，或者是存在某些怪异东西的图片。最终，精简后的数据集包含有 1.1 万张图像，而我们认为这个数据量在现阶段已经足够了。

左边：好样例，中间：包含有太多元素，右边：目标主体太小

Tiramisu 模型

如上所述，我们在 Jeremy Howard 的课程中学习到了 Tiramisu 模型。虽然它的全名 “100 层 Tiramisu” 可能暗示了它是一个巨大的模型，但事实上它非常地经济，甚至只有 900 万个参数，而 VGG16 则拥有超过 1.3 亿个参数。

Tiramisu 模型是参考了 DensNet 的设计而提出的。而 DensNet 是一个最近提出的图像分类模型，在这个模型中所有层都相互连接。此外，Tiramisu 还在上采样层添加了跳跃连接（Skip connections），这一点和 Unet 一样。

如果你还记得，这个架构是符合 FCN 所提出的想法的：使用分类架构、上采样和跳跃连接进行细化。

Tiramisu 的通用架构

DenseNet 模型可以看作是 ResNet 模型的一个自然演变，但是 DenseNet 不是 “记住” 层与层之间的关系，而是记住了模型中的所有层。这种连接被称为高速公路连接（Highway connections）。它导致了过滤器数量激增，这被定义为 “增长率”。Tiramisu 的增长率为 16，因此在每一层我们需要添加 16 个新的滤波器，直到达到某一层具有 1072 个滤波器。你也许会期望 1600 层，因为模型的名字是 “100 层 Tiramisu”，然而上采样层降低了一些滤波器。

Densenet 模型草图，早期的滤波器堆叠在整个模型中

训练

我们按照原始论文所述的方式对模型进行了训练：标准的交叉熵、采用 1e-3 学习率的 RMSProp 优化器和较小的衰减值。我们将精简后的 1.1 万张图像分成了 70% 的训练集和 20% 的验证集，以及 10% 的测试集。接下来展示的所有图片均来自测试集。

为了保证我们的训练过程与论文中的保持一致，我将 Epoch size 设置为每 500 张图像，这使得我们能够随着每个结果的改进而定期保存模型，因为我们对更多的数据进行了训练（原论文使用的 CamVid 数据集包含的图像少于 1 千）。

此外，我们只在两种类别上进行训练：背景与人，而论文中则使用了 12 种类别。我们一开始尝试在 COCO 的类别上进行训练，但是发现这对训练并没有太大帮助。

数据问题

一些数据集的缺陷阻碍了模型的表现能力：

动物——我们的模型有时候分割动物，这将导致较低的 IoU。在我们任务中往主要类别增加动物进去或者是其它东西将导致结果变差。
躯干——由于我们是通过程序自动化地过滤了数据集，所以无法判断一张包含有人物的图片是真的有一个人，还是说仅仅只有一些躯干，比如手或者脚。这些图像并不在我们的目标范围内，但是它们还是出现在了数据集中。

动物，躯干，手持物体

手持物体——数据集中的许多图像都是与运动有关的，图中人物总是伴随着棒球棒、网球拍和滑雪板等出现。我们的模型在某种程度上混淆了应该如何去区分图中的人和物。与在动物的案例中一样，将它们归为主体的一部分或则区分开来都将有助于提高模型的表现力。
粗糙的真实数据（Ground truth）——COCO 数据集不是按像素进行逐个标记，而是使用了多边形进行标注。有些时候，这种程度的标注已经足够了，但是有些时候，这样的真实数据还是过于粗糙，从而阻碍了模型的学习。

原图以及粗糙的真实数据

结果

我们的结果非常令人满意，虽然还有些美中不足：我们在测试集上取得 IoU 是 84.6，而目前最佳的成绩则是 85。这个数据（IoU）也是棘手的，因为它在不同的数据集和类别中会发生波动。有些类别本身就更容易进行分割，例如房子、道路等，许多模型都能很容易取得 90 的 IoU。另一些更加具有挑战性的类别有树和人，在这些类别的分割上大多数模型只能达到 60 的 IoU。为了衡量这一困难，我们帮助网络专注于单一的类别和有限类型的照片。

尽管我们仍然觉得目前的成果尚不足以达到 “可以投入使用（Production ready）” 的标准要求，但是现在是一个暂停下来好好讨论结果的好时机，因为大约有 50% 的照片都给出了不错的结果。

以下是一些比较好的例子。

从左到右分别是：图像，真实数据，输出结果（来自测试集）

调试与记录

训练神经网络的很重要一部分工作就是进行调试。开始一个神经网络的训练非常容易，只要获取数据输入网络，然后就开始进行训练，最终再看看输出结果如何。但是我们发现，跟踪网络训练的每一步都是非常重要的，所以我们为自己制作了相应的工具，以便于能在每个步骤中检查结果。

以下是一些常见的挑战，以及我们所采取的措施：

早期问题——该模型可能无法训练。这可能是因为一些固有的问题，或者由于某种预处理错误，例如忘记归一化数据块导致的。无论如何，将结果可视化出来将非常有帮助。这有一个与该主题有关的贴子。
调试网络——在确定没有重大问题之后，网络将以预定义的损失开始训练。在语义分割领域主要的衡量指标是 IoU——交叠率（Intersect over union）。我们花了一段时间才开始采用 IoU 作为模型训练的主要指标（而不是交叉熵）。另一个有用的做法是在每一个时期（Epoch）展示一些模型的预测结果。这个帖子很好地介绍了如何调试机器学习模型。另外请注意 IoU 在 keras 中并不是标准的损失函数，但是可以很容易在网上找到现成的，比如这个。我们也利用这个 Gist 来绘制每个时期的损失和预测结果。
机器学习的版本控制——通常一个模型都会具有许多参数，有些参数比较难以跟踪。我必须说明我们还没找到完美的方法解决这个问题，除了比较频繁地写出配置信息（并使用 keras 回调自动保存最佳模型，参见下文）。
调试工具——在完成了以上所有措施后，我们将能够在每一个步骤检验我们工作，但是依然无法做到无缝检测。因此，最重要的一步是将上述步骤结合起来，并创建一个 Jupyter 笔记本，我么能够借助它实现无缝地加载每个模型和每个图像，并且能够快速检测它的结果。这样我们可以很容易发现模型之间的差异、模型的缺陷以及其它问题。

以下是我们模型的改进示例，以及参数调整和额外的训练。

保存在验证集上取得最好 IoU 成绩的模型（Keras 提供了一个非常棒的回调函数，使得这件事变得非常简单）：

callbacks = [keras.callbacks.ModelCheckpoint(hist_model, verbose=1,save_best_only =True, monitor= ’val_IOU_calc_loss’), plot_losses]

除了可能的代码错误的正常调试之外，我们还注意到，模型的错误是 “可预测的”。如 “切割” 身体部分超出了正常的躯干范围，没必要的躯干延伸，光照不足，照片质量低和照片中细节过多等。其中一些在添加不同数据集中特定图像时被处理掉了，但是其它的一些则依然是一项还有待处理的挑战。为了在下个版本能够提升结果，我们将使用专门针对 “硬” 图像的扩充进行训练模型。

数据集问题我们在之前已经谈到过了。现在来看看模型面临的一些挑战：

衣服——非常亮或者非常黑的衣服有时会被当成背景
“挖洞”——有些本应该不错的结果却有出现类似于被挖了一个洞情况

衣服和 “挖洞”

3. 灯光——虽然光照不足和图像模糊的情况在照片中很常见，但是 COCO 数据集却不常见这类情况。因此，除了使用标准难度的模型来处理之外，我们还没有为那些富有挑战的图像做好准备。这一问题可以通过获取更多数据，或者增强数据（Data augmentation）来解决。

光照不足的样例

未来

进一步训练

我们的训练数据训练了 300 多个时期，在这之后，模型开始变得过拟合。由于我们的结果已经非常接近于发布的成绩，因此我们没有机会应用数据增强的基本做法。

训练所用的输入图像统一被调整为 224*224 大小。更进一步地，采用更多的数据和更大分辨率的图像（COCO 图像的原始大小约为 600*1000）进行训练模型也对提高最终的结果有帮助。

CRF 和其它增强

在某些阶段，我们发现有些结果的边缘存在噪点。可以采用 CRF 模型来改善这一问题。在这个博客中，作者展示了 CRF 的简单使用样例。

但是该方法对我们的工作来说还不是很有帮助，也许是因为只有当结果比较粗糙的时候该方法才会有所帮助。

抠图（Matting）

即使我们目前取得了这样的结果，但是在实际的分割中还是不够完美。诸如头发、细腻的衣服、树枝和其它精美的物品永远无法被完美分割。事实上，这种非常精细的分割任务被称为抠图（Matting），它定义了一种不同的挑战。这里展示了一个当前最先进的抠图样例，该项工作是在今年早些时候的 NVIDIA 大会上发表的。

抠图样例，输入包含了 Trimap

抠图任务与其它图像相关的任务不太一样，因为它的输入中不仅有原始图像还有三分图（Trimap）。三分图指的是图像边缘的轮廓，这也使得它成了 “半监督” 问题。

我们使用抠图进行了一些实验，就是将我们的分割图像作为三分图使用，但是却没有取得任何有意义的结果。

还有一个问题就是缺乏使用的训练数据集。

总结

正如一开始所说的，我们的目标是开发出一款具有意义的深度学习产品。正如你可以在 Alon 的贴子里看到的那样，部署相对而言更加容易和快速。另一方面，模型的训练才是最棘手的——特别是需要花上一晚的时间进行训练，此时你就需要进行仔细规划、调试和记录结果。

事实证明，要在研究和尝试新事物以及训练和改进模型之间取得平衡是一件不容易的事情。由于使用了深度学习，我们总是觉得最好的模型或者最适合我们的模型就躲在某个角落，只需要多进行一次 Google 搜索或者多阅读一篇论文就将会引导我们发现它。但是在实践中，我们实际的改进来自于从原始的模型中一点点地 “挤压”。而且如上所述，我们仍然觉得有更多的地方可以继续改进。

总而言之，我们在这个工作中收获了许多乐趣，这几个月的日子过得就像科幻小说一般。我们很乐意同大家讨论和回答任何问题，期待在我们的网站上遇见您。

Via Background removal with deep learning