[论文解读] U-Netmer: U-Net meets Transformer for medical image segmentation

Sheng He, Rina Bao+2 more|arXiv (Cornell University)|Apr 3, 2023

Radiomics and Machine Learning in Medical Imaging被引用 6

一句话总结

U-Netmer 提出了一种新颖的 U-Net 与 Transformer 架构融合方法，用于医学图像分割，通过将 U-Net 应用于局部图像块并利用 Transformer 建模全局上下文，解决了 'token-flatten' 和 'scale-sensitivity' 问题。该方法在 7 个数据集和多种模态下实现了最先进性能，且通过不同尺度间的分割差异作为可靠置信度分数，无需真实标签即可对测试图像按难度进行排序。

ABSTRACT

The combination of the U-Net based deep learning models and Transformer is a new trend for medical image segmentation. U-Net can extract the detailed local semantic and texture information and Transformer can learn the long-rang dependencies among pixels in the input image. However, directly adapting the Transformer for segmentation has ``token-flatten" problem (flattens the local patches into 1D tokens which losses the interaction among pixels within local patches) and ``scale-sensitivity" problem (uses a fixed scale to split the input image into local patches). Compared to directly combining U-Net and Transformer, we propose a new global-local fashion combination of U-Net and Transformer, named U-Netmer, to solve the two problems. The proposed U-Netmer splits an input image into local patches. The global-context information among local patches is learnt by the self-attention mechanism in Transformer and U-Net segments each local patch instead of flattening into tokens to solve the `token-flatten" problem. The U-Netmer can segment the input image with different patch sizes with the identical structure and the same parameter. Thus, the U-Netmer can be trained with different patch sizes to solve the ``scale-sensitivity" problem. We conduct extensive experiments in 7 public datasets on 7 organs (brain, heart, breast, lung, polyp, pancreas and prostate) and 4 imaging modalities (MRI, CT, ultrasound, and endoscopy) to show that the proposed U-Netmer can be generally applied to improve accuracy of medical image segmentation. These experimental results show that U-Netmer provides state-of-the-art performance compared to baselines and other models. In addition, the discrepancy among the outputs of U-Netmer with different scales is linearly correlated to the segmentation accuracy which can be considered as a confidence score to rank test images by difficulty without ground-truth.

研究动机与目标

• 为解决视觉 Transformer 中因 1D tokenization 导致的局部图像块交互信息丢失问题，即 'token-flatten' 问题。
• 克服基于图像块的 Transformer 中因固定图像块尺寸导致的 'scale-sensitivity' 问题，即分割精度随块尺寸变化显著。
• 构建一种统一的模型结构，支持多尺度图像块并共享参数，以提升模型鲁棒性。
• 基于不同尺度下分割的一致性，生成测试图像的置信度分数，实现无需真实标签的图像难度排序。
• 在多种医学影像模态和器官上展示最先进的分割性能。

提出的方法

• 将输入图像在多个尺度（s = 1, 2, 4, 8）下划分为非重叠的局部图像块，其中 s=1 表示完整图像。
• 每个局部图像块独立使用 U-Net 主干网络进行分割，以保留块内空间交互信息。
• 通过 Transformer 编码器建模所有局部图像块之间的全局上下文关系，注意力机制关注所有局部图像块的嵌入表示。
• 通过融合 U-Net 各图像块的输出结果，并结合全局注意力特征，生成最终的分割结果。
• 置信度分数通过计算不同尺度下（如 s=1 与 s=2）的分割图之间的 Dice 系数获得，作为分割精度的代理指标。
• 模型采用端到端训练，各尺度间共享参数，实现在无需微调的情况下进行多尺度推理。

实验结果

研究问题

• RQ1 U-Net 与 Transformer 的全局-局部融合结构是否在医学图像分割任务中优于标准 U-Net 及混合模型？
• RQ2 不同尺度下分割图之间的差异是否与实际分割精度具有相关性，从而实现无需真实标签的置信度估计？
• RQ3 所提出的模型是否能在不重新训练的情况下，保持在不同图像块尺度下的稳定性能？
• RQ4 基于尺度差异的置信度分数在图像难度排序方面是否优于现有方法？
• RQ5 该模型是否能在多种器官和影像模态之间实现良好泛化？

主要发现

• U-Netmer 在涵盖 7 种器官和 4 种影像模态（MRI、CT、超声、内窥镜）的 7 个公开数据集上实现了最先进性能。
• 在所有评估数据集中，该模型在 4 个基线模型和 6 个最先进模型上均表现出一致的性能提升。
• 不同尺度下（如 s=1 与 s=2）分割图之间的差异与实际 Dice 精度呈强线性相关，6 个数据集上的皮尔逊相关系数超过 0.8。
• 基于尺度差异的置信度分数在 6 个数据集中优于 ProtoSeg 的平均分割能力得分，在图像难度排序方面表现更优。
• 低置信度图像通常具有小尺寸、低对比度或模糊的目标区域，而高置信度图像则表现出清晰、大尺寸且结构明确的目标。
• 该模型在不同尺度间生成一致预测的能力，使其能够可靠识别出需要人工介入审查的困难分割案例。