半监督医学图像分割方法新SOTA｜校正对比伪监督学习

Paper: RCPS: Rectified Contrastive Pseudo Supervision for Semi-Supervised Medical Image Segmentation
PDF: https://arxiv.org/pdf/2301.05500v1.pdf
Code: https://github.com/hsiangyuzhao/RCPS

导读

医学图像分割方法通常需要全监督训练方式以保证模型性能，但这需要大量的专家进行数据标注，成本高且费力。半监督图像分割可以利用大量未标记图像和有限的标记图像来缓解该问题。然而，由于伪标签中的潜在噪声和特征空间中的类可分离性不足，从大量未标记图像中学习鲁棒性表示仍然具有挑战性，这会极大影像半监督分割方法的性能。

为了解决上述问题，本文提出了一种新的半监督分割方法，称为校正对比伪监督（RCPS），它结合了校正伪监督技术和体素级别的对比学习来提高半监督分割模型的性能。具体而言：本文设计了一种基于不确定性估计和一致性正则化伪监督方法的新策略，以减少伪标签中的噪声影响；此外，本文还在网络中引入了双向体素对比损失，以确保特征空间中的类内一致性和类间判别性，从而增加了图像分割中的类可分离性。

所提出的 RCPS 分割方法已在两个公共数据集（MRI左心房分割，CT胰腺分割）和一个内部临床数据集上（创伤性脑损伤TBI患者的大脑分割）进行了大量实验，实验结果表明：与半监督医学图像分割中的最新方法相比，所提出的方法具有更好的分割性能。

创新点

1. 提出了一种校正伪监督技术（不确定性估计和一致性正则化），来学习图像空间中不同分割对象的鲁棒性表示，进而提高半监督模型的分割性能。
2. 提出了一种双向体素对比损失和一种置信负采样策略，以提高特征空间中不同类别之间的语义可分离性。即对于正样本的两个不同视图，其距离空间应该相近，且远离特征空间中的负样本。

方法剖析

RCPS网络架构

假设分割数据集中包含 个有标签的数据, 和 个无标签数据 。为了表述方便, 我们假设有标签数据的子集表示如下:

无标签数据的子集则可表示如下:

其中, 表示训练数据, 表示训练标签（有标签数据）。半监督训练的目标就是：利用无标签数据 , 进一步提高只使用标签数据 训练出的模型的性能。

如上图1所示, 本文所提出的RCPS网络基于UNet网络架构, 只不过在UNet网络的上采样模块中加入了一个Projection Head用于计算对比损失。在RCPS网络的前向传播中, 包含四个输入：一个原始数据输入 (可理解为正样本), 原始数据输入 的两个视图增强数据 和 , 以及负样本数据 (这里输入负样本数据的作用是为了双向体素对比学习), 这四个输入数据对应的四个输出分别是: 和 。为了引入双向体素对比学习策略, Proje ction Head模块需要输出特征 (分别对应三个输入数据 的输 出）用于后续计算双向体素对比损失。

对于有标签数据的训练和损失函数, 则采用常用的交叉樀损失函数和Dice损失相结合作为总的分割损失函数即可:

其中, 表示交叉熵损失, 表示Dice损失。

对于末标记数据的训练和损失函数, 则计算校正的伪监督损失 和双向体素对比损失 。所以这部分数据的无监督损失可表示如下:

所以, RCPS网络的总损失函数可表示如下:

其中, 和 是可调节的超参数。

校正伪监督技术

伪监督

我们可以使用随机强度缩放、随机强度偏移和高斯噪声等操作获取原始数据的视图增强数据 和 。而伪监督技术则可以通过在和 之间应用监督损失函数来执行，因此伪监督损失定义如下：

其中, 是 的logits, 是softmax, 是temperature。

上述提到的伪监督方法有时不太可靠（因为模型预测伪标签可能会出现错误-俗称伪监督的标签噪声），所以本文在伪监督中引入了预测不确定性估计和一致性正则化技术，以减轻伪监督中标签噪声对结果产生的影响。

不确定性估计

伪监督容易受到标签噪声的影响，因此对伪标签进行高置信度阈值处理是降低标签噪声的常用方法。然而，简单地通过硬阈值对概率进行阈值化并不适合分割任务，因为不同语义类的分割难度是可变的。硬阈值化会使得困难类更难通过阈值，从而导致模型更加偏向于预测为背景类。

因此我们可以使用不确定性估计来校正伪监督，不确定性校正伪监督损失定义如下：

不确定性估计为伪监督损失引入了自适应体素权重，其中置信度较高的体素具有较高的权重，而置信度较低的体素具有较低的权重。该种技术降低了标签噪声的影响，提高了分割的鲁棒性。

一致性正则化

基于半监督学习中的平滑假设，我们期望 和 的预测应该是相似的，尽管它们进行了数据增强产生了不同的外观。因此，我们引入一致性正则化来进一步纠正伪监督，以最小化 和 。我们采用余弦距离来度量来自两个视图的距离：

最小化伪监督损失可以通过减少来自两个视图预测之间的余弦距离来实现，所以最终的本文提出的伪监督损失是不确定性校正损失和一致性正则化的线性组合：

双向体素对比学习

基于半监督学习中的平滑假设，除了追求标签空间的一致性之外，我们还期望在特征空间也需要有一致性。目前我们的模型在训练阶段无法提供足够的监督，使得正负样本的特征空间间距变大。相反，如果我们利用对比学习来学习类间不同的特征表示，开发一种双向体素对比损失则可以实现这一点。如上图 1 所示。来自 和 的空间对应位置的体素对被视为正值对，负样本 与正样本 和 的空间位置被视为负值对。我们的目标则是使得模型在特征空间中，拉近正值对的距离，推远负值对的距离。所以我们提出了一种双向体素对比损失来实现这一目标：

其中， 表示从特征图 提取出来的正体素， 表示负体素， 表示包含负样本的特征图， 是采样负体素的数量， 是temperature超参数。

实验结果

在LA数据集上，与其它方法相比，所提出的RCPS网络可以生成更平滑完整的分割边界，并在小结构中提供更精细的细节。

与 LA 数据集的结果类似，分割结果可视化对比表明：所提出的 RCPS 能很大程度上避免在分割过程中遗漏一些微小的分割目标，且能为微小结构生成更合理的分割图。

在TBI数据集上，UNet在我们提出的半监督设置下训练，甚至能够分割出全监督 UNet通常遗漏的结构，并生成质量更高的分割图。

结论

本文提出了一种用于医学图像半监督分割的新方法- RCPS（校正对比伪监督）。通过引入两种校正方法（不确定性估计法和一致性正则化），提高了伪监督模型的性能。此外，本文在特征空间中引入了双向体素对比损失，确保了特征空间中较大的类间距离和较小的类内距离，从而进一步提高了分割性能。大量实验表明，所提出的 RCPS 大大提高了分割性能，并优于最先进的半监督分割方法。

公众号后台回复“数据集”获取100+数据集分类汇总资源