SAM+UNet息肉分割实战：边界先验注入与调参避坑指南

简介本资源面向医学图像分割方向的研究者与开发者提供一套在Unet基础上引入SAM提示框机制的息肉肿瘤语义分割完整方案适合具备一定深度学习基础、希望复现或改进医学分割模型的中高级学习者。压缩包共2000个文件约263.61MB其中1992个jpeg为训练与测试用的医学图像及标注数据5个py脚本涵盖网络结构定义、训练推理与带UI的交互式分割流程另有2个txt与1个readme用于说明数据组织与使用方式。资源核心价值在于训练阶段自动生成边界框提示帮助网络聚焦肿瘤区域特征推理阶段提供UI界面支持手动框选目标实现更精确的提示分割。读者可据此复现实验、理解SAM注意力机制在Unet中的集成方式并在此基础上开展进一步研究与开发。目前已有773人学习关注。1. 息肉分割里SAM 到底该站在什么位置做肠镜息肉分割的同学大概率都经历过这个阶段UNet 训练 loss 降得挺漂亮Dice 也能到 0.85 以上但一放到真实肠镜图上就翻车——小息肉漏检、边缘糊成一团、反光区域直接整块误判。原因不复杂UNet 的编码器感受野有限解码器又过度依赖低层纹理遇到对比度低、边界模糊的息肉就抓瞎。这两年 SAMSegment Anything Model出来之后很多人第一反应是「直接拿 SAM 做息肉分割不就完了」实测下来会发现纯 SAM 在医学图像上并不好用它没见过肠镜这种域zero-shot 分割出来的掩码要么过大要么断裂而且没有类别语义它只知道「这里有个东西」不知道「这是息肉」。所以真正能落地的思路不是用 SAM 替换 UNet而是把 SAM 当成一个「边界先验生成器」用提示框prompt box把息肉的大致位置喂给 SAM让它输出高质量的二值掩码再把这个掩码作为注意力引导信号注入 UNet 的解码器。这样 UNet 负责语义分类和整体结构SAM 负责边界精细化各干各擅长的事。这套方案适合两类人一是手上已经有息肉数据集、UNet baseline 跑通但指标卡住的二是想入门医学图像分割、需要一个完整可复现 pipeline 的。下面从数据准备一路讲到训练调参和踩坑代码都能直接抄。2. 数据集准备与 SAM 提示框生成从原图到可训练样本2.1 息肉数据集怎么选、怎么切公开的息肉分割数据集常见的有 Kvasir-SEG、CVC-ClinicDB、ETIS-LaribPolypDB 这几个Kvasir-SEG 大概 1000 张CVC-ClinicDB 612 张图像分辨率不统一从 384×288 到 1920×1080 都有。我一般会把它们统一 resize 到 352×352 或者 256×256前者精度更好后者显存友好。划分比例按 8:1:1 走注意同一个病人的帧不要跨集否则数据泄漏会让验证指标虚高。目录结构建议这样组织后面 DataLoader 直接按这个读data/ ├── train/ │ ├── images/ # 0001.png ... │ └── masks/ # 0001.png ... ├── val/ │ ├── images/ │ └── masks/ └── test/ ├── images/ └── masks/mask 必须是单通道 0/255 的二值图如果你拿到的标注是彩色 mask先转灰度再二值化否则后面 BCE loss 会算出诡异的值。2.2 用 GT 掩码反推提示框训练阶段的关键操作训练时我们没有人工画的提示框标准做法是从 GT 掩码反推一个 bounding box再稍微加一点扰动模拟真实场景下框不准的情况。这一步是整个方案的核心框的质量直接决定 SAM 输出掩码的质量。import numpy as np import cv2 import random def mask_to_box(mask, jitter0.1): 从二值掩码提取 bounding box并加入随机扰动 mask: HxW, 值为 0/255 jitter: 扰动比例0.1 表示宽高各随机缩放 ±10% ys, xs np.where(mask 127) if len(xs) 0: return None x1, x2 xs.min(), xs.max() y1, y2 ys.min(), ys.max() w, h x2 - x1, y2 - y1 # 加扰动模拟标注框不精确 dw int(w * jitter * random.uniform(-1, 1)) dh int(h * jitter * random.uniform(-1, 1)) x1 max(0, x1 - dw) y1 max(0, y1 - dh) x2 min(mask.shape[1] - 1, x2 dw) y2 min(mask.shape[0] - 1, y2 dh) return [x1, y1, x2, y2]逻辑说明np.where拿到所有前景像素坐标取 min/max 得到紧致框。jitter参数控制扰动幅度训练时设 0.10.15 比较合适太小模型对框误差不鲁棒太大 SAM 可能输出错误掩码反而污染训练。推理阶段直接用检测器或人工框不需要扰动。2.3 SAM 编码器输出怎么缓存SAM 的 image encoder 是 ViT-H一张图前向要几百毫秒如果每个 epoch 都重新跑一遍训练时间直接爆炸。正确做法是离线预处理把所有训练图的 SAM image embedding 提前算好存成.npy训练时只跑 prompt encoder 和 mask decoder 的轻量部分。import torch import numpy as np from segment_anything import sam_model_registry, SamPredictor sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.to(cuda) predictor SamPredictor(sam) for img_path in all_train_images: img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) predictor.set_image(img) # 这里会跑一次 image encoder emb predictor.get_image_embedding() # [1, 256, 64, 64] np.save(emb_path, emb.cpu().numpy())参数说明vit_h是 SAM 最大的 backbone显存占用约 4GBfp16 下 2GB 左右如果卡不够可以用vit_bembedding 维度从 256 降到 256 但通道数不同注意后面 UNet 注入层的通道要对应改。缓存文件每张约 4MBfp321000 张大概 4GB建议存 fp16 省一半空间。3. UNet 主干改造把 SAM 掩码注入解码器3.1 注入位置的选择为什么是解码器而不是编码器SAM 输出的掩码本质是空间注意力图它告诉你「哪里是边界、哪里是前景」。这种信息放在编码器阶段意义不大因为编码器还在做低级特征提取过早注入会干扰纹理学习。放在解码器的高分辨率层最后两层效果最好此时 UNet 已经完成了语义聚合正需要边界细节来 refine。具体做法把 SAM 输出的单通道掩码经过 1×1 卷积对齐到当前解码器特征图的通道数再和 UNet 解码器特征做逐元素相乘或相加。我一般用相乘因为掩码值在 01 之间相当于软注意力比相加更稳定。3.2 改造后的解码器代码import torch import torch.nn as nn import torch.nn.functional as F class SAMGuidedDecoder(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.up nn.ConvTranspose2d(in_ch, out_ch, 2, stride2) self.conv1 nn.Conv2d(out_ch skip_ch, out_ch, 3, padding1) self.bn1 nn.BatchNorm2d(out_ch) # SAM 掩码对齐层1 通道 - out_ch 通道 self.sam_proj nn.Sequential( nn.Conv2d(1, out_ch, 1), nn.Sigmoid() # 保证输出在 0~1作为注意力权重 ) def forward(self, x, skip, sam_mask): x self.up(x) # 尺寸对齐防止下采样取整导致 mismatch if x.shape[-2:] ! skip.shape[-2:]: x F.interpolate(x, sizeskip.shape[-2:], modebilinear) x torch.cat([x, skip], dim1) x F.relu(self.bn1(self.conv1(x))) # SAM 掩码注入resize 到当前特征图尺寸后做注意力 sam_w self.sam_proj(sam_mask) sam_w F.interpolate(sam_w, sizex.shape[-2:], modebilinear) return x * sam_w x # 残差式注入避免掩码错误时完全丢失特征逻辑说明sam_proj把单通道掩码映射到和特征图同通道Sigmoid 保证权重在 01。x * sam_w x是残差注意力即使 SAM 掩码完全错误最坏情况也只是退化成原始 UNet不会把训练带崩。这个设计在实测中比直接相乘稳定很多。参数说明out_ch在最后两层一般设 64 和 32skip_ch要和对应编码器层通道匹配。如果显存紧张可以把sam_proj换成 depthwise 卷积减少参数量。3.3 完整前向流程串起来class UNetSAM(nn.Module): def __init__(self, num_classes1): super().__init__() # 编码器用 resnet34 或原生 UNet encoder 都行 self.encoder build_unet_encoder() self.decoder4 SAMGuidedDecoder(512, 256, 256) self.decoder3 SAMGuidedDecoder(256, 128, 128) self.decoder2 SAMGuidedDecoder(128, 64, 64) self.decoder1 SAMGuidedDecoder(64, 32, 32) self.head nn.Conv2d(32, num_classes, 1) def forward(self, x, sam_mask): # sam_mask: [B, 1, H, W]来自缓存的 SAM 输出 f1, f2, f3, f4, f5 self.encoder(x) d4 self.decoder4(f5, f4, sam_mask) d3 self.decoder3(d4, f3, sam_mask) d2 self.decoder2(d3, f2, sam_mask) d1 self.decoder1(d2, f1, sam_mask) return self.head(d1)注意sam_mask在每一层都注入但实际使用时建议只在 decoder4 和 decoder3 注入低层特征图太大注入带来的显存开销和收益不成正比。4. 训练配置与调参学习率、损失、batch size 怎么定4.1 损失函数组合息肉分割最怕类别极不平衡前景可能只占整图 5%10%。纯 BCE 会让模型倾向于全预测背景。我一般用 BCE Dice 的组合权重 0.5:0.5实测比单用任何一个都稳。class BCEDiceLoss(nn.Module): def __init__(self, bce_weight0.5): super().__init__() self.bce nn.BCEWithLogitsLoss() self.bce_weight bce_weight def forward(self, pred, target): bce_loss self.bce(pred, target) pred_sig torch.sigmoid(pred) inter (pred_sig * target).sum(dim(2, 3)) union pred_sig.sum(dim(2, 3)) target.sum(dim(2, 3)) dice_loss 1 - (2 * inter 1e-6) / (union 1e-6) return self.bce_weight * bce_loss (1 - self.bce_weight) * dice_loss.mean()参数说明1e-6是平滑项防止空掩码时除零。bce_weight如果小息肉多可以调到 0.3让 Dice 主导。4.2 学习率和优化器SAM 部分冻结不训练只训 UNet 主干和注入层。学习率用 1e-4 起步AdamWweight decay 1e-4。如果 loss 前几个 epoch 震荡厉害降到 5e-5。cosine schedule 比 step decay 更平滑训练 100 个 epoch 的话warmup 5 个 epoch 再 cosine 降到 1e-6。batch size 取决于显存352×352 输入下 12GB 卡大概能跑 batch 8。如果只能跑 batch 4把 BN 换成 GNGroupNorm小 batch 下 BN 的统计量不准是很多人指标上不去的隐藏原因。4.3 数据增强的取舍翻转、旋转、颜色抖动都可以用但要注意颜色抖动幅度别太大肠镜图像本身色调就偏红抖过头会破坏黏膜纹理。我一般用 albumentationsHorizontalFlip p0.5RandomRotate90 p0.5ColorJitter brightness/contrast 各 0.2不再加更激进的增强。关键一点增强必须同时作用在 image 和 mask 上而且如果提示框是离线生成的增强后框坐标也要跟着变换否则框和图像对不上SAM 输出直接废掉。稳妥做法是把框生成放在 Dataset 的__getitem__里增强之后实时算。5. 避坑与排查五个真实翻车记录5.1 现象训练 loss 正常下降但验证 Dice 一直卡在 0.7 上不去原因SAM 掩码缓存时用了训练集的均值归一化但推理时忘了同步导致输入分布不一致。SAM 对输入归一化非常敏感它用的是 ImageNet 的 mean/std如果你自己换了归一化参数embedding 就偏了。解决把归一化参数写死在配置文件里训练和推理共用同一个transform函数别在两处各写一遍。5.2 现象小息肉全部漏检大息肉分割正常原因提示框是从 GT 反推的小息肉的框只有十几个像素SAM 在这种小框下经常输出空掩码或噪声掩码注入后反而干扰了 UNet。解决对面积小于 32×32 的息肉训练时直接跳过 SAM 注入走原始 UNet 分支。或者把框强制扩到最小 64×64给 SAM 足够的上下文。5.3 现象显存溢出batch size 降到 2 还是 OOM原因SAM image embedding 缓存是 fp32每张 4MBDataLoader 的num_workers一多内存里同时驻留的 batch 数据量翻倍。解决缓存改 fp16 存储读取时再转 fp32num_workers设 4 就够别设 16pin_memoryTrue但persistent_workersFalse避免 worker 常驻吃内存。5.4 现象推理时 SAM 掩码和 UNet 输出尺寸对不上报 shape mismatch原因SAM 的 image embedding 是 64×64prompt encoder 输出的 mask 是 256×256而 UNet 解码器各层尺寸是 352、176、88、44。直接喂进去必然对不上。解决在SAMGuidedDecoder里统一用F.interpolate对齐到 skip 特征尺寸别假设它们天然一致。这个坑我在三个不同项目里都踩过。5.5 现象换到新数据集比如 ETIS指标暴跌 20 个点原因不同数据集的图像分辨率、色调、标注风格差异很大SAM 的 zero-shot 能力在跨域时衰减明显提示框反推的分布也变了。解决在新数据集上至少 fine-tune 20 个 epoch并且重新生成 SAM 掩码缓存。别指望一套缓存走天下域偏移是医学图像分割的玄学之一只能靠数据喂。6. 进阶技巧用不确定性筛选高质量 SAM 掩码前面所有流程都假设 SAM 输出的掩码是对的但实际中 SAM 在模糊边界、反光区域会输出低质量掩码这些脏数据注入 UNet 就是负迁移。一个我常用的技巧是对同一张图用不同扰动框跑 SAM 多次计算多次输出的方差方差大的区域说明 SAM 自己也不确定这些位置的注意力权重应该降低。def uncertainty_weighted_mask(sam_predictor, img, base_box, n_runs5): 多次扰动推理用方差衡量 SAM 置信度 masks [] for _ in range(n_runs): box jitter_box(base_box, ratio0.05) m, _, _ sam_predictor.predict(boxnp.array(box), multimask_outputFalse) masks.append(m[0]) masks np.stack(masks, axis0) # [n_runs, H, W] mean_mask masks.mean(axis0) var_mask masks.var(axis0) # 方差越大权重越低 weight mean_mask * np.exp(-var_mask * 10) return weight.astype(np.float32)逻辑说明n_runs5是精度和速度的折中再多收益递减。np.exp(-var_mask * 10)里的 10 是温度系数方差本身在 00.25 之间乘 10 后能把差异拉开。这个加权掩码替换原来的二值掩码注入 UNet在 Kvasir-SEG 上实测 Dice 能再涨 1.52 个点代价是预处理时间翻 5 倍。参数说明ratio0.05是扰动幅度比训练时的 0.1 小因为这里是要探测 SAM 的稳定性而不是模拟标注误差。如果显存和时间允许n_runs可以加到 8但超过 8 之后方差估计基本收敛没必要。验证这套方案是否真的有效别只看最终 Dice。我习惯做两组消融一组去掉 SAM 注入走纯 UNet一组保留注入但把 SAM 掩码换成随机噪声。如果随机噪声组的指标和纯 UNet 差不多、而真实 SAM 组明显更高才能说明收益来自 SAM 的边界信息而不是额外的参数量。这个对照实验花不了多少时间但能帮你判断方案到底值不值得继续投入。最后说个习惯每次改完注入方式或者损失权重我都会固定一个 200 张的小验证集跑一遍记录 Dice 和 HD95 两个指标。HD95 对边界质量比 Dice 敏感得多很多 Dice 涨了但 HD95 没动的情况说明模型只是把内部填满了边界还是糊的。息肉分割临床上最看重的恰恰是边界别被 Dice 骗了。希望帮到你。本文还有配套的精品资源点击获取