这篇文章是2017年MICCAI上的一篇文章,作者来自印度通用电气全球研究室。文章主要是将FCN网络改进并应用到医学上的图像分割。
文章的主要贡献是:定义了一个新的loss函数,将形状的影响引入到分割中;提出了两种数据增强的方式。
类似工作:16年,有人通过施加平滑性和拓扑结构作为先验概率将局部几何内容加入到分割中;3D分割方面,有人在2016年通过表面投影将多个FCN的输出进行合并,再经过条件随机场进行标记;还有人利用玻尔兹曼进行形状先验概率运算。
本文是采用了级联两个CNN来完成形状的先验概率计算,然后将FCN的输出通过自编码器投影到真实图像的流型上,并更新了loss函数的形式,提出了两种数据增强方法针对标记样本量小的图像分割。
方法:
最传统的FCN分割图像如U-net中定义最小化损失函数为:
其中k=1.2....N,Ik是图像,Sk是对应的分割掩膜Ground truth,是分割得到的结果,训练时最小化上式。
本文中,作者对损失函数进行改进,利用卷积自编码器进行图像形状的表示,以便更好地嵌入到FCN中。作者将Sk定义的有效形状组成的底层空间称为M,假设能够学习一个p维形状投影(E)的编码器和解码器,而且这个投影需要能将各种形状的S投影到M上。作者定义了
作为S在M上的有效投影,新定义的loss函数为:
上式第一项目的是通过最小化投影误差使得分割结果尽可能接近M,上式第二项是真实掩膜和预测掩膜的编码之间的距离,第三项目的是从学习到的M中尽量保持Ground truth。
在传统的FCN分割模型中,loss是基于欧几里得距离,模型参数需要预测一个从输入图像到高维的复杂变化,所以需要有足够的训练数据来实现。而本文的方法预测结果和Ground truth是编码表示完成,所以模型复杂度由AE承担,如下图所示。
模型架构:
为了实现上述的loss函数,文章设置了两个FCN的级联,一个用于分割,一个用于正则。如下图1所示。
分割:分割模型用的是U-net网络,结构如下图2:
形态正则化:这个网络的目的是使得不完整的分割图像掩膜符合训练数据的多样性。本文采用CAE完成这个工作,模型结构如下图3,编码器提供一个紧凑的隐层空间,使得解码器可以重构出完整图像。
正则化的网络中不含有skip layer,然后将分割模型和正则化模型按照图1组合,形成上述的loss函数。其中正则化网络需要进行数据增强。
数据增强:针对正则化网络,文章提出了两种正则化方法。一是随机污染图像,即选择任意的黑块对分割区域进行掩盖,并随机移动,如下图4.
二是中间值重构,将收敛之前的训练图像通过U-net预测,将不准确的预测图像重新作为输入到正则化的网络中。如下图5。
实验:
卷积网络选择relu作为激活函数,batch normalization进行正则化,loss函数中的λ设置为0.5左右,数据集是U-net模型的数据集,一共包含231副图像,作者用100张进行训练,剩余131张进行测试,结果表明模型优于U-net模型。
上图B是Ground truth,C是U-net的分割结果,D是本文的结果,可以看出来本文所提出的模型可以显著提升U-net结果。
作者用骰子系数对分割的结果进行量化比较,如下表:
SR-UNet_1是采用第一种数据增强方法的模型结果,SR-UNet_2是采用第二种数据增强方法的模型结果,可以看出相比较U-Net模型,本文模型的骰子系数提升了4%~5%。证明模型的优异。
我在之前所学习的论文“Deep Feature Learning for Pulmonary Nodule Classification in a Lung CT”就提出了结合原始图像形态特征的想法,但他并没有结合深度模型,而本文结合了形态特征与U-Net网络,使得分割效果进一步提升,这两个的结合是一个很好的研究方向。