眼底血管分割数据集与PyTorch实战:从数据到可视化

📅 发布时间:2026/10/11 13:06:03
眼底血管分割数据集与PyTorch实战:从数据到可视化
简介这份资源面向医学图像分割方向的初学者与算法实践者提供一套可直接上手的眼底血管分割数据集与配套工具帮助解决血管提取任务中数据获取与标签制作的门槛问题。数据集基于DRIVE扩充图像分辨率为500至1000训练集含71张原图与71张对应mask测试集含17张原图与17张maskmask为前景255的二值图像便于直观观察血管结构。压缩包共178个文件以177张png图像和1个py脚本为主整体约45.88MBpng覆盖原图与标签py脚本用于可视化展示。该脚本可随机抽取一张图片将原始图像、GT图像以及GT在原图上的蒙板效果一并输出并保存至当前目录方便快速核验数据质量与标注对齐情况。目前已有1264人学习下载适合用于分割模型训练、数据预处理练习与结果可视化验证目录按训练与测试划分结构清晰便于直接接入常见分割网络流程。1. 眼底血管分割数据集从一张眼底图到可训练的血管掩码眼底血管分割是医学图像分割里少有的「数据够小、任务够难、评价够硬」的方向。一张彩色眼底图里血管像素占比通常不到 10%动静脉交叉、黄斑区无血管带、病灶遮挡都会让模型在细血管上直接断线。很多做医学图像分割的同行卡在第一步不是模型而是数据集图像和掩码对不上、类别标签只有 0/1 却混进了灰度值、可视化代码画出来的掩码和原图错位。这个标题讲的是一套「数据集2 分割 类别标签 可视化代码」的组合核心是二分类语义分割——前景是血管背景是视网膜。它适合刚接触医学图像分割、想跑通第一个血管分割 baseline 的人也适合已经做过 YOLO 系列目标检测、想切到像素级任务的工程师。下面按「数据长什么样 → 怎么读怎么训 → 坑在哪 → 怎么验证」推一遍。2. 眼底血管分割数据集的结构与类别标签设计2.1 2 分割任务里图像和掩码到底怎么对应「2 分割」在医学图像分割语境里就是二分类语义分割输出通道为 2 或者单通道二值。常见做法是原图是一张 RGB 三通道的眼底彩照掩码是一张单通道 8 位灰度图血管像素值 255背景 0。训练时把掩码除以 255 变成 0/1再送进损失函数。这里有个容易忽略的点如果掩码是调色板模式P 模式的 PNG直接读进来可能是 0/1 索引而不是 0/255可视化时看着正常训练时标签全错。我一般会在 Dataset 的__getitem__里强制做一次np.unique检查确认标签集合是 {0,1} 或 {0,255}。类别标签的设计直接决定损失函数怎么选。二分类血管分割常用 BCEWithLogitsLoss 或 DiceLoss前者对类别不平衡敏感后者对细小结构更友好。血管像素占比低如果只用 BCE模型倾向于全预测背景也能拿到很高的准确率但 Dice 会很难看。所以类别标签虽然只有两类评估指标不能只看 accuracy要看 Dice、IoU、 Sensitivity召回和 Specificity。眼底血管分割里 Sensitivity 尤其重要漏掉一段血管在临床上可能意味着漏掉一个缺血区域。2.2 目录组织与文件命名约定一个能直接跑的数据集目录结构通常是这样dataset/ ├── images/ │ ├── 01_dr.png │ ├── 02_dr.png │ └── ... ├── masks/ │ ├── 01_dr.png │ ├── 02_dr.png │ └── ... └── splits/ ├── train.txt ├── val.txt └── test.txtimages和masks里的文件名必须一一对应差一个后缀都会在训练中途报FileNotFoundError。splits里每行写一个不含扩展名的样本 ID这样切分逻辑和文件读取解耦。我见过有人把切分直接写死在代码里换一批数据就要改代码这是典型的后悔药场景。常见做法是把切分文件独立出来训练脚本只读 txt。提示如果数据集里图像是 JPG、掩码是 PNG不要为了统一格式把掩码转成 JPG。JPG 有损压缩会在血管边缘产生伪影标签值可能从 0/255 变成 0/254 这种脏数据。2.3 类别标签的两种存法单通道二值与双通道 one-hot单通道二值掩码省空间、读取快适合 BCE 和 Dice。双通道 one-hot 是背景通道 血管通道适合 CrossEntropyLoss 和 Softmax 输出。两者可以互转import numpy as np def mask_to_onehot(mask, num_classes2): # mask: HxW, 取值 0 或 1 onehot np.zeros((num_classes, *mask.shape), dtypenp.float32) for c in range(num_classes): onehot[c] (mask c).astype(np.float32) return onehot def onehot_to_mask(onehot): # onehot: CxHxW return np.argmax(onehot, axis0).astype(np.uint8)mask_to_onehot里用循环而不是广播是为了在类别数少的时候可读性更好如果类别多改成np.eye(num_classes)[mask].transpose(2,0,1)更快。onehot_to_mask用argmax注意当两个通道概率相等时 argmax 返回 0也就是背景这在血管边界上会轻微腐蚀血管评估时 Dice 会掉一点属于正常现象。3. 用 PyTorch 跑通眼底血管分割的最小训练闭环3.1 Dataset 与 DataLoader 的四个必调参数先写一个能跑的 Datasetimport os import cv2 import numpy as np import torch from torch.utils.data import Dataset, DataLoader class VesselDataset(Dataset): def __init__(self, root, split_file, img_size512, augmentTrue): self.root root self.img_size img_size self.augment augment with open(split_file, r) as f: self.ids [line.strip() for line in f if line.strip()] def __len__(self): return len(self.ids) def __getitem__(self, idx): sid self.ids[idx] img cv2.imread(os.path.join(self.root, images, sid .png)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(os.path.join(self.root, masks, sid .png), 0) img cv2.resize(img, (self.img_size, self.img_size)) mask cv2.resize(mask, (self.img_size, self.img_size), interpolationcv2.INTER_NEAREST) mask (mask 127).astype(np.float32) # 统一到 0/1 if self.augment: if np.random.rand() 0.5: img np.fliplr(img).copy() mask np.fliplr(mask).copy() if np.random.rand() 0.5: img np.flipud(img).copy() mask np.flipud(mask).copy() img img.astype(np.float32) / 255.0 img (img - img.mean()) / (img.std() 1e-6) # 逐图标准化 img torch.from_numpy(img).permute(2, 0, 1) mask torch.from_numpy(mask).unsqueeze(0) return img, mask四个必调参数img_size决定显存和细血管保留程度512 是眼底血管分割的常用折中256 会丢细血管1024 显存吃紧augment控制是否翻转增强血管分割里水平翻转和垂直翻转是安全的旋转要小心因为视盘位置有解剖意义interpolation掩码必须用INTER_NEAREST用双线性会插出 0.5 这种中间值normalize用逐图标准化而不是 ImageNet 均值因为眼底图色彩分布和自然图像差很远。DataLoader 这边loader DataLoader( VesselDataset(rootdataset, split_filedataset/splits/train.txt), batch_size4, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue )batch_size4在 512 分辨率下大约占 6GB 显存卡小就降到 2。num_workers设成 CPU 核数的一半左右设太大反而因为进程切换变慢。drop_lastTrue避免最后一个 batch 只有 1 张图BatchNorm 在 batch 为 1 时会报错或统计量失真。3.2 U-Net baseline 的通道数和损失函数选择一个够用的 U-Netimport torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_c, out_c): super().__init__() self.net nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), nn.Conv2d(out_c, out_c, 3, padding1), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.net(x) class UNet(nn.Module): def __init__(self, in_c3, out_c1, base32): super().__init__() self.d1 DoubleConv(in_c, base) self.d2 DoubleConv(base, base*2) self.d3 DoubleConv(base*2, base*4) self.d4 DoubleConv(base*4, base*8) self.pool nn.MaxPool2d(2) self.up3 nn.ConvTranspose2d(base*8, base*4, 2, stride2) self.u3 DoubleConv(base*8, base*4) self.up2 nn.ConvTranspose2d(base*4, base*2, 2, stride2) self.u2 DoubleConv(base*4, base*2) self.up1 nn.ConvTranspose2d(base*2, base, 2, stride2) self.u1 DoubleConv(base*2, base) self.out nn.Conv2d(base, out_c, 1) def forward(self, x): c1 self.d1(x) c2 self.d2(self.pool(c1)) c3 self.d3(self.pool(c2)) c4 self.d4(self.pool(c3)) x self.u3(torch.cat([self.up3(c4), c3], dim1)) x self.u2(torch.cat([self.up2(x), c2], dim1)) x self.u1(torch.cat([self.up1(x), c1], dim1)) return self.out(x)base32是显存和精度的平衡点血管分割不需要 ResNet 那么大的 backboneU-Net 原始结构在眼底血管上 Dice 能到 0.78 到 0.82。out_c1配合 BCEWithLogitsLoss推理时 sigmoid 再阈值 0.5。损失函数我一般用 BCE Dice 各 0.5 权重bce nn.BCEWithLogitsLoss() def dice_loss(pred, target, eps1e-6): pred torch.sigmoid(pred) inter (pred * target).sum(dim(2,3)) union pred.sum(dim(2,3)) target.sum(dim(2,3)) return 1 - (2 * inter eps) / (union eps) loss 0.5 * bce(out, mask) 0.5 * dice_loss(out, mask).mean()纯 BCE 在血管占比 5% 到 8% 的数据上前几个 epoch 会全预测背景Dice 接近 0加 Dice 项能把梯度拉回来。eps防止除零放在分子分母都加保证数值稳定。3.3 训练循环里必须打印的三个指标for epoch in range(50): model.train() for img, mask in loader: img, mask img.cuda(), mask.cuda() out model(img) loss 0.5 * bce(out, mask) 0.5 * dice_loss(out, mask).mean() optimizer.zero_grad() loss.backward() optimizer.step() # 验证 model.eval() dices, sens, specs [], [], [] with torch.no_grad(): for img, mask in val_loader: img, mask img.cuda(), mask.cuda() pred (torch.sigmoid(model(img)) 0.5).float() tp (pred * mask).sum().item() fp (pred * (1 - mask)).sum().item() fn ((1 - pred) * mask).sum().item() tn ((1 - pred) * (1 - mask)).sum().item() dices.append(2*tp / (2*tp fp fn 1e-6)) sens.append(tp / (tp fn 1e-6)) specs.append(tn / (tn fp 1e-6)) print(fepoch {epoch} dice {np.mean(dices):.4f} sens {np.mean(sens):.4f} spec {np.mean(specs):.4f})Dice 看整体重叠Sensitivity 看血管漏没漏Specificity 看背景有没有被误判成血管。三个指标一起看能区分「模型保守」和「模型激进」。如果 Dice 高但 Sensitivity 低说明模型偏向背景可以把阈值从 0.5 降到 0.4 试试如果 Specificity 低说明背景误检多检查一下掩码里有没有把病灶标成血管。4. 可视化代码把原图、掩码、预测叠在一张图上4.1 三栏对比图的绘制函数import matplotlib.pyplot as plt def visualize(img, mask, pred, save_pathNone): # img: HxWx3, 0-1; mask/pred: HxW, 0/1 img (img * 255).astype(np.uint8) overlay img.copy() overlay[pred 0] [255, 0, 0] # 预测血管标红 overlay[mask 0] [0, 255, 0] # 真实血管标绿 # 红绿重叠处会偏黄说明预测正确 fig, axes plt.subplots(1, 4, figsize(20, 5)) axes[0].imshow(img); axes[0].set_title(original) axes[1].imshow(mask, cmapgray); axes[1].set_title(mask) axes[2].imshow(pred, cmapgray); axes[2].set_title(pred) axes[3].imshow(overlay); axes[3].set_title(overlay) for ax in axes: ax.axis(off) if save_path: plt.savefig(save_path, bbox_inchestight, dpi150) plt.close()overlay里先画预测再画真实真实掩码覆盖预测这样红色区域就是「预测有但真实没有」的假阳性绿色区域是「真实有但预测没有」的假阴性黄色是重叠。这个顺序不能反反了假阳假阴就看不出来。dpi150是论文插图的最低要求再低血管边缘会糊。4.2 可视化时最容易翻车的三个点第一掩码值域没归一化。如果掩码是 0/255mask 0没问题但如果直接imshow(mask)而 mask 是 0/1显示出来几乎全黑会误以为掩码是空的。第二原图通道顺序。OpenCV 读进来是 BGRmatplotlib 按 RGB 显示不转换的话血管和背景颜色会偏。第三预测和原图尺寸不一致。如果模型输出是 512原图是 1024直接叠会错位必须先把预测 resize 回原图尺寸且用INTER_NEAREST。注意可视化代码不要放在训练循环里每步都画IO 会拖慢训练。我一般每个 epoch 存一张验证集样本的对比图训练结束后统一看。5. 避坑与排查眼底血管分割数据集常见的五类翻车现象一训练 loss 一直降但 Dice 卡在 0.3 不动。原因通常是掩码读取时值域没统一比如掩码是 0/255代码里没除以 255BCE 的 target 期望 0/1结果梯度方向错了。解决在 Dataset 里打印np.unique(mask)确认只有两个值且训练前统一到 0/1。现象二验证集 Dice 比训练集高很多。原因一般是数据泄露train.txt 和 val.txt 里有重复样本 ID或者同一张图的不同裁剪块分到了两个集合。解决切分前对样本 ID 去重按患者 ID 切而不是按图像切同一患者的左右眼不能分到不同集合。现象三可视化时预测掩码整体偏移几个像素。原因是原图和掩码在预处理时用了不同的 resize 插值或者一个做了 padding 一个没做。解决原图和掩码用同一套几何变换resize 时掩码用INTER_NEAREST原图用INTER_LINEAR但尺寸必须一致。现象四细血管末端总是断。原因是下采样太多次U-Net 的 pool 把细血管特征丢了。解决减少一次下采样或者把base从 32 提到 48或者在 skip connection 上加注意力模块。另一个低成本做法是把输入分辨率从 256 提到 512。现象五换一台机器跑Dice 掉 5 个点。原因是 OpenCV 和 PIL 读图的值域、通道顺序、EXIF 旋转处理不一致。解决固定用同一种读图库训练和推理用同一套预处理代码不要一个用 cv2 一个用 PIL。6. 进阶技巧用血管骨架化指标补上 Dice 看不到的细节Dice 对粗血管敏感对细血管不敏感。一段直径 2 像素的血管断成两截Dice 可能只掉 0.01但临床上这段血管可能正好是黄斑区供血支。我一般会额外算一个骨架化后的指标先把预测和真实掩码都做形态学细化skeletonization再算骨架的 F1。骨架 F1 对断线非常敏感能补上 Dice 的盲区。from skimage.morphology import skeletonize def skeleton_f1(pred, mask): sk_pred skeletonize(pred 0) sk_mask skeletonize(mask 0) tp (sk_pred sk_mask).sum() fp (sk_pred ~sk_mask).sum() fn (~sk_pred sk_mask).sum() prec tp / (tp fp 1e-6) rec tp / (tp fn 1e-6) return 2 * prec * rec / (prec rec 1e-6)skeletonize要求输入是 bool 数组所以先做 0。这个指标不用来训练只用来做模型选型。我自己的习惯是Dice 相差 0.01 以内的两个模型选骨架 F1 高的那个。另外推理阈值不要固定 0.5可以在验证集上扫 0.3 到 0.7选骨架 F1 最高的阈值通常比 0.5 低一点因为细血管的 sigmoid 输出偏低。这个扫阈值的操作花不了几分钟但经常能白捡 1 到 2 个点的召回。希望帮到你。本文还有配套的精品资源点击获取