改进UNet在肺癌患者颈部淋巴结超声图像分割中应用

宫 霞 吴卫华

(上海市胸科医院 上海 200030) (上海交通大学附属胸科医院 上海 200030)

0 引 言

超声成像具有成本低、便携性、无创伤和无辐射等优点,在现代医学检测中应用广泛;对超声图像结果准确性取决于操作者的经验,具有很大的主观性[1-2],比较耗时。近年来随着深度学习技术的发展,在医疗影像智能辅助诊断上的应用越来越广,例如图像分类[3]、分割和目标检测[4]等领域,通过人工标注好的数据,训练出模型的诊断精度都优于人类[5],极大提高了医生的诊断效率和准确性。因此,将深度学习应用于超声图像分割与分类,区分肺癌转移性淋巴结病变和良性病变具有重要的临床应用前景,超声图像分析具有强大的理论基础支撑。超声图像分辨率较CT、MRI低,图像存在大量伪影和噪声,其自身的局限性导致分类效果难以提升,文献[6]提出了一种综合利用B型超声和超声造影视频来提高分类精度的网络模型,深度学习技术应用于超声图像分析发展空间大,具有重要的理论意义和实际应用价值。

锁骨上淋巴结转移对肺癌转移鉴别诊断非常关键,有无淋巴结的转移直接关联到肺癌的分期、手术方式以及预后等,通过深度学习对肺癌患者颈部淋巴结超声图像分割对正确诊断肺癌转移具有重要意义。目前常见的应用于超声图像分割的深度学习模型FCN(Fully Convolutional Networks)[7]、UNet[8]、MaskRcnn[9]和UNet改进版等,其中UNet使用最为广泛,衍生的网络包括V-Net[10]、W-Net[11]、M-Net[12]、Y-net[13]和FPD-M-net[14]。文献[15]提出一种基于改进U-net网络的甲状腺结节超声图像分割方法,利用带有注意力模块的跳跃长连接部分对特征张量进行边缘轮廓保持操作。

1 相关工作

UNet[8]网络利用了神经网络内在的多尺度特性,浅层输出保存了空域细节信息,深层输出保存了相对抽象的语义信息,利用底层信息补充高层信息,适用于医学图像分割、自然图像生成,在医学图像分割比RPN和FCN网络有更好的分割精确度。文献[16]提出了一种改进UNet卷积网络,卷积块采用了inception block,网络将噪声激励函数NHReLU和NHSeLU(Noisy Hard SeLU function)代替ReLU和NReLU(Noisy Rectified Linear Unit functions)噪声激励函数;网络在两个尺度上预测输出,而不是只在上采样最后层输出,这样很好处理了超声图像中标注区域尺寸变化的问题,提高对淋巴结超声图像分割效果,达到了0.89。CE-Net[17]是一篇将空洞卷积和金字塔池化结合,专门用于2D医学图像分割任务。Fabian等[18]提出的nnUNet是基于UNet和3D UNet的医学影像分割算法框架。R2UNet[19]全称叫作Recurrent Residual CNN-based UNet,其融合了UNet、ResNet、RCNN的结构,在视网膜、肺和血管等多个医学影像分割任务上都取得很好的实验结果。Zhou等[20]提出的UNet++可以用于语义分割和实例分割,主要包括:通过不同深度的UNet的有效集成来缓解未知的网络深度,这些UNet可以部分共享一个编码器,并且可以通过深度监督[21]同时进行共同学习;重新设计跳接以在解码器子网络上聚合语义尺度不同的特征,从而产生高度灵活的特征融合方案;设计一种剪枝方案以加快UNet++的推理速度。UNet3+[22]是基于UNet和UNet++基础上提出的,它改进了UNet++没有直接从多尺度信息中提取足够多的信息的缺陷,UNet3+利用了全尺度的跳跃连接和深度监督。UNet3+中,可以从全尺度捕获细粒度的细节和粗粒度的语义。为了进一步从全尺寸的聚合特征图中学习层次表示法,每个边的输出都与一个混合损失函数相连接,这有助于精确分割,特别是对于在医学图像体积中出现不同尺度的器官。

注意力UNet[23]分为Hard Attention和Soft Attention,增加了Attention的机制,即注意力门(Attention Gate,AG)模型,通过自动学习参数来调整激活值。训练时能抑制模型学习与任务无关的部分,同时增加与任务有关的特征。AG接在每个跳跃连接的末端,对提取的feature实现Attention机制。加权的Res-UNet注意力机制[24]是通过将模型的最后一层的特征图与注意掩模相乘来实现的。利用这种加权注意机制,模型将只关注目标ROI区域,抛弃无关的噪声背景。PsP-Net[25]认为传统的交叉熵损失不适合小目标的医学图像,对损失函数进行了改进,选用Dice损失函数和Reg正则化损失函数结合。

上述的分割网络的分割精度低,对小目标检测、小样本训练存在缺陷,为此本文改进了注意力UNet网络结构,引入了新的激励函数、图像多尺度输入和适合小目标检测的损失函数。

2 方法设计

2.1 改进注意力UNet

改进的注意力UNet网络(A2R2UNet)是在注意力R2UNet基础上实现的,其结构与标准UNet基本相同,区别在于增加了循环残差模块和注意力模块。循环残差模块在考虑前面时间步的状态特征的同时,使得网络层数更深,进而提取更加丰富的特征,注意力机制使得网络更能提取感兴趣区域。首先,x和g都被送入到1×1卷积中,将它们变为相同数量的通道数,在上采样操作后(特征图大小相同)将x和g累加;然后,依次通过NHReLU激活函数、1×1的卷积、sigmoid,得到一个0到1的权重值;最后,注意力权重值与上采样特征图相乘,产生这个注意力块的最终输出。为了更好地提取局部特征,网络的输入图像为多尺寸,每一个尺度的输入对应各自的卷积层,其模型结构如图1所示,图中的输入图像分为四个尺度,每个尺度在本层使用conv和pooling,提取的特征分别作为下采样和右侧注意力门的输入,该输出特征与上采样进行concatenate操作。

图1 A2R2UNet网络结构

2.2 数据增强方法

数据增强(Data Augmentation)是在收集数据准备微调深度学习模型时,经常会遇到某些分类数据严重不足的情况,为了防止过拟合,在对预训练的网络模型参数进行微调之前,采用的一种技术。文献[26]结合了GAN和UNet,提出了生成对抗UNet,该模型是无域的,可以泛化到各种医学图像增广。通过实验,该方法对颈部淋巴结超声图像数据增强效果有效。第20个epoch训练结果如图2所示。

图2 数据增强训练结果

2.3 损失函数改进

(1) Dice损失函数。Dice loss来于Dice系数,是一种用于评估两个样本之间相似性度量的函数,取值范围为0～1,值越大表示两个值的相似度越高,计算公式如下:

(1)

式中:N是像素个数,p(k,i)∈[0,1]为类别k类在第i个像素相应的预测概率值,g(k,i)∈[0,1]为类别k在第i个像素的真实值(GT),K是类别数。

(2) Generalized Dice损失函数。Dice loss对小目标分割误差较大,因为如果小目标有部分像素预测错误,那么Dice值会大幅度的变动,从而导致梯度变化剧烈,训练不稳定。当病灶分割有多个区域时,使用Generalized Dice loss,其计算公式如下:

(2)

(3) Focal Loss。Focal Loss函数是对标准交叉熵损失函数的改进,通过增加难分类样本的权重,使得模型在训练时更专注于难分类的样本,公式如下:

(3)

因此,得出总损失函数为:

L=LGDice+LRg+λLFocal

(4)

式中:λ是LFocal的权重,根据验证集其取值为0、0.1、0.5或1;LRg表示正则化损失。式(4)损失函数可以避免在数据不平衡时,组合Loss会退化为Dice Loss,同时对小目标有较好的分割效果。

3 实证分析

本研究采集上海市胸科医院超声科360例肺癌患者的420幅淋巴结超声图像,其中男性200例,淋巴结300个,女性210例,淋巴结295个。所有淋巴结均进行针吸细胞学及细针穿刺活检检查,所有超声诊断结果均与病理结果相对照。论文中的数据增强的方法是旋转(90°、180°、270°、小于90°)、随机水平翻转、随机竖直翻转、cutout和对抗生成网络生成数据,共4 545幅。实验环境基于深度学习PyTorch框架,操作系统为Ubuntu 18.04,处理器为Intel i7-8700K,内存32 GB,显卡为11 GB显存的NVIDIA RTX 1080Ti。

数据集按照8∶1∶1划分为训练集、验证集和测试集,采用Adam优化器,用来训练的图像、UNet下采样和上采样卷积层指标明确,使用Dice系数来评估图像分割结果和原标签数据的相似程度。

通过图3-图5分析,NHReLU UNet网络Dice系数低于0.9,注意力UNet网络Dice系数为0.91,本文提出的A2R2UNet Dice系数达0.94。因此,通过实验得出,A2R2UNet分割效果好于注意力UNet和NHReLU UNet,为下一阶段良恶性的分类提供了可靠的数据支持。

图3 NHReLU Dice系数和损失函数

图4 注意力UNet网络Dice系数和损失函数

图5 A2R2UNet网络Dice系数和损失函数

图6列出了UNet、AUNet、R2AttUNet和A2R2UNet模型对淋巴结病灶分割结果,第一列为原超声图像,第二列为预测结果,第三列为人工标注病灶区域,预测病灶区域与人工标注的一致。激励函数、多尺度输入和损失函数改进,使得新模型针对小样本、小目标分割边界更加清晰,对病灶的细节信息分割更加精确。

图6 不同网络结构淋巴结病灶分割

表1列出了不同网络结构在平均IOU、Dice系数和Acc三个指标上的比较,可以看出本文提出的方法最优,在MIOU上A2R2UNet较UNet网络提升了12%,Dice系数提升了11%,Acc提升了3%。同样,与其余对比算法相比,本文的模型在各项指标上也均有较大改进。

表1 不同网络结构量化指标比较

A2R2UNet网络增加了循环残差和注意力模块,因此,较其他UNet网络在参数和计算资源消耗上都要大,如表2所示。可通过增加GPU数量进一步提高实时性。

表2 不同网络结构参数和比较

4 结 语

本文基于注意力UNet网络分割超声图像,引入了新激励函数、多尺度输入改进了网络结构,为了提高小目标分割性能和小样本训练模型,改进了损失函数,防止网络的过拟合;采用了数据增强和生成对抗UNet扩充训练数据,实验结果表明该网络提高了Dice系数。对超声图像分割出来的区域,需要分类网络进一步预测淋巴结病灶区域的良恶性。同时,在本文基础上,改进损失函数、对UNet++网络引入注意力机制等方法值得今后进一步研究。UNet++网络在参数和计算资源消耗上都优于R2AttUnet,因此,对该网络的优化将在准确率和实时性上都有提高。