ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

UNet眼底血管分割实战:数据集切片、训练与结果文件复现

UNet眼底血管分割实战:数据集切片、训练与结果文件复现 简介本资源面向医学图像分割初学者与深度学习实践者提供一套基于U-Net的眼底血管二分类分割完整方案解决从数据准备到模型推理的全流程问题。压缩包共216个文件以182张png切片图像、8个py脚本、1个pth权重文件及若干txt日志为主整体约153.92MB涵盖数据集、训练代码与训练结果。项目仅训练10个epochs全局像素准确率即达0.95miou为0.67加大训练轮次后性能可进一步提升。代码支持多尺度训练自动将数据随机缩放至设定尺寸的0.5至1.5倍utils中的compute_gray函数可将mask灰度值保存至txt并自动定义输出通道数学习率采用cos衰减损失与iou曲线、训练日志及最优权重均保存在run_results中可查看各类别iou、recall、precision等指标。推理时只需将图像放入inference目录并运行predict脚本即可无需额外参数。目前已有269人学习适合快速上手眼底血管分割任务。1. 眼底血管分割这件事为什么值得用 UNet 认真做一遍眼底血管分割是医学图像分割里少有的「边界清晰、目标明确、评价标准成熟」的任务。输入是一张彩色眼底照输出是一张二值血管掩膜血管像素为白背景为黑。它直接服务于糖网筛查、动静脉交叉分析、血管宽度测量这些下游场景。很多人第一次接触医学分割就是从 UNet 加眼底血管开始的因为数据规模可控、结构简单、训练几小时就能看到肉眼可见的结果。但真正把它做扎实的人不多大部分卡在三个地方数据集切片方式不对导致血管断裂、损失函数选错导致细血管全丢、训练结果文件不会读导致无法复现。这篇笔记就围绕「切片好的数据集、完整代码、训练结果文件」这三样东西把 UNet 做眼底血管分割的完整路径讲清楚包括每一步的参数含义和翻车点。适合已经会写 PyTorch 基础训练循环、想把这个任务跑通并拿到可用结果的从业者。2. 数据集切片与预处理血管分割的第一道分水岭2.1 眼底血管数据集长什么样切片为什么不能随便切常见的公开眼底血管数据集有 DRIVE、STARE、CHASE_DB1 这几个图像分辨率从 565×584 到 999×960 不等。原始图像数量都不大DRIVE 只有 40 张其中 20 张训练、20 张测试。直接拿整图训练会有一个现实问题显存吃紧而且血管在整图里占比很低正负样本极度不平衡。所以主流做法是先切片再训练。切片不是把图切成 256×256 就完事。眼底图有一个圆形视野区域视野外是黑色背景。如果你按固定网格切会切出大量全黑或者只有视野边缘的块这些块对训练没有正贡献反而让模型学会「预测全黑」。我一般会先做视野掩膜只保留视野内的区域参与切片切片时保证每个 patch 里至少包含一定比例的血管像素比如血管占比低于 1% 的 patch 直接丢弃。这一步做完训练集的有效样本比例会明显上升。另一个关键是切片的重叠。血管是连续结构如果切片之间没有重叠血管在 patch 边界处会被硬切断模型学到的血管在边界处总是断的推理时拼接回去就会出现网格状断裂。常见做法是重叠 25% 到 50%推理时重叠区域取平均或者取最大值。我一般用 50% 重叠拼接时用高斯权重融合边界过渡会平滑很多。2.2 用 Python 做视野提取与重叠切片的完整脚本下面这段代码做三件事提取视野掩膜、按 50% 重叠切片、过滤血管占比过低的 patch。输入是原始眼底图和对应的血管标注图输出是切片后的图像块和掩膜块。import cv2 import numpy as np import os def extract_fov_mask(img): 提取眼底视野掩膜视野内为255视野外为0 gray cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 眼底图视野外接近纯黑阈值分割后取最大连通域 _, thresh cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY) contours, _ cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: return np.zeros_like(gray) largest max(contours, keycv2.contourArea) mask np.zeros_like(gray) cv2.drawContours(mask, [largest], -1, 255, -1) return mask def slice_with_overlap(img, mask, patch_size256, overlap0.5, min_vessel_ratio0.01): 按重叠切片过滤血管占比过低的patch stride int(patch_size * (1 - overlap)) h, w img.shape[:2] patches [] for y in range(0, h - patch_size 1, stride): for x in range(0, w - patch_size 1, stride): img_patch img[y:ypatch_size, x:xpatch_size] mask_patch mask[y:ypatch_size, x:xpatch_size] # 计算血管像素占比血管在mask中为255 vessel_ratio np.sum(mask_patch 127) / (patch_size * patch_size) if vessel_ratio min_vessel_ratio: continue patches.append((img_patch, mask_patch, y, x)) return patches # 使用示例 img cv2.imread(retina.png) vessel_mask cv2.imread(vessel_mask.png, 0) fov extract_fov_mask(img) # 视野外的血管标注置零避免干扰 vessel_mask[fov 0] 0 patches slice_with_overlap(img, vessel_mask, patch_size256, overlap0.5) print(f有效patch数量: {len(patches)})这段代码里patch_size我一般设 256眼底血管在这个尺度下既能保留主干又能看到分支。overlap设 0.5 是经验值再高会让 patch 数量翻倍训练变慢。min_vessel_ratio设 0.01 是过滤掉几乎没血管的块如果数据集血管本身稀疏可以降到 0.005。extract_fov_mask里的阈值 10 是针对常见眼底图的黑边如果你的图黑边不够黑可以适当调高。切片时保留y, x坐标是为了推理时能拼回原图。2.3 数据增强与归一化别让颜色成为干扰项眼底图最大的干扰是光照和颜色差异。不同设备拍的图整体色调可能偏黄或偏红。血管分割只关心结构不关心颜色所以增强时不要用随机颜色抖动那会让模型学到错误的颜色关联。我一般用随机水平翻转、随机垂直翻转、随机旋转 90 度、随机裁剪这些几何变换不改变血管结构。归一化用 ImageNet 的均值和方差就行因为 UNet 编码器通常用预训练权重保持一致的归一化能让预训练特征更好地迁移。提示如果你的数据集里有些图视野外不是纯黑而是有噪声extract_fov_mask的阈值要相应调整否则视野掩膜会不准确切片时会把噪声区域当成有效区域。3. UNet 模型搭建与训练从结构到损失函数的选型理由3.1 为什么眼底血管分割首选 UNet 而不是 DeepLab 或 SegFormerUNet 的核心优势是跳跃连接。编码器逐层下采样提取语义解码器逐层上采样恢复分辨率跳跃连接把编码器的高分辨率特征直接送到解码器。血管是细长结构下采样过程中细血管很容易丢失跳跃连接相当于给解码器留了一条「后悔药」把浅层的边缘信息直接带过来。DeepLab 系列用空洞卷积扩大感受野但空洞卷积对细结构的响应不如跳跃连接直接。SegFormer 用 Transformer 做编码器全局建模能力强但需要更多数据才能训好眼底血管数据集通常只有几十张图Transformer 容易过拟合。我一般用 UNet 加一个预训练的 ResNet34 编码器而不是原版 UNet 的从头训练编码器。预训练权重在自然图像上学到的边缘、纹理特征迁移到眼底血管上依然有效收敛速度会快很多。解码器保持原版 UNet 的结构每层两个 3×3 卷积加 BN 加 ReLU上采样用双线性插值加卷积比转置卷积更稳定不容易出现棋盘伪影。3.2 用 segmentation_models_pytorch 搭一个可训练的 UNet下面代码用segmentation_models_pytorch搭 UNet编码器用 ResNet34 预训练权重输出单通道二值分割。import torch import torch.nn as nn import segmentation_models_pytorch as smp def build_unet(): model smp.Unet( encoder_nameresnet34, # 编码器用ResNet34有预训练权重 encoder_weightsimagenet, # 在ImageNet上预训练 in_channels3, # 输入RGB三通道 classes1, # 输出单通道血管为1背景为0 activationNone # 输出不加激活损失函数里处理 ) return model model build_unet() # 统计参数量 total_params sum(p.numel() for p in model.parameters()) print(f模型参数量: {total_params / 1e6:.2f}M)encoder_name可以换成efficientnet-b0到b7b0 最轻量b4 以上精度更高但显存占用大。encoder_weightsimagenet是关键如果设成 None 就是从头训练收敛会慢很多。classes1对应二分类输出 logits后面用带 sigmoid 的损失函数。activationNone是故意的因为 BCEWithLogitsLoss 内部会做 sigmoid如果这里再加 sigmoid 会重复。3.3 损失函数选 Dice 还是 BCE还是两个一起用眼底血管分割的正负样本比例通常在 1:10 到 1:20 之间血管像素远少于背景。如果只用 BCE模型会倾向于预测背景因为预测背景的 loss 更小。Dice Loss 直接优化预测和标注的重叠度对类别不平衡不敏感。但 Dice Loss 在训练初期梯度不稳定尤其是当预测和标注几乎没有重叠时梯度会很小收敛慢。我一般用 BCE 和 Dice 的加权组合权重各 0.5。BCE 提供稳定的梯度Dice 拉高重叠度。代码实现如下import torch.nn.functional as F class BCEDiceLoss(nn.Module): def __init__(self, bce_weight0.5, dice_weight0.5): super().__init__() self.bce_weight bce_weight self.dice_weight dice_weight def forward(self, pred, target): # pred是logitstarget是0/1掩膜 bce F.binary_cross_entropy_with_logits(pred, target) # Dice计算前先sigmoid pred_prob torch.sigmoid(pred) intersection (pred_prob * target).sum() dice 1 - (2 * intersection 1e-6) / (pred_prob.sum() target.sum() 1e-6) return self.bce_weight * bce self.dice_weight * dice criterion BCEDiceLoss(bce_weight0.5, dice_weight0.5)1e-6是平滑项防止分母为零。如果训练时发现细血管还是丢可以把dice_weight提到 0.7让模型更关注重叠度。如果训练不稳定把bce_weight提到 0.7让梯度更稳。这两个权重没有绝对最优我一般从 0.5/0.5 开始看验证集 Dice 曲线调整。3.4 训练循环与学习率调度什么时候该降学习率训练循环本身不复杂关键是学习率调度和验证时机。我一般用 AdamW初始学习率 1e-4权重衰减 1e-4。学习率调度用 CosineAnnealingLR从 1e-4 余弦降到 1e-6训练 100 个 epoch。每 5 个 epoch 在验证集上算一次 Dice保存 Dice 最高的模型权重。from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR optimizer AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max100, eta_min1e-6) best_dice 0.0 for epoch in range(100): model.train() for img, mask in train_loader: img, mask img.cuda(), mask.cuda() pred model(img) loss criterion(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # 每5个epoch验证一次 if (epoch 1) % 5 0: model.eval() dice_scores [] with torch.no_grad(): for img, mask in val_loader: img, mask img.cuda(), mask.cuda() pred torch.sigmoid(model(img)) pred_bin (pred 0.5).float() intersection (pred_bin * mask).sum() dice (2 * intersection 1e-6) / (pred_bin.sum() mask.sum() 1e-6) dice_scores.append(dice.item()) mean_dice np.mean(dice_scores) print(fEpoch {epoch1}, Val Dice: {mean_dice:.4f}) if mean_dice best_dice: best_dice mean_dice torch.save(model.state_dict(), best_unet_retina.pth)T_max100要和总 epoch 数一致CosineAnnealingLR 会在 100 个 epoch 内把学习率从 1e-4 降到 1e-6。如果训练中途发现验证 Dice 震荡可以把T_max设小一点比如 50让学习率降得更快。best_unet_retina.pth就是训练结果文件里面是模型权重后面推理直接加载这个文件。注意验证时pred 0.5是固定阈值实际推理时这个阈值可以调。如果发现细血管漏检多把阈值降到 0.3 到 0.4能召回更多细血管但会有一些假阳性。这个阈值要在验证集上试不要拍脑袋定。4. 训练结果文件怎么读、怎么复现推理4.1 训练结果文件里到底存了什么torch.save(model.state_dict(), best_unet_retina.pth)保存的是模型参数字典不是整个模型对象。这样做的好处是文件小、加载灵活坏处是加载时必须先有模型结构定义。如果你只拿到一个.pth文件没有模型代码是没法直接用的。所以完整代码和训练结果文件必须配套。加载时先实例化模型再load_state_dictmodel build_unet() model.load_state_dict(torch.load(best_unet_retina.pth, map_locationcpu)) model.eval()map_locationcpu是防止在没有 GPU 的机器上加载时报错。如果要在 GPU 上推理加载后再.cuda()。4.2 推理与拼接把 patch 结果拼回原图推理时对每个 patch 预测然后按切片时的坐标拼回原图。重叠区域用高斯权重融合避免拼接缝。def inference_full_image(model, img, patch_size256, overlap0.5): stride int(patch_size * (1 - overlap)) h, w img.shape[:2] # 初始化输出和权重图 output np.zeros((h, w), dtypenp.float32) weight np.zeros((h, w), dtypenp.float32) # 高斯权重中心高边缘低 gaussian cv2.getGaussianKernel(patch_size, patch_size/4) gaussian gaussian gaussian.T for y in range(0, h - patch_size 1, stride): for x in range(0, w - patch_size 1, stride): patch img[y:ypatch_size, x:xpatch_size] patch_tensor torch.from_numpy(patch).permute(2,0,1).float().unsqueeze(0) / 255.0 with torch.no_grad(): pred torch.sigmoid(model(patch_tensor.cuda())).cpu().numpy()[0,0] output[y:ypatch_size, x:xpatch_size] pred * gaussian weight[y:ypatch_size, x:xpatch_size] gaussian # 归一化避免除零 output output / (weight 1e-6) return outputgaussian是二维高斯核中心权重高边缘权重低这样重叠区域中心预测的贡献更大拼接缝更平滑。weight累加高斯核最后归一化。如果不用高斯权重直接取平均拼接缝会明显一些但也能用。4.3 评价指标Dice、IoU、敏感度、特异度怎么算眼底血管分割常用四个指标Dice、IoU、敏感度召回率、特异度。Dice 和 IoU 衡量整体重叠度敏感度衡量血管召回特异度衡量背景正确率。代码实现如下def compute_metrics(pred_bin, gt): pred_bin和gt都是0/1二值图 tp np.sum((pred_bin 1) (gt 1)) fp np.sum((pred_bin 1) (gt 0)) fn np.sum((pred_bin 0) (gt 1)) tn np.sum((pred_bin 0) (gt 0)) dice (2 * tp 1e-6) / (2 * tp fp fn 1e-6) iou (tp 1e-6) / (tp fp fn 1e-6) sensitivity (tp 1e-6) / (tp fn 1e-6) specificity (tn 1e-6) / (tn fp 1e-6) return dice, iou, sensitivity, specificity1e-6是防止分母为零。敏感度低说明漏检多特异度低说明假阳性多。这两个指标要一起看不能只看 Dice。有时候 Dice 很高但敏感度低说明模型把大部分背景预测对了但细血管漏了不少。5. 避坑与排查眼底血管分割最常见的五个翻车点5.1 现象训练 loss 一直降但验证 Dice 不涨原因通常是过拟合。眼底血管数据集只有几十张图UNet 参数量在 20M 以上很容易记住训练集。解决方法是加数据增强、加权重衰减、早停。我一般会在验证 Dice 连续 10 个 epoch 不涨时停止训练取验证 Dice 最高的权重。另外可以把编码器冻结前几个 stage只训解码器减少过拟合风险。5.2 现象细血管全丢只剩主干原因通常是损失函数对细血管不敏感。细血管像素少BCE 的梯度被背景主导。解决方法是用 Dice Loss 或者提高 Dice 权重另外可以在损失里给血管像素加权血管像素权重设为背景的 5 到 10 倍。还有一个原因是下采样太多细血管在编码器里被池化掉了。可以把 UNet 的下采样次数从 5 次减到 4 次保留更高分辨率的特征。5.3 现象推理结果有网格状拼接缝原因是切片重叠不够或者拼接时没做权重融合。解决方法是把 overlap 从 0.25 提到 0.5拼接时用高斯权重而不是直接平均。如果还有缝检查切片时 stride 是否计算正确stride int(patch_size * (1 - overlap))这个公式里 overlap 是比例不是像素数。5.4 现象加载训练结果文件时报 key 不匹配原因是保存和加载时的模型结构不一致。比如保存时用了encoder_nameresnet34加载时用了encoder_nameefficientnet-b0state_dict 的 key 就对不上。解决方法是保存模型时同时保存一份模型配置加载时按配置实例化。或者直接用torch.save(model, full_model.pth)保存整个模型对象但这样文件会大很多而且依赖模型定义代码。5.5 现象验证集 Dice 很高但实际推理效果差原因是验证集和测试集分布不一致。眼底图在不同设备、不同光照下差异很大如果验证集和训练集来自同一批数据Dice 会虚高。解决方法是留出真正独立的测试集最好来自不同设备或不同批次。另外推理时的预处理要和训练时一致归一化参数、图像尺寸都要对齐。6. 把 Dice 从 0.78 推到 0.82 的三个具体技巧第一个技巧是测试时增强TTA。推理时对同一张图做水平翻转、垂直翻转、旋转 90 度分别预测后再把结果翻转回来取平均。这个操作不增加训练成本但通常能把 Dice 提高 1 到 2 个点。代码实现很简单在推理循环里加几个变换就行。我一般用 4 个变换原图、水平翻转、垂直翻转、水平加垂直翻转。取平均后阈值分割。第二个技巧是后处理去小连通域。血管是连续结构预测结果里会有一些孤立的假阳性小点。用连通域分析把面积小于 50 像素的连通域去掉能减少假阳性提高特异度。但要注意细血管本身可能就只有几十个像素面积阈值不能设太大我一般设 30 到 50 像素根据数据集分辨率调整。第三个技巧是阈值搜索。固定 0.5 阈值不一定最优。在验证集上从 0.3 到 0.7 以 0.05 为步长搜索找 Dice 最高的阈值。这个阈值和 TTA、后处理一起用效果会叠加。我自己的习惯是每次训练完都跑一遍阈值搜索把最优阈值记在训练结果文件的文件名里比如best_unet_retina_thr0.45.pth这样后面复现时不会忘。这三个技巧做完DRIVE 数据集上 Dice 从 0.78 推到 0.82 是可行的。再往上就要改模型结构或者用更复杂的后处理了但那是另一个话题。希望帮到你。本文还有配套的精品资源点击获取
返回列表