尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

UNet+SAM融合:基于Point提示的腹部多脏器分割实践

UNet+SAM融合:基于Point提示的腹部多脏器分割实践 简介这是一套面向医学图像分割研究者的交互式分割项目代码基于UNet融合SAM点提示专注腹部多脏器分割任务通过点击前景/背景点即可引导模型生成精准掩膜适用于科研复现、课程设计或辅助标注场景。资源共1005个文件涵盖Python源码、预训练权重(pth)、JSON配置、训练日志、Readme说明以及近千张分割掩膜png和loss_iou_dice曲线、学习率衰减曲线等jpg图片压缩包约126.47MB已有237人学习。项目形成完整训练-评估-应用闭环支持余弦退火学习率调度采用Dice、IoU、Recall、Precision等指标评估训练100个epoch后平均Dice约0.93并输出loss/dice曲线、学习率衰减曲线便于复盘。推理脚本内置图形界面鼠标点击即可标记前景/背景并获得实时掩膜结果UNet架构支持自定义通道数与特征图尺寸数据加载器内置多种增强策略方便针对不同输入图像做扩展与改进对比。 我最近把一个老项目翻出来重构了一遍用UNet融合SAM模型做腹部多脏器分割核心卖点是让SAM的point提示参与推理过程。项目背景是腹部CT里的肝脏、脾脏、双肾、胰腺、胆囊这类器官分割传统UNet在小器官和模糊边界上总是差点意思而SAM这种通用大模型又缺乏医学分类意识单纯拿过来用会输出一堆无标签的mask。我的方案是把两者串起来——用SAM的point提示生成器官先验再喂给UNet做精细多分类。这篇文章就把这个项目的设计思路、融合细节、推理流程和踩过的坑全部梳理一遍适合正在做医学图像分割、或者想折腾UNetSAM组合的同学参考。1. 多脏器分割的痛点UNet的边界困境与SAM的“通才”错觉1.1 经典UNet在腹部多脏器分割中的表现与局限腹部多脏器分割和单器官分割完全不是一回事。单器官分割你只要把肝脏抠出来模型有充足的感受野去关注相对规整的边界多脏器分割要同时输出七八个标签器官之间相邻、灰度重叠、形变大问题一下子就复杂了。我最早直接用3D UNet跑全套流程在公开数据集上Dice能到0.95的肝脏掉到胰腺就只剩0.72左右胆囊甚至经常漏分割。UNet的瓶颈不在特征提取能力而在“全局上下文”不足。腹部CT每个器官的位置相对固定——肝脏在右上腹脾脏在左上腹双肾在脊柱两侧胰腺在腹腔深处——但UNet本质上是卷积堆叠的局部感受野组合即使加了skip connection它对“这个区域大概率是什么器官”的全局先验建模还是偏弱。训练数据不足时模型很容易把脾脏和胰腺混淆因为两者在窗宽窗位调过之后灰度分布高度重叠。另一个问题是边界模糊。腹部CT中很多器官之间的分界是靠筋膜、脂肪间隙和血管走的对比度非常低。普通UNet输出概率图在边界附近往往是“糊成一团”的后处理无论用条件随机场还是连通域分析都很难把细小的胰尾修出来。传统解法是多尺度输入、深监督、或者换Transformer架构但这些改动都治标不治本——模型还是不知道“胰腺应该长在哪个位置”。1.2 SAM模型引入医疗场景的适配问题SAM是视觉分割大模型输入一个点、一个框或者一键全图它就能输出高质量的二值mask。泛化能力确实强我拿腹部CT的window level调整后的切片试过SAM对肝脏、肾脏这种轮廓明显的器官能给出像模像样的mask但问题也很明显。第一SAM输出的是“通用物体”mask没有类别标签。你让它同时分割多脏器它默认会把图像里所有看起来独立的东西都切成mask分不清哪个是脾脏哪个是胰腺。第二SAM的医学图像迁移能力有限。它在自然图像上训练对于低对比度、高噪声的CT切片经常出现过度分割或欠分割胰腺这类器官的mask质量比UNet差不少。第三SAM的image encoder会把输入resize到1024x1024腹部CT原始切片如果直接喂进去解剖结构容易变形边界细节损失严重。但SAM有一个UNet不具备的优势prompt机制。你给它一个point它会优先围绕这个点展开分割。这个能力在医学场景里极有价值因为医生看CT时本来就会“指着一个区域说这是病灶”point提示正好能把这句人话翻译成模型能理解的东西。所以我的核心想法不是让SAM替UNet干活而是让SAM用point提示把“器官在哪里”这个先验给出来再由UNet负责“这个器官的精细边界和类别”。2. point提示融合的核心逻辑为什么一个“点”能改变分割结果2.1 SAM点提示机制与医学图像坐标对齐SAM的point prompt实现起来其实很简单你把一个二维坐标归一化到[0,1]区间和图像一起输入模型SAM的prompt encoder会把这个点编码成位置嵌入并在解码的时候用它来引导mask生成。关键在于每一个point都会影响最终mask的完整度。我做过一组实验同一个胰腺区域点在中心时SAM输出的mask还算完整点稍微挪到边缘mask立刻少了一半点一旦点到脂肪间隙SAM直接给出一个完全无关的区域。所以在医学图像上point的质量几乎决定了SAM先验的质量。坐标对齐这里有三个容易被忽略的细节SAM的输入图像会resize到1024x1024坐标也要按同样比例缩放不能直接拿原图坐标。CT图像要先做windowing把HU值映射到[0,255]否则SAM对灰度分布非常敏感。point所在的slice必须是原始体数据中对应层的切片如果做了插值重采样坐标要跟着spacing变化。我最终的做法是先把腹部CT体数据重采样到统一spacing比如1.5mm x 1.5mm x 3mm然后逐片推理。每个器官的point不是在切片上一个一个点的而是从体数据中器官的质心坐标出发映射到对应slice上再在质心周围做一个小范围的随机扰动生成3-5个候选点逐一输入SAM取最佳mask。2.2 UNet与SAM融合的三种做法及选型我调研了一圈UNet和SAM的融合大致有三种路子融合方式原理优点缺点Mask先验输入SAM生成二值mask作为额外通道与原图拼接输入UNet实现简单训练稳定UNet改动小需要额外维护SAM推理流程mask质量直接影响上限特征图融合SAM image encoder输出的特征图与UNet encoder特征做concat或attention能利用SAM的深层语义特征理论上更强显存占用极高训练难度大容易过拟合提示嵌入解码器将point embedding注入UNet decoder的每一层交互式推理最灵活需要修改UNet结构复现成本高我最后选了第一种Mask先验输入。原因是这个项目定位是“point提示推理”应用场景偏向批量处理预先标注好的器官质心而不是实时交互分割。Mask先验方案改动最小UNet还是那个UNet只是在输入层从单通道变成了“1 K”通道K是目标器官数量而且推理的时候可以随时替换SAM生成的mask工程上可控。注意力机制也值得提一句我没有在UNet内部额外加attention因为Mask先验已经扮演了“空间注意力”的角色。UNet的卷积会自动学到“当前像素靠近哪个器官先验区域”这个信息并把它和原始灰度特征结合起来。这样做的好处是SAM对小器官的错误mask不会被当作硬约束UNet有能力修正它。3. 项目整体架构与推理流程落地3.1 数据与预处理管线设计实验数据来自腹部CT一共用了80例手工标注了6个器官肝脏、脾脏、左肾、右肾、胰腺、胆囊。这个规模在医学图像领域算小的所以更依赖先验融合来提升泛化能力。预处理管线按这个顺序执行spacing重采样到统一分辨率我选的是1.5mm x 1.5mm x 3mm。太精细会导致切片数量爆炸太粗糙会丢失小器官细节。窗宽窗位设置成400/40 HU专门看腹部软组织。这个操作能把CT的16bit信息压缩到8bit同时把不同扫描设备的灰度差异拉齐。归一化到[0,1]减去均值除以标准差。生成器官质心坐标和切片映射表。这一步是为point提示准备的每个器官对应一组质心坐标推理时用。预处理时有个关键点不要把全腹部范围都塞进去。训练和推理都只取包含所有目标器官的slice范围一般是横膈膜到肾脏下缘这一段大概20-30张切片。这样UNet任务空间更集中SAM的point提示也不会跑到无意义的区域。3.2 融合模型的训练与微调策略SAM部分我选择完全冻结只当作一个“先验生成器”。原因很实在SAM参数量太大在80例小数据集上微调风险极高很容易灾难性遗忘而且训练时间消耗惊人。我不需要SAM变得更准只需要它老老实实根据point给出器官的粗定位mask。UNet部分我用的是标准U-Net编码器-解码器结构输入通道数改成了“1灰度 K器官先验”。训练时为了模拟推理阶段SAM生成的真实先验我没有直接用真实标签当先验通道而是对真实标签做随机形态学腐蚀和膨胀再加上高斯噪声。这个“带噪先验”训练策略非常关键它让UNet学会的不是“无脑复制先验”而是“结合灰度图像修正先验”。实测下来推理时SAM的噪声先验输入进去UNet输出依然稳定。训练损失用的是Dice Loss加Cross Entropy权重0.7:0.3。优化器选AdamW初始学习率1e-4配合余弦退火。batch size设成8因为先验通道是稀疏的显存占用比同分辨率纯CT图像高了不少。3.3 point提示推理的具体实现推理阶段是整个项目的核心环节我把它拆成三步走第一步生成point提示。对于批处理场景用预置的器官质心坐标直接映射对于交互场景让用户在某一层切片上点一下器官中心程序自动把该坐标映射到相邻切片。我做的系统支持批量模式比如加载完CT序列后程序在关键slice上为每个器官生成5个候选点每个点分别让SAM快速推理一次选mask面积最大的那个。第二步SAM生成先验mask。用sam_model_registry加载vit_b权重半精度模式推理。输入是windowed切片和归一化坐标输出二值mask。把K个器官的mask按顺序堆叠成K通道的tensor并resize到UNet期望的分辨率比如512x512。第三步UNet多类分割。把原始灰度图和K通道先验拼起来喂给UNet得到K1类概率图取argmax得到最终多脏器分割结果。核心代码如下# 伪代码基于PyTorch image windowing(ct_slice, window_width400, window_level40) image_tensor torch.from_numpy(image_normalize(image)).unsqueeze(0) # [1, H, W] prior_masks [] for organ_name, point in points.items(): mask sam_predict(image_tensor, point) # [1, H, W] prior_masks.append(mask) prior_tensor torch.cat(prior_masks, dim0) # [K, H, W] model_input torch.cat([image_tensor, prior_tensor], dim0).unsqueeze(0) # [1, 1K, H, W] logits unet(model_input) # [1, K1, H, W] seg torch.argmax(logits, dim1).squeeze(0)整个流程跑一遍单个体数据大约耗时7秒其中SAM占了5秒UNet只有2秒。如果追求实时可以把SAM的image encoder输出缓存下来只交换point embedding但这需要改SAM源码我没在这版做。4. 实测效果与调参经验4.1 多脏器Dice对比与边界指标我在20例测试数据上做了对照实验一组是纯UNet无先验输入另一组是融合SAM mask先验其余训练配置完全一致。结果差距比预期大器官纯UNet Dice融合方案 DiceHD95(mm)提升肝脏0.9580.96312.4 - 8.7脾脏0.9290.93816.1 - 11.2左肾0.9420.9519.6 - 6.4右肾0.9380.94710.2 - 7.5胰腺0.7210.80123.8 - 15.1胆囊0.7020.76918.5 - 12.9提升最大的是胰腺和胆囊这两个器官恰好是纯UNet最容易出错的地方。原因也很清晰胰腺体积小、周围脂肪多SAM虽然给不出精细边界但point提示能锁定大致位置UNet在小区域范围内的分类压力大大降低胆囊虽然边缘模糊但先验通道提供了明显的空间约束。边界质量方面融合方案在HD95上平均减少了35%说明模型输出的分割表面离真实边界更贴近了。目测效果也直观胰腺尾部的细长结构不再被切断胆囊和肝脏粘连处的分界更干净。4.2 显存、推理速度与病例差异显存是这个方案最需要关注的硬件指标。SAM的vit_b在1024x1024输入下半精度推理需要约6GB显存UNet在512x512输入下约3GB两者同时跑在单卡上显存峰值接近10GB。我是在NVIDIA RTX 3090上调试的24GB显存完全够用如果换成16GB卡需要把SAM输出mask的尺寸直接降到256x256再放大回512先验细节损失不大。推理速度上我做了两个优化用torch.cuda.amp.autocast()把SAM和UNet都跑在半精度整体提速约40%。对同一CT序列把多个slice的point推理打包成batchSAM的image encoder一次处理多张切片比逐slice推理快一倍。病例差异方面扫描设备不同导致CT值分布略有偏移windowing之后影响不大。比较麻烦的是增强扫描和平扫混在一起的情况血管和实质的灰度关系完全不一样模型泛化会变差。我现在只在这个项目里处理平扫数据增强扫描需要单独扩充训练集。4.3 值得注意的坑第一坑point位置对SAM mask质量太敏感了。如果只是用器官质心做point对偏心的器官比如胃旁边的胆囊经常给出错误先验。我的解决方案是生成多个候选点选SAM输出mask面积最大的那个这个简单策略能显著提高先验质量。第二坑UNet训练时的先验通道如果直接用真实mask推理时换成SAM输出会出现domain shift模型在推理时明显退化。后来改成“带噪先验”训练才解决。具体做法是对真实mask做随机腐蚀和膨胀半径随机取1到5像素再叠加零均值高斯噪声。第三坑多类先验输入时如果两个器官的mask存在重叠UNet输出在重叠区域容易混乱。我在生成SAM先验时做了处理按器官优先级排序后生成的器官只保留“不属于前面所有器官”的部分。这样先验通道天然互斥UNet分类就稳定很多。第四坑必须保证point和切片对齐。训练和推理如果用了不同的spacing重采样参数coordinates就会错位。我最后把质心坐标全部在原始坐标系下计算再通过affine变换映射到重采样后的空间避免手写映射。5. 从项目复盘看可复用方案5.1 哪些场景值得用这套融合路线这个方案不是万能药我复盘下来比较适合三类场景。一类是训练数据少但器官位置相对固定的项目例如腹部CT多脏器、头部结构、盆腔器官分割。这些器官的质心坐标解剖学先验明确point提示能发挥最大作用。二类是交互式分割工具。如果要做一套医生点击器官中心就自动完成精细分割的工具UNet加SAM的融合几乎是零成本起步方案不需要让医生逐像素标注。三类是难例修正。纯UNet模型对某些病例的某些器官分割特别差这时候可以在推理阶段单独为这个器官生成SAM point提示把先验通道叠加进去只修正难例不改动原有模型。5.2 后续可以怎么扩展后续可以再进一步的点包括把SAM的image encoder特征图也引入UNet而不是只限制在mask先验层面尝试用更轻量的SAM变体来降低显存占用还可以把point提示扩展为box提示因为有些临床场景医生更习惯画框框选器官。还有一个我特别想做的方向把这套融合推理接到大模型后面做成“自动规划point”的Agent。让大模型先分析CT影像结构自动推理出应该在哪个位置生成point再调用SAM和UNet完成分割。这样整个推理链路就真的串起来了。回到项目本身我最大的体会是UNet这种经典模型不需要推翻重来借助SAM的point提示机制一样能获得接近大模型的分割能力。关键在于把先验信息用对地方并让训练过程模拟推理时的“噪声环境”这样的融合才真正能落地到实际数据上。本文还有配套的精品资源点击获取
返回列表