融合多注意力机制的自监督小样本医学图像分割*

要媛媛,刘宇航,程雨菁,彭梦晓,郑 文,2

(1.太原理工大学计算机科学与技术学院(大数据学院),山西 晋中 030600; 2.长治医学院山西省智能数据辅助诊疗工程研究中心,山西 长治 046000)

1 引言

图像分割技术旨在将感兴趣区域(Region of Interest)从某视野内分离出来。医疗图像中对感兴趣区域(如肝脏、肿瘤)的研究具有极其重要的意义。将图像分割技术应用于医疗领域(如CT(Computed Tomography)、MRI(Magnetic Resonance Imaging)等图像),能够帮助医生更加清晰地认识各内部器官或病理结构,达到智能辅诊、高效诊断的效果。

主流的基于全监督的深度学习分割模型在标记丰富的数据上训练时可以取得良好的效果。这些有监督的神经网络的性能在很大程度上取决于训练数据集的质量和数量,例如注释良好的数据的数量、类分布的平衡度以及样本表示[1,2]。就自然图像而言,大量含标注的公开数据集(如VOC(Visual Object Classes)等)使得研究人员极易获得预训练模型,便于开展后期研究。然而,对医疗图像的认知往往基于丰富的临床专业知识,获取含标注的医疗图像会耗费大量的人力成本及时间成本。此外,疾病多种多样,影像表现复杂多变,以某一训练好的特定模型认知全部未知类别的特征极不现实。受人类仅使用少量输入数据来区分对象的认知能力的启发,小样本学习技术应运而生,探求模型是否拥有从少量样本中学习和概括的能力。

传统的深度学习模型用小样本训练时易导致过拟合或对目标任务的欠拟合。为避免此类问题,需要依赖于具有许多注释训练类的大型训练数据集,但这样又违背了小样本学习的初衷。自监督学习主要是利用图像着色、超像素生成伪标签等辅助任务从原始的无监督数据中提取有效信息。这些有效信息用于训练网络,从而学习到对下游任务有价值的表征。其特征在于,模型训练过程中使用的信息源于自身,而非受人为因素影响的外部标注,以此为模型提供丰富的内部信息。

小样本分割FSS(Few-Shot Segmentation)的目标是使用相应的注释掩码分割查询图像中选定类别的目标区域。FSS最流行的方法是基于度量的原型学习。通过掩码平均池化,将每个类别在空间上平均为一个一维代表原型或线性分类器的权重向量,例如,PANet(Prototype Alignment Network)[3]、PFENet(Prior guided Feature Enrichment Network)[4]、SG-One(Similarity Guidance network to tackle the One-shot)[5]、CANet(Class-Agnostic segmentation Network)[6]。然而,因人体构造复杂,且器官、病灶或其他兴趣目标在医学图像中占比较小,存在前景-背景极其不平衡的问题。如果背景中的无关外观信息被错误地识别和破坏,就可能使得前景-背景边界模糊,增加预测的难度。因此,通过掩码平均池化操作,此类原型网络可能会丢失图像的详细空间信息。

针对医学图像复杂、可能存在多个感兴趣区域、预训练模型难获得,以及前景-背景极其不平衡的问题,本文使用超像素分割方法为图像生成伪标签,为模型提供从图像本身获取信息的可能;同时,摒弃掩码平均池化操作,聚焦于局部特征及全局特征,引入注意力机制以明确前景-背景边界,达到对感兴趣的区域进行分割的目的。

本文提出的SSF-MANet(Self-Supervised Few-shot learning with Multi-Attention mechanism Network)旨在通过构造融合样本自身通道信息、空间信息及样本间信息等多类信息的模型,探寻小样本数据的内在关联,以便在综合考量局部及全局信息后完成对兴趣目标的分割。首先,SSF-MANet模型利用依赖于超像素经旋转及亮度改变后生成的伪标签得到特征图谱,自适应局部原型池化网络ALPNet(Adaptive Local Prototype pooling Network)分别接收support特征及query特征。之后,将得到的support特征传入外部注意力机制部分,以捕获样本间的联系;query特征则被传入通道-空间双注意力机制中,以便于图像内部多尺度特征的融合。最后,基于相似度的分类器对原型向量和query的特征图谱进行对比,并进行聚类融合得到最后的分割结果。值得注意的是,在整个过程中,通过PANet模型中的原型对齐正则化更新损失,以便最终的分割结果更接近真实的语义标签而非伪标签。

2 相关工作

2.1 小样本语义分割

是否拥有从少量样本中学习和概括的能力,是将人工智能和人类智能进行区分的明显分界点。小样本学习更倾向于模拟人类从少数样本中认知事物,受到了越来越多研究者的关注。

小样本图像语义分割的关键点在于构建query集与仅有的少数support集样本间的关系。多分支方式[6-8]的引入使得模型在充分获取support及query集各自的特点后进行融合,进而对query图像的分割结果进行动态调整,以递进的方式对图像分割模型进行改进和完善。CANet 在多分支结构的基础上,引入密集比较模块,开始注重support图像和query图像间的多层次特征比较。之后的过程中,研究人员常常引入余弦相似度以匹配support和query图像间的特征。Zhang等人[5]首先采用掩码平均池化(Masked Average Pooling)策略,通过只考虑support图像的像素来生成引导特征,然后利用余弦相似度度量引导特征与query图像特征之间的关系。

为分析query图像和support图像之间的细粒度对应关系,Min等人[9]提出利用多层次特征相关和高效四维卷积的超相关挤压网络,旨在从中间卷积层的不同层次中提取多样特征,以从粗到细的方式逐步将高相关的高级语义和低级几何线索压缩到精确的分割掩码中。在医学图像分割领域,Roy等人[10]提出了一种新颖的Few-Shot框架,用于只有少量注释切片的医学图像分割,条件分支和分割结构通过“通道挤压和空间激励”模块增强交互。

这些研究应用元学习的思想动态调整分割结果,并在此基础上进行图像间的多层次特征比较,证明了元学习获得共同知识并概括到新任务的能力。然而,这些研究局限于搜索匹配小样本间的信息。本文将重点转移至建立样本自身空间及通道信息,以便准确捕获数据特征,从而提高目标任务中的分割性能。

2.2 注意力机制

获取多个尺度丰富的上下文信息往往需要依赖于注意力机制。Vaswani等人[11]率先提出融合注意力机制的Tranformer模型以捕获长距离的依赖关系,为图像分割工作提供了异于RNN(Recursive Neural Network)和CNN(Convolutional Neural Network)的工作思路。

受全卷积网络有效感知域较小、无法充分捕获长距离信息等特点的限制,研究人员将注意力机制引入到多尺度输入的语义分割网络中[12,13],以缓解局部邻域约束。SENet(Squeeze-and-Excitation Network)[14]则更关注图像通道间的关系,通过计算各特征通道的重要性,然后针对不同的任务增强或者抑制不同的通道,达到提取特征权重的目的。

为了捕捉丰富的上下文信息,Fu等人[15]与基于多尺度特征融合的方法不同,提出了位置注意模块和通道注意模块,分别对空间和通道维度上的语义相互依赖关系进行了建模,利用自注意力机制,自适应地整合局部特征及其全局依赖关系。

注意力机制在捕获样本内部各空间、通道间长距离信息的应用中逐渐深入人心,如何建模样本间的特征关系成为研究人员关心的问题。EA(External Attention)[16]考虑到自注意力机制具有二次复杂度,并且忽略了不同样本之间的潜在相关性,提出外部注意力,依靠2个外部的、小的、可学习的、共享的内存,隐式地考虑所有数据样本之间的相关性。

DANet(Dual Attention Network)[15]及EA[16]分别从不同的角度出发,进行样本内及样本间关系的建模,隐式地学习了整个数据集的特征。本文利用以上注意力机制模块,尽可能多地捕获样本自身通道信息、空间信息及样本间信息等多类信息。

2.3 自监督学习

自监督学习的出发点是考虑在缺少标签或完全没有标签的情况下,依然能够学习到表示原始图像的有意义的特征,其从数据本身提取信息的思想被越来越多的人接受。

为了缓解数据稀疏问题,研究人员基于图像修复的思想,通过填补洞的方式为图像增添标注[17,18]。然而,Context encoders等工作[17-24]从自我监督中学习到的参数对本文的分割任务不具有强相关性,往往需要进行第二阶段的微调。与此不同的是,Li 等人[25]提出了一个新的自我监督的方法,即在查询图像的背景中通过超级像素生成伪类,以提供额外的训练数据。伪类的引入使得研究人员必须关注其与真实分割的相关性,PANet[3]在训练过程中进一步引入了一种原型对齐正则化,通过反向执行few-shot分割,从查询图像和支持图像中提取原型,并在嵌入空间中对齐。

本文沿用文献[25]的思路,通过使用超像素分割方法为图像生成伪标签,对图像内部的相似颜色、纹理等特征进行聚类。超像素作为真实物体语义掩码的紧凑构建块,与分割真实物体直接相关。

2.4 超像素分割

超像素是由一系列位置相邻且颜色、亮度、纹理等特征相似的像素点组成的小区域。这些小区域大多保留了进一步进行图像分割的有效信息,且一般不会破坏图像中物体的边界信息,用少量的超像素代替大量像素表达图像特征,降低了图像处理的复杂度。分割趋于使一个区域内的元素相似,不同区域内的元素不相似。EGBIS(Efficient Graph-Based Image Segmentation)算法是基于图的贪心聚类算法。为了解决EGBIS算法容易过分割的问题,Felzenszwalb等人[26]结合流行的超像素方法SLIC(Simple Linear Iterative Cluster)[27],在构造无向图前,先对原图进行预分割,得到形状大小基本一致且贴合图像边界的SLIC超像素;然后以超像素的颜色平均值对其进行特征描述,将每个超像素视为无向图节点,构造无向权重图,进而实现基于图的无监督分割。由此生成的超像素在形状上更加多样化,更易分辨出未知类别的形态。

本文采用文献[26]中的方法对图像进行超像素预分割。

Figure 1 Structure of SSL-MANet图1 SSL-MANet结构

3 模型介绍

本节首先介绍小样本自监督学习研究的问题表述,阐述研究过程中的N-way K-shot问题。然后,介绍架构中的MAP(Mixed Attention Pooling)模块,重点是通过自适应池化网络捕获的原型向量如何融合多注意力机制,以突出表现其多尺度的特征信息。最后,介绍如何将依赖于伪标签生成的分割结果与真实的语义标签相匹配。

3.1 问题的提出

小样本学习的研究目的是用少量的标记样本迅速完成训练,建立泛化性能较好的模型。在小样本图像分割中,给定一个包含训练语义类别Ctr的训练集Dtr和一个包含测试未见类别Cte的测试集Dte,其中Ctr∩Cte=∅。任务是给定几个Cte的注释例子(如参照几个肾脏标记图像来分割肾脏),在Dtr上训练一个分割模型,该模型可以在Dte的图像中分割语义类Cte,而无需重新训练。在每个推理过程中,都会给出一个支持集S= {(xs,ys(cj))}(其中s=1,2,3,…,N代表支持集中第s个图像及其对应的标注,cj代表第j个类)和一个查询集Q={xq}(其中,q=1,2,3,…,N代表查询集中第q个图像)。一个支持-查询对(S,Q)包括一个episode对。如果有N个要分割的类(也称为N个任务)和每个类的S中的K个标记的图像,则每个episode对都定义了一个N-way-K-shot的分割子问题。背景类被表示为c0,不计入Ctr或Cte。

3.2 网络架构

(1)小样本图像利用特征提取器将其表征为向量。fθ(·):χ→ε,其中θ为参数,ε代表分割操作的特征空间。

(2)提议的融合多注意力机制的原型池化模块,用于从支持标签及特征中提取原型向量,g(·,·):ε×Y→ε。

(3)基于相似度的分类器通过比较原型和查询特征进行前景-背景的分割,sim(·,·):ε×ε→Y。

(4)原型对齐正则化模块,通过反向执行few-shot分割,从查询图像和支持图像中提取原型,并在嵌入空间中对齐。

网络结构如图1所示。Adaptive local prototype pooling模块用于从支持特征和标签中提取表示原型;FSS model模块用于原型从嵌入空间中支持和查询图像对齐。过程中,特征提取器网络fθ(·)通过将xs和xq映射到特征空间ε,为ALP(Adaptive Local Prototype pooling)提供了特征图。ALP将每个(fθ(xs),ys(cj)) 对作为输入,计算语义类cj和背景c0的局部原型及类级原型。这些原型以后将被用作每个类别的参考,辅助分割查询图像。所有cj的原型形成一个原型集合P={pk(cj)},j= 0,1,2,…,N,其中k是原型索引,每个cj的k≥1。这个原型集合被分类器sim(·,·)用来预判查询图像的分割,即yq=sim(P,fθ(xq))。即首先测量每个pk(cj)和查询特征图fθ(xq)之间的相似性,然后将这些相似性融合在一起。

MAP模块中,自适应局部原型池模块(ALP)保留了原型中的局部信息。受Fu等人[28]的工作启发,每一个兴趣目标向量由多个子目标组成。在ALP中,每个局部原型只在覆盖在支持物上的局部池化窗口内计算,即只代表对象的一部分支持,并且只代表兴趣目标的一个部分。具体来说,对每个fθ(xs)∈RD×H×W进行平均池化,池化窗口大小为(LH,LW),其中(H,W)是空间大小,D是通道深度。值得注意的是,(LH,LW)决定了在表示空间E中计算每个局部原型的空间范围。在平均池化特征图的空间位置(m,n)上获得的具有未决定类c的局部原型pmn(c)由式(1)给出:

pmn(c)=avgpool(fθ(xs))(m,n)

(1)

其中,mLH≤h<(m+1)LH,nLW≤w<(n+1)LW。

考虑到超像素分割可能将兴趣目标分割为多个部分,使用掩蔽平均池化法计算一个类级原型pg(cj)(如式(2)所示),以避免类内信息的丢失。

(2)

其中,g代表全局类别。

之后,为捕获图像内自身多尺度信息及图像间信息,分别将Q集(查询集)及S集(支持集)中的图像传入通道-空间混合注意力模块及外部注意力模块中,以充分利用样本自身的特征及少数样本间可能存在的关联性。

针对查询集中的图像特征A∈RC×H×W,首先输入一个卷积层,分别生成2个新的特征图B和C,其中B,C∈RC×H×W;然后进行向量化操作,将RC×H×W转化为RC×N(C代表图像的通道数,N=H×W);接下来,将C的转置与B矩阵进行乘法操作,得到图像内像素的相似度矩阵。空间注意力图S∈RN×N的计算公式如式(3)所示:

(3)

其中,sji表示第i个位置对第j个位置的影响。

2个位置的特征表征越相似,它们之间的相关性就越大。将通过空间注意力模块的特征图D∈RC×H×W转化为RC×N形式后,将其乘以一个尺度参数α,并对特征A进行元素求和运算,得到最终输出Ej∈RC×H×W,如式(4)所示:

(4)

与空间注意力模块不同,通道注意力不需要通过卷积重新生成的2个新特征图,而是直接从输入的特征图A计算通道注意力图(X∈RC×C),如式(5)所示:

(5)

其中,xji表示第j个通道对第i个通道的影响。与空间注意力一样,将其乘以一个尺度参数β,并对特征A进行元素求和运算,得到最终输出Ej∈RC×H×W,如式(6)所示:

(6)

针对支持集中的图像特征,设计随机初始化的共享矩阵M∈RT×d,M是随机初始化的,d是特征向量的维数,T是实验设置的超参数,将其设置为64,参照自注意力机制的公式,将DA记为一种M对输入特征F的注意力矩阵,如式(7)所示:

DA=(γ)i,j=Norm(FMT)

Fout=DAM

(7)

其中,(γ)i,j表示第i个像素和第j行之间的相似度。M是共享的,所以能够隐式地考虑不同样本之间的关联。

3.3 基于超像素分割的自监督学习

本文所用的语义类标注伪标签依赖于超像素分割生成结果,模型训练前就已得到了这些分割好的超像素伪标签。为了使得模型训练所得的分割结果更贴近真实的语义标签而非伪标签,引入原型对齐正则化模块,在余弦相似度对比分割结果的过程中,通过反向执行few-shot分割,不断提高图像的分割效果。

(8)

4 实验与结果分析

4.1 数据集

本文实验在CHAOS(Combined CT-MR Healthy Abdominal Organ Segmentation)健康腹部器官分割比赛中的Abdominal-MRI和Abdominal-CT数据集上进行,包括20位不同患者的图像。该数据库不包括注释器官(即肝脏、肾脏、脾脏)边界处的任何肿瘤或病变。

下载数据集后,将数据转换为nii文件并以病人ID命名,以便于查看。为了统一实验设置,所有图像都被重新制作成二维轴向切片,并调整为256×256像素,然后生成超像素伪标签。每个二维切片在通道尺寸上重复3次,以适应网络。

4.2 评价指标及实验设置

为了评估改进模型的效果,本文采用了常用的图像分割指标DSC(Dice Similariy Coefficient)和Recall对图像分割结果的精准度进行量化。DSC是一种集合相似度度量指标,通常用于计算2个样本的相似度,其计算如式(9)所示:

(9)

其中,T1和T2分别表示2个样本,T1∩T2表示2个样本的交集,|T1|和|T2|分别表示2个样本的元素个数。

召回率(Recall)又称查全率,是一个类别被预测正确的概率,如式(10)所示:

(10)

其中,TP表示模型预测结果为正,并且实际结果也为正的数量;FN表示模型预测结果为负,但实际结果为正的数量。

本文中沿用SSL-ALPNet(Superpixel-based Self-supervised Learning Adaptive Local Prototype pooling enpowered prototypical Network)[30]中的实验设置,设置1(set 1)使用少量样本建立标准分割标注,这与多数传统方式相同,为了模拟极端的标注稀疏情况,仅采用1-shot设置。为了评估对未见过的测试类的泛化能力,设置2(set 2)通过从训练数据集中删除包含测试类的图像,强制测试类(甚至是未标记的)完全不被看到。

4.3 对比实验

本文使用DSC和Recall评价指标,将SSF-MANet与医学图像FSS方法SE-Net及SSL- ALPNet模型进行比较,结果如表1和表2所示。其中,表2显示的是强制测试类完全不被看到的情况。图2展示的是Abdominal-CT数据集下分割结果的定性对比图。在不使用任何人工注释的情况下,本文提出的模型在4个目标器官上整体DSC提升了3%左右。

相较于本文的baseline模型SSL-ALPNet,融合多注意力机制的模型SSF-MANet在肾脏及肝器官上的分割DSC平均提升了约4.5%,而脾脏器官部分的分割DSC降低1.9%左右。

综合脾脏器官分割指标DSC下降的原因,注意力机制将重心更多转向于空间等多维度信息。而脾脏在影像学中大而集中,且其在成像中偏独立,周围可能对其产生影响的器官较少,过多在意捕获远距离信息反而会给脾脏的分割带来噪声,影响分割结果。同样成像较大而密集的器官,肝脏通常在图像中呈现更为明显的边界,并且其组织结构相对均匀。这种特性在初始生成超像素的过程中起到了关键作用,因为它使肝脏更容易被清晰地划分为一块完整的超像素,为模型分割提供了良好的数据基础。

同时还注意到,同为肾脏,右肾的分割效果明显好于左肾的。从器官成像角度考虑,由于右肾上方有肝脏,所以右肾通常较低,造成右肾较左肾小5%左右。因此,考量右肾分割效果好于左肾的原因在于,前者成像特点突出,且与大而集中的肝相邻,边界清晰;前者成像区域较小,在超像素聚类融合时优势更明显,可以在一定程度上免去生成超像素带来的误差。

Table 1 DSC comparison of SSF-MANet model and other models on CHAOS dataset (set1)表1 SSF-MANet模型与其他模型在CHAOS数据集上的DSC对比(set1)

Table 2 DSC comparison of SSF-MANet model and other models on CHAOS dataset (set2)表2 SSF-MANet模型与其他模型在CHAOS数据集上的DSC对比(set2)

Figure 2 Segmentation results under the Abdominal-CT dataset图2 Abdominal-CT数据集下分割结果

4.4 消融实验

为了探究本文引入多注意力机制对模型分割结果的整体影响,分别对通道-空间注意力模块及外部注意力模块进行消融实验,结果如表3所示。

很多情况下,仅加入通道-空间注意力机制或外部注意力机制都可能导致分割性能降低。主要原因在于,Abdominal-MRI图像中的大多数器官与周围的组织有明显的对比,对某一图像内进行单一长距离建模时,可能会给图像的分割带来一定程度的噪声。此时,如果能将support中仅有的数据关联起来,就能够在整体上对图像的分割起到更积极的作用。

图3为支持集内不同样本数量对分割效果影响程度的曲线图。可以看到,在样本数量达到10个左右时,曲线逐渐平缓,DSC在0.81附近。Kavur等人[30]在该数据集上的全监督工作,Dice>0.8,与本文结果的差距较小,足以证明小样本学习工作的可信性。由此说明,本文模型克服了图像标注的困难,解决了标注稀缺的问题,在可接受误差范围内,小样本学习极大程度上解决了医疗领域的图像分割难题。

Table 3 Results of ablation experiment (set2)表3 消融实验结果(set2)

Figure 3 Effect of the number of support set samples on the segmentation effect using DSC as an indicator图3 以DSC为指标探究支持集样本数量 对分割效果的影响

5 结束语

本文提出将多种注意力机制模块加入到小样本自监督学习网络中,在利用超像素分割等方法从图像自身获取信息的基础上,以较小的成本,尽可能多地挖掘小样本数据的内在关联信息,以便在综合考量局部及全局信息后完成对感兴趣目标的分割。本文通过使用位置和通道双注意力模块对单一图像内部的多尺度特征进行融合,使用外部注意力模块显著突出样本间的联系。实验结果表明,本文模型在小样本数据情况下,具有较高的准确度,在多个评估指标上均优于其他模型的。这表明,多注意力机制在小样本医学图像分割任务中是一种有效的解决方案。但是,在实验结果分析中还发现,SSF-MANet在分割大而集中的感兴趣目标时稍显逊色,后续工作会不断改进模型结构,使其能在多变的病灶及其他可能存在的感兴趣目标分割任务中,获得稳定的提升效果。

小样本学习因数据稀疏问题,相较传统深度学习模型,可能会带来一定的分割精度损失。与传统深度学习模型对比后,通过调整N-way-K-shot任务数,寻找分割性能提升临界值,评估小样本学习的意义。研究结果表明,本文模型在小样本数据支撑下的精度损失在可接受范围内。

本文实验在CHAOS分割比赛的数据集上进行,分割的目标为健康者的器官,这与医疗领域病灶分割需求不相符合。因此,后期会在含病灶的其他数据集上进行进一步的测试及模型调优,使其在医疗图像分割领域中具有普适性,同时凭借其较低的计算复杂度和较少的训练样本需求的优势,更广泛地应用于智能辅诊领域。