[医学小样本分割]Self-supervision with Superpixels: Few-shot Medical Image Segmentation without Annotation

基于超像素的自监督:训练无注释的Few-Shot医学图像分割

论文地址
开源代码

摘要

Few-Shot语义分割(FSS)在医学影像领域具有广阔的应用前景。大多数现有的FSS技术都需要大量的注释语义类来进行训练。但由于缺乏标注，这些方法可能不适用于医学图像。为了解决这一问题，我们做出了以下几点贡献：(1)提出了一种新的自监督医学图像FSS框架，消除了训练过程中对标注的要求，生成基于超像素的伪标签进行监督；(2)将自适应局部原型池化模块插入原型网络，解决医学图像分割中常见的前景-背景不平衡问题；(3)我们证明了所提出的方法的一般适用性，医学图像使用三个不同的任务：腹部器官CT和MRI分割，以及心脏的MRI分割。结果表明，在医学图像分割方面，该方法优于传统的需要人工标注训练的FSS方法。

存在的问题及解决方案

现有的大多数few-shot语义分割模型在医学图像上并不算成功，因为大多数的FSS方法依赖于大量的训练数据集与许多注释的训练类来避免过拟合。为了绕过这种难以满足的注释需求，我们提出在未标记的图像上训练FSS模型，而不是通过单纯的自监督学习，这是一种无监督技术，通过精心设计的任务来学习可泛化的图像表示。对于许多最先进的FSS网络体系结构来说，另一个挑战是在一个空间变化的类上学习的表示容易丢失局部信息。这一问题在医学图像中尤其突出，因为医学图像中普遍存在着极端的前景-背景不平衡。

如图1(b)所示背景类较大且空间上不均匀，而前景类(紫色)较小且均匀。在这种情况下，如果不合理地平均出背景中不同局部区域(或部分)的不同外观信息，可能会导致前景背景边界预测的模糊性。不幸的是，这种类内局部信息的丢失在最近的许多工作中都存在，在这些工作中，每个类都被空间平均为一个一维表示原型或线性分类器的权重向量。为了解决这个问题，我们通过为每个类提取局部表示的集合来鼓励网络保留类内部的局部信息。

同时为了打破训练数据稀缺的僵局，提高分割精度，我们提出了一种自监督的Few-Shot医学图像语义分割框架SSL-ALPNet。该框架利用基于超像素的自监督学习(SSL)，使用超像素来消除手工标注的需要，以及自适应局部原型池化支持的原型网络(ALPNet)，通过在学习的表示中保留局部信息来提高分割精度。

如图1 (a)所示，为了确保通过自监督学习到的图像表示可以很好地推广到真实的语义类，我们使用超像素生成伪语义标签，超像素是语义对象的紧凑构建块。此外，为了提高学习到的图像表示的辨别能力，我们提出了一种基于一超像素的自监督学习方法，用于Few-Shot医学图像的分割。

此外，为了加强支持和查询之间表示的不变性，这对现实世界中的Few-Shot分割是至关重要的，我们通过应用支持和查询之间的随机几何变化和强度转换来合成新的形状和强度图像。此外，为了提高分割的精度，如图1所示，我们设计了自适应局部原型模块(adaptive local prototype module, ALP)，在每个类的原型表示中保留其局部信息。这是通过提取一个局部表现原型来实现的，每个原型关注不同的区域。值得注意的是，原型的数量由网络根据每个类的空间大小自适应分配。通过这种方式，ALP可以有效地解决局部信息不足导致的分割模糊问题。

方法

自适应局部原型池化（Adaptive local prototype pooling）

我们提出通过引入自适应局部原型池化模块(ALP)来保留原型中的局部信息。在ALP中，每个局部原型仅在覆盖在支持上的局部池窗口中计算，并且只表示感兴趣对象的一部分。具体来说，我们在每个$f_{\theta }\left({\mathbf{x}}_l^s\right)\in {\mathbb{R}}^{D\times H\times W}$上使用一个大小为$(L_H,L_W)$的池化窗。需要注意的是，$(L_H,L_W)$决定了在表示空间$\mathcal{E}$中计算每个局部原型所依据的空间范围。平均池化特征图在空间位置$(m,n)$上类$c$上的局部原型$p_{l,mn}\left(c\right)$可以表示为：

为了决定每个$p_{l,mn}\left(c\right)$的类$c$，我们平均池化前景类$c^{\hat{j}}$的二值掩码${\mathbf{y}}_l^s\left(c^{\hat{j}}\right)$到相同大小$\left(\frac{H}{L_H},\frac{W}{L_W}\right)$。让$y_{l,mn}^a$表示${\mathbf{y}}_l^s\left(c^{\hat{j}}\right)$在位置$(m,n)$上使用平均池化后的值，$c$被赋值为：

T为前景的下界阈值，根据经验设置为0.95。

为确保小于池化窗$(L_H,L_W)$的小物体可能被分配至少一个原型，我们同样使用掩码平均池化计算了类级(class-level)的原型$p_l^g\left(c^{\hat{j}}\right)$：

最后，为了方便起见，$p_{l,mn}\left(c^j\right)$和$p_l^g\left(c^{\hat{j}}\right)$用下标$k$重新索引，从而组成原型表示集合$\mathcal{P}=\left\{p_k\left(c^j\right)\right\}$。通过将不同的局部区域显式地表示为单独的原型，保留了更多的类内局部区别。

基于相似的分割（Similarity-based segmentation）

设计基于相似度的分类器$sim(\cdot ，\cdot )$，利用$\mathcal{P}$中的局部图像信息对查询进行密集预测，首先将每个原型与查询中的相应局部区域进行匹配，然后将局部相似性融合在一起。

为了在查询中分割一个大的肝脏，在第一阶段，一个类为$c^L=liver$的局部原型$p_k(c^L)$，当池化窗落在肝右叶上时，会发现一个类似于右叶的区域(而不是与整个肝脏匹配)。然后，要得到一个完整的肝脏，将右肝和左肝融合在一起，形成一个肝脏。

具体地，$sim(\cdot ，\cdot )$首先用查询特征$f_{\theta }\left({\mathbf{x}}^q\right)$和原型集合$\mathcal{P}=\left\{p_k\left(c^j\right)\right\}$作为输入来计算它们之间的局部相似图$S_k(c^j)$。与$f_{\theta }\left({\mathbf{x}}^q\right)$对应的空间位置$(h,w)$上的每个输入$S_k(c^j)(h,w)$由：

其中，$\odot$表示余弦相似，例如$a\odot b=\frac{⟨a,b⟩}{\Vert a{\Vert }_2\Vert b{\Vert }_2}$ $a,b\in {\mathbb{R}}^{D\times 1\times 1}$，α是一个乘法器，它有助于训练中的梯度反向传播。

然后，为了得到每个类$c^j$整体的相似图(未归一化)，将局部相似图$S_k(c^j)$分别融合到每个类中成为class-wise的相似$S^{\prime }(c^j)$：

${\mathrm{softmax}}_k\left[S_k\left(c^j\right)(h,w)\right]$表示先将所有$S^{\prime }(c^j)(h,w)$沿通道尺寸叠加，再计算沿通道的softmax函数。

为了得到最终的稠密预测，最后将类间相似度归一化为概率:

基于超像素的自监督学习（Superpixel-based Self-supervised Learning）

为了获得准确和鲁棒的结果，基于相似度的分类器需要两个特性。对于每个类，其表示应该是聚类的，这是为了再相似度量下拥有判别性；同时，这些表示在图像之间应该是不变的（在我们的例子中是指支持和查询的任何组合），以确保预测的鲁棒性。

这两个特性受到了基于超像素的自我监督学习(SSL)的鼓励。由于真实语义类的注释不可用，SSL利用伪标签在超像素级强制集群。这是通过基于余弦相似度的分类器反向传播分割损失自然实现的。在这里，超像素级的聚类属性可以转移到真实的语义类，因为一个语义掩码通常由几个超像素组成。此外，为了鼓励图像特征空间表示对图像之间的形状和强度差异保持不变，我们在支持和查询之间执行几何和强度转换。这是因为形状和强度是医学图像变化的最大来源。

Unspervised pseudolabel generation：
为了获得伪标签的候选对象，对每幅图像$x_i$生成超像素集合${\mathcal{Y}}_i=\mathcal{F}\left({\mathbf{x}}_i\right)$。使用$\mathcal{F}$表示的无监督算法可以有效地做到这一点。

Online episode composition：
对于每一个episode$i$，一副图像$x_i$和一个随机选择的超像素${\mathbf{y}}_i^r\left(c^p\right)\in {\mathcal{Y}}_i$用于形成支持集 S i = { ( x i , y i r ( c p ) ) } \mathcal{S}_{i}=\left\{\left(\mathbf{x}_{i}, \mathbf{y}_{i}^{r}\left(c^{p}\right)\right)\right\} Si​={(xi​,yir​(cp))}，这里${\mathbf{y}}_i^r\left(c^p\right)$表示带索引$r=1,2,3,\dots ,\left|{\mathcal{Y}}_i\right|$的二值掩码，$c^p$表示伪标签(pseudolabel)类（对应的背景掩码${\mathbf{y}}_i^r\left(c^0\right)$用$1-{\mathbf{y}}_i^r\left(c^p\right)$表示）。同时，查询集 Q i = { ( T g ( T i ( x i ) ) ) } , T g ( y i r ( c p ) ) ) \left.\mathcal{Q}_{i}=\left\{\left(\mathcal{T}_{g}\left(\mathcal{T}_{i}\left(\mathbf{x}_{i}\right)\right)\right)\right\}, \mathcal{T}_{g}\left(\mathbf{y}_{i}^{r}\left(c^{p}\right)\right)\right) Qi​={(Tg​(Ti​(xi​)))},Tg​(yir​(cp)))通过对支持集应用随机几何和强度变换构建。这样，每一个$(S_i,Q_i)$就形成了1-way 1-shot分割问题。在实际中，${\mathcal{T}}_g\left(\cdot \right)$包括仿射变换和弹性变换，${\mathcal{T}}_i\left(\cdot \right)$是gamma变换。

End-to-end training：
该网络是端到端的训练，其中每次迭代$i$以一个episode$(S_i,Q_i)$作为输入。采用交叉熵损失，每次迭代的分割损失${\mathcal{L}}_{seg}^i$为:

其中${\hat{\mathbf{y}}}_i^r\left(c^p\right)$表示查询伪掩码${\mathcal{T}}_g\left({\mathbf{y}}_i^r\left(c^p\right)\right)$的预测，我们还使用了原型对齐正则化：将预测反过来作为support，也就是${\mathcal{S}}^{\prime }=\left({\mathcal{T}}_g\left({\mathcal{T}}_i\left({\mathbf{x}}_i\right)\right),{\hat{\mathbf{y}}}_i^r\left(c^p\right)\right)$，它应该正确分割原始支持图像$x_i$，它是这样表示的：

其中${\overline{\mathbf{y}}}_i^r\left(c^p\right)$表示以$x_i$为查询的${\mathbf{y}}_i^r\left(c^p\right)$的预测。

总的来说，每次epioside的损失函数为：

实验结果

总结

在这项工作中，我们提出了一种自监督的少镜头医学图像分割框架。该方法在不需要任何手工标记训练的情况下，成功地超越了自监督学习在少镜头医学图像分割中的SOTA方法。
此外，在我们的实验中，它证明了对看不见的语义类的强泛化。此外，所提出的基于超像素的自我监督技术为图像表示学习提供了一种有效的方法，为未来半监督和无监督图像分割工作开辟了新的可能性。