简介面向眼底血管分割的 U-Net 深度学习项目基于 DRIVE 公开数据集训练模型可自动识别并分割眼底图像中的血管适合医学影像处理初学者、研究生及算法工程师参考。DRIVE 数据集包含 40 张带专业医生标注的彩色眼底图像按训练、测试划分覆盖不同年龄、性别及疾病状态具有较强代表性与挑战性。资源包大小约 120.72MB以 Python 脚本和 PyTorch 权重文件为主涵盖模型定义、训练流程、数据预处理与增强、预测推理等完整工程模块并附有训练完成的权重文件和相关输出目录结构清晰便于对照学习。目前已有 334 人学习下载。通过本项目可系统掌握 U-Net 架构、跳跃连接原理、医学图像预处理与增强方法、训练验证流程及模型评估思路还能体会深度学习在辅助眼科诊断中的应用价值快速上手同类医学图像分割任务。1. U-net在Drive眼底数据集上做分割20张训练图凭什么成为医学分割的入门标配很多人第一次接触医学图像分割时会被告知“U-net Drive眼底数据集”是必须跑通的项目。Drive只有20张训练图每张565×584像素目标是从彩色眼底照片里把细如发丝的视网膜血管分割出来。这个组合反直觉的地方在于数据这么少一个从零训练的U-net却能拿到接近0.95的AUC原因是血管结构高度重复且U-net的跳跃连接恰好保留了细节。它适合刚入门语义分割算法的人快速验证模型实现是否正确也适合想把医学分割落到业务里的工程师做基线。2. Drive数据集与预处理把565×584的眼底图变成U-net能吃到的张量2.1 Drive数据集的真实构成20张训练图、20张测试图与FOV掩码DriveDigital Retinal Images for Vessel Extraction是针对糖尿病视网膜病变研究发布的公开眼底数据集。核心组成是20张训练原图、20张测试原图每张都是8位RGB、565×584像素。训练集附两张专家标注文件名带manual1和manual2测试集只有一张标注。实际做项目时常见做法是把manual1当成标准答案manual2留着做标注差异分析如果只是想跑通U-netmanual1就够了。一个容易被忽略的细节是Drive的划分是官方写死的不要自己再随机切训练集和测试集。2020的划分和论文对比的价值就来自这里。你在网上看到有些人把40张图重新洗牌跑出来AUC特别高因为这个划分用到的训练样本变了跟官方结果没有可比性。图像内容上眼底照相的FOV是一个圆形区域直径大约540像素圆外是纯黑背景。任务里的“血管”只在圆内存在圆外不仅没有血管还经常带暗角、反光和成像噪声。这些区域如果不提前处理U-net会花大量参数去学“黑色背景不是血管”压缩了真正有效区域的表达能力。文件组织上不同下载渠道解压后的目录名不太一样常见做法是把原图和mask分开放。我会在项目根目录按下面的约定组织drive/ images/ # 训练原图 mask/ # 训练 manual1 标注 test/ # 测试原图 test_mask/ # 测试标注这个目录不是Drive官方结构只是个人习惯。好处是后续训练脚本和评估脚本拿到路径就能直接跑不用每次改文件名匹配逻辑。2.2 预处理三件套FOV掩码、绿色通道CLAHE、归一化眼底图像最难处理的问题是光照不均匀。同一个眼球中心亮、边缘暗血管在暗区的对比度远低于亮区。直接把RGB转灰度丢给U-net模型会优先学习光照分布而不是血管形态表现为边缘区域的血管预测不出来或者把暗区噪声当成血管。常见做法是用绿色通道而不是灰度图。原因很简单眼底照片里血管和背景对比最强的波段是绿色通道红色通道里血管几乎隐形蓝色通道噪声大。很多入门教程把三通道求平均转灰度能用但血管细节会丢尤其细血管的召回率会明显下降。如果你坚持用三通道输入预处理也可以做但要意识到绿色通道的信息会被卷积核摊薄效果不一定更好。代码import cv2 import numpy as np def preprocess_drive_image(img_path, mask_path, size(512, 512)): # 读取原始RGB眼底图 img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # FOV掩码原图里圆形区域是有效观测面 mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) mask cv2.resize(mask, size, interpolationcv2.INTER_NEAREST) mask (mask 0).astype(np.uint8) # 绿色通道是血管对比度最高的波段 g img[:, :, 1] g cv2.resize(g, size, interpolationcv2.INTER_AREA) # 只在FOV内做CLAHE避免圆外的黑色背景也被增强 clahe cv2.createCLAHE(clipLimit2.0, tileGridSize(8, 8)) g_clahe clahe.apply(g) g_clahe cv2.bitwise_and(g_clahe, g_clahe, maskmask) # 归一化到[0,1]单通道输入给U-net第一层 g_norm g_clahe.astype(np.float32) / 255.0 return g_norm[None, ...], mask这里面的参数值得解释。clipLimit2.0是CLAHE的经典取值越大对比度增强越强血管会更显眼但噪声也会同步放大。tileGridSize(8,8)表示把图像切成8×8的小块做直方图均衡块数越多局部细节越好但出现伪影的概率也越高。bitwise_and配合mask是为了让FOV外的背景彻底变成0这一步不是可有可无它告诉模型“圆外的黑色不参与语义判断”避免网络花参数去学习背景边界。resize到512×512同样是关键决定不是随便选的。U-net要做4次下采样输入尺寸必须是16的倍数。565×584在两次下采样后变成142×147左右第三次就会出现不对称上采样时和编码器特征拼接会错位。512×512是尽量贴近该比例又满足16整除的常用选择再小就会裁掉边缘血管。2.3 数据增强的边界旋转翻转是安全区弹性形变要克制20张图靠原始样本直接训练U-net一定过拟合。常见的增强方式包括随机旋转±30°、水平/垂直翻转、随机平移和缩放这些对血管分割都是安全的因为血管方向本来就有很强的随机性翻转旋转不会改变“血管”这个语义。用albumentations做同步增强很方便核心是把原图和mask传进同一个transform管道保证同一随机种子如果不小心对原图和mask分别调用随机旋转标注就会错位模型学到的是错误监督信号。危险的是弹性形变。很多自然图像分割的增强库默认开elastic deformation放在血管上会制造出根本不存在的弯曲。视网膜血管的形态受解剖结构约束过度形变等于给模型灌入错误监督训练过程中Dice会震荡最终预测图会出现很多S形伪血管。我的建议是如果一定要用弹性形变把sigma压在2以下、alpha压在10以下或者干脆不用只靠几何增强和随机裁剪。另外增强必须在FOV掩码上同步进行比如随机旋转时原图和mask用同一个随机角度否则标注和图像对不齐。这一点属于做多了就有体感的部分。3. U-net架构拆解编码器-解码器与跳跃连接手写一个可复现的最小实现3.1 为什么U-net是血管分割的天然选择小样本与精细结构的平衡U-net最早是2015年为细胞分割比赛设计的当时的训练数据也只有30多张图。它先用编码器逐层下采样把感受野从局部像素扩展到上下文区域再用解码器逐步恢复分辨率最后用跳跃连接把编码器每一层的特征拼回解码器对应层。这个结构对血管分割特别合适因为血管是又细又长的拓扑结构既需要局部像素判断“这里是不是血管”又需要上下文判断“这条细线是否属于一条连续的血管”。和FCN这类早期语义分割算法相比FCN的融合方式比较粗糙U-net在每个尺度都做拼接参数多了一些但在小样本上的稳定性好很多。这也是为什么医学影像项目的基线几乎都用U-net而不是直接上一个深层的分类骨干替换成分割头。不过要清醒一点U-net好用的是结构不是名字。如果你把编码器换成ResNet50的预训练权重在Drive上不一定会更好因为预训练权重学习的是自然图像纹理眼底图的绿色通道和自然RGB差异太大迁移收益会被领域差距抵消。3.2 用PyTorch手写U-net双卷积块、四次下采样、四次上采样网上流传的U-net实现有很多版本有的没有padding、有的把BatchNorm放在激活之后、有的忘记拼接跳跃连接直接抄会踩到莫名其妙的坑。下面这个版本是在Drive上验证过的可复现实现import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch1, out_ch1, base64): super().__init__() self.inc DoubleConv(in_ch, base) self.down1 nn.Sequential(nn.MaxPool2d(2), DoubleConv(base, base*2)) self.down2 nn.Sequential(nn.MaxPool2d(2), DoubleConv(base*2, base*4)) self.down3 nn.Sequential(nn.MaxPool2d(2), DoubleConv(base*4, base*8)) self.down4 nn.Sequential(nn.MaxPool2d(2), DoubleConv(base*8, base*8)) self.up1 nn.ConvTranspose2d(base*8, base*4, 2, stride2) self.conv1 DoubleConv(base*8, base*4) self.up2 nn.ConvTranspose2d(base*4, base*2, 2, stride2) self.conv2 DoubleConv(base*4, base*2) self.up3 nn.ConvTranspose2d(base*2, base, 2, stride2) self.conv3 DoubleConv(base*2, base) self.up4 nn.ConvTranspose2d(base, base, 2, stride2) self.conv4 DoubleConv(base*2, base) self.out nn.Conv2d(base, out_ch, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5) x torch.cat([x, x4], dim1) x self.conv1(x) x self.up2(x) x torch.cat([x, x3], dim1) x self.conv2(x) x self.up3(x) x torch.cat([x, x2], dim1) x self.conv3(x) x self.up4(x) x torch.cat([x, x1], dim1) x self.conv4(x) return self.out(x)代码有几个值得注意的设计。DoubleConv内部是两轮“卷积-BN-ReLU”3×3卷积都带padding1保证特征图尺寸在每一层不缩水。下采样用MaxPool2d而不是带stride的卷积虽然两者效果接近但MaxPool更直观显存占用也更小。通道数从64增长到512最底层保持512不再翻倍因为20张训练图的容量摆在那里翻到1024只会加速过拟合。上采样用ConvTranspose2dkernel2、stride2刚好把空间尺寸翻倍。跳跃连接用torch.cat在通道维度拼接而不是加法这是U-net原文的做法信息容量更大。输出层是1×1卷积输出单通道logits不接sigmoid这样训练时可以用BCEWithLogitsLoss数值上比先sigmoid再算BCE稳定很多。3.3 关键设计参数base通道数、深度、BatchNorm与输入通道base64是U-net原文的起点显存紧张可以降到32血管分割这种细结构对通道数不敏感降到32损失不大。深度固定为4次下采样512分辨率经过4次变成32×32最低层特征图仍保留足够空间信息。如果输入是原始565×584建议先resize到512或做padding避免最后两层特征图尺寸不对称导致上采样拼接错位。BatchNorm在这个任务里是必须的。网络内部经过多次卷积后输出分布会漂移BN能把每层输入拉回稳定范围训练更稳。但要注意batch size设为1时BN的统计量非常不稳Drive训练时batch太小的话可以换GroupNorm不过基线阶段不用折腾batch为2到4时BN表现正常。另外forward里不要在最后接sigmoid。训练时用logits直接算损失预测时才做sigmoid转概率再配合阈值转二值图。这样既省一次运算也避免概率上的非线性影响梯度传播。4. 训练与评估损失函数、超参数和Drive官方指标的取舍逻辑4.1 损失函数怎么选BCE、Dice Loss与组合损失血管在眼底图像中大约只占12%的像素也就是正负样本比接近1:9。如果直接用二分类交叉熵模型学到的最优策略大概率是全部输出背景因为这样分类错误率只有约12%但血管一个都出不来。这个任务不能只用BCE。Dice Loss的想法是直接用预测区域和真实区域的交叠程度做优化目标公式是1减去Dice系数。它对小目标更友好能缓解正负样本比例失衡。但纯Dice Loss在训练早期梯度不稳定网络初始权重接近随机时预测概率几乎全在0.5附近Dice的分母和分子都很小梯度波动很大loss曲线会像锯齿一样。常规做法是BCE和Dice按1:1相加def bce_dice_loss(pred, target, eps1e-7): # pred: 网络输出的 logits形状 B,1,H,W # target: 二值掩码和 pred 同形状取值为 0 或 1 bce nn.functional.binary_cross_entropy_with_logits(pred, target) prob torch.sigmoid(pred) inter (prob * target).sum() union prob.sum() target.sum() dice 1 - (2 * inter eps) / (union eps) return bce dice代码逻辑不复杂。BCE直接在logits上算Dice部分先sigmoid成概率再把所有像素拉平当成一个整体算交集和并集。eps1e-7防止早期分母为0。如果训练震荡明显一种调法是减小Dice的权重比如把损失改成bce 0.5*dice而不是大改学习率。原因在于Dice Loss的梯度方向在类别极度不平衡时偏向稀有类别权重太大会打破BCE提供的稳定梯度。4.2 训练超参数lr、batch size、epoch与Adam的默认陷阱U-net Drive的“不翻车参数组合”我一般是这一套Adam优化器学习率1e-4weight_decay1e-5batch size设为2到4训练200到300个epoch配合验证集Dice早停。组合在大多数消费级显卡上都能跑batch2时显存占用大概2GB左右。为什么batch不能大因为输入至少是512×512batch8时显存会到3GB以上但这不是关键。关键是小样本下大batch会加快过拟合模型过早记住训练图的光照纹理。batch2时BN统计量虽然不太稳定但在20张图这个规模下影响不大配合数据增强反而更稳。Adam的学习率不建议用默认的1e-3。Drive任务噪声强、样本少1e-3下Dice Loss很容易震荡loss在0.2附近来回跳。1e-4虽然收敛慢但曲线平滑。如果你用SGD学习率通常要放大到0.01附近并配合momentum新手还是Adam省心。训练循环model UNet(in_ch1, out_ch1).cuda() opt torch.optim.Adam(model.parameters(), lr1e-4, weight_decay1e-5) for epoch in range(300): model.train() for imgs, masks in train_loader: # imgs: B,1,512,512masks: B,1,512,512 imgs imgs.cuda() masks masks.cuda() pred model(imgs) loss bce_dice_loss(pred, masks) opt.zero_grad() loss.backward() opt.step() # 每个epoch后在验证集上算Dice保存最优权重这段代码里没有对预测做阈值因为训练阶段只需要logits和损失。有两点要注意第一mask在载入时要和输入一样是float张量取值为0和1不要用uint8直接喂进损失函数第二每个epoch之后要保存验证集Dice最高的权重而不是保存最后一个epoch的权重。很多新手直接拿最后一轮模型去测结果下降好几个点。4.3 评估指标AUC、ACC、Dice与Drive官方评价逻辑ACC在血管分割上是最没有参考价值的指标。背景像素占88%模型全输出背景也能拿到88%的ACC。真实业务里更关注Dice、IoU以及Drive官方论文里的AUC、Sensitivity、Specificity。评估代码from sklearn.metrics import roc_auc_score def evaluate(model, loader): model.eval() prob_list, true_list [], [] with torch.no_grad(): for imgs, masks in loader: prob torch.sigmoid(model(imgs)).cpu().numpy() prob_list.append(prob) true_list.append(masks.cpu().numpy()) prob np.concatenate(prob_list) true np.concatenate(true_list) # 输入已经用FOV mask归零直接拉平计算即可 p prob.ravel() t true.ravel().astype(np.float32) auc roc_auc_score(t, p) pred_bin (p 0.5).astype(np.float32) inter (pred_bin * t).sum() dice 2 * inter / (pred_bin.sum() t.sum() 1e-7) return auc, dice评估有几个容易踩的点。AUC计算用sigmoid之后的概率不要用logits。Dice的默认阈值0.5只是约定不是最优解后面会讲怎么扫描阈值。另一个细节是如果训练时输入已经用FOV mask把圆外归零评估时直接拉平全图就行如果用的是带原始背景的输入评估时需要额外排除FOV外的预测否则Dice会被大量背景像素稀释。Drive论文里报告的良好结果一般AUC在0.95左右、Dice在0.75上下不要强求完全复现显卡、预处理细节、PyTorch版本都会带来零点几个百分点的浮动。真正要盯的不是绝对数字而是训练集和测试集之间有没有出现过大的差距。5. 避坑指南U-net分割眼底血管时常见的5个翻车现场跑通U-net和跑好U-net是两回事。20张训练图、512分辨率、一个标准U-net看起来不难但实际操作里能翻车的地方一点不少。下面5条按出现频率排每一条都是从真实日志里能看到的场景。5.1 现象loss降不下去预测图全是黑的训练了50个epochloss停在0.5以上不动预测图输出全是0血管完全没出来。原因往往是网络初始化后输出偏向某一侧加上BCE在类别不平衡下给模型一个极强的“全背景”梯度模型直接收敛到平凡解。解决第一按第4章改成BCEDice组合损失Dice能保证模型即使预测得不好也在往交叠区域走第二检查标签和输入有没有对齐很大概率是增强时原图和mask用了不同的旋转角度导致监督信号错位第三观察训练前几个batch的loss如果从0.5一次性跳到0.7以上多半是输入里有NaN检查归一化是否有除零。5.2 现象血管断成一截一截像虚线模型能预测出血管但细血管经常断粗血管轮廓也毛糙。原因有两层一是下采样次数太多最底层特征图只有32×32细血管的拓扑信息在这个尺度几乎消失上采样恢复时只能靠跳跃连接补二是损失函数按像素独立计算不惩罚断连。解决优先换Dice Loss它对区域连续性比BCE更敏感换了之后断连会明显改善。其次可以在后处理里对预测图做一次形态学闭运算但这是事后补救会同时把相邻背景连起来容易误报。更彻底的办法是减少下采样次数比如去掉最深层改成3次下采样或把输入放大到640×640给最底层多留一些空间细节。5.3 现象训练集Dice很高测试集掉得措手不及20张训练图只要网络容量足够训练集Dice可以冲到90%以上测试集只有70%多。这是典型过拟合。原因不用怀疑样本太少模型记住了训练图的血管纹理、光照分布甚至背景噪声。数据增强只在epoch内有随机性模型看到的仍是同一批底图的变体。解决第一把数据增强的旋转范围从±15°扩大到±30°并加入随机裁剪到448×448再resize回512让模型不靠固定纹理识别第二weight_decay从0加到1e-5或1e-4限制权重幅度第三如果条件允许用更大的预训练骨干做迁移。但要记住Drive的测试集和训练集来自不同人眼数据分布差异本来就在测试集比训练集低5个点是正常的低太多才需要干预。5.4 现象预测图边缘出现一圈圆形边界像贴了个圆环有些跑出来的预测图在FOV边缘有一条清晰的环形假阳性带。原因是预处理时把FOV外的区域直接归零让圆内外的像素值在边界处出现巨大跳变模型学到这个跳变是一个“血管特征”于是沿着边界画圈。解决最直接的办法是在损失函数里把FOV外的预测排除让模型根本不需要预测圆外区域。更平滑的做法是对FOV边界做羽化用高斯模糊把硬边界变成软过渡比如把mask膨胀5像素再模糊后用于归一化。我一般会两者都做输入用羽化边界损失用严格mask这样能消掉圆环又不丢FOV边缘的细血管。5.5 现象同一份代码跑两次Dice差两个点数据增强的随机裁剪、PyTorch的CUDA非确定性、Adam的随机初始化叠加起来导致每次结果有百分之二上下的浮动。这属于随机性翻车不是模型写错了。解决第一固定所有随机种子包括Python、NumPy、PyTorch、CUDA代码开头用torch.manual_seed(0)和torch.backends.cudnn.deterministic True第二固定数据增强的随机种子或固定验证集做模型选择第三如果两次结果仍然有明显差异多半是batch太小导致BN统计量抖动可以试GroupNorm或把batch加大配合梯度累积。最后要接受一点血管分割模型本身就有随机性多次实验取均值才是靠谱的对比方式。6. 进阶用预测概率图的阈值扫描换更好的Dice再往“检测分割”迁移模型训练完之后还有一个低成本技巧阈值扫描。训练时默认阈值是0.5但模型输出的概率分布不一定以0.5为中心。在验证集上把阈值从0.3扫到0.7每隔0.05算一次Dice选最优阈值再用于测试集通常能白拿1到2个点的Dice。这个技巧不改变任何模型参数只改后处理决策边界是血管分割上性价比最高的优化。best_th, best_dice 0.5, 0.0 for th in [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7]: pred_bin (prob th).astype(np.float32) d 2 * (pred_bin * true).sum() / (pred_bin.sum() true.sum() 1e-7) if d best_dice: best_th, best_dice th, d阈值扫描的道理很简单模型预测概率偏低时降阈值能召回更多细血管背景噪声偏多时升阈值能压掉假阳性。具体最优阈值随训练结果漂移所以每次都扫不要套用一个固定值。做完这个项目后我最大的教训是不要只盯Dice一个数。有一版模型Dice比另一版高0.3个点但血管连续性明显更差因为Dice只算像素交叠不管拓扑。后来我把评价拆成粗血管和细血管分别看才发现细血管召回率才是眼底分割真正难啃的地方。如果你想把这套流程迁移到息肉分割数据集或其他医学语义分割任务U-net这个底座和BCEDice损失可以直接复用但预处理和阈值必须重新调不能照搬参数。希望帮到你。本文还有配套的精品资源点击获取