UNet眼底血管分割实战：从切片数据集到PyTorch训练推理-北京尧图网络

简介面向眼底血管分割任务的Unet完整实践项目包含已切片好的眼底图像数据集、可直接运行的训练与推理代码以及训练得到的结果文件适合医学图像分割入门和需要快速跑通完整流程的开发者。压缩包共216个文件、约153.92MB以182个png眼底图像为主数据另有8个py训练/预测脚本、14个pyc编译文件、txt标签/日志、xml配置及pth权重文件目录结构清晰已有269人学习下载。项目仅训练10个epochs就在眼底血管二分割上取得全局像素准确率0.95、miou 0.67加大轮次后性能仍有提升空间。代码支持数据随机多尺度缩放与cos学习率衰减并能根据mask灰度值自动配置Unet输出通道训练日志及损失/IOU曲线保存于run_results便于逐类分析指标。推理时仅需将图像放入inference目录并运行predict脚本小白也能快速上手。1. 说清眼底血管分割这件事眼底血管分割是医学图像分析中的一个经典任务输入一张彩色眼底照片输出每个像素属于血管还是背景的概率图。血管的粗细差异很大末梢血管只有几个像素宽对比度又低因此落地方案通常不像自然图像分割那样直接扔给大模型而是先裁剪成patch再让UNet去学。标题里的“切片好的数据集”指的就是这类已经按固定窗口切好的训练样本配合完整代码和训练结果文件能少走很多弯路。它面向要跑通UNet眼底血管分割的工程师和研究生覆盖从数据集加载、模型定义、训练指标到推理后处理的完整链路。如果你有DRIVE、CHASE或本地眼底图按文中的文件组织就能直接训练。2. UNet结构解析与眼底血管分割的适配点2.1 编码器下采样血管分割需要多大的感受野UNet从FCN发展而来主体分编码器和解码器。编码器通过4次下采样把输入从HxW降到H/16xW/16特征通道数从64增加到512。下采样对眼底血管分割有三个实际作用一是让卷积核看到更大范围的视网膜背景从而区分血管和出血点二是减少后续计算量三是迫使模型学到不同尺度的血管响应。血管在眼底图像上既有跨越半个视野的动脉主干也有只有两三个像素宽的毛细血管单靠单一尺度卷积无法同时覆盖两种目标。阶段操作序列输出分辨率关注特征C13x3 conv, BN, ReLU x2H/1 x W/1 x 64血管边缘、纹理P1D2maxpool conv blockH/2 x W/2 x 128局部血管走向P2D3maxpool conv blockH/4 x W/4 x 256分叉与交叉P3D4maxpool conv blockH/8 x W/8 x 512大血管区域bridgeconv blockH/16 x W/16 x 512全局上下文注意这里没有继续池化到H/32。眼底血管分割不是前景分类过大的下采样会让最细的血管在特征图上直接消失保留1/16分辨率作为编码器最深处是常见折中。如果显存比较紧张可以把最深处设为H/8解码器也相应减少一层但感受野变小后大出血块和血管容易混在一起。改结构时要同步修改后面模型的forward内skip数量不能只改一处。2.2 解码器与跳跃连接薄血管恢复的关键解码器每次先对底层特征做2倍上采样然后与编码器同分辨率特征拼接再进行两个3x3卷积。跳跃连接的贡献不只是补细节它把编码器前期的空间坐标信息直接传给解码器。血管末梢只有2-3个像素时仅靠深层语义无法定位精确边界所以concat比sum更常用。在UNet的眼底血管分割实现里不要轻易删除早期跳跃连接如果为了减少参数至少保留C1或C2层否则边缘预测会明显变粗。上采样可以选用转置卷积或双线性插值。转置卷积带可学习参数能恢复更多纹理但也更容易产生棋盘伪影我通常用双线性插值加上后面的卷积层血管边缘反而更平滑。PyTorch里使用F.interpolate还是ConvTranspose2d会直接影响最后输出的像素级精度。若发现预测图出现一格一格的纹理优先把解码器的转置卷积换成双线性上采样。2.3 一个能跑的UNet定义训练时可改的3个参数下面是经典UNet的PyTorch实现。代码省去了注意力、空洞卷积等改造先保证能作为基线跑通后面要换backbone时只需要调整Encoder部分。# unet.py 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_channels3, out_channels1, features(64, 128, 256, 512)): super().__init__() self.downs nn.ModuleList() self.ups nn.ModuleList() self.pool nn.MaxPool2d(2) # 编码器逐级提取特征 for f in features: self.downs.append(DoubleConv(in_channels, f)) in_channels f # 最深层 bridge self.bridge DoubleConv(features[-1], features[-1] * 2) # 解码器转置卷积上采样跳跃连接 for f in reversed(features): self.ups.append(nn.ConvTranspose2d(f * 2, f, kernel_size2, stride2)) self.ups.append(DoubleConv(f * 2, f)) self.out_conv nn.Conv2d(features[0], out_channels, 1) def forward(self, x): skip [] for down in self.downs: x down(x) skip.append(x) x self.pool(x) x self.bridge(x) skip skip[::-1] for i in range(0, len(self.ups), 2): x self.ups[i](x) x torch.cat([x, skip[i // 2]], dim1) x self.ups[i 1](x) return self.out_conv(x)这段网络不包含任何注意力机制但足够作为眼底血管分割的基线。三个能直接改的参数分别是in_channelsout_channelsfeatures。in_channels对应输入图像通道数彩色眼底图是3灰度眼底图改为1out_channels是分割输出的通道数二分类血管分割用1而不是2features代表编码器各层通道数默认是(64,128,256,512)。如果显存不高可以改成(32,64,128,256)模型参数会明显下降但通常需要更多epoch才能追平效果。这里有个容易忽略的地方解码器里bridge的输出通道是features[-1]2而ups列表里第一个转置卷积的输入通道也必须是f2。修改features时只要保持f*2一致即可否则会在forward阶段出现channels mismatch。这个错误在PyTorch中要等到实际跑数据时才报出来建议写完模型后用随机输入跑一次前向传播再开始训练。3. 切片数据集的整理与加载从原图到训练样本3.1 切片好的数据集长什么样目录约定与样本对应关系标题里的“切片好的数据集”通常不是一张完整眼底图而是一批固定size的patch。常见目录结构如下images放彩色patchmasks放黑白掩膜两个目录的文件名一一对应。data/ ├── images/ │ ├── 01_00000.png │ ├── 01_00001.png │ └── ... ├── masks/ │ ├── 01_00000.png │ ├── 01_00001.png │ └── ... ├── train.txt └── val.txt文件名前缀由原图编号和patch坐标组成例如01_00000表示第1张原图的左上角。train.txt和val.txt每行一个文件名前缀模型读取时按这个列表加载切片。训练集和验证集必须分到不同的原图不能只分patch。如果同一张原图的不同patch同时出现在训练和验证中模型会通过感受野外的周边信息记住图像导致验证结果虚高。切片好的数据集和直接输入原图的区别在于模型看到的是局部视野。视网膜血管横向跨度很大一个256x256的patch有时只包含一根血管主干的一部分。遇到这种情况不要在patch里硬分前景和背景而是依赖推理时滑动窗口的重叠预测。切片参数没有绝对标准下面是我常用的参考值。参数小显存中等显存说明patch_size128256需要能被模型下采样次数整除建议是16的倍数stridepatch_size - 32patch_size - 64重叠越多训练样本越多但训练更慢过滤阈值1%1%全背景切片保留少量即可存储格式pngpng不要用jpeg血管边缘会失真设置patch_size时要避开一个坑UNet每层做2倍下采样一共4次所以patch的长宽最好是16的倍数。如果尺寸不是16的倍数上采样后的特征图尺寸会和跳跃连接层差1个像素虽然有些情况仍能跑但整体分割结果会在边缘处错位。3.2 切片代码固定步长和重叠切片如果没有拿到别人切好的数据自己从原始眼底图制作也很简单。下例以256x256的patch、192的步长把大图和掩膜同步切片。# make_patches.py import os import random from PIL import Image import numpy as np def extract_patches(image_path, mask_path, save_dir_img, save_dir_mask, patch_size256, stride192, keep_ratio0.1): image np.array(Image.open(image_path).convert(RGB)) mask np.array(Image.open(mask_path).convert(L)) h, w mask.shape name os.path.splitext(os.path.basename(image_path))[0] idx 0 for y in range(0, h - patch_size 1, stride): for x in range(0, w - patch_size 1, stride): img_patch image[y:ypatch_size, x:xpatch_size] mask_patch mask[y:ypatch_size, x:xpatch_size] foreground_ratio (mask_patch 0).mean() * 100 # 过滤完全无血管的切片按比例保留背景样本 if foreground_ratio 1 and random.random() keep_ratio: continue Image.fromarray(img_patch).save( os.path.join(save_dir_img, f{name}_{idx:05d}.png)) Image.fromarray(mask_patch).save( os.path.join(save_dir_mask, f{name}_{idx:05d}.png)) idx 1代码以stride192在256的patch上移动patch与patch之间有64像素重叠。重叠的作用是让血管在多个patch中都完整出现推理拼接时也不会在patch边界留下折痕。keep_ratio控制全背景切片的保留比例如果全部丢掉模型容易把视盘周围的暗区误判为血管所以保留10%左右的背景样本会更稳。如果你的数据集切片后样本数已经很大可以去掉这个过滤条件直接让模型学习背景分布。保存时统一用PNG。PNG是无损格式掩膜的边界不会有JPEG压缩带来的灰边。若原始掩膜本身就是JPEG建议在切片前用形态学闭运算把断裂的细小血管连接一次否则模型会把断裂当作学习目标。3.3 Dataset加载与增强血管分割里哪些增强可以开切片做好后写一个torch.utils.data.Dataset把patch读进来。常见做法是加载img和mask后一起做随机翻转和旋转。对于眼底血管分割颜色抖动要谨慎血管颜色是重要特征过度改变颜色会让模型不鲁棒。推荐开启的增强包括水平翻转、垂直翻转、90度旋转和轻度随机仿射弹性形变对医学图像很有效但注意位移不要超过3个像素。# dataset.py import os import random import numpy as np import torch from torch.utils.data import Dataset from PIL import Image class VesselDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, trainTrue): self.img_dir img_dir self.mask_dir mask_dir self.samples [line.strip() for line in open(file_list)] self.train train def __len__(self): return len(self.samples) def __getitem__(self, idx): name self.samples[idx] img np.array(Image.open( os.path.join(self.img_dir, name .png)).convert(RGB)).astype(np.float32) mask np.array(Image.open( os.path.join(self.mask_dir, name .png)).convert(L)) mask (mask 127).astype(np.float32) if self.train: # 翻转时图像和掩膜必须同步 if random.random() 0.5: img img[:, ::-1, :] mask mask[:, ::-1] if random.random() 0.5: img img[::-1, :, :] mask mask[::-1, :] k random.choice([0, 1, 2, 3]) if k: img np.rot90(img, k, axes(0, 1)) mask np.rot90(mask, k, axes(0, 1)) img img / 255.0 img torch.from_numpy(img.transpose(2, 0, 1)) mask torch.from_numpy(mask).unsqueeze(0) return img, mask加载时把mask二值化大于127视为血管其余视为背景。增强顺序是先翻转再旋转顺序不能调换否则坐标对应关系会乱。这个类没有做resize要求所有切片已经统一成patch_size。如果数据集来自多个来源尺寸不一致需要在__getitem__里补上Resize但mask必须用最近邻插值不能用线性插值否则血管边缘会出现中间灰值。验证时trainFalse只做标准化不做增强。4. 训练UNet损失函数、评价指标与训练结果文件的输出4.1 损失函数为什么用BCE加Dice组合眼底血管分割是像素二分类。BCE容易优化但正负样本不平衡血管像素通常只占10%左右训练初期模型会倾向把所有像素预测为背景。Dice Loss直接优化前景/背景重叠对小目标更敏感但单独使用梯度不平滑。常见做法是让两者相加total_loss bce_loss dice_loss。# losses.py import torch import torch.nn.functional as F def mixed_loss(pred, target): pred torch.sigmoid(pred) bce F.binary_cross_entropy(pred, target, reductionmean) smooth 1e-5 pred_flat pred.reshape(pred.size(0), -1) target_flat target.reshape(target.size(0), -1) intersection (pred_flat * target_flat).sum(dim1) dice 1 - (2 * intersection smooth) / ( pred_flat.sum(dim1) target_flat.sum(dim1) smooth) return bce dice.mean()smooth的作用是防止分母为0一般取1e-5。Dice部分没有做one-hot因为血管分割只有一个前景类。这里返回的是bce dice.mean()如果只想对前景加权可以在BCE里给pos_weight传一个小数比如1.2但通常组合损失已经够用。一个关键点是pred要先经过sigmoid不能用logits直接算Dice。训练时如果发现dice loss出现nan先检查标签是不是0/255而不是0/1。Dataset里已经做了mask 127所以target_flat和pred_flat数值尺度一致。另一个来源是batch里某张mask全为零smooth会阻止除零但如果smooth忘了加loss就会变成无穷大。4.2 评价指标计算F1、敏感度和特异度训练时只看loss不够还需要每几个epoch在验证集上评估。眼底血管分割常用指标有accuracy、F1、sensitivity和specificity。sensitivity反映血管漏检率specificity反映背景误判率两个指标都在固定阈值下计算。如果验证集F1很高但sensitivity低说明模型只分割了大血管漏掉了末梢小血管这时要降低预测阈值或增加薄血管切片权重。# metrics.py def compute_metrics(pred_sigmoid, label, threshold0.5): pred (pred_sigmoid threshold).float() label label.float() tp (pred * label).sum().item() fp (pred * (1 - label)).sum().item() fn ((1 - pred) * label).sum().item() tn ((1 - pred) * (1 - label)).sum().item() sensitivity tp / (tp fn 1e-6) specificity tn / (tn fp 1e-6) precision tp / (tp fp 1e-6) f1 2 * precision * sensitivity / (precision sensitivity 1e-6) return { f1: f1, sensitivity: sensitivity, specificity: specificity, accuracy: (tp tn) / (tp tn fp fn) }计算指标时先把预测值转成0/1再和标签做逐像素比较。这里没有用torchmetrics是为了减少训练脚本的额外依赖。小批量验证时可以在batch维度上累加tp/fp/fn/tn最后再算指标。如果发现验证指标剧烈抖动先确认验证集是不是只有几十张patch样本太少时sensitivity会受单张图影响。建议至少保留200张patch做验证集。4.3 训练循环与checkpoint训练结果文件怎么保存与恢复训练循环可以分成train_one_epoch和evaluate两个函数。训练时model.train()验证时model.eval()并包在torch.no_grad()里。保存模型时不要只存权重建议把epoch、optimizer和best_f1都放进一个字典。# train.py def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0.0 for img, mask in loader: img img.to(device) mask mask.to(device) pred model(img) loss criterion(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * img.size(0) return total_loss / len(loader.dataset) # 保存与恢复逻辑 best_f1 0.0 for epoch in range(start_epoch, epochs): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) metrics evaluate(model, val_loader, criterion, device) if metrics[f1] best_f1: best_f1 metrics[f1] torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_f1: best_f1, }, best_model.pth) torch.save(model.state_dict(), last_model.pth)每个epoch结束后用验证集F1判断是否保存best_model.pth这是最常见的做法。训练结果文件里包含optimizer状态是为了之后从断点恢复如果只做推理用torch.load加载后只取model_state_dict即可。建议把每个epoch的loss写入csv曲线会显示模型是真正收敛还是在震荡。如果训练时用了DataParallel保存的model_state_dict会带module前缀。推理加载时把权重名字里的module.去掉再load_state_dict否则会报告unexpected key。这个问题通常在多卡训练完、单卡推理时出现训练结果文件越大越容易忽略。4.4 训练环境与参数速查下面这套参数是从实际项目里提炼出来的起点。GPU显存不同需要优先改patch_size和features而不是只改batch_size。配置项小显存6G中显存12G说明patch_size128256显存不够优先减小patchbatch_size168数值受数据加载速度影响features(32,64,128,256)(64,128,256,512)模型宽度减半显存约降至1/4学习率1e-31e-3Adam分割任务常用1e-4到1e-3调度器CosineAnnealingLRCosineAnnealingLR避免后期loss震荡epoch200150数据量小时需要更多轮次训练环境只要支持CUDA即可PyTorch 2.x和1.x在这组代码上没有本质差别。显存不够时先把batch_size调成1再不行就降patch_size。patch从256降到128同batch下特征图计算量约降为原来的1/4因为长宽各减半显存下降会更明显所以调整优先级很高。5. 用训练好的模型做推理验证与误分割处理5.1 重叠滑动窗口拼接模型训练完成后验证整图不能把原图直接输入网络因为显存和patch训练分布都不允许。常见做法是复用训练时的patch_size以patch_size的一半作为步长滑动推理重叠区域取多次预测的平均值。拼接时维护一张prob_map和一张weight_map每个位置累加预测值和计数最后prob_map除以weight_map就得到整图概率图。这样做的效果是patch边界不会出现一字折痕小血管在重叠区域也会被预测两到三次结果更连续。5.2 形态学后处理去碎屑、填孔洞概率图转二值图后最常出现的两个问题是视盘周围被识别为血管以及细血管内部出现断裂。用OpenCV做两步处理先按连通域面积删除小碎屑再做一次闭运算填补孔洞。import cv2 import numpy as np def postprocess(prob, threshold0.5, min_area10): binary (prob threshold).astype(np.uint8) * 255 n, labels, stats, _ cv2.connectedComponentsWithStats(binary, connectivity8) result np.zeros_like(binary) for i in range(1, n): if stats[i, cv2.CC_STAT_AREA] min_area: result[labels i] 255 return cv2.morphologyEx( result, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)))min_area在普通眼底图分辨率下取10能去掉零星噪点分辨率更高时建议按连通域外接圆半径过滤。闭运算内核不要超过7否则会把相邻血管粘连。后处理只作用于二值图不要对概率图做开闭运算否则概率值会整体偏移。5.3 阈值扫描与TTA的选择最终概率图不一定用0.5做阈值。在验证集上扫描0.3到0.7找到F1最高的阈值保存到模型目录推理时读取。这一步比开TTA更简单提升也更直接敏感度不足就调低阈值背景噪声太多就调高阈值。如果离线分析时间充足再叠加一次TTA即对输入patch做水平翻转、垂直翻转把三次预测翻转回原方向取平均。TTA不改变模型参数只降低单次预测方差对末梢血管连续性有可见改善。部署到在线服务时优先保留重叠推理不开TTA。把最优阈值写进模型配置文件后续复现时不需要重新扫。本文还有配套的精品资源点击获取

UNet眼底血管分割实战：从切片数据集到PyTorch训练推理

相关推荐

React Native列表在OpenHarmony上的高性能封装实践

深入解析 ESLint `space-before-keywords` 规则：关键字前置空格的强制规范与 `keyword-spacing` 演进之路

Crayfish容器引擎：桌面智能体安全运行的核心底座

更多资讯

鳑鲏鱼优化算法结合BP神经网络实现分类预测的Matlab源码详解

改进型YOLOv7实战解析：从结构优化到部署

开源免费、轻量高效的mdput：能否成为Typora的可靠平替

TDengine 零代码接入 SparkplugB：基于 taosExplorer 的 IIoT 数据同步任务配置指南

mdput评测：一款免费轻量级Markdown编辑器，Typora的替代之选

使用 AI SDK 接入 GMI Cloud：OpenAI 兼容协议下的开源权重模型推理实战

今日更新

SpacetimeDB C++ 模块库实战：用 C++20 在数据库内构建 WebAssembly 模块

别踩雷！不是所有 AI 写作工具都靠谱，2026 高校认可工具精选

SSM+MySQL幼儿园管理系统源码解析：从部署到优化

本周精选

超人会飞不算本事：系统稳定依赖清晰规则与边界设计

超人VS蜘蛛侠：拆解超级IP的影响力与传播方法论

基于CNN的调制信号识别：MATLAB实现时频图分类实战

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？