[论文解读] Push the Boundary of SAM: A Pseudo-label Correction Framework for Medical Segmentation

Ziyi Huang, Hongshan Liu+8 more|arXiv (Cornell University)|Aug 2, 2023

Advanced Neural Network Applications被引用 5

一句话总结

本文提出了一种伪标签校正框架，通过提升基于Segment Anything Model（SAM）生成的伪标签质量，增强了医学图像分割性能。通过集成标签质量评估模块和基于不确定性的自我校正机制，该方法迭代地优化噪声标签并重新训练分割网络，在无需专家标注掩码的情况下实现了最先进性能。

ABSTRACT

Segment anything model (SAM) has emerged as the leading approach for zero-shot learning in segmentation tasks, offering the advantage of avoiding pixel-wise annotations. It is particularly appealing in medical image segmentation, where the annotation process is laborious and expertise-demanding. However, the direct application of SAM often yields inferior results compared to conventional fully supervised segmentation networks. An alternative approach is to use SAM as the initial stage to generate pseudo labels for further network training. However, the performance is limited by the quality of pseudo labels. In this paper, we propose a novel label correction framework to push the boundary of SAM-based segmentation. Our model utilizes a label quality evaluation module to distinguish between noisy labels and clean labels. This enables the correction of the noisy labels using an uncertainty-based self-correction module, thereby enriching the clean training set. Finally, we retrain the segmentation network with updated labels to optimize its weights for future predictions. One key advantage of our model is its ability to train deep networks using SAM-generated pseudo labels without relying on a set of expert-level annotations while attaining good segmentation performance. We demonstrate the effectiveness of our proposed model on three public datasets, indicating its ability to improve segmentation accuracy and outperform baseline methods in label correction.

研究动机与目标

• 解决零样本SAM与全监督分割模型在医学影像中的性能差距。
• 克服SAM生成的伪标签质量低下导致下游分割性能下降的局限性。
• 开发一种弱监督框架，提升分割精度而不依赖于专家标注的像素级真实标签。
• 通过迭代式标签校正，仅使用SAM生成的伪标签实现深度分割网络的鲁棒训练。
• 提升医学图像分割任务中对标签噪声的泛化能力与鲁棒性，适用于多种成像模态。

提出的方法

• 利用Segment Anything Model（SAM）在无需任何像素级标注的情况下，为医学图像生成初始伪标签。
• 引入标签质量评估模块，基于预测置信度和不确定性区分干净像素与噪声像素。
• 采用基于不确定性的自我校正模块，通过置信度感知的重加权策略，迭代校正噪声伪标签。
• 应用多层级重加权策略，在像素级和图像级并行操作，以控制过拟合并提升鲁棒性。
• 通过校正后的伪标签重新训练分割网络，形成自我改进的循环，以增强模型泛化能力。
• 集成重训练模块，支持使用更新后标签的多次优化循环，提升收敛速度与性能。

实验结果

研究问题

• RQ1 伪标签校正框架能否显著提升SAM生成的医学图像分割伪标签质量？
• RQ2 所提方法在无需专家标注的情况下，如何优于现有弱监督学习方法，实现对噪声伪标签的校正？
• RQ3 迭代式标签校正与重训练在多种医学影像数据集上能将分割精度提升到何种程度？
• RQ4 与顺序方法相比，并行多层级重加权策略是否能减少误差累积并提升鲁棒性？
• RQ5 该框架能否在X光、CT和超声等多种成像模态中实现有效泛化？

主要发现

• 所提框架在三个公开医学数据集（JSRT X光、肺部CT和超声乳腺癌数据集）上实现了最先进分割性能。
• 与基线方法（如SEPL、ADELE和PINT）相比，该方法在提升分割精度方面表现显著，尤其在修正低置信度和噪声像素方面优势明显。
• 通过校正并保留低权重像素而非直接丢弃，该框架丰富了训练集，增强了模型泛化能力。
• 增加一次重训练迭代带来的计算开销极低，使该框架适用于实时临床应用。
• 与PINT相比，该框架通过并行重加权而非顺序阶段，减少了误差传播，提升了效率。
• 该方法在不同成像模态中对多种解剖结构（包括肺组织和肿瘤区域）均表现出强大的泛化能力。