简介本资源面向医学图像分割方向的研究者与开发者提供一套在Unet基础上引入SAM提示框机制的息肉肿瘤语义分割完整方案用于提升肿瘤边界识别精度。压缩包共约2000个文件以1992张jpeg医学图像及分割标注为主另含5个Python源码文件、2个txt说明与1份readme整体约263.61MB覆盖数据集、网络结构定义、训练与推理流程以及带UI的交互式infer脚本。训练阶段系统自动生成边界框提示帮助网络聚焦肿瘤区域特征推理阶段用户可在界面中手动框选目标实现更精确的提示分割。源码包含数据加载、模型搭建、训练测试与界面设计等关键模块便于复现实验结果或在此基础上继续改进。目前已有773人学习下载适合具备一定深度学习基础、希望快速上手医学图像分割实践的研究者参考。1. 息肉分割只靠 U-Net 为什么总在边缘翻车做肠镜息肉分割的同行大概率都经历过这个场景U-Net 在 Kvasir-SEG 上 Dice 能刷到 0.90 以上可一放到边缘模糊、反光强烈、或者只有几帧清晰画面的小息肉上掩码就开始「糊边」——要么多切一块正常黏膜要么把息肉头部啃掉一角。问题不在 U-Net 的编码器不够深而在于它没有「提示」这个概念模型不知道你关心的是哪一块区域只能靠卷积核自己猜。SAMSegment Anything Model的出现给了另一条路用提示框bounding box prompt把目标位置告诉模型让分割从「全图盲猜」变成「框内精修」。把 SAM 的提示机制接进 U-Net 改进流程本质上是给语义分割加了一个轻量的空间先验让息肉肿瘤这类边界敏感任务少走弯路。这篇笔记面向已经跑过 U-Net、手里有息肉数据集、想试提示式分割的从业者从数据准备、提示框生成、模型改造到训练排查一步步拆开讲源码和数据集的组织方式也会给到可复现的路径。2. 把 SAM 提示框接进 U-Net结构改造与数据流2.1 为什么不是直接微调 SAM而是拿它当提示生成器直接微调 SAM 做息肉分割显存和标注成本都劝退SAM 的 ViT-H 主干在 1024×1024 输入下单卡 24G 只能塞下 batch size 1而且息肉数据集通常只有几百到几千张微调大模型极易过拟合。更务实的做法是把 SAM 当成「提示框生成器」或「伪标签精修器」U-Net 仍然是主力分割网络。具体有两种接法。第一种是离线提示用目标检测器YOLOv8 或 GroundingDINO先出息肉候选框把框喂给 SAM 得到高质量掩码再把这些掩码作为软标签去训练 U-Net。第二种是在线提示在 U-Net 的跳跃连接或解码器输入端把提示框编码成一个空间注意力图和特征图做逐元素调制。前者落地快、对训练框架改动小后者端到端但需要改网络结构。我一般先用离线方案验证提示框对 Dice 的增益确认有效再考虑在线融合。提示如果数据集本身带边界框标注如 Kvasir-SEG 的 JSON 里有 bbox 字段离线方案几乎零额外标注成本只有掩码没有框的数据集可以用掩码外接矩形自动生成框。2.2 提示框编码从 bbox 到空间注意力图的最小实现在线方案的核心是把[x1, y1, x2, y2]变成一个和特征图同尺寸的权重图。最直接的做法是生成一个二值掩码框内为 1、框外为 0再经过几层卷积得到软注意力。下面是一个可嵌入 U-Net 解码器的提示编码模块import torch import torch.nn as nn import torch.nn.functional as F class BoxPromptEncoder(nn.Module): 把归一化 bbox 编码成与特征图同尺寸的空间注意力图 def __init__(self, in_channels, feat_size64): super().__init__() self.feat_size feat_size # 用 3x3 卷积平滑框边界避免硬边导致训练震荡 self.smooth nn.Sequential( nn.Conv2d(1, 8, kernel_size3, padding1), nn.BatchNorm2d(8), nn.ReLU(inplaceTrue), nn.Conv2d(8, 1, kernel_size3, padding1), nn.Sigmoid() ) # 把提示图映射到特征通道维度做通道级调制 self.channel_gate nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(1, in_channels, kernel_size1), nn.Sigmoid() ) def forward(self, boxes, feat): boxes: (B, 4) 归一化坐标 [x1, y1, x2, y2]范围 0~1 feat: (B, C, H, W) 待调制的解码器特征 B, C, H, W feat.shape device feat.device # 生成二值框掩码 prompt torch.zeros(B, 1, H, W, devicedevice) for i in range(B): x1, y1, x2, y2 boxes[i] # 坐标反归一化到特征图尺度 ix1, iy1 int(x1 * W), int(y1 * H) ix2, iy2 int(x2 * W), int(y2 * H) # 防止越界和空框 ix1, iy1 max(0, ix1), max(0, iy1) ix2, iy2 min(W, ix2), min(H, iy2) if ix2 ix1 and iy2 iy1: prompt[i, 0, iy1:iy2, ix1:ix2] 1.0 # 平滑 通道调制 prompt self.smooth(prompt) # (B,1,H,W) gate self.channel_gate(prompt) # (B,C,1,1) return feat * gate feat # 残差式调制保留原始特征这段代码的关键设计有三点。第一框掩码用smooth做软边硬边会让梯度在框边界处突变训练 loss 容易抖。第二channel_gate把空间提示压成通道权重相当于告诉网络「这个框内的通道响应要放大」而不是直接替换特征。第三最后用feat * gate feat的残差形式即使提示框不准网络也能退回到原始特征不会因为一个坏框把整张图带偏。参数上feat_size不需要显式传入实际由feat的 H、W 决定smooth里的 8 通道是经验值息肉数据量小于 1000 张时可以降到 4减少过拟合。如果框的坐标是像素值而非归一化值记得在 dataloader 里先除以原图宽高。2.3 数据集组织Kvasir-SEG 与 CVC-ClinicDB 的目录规范息肉分割最常用的公开数据集是 Kvasir-SEG1000 张和 CVC-ClinicDB612 张两者都提供原图、掩码Kvasir-SEG 额外有 bbox 的 JSON。为了让提示框方案能跑通建议统一成下面的目录结构polyp_dataset/ ├── images/ │ ├── train/ │ └── val/ ├── masks/ │ ├── train/ │ └── val/ ├── boxes/ │ ├── train/ # 每张图一个 .txt内容 x1 y1 x2 y2归一化 │ └── val/ └── splits/ ├── train.txt # 每行一个样本 ID └── val.txt如果只有掩码没有框用下面这段脚本从掩码反推外接矩形注意要过滤掉面积过小的连通域否则会把噪点也框进去import cv2 import numpy as np import os def mask_to_bbox(mask_path, min_area30): mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) h, w mask.shape _, binary cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY) contours, _ cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) boxes [] for cnt in contours: if cv2.contourArea(cnt) min_area: continue x, y, bw, bh cv2.boundingRect(cnt) # 归一化并做 2% 外扩给 SAM 留上下文 pad_x, pad_y 0.02 * w, 0.02 * h x1 max(0, (x - pad_x) / w) y1 max(0, (y - pad_y) / h) x2 min(1, (x bw pad_x) / w) y2 min(1, (y bh pad_y) / h) boxes.append([x1, y1, x2, y2]) return boxesmin_area30是针对 512×512 以下分辨率的经验阈值高分辨率图像要按比例放大。外扩 2% 是为了让 SAM 看到息肉周围的黏膜纹理框贴太紧反而会让 SAM 把边界判到息肉内部。生成完框后建议抽查 20 张可视化确认框没有把整个画面框住——这种情况通常来自掩码反相或阈值选错。3. 训练配置损失、学习率与提示框的配合3.1 Dice BCE 的组合权重怎么定息肉分割的类别极不平衡背景像素通常是前景的 5 到 10 倍单用 BCE 会让模型偏向预测背景。常见做法是0.5 * BCE 0.5 * Dice但在提示框方案下我一般把 Dice 权重提到 0.7因为提示框已经提供了位置先验模型更需要被「逼」着去贴合边界而 Dice 对区域重叠更敏感。class DiceBCELoss(nn.Module): def __init__(self, dice_weight0.7): super().__init__() self.dice_weight dice_weight self.bce nn.BCEWithLogitsLoss() def forward(self, pred, target): # pred: (B,1,H,W) logits; target: (B,1,H,W) 0/1 bce_loss self.bce(pred, target) pred_sig torch.sigmoid(pred) intersection (pred_sig * target).sum(dim(2, 3)) union pred_sig.sum(dim(2, 3)) target.sum(dim(2, 3)) dice (2 * intersection 1e-6) / (union 1e-6) dice_loss 1 - dice.mean() return self.dice_weight * dice_loss (1 - self.dice_weight) * bce_loss平滑项1e-6不能省小息肉在 batch 内可能只占几十个像素没有平滑项时 dice 会算出 0/0。如果验证集 Dice 在 0.85 附近震荡上不去先把dice_weight降到 0.5 试两个 epoch观察是边界问题还是定位问题。3.2 学习率与提示框噪声的关系提示框来自检测器或掩码反推本身带噪声。学习率设太大时模型会过度信任提示框把框外的正确预测也压掉。我一般用1e-4起步配合 cosine 退火前 5 个 epoch 做 warmup。如果用的是离线 SAM 伪标签方案伪标签边缘可能有 1 到 2 像素的系统性偏移这时可以把学习率再降到5e-5让模型有足够时间从原图里纠正。注意提示框编码模块的参数是随机初始化的如果和 U-Net 主干用同一个学习率前期提示模块的输出会很不稳定。稳妥做法是给提示模块单独设1e-3的学习率让它快速学会「框内加权」这个基本行为再和主干一起退火。3.3 一个可复现的训练循环骨架import torch from torch.optim import Adam from torch.optim.lr_scheduler import CosineAnnealingLR def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for imgs, masks, boxes in loader: imgs imgs.to(device) masks masks.to(device) boxes boxes.to(device) optimizer.zero_grad() # 模型 forward 接收 boxes内部调用 BoxPromptEncoder preds model(imgs, boxes) loss criterion(preds, masks) loss.backward() # 梯度裁剪防止提示模块初期梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() return total_loss / len(loader) # 提示模块和主干分组设学习率 prompt_params list(model.prompt_encoder.parameters()) backbone_params [p for n, p in model.named_parameters() if prompt_encoder not in n] optimizer Adam([ {params: prompt_params, lr: 1e-3}, {params: backbone_params, lr: 1e-4} ]) scheduler CosineAnnealingLR(optimizer, T_max50, eta_min1e-6)clip_grad_norm_的max_norm1.0是息肉分割的常用值如果训练日志里 grad norm 经常超过 5说明提示框坐标有异常值比如归一化时除了零要先检查数据管道。分组学习率在 PyTorch 里用 param groups 实现注意named_parameters过滤时字符串匹配要准确否则会把提示模块的参数也放进主干组。4. 避坑与排查提示框方案最容易翻车的五个地方4.1 现象训练 loss 正常下降但验证 Dice 卡在 0.80 不动原因通常是提示框在验证集上分布和训练集不一致。比如训练集用检测器生成的框验证集用掩码反推的框两者外扩比例不同模型学到的「框内权重」在验证时对不上。解决方法是统一框的生成方式或者在验证时也走同一套检测器。另一个隐蔽原因是验证集的框有越界ix2 ix1判断失败导致提示图全零模型退化成无提示 U-Net。4.2 现象小息肉面积小于图像 1%的 Dice 明显低于大息肉提示框对小目标本来应该帮助更大但如果框的外扩比例固定为 2%小息肉的框可能只比目标大几个像素SAM 或编码器拿到的上下文不足。解决方法是按目标面积自适应外扩面积小于 1% 时外扩 5%1% 到 5% 外扩 3%大于 5% 外扩 2%。这个策略在 Kvasir-SEG 上能把小息肉 Dice 拉高 3 到 5 个点。4.3 现象训练到 20 epoch 后 loss 突然变成 NaN先查提示框坐标里有没有 NaN 或 Inf。掩码反推框时如果掩码全黑cv2.findContours返回空列表后续计算可能产生除零。在 dataloader 里加一层过滤没有有效框的样本直接跳过或者给一个覆盖全图的默认框。另外BCEWithLogitsLoss在 logits 绝对值很大时也会溢出梯度裁剪能缓解但不能根治必要时把 logits 做 clamp。4.4 现象推理时单张图耗时比原 U-Net 增加 3 倍以上提示编码模块本身计算量很小耗时通常出在框的生成环节。如果在线方案里每张图都要跑一次检测器推理时间自然翻倍。落地时建议把检测器和分割网络解耦检测器可以量化或换小模型分割网络保持 FP16。如果业务允许把检测结果缓存下来同一视频帧序列里相邻帧复用框能省掉大量重复计算。4.5 现象换到自己的数据集后提示框方案还不如原 U-Net这通常不是提示机制的问题而是框的质量太差。自建数据集如果只有粗标注的掩码反推的框可能把多个息肉连在一起提示图变成一个巨大的矩形反而干扰分割。解决方法是先做连通域分析确保一个框只对应一个息肉多息肉图像要么拆成多个框分别编码要么在损失里对每个框单独算 Dice。我见过最极端的案例是掩码里混进了器械反光反推的框把整个画面框住模型直接学成「全预测前景」。5. 验证提示框是否真的起作用消融与可视化技巧判断提示框方案值不值得留在你的 pipeline 里不能只看最终 Dice要做两组对照。第一组是「有提示 vs 无提示」把BoxPromptEncoder的输出置零或直接去掉其他配置完全一致跑 3 个随机种子取均值。如果 Dice 提升小于 1 个点而推理链路多了检测器那这个方案在你的场景里可能不划算。第二组是「提示框扰动测试」把验证集的框随机平移 10% 或缩放 20%观察 Dice 下降幅度。下降越小说明模型对框的依赖越健康如果平移 10% 就掉 10 个点说明模型把提示框当成了硬约束框一偏就崩这种模型在真实场景里很脆弱。可视化上我习惯把提示图、预测掩码、真值掩码叠成三通道图提示图放红色通道预测放绿色真值放蓝色。框内预测准的地方会呈现青色绿蓝框外误检会偏绿漏检会偏蓝。下面这段代码可以直接生成这种叠加图import cv2 import numpy as np def visualize_prompt_result(img, prompt, pred, gt, save_path): img: 原图 HxWx3 prompt: 提示图 HxW0~1 pred: 预测掩码 HxW0/1 gt: 真值掩码 HxW0/1 h, w pred.shape overlay np.zeros((h, w, 3), dtypenp.uint8) overlay[..., 2] (prompt * 255).astype(np.uint8) # 红提示 overlay[..., 1] (pred * 255).astype(np.uint8) # 绿预测 overlay[..., 0] (gt * 255).astype(np.uint8) # 蓝真值 # 和原图做半透明混合保留解剖结构 blended cv2.addWeighted(img, 0.6, overlay, 0.4, 0) cv2.imwrite(save_path, blended)addWeighted的 0.6/0.4 是经验值原图太暗时可以调到 0.5/0.5。生成一批验证集的可视化后重点看三类样本框贴边但预测正确的、框偏移但预测仍正确的、框正确但预测崩掉的。第三类样本是排查重点通常是该区域纹理和训练集差异大或者提示编码器的感受野不够。还有一个容易被忽略的验证点把提示框方案和纯 U-Net 的失败样本做交集。如果两者失败的图高度重合说明瓶颈在数据或主干不在提示机制如果提示方案修好了纯 U-Net 的一部分失败样本但引入了新的失败样本就要权衡净收益。我自己的习惯是任何新模块上线前先在一个 200 张的固定验证子集上跑三遍记录每张图的 Dice 变化只保留那些「稳定提升且没有明显恶化」的配置。息肉分割这个方向数据噪声大、边界主观性强别指望一个提示框模块解决所有问题但它确实能把边缘那 2 到 3 个点的 Dice 抠出来而这几个点往往就是临床可用和不可用的分界线。希望帮到你。本文还有配套的精品资源点击获取