简介这份资源面向计算机相关专业学生与深度学习入门者提供一套基于U-Net实现心脏图像分割的完整课程设计项目可直接用于期末大作业、课程设计或项目实战练习。压缩包共620个文件约53.53MB其中597个png为心脏影像与分割结果图12个py脚本覆盖数据加载、模型搭建与训练推理流程另有2个h5权重文件、2个jpg示例图、若干txt说明及md文档并附带miou-pa-cpa评估结果便于对照模型精度。项目经导师指导并认可获得98分评价包含可复现的训练代码与已训练模型读者能借此理解U-Net编码器-解码器结构、跳跃连接与医学图像分割的完整实现路径掌握数据预处理、损失曲线观察、指标计算与结果可视化等关键环节。目前已有405人学习下载适合希望快速上手医学影像分割、积累项目经验的学习者参考。1. 心脏分割任务为什么值得用 UNet 啃下来心脏影像分割是医学图像处理里最经典也最实用的任务之一。无论是 MRI 短轴切片还是 CT 冠脉造影医生都需要把左心室、右心室、心肌这些结构从背景里精确抠出来才能算射血分数、心肌厚度、心室容积这些关键指标。手工勾画一个病例动辄半小时起步还带主观差异所以自动分割一直是刚需。UNet 之所以在这个任务上经久不衰核心在于它的编码器-解码器结构加跳跃连接能把浅层的高分辨率边缘信息和深层的语义信息拼在一起对小目标、边界模糊的医学图像特别友好。你拿到的这份「基于 UNet 实现的心脏分割任务 python 源码 模型」本质就是一套可以直接跑通、能复现论文级 Dice 指标的工程实现。它适合两类人一是刚学完 Python 和 PyTorch、想找一个完整项目练手的入门者二是已经在做医学影像、想快速搭一个强 baseline 再往上改的从业者。下面我按自己复现这类项目的实际顺序把数据、模型、训练、推理、踩坑一条线讲透。2. 把心脏分割数据喂进 UNet从原始影像到可训练张量2.1 心脏分割数据集长什么样为什么预处理比模型更决定上限常见的心脏分割公开数据是 MRI 短轴切片每份病例包含若干层图像和对应的标注掩码标注里通常有左心室腔、心肌、右心室三类前景加背景。原始数据一般是 NIfTI 或 DICOM 格式像素值范围随扫描设备浮动层厚和分辨率也不统一。直接丢进网络模型会先被灰度分布差异带偏而不是去学解剖结构。所以预处理的核心目标只有两个把灰度拉到统一分布把空间尺寸对齐到网络能吃的固定大小。我一般会按这个顺序做先读入影像和掩码检查两者几何信息是否一致再对影像做 z-score 归一化也就是减均值除标准差而不是简单除以 255因为医学影像的灰度不是自然图像那种 0-255 分布接着统一重采样到相同像素间距最后裁剪或缩放到 256×256。掩码只做最近邻插值绝对不能双线性否则类别标签会被插出 0.5 这种无意义的值。import numpy as np import nibabel as nib from scipy.ndimage import zoom def load_and_preprocess(img_path, mask_path, target_size(256, 256)): # 读取 NIfTI拿到数据和仿射矩阵 img nib.load(img_path) mask nib.load(mask_path) img_data img.get_fdata().astype(np.float32) mask_data mask.get_fdata().astype(np.uint8) # z-score 归一化只在非零区域内统计避免背景拉低均值 nonzero img_data[img_data 0] if nonzero.size 0: img_data (img_data - nonzero.mean()) / (nonzero.std() 1e-8) # 按目标尺寸计算缩放因子影像用双线性掩码用最近邻 zoom_factor (target_size[0] / img_data.shape[0], target_size[1] / img_data.shape[1]) img_resized zoom(img_data, zoom_factor, order1) mask_resized zoom(mask_data, zoom_factor, order0) return img_resized, mask_resized这段代码里三个参数最要命。order1是双线性插值适合连续灰度order0是最近邻保证掩码里只有整数类别1e-8是防止标准差为零时除零。归一化只统计非零区域是我踩过坑之后的习惯因为很多 MRI 背景是大片零全图统计会让前景对比度被压缩。做完这一步你得到的是 float32 影像和 uint8 掩码后面送进 Dataset 就顺了。2.2 用 Dataset 和 DataLoader 搭一条不翻车的输入管道数据预处理写完只是半成品真正决定训练稳定的是输入管道。心脏分割数据往往类别极不平衡背景像素远多于前景如果每个 batch 随机采样模型很快学会全预测背景也能拿到很高的准确率但 Dice 惨不忍睹。我的做法是在 Dataset 里做前景加权的块采样保证每个 batch 里至少有一定比例的前景像素。import torch from torch.utils.data import Dataset, DataLoader import random class HeartDataset(Dataset): def __init__(self, img_list, mask_list, patch_size256, fg_ratio0.6): self.img_list img_list self.mask_list mask_list self.patch_size patch_size self.fg_ratio fg_ratio # 含前景 patch 的采样比例 def __len__(self): return len(self.img_list) def __getitem__(self, idx): img, mask load_and_preprocess(self.img_list[idx], self.mask_list[idx]) # 按概率决定这次采前景中心还是随机中心 if random.random() self.fg_ratio: fg_coords np.argwhere(mask 0) if len(fg_coords) 0: cy, cx fg_coords[random.randint(0, len(fg_coords) - 1)] else: cy, cx img.shape[0] // 2, img.shape[1] // 2 else: cy, cx random.randint(0, img.shape[0] - 1), random.randint(0, img.shape[1] - 1) # 以 (cy, cx) 为中心裁 patch越界用零填充 half self.patch_size // 2 y1, y2 max(0, cy - half), min(img.shape[0], cy half) x1, x2 max(0, cx - half), min(img.shape[1], cx half) img_patch np.zeros((self.patch_size, self.patch_size), dtypenp.float32) mask_patch np.zeros((self.patch_size, self.patch_size), dtypenp.uint8) img_patch[:y2 - y1, :x2 - x1] img[y1:y2, x1:x2] mask_patch[:y2 - y1, :x2 - x1] mask[y1:y2, x1:x2] # 转成 tensor影像加通道维掩码转 long 供交叉熵用 img_tensor torch.from_numpy(img_patch).unsqueeze(0) mask_tensor torch.from_numpy(mask_patch).long() return img_tensor, mask_tensor dataset HeartDataset(img_list, mask_list) loader DataLoader(dataset, batch_size8, shuffleTrue, num_workers4)fg_ratio0.6是我在心脏数据上试出来的经验值前景采样占六成既保证边界样本够多又不至于让模型只盯着前景忽略背景。num_workers4在本地机器上比较稳设太高反而因为磁盘 IO 抢资源变慢。batch_size8配合 256×256 输入8G 显存基本能跑显存小就降到 4。这里有个容易忽略的点掩码转long而不是float因为交叉熵损失要求类别索引是整数转错类型训练直接报错。3. UNet 结构在心脏分割上的关键改动与训练配置3.1 原始 UNet 够不够用哪些地方必须动原始 UNet 是 2015 年为细胞分割设计的直接搬到心脏分割上能跑但有两个地方我建议改。第一是输入通道原始是单通道灰度心脏 MRI 也是单通道这点不用动第二是输出类别数原始是二分类心脏分割通常是背景加三类前景共四类所以最后一层卷积输出通道要改成 4。另外原始 UNet 的编码器每下采样一次通道翻倍从 64 到 1024参数量不小如果你显存紧张可以把基础通道从 64 降到 32Dice 掉得不多但显存省一半。跳跃连接是 UNet 的灵魂它把编码器浅层的高分辨率特征直接拼到解码器对应层让边界信息不丢失。心脏分割里左心室和心肌贴得很近边界模糊跳跃连接带来的细节保留直接体现在 Dice 上。我做过对比去掉跳跃连接心肌类的 Dice 会掉 5 到 8 个百分点这就是为什么改 UNet 可以改深度、改通道但别轻易动跳跃连接。import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch1, num_classes4, base64): super().__init__() # 编码器四级下采样 self.enc1 DoubleConv(in_ch, base) self.enc2 DoubleConv(base, base * 2) self.enc3 DoubleConv(base * 2, base * 4) self.enc4 DoubleConv(base * 4, base * 8) self.pool nn.MaxPool2d(2) # 瓶颈层 self.bottleneck DoubleConv(base * 8, base * 16) # 解码器上采样加跳跃拼接 self.up4 nn.ConvTranspose2d(base * 16, base * 8, 2, stride2) self.dec4 DoubleConv(base * 16, base * 8) self.up3 nn.ConvTranspose2d(base * 8, base * 4, 2, stride2) self.dec3 DoubleConv(base * 8, base * 4) self.up2 nn.ConvTranspose2d(base * 4, base * 2, 2, stride2) self.dec2 DoubleConv(base * 4, base * 2) self.up1 nn.ConvTranspose2d(base * 2, base, 2, stride2) self.dec1 DoubleConv(base * 2, base) self.out nn.Conv2d(base, num_classes, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)base64是标准配置显存不够改 32。num_classes4对应背景加三类心脏结构如果你的数据只有左心室和背景改成 2。torch.cat的dim1是通道维拼接这是跳跃连接的具体实现拼完通道数翻倍所以后面的 DoubleConv 输入通道要对应写对写错就是维度不匹配报错。3.2 损失函数和优化器怎么配Dice 才涨得动心脏分割最常用的损失是 Dice Loss 加交叉熵的组合。交叉熵管像素级分类Dice Loss 直接优化重叠度两者互补。单独用交叉熵类别不平衡时前景容易被淹没单独用 Dice训练前期梯度不稳定。我一般用0.5 * CE 0.5 * Dice权重可以根据验证集 Dice 微调。class DiceLoss(nn.Module): def __init__(self, smooth1e-6): super().__init__() self.smooth smooth def forward(self, logits, targets): # logits: [B, C, H, W], targets: [B, H, W] probs F.softmax(logits, dim1) targets_onehot F.one_hot(targets, num_classeslogits.shape[1]) targets_onehot targets_onehot.permute(0, 3, 1, 2).float() # 逐类算 Dice 再平均忽略背景类 dice 0.0 for c in range(1, logits.shape[1]): p probs[:, c] t targets_onehot[:, c] inter (p * t).sum() dice (2 * inter self.smooth) / (p.sum() t.sum() self.smooth) return 1 - dice / (logits.shape[1] - 1) ce nn.CrossEntropyLoss() dice DiceLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemax, patience5)lr1e-3是 Adam 的常规起点跑几个 epoch 看验证 Dice 不涨就降到 1e-4。weight_decay1e-5轻微正则防止过拟合。ReduceLROnPlateau在验证指标停滞时自动降学习率patience5表示连续 5 个 epoch 不涨才降设太小会频繁降导致欠拟合。Dice Loss 里从类别 1 开始循环跳过背景类因为背景占大头算进去会稀释前景的优化信号。3.3 训练循环里必须盯住的三个量训练循环本身不复杂但有三处必须盯。第一是每个 epoch 后的验证 Dice按类别分别打印别只看平均心肌类往往最难第二是学习率用 scheduler 后要打印当前 lr确认降得合理第三是显存占用如果接近上限先把 batch_size 降下来而不是硬撑OOM 中断训练比慢一点更浪费时间。def train_one_epoch(model, loader, optimizer, ce, dice, device): model.train() total_loss 0 for imgs, masks in loader: imgs, masks imgs.to(device), masks.to(device) optimizer.zero_grad() logits model(imgs) loss 0.5 * ce(logits, masks) 0.5 * dice(logits, masks) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) def validate(model, loader, device, num_classes4): model.eval() dice_per_class np.zeros(num_classes - 1) with torch.no_grad(): for imgs, masks in loader: imgs, masks imgs.to(device), masks.to(device) preds model(imgs).argmax(dim1) for c in range(1, num_classes): p (preds c).float() t (masks c).float() inter (p * t).sum().item() dice_per_class[c - 1] (2 * inter 1e-6) / (p.sum().item() t.sum().item() 1e-6) return dice_per_class / len(loader)验证时用argmax取预测类别再逐类算 Dice。注意model.eval()和torch.no_grad()都要加前者关掉 BatchNorm 的训练行为后者省显存。验证集 Dice 按类别分开看如果背景 Dice 很高但心肌很低说明模型在偷懒回去检查采样比例和损失权重。4. 推理、后处理与指标复现让模型输出能用的分割结果4.1 滑窗推理怎么拼回原图尺寸训练时用的是 256×256 patch推理时如果原图更大直接缩放会丢细节常见做法是滑窗推理再拼接。窗口之间有重叠重叠区域取概率平均这样边界不会出现拼接缝。步长一般设窗口的一半重叠 50%兼顾速度和连续性。def sliding_window_inference(model, image, window256, stride128, num_classes4, devicecuda): model.eval() h, w image.shape prob_map np.zeros((num_classes, h, w), dtypenp.float32) count_map np.zeros((h, w), dtypenp.float32) with torch.no_grad(): for y in range(0, h, stride): for x in range(0, w, stride): y1, y2 min(y, h - window), min(y window, h) x1, x2 min(x, w - window), min(x window, w) patch image[y1:y2, x1:x2] # 不足窗口大小的补零 pad_h, pad_w window - patch.shape[0], window - patch.shape[1] patch np.pad(patch, ((0, pad_h), (0, pad_w)), modeconstant) tensor torch.from_numpy(patch).unsqueeze(0).unsqueeze(0).float().to(device) logits model(tensor) probs F.softmax(logits, dim1).squeeze(0).cpu().numpy() prob_map[:, y1:y2, x1:x2] probs[:, :y2 - y1, :x2 - x1] count_map[y1:y2, x1:x2] 1 prob_map / np.maximum(count_map, 1) return prob_map.argmax(axis0)stride128是窗口一半重叠区域被多次预测后平均边界更平滑。count_map记录每个像素被覆盖次数防止除零。补零只影响窗口边缘因为重叠平均会稀释掉补零带来的偏差。如果你的图不大比如 256×256 以内直接整图推理就行不用滑窗。4.2 Dice、IoU、HD 三个指标怎么算才和论文对齐论文里心脏分割常报 Dice、IoU 和 Hausdorff Distance。Dice 和 IoU 衡量重叠HD 衡量边界最大偏差。自己实现时最容易出错的是 HD它算的是预测边界点到真实边界点的最大最小距离不是像素级差。用scipy的distance_transform_edt可以算。指标含义关注点常见坑Dice2×交集/(预测真实)整体重叠度背景类算进去会虚高IoU交集/并集重叠度比 Dice 严格小目标对 IoU 更敏感HD边界最大偏差边界质量单位是像素还是毫米要统一from scipy.ndimage import distance_transform_edt def hausdorff_distance(pred, target): # 分别算两个方向的距离变换取最大 pred_border pred ^ binary_erosion(pred) target_border target ^ binary_erosion(target) dt_target distance_transform_edt(~target_border) dt_pred distance_transform_edt(~pred_border) hd1 dt_target[pred_border].max() if pred_border.any() else 0 hd2 dt_pred[target_border].max() if target_border.any() else 0 return max(hd1, hd2)算 HD 前先把预测和真实都二值化到同一类别比如只算心肌。binary_erosion提取边界distance_transform_edt算每个点到最近边界点的距离。两个方向取最大才是对称 HD。如果预测里某一类完全为空pred_border.any()为假直接返回 0 并记录别让它污染平均指标。5. 心脏分割 UNet 避坑清单五条血泪经验5.1 掩码插值用了双线性类别标签被插成小数现象训练 loss 一直不降打印掩码发现里面有 0.3、0.7 这种值。原因预处理时对掩码用了双线性插值类别标签被混合。解决掩码重采样一律用最近邻order0转 tensor 前再确认一遍mask.max()是整数且不超过类别数。5.2 验证 Dice 很高但可视化全是噪点现象验证集 Dice 0.9 以上但把预测图叠回原图看前景碎成一片。原因Dice 在类别极不平衡时会被背景主导前景小目标即使预测很差整体 Dice 依然高。解决按类别分别算 Dice重点看心肌和右心室同时加 IoU 和 HD 交叉验证别只信一个指标。5.3 显存够但训练速度奇慢GPU 利用率上不去现象nvidia-smi 看 GPU 利用率只有 20%训练一个 epoch 要很久。原因DataLoader 的num_workers设太小或者预处理在__getitem__里做了太重的计算CPU 成了瓶颈。解决把预处理结果提前缓存成 npy 文件num_workers调到 CPU 核数的一半左右用pin_memoryTrue加速数据传输。5.4 学习率没调模型卡在局部最优现象训练 loss 降到某个值就不动了验证 Dice 也不涨。原因初始学习率太大导致震荡或者太小导致收敛慢。解决用ReduceLROnPlateau自动降初始 lr 从 1e-3 试起如果前几个 epoch loss 就爆 NaN降到 1e-4如果 loss 下降极慢升到 3e-3 试。5.5 推理时忘了切 eval 模式BatchNorm 用训练统计量现象训练时验证 Dice 正常单独跑推理脚本结果差很多。原因推理时没调model.eval()BatchNorm 还在用当前 batch 的均值和方差单张图推理时统计量偏差大。解决推理前固定加model.eval()和torch.no_grad()这两个是标配别省。6. 把 UNet 心脏分割推到能用的水平两个进阶技巧第一个技巧是深监督。在解码器每一级输出都接一个辅助分类头算辅助损失再加权到总损失。这样梯度能直接传到浅层缓解深层网络梯度消失心脏分割这种边界要求高的任务上Dice 通常能再涨 1 到 2 个点。实现上就是在dec4、dec3、dec2后面各加一个 1×1 卷积输出num_classes通道上采样到原图尺寸算损失权重从 0.1 到 0.3 递减。# 在 UNet forward 里返回多级输出 def forward(self, x): # ... 前面编码解码不变 out_main self.out(d1) aux4 F.interpolate(self.aux4(d4), sizex.shape[2:], modebilinear) aux3 F.interpolate(self.aux3(d3), sizex.shape[2:], modebilinear) return out_main, aux4, aux3 # 损失里加权 loss ce(out_main, masks) 0.3 * ce(aux4, masks) 0.1 * ce(aux3, masks)第二个技巧是测试时增强。推理时把图像水平翻转、旋转 90 度各预测一次再把概率图平均。这个不增加训练成本纯推理端操作心脏结构近似对称翻转增强通常能带来 0.5 到 1 个 Dice 提升。代价是推理时间翻几倍如果做实时应用要权衡。def tta_inference(model, image, device): probs [] # 原图 probs.append(predict_prob(model, image, device)) # 水平翻转 probs.append(np.fliplr(predict_prob(model, np.fliplr(image), device))) # 旋转 90 度 probs.append(np.rot90(predict_prob(model, np.rot90(image, 1), device), -1)) return np.mean(probs, axis0).argmax(axis0)这两个技巧我一般先上深监督因为它改的是训练一次到位TTA 留到模型定型后冲指标用。最后说个我自己的习惯每次改完结构或损失先跑 5 个 epoch 看验证 Dice 趋势趋势不对立刻停别等跑完 100 个 epoch 才发现方向错了。心脏分割这活儿数据质量和预处理占七成模型结构占三成把前面数据管道搭稳UNet 本身不会让你失望。希望帮到你。本文还有配套的精品资源点击获取