TMI2024 | 阿大、同济等提出TraCoCo，用于3D医学图像半监督分割

编辑 | ScienceAI

3D 医学图像分割方法已经取得了成功，但它们对大量体素级标注数据的依赖是一个需要解决的缺点，因为获取这些标注的成本很高。

半监督学习（SSL）通过使用大量未标注数据和少量标注数据进行模型训练，解决了这一问题。

最成功的 SSL 方法基于一致性学习，即通过最小化从扰动视图中获得的模型响应之间的距离来实现的。这些扰动通常在视图之间保持空间输入上下文的一致性，这可能导致模型从空间输入上下文中学习分割模式，而不是从前景对象中学习。

在最新的研究中，澳大利亚阿德莱德大学（The University of Adelaide）、同济大学等机构的研究人员提出了 TraCoCo，这是一种一致性学习 SSL 方法，它通过改变输入数据视图的空间输入上下文来进行扰动，使模型能够从前景对象中学习分割模式。

此外，该团队提出了一种新的“Confident Regional Cross entropy (CRC)”损失，该损失提高了训练收敛性，并保持了对共训练伪标签错误的鲁棒性。该方法在多个3D数据基准上达到了最先进的结果，例如左心房（LA）、胰腺CT（Pancreas）、和脑肿瘤分割（BraTS19）。该方法在一个 2D 切片基准——自动心脏诊断挑战（ACDC）上也取得了最佳结果，进一步证明了其有效性。

文章地址：https://ieeexplore.ieee.org/abstract/document/10695462

代码地址：https://github.com/yyliu01/TraCoCo

背景

3D医学图像分割方法虽然已取得成功，但它们依赖于大量体素级别标注数据，这是一个亟需解决的问题，因为获得这些标注的成本非常高。为了克服这一挑战，半监督学习（SSL）通过结合大量未标注数据和少量标注数据来训练模型，从而减少对标注数据的需求。

最成功的 SSL 方法基于一致性学习，通过最小化模型在未标注数据的不同扰动视图下的响应差异来实现一致性。然而，这些扰动通常保持视图之间的空间输入上下文较为一致，这可能导致模型从空间上下文中学习分割模式，而不是从前景对象中学习。

为解决这一问题，研究人员提出了一种新的 SSL 方法，称为 Translation Consistent Co-training（TraCoCo）。TraCoCo 通过改变输入数据视图的空间输入上下文来扰动视图，使模型能够从前景对象中学习分割模式。

同时，该研究还提出了一种新的 Confident Regional Cross-Entropy（CRC）损失，该损失旨在提高训练的收敛性并保持对共训练伪标签错误的鲁棒性。

Translation Consistent Co-training (TraCoCo)

TraCoCo是一种基于互助学习（Co-training）框架的方法，该框架通过翻译一致性（Translation Consistency）来扰动输入数据的空间上下文，从而减少模型对背景模式的“记忆”，确保模型专注于前景对象的分割。

模型框架

TraCoCo 包含两个初始参数不同的独立网络（通常为 VNet或 3D-UNet）。输入数据通过随机裁剪生成两个子体积（sub-volumes），这两个子体积的空间上下文不同，但在体素网格中存在重叠区域。

模型目标

为了确保模型在不同空间上下文下的前景分割结果一致，TraCoCo 引入了多种损失函数，包括监督学习损失、半监督学习损失和翻译一致性损失。本文的方法实现通过最小化以下损失函数来进行训练：

其中，分别由监督学习损失，半监督学习损失，和翻译一致性损失组成。

监督学习损失（Supervised Learning Loss）

监督学习损失通过带标注的数据集 D_L 进行计算，包括体素级交叉熵损失（Cross-Entropy Loss）和 Dice损失，用来优化模型的分割性能。损失函数定义为：

Translation Consistency Loss

Translation Consistency Loss 主要目标是确保模型在不同空间上下文下的分割结果一致。具体实现过程如下：从训练体积中随机提取两个中心不同的子体积

损失函数定义为：

其中，公式 4 是 Kullback-Leibler（KL）散度，用于计算两个子体积在重叠区域的分割结果之间的差异；公式 5 是基于负熵的正则化损失，用于平衡训练体素中的前景和背景类。

半监督学习损失（Semi-supervised Learning Loss）

半监督学习损失用于强化两个模型之间的分割一致性，Confident Regional Cross-Entropy (CRC) Loss, 其定义为：

3D CutMix

为了进一步提高训练的泛化性，本文采用了 3D CutMix 技术。实现为随机生成一个 3D 二值掩码，该掩码包含一个随机定义位置和尺寸的“1”方块。掩码应用在未标注数据和伪标签上，定义为：

实验

本文在四个公开的 3D 医学图像半监督分割数据集上进行了实验，包括：

左心房（Left Atrium, LA）：100 个 3D MRI 体积，其中 80 个用于训练，20 个用于测试。

胰腺 CT（Pancreas-CT）：82 个对比增强的 CT 扫描，采用特定的预处理步骤。

脑肿瘤分割 2019（BraTS19）：包含 335 个脑部 MRI 样本，每个样本包含四种扫描类型（T1, T1-ce, T2, FLAIR）。

自动心脏诊断挑战（ACDC）：该数据集包含 100 个心脏 MRI 扫描，本文使用了其 2D 切片进行半监督学习。

实验设置：采用了 VNet 和 3D-UNet 作为模型架构，实验评估指标包括Dice、Jaccard、平均表面距离（ASD）和 95% Hausdorff距离（95HD）。实验结果显示，TraCoCo 在这些基准测试上均优于现有的最先进（SOTA）方法，特别是在标注数据较少的情况下。

总结

本文提出了 TraCoCo 和 CRC 损失在 3D 和 2D 医学图像分割任务中的有效性，证明了其在减少对背景模式的“记忆”以及提高前景对象分割精度方面的能力。

未来的工作计划包括扩展 TraCoCo 到更多复杂的医学图像任务，并在多模态医学图像中探索其半监督学习能力 。