0-ESP-MedSAM: Efficient Self-Prompting SAM for Universal Domain-Generalized Medical Image Segmentation

ESP-MedSAM: Efficient Self-Prompting SAM for Universal Domain-Generalized Medical Image Segmentation

ESP-MedSAM

arXiv:2407.14153v4 [eess.IV] 18 Aug 2024 IEEE TRANSACTIONS ON MEDICAL IMAGING ESP-MedSAM： 用于通用域广义医学图像分割的高效自提示 SAM

背景

SAM给医学图像分割提供了新的可能，但是SAM巨大的计算成本、对人工注释作为提示的需求以及容易发生冲突的解码过程降低了其在临床场景中的泛化性和适用性。为了解决这些问题，本文提出了一种高效的自提示 SAM，用于通用域广义医学图像分割，名为 ESP-MedSAM。首先设计了多模态解耦知识蒸馏 （MMDKD） 策略来构建一个轻量级的半参数共享图像编码器，为不同模态产生判别性视觉特征。此外引入了 Self-Patch Prompt Generator （SPPG） 来自动生成高质量的密集提示嵌入，用于指导分段解码。最后设计了查询解耦模态解码器 （QDMD），它利用一对一的策略为每种模态提供独立的解码通道。

  • MMDKD 策略，将模态特定和模态常识从基础模型分别提炼到模态控制器和模态聚合器中。两者都构成了一个轻量级的半参数共享图像编码器，从而为不同的医疗模式生成了独特的特征
  • 设计了 SPPG 来自动生成高质量的补丁提示，而无需手动注释。这些提示用于指导分割掩码的预测。
  • 用于分段解码的 QDMD。它利用一对一策略为每种模式提供私有分段工作流，防止不同模式相互干扰。
  • 提炼的轻量级半参数共享图像编码器 SPPG 和 QDMD 来构建ESP-MedSAM 具有显著的泛化-效率权衡。对各种医学成像模式进行了广泛的实验，证明的 ESP-MedSAM 在通用域广义医学图像分割方面优于最先进的技术

MedSAM 和 SAMMI 收集了超过 1M 的公共医学图像，以通过框和点提示对 SAM 进行全面微调，以实现域广义通用医学图像分割。然而，这种方法会迅速增加数据和计算成本，这在临床场景中既昂贵又不切实际。为了减轻迁移学习过程中对数据大小和计算资源的依赖，SAM 中引入了参数高效的微调技术。具体来说，Adapter 已被广泛用于集成到 SAM 的图像编码器中，以改进医学成像中的特征表示。但是无论那种方法都基于巨大的ViT编码器，有着复杂的计算，本文的方法克服了挑战，并说明了各种医学成像模式中卓越的泛化-效率权衡。

本文的方法还使用了知识蒸馏，

实验方法

ESP-MedSam

ESP-MedSAM 的端到端架构，用于通用医疗图像分割。

（a） 多模态解耦知识蒸馏。

  • 教师模型

多模态教师模型，该模型涉及 SAM 的自然图像编码器 FViT-SAM 和医学图像编码器 FViT-Med。在 SAM 中冻结图像编码器的权重，并额外创建一组可学习的标记。它们与多头注意力层的键和值通道连接，以学习医学图像中的新模式，从而更新注意力图。此外，将可学习的多层感知器并行到 ViT 中的前馈网络 （FFN） 中，以存储不同的模态信息。利用这个 FViT-Med 、 SPPG 和 QDMD 在医学源域 S 上实施微调。教师模型由焦点损失和骰子损失的组合进行监督，FViT-SAM 和 FViT-Med 能够分别提供模态常识和模态特异性知识。

  • 解耦知识蒸馏

分割掩码的质量取决于从图像编码器中提取的特征。从分而治之算法中汲取灵感，作者提出 MMDKD 方法将 KD 过程解耦为两个子任务：模态公共特征蒸馏和模态特定特征蒸馏。在模态共同特征蒸馏过程中，目标是将知识从 FViT-SAM 转移到模态聚合器 FMA，MMDKD 策略为每种医学成像模式提供了一个私有蒸馏通道，防止不同模式相互干扰。

  • 轻量级半参数共享图像编码器

为了无缝协调不同模态的特征表示，采用提炼的模态控制器和模态聚合器来构建一个轻量级的半参数共享图像编码器，总体而言，与 SAM 相比， MMDKD 方法有效地降低了图像编码器的计算成本，同时保留了生成可泛化特征图的能力。

（b） Self-Patch Prompt 生成器。

目前的医疗SAM和轻量级SAM主要利用手动提示（例如，点和框）来指导模型提供满意的分割掩码。然而，这些方法依赖于病理学家在医疗场景中的经验，这既昂贵又耗时。为了消除对手动注释的需求，作者使用包含补丁生成器和密集提示（DP）编码器的 SPPG，以自动生成一组高质量的补丁提示，以协助分割解码：

Self-Patch Prompt Generator

通过这种方式，SPPG 模块自动生成一组高质量的密集提示 p 来指导分割掩码的预测，提高了 ESP-MedSAM 在临床场景中的适用性。

（c） 查询解耦模态解码器。

SAM的掩码解码器利用与模态无关的查询标记来处理自然图像中的所有分割任务。但是，这并不是医学图像分割的最佳选择。由于各种医学成像模态存在固有的异质性，这种常见的预测通道存在解码冲突，降低了模型的泛化能力，ESP-MedSAM 框架提出了 QDMD。 QDMD 为每种模态提供了一个独立的解码过程，避免了不同模态之间相互冲突的固有异质性，提高了 ESP-MedSAM 框架的泛化能力。

实验结果

与通用医学图像分割（源域）中最先进的框架进行比较。

与域广义医学图像分割（看不见的域）中最先进的框架进行比较。

总结

通用医学图像分割的 ESP-MedSAM 框架引入了 MMDKD 策略，将基础模型中的知识提炼到模态控制器和模态控制器中，以创建一个用于判别特征生成的半参数共享图像编码器。然后，SPPG 被设计出来自动生成一组高质量的补丁提示，以协助分割解码。最后，QDMD 为每种模态定制了特定的分割工作流程。大量实验表明，ESP-MedSAM 的计算复杂度低于标准 SAM，在各种医学成像分割任务中优于 SOTA 任务特定架构和轻量级 SAM，表现出卓越的零样本泛化和通用能力。

代码分析

下载了两个有关眼底的数据集stare和drive（这里说一下stare数据集，给的tar文件里是gz文件然后图片是ppm格式，要先tar解压一次，在gunzip解压一次，在convert换格式成ppm，给数据集作者点👍）

1
2
3
4
#sudo bash this.sh
for ppm in *.ppm; do
    convert "$ppm" "${ppm%.ppm}.png"
done

作者给出的checkpoint并没有有关这两个数据集的，于是使用了ESP_retinal_best.pth这个PTH文件，但是仍然会报模型和权重不匹配的错误

1
2
Missing key(s) in state_dict: "module.SemiTViT.patch_embed.seq.0.c.weight", "module.SemiTViT.patch_embed.seq.0.bn.weight", "module.SemiTViT.patch_embed.seq.0.bn.bias", "mo......
Unexpected key(s) in state_dict: "SemiTViT.patch_embed.seq.0.c.weight", "SemiTViT.patch_embed.seq.0.bn.weight", "SemiTViT.patch_embed.seq.0.bn.bias", "SemiTViT.....

看着像是前面多了一个module.对吧，那么给他改一下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
 model = ESPMedSAM()
 if torch.cuda.device_count() >= 1:
     print("Let's use", torch.cuda.device_count(), "GPUs!")
     model = nn.DataParallel(model)

 state_dict = torch.load(args.model)
 new_state_dict = {}
 for k, v in state_dict.items():
     if k.startswith('module.'):
         new_state_dict[k[7:]] = v  # 去掉 'module.' 前缀
     else:
         new_state_dict[k] = v

 model.load_state_dict(new_state_dict, strict=True)
# model.load_state_dict(torch.load(args.model), strict=True)

 model = model.cuda()

再对数据集结构dataloder进行更改，最后运行结果如下：

1
2
3
4
5
6
LR SCALES: [0.08589934592000005, 0.10737418240000006, 0.13421772800000006, 0.1677721600000001, 0.20971520000000007, 0.2621440000000001, 0.3276800000000001, 0.4096000000000001, 0.5120000000000001, 0.6400000000000001, 0.8, 1.0]
Let's use 1 GPUs!
  0%|          | 0/2 [00:00<?, ?it/s]Could not load library libcudnn_cnn_infer.so.8. Error: libcuda.so: cannot open shared object file: No such file or directory
Please make sure libcudnn_cnn_infer.so.8 is in your library path!
/tmp/tmp3j0m_l6_: line 3: 42326 Aborted                 (core dumped) python /home/wuyueyu/git/ESP-MedSAM/test.py
ERROR conda.cli.main_run:execute(125): `conda run python ***/ESP-MedSAM/test.py` failed. (See above for error)

我确定我的 CUDA 和 cuDNN 是正确安装