A Deep Learning Design for improving Topology Coherence in Blood Vessel Segmentation
改善血管分割中拓扑一致性的深度学习设计

from MICCAI2019

Abstract

医学图像中广泛研究的血管分割任务，对许多临床应用都产生了重大影响。在可以获得手动标注信息时许多监督学习方法应用在这上面，主要是通过交叉熵损失函数来进行训练，但这样没有对拓扑结构的错误进行惩罚。比如血管树缺失某一部分或者某一分支不连贯，这样可能会导致网络的拓扑结构严重改变，影响血管通路以及血管特征的提取。

本文就提出一种端到端的设计，通过将一个传统的分割网络与一个变分自编码器级联，通过学习丰富且紧凑的隐空间特征，来修正这种拓扑结构的不一致性。本文在3类经典的视网膜血管数据集上进行了测试，分别是DRIVE，CHASEDB1,STARE，验证了本文的方法可以在不损伤分割像素级精度的前提下分割出更好的拓扑结构。

Section I Introduction

血管分割在临床诊断中有重要作用，如动脉粥样硬化和动脉瘤等疾病的真算以及评估器官移植手术情况等。但是分析这些图像需要耗费医生大量的时间和精力，也需要大量的数据，因此借助计算机进行自动分析方面开战了广泛研究并且取得了SOTA的结果。随着一些带标注数据的数据库的开源，出现了基于监督学习的分支。随着深度学习的发展，监督学习进行血管分割的效果又更上一层楼。

但是这些系统并不是完美无缺的，毛细血管就常常被遗漏，以及分割出的血管树的拓扑结构常常出错，比如由于中央反射效果等使得一些血管断裂或者出现分段，具体例子参考Fig 1.

这些误分割可能会导致血管通路的或者血管特征的提取出错。目前一些SOTA的方法容易出现这种拓扑错误，因为他们使用的是交叉熵损失函数，这种只惩罚像素级的分类错误，并不会意识到拓扑结构上的不一致。

近年来，已经有工作意识到应该将拓扑结构包含进不同的深度学习任务中，比如[2][3][4]以及由于信息的丢失，这种拓扑结构的不一致性不仅出现在分支中还有可能出现在分支的中间。


因此本文就考虑能否在一个典型的分割网络后增加一个经理按步骤，寻找一种端到端的设计来从这些噪声和拓扑不一致中学习到真正的血管拓扑结构。为此本文在一个常规的分割网络后级联了一个变分自编码器（VAE），通过这个VAE接受分割网络的输出，并重建真正的GT，通过学习一个隐空间的映射，最终达到拓扑结构一致的分割效果。

Section II Methodology

本节将详细介绍这种端到端深度神经网络的设计，包括一个典型的分割网络以及一个精炼模型，这个模型主要负责从已经破坏了的数据中学习一些有意义的特征，还会讨论应该如何设计能够减少这种拓扑结构的失误。
其中x和y分别代表输入的灰度图和GT血管分割图（分割掩膜mask),而y’和y’‘分别代表分割网络的输出和精炼模块的输出。



Part A Auto-encoding for learning local topology




我们可以将分割网络的错误看做是影响真实血管信号y的一种隐藏噪声，那么我们就希望y’分割网络不能建模这种噪声，而能够让y’'来更好的描述y（真实的GT）。




通常自编码器就会学习一种将整幅图到一个特征向量z的映射，这样就可以学习到复杂的较大尺度上空间上的交互。但对于血管这种会反复出现的图形结构这不是最合适的编码方式。对于这种数据最好是像[6]中的那样，通过输入一个3D的特征向量z,是一系列特征图谱的堆叠，来学习捕获空间信息。





下面让我们来考虑分割掩膜y是如何生成的。它需要从一个先验分布找那个采样然后根据条件分布产生掩膜，我们假设这些分布属于pθ(Z)和pθ(y|z)的参数簇。





对于给定的GTy我们希望通过：


来获得隐空间分布上的不同表述。
这就产生了一个问题：
如果要评估pθ(Y)的相似性需要在整个隐空间上进行计算。
而可以通过变分自编码器来规避这一问题，通过利用一簇qλ(Z)来逼近这一后验概率，其中最优参数就是使得两个分布之间KL散度最小的那个参数：


但是这也需要计算边界可能性，通过注意到D KL 是非负值因此将log缩放到：


这样就获得了ELBO这是可以求最大值的，因此可以用来做近似的推理。
变分自编码器VAE就是这种近似，两个分布都是通过神经网络参数化的，通常被称为识别模型和生成模型。这两个网络的权重通过随机梯度变分贝叶斯估计器进行学习。
Part B Refinement model as a Denoising VAE

将精炼模块作为去噪VAE
去噪VAE(DVAE)[7]将
 噪声建模为条件分布pγ(y0|y)。 本文中分割网络的输出y0被解释为真实信号y的损坏版本。因此在DVAE中识别模型定义为：

修改后近似为：

与训练常规的DVAE是等价的，主要就是从噪声中学习到有意义的特征，从而获得隐空间的表述，生成模型根据这一表述可以尽可能接近无损的数据。在Fig2中我们认为这样获得的最终的mask在拓扑结构上具有更好的一致性。



Section III Experiments and Discussion

UNet是医学图像分割经典的模型，得益于能够同时学习不同尺度的特征，因此本文中也选用UNet作为分割网络，而本文的网络主要与以下两种设定进行对比：


（1）单一UNet的分割结果；


（2）两个级联的UNet（DUNet）。


两种模型的损失函数分别为：

其中L1=BCEw，L2=FL,还测试了不同损失函数的影响，比如BCEw倾向于惩罚假阴性多于假阳性，而FL更关注于误分类的样例。


Part A Datasets



DRIVE



STARE



CHASEDB1



DRIVE数据集中包含7例早期视网膜糖尿病变的样例，而STARE数据集中一半是病理图像，而CHASEDB1数据集中的图像中央血管反射丰富。



本文DRIVE数据集即作为训练集也有作为测试集 = 20:2剩下两个数据集中后10张用于测试，其余作为训练样本。




Part B 评价矩阵





ROC、
AUC 
以及为了测试分割掩膜的拓扑一致性，我们采用[13]相近的方法，随机从GT中选择一条连接的路径，与mask中等价的路径进行对比，如果找不到对应的就认为是不可行的，或者如果二者长度差超过10%就认为是不一致的。我们从测试图像中选取1000条路径。





Part C 实施细节
DRIVE数据集划分：train:val:test = 15:5:20
L1 = BCEw
L2 = FL
patch = 64*64

其中refine step体现在 alpha随着epoch从1减小到0.3，每次减少5e-3,此时两l1=l2；随后出于训练稳定性时在focal loss中会将L1设为BCE，L2= FL。
total epoch = 150
图像均取绿色通道，采取的图像变换有：水平、垂直翻转、旋转以及强度的变换。

原始的UNet结构包含4层编解码结构，本文则使用2层是最佳的，可以达到最佳的AUC。随后本文将基础框架换成了Double UNet以及本文提出的结构，基本上参量在4M左右。
模式识别部分包含4层卷积结构，每层kernel = 3x3，padding = 1，通道数分别是64,64,256,256.
前三层都使用的最大池化，最后一层就卷积输出会参数华为对角阵编码到隐空间中100维的latent code。
对应的在解码网络逐渐通过3层4x4大小的转置卷积恢复分辨率，最后实现的是二分类。
VAE中则使用的是ReLU**函数，最后一层使用的是Sigmoid。

Part D Results and Discussion

Table I展示了不同设定下的结果对比，其中使用focal loss对AUC略有提升，说明更适合二分类，但是对最终mask的质量并没有进一步的提升，这也不足为奇，因为更多的关注与haid cases并不能保证会给予哪些导致拓扑出错的像素更高的权重。反而是使用BCEwloss对血管部分施加更多的权重，可以达到我们预料中的对Se和拓扑结构的提升。

可以看到本文的方法可以显著减少不可行路径的数量，主要在于找到了正确的拓扑结构，并且并没有损失像素精度，从AUC指标上可以看出。
而且鉴于总是需要在敏感性Se和特异性Sp之间进行折中因此选择他们作为比较往往是不全面的。
通过与同样复杂的DUNet对比我们可以得出结论：文本的方法学到了更好的特征表述从而保证了拓扑一致性。Fig3就展示了使用BCEw的一些可视化结果，从左至右依次是：原图，GT，UNet、DUNet和本文的分割结果。

Section IV Conclusion

本文提出了在分割网络后级联VAE的设计，从而提高了血管分割的拓扑一致性，实验说明了这种方法是通过预测更正确的路径以及更不可行的路径达到的，同时还不会损害分割网络的精度，与同样两个UNet级联的情况相比，在不进一步增加模型复杂度的前提下进一步提升了精度。 未来我们还将进一步研究对拓扑结构更敏感的可微损失函数，此外我们还认为本文的这种设计对血管的半监督分割也是十分有用的。