基于U-Net的COVID-19病灶医学影像ZMINet分割模型

谷辛稼 陈一民

(上海大学计算机工程与科学学院 上海 200444) 2(上海建桥学院信息技术学院 上海 201306)

0 引 言

新型冠状病毒肺炎(COVID-19)是SARS-CoV-2病毒引起的一种传染病[1]。COVID-19对人们的健康、生活和经济造成严重的破坏。根据国家卫健委印发的《新型冠状病毒肺炎诊疗方案(试行第九版)》,肺部影像学特征也被作为判别COVID-19疑似病例的特征之一。

基于卷积神经网络的方法可将肺部影像信息转化为深层的高级特征,从而对病变区进行高效准确的定量分析[2]。Long等[3]提出了Fully Convolutional Networks(FCN)。FCN使用编码-解码的网络结构并使用跳跃连接传递特征,但FCN对于像素类别的预测没有考虑像素与像素之间的关系,上下文信息丢失,并直接对特征图进行高倍数放大导致预测部分结果模糊与细节丢失。为此Ronneberger等[4]提出U-Net,U-Net在解码器上使用参数可学习的反卷积进行上采样,并将编码器的所有层与解码器对应层进行跳跃连接以取得比FCN好的分割效果。但U-Net在不断下采样的过程中小物体信息和边缘细节容易丢失。对此Zhou等[5]提出U-Net++,在U-Net基础上使用层与层之间更为密集的连接来取代直接连接编码与解码对应层的跳跃连接,使小目标和边缘预测效果得到改善。Dense U-Net[6]将编码器上的卷积层替换为DenseNet的DenseBlock,使编码器层与层的连接更为紧密让特征传递变得更加有效,并且改善了训练过程中梯度消失的问题,但更密集的连接会带来更大的计算负担。Attention U-Net[7]并没采用密集连接的策略,而是在跳跃连接后加入了Attention Gate用来抑制与目标不相关区域的特征响应,实现在不大量增加参数的情况下提高性能。

不少研究针对于COVID-19病灶的分割在类U-Net的编码器-解码器结构网络的基础上进行改进。Zhou等[8]参考膨胀卷积设计了Res_dil block使U-Net获取更多来自不同区域的语义信息,同时在通道与空间维度引入注意力机制提升分割效果,并使用基于Tversky index的Focal loss损失函数(FTL)[9]以缓解小病灶分割效果不佳的问题,但是分割的精度不够高。Fan等[10]参考临床医生先确定病变大致方位,再通过观察局部准确提取病灶轮廓的思路,设计了反向注意力与边缘注意力机制建立Inf-Net分割模型,并利用半监督学习扩充训练数据,建立了优于当时大多数模型的Semi-Inf-Net,但该研究当时未考虑到实际中存在的无病灶CT影像,使模型在实际使用中性能可能会下降。Wang等[11]结合了Dice损失和MAE损失设计抗噪声的Dice损失函数,并提出了COPLE-Net结合自适应自集成学习框架从充满噪声的CT切片中准确找出病变区域,但是仅假设了噪声满足高斯分布,如果假设没有被满足,那么模型的性能也会下降。Ma等[12]从高效利用公开肺部数据集的角度出发,充分利用7个不同肺部CT数据集结合有限的COVID-19 CT数据集,分步实现对病灶的分割,但如果进一步改进模型结构可以取得更好的性能表现。

通过这些研究可以看出COVID-19病变区域分割的难点在于极大极小的病灶同时存在,而小病灶容易漏检;病变区灰度与周边组织接近并且背景噪声大,难以准确分割病灶边缘,并且带标签的公开数据集少不利于深度学习的训练。从深度卷积神经网络方法的特点来看,造成小病灶与边缘分割不佳的原因是多次卷积使特征图的分辨率下降,导致像素的空间信息丢失,致使网络无法对小目标和边缘细节进行精确还原。

为了改善这个问题,大量研究通过注意力模块、空洞卷积或空洞空间金字塔池化(Atrous Spatial Pyramid Pooling,ASPP)通过不同大小的感受野组合来获取尺度特征以改善这个问题。虽然在模型中加入多尺度感受野,如ASPP、金字塔池化模块[13](Pyramid Pooling Module,PPM)、它们的变体等可以提升模型对于多尺度特征的提取,但这类模块通常配置在网络的深层。然而随着网络深度的增加包含的小尺度目标的信息会减少[14],因此这些常用的方法并不能有效处理尺度剧烈变化的目标,而COVID-19病灶面积可以极大,比方说重症患者,也可以是极小的,因此仅靠空洞卷积、ASPP等方法对于小病灶检测改善有限。另一方面,少有研究针对从如何有效地将浅层特征有效传递到深层的角度去设计COVID-19病灶分割模型,但这对于COVID-19病变检测非常重要,因为病变区域是通过表面纹理特征来定义,如磨玻璃结节斑片、铺路石征、晕征等[15]。另一方面,病变区与其灰度相近的周边组织的差异同样需通过纹理来区分。密集的连接虽然能够用来传递浅层特征,但是模型的参数和训练时间也会随着连接数量的提升而增加。过于紧密的连接可能将浅层噪声也传入深层,导致模型的预测性能和执行效率同时降低。

针对COVID-19病灶边界模糊、尺度变化剧烈、梯度复杂、数据少的特点,我们提出一种基于U-Net的多尺度多层级特征融合的COVID-19病灶医学影像ZMINet分割算法。与其他基于U-Net的COVID-19病灶分割模型相比,我们使用多尺度多层级特征融合策略针对U-Net进行改进。在输入层面上,相较于传统单图输入,我们对图进行缩放增加两个额外尺度输入网络。在编码器层面上,我们使用ResNet50为主干网络提取特征。在特征传输层面上,参考文献[14,16-18]设计SIU-AIM模块取代U-Net的跳跃连接,并选择SIM作为解码器。首先多尺度整合模块(Scale Integration Unit,SIU)对三种尺度输入的特征进行逐层融合,并利用注意力机制使模型能更好地关注不同尺度下关键特征,接着我们使用聚合交互模块(Aggregate Interaction Module,AIM)融合相邻多层特征使模型减少特征的丢失进一步提升网络的性能。在解码器层面上,我们使用自交互模块(Self-Interaction Module,SIM)在网络内生成两个尺度的特征图丰富多尺度特征,并使用类似FPN的设计自下而上集成特征以提高模型性能。此外,我们提出一种相对于基于深度学习的方法更为轻量的数据扩充方法以改善模型的训练效率。因此,开发COVID-19病灶自动分割技术,可以提高COVID-19临床诊断过程的效率,对后续诊疗具有重要意义。另外,我们的研究对COVID-19的诊断与症状严重程度的判别具有重要意义。

1 模型设计

1.1 ZMINet模型

U-Net网络模型首次由Ronneberger等提出,该网络模型基于FCN网络,在结构上二者较为相似[19]。U-Net网络和FCN网络一样,都应用了编码器和解码器,都存在跳跃连接的拓扑结构。不同于FCN网络,U-Net网络在结构形态上是左右对称的。我们参考了在医学影像分割任务中常用的U-Net并对其进行改进,使其满足COVID-19病变区域分割任务。U-Net为编码器-解码器结构,如图1所示。其左侧是用来捕获上下文信息的编码器,右侧则是用于执行精确定位的解码器。首先编码器对输入的图像进行下采样,通过卷积、池化和归一化等操作对图像的特征进行提取。接着解码器通过反卷积将下采样中的深层特征恢复成原图大小,并利用跳跃连接将编码器与解码器对应层相连,将不同层的特征传递至解码器,最后通过训练进行像素级分类生成分割结果。

图1 U-Net网络结构

尽管类U-Net的方法在医学影像分割领域已经得到广泛的应用,但是COVID-19肺炎病变区域形状大小不一,灰度与周边组织接近,分割难度较大,这使得传统基于U-Net改进的方法预测结果边缘模糊,病变区域不完整以及边缘不连续。为了改善分割结果,我们首先将原图缩放生成额外两个尺度。这两个尺度分别是缩小至原图二分之一与放大至原图1.5倍,分别输入主干网络进行特征提取以获取更丰富的多尺度信息。小尺度图像输入能更好地捕捉目标与全局图像的关系,而大尺度图像则保留了局部特征与小尺度对象的细节。

编码器以ResNet50[20]为主干网络,在ResNet50的一至五层逐层加入了用于融合每层三个尺度的输出以及提取层间特征的SIU-AIM,作为传输层取代跳跃连接起到传递特征的功能,并使用SIM逐层融合到解码器的对应各层,改进后的模型结构如图2所示,我们将其命名为ZMINet分割模型。

图2 ZMINet分割模型

1.2 传输层SIU-AIM模块

在传输层部分,SIU-AIM模块由SIU和AIM组成。SIU的功能是对主干网络三种尺度的图像输入在主干网络每层的输出进行融合,并计算注意力使模型能自适应的获取不同尺度中的关键信息,其结构如图3所示。

图3 传输层SIU-AIM中SIU结构

SIU的整个计算过程可以用式(1)和式(2)表示:

(1)

(2)

式中:Ψ(▲,φ)表示注意力权重计算过程中的三个CBR层,φ表示其中的参数,[▲]表示通道拼接操作;U表示使用双线插值上采样;D表示最大池化与平均池化的操作。

尽管SIU能通过多尺度输入提取单层内的多尺度信息,然而单层特征只能表征该层尺度的信息,获取的多尺度信息有限,随着网络加深浅层特征依旧会逐渐减少。为了解决此问题,我们给出了AIM模块,融合不同分辨率相邻层的信息并传递至解码器,使用相邻层进行连接避免更为密集的策略带来的额外计算量,并且保留了层与层间的细节语义混合特征。其结构如图4所示。

图4 传输层SIU-AIM中AIM结构

AIM的整个计算过程可以用式(3)和式(4)表示。

(3)

(4)

1.3 解码器SIM模块

在解码器部分,SIM模块也采用类似AIM模块的辅助分支交互式学习方式。对输入的特征图进行缩放获取不同尺度下的信息表征。其结构如图5所示。

图5 解码器SIM结构

SIM使用从SIU-AIM输出的作为主路径,并对主路径进行一次下采样,作为分支路径。在对输入进行一次CBR后,SIM对低分辨率路径特征进行上采样,接着对高分辨率路径特征下采样,使两条路径均获得另一条路径同分辨率的特征并相互传递通过相加进行融合。最后低分辨率路径经过上采样、归一化和非线性处理(ReLU)后与另一条路径的特征,和通过残差路径传递的特征进行融合。

SIM的整个计算过程可以用式(5)表示。

(5)

(6)

式中:Fi+1(·)表示对残差分支与合并后的分支路径进行融合;Ui+1(·)表示上采样操作。

1.4 损失函数

参考相关的研究,如果COVID-19病灶仅以二元交叉熵(BCE,式(7))作为损失函数(BCEL)会导致预测结果模糊。我们通过研究针对预测结果模糊的问题选择了不确定感知损失(Uncertainty-Aware Loss,UAL)(式(9))作为损失函数,并采取与原文同样使用余弦变化来设定λ以达到最佳效果。经过实验,UAL的良好效果得到了验证。

(7)

式中:g代表实际为1的概率;p代表预测为1的概率。

(8)

L=LBCEL+λLUAL

(9)

2 针对非显著物体的肺部医学影像数据扩充

一项显著性目标检测领域的研究[21]提出并解决了一个现有显著物体检测(Salient Objective Detection,SOD)数据集中存在的设计问题。数据集假设每幅图都至少应包含一个清晰明显的物体。这种研究的设计偏差导致当前最先进的SOD模型在现有数据集上进行评估时出现性能优异,但在实际应用场景下表现不佳的问题。Xia等[22]基于对于显著物体的判别提出了先进的SOD模型,这体现了数据集中非显著物体的存在对准确推理显著物体至关重要。当前的大部分肺部医学影像数据集具有同样的设计偏差,假设每幅图像中都至少应包含一个病变目标。我们定义非显著物体的图像为不包含COVID-19所有种类病变的CT影像,最理想的图像是健康的肺。然而这类CT影像更为稀缺,因为少有肺部健康的人会主动进行胸部CT检测。

为了改善这个问题我们根据肺部医学影像的特点提出一种利用肺部的CT影像特征,在不依赖额外的数据的前提下,针对非显著物体样本的快速数据扩充方法。

基本原理是利用左右肺部近似轴对称的特点,将不包含病变区域的沿对称轴进行镜像复制,从而获得一幅没有病变区域的左右肺图片作为非显著物体的样本。我们首先判断肺部图像的方向,我们使用的数据集肺部影像为标准轴向CT切片顺时针旋转90°。因此,以水平方向的对称轴作为图像分割线,选择训练集真值图中按水平对称轴裁切后存在半边全黑的掩膜对应的训练集中的图,并将全黑半张掩膜对应的半边图像作为数据扩充的对象。将该半边图像按照原图的水平对称轴镜像复制,作为扩充后的非显著物体的样本,并生成与原图大小一致的全黑图作为该样本图片对应的真值图。

按照此流程,总共找到5幅满足条件的图片,将其按照流程处理后生成5幅不带显著目标的图,具体如图6所示。

图6 基于非显著物体数据扩张新增的样本与真值图

3 实验与分析

3.1 实验环境以及实验数据集

本文实验环境的CPU为i5-8600k,内存为32 GB RAM,显卡为GeForce GTX3090。所有模型均在PyTorch框架中执行,程序基于Python编写。

我们使用小批量随机梯度下降算法(Stochastic Gradient Descent,SGD)优化器来最小化损失函数。初始学习率为0.05,动量(Momentum)设置为0.9,权重衰减(Weight decay)设置为5E-4,epoch设置为50,Batch size设置为8。我们使用旋转、翻转和缩放,增强训练数据的多样性,以提高模型的泛化能力。模型使用ImageNet预训练后的ResNet50作为主干网络以缩短训练时间。

我们使用的数据集为COVID-19 CT Segmentation Dataset[23]。该数据集包括98幅来自不同COVID-19肺炎患者的轴向二维CT图像切片(数据样例见图6),由意大利医学放射学会(SIRM)收集,并由专业的放射科医生对肺部病变区域进行标注。我们在数据集中随机选取45幅二维CT图像作为训练集,5幅二维CT图像作为验证集,其余48幅二维CT图像用于测试。

3.2 评价指标

我们选取了6个医学分割中常用的度量指标,即:相似系数、敏感性、特异性、精确率、交并比、平均绝对误差。

相似系数Dice Score(DSC)是用来评估真值图(Ground Truth,GT)与预测结果(SR)的重叠率,DSC越接近1,预测效果越好。相似系数(Dice)的计算过程可以用式(10)表示。

(10)

敏感性反映模型对于病灶的像素的正确分类的能力,该值越高对于病灶的漏检率就越低。敏感性(Sen.)的计算公式如下:

(11)

特异性是用来衡量模型区分反例的能力,该值高说明模型对于病灶误检率低。特异性(Spec.)的计算公式如下:

(12)

精确率反映了在所有被预测为正例的结果中,真正的正例的占比。精确率(Prec.)的计算公式如下:

(13)

交并比(Intersection over Union,IoU)反映了预测的结果与真实值的交集和并集的比值。交并比(IoU)的计算公式如下:

(14)

式中:TP(True Positive)为真阳,即预测为真,真值图中的标签也为真;TN(True Negative)为真阴,表示预测与GT标签均为假;FN(False Negative)表示假阴,代表预测为假,但GT中的标签为真;FP(False Positive)即预测为真,GT标签却为假。

平均绝对误差(Mean Absolute Error,MAE)反映预测图Sp和真值图G之间像素的平均误差。平均绝对误差(MAE)的计算公式如下:

(15)

式中:w表示图像宽度;h表示图像高度。

综合评价指标精确率-召回率(Precision-Recall,PR)曲线是以召回率(同敏感性)为横坐标、精确率为纵坐标的曲线,表示在不同召回率取值下精确率的大小,曲线下面积越大表示模型性能越好。

综合评价指标F-beta Measure是准确率和召回率的调和平均值。Beta是用来调节精确率和召回率(敏感性)的偏重,当Beta为1的时候,就是统计学里的F1-Measure。F值越高表示模型性能越好。F值的计算公式如下:

(16)

F-beta Measure曲线代表了β不同取值下Fβ的大小。曲线下面积越大表示模型性能越好,模型的预测结果与真值图有更高的空间一致性。

4 实验结果分析

4.1 模型性能

我们选择U-Net、Attention U-Net、U-Net++和Dense U-Net为基准进行比较。Attention U-Net利用注意力机制去提升网络性能,而U-Net++和Dense U-Net则使用层间密集连接的方式传递特征提升性能。它们分别代表了提高关键特征的权重与传递因网络加深而丢失的细节这两种不同的改进思路。本文方法基于多尺度多层级特征密集集成与多分支交互式学习策略,由表1可以看出,相比这些模型,使用多种多尺度多层融合策略的模型ZMINet所有性能指标均有提高。

表1 ZMINet与基准模型的性能比较

表1中所有模型的Sen.都低于Spec.,结合MAE各个模型的变化来分析,像素预测错误的降低主要来自于模型区分反例能力的增强,并非对于病灶漏检率的降低,图7为分割结果。

图7 ZMINet与基准模型预测结果比较

可以看到,ZMINet的预测结果边缘更为清晰,对于小病灶的敏感性较强,预测目标的整体形状也与真值图(Ground Truth)更为接近。

4.2 针对非显著目标的数据扩充

表2为将扩充后的数据加入模型后的表现。从结果来看,加入了5幅不包含任何显著目标(COVID-19病灶)图片后,ZMINet的大部分体现模型整体分割性能指标如Dice、IoU和MAE(值越低越好)均有明显提升,说明这个数据扩充策略是有效的。我们观察到这个扩充策略会带来Sen.的提升与Prec.指标的下降和Spec.的轻微下降,这代表这个数据扩充的策略帮助模型识别出更多的病灶的同时会带来一定的误检,但在总体上模型的性能是变好的。

表2 数据扩充后的ZMINet与扩充前模型性能比较

图8为使用数据扩充策略后的ZMINet与使用前的预测图比较。结果显示这个数据扩充的策略能在一定程度上让模型更好地区分背景与小目标,改善小病灶的预测效果。

图8 扩充数据后ZMINet与扩充前的预测结果比较

4.3 消融实验

与U-Net、U-Net++等模型相比,在改进编码器,加入多尺度多层级特征融合组合模块SIU-AIM,并使用SIM集成特征后,对像素分类的性能得到了有效的提升。通过表3可以发现,无论是使用交叉熵(BCE)作为损失还是使用UAL,相对于U-Net、U-Net++等模型,各项指标都有了较明显的提升,与小病灶漏诊相关的敏感性指标得到了较大幅的提升,这意味着小病灶的漏检率的降低。在使用UAL后,Dice、IoU得到了提升。结合图7可以发现,对于边缘像素的预测相对于U-Net、U-Net++等模型变得更为准确。

表3 消融实验结果

为了更直观地展现基于U-Net改进的ZMINet的分割性能,我们绘制了消融实验中各模型的性能图,分别包括Precision-Recall(PR)曲线与F-beta Measure曲线用来衡量模型预测正负类的性能,具体如图9(a)和图9(b)所示。

(a) ZMINet消融实验Precision-Recall(PR)曲线

通过PR曲线可以看到ZMINet模型相对于U-Net、U-Net++等模型的AUC有显著的提高,这说明模型对于正类和负类的预测准确性均好于U-Net、U-Net++等。在使用UAL后,与使用BCEL相比,ZMINet可以在同样的召回率下取得略高的精确度。从F-Beta曲线可以看到,在各个阈值下,ZMINet的F分数均高于U-Net、U-Net++等,这体现了ZMINet的预测结果与真值图有更高的空间一致性。

5 结 语

我们在U-Net网络结构上建立一种精细高效的改进网络ZMINet,用于更好地分割COVID-19 CT影像中的病变区域。首先使用ResNet50提取全局信息,并在ResNet50的一至五层引入可以有效融合多尺度特征与多层特征的SIU-AIM,利用SIU-AIM模块与SIM模块融合多尺度与多层级信息以改善小病变区域的识别效果以及在网络加深的过程中保留更多细节,以提高模型的性能。为了证明本文方法的性能,我们使用不同的模型与本文方法进行比较。实验结果证明了本文方法的有效性。此外,我们使用了UAL作为损失函数来改善预测结果模糊的问题。我们使用了一种高效的、且不依赖额外数据的肺部CT影像的数据扩充方法,进一步改善模型处理图像的性能。但该方法尚有一些不足,首先是模型的泛化性能力还有待进一步加强,另外,由于ZMINet相较于U-Net等模型更为复杂、参数更多,因此训练时间相对长一些,但作为医学影像分割任务训练,还是在可接受范围之内。未来我们将在模型的泛化性能力提升及训练时间缩短方面做更为深入的研究。