基于跨任务一致性的半监督肝脏CT图像分割

李明漾,王庆凤,陈立伟,黄 俊,周 莹

(1.西南科技大学 计算机科学与技术学院,四川 绵阳 621000;2.绵阳市中心医院 放射科,四川 绵阳 621000)

0 引 言

肝脏作为人体非常重要的器官发病率非常高,肝脏是容易病变产生肿瘤的高发部位,肝癌属于恶性肿瘤,死亡率极高。目前计算机断层扫描(Computed Tomography,CT)已经成为了医学成像的主要方法之一,快速精准地从CT图像中分割出肝脏区域对医生诊断肝癌、定性分析以及制定治疗方案具有重要意义,肝脏分割结果可以为医生临床诊断提供可靠的诊断依据。目前临床诊断中,基于CT图像对肝脏的分割是由经验丰富的医生手动进行,分割结果容易受医生主观判断和分析的影响,增加误判或漏判的可能性。

近几年,研究人员提出了一些基于深度学习的医学图像分割方法,在肝脏分割领域中取得了显著成果,分割精度大大提高。现有大多数方法使用以深度卷积神经网络为代表的深度学习方法,它们主要是采用全监督的设置来完成肝脏CT图像分割任务,作为全监督学习的一个重要特性,其性能在很大程度上依赖于足够数量的标记数据。然而在医学领域标注可靠的注释图像非常耗时耗力。据统计,手动分割一个肝脏序列平均需要30 min[1]。标记数据的缺乏引导了许多研究者对注释效率的研究,例如弱监督学习、无监督学习、半监督学习等。由于获取少量带有标记数据和大量未标记数据在临床上是可行的,所以该文将基于半监督学习的方法分割肝脏CT图像。

主要贡献:为了解决医学图像中肝脏标注数据缺乏的问题,提出了一种基于跨任务一致性的半监督肝脏CT图像分割方法。该方法具有多解码器的网络模型以增加模型扰动的多样性,同时引入辅助预测任务学习肝脏的几何信息提高模型分割性能。

1 相关工作

基于深度学习的图像分割已经在医学图像领域取得了显著的成功,所以首先需要回顾一些经典的医学图像分割方法和基于半监督学习的医学图像分割方法。

1.1 用于医学图像分割的深度学习方法

深度神经网络对医学图像分割获得了较好的效果,许多优秀的医学图像分割算法已经得到了广泛的应用。全卷积神经网络(FCN)[2]是医学图像分割领域中极其重要的一部分,FCN能够提取图像中目标区域的高层次特征与低层次特征并将两者进行结合,实现了图像的自动分割。U-Net[3]已经成为一种经典医学图像分割框架,具有编码-解码结构成对称形状,并能够通过跳跃连接将编码器和解码器得到的不同分辨率的特征信息进行融合。He等人[4]引入了残差结构思想,提出了ResNet用于医学图像分割,该结构将输入特征图与输出特征图进行特征融合,在增加网络深度的同时避免了梯度爆炸或消失等问题。李秀华等人[5]将U-Net中编码阶段采用混合空洞卷积替换原始卷积并在跳跃连接中添加注意力机制,改进后提高了肝脏肿瘤分割精度。Cicek等人[6]将U-Net扩展到3D医学图像分割领域用于体积分割,其思想是采用3D卷积核提取特征,捕获图像丰富的空间信息。戴振晖等人[7]通过在3DU-Net添加残差块并结合Transformer模块提出了Res-Swim-UNet模型,实现了肝脏及肿瘤的自动分割。Milletari等人[8]提出了V-Net,他们将残差结构与3DU-Net结合并将下采样的池化操作用步长为2的卷积代替,上采样使用转置卷积,并使用3D图像输入网络模型。张宇啸等人[9]提出了一种轻量级V-Net并结合深监督策略的肝脏及肿瘤分割算法,在减少模型的参数量和计算量的同时提高了分割精度。

1.2 半监督医学图像分割

半监督学习在医学图像分割领域已经得到了广泛的研究和应用。半监督学习方法是需要将标记数据和未标记数据都用来训练模型,目的是充分利用无标记数据来提高模型性能。Bai等人[10]采用了自训练式半监督学习的思想,利用标记数据训练初始模型,然后利用训练好的初始模型预测未标记数据生成伪标签,挑选出高质量伪标签从而扩展带标记的数据集。刘清清等人[11]提出的自训练式半监督肝脏分割框架是通过改进的注意力模块的3D scSE-Unet作为分割网络进行循环迭代训练。另外一种基于图的算法假设两个相似的样本具有相同的输出[12]。对抗性学习也是半监督学习的一种流行方法。Fang等人[13]通过结合对抗性训练策略区分像素是来自预测分割结果还是标签样本。Zhang等人[14]采用对抗性学习使注释图像与未注释图像产生相似的分割输出。

最近,越来越多的基于一致性正则化的半监督学习方法用来进行医学图像分割。其利用未标记数据的主要思想是基于假设,即通过将输入数据添加小扰动,模型输出产生的预测结果应该是一致的。例如,Sohn等人[15]在执行训练样本强弱条件下,检验数据扰动预测的一致性。Laine等人[16]提出的π模型方法简单有效,为了执行可靠的预测,通过时间集合策略用前几个时期的预测来扩展π模型。Tarvainen等人[17]提出了均值教师模型(MT),使用教师模型来产生学生学习时的目标,学生模型利用教师模型产生的目标学习,并讨论模型参数一致性。Yu等人[18]进一步在MT模型中引入不确定性估计指导学习。杨星宇等人[19]提出了一种多时期一致的半监督模型(MPC-MT)实现无造影剂图像的肝脏肿瘤分割,其思想是基于均值教师模型,该模型在多时期磁共振图像之间建立一致性正则化。Li等人[20]提出的SASSNet模型引入了带符号距离图(SMD)[21],进一步学习器官的几何形状信息。Luo 等人[22]提出了基于双任务一致性的半监督学习框架,通过在分割任务和水平集函数回归任务之间建立一致性正则化约束,实现了任务级的正则化。与以往的方法不同,该文提出的半监督学习框架包含多个解码器,同时添加辅助回归任务学习肝脏的几何形状信息。该框架在主副解码器的不同任务之间执行跨任务一致性正则化以利用大量未标记数据,提高模型分割肝脏的性能。

2 文中方法

2.1 模型框架

本节将介绍提出的基于跨任务一致性的半监督肝脏分割模型框架,总体结构如图1所示。该文采用V-Net作为基础网络。它由一个共享编码器、一个主解码器和两个副解码器组成,每个解码器有两种预测任务。这三个解码器在结构上是相同的,仅仅在上采样方法上不同。定义副解码器D1,D3的上采样方法分别采用三线性插值、最近插值。主解码器D2的上采样方法与V-Net保持一致,采用转置卷积。在主副解码器中使用Dropout层以实现模型层面的扰动。

图1 基于跨任务一致性的半监督学习肝脏分割框架

2.2 跨任务一致性正则化

在一般的基于一致性的半监督学习方法中,一致性损失设计用于数据级上的平滑预测。与之前方法不同,该文利用预测回归任务得到的有符号距离图(SDM)训练网络模型,该任务可以使模型学习目标特征的几何轮廓信息,其中回归任务预测的SMD可由解码器后添加一个1×1×1卷积块和Tanh激活函数实现。可以通过T(·)转换函数将分割标签转换成SMD,转换函数定义如下:

(1)

在SMD中的每个点(体素)表示其之间的距离和最近边界,所以‖x-y‖2表示体素x,y之间的欧氏距离,∂C表示目标对象的表面,SMD中点的值为0则表示位于目标对象的表面。Cin和Cout分别表示为目标对象的内部和外部。对于目标对象的内部区域体素取负值,外部区域取值反之。由于有符号距离图(SMD)与像素分割概率图不在同一个域内,且转换式不可微的性质,为了建立主副解码器之间的跨任务一致性,该文将SMD通过T-1(·)函数转换成近似的像素分割概率图。采用Luo等人[22]提出的方法中的平滑近似函数来实现转换,转换公式定义为:

(2)

其中,z是体素x的有符号距离图的距离值,k是一个超参数尽可能地设置较大的值来近似变换接近分割真实值。此公式的目的是将SDM映射到与像素分割概率图同一域内,SDM变换成近似分割概率图使得建立跨任务一致性正则化有意义。

2.3 半监督训练损失函数

在上述半监督学习方法描述中,如图1所示,文中的损失函数由无监督损失Lmrc和监督损失Lsup组成。对于监督损失:文中方法的监督损失由标记图像集合Dl上对应在主副解码器的双任务优化目标组成,采用Dice loss和Cross-entropy loss作为三个解码器像素级分割预测任务的混合监督损失,计算如下:

(3)

式中,fseg表示像素级分割图,Y表示分割标签图像。

该文引入的回归任务作为模型中的重要部分,也将计算SDM产生的预测误差,其监督损失定义如下:

(4)

式中,fsdm表示有符号距离图(SDM),T(y)是分割标签图像Y通过T(·)函数转换成的符号距离图真实标签。总的监督损失由以上两个部分的损失加权求和所得。

Lsup=Lseg+λLsdm

(5)

式中,λ是平衡分割任务和回归任务损失项中的权重系数,将其设置为0.3。对于整体训练集D的无监督损失:模型在主副解码器上采样操作的不同以及Dropout层的作用,形成了模型扰动,在主解码器与副解码器之间建立跨任务一致性损失,使模型能够利用大量未标记数据进行学习。在训练过程中,在主解码器D2的像素级分割图和副解码器D1,D3输出的有符号距离图,以及主解码器的有符号距离图和副解码器D1,D3输出的像素级分割图之间建立一致性约束进行优化,具体一致性无监督损失计算如下:

(6)

Ltotal=Lsup+λ2Lmrc

(7)

继前面的研究工作,该文采用高斯升温函数策略λ2=0.1×e-5(1-t/tmax)来控制监督损失和无监督的一致性损失之间的平衡。t表示当前训练步骤,tmax表示最大训练步骤。高斯升温函数可以使模型在训练时更加高效,因为在训练的前期阶段模型预测结果相对来说可靠度低。

3 实验结果与分析

3.1 数据集及预处理

该文采用肝脏肿瘤分割公开挑战赛LiTS2017数据集进行方法验证,该数据集包含131例含有肝脏部位CT扫描图像,并且由3名放射科医生手动注释标签。在数据集中随机挑选104例CT图像用于训练,剩余27例CT图像用于测试。最后,对数据集做预处理操作,将原始图像的CT值(HU)调整到[-200,200]的范围,并归一化到[0,1]。为了降低计算机性能的要求,去除大量无关背景,将原始图像进行中心裁剪,最终得到尺寸大小为256×256×256的CT图像。

3.2 实施细节及评价指标

文中的实验是基于Pytorch框架实现。实验使用24 GB内存的NVIDIA RTX 3090 GPU进行训练。该网络框架采用SGD优化器经过6 000次的迭代训练,将初始学习率设置为0.01,每经过2 500次迭代学习率衰减0.1,每次迭代的Batch size设置为4,其中包含两个标记图像和两个未标记图像。该文将采用以下评价指标来验证提出方法的有效性:Dice相似系数(Dice Similarity Coefficient)、Jaccard系数(Jaccard Coefficient)和平均表面距离(Average Surface Distance,ASD)。各项评价指标计算公式如下:

(8)

(9)

(10)

式中,TP(True Positive)表示肝脏分割像素正确分类的预测结果;FP(False Positive)表示肝脏分割像素错误的预测结果;FN(False Positive)表示预测像素分割结果分类为肝脏实际为背景的像素;Aseg表示预测分割结果的边界;Agt表示真实分割标签的边界;p1,p2表示边界上的像素点。

3.3 对比实验与定量分析

该文在LiTS2017数据集上验证了该方法的有效性。与UA-MT[18],SASSNet[20],DTC[22],SMTL[23]等半监督方法进行比较。上述半监督方法都是以V-Net模型作为主干网络,为了体现与其他半监督方法对比的有效性,控制变量,文中半监督框架也采用V-Net网络模型作为基础网络。该框架在V-Net基础上添加了两个并行解码器以增加模型扰动的多样性。同时以上所述的半监督方法都是在标记数据和未标记数据设置数量相同比例下进行比较,在LiTS训练集样本中分别设置10%带标注样本和20%带标注样本的实验组,如表1和表2所示,其中V-Net[8]方法采用全监督设置进行训练作为基线对比。

表1 在10%标记数据训练下与其他半监督方法的对比实验

表2 在20%标记数据训练下与其他半监督方法的对比实验

首先,通过表1和表2的实验结果可知,所有的半监督方法与利用相同标记数据数量设置下的全监督V-Net方法相比都得到了更优的分割结果,达到不错的性能增益。其中表1显示了使用10%标记数据训练时, SASSNet,UA-MT,DTC的Dice系数、Jaccard系数、ASD比较接近,性能增益不明显,其中SMTL在ASD评判标准中表现最好。文中方法在Dice系数和Jaccard系数评判标准上达到了93.13%和87.43%,相比于V-Net[8]方法全监督设置下,该方法通过利用未标记数据训练模型将Dice系数提高了4.41百分点,Jaccard系数提高了6.34百分点。表2显示了在使用20%标记数据训练时,文中方法在Dice系数和Jaccard系数相比其他半监督方法产生分割精度增益更明显。该方法相比SMTL将Dice系数提高了1.34百分点,Jaccard系数提高了2.38百分点,且在ASD评判标准下比UA-MT,SASSNet,DTC拥有更优的分割结果。

综合两个表格实验结果表明,文中方法的分割结果和真实结果拥有更高的重叠率和相似度,在Dice和Jaccard评价标准上都优于其他半监督方法。增加了辅助任务使模型在肝脏的某些边界区域能够产生更好的分割结果,并且添加两个并行解码器建立一致性约束能充分利用未标记数据学习,提高分割结果质量。值得一提的是,文中方法在仅使用20%带标记数据时就能达到接近100%带标记数据全监督训练时分割结果的Dice,Jaccard精度,证明了该方法能够有效地利用大量无标记数据获取信息,提高分割性能。

3.4 实验结果可视化

最终,随机选取两名代表性患者的分割结果进行可视化分析。图2展示了患者在V-Net网络使用不同标记样本数量全监督设置下的分割结果与文中方法分割结果的2D可视化对比。

图2 与V-Net基础网络在不同标记数量的全监督设置下对比

其中(a)是原始图像,(b)专家标签,(c)和(d)分别是V-Net网络使用20%和100%标记数据在全监督情况下的分割结果,(e)是文中方法可视化结果。在使用相同标记数据数量训练网络模型的情况下,提出的半监督方法明显比全监督设置下的分割结果更好,尤其对肝脏的边界区域分割更准确更接近于真实标签,甚至部分细节更优于100%标记数据全监督设置下的分割结果。这是因为引入的辅助预测任务来学习肝脏的全局几何信息,它能够使模型学习到更多的肝脏边界信息,从而提高模型的分割性能。

4 结束语

提出了一种基于跨任务一致性的半监督肝脏分割方法。该方法以V-Net作为基础网络,添加两个副解码器以形成一定程度上的模型扰动,同时在每个解码器引入回归任务,利用主副解码器之间跨任务一致性约束建立损失函数,从而使模型能够利用大量无标记数据学习。同时,该方法与其他半监督方法相比,在LiTS数据集上肝脏分割结果精度更优,并且在仅使用20%带标记数据的情况下,肝脏分割结果在Dice和Jaccard的评价指标几乎接近100%带标记数据全监督设置下的分割结果。这意味着该方法在实用性方面具有一定优势。