使用条件随机场（CRF）来提升图像分割的表现

点击上方“AI公园”，关注公众号，选择加“星标“或“置顶”

---

作者：Ihor Shylo

编译：ronghuaiyang

导读

一个基于TensorFlow的CRF用法和实现的简单介绍。

在一个理论上计算能力无限的现代世界，语义图像分割已经成为许多应用的关键方法，如自动驾驶、高级医学图像分析、目标检测和许多其他应用。通常，一个基本的U-Net神经网络在大多数时候都可以得到良好的结果。Dice系数是一个流行的图像分割度量。然而，在仔细检查预测mask之后，发现了错误预测像素的小的“孤岛”。因此，问题出现了：如何改进这些微小的不一致？

条件随机场也被称为CRF，它经常被用作后处理工具来提高算法的性能。然而，在推断过程中，这种操作的计算成本可能很高，特别是在移动设备上。它使用了一组需要硬编码的参数，这使得它很难适用于整个测试集。这个问题的一个可能的解决方案是以循环神经网络(RNN)的形式添加一个CRF算法作为神经网络的附加层，并使其具有可训练性。虽然有很多关于这种方法的科学论文，但在大多数深度学习框架中还没有开箱即用的CRF-RNN实现。因此，这篇博文的主要目的是演示如何使用Tensorflow将这个额外的层插入到原来的U-Net模型中。

用例描述

用例是对文档(特别是收据)的像素级检测。常规检测无法捕捉到收据形状的变形。图1展示了其中一个标注图像，它是测试集的一部分。

图1：用例的真实标注

本质上，它是像素级的二元分类，其中0类定义为图像的背景，1类定义为文档。

只使用 U-Net 方法

数据集分为3个子集：训练集、开发集和测试集。后者用于模型的最终验证。用于图像像素分类的模型就是所谓的U-Net模型。

图2：典型的U-Net结构

它由两个主要部分组成：编码器和解码器。该技术的实现是受到Tensorflow图像分割的例子的启发。为了保持简短，可以在图3查看模型的大致结构。

图3：U-Net模型的总体结构

在通过网络之前，图像被缩放为224x224并归一化。请注意softmax层被省略了。在500张图像上训练30个epoch已经显示出相当不错的测试数据集的结果：Dice coefficient 0.983。图4中可以看到一个预测mask的例子。

图4：U-Net预测Mask

值得注意的是，正如前面所描述的，有一些用蓝色圈出的错误分类的小“岛”。这就是CRF-RNN层会派上用场的地方。在开始之前，有必要注意的是，由于特征已经训练好了，在添加新的CRF-RNN层之前，网络的权重应该被固定和设置为不可训练。

CRF-RNN

一旦你训练好了特征，就可以开始加入CRF-RNN层，并再次训练网络。不幸的是，Tensorflow中没有一个预训练的CRF层。经过一番彻底的搜索，我偶然发现了Sadeep Jayasumana的GitHub仓库。他为Keras创建了一个自定义类，并将其公开。唯一的限制是批大小必须为1，这使得训练有点慢。然而，考虑到特征已经训练过的事实，这个限制似乎不是什么大问题。第一步，克隆git仓库并遵循安装说明。一旦完成，你就可以将定制的CRFRNNLayer加到你的网络上。下面的代码片段演示了它是如何在我们的用例中实现的：

from crfrnn_layer import CrfRnnLayer

def add_crf_layer(original_model):
    original_model.trainable = False

    crf_layer = CrfRnnLayer(image_dims=(224, 224),
                            num_classes=2,
                            theta_alpha=3.,
                            theta_beta=160.,
                            theta_gamma=3.,
                            num_iterations=10,
                            name='crfrnn')([original_model.outputs[0], original_model.inputs[0]])

    new_crf_model = tf.keras.Model(inputs = original_model.input, outputs = crf_layer)

    return(new_crf_model)

有几个参数需要指定，如images_dims和迭代次数。第一个需要匹配特征提取中最后一层的输出维度。迭代次数是任意参数。对于所有其他的，它们服从超参数优化。新模型编译完成后，模型摘要如下：

图5：自定义CRF-RNN Layer模型结构

作为最后一步，在重新训练模型时，需要将EPOCHS参数设置为1，因为已经在自定义层中指定了迭代次数。

在通过网络并在测试集上运行验证之后，我们观察到性能指标略有增加，显示Dice coefficient为0.9857。更重要的是，那些被错误分类的“小岛”消失了：

图6：CNN CRF-RNN Mask预测

现在你的解决方案离生产部署更近了一步！

总结

综上所述，Tensorflow中还没有实现开箱即用的CRF-RNN层。然而，感谢开源社区，有一个自定义的实现。要做到这一点，需要遵循以下步骤：

1. 使用cnn训练你的特征。

2. 使特征无法训练。

3. 插入自定义CRF-RNN层。

4. 重新训练网络。

5. 使用新模型进行推理。

还有几个要点需要进一步改进该层。首先，目前批量大小只有1的可能性。这是可以改进的，但是需要对层进行一些重构。其次，这个自定义方法不适用于Tensorflow Lite，因为这个操作符还没有在其中注册。也可以使用定制的内核，但这需要Tensorflow核心Lite库中的一系列c++实现。

—END—

英文原文：https://medium.com/@ihor.shylo/improving-performance-of-image-segmentation-with-conditional-random-fields-crf-8b93f7db396c

请长按或扫描二维码关注本公众号

喜欢的话，请给我个在看吧！