CV | Medical-SAM-Adapter论文详解及项目实现

本文主要讲解Medical-SAM-Adapter论文主要方法及项目实现~

2023.12.29第七版_Medical SAM Adapter:Adapting Segment Anything Model for Medical Image Segmentation

论文地址：https://arxiv.org/pdf/2304.12620v7

论文代码：https://github.com/KidsWithTokens/Medical-SAM-Adapter

1.论文详解

1.1.摘要

Segment Anything Model （SAM） 在图像分割领域广受欢迎，因为它在各种分割任务中具有令人印象深刻的功能和基于提示的界面。然而，最近的研究和个别实验表明，由于缺乏医学特定知识，SAM在医学图像分割方面表现不佳。这就提出了一个问题，即如何增强SAM对医学图像的分割能力。在本文中，没有对SAM模型进行微调，而是提出了医疗SAM适配器（Med-SA），它使用一种轻巧而有效的适应技术将特定领域的医学知识整合到分割模型中。在Med-SA中，提出了空间深度转置（SD-Trans）来实现2D SAM与3D医学图像的适应，以及超提示适配器（HyP-Adpt）来实现快速条件适应。对不同图像模态的 17 个医学图像分割任务进行了综合评估实验。Med-SA 的性能优于几种最先进的 （SOTA） 医学图像分割方法，而仅更新了 2% 的参数。

1.2.论文架构

主要讲解方法部分

论文基于SAM,依旧使用图像编码器、提示编码器和mask解码器。

• 图像编码器：基于标准ViT被MAE训练，这里使用ViT_h/16变种，有14×14窗口的注意力和4个等间距的全局注意力。图像编码器的输出是对输入图像的16倍下采样嵌入。

• 提示编码器：可以是稀疏的点、框、文本或密集的mask；本文中只关注稀疏提示编码器，它将点和框表示为位置编码，并对每个提示类型进行学习嵌入。

• mask解码器：是一个Transformer解码器，包括动态掩模预测头。

SAM 使用双向交叉注意，一个用于提示到图像的嵌入，另一个用于图像到提示的嵌入，在每个块中学习提示和图像嵌入之间的相互作用。在运行了两个块之后，SAM 上采样嵌入图像，MLP 将输出标记映射到一个动态线性分类器，从而预测给定图像的目标掩码。

Med-SA架构(Med-SA architecture)

为了对 SAM 架构进行微调以适应医学图像分割，本文没有完全调整所有参数，而是冻结了预先训练好的 SAM 参数，并在架构的特定位置插入了 Adapter 模块。

Adapter 是一个 bottleneck 结构，它依次包括：下投影、ReLU 激活和上投影。下投影使用简单的 MLP 层将给定的嵌入压缩到较小的维度；上投影使用另一个 MLP 层将压缩的嵌入扩展回其原始维度.

(在第一版里叫MSA,第七版里叫Med-SA)

2D Medical Image Adaption

在 SAM 编码器中，本文为每个 ViT 块部署了两个 Adapter。

修改标准 ViT block (a)，得到 2D Medical Image Adaption (b)

• 将第一个 Adapter 放在多头注意力之后、残差连接之前

• 将第二个 Adapter 放在多头注意力之后 MLP 层的残差路径上

紧接着第二个 Adapter 之后，按照一定的比例系数对嵌入进行了缩放，引入缩放因子 s 是为了平衡与任务无关的特征和与任务有关的特征。

Decoder Adaption

在 SAM 解码器中，本文为每个 ViT 块部署了Adapter，把这个叫做Hyper-Prompting Adapter

• 第一个 Adapter 部署在 prompt-to-image 嵌入的多头交叉注意之后，并添加了提示嵌入的残差。本文使用了另一种向下投影来压缩提示嵌入，并在 ReLU 激活之前将其添加到 Adapter 的嵌入上，有助于 Adapter 根据提示信息调整参数，使其更加灵活和通用于不同的模式和下游任务

• 第二个 Adapter 的部署方式与编码器完全相同，用于调整 MLP 增强嵌入

• 第三个 Adapter 部署在图像嵌入的残差连接之后，以提示交叉注意

• 另一个残差连接和层归一化在自适应后连接，以输出最终结果

SD-Trans

尽管 SAM 可以应用于病灶的每个切片以获得最终的分割，但是它没有考虑深度维中的相关性。

本文提出了一种新的适配方法，其灵感来源于 image-to-video adaptation，具体架构如 (c)。

在每个 block 中，本文将注意力操作分成两个分支：空间分支和深度分支。

2.项目实现

2.0.环境设置

Python3.8+docker容器（Ubuntu）

git clone https://github.com/KidsWithTokens/Medical-SAM-Adaptercd Medical-SAM-Adapterwget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
# 建议使用conda,如果自定义，需要重新设置要求文件 -> requestion.txtpip install -r requestion.txt

因为我没有使用conda新建虚拟环境，自己新建了一个requestion.txt,按照作者要求是torch版本是1.12，我是1.14，并删掉了安装torch的命令。(大家可以根据自己的环境进行自定义)

2.1.数据集设置

数据准备

ISIC数据集官网：https://challenge.isic-archive.com/

数据为医学影响类黑色素瘤的数据，对于黑色素瘤进行分割。

 新建文件夹data/isic

#下载数据集wget https://isic-challenge-data.s3.amazonaws.com/2016/ISBI2016_ISIC_Part1_Training_Data.zip
#wget https://isic-challenge-data.s3.amazonaws.com/2016/ISBI2016_ISIC_Part1_Training_GroundTruth.zip
wget https://isic-challenge-data.s3.amazonaws.com/2016/ISBI2016_ISIC_Part1_Test_Data.zip
#wget https://isic-challenge-data.s3.amazonaws.com/2016/ISBI2016_ISIC_Part1_Test_GroundTruth.zip

#下载csv文件https://github.com/KidsWithTokens/MedSegDiff/blob/master/data/isic_csv/ISBI2016_ISIC_Part3B_Test_GroundTruth.csv
https://github.com/KidsWithTokens/MedSegDiff/blob/master/data/isic_csv/ISBI2016_ISIC_Part3B_Training_GroundTruth.csv

进入文件夹解压缩： 

unzip '*.zip'

解压后，格式如图

2.2.训练

训练数据集1：ISIC2016

python train.py -net sam -mod sam_adpt -exp_name msa_isic -sam_ckpt ./checkpoint/sam/sam_vit_b_01ec64.pth -image_size 1024 -b 32 -dataset isic -data_path ./data/isic

一张单卡24GGPU的情况，batch size为2,17929MiB,如果现存较小，改小batch size或者image size.

100个epoch时，IOU:0.85,DICE:0.91 

采样结果

训练数据集2

python train.py -net sam -mod sam_adpt -exp_name msa_kvasir -sam_ckpt ./sam_vit_b_01ec64.pth -image_size 1024 -b 32 -dataset Kvasir-SEG -data_path /workspace/SAM/datasets/Kvasir-SEG

3.代码详解

论文的核心块adapter_block.py

位置：Medical-SAM-Adapter/models/ImageEncoder/vit/adapter_block.py

基于原本的transformer块添加窗口注意力和残差传播块，流程如下：

• 对输入张量 x 进行归一化操作。

• 进行注意力计算，包括多头注意力机制和相对位置嵌入。

• 2D图像的情况使用 Space Adapter 处理注意力计算的输出。

• 如果数据是 3D ，则对输入进行额外的处理，并将其与 Space Adapter 处理后的结果相加。

• 对处理后的张量再次进行归一化操作。

• 将张量传递给 MLP 模块，并加上 MLP Adapter 处理后的结果的 scaled version。

• 将上述两项相加，并返回最终的输出张量。

可能遇到的问题及解决PS

[PS1] ValueError: num_samples should be a positive integer value, but got num_samples=0

错误分析：csv文件问题

解决方案：

重新下载csv文件就可以啦，不能使用wget方式下载

[PS2] TypeError: unsupported operand type(s) for %: 'int' and 'NoneType'

或者出现 ZeroDivisionError: integer division or modulo by zero错误时：

解决方案：

训练命令需要填写或修改vis自定义参数，默认是None,不能填写0，修改为1

想要了解更多内容，可在小程序搜索🔍AI Pulse,获取更多最新内容。