采用双支路和Transformer的视杯视盘分割方法

王甜甜, 史卫亚, 张世强, 张绍文

(1.河南工业大学信息科学与工程学院, 郑州 450001; 2.河南工业大学粮食信息处理与控制教育部重点实验室, 郑州 450001;3.河南工业大学人工智能与大数据学院, 郑州 450001)

青光眼在早期阶段没有病症,但可导致神经受损,进而导致视力下降[1]。视盘(optic disc,OD)是眼底图像中较明亮的区域,近似椭圆形;视杯(optic cup,OC)是视盘区域上存在的可变尺度的中央凹陷[2]。随着眼底图像分割技术的发展,基于OC和OD比率(cup-to-disc ratio,CDR)的客观评估方法成为可能[3]。以往的研究[4]表明,CDR的值越大,患青光眼的概率就越高。该评估方法要求能够精确分割出视盘和视杯,因此得到一个精确的自动分割视杯和视盘的算法至关重要[5]。

病变区域会对分割结果产生影响,可以通过预处理操作裁剪感兴趣区域(region of interest, ROI)来缓解该问题。因此,学者们将多种方法应用于视盘定位。李宁等[6]利用空洞空间金字塔模块和金字塔池化模块对视盘中心进行定位。Park等[7]使用重复阈值技术来检测图像中最亮的粒子,然后通过评估这些颗粒的圆度获取视盘中心。Wang等[8]构建了一个提取网络E对OD区域进行粗提取,从而确定OD的中心。通过定位的视盘中心,裁剪出指定的感兴趣区域。

基于深度学习的OD和OC分割算法主要包括全卷积神经网络(fully convolutional networks, FCN)[9]和U-Net[10]及其改进。由于卷积运算固有的局部性,基于CNN(convolutional neural networks)的网络模型无法很好地学习全局和远程语义信息交互。Transformer本身存在的自注意力机制利用了全局信息,能够有效解决上述问题。随着视觉Transformer(vision transformer, VIT)[11]的提出,Transformer首次应用于图像分类任务。VIT的补丁嵌入过程忽视了语言和图像之间的差距,导致了图像的局部特征遭到破坏。此外,VIT冗余的注意力机制设计,导致VIT模型计算量大。Wang等[12]引入金字塔视觉Transformer,逐步缩小金字塔以减少大型特征图的计算。同时,使用空间缩减注意力(SRA)层来进一步降低学习高分辨率表示的资源消耗,从而使得Transformer架构更适合密集预测任务。Wang等[13]在PVTv1的基础上提出了PVTv2,将PVTv1的计算复杂度降低到线性。

图1 PVT模型的特征提取过程Fig.1 Feature extraction process for PVT model

模型提取的低层次特征感受野较大,几何细节信息表征能力强,但是语义信息表征能力弱;高层次特征则相反[14]。为了更好地利用多尺度特征信息,研究人员通过不同的特征融合方法对多尺度特征进行融合。张文博等[15]提出了DeepLabV3+,该模型使用空洞金字塔池化结构来提取不同尺度的特征图并结合。Yu等[16]提出了BiSeNet,使用双分支结构,同时解决空间信息丢失和感受野过小的问题。傅双杰等[17]同样使用双分支结构。在语义路径嵌入上下文捕获模块和自适应特征融合模块,捕获多尺度的高语义上下文信息。

现提出一种采用双支路和Transformer的视杯视盘分割方法,首先使用Transformer编码器对图像特征进行提取,解决卷积运算的局部性;尺度感知-特征融合模块(scale-aware-feature fusion module,SCA-FFM)对高层次特征进行融合,使用注意力机制动态分配不同尺度特征的重要程度,收集高层次特征的语义特征;识别模块(identification module,IM)使用注意力机制减少低层次特征中存在的噪声;图卷积域-特征融合模块(graph convolution domain-feature fusion module,GCD-FFM)使特征图同时具有不同尺度的信息。

1 相关方法

1.1 金字塔视觉Transformer(PVT)

最近的研究表明[18-19],视觉Transformer比CNN表现出更优的性能。受到PVTv2模型的启发,本文研究使用PVTv2作为主干网络提取视杯和视盘的特征信息。PVTv2包含4个阶段,每个阶段生成不同尺度的特征图。每个阶段的结构相似,都是由补丁嵌入层和Transformer编码器组成,如图1所示。

其中,X1显示出视杯和视盘的外观信息,X2、X3和X4提供了高级语义信息。由于PVTv2原本是用于分类任务,为了将其适用于视杯和视盘的分割任务,将最后一层的分类器移除,在4个多尺度特征图之上设计一个视杯和视盘分割头。

1.2 空间和通道注意力机制

通道注意力机制[20]Attc将长度为H′、宽度为W′、通道为C′的输入特征图x∈RH′×W′×C′分别进行自适应平均池化和最大池化,将输入特征x变为1×1×C′大小,并得到两个特征图。然后分别将所得到的特征图送入两层的全连接神经网络,得到两个特征图并相加。随后,将相加得到的特征图通过Softmax函数得到权重系数,将其与输入的特征图x相乘,得到最终的特征图。

Attc(x)=σ{H1[Pmax(x)]+H2[Pavg(x)]}⊙x

(1)

式(1)中:x为输入的特征图;σ为Softmax函数;Pmax和Pavg分别为最大池化函数和自适应平均池化函数;H1(·)和H2(·)为共享参数,由1×1卷积核的卷积层组成,用于降低通道维数,然后是一个ReLu层和另一个1×1卷积层,恢复到原始通道数;⊙为哈达玛积。

空间注意力机制Atts通过两个池化操作分别获得通道维度的最大值和平均值,并进行拼接。然后经过一个卷积层,生成一个新的特征图。将得到的特征图通过Softmax函数生成空间权重系数,然后再与输入的特征图相乘得到最终的特征图x′∈RH′×W′×C′。

Atts(x)=σ(G{Concat[Rmax(x),Ravg(x)]})⊙x

(2)

式(2)中:Rmax和Ravg分别为最大值和平均值;G(·)为7×7的卷积层。

图2 模型的总体结构图Fig.2 Framework of the model

2 本文算法

2.1 整体框架

如图2所示,本文模型分为两个支路:①支路一使用识别模块获取空间语义信息,同时设计空间注意力机制和通道注意力机制来增强空间细节特征的提取;②支路二使用尺度感知模块-特征融合模块将多个阶段具有强语义信息的多尺度特征表示相结合,从而提高图像分割能力。

此外,为了提高分割精度,本文研究使用主损失函数Lmain和辅助损失函数Laux对模型进行优化。

2.2 双分支结构

2.2.1 尺度感知-特征融合模块

受采集设备和其他不可控因素的影响,图像目标区域存在大小不同的问题,因此本文研究设计了SCA-FFM作为分支一。即不再采用简单的解码器以达到较高的预测速度,而是采用SCA-FFM来捕获多尺度特征信息,从而平衡计算资源和分割精度。

SCA结构如图3所示,它引入注意力机制,可以动态选择合适的尺度特征并进行融合。输入两个不同尺度的特征Fa和Fb,将更高层次的特征图Fb上采样到与Fa相同的大小。接下来将它们进行拼接,经过一系列卷积操作后得到两个特征图A和B。然后经过Softmax运算生成像素级注意力特征A和B,将其分别与原始输入特征Fa和Fb进行哈达玛积运算得到对应的特征图并相加,得到最后的融合特征:

(3)

SCA(Fa,Fb)=A⊙Fa+B⊙Fb

(4)

式中:H为特征图的高度;W为特征图的宽度;eAi和eBi分别为将特征图A和B的第i个像素转化到指数函数上,用于保证概率的非负性。

SCA-FFM的最终输出为

T1=SCA[X′2,SCA(X′4,X′3)]

(5)

2.2.2 识别模块

低层次特征主要包含边缘、颜色、纹理和形状等信息,能够定位目标的位置。由于眼底图像的血管复杂,且背景与视杯和视盘相似。因此,本文研究中设计IM作为分支二来识别低层次特征中的视杯和视盘细节。

由图4可知,IM由通道注意操作Attc和空间注意操作Atts组成。该模块受到CBAM注意力机制模块的启发,能够使得模型更关注目标物体。

IM的最终输出为

T2=Atts[Attc(X1)]

(6)

2.3 图卷积域聚合模块

(7)

图3 SCA-FFM结构图Fig.3 Framework of SCA-FFM

图4 IM模块Fig.4 IM model

V=AP{K⊙F[Wg(T2)]}

(8)

式(8)中:Wg(·)为卷积单元;F(·)操作在通道维度上应用Softmax函数,并选择第二个通道作为注意图;AP(·)为池化和裁剪操作。

a=σ(V⊗KT)

(9)

Y=aT⊗GCN(a⊗Q)

(10)

式中:⊗为矩阵乘法;a为K和V的相关注意图;GCN(·)为图卷积层。

F=T1+Wz(Y)

(11)

图5 GCN-FFM细节图Fig.5 Details of the GCN-FFM

2.4 损失函数

本文研究的损失函数包含两部分,Lmain主要用于P2和手工标注真实值G之间计算。Laux用于监督SCA-FFM生成P1。计算公式为

L=Lmain+Laux

(12)

(13)

(14)

3 实验

先将OD分割出来,而后通过OC处于OD内部这一先验知识分割OC。分别在Drishti-GS1[25]、RIM-ONE-r3[26]和REFUGE[27]数据集上进行实验。

本文提出的OD和OC分割算法步骤为:①使用E网络,裁剪出512×512大小的图片作为OD的ROI;②将ROI区域送到OD分割网络,而后输出分割图;③将步骤②中得到的OD分割二值图与步骤①中的ROI区域图进行像素“与”操作,得到mask ROI(masked by optic disc)区域;④将步骤③中的mask ROI区域送到OC分割网络,得到分割图作为输出;⑤将步骤②和步骤④中输出的分割图进行合并,得到合并的单张图像,完成OD和OC的分割。

3.1 实验细节

本文提出的模型基于预训练的PVTv2。本文实现基于公共平台PyTorch、GeForce RTX 2080Ti GPU和12 G内存。使用Adam优化器并将学习率和权重衰减均设置为1×10-4。Batch大小设置为8,epoch设置为100。

3.2 数据集与数据预处理

在Drishti-GS1[25]、RIM-ONE-r3[26]和REFUGE[27]三个公共数据集上进行实验。每个数据集的相关信息如表1所示。

Drishti-GS1数据集中包含101张图像,其中包含51张训练集,25张验证集和25张测试集,分辨率均为2 896×1 944。RIM-ONE-r3数据集包含159张图像,其中包含99张训练集,30张验证集和30张测试集,分辨率均为2 144×1 424。REFUGE数据集包含1 200张图像,其中包含400张训练集,分辨率为2 124×2 056;400张测试集和400张验证集,分辨率为1 634×1 634。

表1 数据集描述

为降低计算成本,同时降低病变区域对分割结果的影响。采用Wang等[8]提出的视盘定位方法裁剪出ROI。首先定位视盘的中心,然后以该中心点裁剪出512×512像素大小的视盘区域,该区域完整包含视杯和视盘,如图6所示。

由于医学图像数据涉及病人的隐私,同时又需要专业人士对数据集进行标注处理,因此医学图像数据的获取困难且数量有限。本文研究对所用数据集进行数据增强来增加眼底图像的数量,所采用的数据增强方法包括眼底图像的平移、旋转、拉伸和缩放。

图6 裁剪ROI区域Fig.6 Cropping the ROI area

3.3 评价指标

采用Dice[28]、IoU[29]和Fβ[30]对OD和OC分割模型进行评价。Dice和IoU都是用于衡量预测结果与mask之间的区域相似性。

(15)

(16)

式中:TP、TN、FP、FN分别为真阳性、真阴性、假阳性、假阴性。

Fβ综合考虑了精确率和召回率,为不同位置的不同误差分配不同的权重。

(17)

(18)

文献[30]表明β2=0.3,精确度更高。因此本文设β2=0.3。

3.4 学习能力的定量评估

本文研究添加了3个模块提高模型的性能,但是造成了模型参数量的增加。原始的PVTv2模型的参数量为22.6 M,本文模型的参数量为28.11 M。为了进一步证明本文算法的学习能力,在Drishti-GS1、RIM-ONE-r3和REFUGE数据集上按照表1的数据划分方式进行实验。将本文算法与M-Net、U-Net、Robust[31]、CA-Net[32]、pOSAL、CE-NET[33]和CDED-Net[34]进行对比。在Drishti-GS1数据集上,本文OD和OC分割结果的Dice比CA-Net分别高0.45%和0.11%,如表2所示。在RIM-ONE-3数据集上,本文OD分割结果的Dice比CA-Net低0.76%;但是本文OC分割结果的Dice比其高2.65%,如表3所示。在REFUGE数据集上,本文OD分割结果的Dice比CDED-Net高2.39%;但是在OC分割任务中略逊于CDED-Net,如表4所示。

表2 不同方法在Drishti-GS1数据集上的分割结果对比

表3 不同方法在RIM-ONE-r3数据集上的分割结果对比

表4 不同方法在REFUGE数据集上的分割结果对比

分析可得,本文模型在视杯和视盘分割任务中能够达到与其他算法持平或更优的效果。

对比实验表明,本文算法具有更好的学习能力。这主要得益于Transformer架构强大的特征提取能力,从而获得更强的特征表达能力。同时通过融合模块对不同层次的特征进行融合,充分利用不同层次的特征。

3.5 泛化能力的定量分析

为进一步验证模型的泛化能力,在Drishti-GS1、RIM-ONE-r3和REFUGE数据集上进行实验。使用Drishti-GS1和RIM-ONE-r3数据集进行训练,REFUGE数据集进行测试。值得注意的是,Drishti-GS1、RIM-ONE-r3和REFUGE的采集设备不同,因此,图像的颜色和纹理均存在差异(图7)。

图7 不同数据集的图像示例Fig.7 Examples of images from different datasets

绿色区域代表视盘区域;蓝色区域代表视杯区域图8 不同方法在DRISHTI-GS1数据集上的分割结果Fig.8 Segmentation results of different methods on the DRISHTI-GS1 dataset

从表5可得,本文方法分割视盘的Dice和IoU分别为0.959 8和0.922 5,分割视杯的Dice和IoU

分别为0.879 6和0.778 9。由于训练集和测试集的图片差异较大,不同的方法都存在一定程度的精度下降。但是,本文提出的方法包含IM模块,能够准确识别图像的颜色、纹理等信息,因此在交叉实验中仍然比其他算法的分割性能更优。综合上述可得,本文提出的方法具有良好的泛化能力。

表5 在REFUGE数据集上进行泛化实验

绿色区域代表视盘区域;蓝色区域代表视杯区域图9 不同方法在RIM-ONE-r3数据集上的分割结果Fig.9 Segmentation results of different methods on the RIM-ONE-r3 dataset

绿色区域代表视盘区域;蓝色区域代表视杯区域图10 不同方法在REFUGE数据集上的分割结果Fig.10 Segmentation results of different methods on the REFUGE dataset

3.6 定性分析

图8、图9和图10分别为不同方法在3个数据集上的分割结果。

从图8可得,M-Net和U-Net的视盘和视杯分割效果均略次于其他算法;Robust、pOSAL和CE-NET算法对视盘的分割效果较好,但是在视杯分割任务中表现欠佳。尤其是CE-NET方法中视盘和视杯的Dice相差最大,达到10.9%,这与图5的可视化结果一致。在Drishti-GS1数据集中,本文算法的视盘和视杯分割效果均达到最优,展现了良好的分割效果。

由图9可知,RIM-ONE-r3数据集中的图像清晰度低于Drishti-GS1数据集,导致不同算法的分割效果均存在下降现象。U-Net和M-Net模型分割出的边界不规则,且错分严重。CE-NET网络在Drishti-GS1和RIM-ONE-r3数据集的实验结果相差不大,但对视杯和视盘的分割效果均次于Robust、pOSAL这两种联合分割网络。此外,Robust、pOSAL和本文模型对视盘的分割效果持平,但是本文算法在视杯分割任务中体现出优越性。在REFUGE数据集上,本文模型分割边界平滑,且分割效果良好。

综上所述,本文方法可以适应不同的数据集,同时分割的边缘更加接近真实标签。

3.7 消融实验

为了验证ROI区域、SCA-FFM、IM和GCD-FFM的有效性,本文研究进行了一系列的实验。如表6所示,使用PVTv2在DROSHTI-GS1数据集上进行实验,分割视盘的Dice和IoU分别为97.05%和94.39%。在PVTv2基础上添加的3个模块,均对分割结果产生不同程度的提升。其中,SCA-FFM使得分割视盘的Dice和IoU分别提升0.73%和1.12%,同时使得分割视盘的Dice和IoU分别提升0.73%和1.06%,对分割性能的提升效果最佳。3个模块均添加之后,效果达到最优,综上所述,本文添加的3个模块均有助于优化视盘和视杯的分割结果。

表6 DRISHTI-GS1数据集中的消融实验

4 结论

提出了一种采用双分支和Transformer的视杯和视盘分割模型。双分支结构分别对低层次特征和高层次特征进行处理,使得提取的特征更加稳健,从而提高分割性能。本文算法在3个经典数据集上均体现出其优越性,但是引入的3个模块增加了计算的工作量。此外,由于数据集样本较少,影响了研究的结果。在未来的工作中,将从降低计算成本和小样本分割方面入手,进一步完善实验结果。