MIS微调SAM模型实时交互UI界面

前言

SAM模型的基本介绍可见SAM（Segment Anything Model）大模型使用--point prompt_sam大模型-CSDN博客

针对Meta团队去年发布的SAM大模型在医学图像分割领域表现性能较差的情况，笔者收集了一些MIS领域的数据集对SAM的架构进行fine tune，目前已完成肺部，心脏，脾脏，结肠数据集的微调，于是就写了一个UI交互的界面脚本，目前已经发布在github上。

github地址

GitHub - jaychan0/fine_tune_SAM_MIS: This is a UI project, which uses medical image to fine tune SAM in hugging face structure.

训练说明

数据集主要是来自MSD开源数据集的CT扫描文件，处理的流程为：

把CT扫描文件分为多个切片，取有mask的切片作为备选数据集；

对于不同的数据集，根据文件的大小随机生成指定序号的数量，从备选数据集中获得使用数据集；

对于使用数据集，按照8:1:1的比例划分为训练集、验证集和测试集；

使用3060显卡训练，lr为1e-5，每个数据集训练60个epoch；

训练得到的预训练数据权重都已经上传到网盘中，github上有相应的链接和提取码。

依赖环境

需要说明的是，本项目对于SAM模型的微调不同于现在常见的直接在SAM源码上进行fine tune，而是使用hugging face封装好的系列接口，选择这个微调方式的原因是hugging face中的数据集预处理直接可以通过processor函数一条龙做好，自己再封装一个class就可以直接丢进model里进行训练，再者，它也提供了非常方便地config接口直接调整或更改SAM的结构，对于后续进一步要提高SAM对于下游任务的性能提供了非常方便地途径。

运行交互UI界面，最好本地有独显，笔者本地的是1650Ti显卡（cpu没跑过，估计速度会慢一些）

pytorch架构、transformer库依赖环境，直接从transformer库下载pretrain的权重需要hugging face的API token，在这里我也把原始的轻量化权重vit_b上传到网盘中了。

运行界面及结果

功能在按钮里已经明显了，这里就不赘述了

这里选择肺部的图片进行演示，注意，图片和GT的大小必须是 $256\times 256$ ！

选择完待分割图片后点一下UI界面上对应的矩形区域，图片才会显示，这样便于显示后续实时交互box prompt

直接用鼠标在待分割图像上进行box的输入，命令行会实时输出box的信息，即方框左上角的坐标和右下角的坐标

下一步，选择权重类型，我们先来选择原始权重

之后点击“生成掩膜”，就会看见SAM生成的图片，如果我们想保存生成的图片，就可以点击“保存图片”的按钮

可见原权重对于肺部的分割效果不太好，那我们再使用经过微调之后的肺部权重看看效果如何

可见，微调后的SAM对于肺部图像的分割性能有了显著的提升。

扩展

如果大家有自己的下游任务微调好的权重想要有直观的交互显示，也可以直接在“自定义权重选择”按钮中进行选择，前提是必须是使用hugging face提供的API接口进行微调的vit_b。

最后，如果大家也有在MIS领域微调好的权重，非常希望可以联系我，我可以把大家的权重加到这个体系当中，不断地完善这个系统；对于个人来说，可能无法做到像Meta团队那样构建庞大的数据集进行训练从而获得在通用图像领域高性能的大模型SAM，但是我觉得我们可以汇集大家的力量，不必要只限于单个checkpoint，把不同的MIS分割任务汇集起来，对应不同的权重，那么在MIS领域也可以有我们自己的Medical Segment Angthing Model。

欢迎大家讨论交流~