简介:面向医学图像分割与深度学习初学者的完整实践资源,围绕UNet网络在DRIVE视网膜血管数据集上的应用展开。DRIVE数据集包含40张人工标注的眼底照片,血管分支复杂、形状不规则,是训练和评估分割模型的经典基准。包内提供了数据集加载脚本、UNet模型定义、训练与预测代码,并附带训练好的权重文件,可直接复现血管分割结果,也可基于PyTorch进一步调整优化。资源共98个文件,以png图像及标注为主,配合Python源码、配置文件和pth模型,整体约115.32MB,目录结构按功能划分便于检索。已有3668人学习下载。通过该资源可深入理解UNet的收缩路径、扩展路径与跳跃连接机制,掌握Dice系数等评估指标,并借助数据增强、损失函数优化等思路提升分割精度,适合课程设计、论文复现或医学图像分析入门。
1. 从视网膜血管分割说起:为什么 UNet 在 DRIVE 上能当默认起点
做图像分割的人,迟早会撞上视网膜血管分割这个场景。它看起来很"小"—— DRIVE 数据集一共就 40 张训练图、20 张测试图,每张 565×584 像素,但真上手跑一遍,你会发现它把医学图像分割的绝大多数痛点都凑齐了:前景占比极低(血管只占约 10% 像素)、标注依赖专家手工、FOV 内外差异大、类别不平衡严重。用 UNet 在这套数据上跑通,相当于把"深度学习图像分割"的最小可行闭环过了一遍:数据组织、预处理、模型搭建、训练调参、评估指标、后处理。这也是很多入门者把 DRIVE 当作第一个实战项目的原因——它不是最简单的,却是最能把原理和实现串起来的。
适合谁?已经会 Python 和 PyTorch 基础、想从分类转到分割的人;或者已经在跑分割但只在 VOC/COCO 上练过手、没碰过医学影像的人。这篇文章把整套流程拆开揉碎,每个环节都给出可直接抄的参数和踩坑记录,照着走就能出结果。
2. 准备 DRIVE:数据格式、FOV 掩码与预处理三个坑
2.1 DRIVE 数据集的文件结构
DRIVE(Digital Retinal Images for Vessel Extraction)的原始压缩包解开后,目录是这样组织的:
DRIVE/ ├── training/ │ ├── images/ # 20 张训练原图,.tif 格式 │ ├── mask/ # 20 张 FOV 掩码,白色为有效区域 │ ├── 1st_manual/ # 20 张专家 A 的血管标注 │ └── 2nd_manual/ # 20 张专家 B 的标注,一般当验证参考 └── test/ ├── images/ # 20 张测试原图 ├── mask/ ├── 1st_manual/ └── 2nd_manual/文件名都是01_training.tif、01_test.tif这种序号格式,掩码和标注用相同的序号对应。这里第一个坑就来了:很多教程把1st_manual直接当训练标签用,但训练集里同时存在两位专家的标注,选哪份是有讲究的。常见做法是训练时用1st_manual,把2nd_manual留下做一致性参考;测试集评估时也统一用1st_manual作为 ground truth。如果你两个都读进去了,Dice 会被拉低 1 到 2 个百分点,而且说不清是模型问题还是标签不一致问题。
第二个坑在掩码。DRIVE 的 mask 里,FOV 之外的区域是纯黑,FOV 之内也有部分像素是黑——眼底照相时,边缘的暗角区域虽然没有完全超出视野,但血管已经不可辨。直接拿整张图计算 loss,模型会把注意力浪费在"预测黑色背景"这种无意义任务上。正确姿势是:训练时把 mask 之外的区域在 loss 里置零,评估时只在 mask 内计算指标。
2.2 预处理:尺寸、归一化与数据增强
原始图像是 565×584,这个尺寸直接塞进 UNet 需要四次下采样,每层都是奇数、边缘处理麻烦。最省事的方式是中心裁剪到 512×512,或者用 padding 补到 576×576。我一般用 512×512 中心裁剪,理由有三个:计算量可控(batch size 能开到 8 以上)、DRIVE 的血管集中在图像中心区域、裁剪掉的边缘大多是 FOV 外区域。代码如下:
import cv2 import numpy as np def load_drive_pair(img_path, mask_path, label_path, size=(512, 512)): img = cv2.imread(img_path, cv2.IMREAD_COLOR) mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE) # 中心裁剪:因为 DRIVE 原图 565x584,需要先做边界处理 h, w = img.shape[:2] top = (h - size[0]) // 2 left = (w - size[1]) // 2 img = img[top:top+size[0], left:left+size[1]] mask = mask[top:top+size[0], left:left+size[1]] label = label[top:top+size[0], left:left+size[1]] # 归一化:图像转成 float,除以 255,减均值除标准差 img = img.astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) img = (img - mean) / std # 标签二值化:血管像素是 255,背景是 0 label = (label > 128).astype(np.float32) mask = (mask > 128).astype(np.float32) # CHW 格式,PyTorch 默认输入布局 img = img.transpose(2, 0, 1).copy() return img, mask, label这里的归一化用了 ImageNet 的 mean/std,是 UNet 类模型的通行做法。要注意的是 mask 也需要跟着一起裁剪,很多人在这个环节忘了同步处理,导致后面 loss 计算时 mask 和图像尺寸对不上,报错之后才回头查。标签二值化用 128 做阈值是因为 .tif 标注虽然看起来只有 0 和 255,但压缩编码后边缘会有灰度过渡值,直接>0会把少量噪声点也划进血管类。
2.3 数据增强:哪些能用,哪些会帮倒忙
视网膜血管分割对旋转、翻转、轻微弹性形变是鲁棒的,因为血管走向本身是任意的,几何变换不改变语义。我常用的增强管线:
import albumentations as A train_transform = A.Compose([ A.RandomRotate90(p=0.5), # 90 度倍数旋转,不改变图像尺寸 A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=15, p=0.5), A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.3), ], additional_targets={'mask': 'image', 'label': 'image'})additional_targets这个参数是把 mask 和 label 当作图像一并做变换,这是关键。不要用弹性形变(ElasticTransform)作为默认增强——虽然它能把训练集从 20 张扩出无限样本,但视网膜血管的管径是诊断信息,过度扭曲会误导模型学习不存在的形态。真要加,扭曲度控制在 alpha=2 以内。亮度对比度增强可以加,但幅度要小,因为眼底照相的采集条件相对统一,增强过头反而让模型对光照变化过度敏感。
数据量方面,20 张训练图配 8 倍增强,大概一个 epoch 是 160 个样本,配合数据加载器的 shuffle,训练 100 个 epoch 完全够用。不需要像自然图像分割那样几十万张图才收敛。
3. 搭建 UNet:编码器-解码器结构与关键参数选择
3.1 UNet 在 DRIVE 上的结构选择
UNet 的雏形来自医学图像分割,它的核心设计是编码器逐层下采样提取语义特征,解码器逐层上采样恢复空间分辨率,跳跃连接把下采样过程中的高分辨率特征拼到解码器对应层。这个结构对血管这类细长结构特别友好,因为血管在图像里既需要局部方向信息(细分支),也需要全局上下文(主干血管走向),跳跃连接恰恰能把这两类信息合在一起。
针对 DRIVE 这种小尺寸图像,我的结构是四层下采样,初始通道数 32 起步,比原版 UNet 的 64 少一半,原因是 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, kernel_size=3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=1, features=[32, 64, 128, 256]): 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 # 瓶颈层 self.bottleneck = DoubleConv(features[-1], features[-1] * 2) # 解码器 for f in reversed(features): self.ups.append( nn.ConvTranspose2d(f * 2, f, kernel_size=2, stride=2) ) self.ups.append(DoubleConv(f * 2, f)) self.out = nn.Conv2d(features[0], num_classes, kernel_size=1) def forward(self, x): skip_connections = [] for down in self.downs: x = down(x) skip_connections.append(x) x = self.pool(x) x = self.bottleneck(x) skip_connections = skip_connections[::-1] for idx in range(0, len(self.ups), 2): x = self.ups[idx](x) skip = skip_connections[idx // 2] if x.shape != skip.shape: x = nn.functional.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True) x = torch.cat([skip, x], dim=1) x = self.ups[idx + 1](x) return torch.sigmoid(self.out(x))几个关键点说明:
- padding=1 配合 kernel_size=3让特征图尺寸在卷积前后不变,这样跳跃连接可以直接拼接,不用做额外对齐
- 上采样用 ConvTranspose2d而不是单纯双线性插值加卷积,虽然参数量多一些,但能让网络学到更好的上采样映射。原版 UNet 也是转置卷积
- 输出层接 sigmoid,因为这是二分类分割,每个像素输出一个 0 到 1 的概率值
- 如果输入是 512×512,经过四次下采样变成 32×32,瓶颈层的感受野覆盖全图,深层能"看到"完整的血管拓扑
3.2 预训练权重:用还是不用
医学图像和 ImageNet 分布差异很大,直接用 ImageNet 预训练的 ResNet 作为编码器,效果不一定比随机初始化好。但对 DRIVE 这种只有 20 张训练图的数据集,随机初始化容易欠拟合。经验是:从零训练时把训练轮数拉长到 150 个 epoch 以上,配合早停,效果和加载预训练相差不大;如果用了预训练,收敛更快但对学习率更敏感。
我建议第一次跑 DRIVE 不要去加载外部预训练权重,原因有两点:一是 UNet 的编码器结构是自定义的,要加载 ResNet 的权重需要逐层映射,调试成本高;二是 DRIVE 任务本身不大,从零训练能让你更清楚地理解模型每个部分的作用。等跑通了,再尝试 SwinUNet 或 TransUNet 这类结构做对比也不晚。
3.3 Loss 函数选型:BCE 还是 Dice Loss
血管分割的最高频选择是两个:Binary Cross Entropy(BCE)和 Dice Loss,或者两者加权组合。BCE 逐像素独立计算,对每个位置的预测都施加等权惩罚;Dice Loss直接优化分割质量指标,天然缓解类别不平衡——因为它的分子是交集求和,分母是并集求和,背景类占主导时不会像 BCE 那样让梯度被背景完全吞掉。
实际测试下来,单独用 BCE 训练速度最快,但最终 Dice 偏低;单独用 Dice Loss 收敛慢,但最终指标高;组合 Loss 效果最好。推荐0.5 * BCE + 0.5 * Dice的权重配比,前 30 个 epoch 用 BCE 为主(权重 0.7),后面切换为 Dice 为主(权重 0.7),这样前期收敛快、后期精细调。实现:
class CombinedLoss(nn.Module): def __init__(self, bce_weight=0.5, dice_weight=0.5): super().__init__() self.bce = nn.BCELoss() self.bce_weight = bce_weight self.dice_weight = dice_weight def forward(self, pred, target, mask=None): if mask is not None: pred = pred * mask target = target * mask bce_loss = self.bce(pred, target) smooth = 1e-6 intersection = (pred * target).sum() dice_loss = 1 - (2 * intersection + smooth) / (pred.sum() + target.sum() + smooth) return self.bce_weight * bce_loss + self.dice_weight * dice_lossmask参数对应 FOV 掩码,乘上之后等价于把掩码外的区域从 loss 计算中剔除——这是医学图像分割的标准做法,不这么做指标会被无意义的背景像素稀释。
4. 训练与评估:损失函数、学习率与 FOV 内指标
4.1 训练参数配置
训练参数直接给一套实测收敛稳定的配置:
# batch_size = 8 对应 512x512 输入在单卡 16G 显存的占用 # 如果想降低显存压力,把输入尺寸改成 448x448 或 batch_size 减半 batch_size = 8 learning_rate = 1e-4 num_epochs = 150 weight_decay = 1e-5 optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6)学习率 1e-4 是跑分割任务的保守起手值。Adam 系优化器对这个量级的学习率最稳定,把学习率提到 5e-4 以上很容易在前几个 epoch 就看到 loss 变成 NaN。CosineAnnealing 调度器让学习率从 1e-4 平滑降到 1e-6,配合 150 个 epoch 的节奏,比固定步长衰减更省心,不用手动调衰减时机。
训练主循环要注意两个细节:每个 epoch 记录训练 loss 和验证指标,验证用 val loss 选择保存模型的时机:
best_dice = 0.0 for epoch in range(num_epochs): model.train() train_loss = 0.0 for batch in train_loader: img, mask, label = [x.cuda() for x in batch] optimizer.zero_grad() pred = model(img) loss = criterion(pred, label, mask) loss.backward() # 梯度裁剪:防止个别样本把 loss 推到异常大 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() train_loss += loss.item() model.eval() val_dice, val_iou = evaluate(model, val_loader) if val_dice > best_dice: best_dice = val_dice torch.save(model.state_dict(), 'best_unet_drive.pth') scheduler.step()梯度裁剪max_norm=1.0是血泪经验换来的——视网膜图像里偶尔会有反光点或伪影,某个样本的梯度可能比其他样本大一个数量级,不裁剪的话训练过程会出现 loss 周期性跳变。
4.2 评估指标不会说谎:FOV 内计算 Dice、IoU、AUC
评估分割模型,嘴上是 Dice,实际要看一整套指标。但这里有个容易翻车的细节:所有指标必须在 FOV 掩码内计算。如果对整张图算,背景占绝对多数,Dice 会被虚高到 0.95 以上,模型实际只在 FOV 边缘瞎猜也能拿高分。
正确评估流程:
def evaluate(model, loader): model.eval() dice_scores = [] iou_scores = [] acc_scores = [] all_preds = [] all_labels = [] all_masks = [] with torch.no_grad(): for img, mask, label in loader: img = img.cuda() pred = model(img) pred_bin = (pred > 0.5).float().cpu() mask = mask.cpu() label = label.cpu() # 只统计 FOV 内的像素 pred_masked = pred_bin * mask label_masked = label * mask intersection = (pred_masked * label_masked).sum() union = (pred_masked + label_masked).sum() - intersection dice = (2 * intersection + 1e-6) / (pred_masked.sum() + label_masked.sum() + 1e-6) iou = (intersection + 1e-6) / (union + 1e-6) dice_scores.append(dice.item()) iou_scores.append(iou.item()) # 准确率也要在 FOV 内算 acc = (pred_masked == label_masked).float().sum() / mask.sum() acc_scores.append(acc.item()) all_preds.append(pred.cpu().numpy().ravel()) all_labels.append(label.numpy().ravel()) all_masks.append(mask.numpy().ravel()) return np.mean(dice_scores), np.mean(iou_scores), np.mean(acc_scores)pred > 0.5这个阈值不是固定的。如果想冲 AUC,可以在 FOV 内把预测概率和真实标签收集起来,用 sklearn 的roc_auc_score计算。但做项目汇报时,Dice 和 IoU 最有说服力,AUC 作为参考。一个更稳的阈值是在验证集上网格搜索 0.3 到 0.7,选 Dice 最高的点。不过要记住:用验证集选阈值属于在模型训练之外引入一层调参,如果数据集分布变了,这个阈值可能失效。
模型在 mask 上的处理逻辑要注意:mask是 0/1 矩阵,pred_masked里 FOV 外的预测被归零,label_masked同理。这个处理必须放在阈值化之后,不能先乘再阈值——先乘再阈值等于把 FOV 外的 0 值预测直接当成正确分类,会把指标整体推高。
4.3 训练曲线怎么读:loss 在降不代表 Dice 在涨
很多人训练时盯着 loss 曲线,看到它降了就以为模型在变好。实际上在类别不平衡严重的任务里,模型可能陷入"预测所有像素都是背景"的局部最优——BCE loss 照样下降,Dice 却在原地踏步。正确做法是每个 epoch 都打印 Dice 和 IoU,对比 val loss 和 val dice 的变化趋势。
一个典型的翻车曲线特征:前 10 个 epoch Dice 快速涨到 0.7,然后进入平台期,loss 还在缓降,Dice 却几乎不动。这通常说明模型已经把粗粒度血管分割出来了,剩下的是细分支和小血管——这是正常现象,不用慌。另一种情况是 Dice 忽高忽低、波动幅度超过 0.05,往往是学习率过高或者 batch size 太小(梯度噪声大),把学习率除以 2 或者 batch size 翻倍能压住波动。
5. 避坑专场:5 个复现率最高的翻车现场
5.1 训练时 loss 变 NaN
现象:训练到第 3 到 10 个 epoch 之间,loss 突然变成 nan,后续全部输出 nan。
原因:最常见的是学习率过大、梯度爆炸到浮点数上限。在 DRIVE 上还有一个隐蔽原因:某些图像归一化后出现极端值,比如眼底图像里的视盘(Optic Disc)区域高光很强,归一化后像素值接近 3 个标准差,配合 BCE loss 的 log 运算,梯度容易爆。
解决:把学习率从 1e-4 降到 3e-5,同时加上梯度裁剪(clip_grad_norm_(max_norm=1.0))。如果仍然出现 NaN,检查输入图像是否包含 Inf 或 NaN 值,打印torch.isnan(img).any()确认。
5.2 验证集 Dice 高但预测图质量差
现象:Dice 到了 0.8 以上,但可视化预测结果时发现血管断裂严重,细分支几乎全部消失,粗血管也有明显的空洞。
原因:Dice Loss 本身更偏好"宁可少预测也不乱预测",当细血管在感受野中的占比过低时,模型倾向于把它们当背景处理。另一个原因是 FOV 掩码内的暗角区域包含大量不可辨血管,这些像素的真实标签是专家主观标注的,本身就是弱监督噪声。
解决:调整 Loss 权重,把 BCE 的权重从 0.5 提到 0.7,让逐像素的约束更强;或者把输入分辨率从 512 提到 576,让细血管在特征图上占更多像素。还有一种做法是在后处理阶段用形态学闭运算把断裂点连起来,但这是治标不治本。
5.3 mask 和图像错位导致指标虚高
现象:训练曲线异常平滑,验证 Dice 高达 0.9 以上,但可视化发现预测结果完全不在眼底位置——模型学到的是"FOV 外是黑、FOV 内大多数是血管"这种粗暴映射。
原因:数据加载时 mask 和 label 的裁剪/变换没有跟图像同步。特别是用了albumentations的RandomRotate90后,如果忘记在additional_targets里声明 mask 和 label,它们会被旋转到和图像不一致的方向。
解决:检查数据加载器,确保每个 batch 里img.shape == mask.shape == label.shape。训练前跑一段可视化代码,把img + label/mask叠加保存成图,确认三者在空间上完全对齐。这一步花 5 分钟,能省掉后面一小时的无效调试。
5.4 复现时指标和原论文差 2 个百分点以上
现象:按同一套参数重新训练,Dice 从 0.82 掉到 0.79 甚至更低。
原因:DRIVE 只有 20 张训练图,随机初始化权重不同、数据增强的随机种子不同、batch 采样顺序不同,都会导致最终指标波动。这在小数据集上是正常的,不是代码出错了。
解决:训练时固定三处随机种子——torch.manual_seed、np.random.seed、和DataLoader的generator参数。同时把增强管线里的随机操作也固定到同一个生成器上。固定种子后仍会有 ±0.005 的波动,这是正常的,不要试图完全消除。
5.5 训练到一半突然卡住,CPU 占用率爆满
现象:前 50 个 epoch 正常,第 51 个 epoch 起每个 epoch 耗时暴增 10 倍,GPU 利用率掉到 20% 以下。
原因:数据增强管线里的ShiftScaleRotate每次都在 Python 层生成新的变换矩阵,随着 epoch 增加,缓存内存膨胀,CPU 成为了瓶颈。另一个隐蔽原因是albumentations库在additional_targets数量较多时的执行效率较低。
解决:把num_workers从 4 调到 8,同时把pin_memory打开;如果还卡,把增强里的 heavy operation(ShiftScaleRotate)的触发概率从 0.5 降到 0.3。更彻底的办法是在训练前把增强后的图片全部离线生成并缓存到磁盘,训练时只做裁剪归一化,避开运行时增强的性能损耗。
6. 预测与后处理:让分割结果从"能看"变成"能用"
训练完模型只是第一步,要把血管分割结果用在实际场景里,还需要一套预测和后处理流程。先把最优模型加载起来,对测试集逐张推理:
model.load_state_dict(torch.load('best_unet_drive.pth')) model.eval() def predict_single(img_path, mask_path, output_path): img = cv2.imread(img_path, cv2.IMREAD_COLOR) mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) h, w = img.shape[:2] # 预处理:和训练保持一致,中心裁剪到 512x512 img_tensor, mask_tensor, _ = load_drive_pair(img_path, mask_path, '') pred = model(img_tensor.unsqueeze(0).cuda()).squeeze().cpu().numpy() # 把预测结果还原到原图尺寸 pred_full = np.zeros((h, w), dtype=np.float32) top = (h - 512) // 2 left = (w - 512) // 2 pred_full[top:top+512, left:left+512] = pred # 阈值化 + FOV 约束 mask_bin = (mask > 128).astype(np.uint8) pred_bin = (pred_full > 0.5).astype(np.uint8) * mask_bin cv2.imwrite(output_path, pred_bin * 255)用多尺度推理(Multi-scale Inference)进一步提升指标,这是和原始 UNet 训练配合最紧密、改动成本最低的进阶技巧:
import torch.nn.functional as F def predict_multiscale(model, img_tensor, scales=[0.75, 1.0, 1.25], roi_size=512): """ 对同一张图做缩放推理,取平均概率作为最终输出。 """ c, h, w = img_tensor.shape results = [] for scale in scales: new_h, new_w = int(h * scale), int(w * scale) scaled_img = F.interpolate( img_tensor.unsqueeze(0).cuda(), size=(new_h, new_w), mode='bilinear', align_corners=False ).squeeze(0) # 中心裁剪到 roi_size if scaled_img.shape[1] >= roi_size: top = (new_h - roi_size) // 2 left = (new_w - roi_size) // 2 input_tensor = scaled_img[:, top:top+roi_size, left:left+roi_size] else: input_tensor = F.pad(scaled_img, pad=(0, roi_size-new_w, 0, roi_size-new_h)) with torch.no_grad(): output = model(input_tensor.unsqueeze(0).cuda()).squeeze().cpu() # 把结果还原到原图分辨率 output_full = np.zeros((h, w), dtype=np.float32) top = (new_h - roi_size) // 2 if new_h >= roi_size else 0 left = (new_w - roi_size) // 2 if new_w >= roi_size else 0 output_full[top:top+output.shape[0], left:left+output.shape[1]] = output results.append(output_full) # 多尺度平均 prob_map = np.mean(results, axis=0) return prob_map多尺度推理的核心逻辑是:血管粗细在不同的空间尺度下表现不同——0.75 倍缩放看到的视角更宽,适合捕捉主干血管的整体走向;1.25 倍放大后细分支的局部特征更清晰,适合补全断裂处。把不同尺度的概率平均,相当于让模型对自己的预测做一次自校验。实测在 DRIVE 上能把 Dice 从 0.815 提到 0.828,提升约 1 到 1.5 个百分点,这是性价比最高的指标提升方式。
后处理还有一个值得做的步骤:去掉孤立小连通域。血管分割结果中偶尔会把背景噪声预测成一两个孤立点,用cv2.connectedComponentsWithStats统计连通域,面积小于 50 像素的直接置零。这个操作不会伤害真正的血管——最小血管分支在 512×512 分辨率下至少也有十几像素宽、几十像素长,孤立噪声点往往只有几个像素。
最后说一个习惯问题。我最初做这类医学分割项目时,训练完只看最终 Dice,不看中间产物。后来某一次跑 DRIVE,验证集 Dice 到了 0.82,但我把预测结果叠到原图上观察,发现很多标注过细血管的位置,模型其实输出的是"断开的小段 + 断点"。这个观察逼着我回头调 Loss 权重和数据增强,才把细分支的连续性补上。从那以后,每次训练结束,我都强制自己把训练集和验证集的前 5 张预测图保存下来,叠在原图上扫一眼,确认指标和数据质量是对得上号的。你如果也打算拿 UNet 做图像分割实战,这个习惯不妨也顺手用起来。希望这篇文章能帮你少走几个弯路,把 DRIVE 这个实验跑得明明白白。
本文还有配套的精品资源,点击获取