
简介面向眼底血管分割任务的Unet完整资料包包含已切片好的数据集、训练代码、推理脚本及训练结果文件。数据集对应眼底血管二分割任务模型仅训练10个epochs全局像素准确度达0.95miou为0.67若增大训练轮数性能还有提升空间。资源共216个文件以182个png切片图像为主另有8个py与14个pyc组成的Python工程、1个pth权重文件、5个xml工程配置及3个txt/readme说明文档整体压缩包约153.92MB。目前已有269人学习/下载。代码支持多尺度随机缩放训练utils中的compute_gray函数可自动保存mask灰度值并自适应输出通道方便扩展多分割项目训练采用cos学习率衰减run_results内可查看损失与iou曲线、训练日志及每类指标推理时只需将图像放入inference目录后运行predict脚本无需额外参数适合医学图像分割入门者、需要快速复现Unet分割流程或想训练自定义数据集的研究人员。1. 眼底血管分割为什么上来就选Unet给眼底照片做血管分割时医生最先看的就是血管网络。血管是否变细、迂曲、有没有微动脉瘤直接对应糖尿病视网膜病变的分级但血管末端常常只有两三个像素宽边缘和背景的对比度又低这类任务必须同时依赖全局语义和局部纹理。Unet几乎是这个任务的默认起点编码器把感受野逐层扩大跳跃连接又把浅层边界带回解码器正好同时保住语义与细节。这也是我拿到一个带切片好的数据集、完整代码和训练结果文件的项目包时会先去确认数据加载方式和checkpoint落盘策略的原因——这两处决定了换机器后能不能复现原有性能。2. 切片好的数据集怎么组织目录约定与 Dataset 读取拿到项目包先别急着打开训练脚本。这类项目的目录结构通常逃不出一个套路train和val下面各有一个img目录和一个mask目录图片和标签靠文件名一一对应。先把目录结构读明白后面训练、推理、评估都能少踩一半的坑。这一章就从切片尺寸、同名约定和Dataset实现三个角度把它讲透。2.1 切片行为什么重要分辨率、感受野和显存原始眼底图的分辨率差异很大从公开数据集的几百像素见方到竞赛里常见的2000像素以上。如果整张图直接丢给Unet显存会立刻吃紧更关键的是感受野问题Unet做五次下采样后深层特征图分辨率只有输入图像的1/32一根两三个像素宽的细血管在深层基本只剩一个响应点边缘信息全靠跳跃连接从浅层抄回来所以输入尺寸不能随意缩小但也不能盲目加大。因此在拿到“切片好的数据集”时先确认切片尺寸和切的时候有没有重叠。常见做法是切成512×512或256×256的patch大图上按固定步长滑窗切出来再按坐标存回同一命名空间。两种尺寸的取舍见下表patch 尺寸显存需求8G/16G卡可跑batch细血管保留训练速度256×25616 / 32一般末梢血管容易断快512×5124 / 8好血管连续性明显更稳慢我的建议是只为了快速把代码跑通256够用如果目标是要给医生当辅助工具512起步。切片时的重叠不能省至少在patch边界的1632像素要重叠否则后面推理时会在patch交界处出现假阳性。2.2 目录约定与同名规则图片和掩膜如何一一对应大多数此类项目包的目录结构长这样dataset/ ├── train/ │ ├── img/ │ │ ├── 001.png │ │ └── 002.png │ └── mask/ │ ├── 001.png │ └── 002.png └── val/ ├── img/ └── mask/规则就三条图片和掩膜严格同名掩膜是8位灰度PNG背景0、血管255不要用JPEG存掩膜JPEG会把0/255的硬边界压出过渡带。如果包里给的是.npy文件先确认存的是完整大图还是切片很多旧脚本会把整张图转成float32的npy体积比PNG大好几倍读取时还要额外做一次维度调整。另一个容易被忽略的文件是fov_mask也就是视盘区域的掩膜。DRIVE这类公开数据集通常会带一个只覆盖眼底圆形区域的掩膜评估指标只在掩膜内计算。如果包里正好有这个目录训练阶段可以不管它但评估阶段一定要把它当作乘法因子加进去否则眼底照片四角的黑色区域会把Specificity拉高Dice也会虚高。2.3 一个可直接改的torch Dataset读取切片、归一化和增强下面这段Dataset是我最常用的一套写法核心就四点glob搜集路径、统一读成RGB和灰度、归一化到0~1、在概率上做掩膜二值化放在单卡训练脚本里能直接跑import os from glob import glob import random import numpy as np from PIL import Image import torch from torch.utils.data import Dataset class VesselPatchDataset(Dataset): def __init__(self, img_dir, mask_dir, augmentTrue): self.img_paths sorted(glob(os.path.join(img_dir, *.png))) self.mask_paths sorted(glob(os.path.join(mask_dir, *.png))) assert len(self.img_paths) len(self.mask_paths), \ fimg和mask数量不一致: {len(self.img_paths)} vs {len(self.mask_paths)} self.augment augment def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(RGB) mask Image.open(self.mask_paths[idx]).convert(L) if self.augment: # 只用90度的整数倍旋转避免插值破坏单像素细血管 if random.random() 0.5: img img.transpose(Image.ROTATE_90) mask mask.transpose(Image.ROTATE_90) img np.asarray(img, dtypenp.float32) / 255.0 mask np.asarray(mask, dtypenp.float32) / 255.0 mask (mask 0.5).astype(np.float32) img torch.from_numpy(img.transpose(2, 0, 1)) mask torch.from_numpy(mask[None, ...]) return img, mask代码里img_paths和mask_paths都用sorted(glob(...))匹配顺序是稳定的掩膜用convert(L)读成单通道再除以255并做0.5的二值化是为了避免PNG压缩残留的126、254这类脏像素进损失函数。旋转只用90度整数倍而不是随机小角度因为小角度旋转要插值对单像素宽的血管是致命的。这个Dataset默认输入已经是切好的patch如果拿到的是原始大图需要再补一个随机裁剪步骤。如果训练时发现iteration跑得飞快但验证指标不动优先检查mask路径是不是串到了img目录——这是同名目录最容易犯的错。3. Unet结构要点与训练配置把完整代码调成能收敛的模型所谓“完整代码”通常指三样东西网络定义、损失函数、训练循环。跑通之前先确认三件事网络用了几层下采样、损失怎么算、checkpoint按什么频率保存。这一章按这三个顺序讲清楚顺便把训练环境相关的坑点一起说掉。3.1 从Unet网络结构图看血管任务的三个关键设计如果没看过Unet网络结构图可以先把它理解成一个U形的对称结构左边编码器每一层做两个3×3卷积加一次池化通道数从64逐步翻到512右边解码器先上采样再和编码器对应层拼接通道数再降回来。血管分割真正吃到红利的是三个设计。第一编码器给了解码器足够的上下文。一根小血管断了只看局部很难判断该不该连上编码器把32倍下采样之后的语义带回来模型才能“猜”出这条血管大概从哪来、到哪去。第二跳跃连接把第一层和第二层的边缘响应原样拼给解码器末梢血管的细节就是这样保住的。第三Unet的参数量对医学图像任务很合适小数据集上不轻易过拟合一张8G显存的卡就能训练。如果你拿到的是带ResNet骨干的变体注意骨干预训练权重要么有、要么彻底不用随机初始化一半再冻结一半的训练方式最不稳定血管分割这种强局部特征任务往往还不如标准Unet好用。3.2 损失函数选择BCE、Dice还是两者相加眼底血管分割的类别不平衡公开数据集上大约是1:9也就是说背景像素数量接近血管的十倍。直接上标准BCE时模型会把几乎所有像素判成背景因为这样loss已经足够低。所以常见做法是把Dice损失和BCE加在一起import torch import torch.nn.functional as F def bce_dice_loss(logits, targets, smooth1.0): probs torch.sigmoid(logits) bce F.binary_cross_entropy_with_logits(logits, targets) inter (probs * targets).sum(dim(1, 2, 3)) union probs.sum(dim(1, 2, 3)) targets.sum(dim(1, 2, 3)) dice 1 - ((2 * inter smooth) / (union smooth)).mean() return 0.5 * bce 0.5 * dicebce用binary_cross_entropy_with_logits而不是先过sigmoid再算交叉熵数值上更稳定dice的smooth取1.0只在训练早期起作用不影响最终收敛0.5和0.5的权重可以按结果微调。我一般会让dice权重大一点比如0.6 dice、0.4 bce对小血管更友好。如果预测偏保守就把dice权重调高如果预测全是噪声点把bce权重拉回0.5。损失特点适合场景BCE逐像素独立梯度平稳作为辅助项稳定早期训练Dice对类别不平衡不敏感直接优化目标血管这类小目标BCEDice折中工程上最稳绝大多数血管分割项目3.3 训练入口参数表、batch调节和checkpoint保存把Dataset和损失函数接进训练脚本后命令行入口我一般长这样python train.py \ --img_dir dataset/train/img \ --mask_dir dataset/train/mask \ --val_img_dir dataset/val/img \ --val_mask_dir dataset/val/mask \ --batch_size 8 \ --lr 1e-3 \ --epochs 100 \ --patch_size 512 \ --val_every 5 \ --save_dir checkpoints这些参数来自常见的训练环境配置patch为512时batch 8大约吃16G显存8G显存的卡就把batch降到4同时把patch降到320效果差距不大。学习率用1e-3配Adam跑前20个epoch是血管分割里最稳的组合之一。参数推荐值注意--lr1e-3Adam下偏大超过50个epoch没降就再降一半--batch_size48显存不够时优先降batch再降patch--epochs100大多数公开数据集60100个epoch足够--val_every5间隔太久容易错过最佳checkpoint--patch_size512256能跑但末梢血管会明显变差checkpoint保存策略坚持两条始终保留验证集Dice最高的模型单独存成best_dice.pth再存一个last.pth用于事故现场排查。你拿到的“训练的结果文件”里如果只有一个模型先看脚本里保存条件是什么。另外多看一眼train.log里的验证曲线如果验证Dice在中间epoch就开始下滑说明要么数据增强太强要么lr需要在40个epoch附近做一次0.1倍衰减。提示换机器复现时先确认torch版本和CUDA版本跟原log一致很多加载报错不是代码问题是版本差异造成的。4. 加载训练结果文件做推理切片拼接、阈值和评估训练跑完手里有best_dice.pth、last.pth和train.log。这一章讲怎么把这些结果文件变成能看的血管图再变成能写进文档的数字。最常见的坑是直接把训练脚本里的验证代码拿来做推理那样往往会做随机增强还拼了batch输出顺序对不上原图。4.1 单张切片推理脚本从checkpoint加载模型import torch import numpy as np from PIL import Image model UNet(in_channels3, out_channels1) # 你训练时用的模型类 ckpt torch.load(checkpoints/best_dice.pth, map_locationcpu) state ckpt.get(model_state_dict, ckpt) # 去掉DataParallel加的前缀 state {k.replace(module., ): v for k, v in state.items()} model.load_state_dict(state) model.eval() img Image.open(dataset/val/img/001.png).convert(RGB) x np.asarray(img, dtypenp.float32) / 255.0 x torch.from_numpy(x.transpose(2, 0, 1)).unsqueeze(0) with torch.no_grad(): logits model(x) prob torch.sigmoid(logits)[0, 0].numpy() Image.fromarray((prob * 255).astype(np.uint8)).save(out_prob.png)这里有几个关键点。加载前先检查ckpt里的key是不是带module.前缀用了DataParallel就必须去掉map_location要跟当前环境一致半精度权重加载时尤其容易报错推理前必须调model.eval()否则带BatchNorm的模型结果会抖动。输出是sigmoid概率而不是阈值图阈值后处理要放到统计指标时统一做不要在推理阶段提前二值化。4.2 滑窗重叠拼接消除切片边界断裂训练和推理的区别在拼接这里体现得最明显。训练时可以随机裁剪因为数据增强了多样性推理时必须保证每个patch的位置能映射回原图。直接用非重叠滑窗会把血管在patch边界处切成两段重叠拼接是常见解法重叠区域被多个patch推理概率做平均边界自然平滑。def sliding_predict(model, img, patch_size512, overlap32, devicecuda): h, w img.shape[:2] step patch_size - overlap # 把图pad到能被step整除避免漏掉右边界和下边界 pad_h (patch_size - h % step) % step pad_w (patch_size - w % step) % step img np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), modereflect) prob_sum np.zeros_like(img[..., 0], dtypenp.float32) weight np.zeros_like(img[..., 0], dtypenp.float32) for y in range(0, img.shape[0] - patch_size 1, step): for x in range(0, img.shape[1] - patch_size 1, step): crop img[y:y patch_size, x:x patch_size].astype(np.float32) / 255.0 x_t torch.from_numpy(crop.transpose(2, 0, 1)).unsqueeze(0).to(device) with torch.no_grad(): out torch.sigmoid(model(x_t))[0, 0].cpu().numpy() prob_sum[y:y patch_size, x:x patch_size] out weight[y:y patch_size, x:x patch_size] 1 prob prob_sum / np.maximum(weight, 1.0) return prob[:h, :w]重叠区会累加多次预测weight矩阵记录每个像素被预测了多少次最后做除法取平均。overlap设16到32比较合适太小抑制不了接缝太大计算量翻倍。如果推理显存紧张可以把patch_size调小同时把overlap等比调小但不要只改patch_size不改overlap。4.3 评估指标怎么算Dice、Sensitivity和Specificity评估时要注意阈值的选择。不要默认0.5常见做法是在验证集上画出阈值和指标的关系曲线再取Youden index对应的阈值。四个最常用的指标定义如下指标公式关心的问题Dice2TP / (2TP FP FN)血管区域整体重合度IoUTP / (TP FP FN)与Dice类似分母略大SensitivityTP / (TP FN)末梢小血管漏检多少SpecificityTN / (TN FP)背景噪声有多少计算指标的代码可以直接把四个指标合在一起def compute_metrics(pred_mask, gt_mask): pred_mask pred_mask.astype(bool) gt_mask gt_mask.astype(bool) tp (pred_mask gt_mask).sum() tn (~pred_mask ~gt_mask).sum() fp (pred_mask ~gt_mask).sum() fn (~pred_mask gt_mask).sum() eps 1e-6 return { dice: 2 * tp / (2 * tp fp fn eps), iou: tp / (tp fp fn eps), sens: tp / (tp fn eps), spec: tn / (tn fp eps), acc: (tp tn) / (tp tn fp fn eps), }如果数据集带fov_mask在调用这个函数之前先对pred_mask和gt_mask做一次 fov操作。没有fov就按全图算但报告里要注明因为背景黑色区域会把Specificity顶得很高不同数据集之间的数字不具备可比性。5. 让血管不断裂的距离图监督给Unet加辅助头如果前面几步都跑通了最常见的残留问题是末梢血管断成一截截。BCE和Dice都是像素级损失它们不惩罚结构上的“断点”只要断裂的像素数量占比低整体损失看起来依然正常。这时候可以上一个成本很低的Unet模型改进距离图监督。先用骨架提取把血管压成中心线再计算每个骨架像素到血管边缘的距离import numpy as np from scipy.ndimage import distance_transform_edt from skimage.morphology import skeletonize def vessel_dist_map(mask): bin_mask (mask 0.5).astype(np.uint8) skeleton skeletonize(bin_mask) dist distance_transform_edt(bin_mask) return (dist * skeleton).astype(np.float32)mask是二值标签skeleton是骨架提取结果distance_transform_edt算出每个血管像素到背景的最近距离乘上skeleton之后只有中心线上的点保留距离值背景全是0。网络在分割头之外再接一个单通道输出预测这个距离图相当于给每个血管像素加了一个“你离边缘有多远”的连续目标监督信号比二值掩膜强得多。训练时总损失这样拼seg_loss bce_dice_loss(logits, mask) dist_loss F.smooth_l1_loss(dist_logits, dist_gt) loss 0.5 * seg_loss 0.5 * dist_losssmooth_l1_loss对离群点比MSE稳适合距离图这种偶尔出现大值的回归目标。分担比例从0.5:0.5起步如果分割Dice掉得厉害把dist_weight降到0.3再试。需要注意骨架提取本身可能把血管末端磨掉导致末端距离值失真。训练前把距离图可视化一次看骨架末端是不是缩了一截skimage的skeletonize对45度走向的血管保留更完整scipy的morphological_thinning对斜向血管偏保守两种都可以试。评估时只看分割结果在断裂处有没有变少不看距离图本身的指标。如果距离图loss降了但分割没改善先确认辅助头是在同一个decoder上多接了一层卷积还是从encoder最深层直接引出来——前者通常更稳。本文还有配套的精品资源点击获取