CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection论文解读和实验复现

背景

1. 医学图像数据集的增长：随着公开的医学图像数据集数量的增加，自动化器官分割和肿瘤检测技术得到了显著发展。
2. 数据集的局限性：尽管数据集数量增加，但每个数据集通常规模较小，且存在部分标注问题。这意味着不是所有的图像都有完整的标注信息，这对于训练全面的AI模型构成了挑战。
3. 肿瘤类型的多样性：现有模型通常只能处理特定类型的器官或肿瘤，缺乏对不同类型肿瘤的广泛研究。
4. 模型的泛化能力：现有模型往往只能处理特定数据集中的特定任务，缺乏跨数据集、跨领域的泛化能力。

要解决的问题

数据集相关问题

1. 标签不一致性
   • 索引和名称差异：不同数据集对同一器官标注的索引和名称不同，如胃在BTCV中标记为“7”，在WORD中为“5”；“postcava”（AMOS22）和“inferior vena cava”（BTCV）名称不同但指代同一结构。
   • 背景标注差异：在组合数据集时，某些器官在不同数据集中的背景标注不一致，如胰腺在Pancreas - CT和MSD - Spleen数据集中的标注情况。
   • 器官和数据重叠问题：器官之间存在重叠关系，如“Hepatic Vessel”属于“Liver”，“Kidney Tumor”是“Kidney”的一部分；部分数据集之间存在数据重叠但标注不同，如KiTS与AbdomenCT - 1K。
2. 标签正交性：传统的one - hot标签无法体现器官和肿瘤之间的语义关系，few - hot标签虽有所改进，但仍不能充分表达器官间的复杂关系。

模型性能问题

1. 分割准确性：提高模型对腹部器官的分割准确性，尤其是在多器官和肿瘤同时存在的复杂情况下，使模型能够准确识别和分割25种器官和6种肿瘤。
2. 肿瘤检测能力：增强模型对肿瘤的检测能力，在保证高灵敏度的同时减少假阳性，准确识别肿瘤病例，平衡检测的敏感性和特异性。

模型效率与适应性问题

1. 计算效率：提高模型的计算效率，使其在临床应用中能够更快地处理图像数据，实现快速推理，相比数据集特定模型有显著的速度提升（>6倍）。
2. 泛化能力：增强模型对不同来源CT扫描数据的泛化能力，使其能够在未见过的医院数据上表现良好，无需额外调整和适应。
3. 扩展性：使模型框架能够灵活应用于不同的骨干网络，如CNN和Transformer，便于在不同架构下实现高效的器官分割和肿瘤检测。

要做的任务，输入和输出，工作流程

任务

实现对25种器官和6种肿瘤的准确分割与检测，提高模型在腹部器官分割和肿瘤检测任务中的性能，包括准确性、效率、泛化性、可扩展性等方面。

输入

1. CT扫描数据：来自14个公共数据集的3410个CT扫描，以及用于测试的其他公共和私有数据集的CT扫描。
2. 文本提示：用于生成CLIP嵌入的医学提示模板，如“A computerized tomography of a [CLS]”等。

输出

1. 器官和肿瘤分割掩码：对输入CT扫描中的25种器官和6种肿瘤进行分割，生成对应的二值掩码，指示每个体素是否属于相应的器官或肿瘤。
2. 肿瘤检测结果：判断CT扫描中是否存在肿瘤，并提供患者级别的敏感性（Sensitivity）和特异性（Specificity）指标，用于评估肿瘤检测的准确性。

工作流程

1. 数据预处理
   • CT扫描处理：对CT扫描进行各向同性间距调整和强度归一化，将其重新切片为相同的体素大小（$1.5\times 1.5\times 1.5m{m}^3$），并截断强度范围至[-175, 250]后线性归一化到[0, 1]。同时，根据前景对象裁剪图像，在训练时随机裁剪固定大小（96×96×96）的区域。
   • 标签处理：统一14个数据集的标签索引，处理器官和数据重叠问题，将每个器官的分割结果表示为二值掩码。对于未区分左右器官的数据集（如KiTS、WORD、AbdomenCT - 1K、CT - ORG），通过脚本将器官（如肾、肾上腺、肺）拆分为左右部分。
2. 模型训练
   • 文本分支：使用预训练的CLIP文本编码器（如“ViTB/32”）和医学提示模板生成器官和肿瘤的CLIP嵌入，将CLIP嵌入与全局图像特征连接后输入多层感知器（MLP），生成用于控制视觉分支的参数${\theta }_k$。
   • 视觉分支：采用Swin UNETR作为视觉编码器，对预处理后的CT扫描进行特征提取，然后通过三个连续的$1\times 1\times 1$卷积层（text - driven segmentor）处理图像特征，结合文本分支生成的参数${\theta }_k$计算每个类别的预测概率$P_k$。
   • 损失计算与优化：采用掩码反向传播技术，计算预测掩码与真实标签之间的二进制交叉熵（BCE）损失，但仅对有标签的类别进行损失计算和反向传播，使用AdamW优化器和warm - up cosine调度器进行50个epoch的训练。

详解框架

CLIP-Driven Universal Model的框架结合了文本和图像信息，以提高医学图像分割和肿瘤检测的准确性和泛化能力。

1. 框架组成

文本分支（Text Branch）

• 预训练文本编码器：使用CLIP模型的预训练文本编码器来处理与医学图像相关的文本描述（如器官或肿瘤的名称），生成文本嵌入。
• 医学提示（Medical Prompts）：将文本描述格式化为特定的提示模板，例如“一个[类别]的计算机断层扫描图像”，其中[类别]是具体的器官或肿瘤名称。
• 文本嵌入生成：通过预训练的文本编码器和医学提示，为每个器官和肿瘤生成文本嵌入，这些嵌入捕获了文本描述的语义信息。

视觉分支（Vision Branch）

• 图像预处理：对CT扫描图像进行标准化处理，包括调整方向、体素大小标准化以及强度归一化。
• 视觉编码器：使用如Swin UNETR这样的网络结构作为视觉编码器，提取图像特征。
• 特征处理：通过一系列卷积层进一步处理视觉编码器提取的特征。

2. 特征融合与分割

• 特征对齐：将文本嵌入与图像特征进行对齐，通过连接操作实现，以便模型能够同时考虑文本和图像信息。
• 参数生成：使用多层感知机（MLP），将对齐的特征转换为分割任务的参数。
• 分割掩码预测：使用文本驱动的分割器（如1×1×1卷积层）来预测每个像素点属于特定器官或肿瘤的概率，形成分割掩码。

3. 训练与优化

• 损失函数：使用二元交叉熵（BCE）损失函数进行监督学习，并通过掩码反向传播技术只对有标注的类别进行损失计算和反向传播。
• 优化器：使用AdamW等优化器进行模型参数的优化。

4. 泛化与迁移能力

• 多数据集训练：模型在多个数据集上进行训练，以提高其泛化能力。
• 迁移学习：模型可以在新任务上进行微调，以利用在大规模数据集上预训练得到的特征表示。

5. 性能评估

• 评估指标：使用Dice相似系数（DSC）、归一化表面距离（NSD）、敏感性、特异性等指标来评估模型性能。
• 跨数据集验证：在多个外部数据集上进行测试，以验证模型的泛化能力。

6. 开源框架支持

• MONAI集成：模型开发和实验过程中使用了MONAI框架，这是一个为医学成像研究提供标准化、用户友好、可复现的深度学习工具的开源项目。

怎么解决这些问题

1. 整合CLIP文本嵌入：

   • 利用CLIP预训练模型生成的文本嵌入来编码器官和肿瘤的标签，这些嵌入能够捕捉图像和文本之间的语义关系，从而克服了传统one-hot编码的局限性。

2. 构建通用模型框架：

   • 提出了一个可以从多个部分标注的数据集中学习的通用模型，这个模型不仅能够处理特定的器官或肿瘤分割任务，还能够泛化到新的领域和任务。

3. 解决标签不一致性：

   • 通过统一不同数据集中的标签索引，将不同数据集的标注标准化，使得模型能够在合并的数据集上进行训练。

4. 掩码反向传播技术：

   • 为了解决部分标签问题，提出了一种掩码反向传播技术，该技术只对那些在数据集中有标注的类别进行损失计算和反向传播，从而优化模型参数。

实验

数据集和评估指标

• 训练数据集：使用了14个公开的CT扫描数据集，共计3410个扫描图像。
• 测试数据集：在3个额外的数据集上进行评估，包括6162个外部CT扫描图像。
• 评估指标：使用了Dice Similarity Coefficient (DSC) 和 Normalized Surface Distance (NSD) 来评估器官/肿瘤分割的效果；使用Sensitivity（敏感性）和Specificity（特异性）来评估肿瘤检测的性能。

实施细节

• 数据增强：包括改变CT扫描的方向、使用等间距重新切片每个扫描以统一体素大小，以及强度截断和归一化。
• 网络结构：详细描述了文本分支和视觉分支的网络结构，包括使用的预训练模型和自定义的网络层。
• 损失函数和优化器：使用了二元交叉熵损失函数，并采用AdamW优化器进行训练。

性能比较

• MSD和BTCV挑战：模型在Medical Segmentation Decathlon (MSD) 和Beyond The Cranial Vault (BTCV) 挑战中的表现。
• 与其他模型的比较：与现有的SwinUNETR模型（之前在MSD排行榜上排名第一）进行了性能比较。

器官分割性能

• MSD任务：提供了5折交叉验证的结果，并展示了在MSD各个任务上的DSC分数。
• BTCV任务：在BTCV验证集上与其他方法进行了比较，并提供了详细的DSC分数。

肿瘤检测性能

• 五个数据集上的肿瘤检测：在包含肝脏、肾脏和胰腺肿瘤的数据集上评估了模型的肿瘤检测性能，并提供了敏感性和特异性的调和平均值。

CLIP嵌入的有效性

• 嵌入空间的可视化：通过t-SNE可视化展示了使用CLIP嵌入与one-hot编码相比，如何更好地捕捉特征之间的聚类和解剖结构关系。
• 不同提示模板的效果：评估了不同文本提示模板对CLIP嵌入性能的影响。

模型的临床相关特性

• 效率：比较了模型的计算效率，使用FLOPs（浮点运算次数）作为指标。
• 扩展性：展示了模型可以扩展到不同的网络骨干，如CNN和Transformer。
• 泛化能力：在多个外部数据集上评估了模型的泛化能力。
• 迁移学习能力：评估了模型在不同下游任务上的迁移学习能力。

指标

1. Dice相似系数（DSC）：

   • DSC是一个统计量，用于比较两个样本集合的相似度，通常用于评估图像分割任务中预测的分割掩码与真实标注之间的一致性。
   • 它的计算公式为：
     $$\text{DSC}=\frac{2|X\cap Y|}{|X|+|Y|}$$
     其中 $X$ 是预测的分割掩码，$Y$ 是真实的分割掩码，$|X\cap Y|$ 是两者交集的体积，$|X|$ 和 $|Y|$ 分别是预测和真实分割掩码的体积。
   • DSC的值范围从0到1，值越高表示预测的分割与真实标注越相似。

2. 归一化表面距离（NSD）：

   • NSD是衡量两个表面之间距离的指标，通常用于评估图像分割任务中预测的分割边界与真实边界之间的接近程度。
   • 它的计算公式为：
     $$\text{NSD}=\frac{\sqrt{\frac{1}{N}{\sum }_{i=1}^N{d}_i^2}}{\text{spacing}}$$
     其中 $d_i$ 是预测表面和真实表面之间的距离，$N$ 是表面点的数量，$\text{spacing}$ 是图像的体素间距。
   • NSD的值越小表示预测的分割边界与真实边界越接近。

3. 计算敏感性 ( S )：
   $$S=\frac{\text{True Positives (TP)}}{\text{True Positives (TP)}+\text{False Negatives (FN)}}$$

4. 计算特异性 ( Sp )：
   $$Sp=\frac{\text{True Negatives (TN)}}{\text{True Negatives (TN)}+\text{False Positives (FP)}}$$

5. 计算和谐均值 ( H )：
   $$H=2\times \frac{S\times Sp}{S+Sp}$$

解释：

• 敏感性 ( S )：衡量模型检测出真正阳性（即实际有肿瘤且模型检测到肿瘤）的能力。
• 特异性 ( Sp )：衡量模型正确识别出实际没有肿瘤的案例的能力。
• 和谐均值 ( H )：提供了一个综合的效能评估，它考虑了模型的敏感性和特异性，给出一个单一的指标来衡量模型的整体性能。

和谐均值特别有用，因为它能够平衡敏感性和特异性，确保模型在两个方面都有良好的表现，而不是只在一个方面表现得很好。这对于医学诊断来说非常重要，因为过度诊断（增加FP）和漏诊（增加FN）都可能导致严重的后果。

复现

目录结构

pred_pseudo.py跑自己的数据集得到硬标签

复现模型在xx数据集的表现（论文提到的）

我选用的是MSD的task 09 Spleen

下载好模型权重和数据集，并放置在对应的文件夹下，再运行validation.py即可同时，获得DSC分数