SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

论文阅读:SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation 用于视网膜血管分割的空间注意力U-Net
      

摘要:视网膜血管的精确分割(精确分割precise segmentation)对早期诊断与眼病有关的疾病,如糖尿病和高血压具有重要意义。在这项工作中,我们提出了一个轻量级的网络(lightweight network),称为空间注意力U-Net(SA-UNet),它不需要数千个带注释的训练样本,并且可以以数据增强(data augmentation)的方式利用(utilize)它来更有效地使用可用的带注释的样本。SA-UNet引入了一个空间注意模块,该模块可沿空间维度(spatial dimension)推断注意图,并将注意图与输入特征图相乘以进行自适应特征细化(自适应特征细化adaptive feature refinement)。另外,所提出的网络采用结构化的dropout卷积块代替U-Net的原始卷积块,以防止网络过拟合(overfitting)。我们基于两个基准视网膜数据集(benchmark retinal datasets)评估SA-UNet:血管提取(DRIVE)数据集和儿童心脏与健康研究(CHASE_DB1)数据集。结果表明,所提出的SA-UNet在两个数据集上均实现了最新的性能(最新的,最先进的state-of-the-art)。


Keywords:Segmentation; retinal blood vessel; SA-UNet; UNet; spatial attention


I. INTRODUCTION

      通过观察眼底血管系统(the fundus vascular system)可以很容易地诊断和追踪(diagnosed and tracked)许多疾病,因为这些疾病(例如糖尿病和高血压)会导致视网膜血管的形态变化。全身性微血管和小血管疾病(Systemic microvascular and small vessel diseases)是糖尿病(diabetes)引起的常见病理变化,尤其是眼底视网膜血管疾病最为脆弱。糖尿病性视网膜病(DR)是由糖尿病引起的。如果观察到糖尿病患者视网膜中血管的肿胀,则需要特别注意。患有长期高血压的患者可能会因动脉血压升高或血管狭窄而观察到血管弯曲,这被称为高血压性视网膜病(HR)。视网膜血管分割是眼底图像定量分析的关键步骤。通过分割视网膜血管,我们可以获得视网膜血管树的相关形态信息(例如血管的曲率,长度和宽度)。此外,视网膜血管的血管树(the vascular tree)具有独特的特征,也可以应用于生物特征识别。因此,准确分割视网膜血管具有重要意义。
      然而,视网膜血管具有许多小而脆弱并紧密相连的血管,因此视网膜血管树的结构相当复杂。另外,血管面积和背景之间的差异不明显,并且眼底图像也容易受到不均匀照明(uneven lighting)和噪声的影响。上述原因导致视网膜血管分割是一项艰巨的任务。
      在过去的几十年中,已经提出了大量的视网膜血管分割方法,主要分为手动(manual)和自动(automatic)分割方法。前者既费时又费力(time-consuming and labor-intensive),并且要求从业人员(practitioners)具有很高的专业技能。后者可以减轻手工分割的负担,因此对自动分割算法(algorithms)的研究具有重要的意义。近年来随着深度学习的发展(With the advancement of deep learning in recent years),深度学习逐渐成为视网膜分割的主流技术。
      在医学图像分割领域,U-Net 是一个常见且广为人知的骨干网络(backbone network)。基本上,U-Net由典型的下采样编码器和上采样解码器结构以及它们之间的“跳跃连接”组成。它通过编码和解码过程结合了局部(local)和全局上下文信息(context information)。由于U-Net的出色性能,许多用于视网膜血管分割的最新方法都基于U-Net。 Wang等提出了(reported)双重编码U-Net(DEU-Net),该网络显着增强了网络以端对端和像素对像素的方式分割视网膜血管的能力。 Wu等提出了Vessel-Net,它首次结合了初始方法(initial method)和残差方法(residual method)的优势来进行视网膜血管分割。Zhang等提出了AG-Net,AG-Net设计了一种称为“ Attention Guide Filter”的注意力机制,以更好地保留结构信息。尽管这些U-Net变体表现良好,但它们不可避免地使网络变得更加复杂且难以解释。
      为了解决这些问题,我们在U-Net中引入了空间注意力,并提出了一个轻量级的网络模型(lightweight network model),我们将其命名为Spatial Attention U-Net(SA-UNet)。如SD-Unet 所示,使用DropBlock 可以有效地防止网络的过度拟合,因此即使是很小的样本数据集(例如视网膜眼底图像)也可以得到很好的训练。另外,批标准化(BN)可以提高网络的收敛速度(convergence speed)。因此,SA-UNet首先采用集成DropBlock和批标准化(BN)的结构化丢弃卷积块的变体来替换原始的U-Net卷积块。更重要的是,视网膜眼底图像中的血管特征和非血管特征之间的差异并不明显,尤其是小而边缘(marginal)的血管区域(vascualar arears)。通过引入少量其他参数(attention parameters),空间注意力可以增强重要特征(例如血管特征)并抑制(supress)不重要的特征,从而提高网络的表示能力。我们在两个公共视网膜眼底图像数据集上评估SA-UNet:DRIVE和CHASE_DB1。我们首先通过消融实验评估网络新引入的部分。实验结果表明,引入的结构化dropout卷积块和空间注意力 是有效的,并且与原始的U-Net和AG-Net相比,所提出的SA-UNet非常轻巧。最后,与其他现有的视网膜血管分割技术相比,我们提出的SA-UNet可以实现最新的(state-of-art)性能。


SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

II. METHODOLOGY

A.Network Architechure

      如图1所示,我们提出了一个具有U形编码器(左侧)-解码器(右侧)结构的SA-UNet。编码器的每一步都包括一个结构化的dropout卷积块和一个2×2 maxpooling 操作。每个卷积块的卷积层之后是一个DropBlock,一个批标准化(BN)层和一个ReLU激活函数(ReLU:the Rectified Linear Unit),然后将最大池化操作用于步长为2的下采样。通过每个下采样后,特征通道的数量增加了一倍(double the number)。解码器中的每个步骤都包括用于上采样的2×2反卷积操作,使得特征通道的数量减半,将自编码器中相应特征图串联在一起,然后是结构化的dropout卷积块。空间注意模块添加在编码器和解码器之间。在最后一层,使用1×1卷积和Sigmoid激活函数来获取输出分割图。

B.Structured Dropout Convolutional Block

      尽管对原始数据集进行了数据扩充,但是在原始UNet训练期间仍然会出现严重的过度拟合,如图2所示(左)。 因此,具有18个卷积层的轻量级U-Net被用作我们的基本体系结构,但它仍然存在过拟合问题,如图2(中)所示。 在最近的计算机视觉作品[10],[11],[13],[19]中成功应用DropBlock结构,我们采用DropBlock来规范网络。
      DropBlock是一种结构化的Dropout形式,可以有效地防止卷积网络中的过拟合问题。它与Dropout的主要区别在于,它从图层的特征图中丢弃连续的区域,而不是丢弃独立的随机单位。基于此,我们构造了一个结构化的Dropout卷积块,即每个卷积层之后是一个DropBlock,一个批标准化(BN)和一个ReLU激活单元,如图3的右侧所示。在SD-Unet的卷积模型中(如图3的中间所示),结构化的dropout卷积块引入了批标准化(BN)以加速网络收敛。我们使用此结构化的dropout卷积块代替U-Net的原始卷积块来构建U形网络作为“主干网”。与原始U-Net的23个卷积层相比,我们的Backbone仅具有18个卷积层,如图2所示(左图),过度拟合问题得到了完美解决,并加快了网络的收敛速度。


SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

C. Spatial Attention Module (SAM)

      空间注意力模块(SAM)是卷积块注意力模块的一部分,用于分类和检测。 SA使用要素之间的空间关系来生成空间注意图。 为了计算空间注意力,SA首先沿通道轴应用最大池化和平均池化操作,并将它们连接起来以产生有效的特征描述符,如图4所示。形式上,输入特征

通过通道的方式最大极化而向前移动和平均池化以生成输出SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
分别。然后,使用卷积层,在级联特征描述符上使用Sigmoid激活函数,以生成空间注意图
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
简而言之,空间注意力模块的输出特征
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

的计算公式为:
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
其中
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

表示内核大小为7的卷积运算,而
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

表示Sigmoid函数。

III. EXPERIMENTS

A. Datasets

      我们在两个公共视网膜眼底图像数据集上评估我们提出的SA-UNet:DRIVE和CHASE DB1。有关这两个数据集的具体信息,请参见表I。应注意,这两个数据集的原始大小不适合我们的网络,因此我们通过在其周围填充零填充来调整其大小,但将大小裁剪为评估过程中的初始大小。为了扩充数据,我们对两个数据集采用表I最后一列所示的四种数据扩充方法,每种方法都从原始图像生成了三张新图像,也就是说,我们从原始的20张训练图像中扩充了两个原始数据集到256张图像。

SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

B. Evaluation Metrics

      为了评估我们的模型,我们将分割结果与相应的标签进行比较,并将每个像素比较的结果分为真正(TP),假正(FP),假负(FN)和真负(TN) 。 然后,使用灵敏度(SE),特异性(SP),F1-分数(F1)和准确性(ACC)评估模型的性能。
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

      在视网膜血管分割中,只有9%-14%的像素属于血管,而其他像素被视为背景像素。 Matthews相关系数(MCC)适用于对具有不同大小的两个类别的二进制分类进行性能测量。 因此,MCC值可以帮助找到血管分割算法的最佳设置。 MCC定义为:
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
      ROC曲线(AUC)下的面积可用于测量分割的性能。 如果AUC值为1,则表示完美分割。
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

C.Imlpementation Details

      为了监视我们的网络是否过度拟合,我们在DRIVE和CHASE_DB1扩充数据集中随机选择了26和13张图像作为验证集。如前所述,图2显示了在DRIVE数据集上训练100个轮的情况。使用增强训练集从头开始训练SA-UNet。对于这两个数据集,均使用了Adam优化器和二进制交叉熵损失函数,并且为了使参数数量保持较小,将第一卷积层之后的通道数设置为16。训练次数为150,且前100轮的学习率是0.001,后50个轮的学习率是0.0001。 DropBlockd discard blocks的大小设置为7。
      分别对于DRIVE数据集,训练的批处理大小设置为8,DropBlock的dropout rates设置为0.18。对于CHASE_DB1,批处理大小设置为4,dropout rates为0.13。
      该实验基于以Tensorflow为后端的公共Keras,所有实验均在具有12 GB显存的NVIDIA TITAN XP GPU上运行。

IV. RESULTS

A.Ablation Experiments消融实验

      为了证明所提出的SA-UNet的每一部分都可以改善视网膜血管分割的性能,分别在DRIVE和CHASE_DB1上进行了消融实验。表II和III显示了U-Net,U-Net + SA,SD-Unet(即U-Net + DropBlock),主干网(即SD-Unet + BN)和SA-UNet(即主干+ SA)的分段性能)分别从顶部到底部。另外,表IV列出了不同型号的参数数量。
      从结果中,我们可以获得一些有用的观察结果:(1)仅增加了98个参数,U-Net + SA的性能就比U-Net更好,证明了引入空间注意力的策略是有效的。 (2)在使用基于U-Net的结构化Dropout卷积块的情况下,主干的ACC,AUC,F1和MCC分别为0.28%/ 0.22%,0.73%/ 0.59%,2.42%/ 2.48%和2.48比DRIVE和CHASE_DB1上的U-Net高出%/ 2.64%,这表明采用新构造的结构化Dropout卷积块来构建骨干网的有效性。 (3)虽然参数数量略有增加,但与原始SD-Unet相比,骨干网具有更好的性能,这表明添加批标准化(BN)可以在一定程度上改善网络性能。 (4)最后,提出的SA-UNet在大多数指标上均达到最佳性能,并且与AG-Net和具有23个卷积层的原始U-Net相比,我们的SA-UNet的参数量要少得多,因此对于视网膜血管分割的任务,SA-UNet是一个轻量而有效的网络。

SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

      在图5中,我们显示了一个关于DRIVE数据集的测试示例,包括通过U-Net,U-Net + SA,AG-Net,SD-Unet,Backbone和提出的SAUNet取得的分割结果,以及相应的标签。与UNet和U-Net + SA相比,AG-Net在边缘结构的分割方面确实具有一定优势,但是在小血管的交点处,AG-Net是不足以强大的。 SDUnet忽略了一些边缘和小的血管结构,甚至存在不正确的分割。与U-Net和SD-Unet相比,骨干产生的小血管分割更为准确,这证明了使用结构化的dropout卷积块构造的骨干的有效性。与骨干相比,本文提出的SA-UNet可以对小边界血管产生更准确的分割结果,并保留更多的视网膜血管空间结构,证明空间注意力机制可以突出血管并减少背景的影响。为了更好地观察测试结果,我们分别在图6和图7中在DRIVE和CHASE_DB1上显示了更多U-Net,Backbone和SA-UNet的分割示例。

B. Comparisons with state-of-the-art methods

      最后,我们将SA-UNet的性能与当前在视网膜血管分割任务中使用的其他最新方法进行了比较。在表V和VI中,我们总结了不同方法的发布年份以及DRIVE和CHASE_DB1数据集的性能。从结果可以得出结论,SA-UNet在DRIVE和CHASE_DB1上均达到了最佳性能。它的最高灵敏度(sensitivity)为0.8212 / 0.8573,最高准确度为0.9698 /0.9755,最高AUC为0.9864 / 0.9905,而特异性可与其他方法媲美。另外,与先前方法中性能最好的AG-Net相比,SA-UNet在小血管交叉处具有更好的分割性能,如图5所示。注:SA-UNet的参数量要小得多。比AG-Net的要好。以上结果表明,我们提出的SA-UNet在视网膜血管分割挑战中达到了最先进的性能(state-of-the-art performance)。

V. CONCLUSION

      大多数视网膜眼底图像数据集是典型的小样本数据集,这可能会使训练深层神经网络有问题。为了实现学习,以雄心勃勃的方式应用了数据增强,然后使用了轻量级的U-Net,但是仍然观察到过拟合。受到DropBlock和批处理规范化在卷积神经网络中成功应用的启发,我们将U-Net的卷积模块替换为结构化的辍学卷积模块,该结构将DropBlock和批处理规范化集成为我们的骨干网。另外,在眼底图像中,血管区域和背景之间的差异不明显,尤其是边缘和小血管。为了帮助网络学习这些知识,我们在主干的编码器和解码器之间添加了一个空间注意模块,并提出了空间注意U-Net(SAUNet)。空间上的关注可以帮助网络专注于重要特征,并抑制不必要的特征,从而提高网络的表示能力。我们在两个公开可用的视网膜眼底图像数据(包括DRIVE和CHASE_DB1)上评估SA-UNet。实验结果表明,使用卷积块的结构化缺失和引入空间注意力是有效的,并且通过与其他用于视网膜血管分割的最新方法进行比较,我们的轻型SA-UNet实现了最新技术表现。由于视网膜图像的血管结构特征相似,因此我们得出结论,SA-UNet是一个通用网络,可以应用于其他视网膜血管分割任务。

SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation
SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation

论文地址:https://arxiv.org/pdf/2004.03696.pdf
代码地址:https://github.com/clguo/SA-UNet

上一篇:论文笔记二:Weakly Supervised Learning of Instance Segmentation with Inter-pixel Relations


下一篇:论文翻译:U-Net: Convolutional Networks for Biomedical Image Segmentation