1、前言
最近SAM 模型复现的多了,看了不少官方的源码,尝试下SAM和Unet的结合
SAM的提示分割,其实就是在分割的时候,为数据增加一个提示信息,可以是框,点,或者文本等等。这样大模型网络就可以根据提示信息进行分割
SAM模型的复现可以参考:视觉大模型学习笔记_听风吹等浪起的博客-CSDN博客
如下图,就是手动的给出提示框,然后网络分割出框里的内容
需要注意的是,本文不对SAM的原理进行讲解,只是借鉴SAM的提示分割思想,将其和unet结合
其实,思想很简单,一般的图像分割网络是3通道的,因为图像都是RGB的。
而这里需要传入sam的提示信息,所以还要增加一个维度,也就是说最开始输入神经网络的的通道是4维度的!
神经网络输入四维度很好实现,就是更改卷积核的个数罢了。毕竟正常神经网络隐藏层都是上百个,这里一开始和输入和中间隐藏层没啥区别,就是更改个数字罢了
实现了输入,最重要的是怎么把数据传进去!!!这里非常重要!!!
看了 MedSAM 源码实现,发现实现的思路确实很巧妙
大概思路就是,在图像上加上一层canvas,是全黑的。然后对mask的前景分割类别随机挑选一个,变成二值数据,这样就可以求取min和max,就可以得到box,将这个box传给canvas就可以达到提示框的目的
这里后面会根据代码介绍
2、SAM 提示框分割
SAM的提示信息分割有好几种,因为本人是医学图像专业,这里只实现了提示框分割。
其实点提示分割也很简单,明白了提示框分割的思路后,点分割就是怎么对数据处理的问题。或者就是把canvas的提示框换成其他mask模板。
点提示的话,可以用区域生长进行掩膜?
提示框的作用是什么呢?
说白了就是让网络更集中在一块区域, 对指定区域进行分割。
一般的语义分割,就是image和mask 一一对应,然后交叉熵、优化。这样网络同时要预测好几个类别,这也是为啥多类别分割效果不如二值分割好的原因
而提示框的box选定后,不管mask中有几个类别,这样mask都会变成二值的。
就像下面,分割的前景好十来个,但是我们绘制了box圈中小车后,那么对于网络来说,就是二值的分割。也就是前景车1 ,背景是除车的所有区域
代码实现的思路很简单,mask是0 1 2 3 4等等阈值图像,我们随机选中一个值,只取出这个值的数据。例如是3,那么除了3位置(x,y)得其他区域全部是0,3变成了新的阈值前景,这样不就是二值图像了嘛!!
至于二值图像,计算前景的min,max,就可以绘制边界框,这样就是sam提示框分割的精髓。代码实现如下
label_ids = np.unique(mask)[1:]
label_id = random.choice(label_ids.tolist())
mask = np.uint8(mask == label_id) # only one label
自动添加边界框:
y_indices, x_indices = np.where(mask > 0)
x_min, x_max = np.min(x_indices), np.max(x_indices)
y_min, y_max = np.min(y_indices), np.max(y_indices)
这样输入的数据就不是RGB图像,而是RGB+canvas(矩形边界框)的四通道数据
注意:这里的操作都是在dataset脚本实现的
3、Unet 网络
unet网络分割思想很简单,之前写了很多关于unet分割的介绍,自己参考就行
因为datase已经产生了四通道的数据,所有unet的输入是4维的,而输出是二值图像,所有outout是2维的
这里用1,因为后面用了 DiceCELoss 损失,而非交叉熵
Tips:
虽然Unet网络略显过时,什么DenseUne、TransUnet、SwinUnet等等都很好。
这里为了方便只是将sam和Unet结合,如果需要把sam和其他分割网络结合,把上面的model换成想要的分割模型,定义输入4,输出1就行
其他的代码不需要更改!
4、实验准备
环境如下,pip安装即可
matplotlib==3.9.2
monai==1.3.1
numpy==2.1.0
opencv_python==4.10.0.82
Pillow==10.4.0
torch==2.0.0
torchvision==0.15.1
tqdm==4.66.4
4.1 数据集
数据集使用的是自动驾驶车道线的分割,标签如下,可视化如下
0 background
1 white_lines
2 yellow_lines
数据集摆放如下:
--data--train---images 训练集的图像 --data--train---masks 训练集的图像标签 --data--val---images 验证集的图像 --data--val---masks 验证集的图像标签
4.2 训练
因为SAM提示框每次都是挑选一个前景分割,对于类别多的话,epoch最好设置的多点
参数如下:
parser = argparse.ArgumentParser(description="SamUnet segmentation")
parser.add_argument("--batch-size", default=8, type=int)
parser.add_argument("--epochs", default=200, type=int)
parser.add_argument('--lr', default=0.01, type=float)
parser.add_argument('--lrf',default=0.0001,type=float) # 最终学习率 = lr * lrf
parser.add_argument("--img_f", default='.jpg', type=str) # 数据图像的后缀
parser.add_argument("--mask_f", default='_mask.png', type=str) # mask图像的后缀
args = parser.parse_args()
训练集和验证集的比例约为3:1
数据的输入效果:
注意,这里前景应该是白色的,但由于白色不太好区别,这里填充为红色
loss和dice的曲线:
4.3 推理
这里编写了简单的ui推理界面,如下
这里会打印矩形框的坐标:
5、1 项目下载
项目为付费资源,包含数据集、完整代码、训练权重等等