深度网络图像分割通俗指南

作者：Rohan Relan

译者：夏天

审校：主题曲

给定一张包含家具的照片，你可以让程序自动将家具与背景分开吗？

在这篇文章中，我将介绍如何使用当前最先进的深度学习来尝试解决这个问题。我不是机器学习的专家，所以我希望这个帖子对于其他希望使用这个强大新工具的非专家们有一定的帮助作用。

这个问题称为分割。也就是说，从这张图开始：

到这张图：

我们可以将这个遮罩应用到源图像上，获得没有背景的椅子。我们将使用一些工具来简化这个工作：

keras - 一个非常棒的用于创建神经网络的库。 Keras是像Tensorflow这样的较低级别库的前端，它能为用户处理构建神经网络过程中存在的大量繁琐的细节。。

U-Net - 用于图像分割的神经网络架构。 U-Net最初被设计用于生物医学图像分割（例如，在CT扫描中识别肺结节），但它也可用于分割常规2D图像。在下文将看到，即使没有大数据集，U-Net的强大功能也能让你大吃一惊。

Brine - 一个数据集管理器，可以利用该管理器轻松地共享和管理图像数据集。构建模型最令人讨厌的部分就是获取和选择用于训练模型的数据集。我创建了brine来轻松共享数据集，使之能应用在PyTorch/Keras模型上。我们将使用它来下载数据集并将其与Keras进行交互。

一个Github代码库 - Carvana图像遮罩挑战赛是Kaggle的一项赛事，它提出了类似的问题：将汽车从背景中扣出来。人们经常在Kaggle比赛中分享他们的解决方案，而在这个代码库中，有人分享了一个使用Keras和U-Net的解决方案。我们的目标是利用这个解决方案来解决我们当前这个家具分割问题。

一个数据集 - 这是一个朋友提供的数据集。请注意，它非常的小，只包含了97张椅子和相应的遮罩。一般来说我不会指望通过这么少的数据来做很多的工作（Carvana挑战赛中提供了数千个样本），但是让我们来看看最终到底可以做到何种程度吧。

代码

这里有一个jupyter笔记，其中包含了建立模型的所有代码。我将重点介绍其中最重要的部分，并解释它的原理。

第一步是安装数据集。由于它托管在Brine上，所以可以用一个简单的命令来实现：brine install rohan/chairs-with-masks。

下一步是加载数据集。可以通过Brine的函数load_dataset来执行此操作chairs = brine.load_dataset('rohan/chairs-with-masks')。

该数据集包含了97个样本，每个样本是图像及其遮罩。遮罩是一个只有两种颜色的图像，蓝色代表背景，红色代表前景。

数据集加载了，现在来加载U-Net网络。把“Kaggle-Carvana-Image-Masking-Challenge”代码库中的目录复制下来。导入这个网络，执行model = unet.get_unet_256() 。感谢 petrosgk 的工作，只需调用这一个函数即可返回一个Keras内置的U-Net网络。 Keras提供了model.summary()方法来查看网络的结构，虽然从中可以看到大量的信息，但最重要的是第一个和最后一个，它告诉了我们网络期望的输入和输出的形状。

我们可以看到输入的形状是(None, 256, 256, 3)，输出的形状是(None, 256, 256, 1)。元组的第一个元素是批量的大小，所以我们现在可以忽略它。这告诉我们，网络期望的输入是一批256x256的三通道图像，并将输出一批256x256个单通道遮罩。我们的遮罩也需要匹配这个形状。

下一步是准备样本，使之与网络一起使用。我们将为训练数据定义一个处理函数，在样本传给网络之前需应用于每个样本。

这里做了很多事情，我会一步一步进行解释。样本作为一个元组被传递进去，所以首先要进行解包。接下来的两行使用numpy切片来确保图像只有3个通道，如果有第四个alpha通道，则忽略。然后，使用cv2（OpenCV的python绑定）将遮罩转换为灰度图，这样，我们现在就有了一个单通道遮罩，这是网络所期望的。这里没有为两种颜色使用两个随机的灰度数字，而是强制使用fix_mask函数将掩码设置为0和255以表示背景和前景。然后，我们将图像和掩码的大小调整为256x256，以匹配网络指定的大小。

由于数据集较小，因此我们将使用数据扩充。数据扩充是指在训练期间在保留原始信息的基础上随机修改图像，人为地生成更多的数据。例如，旋转5度的椅子仍然是椅子，所以网络应该能够正确地识别出来。在代码中，我们使用petrosgk的Carvana示例中的3个函数来随机改变图像的色相、饱和度和值，并随机旋转和翻转图像。如果旋转或翻转了图像，则必须对遮罩执行相同的操作，以使遮罩与原始图像保持一致。

最后，我们通过将所有像素值除以255来对数据进行归一化操作，这样，所有的值都在0和1之间了。如果在此时打印image.shape，结果是256x256x3，这正是网络所需要的。尽管mask.shape是256x256，但网络是256x256x1，所以我们使用np.expand_dims来让遮罩匹配这个形状。最后，返回新的图像和遮罩。

在开始训练网络之前，还需要用一些样本进行验证。验证集不能用于训练，只能用于检查模型的性能。我们可以使用Brine的create_folds创建：

我们将样本集中的20个图像作为验证集，剩下的77个图像作为训练集。

我们还定义了用于样本验证的处理函数，这与用于训练的处理函数非常相似，除了在验证时不使用数据扩充之外。

最后，Brine返回一个可用于训练和校验集的图片生成器：

这些生成器将返回样品的批处理，以及之前定义的处理函数。它们可以被直接传给Keras的fit_generator方法来训练模型。

现在，准备开始训练模型了：

这些回调会改变Keras的训练过程。 EarlyStopping将在验证损失的改进停止后停止训练，ReduceLROnPlateau将降低学习率，ModelCheckpoint将对在验证集上表现最佳的模型版本进行保存。

最后，可以开始训练了：

Keras开始训练模型了，它将在数据集上运行多次（由于数据增加和混洗，每次运行的结果会略有不同），并输出训练和验证集的损失和DICE分数。在某些时候，运行会停止，可能是因为EarlyStopping回调，或是因为达到了100次迭代。

好了，就这些！我们训练了一个U-Net网络，并尝试在图像中把椅子分割出来。让我们来看看它的实际表现如何吧。

生成预测

要产生一些预测，首先使用model.load_weights('weights/best_weights.hdf5')来加载最佳模型的权重，然后在验证集中的某些图像上使用模型的predict方法：

以下是一些样本的结果：

在原始图像上覆盖了网络预测出来的遮罩。颜色越黄，则说明网络对该像素的预测具有更高的置信度。

虽然这对于生产系统来说还不够好，但网络已经学到了有关前景与背景、椅子与非椅子的一些内容。我个人觉得这个网络在没有进行预先训练，并且只有77张图片的情况下表现得已经非常不错了。如果有更大的数据集，网络的准确性可能会大大提高，或许就能用于生产系统了。

本文由北邮@爱可可-爱生活老师推荐，阿里云云栖社区组织翻译。

文章原标题：《A Non-Expert’s Guide to Image Segmentation Using Deep Neural Nets》

*推荐文章*