简介本资源面向计算机、人工智能、通信工程等专业的在校学生与教师以及从事深度学习入门进阶的开发者提供一套基于U-Net网络实现心脏图像分割任务的完整Python源码与训练模型可作为毕业设计、课程设计、作业或项目立项演示的参考方案。压缩包共620个文件约53.4MB其中597个png为训练样本与分割结果可视化图像12个py脚本负责数据加载、模型搭建与训练推理另有h5权重文件、txt说明及md文档便于快速复现与二次开发。资源内代码均经过测试运行成功功能完整已有460人学习下载。读者可从中获取U-Net编解码结构的具体实现、心脏区域分割的训练流程、损失与验证指标记录方式以及mIoU、像素准确率等评估结果并可在现有代码基础上修改以适配其他医学图像分割任务适合作为入门深度学习分割方向的实践素材。1. 心脏分割任务里Unet 为什么还是那个绕不开的基线心脏磁共振影像里把左心室、右心室和心肌逐像素标出来是心功能定量分析绕不过去的一步。射血分数、心室容积、心肌质量这些指标全都建立在分割掩膜准确的前提上。手工勾画一层短轴切片要几分钟一个病人十几层一个队列几百例人力根本扛不住。所以心脏分割一直是医学图像分割里被反复做的任务而 Unet 是这条线上最稳的起点。你拿到的这个「基于 Unet 实现的心脏分割任务 python 源码 模型」本质是一套能跑通「数据读取 → 预处理 → Unet 前向 → 损失回传 → 推理出掩膜」全流程的工程骨架。它解决的不是发论文刷点而是让你在本地把心脏分割这条链路真正跑起来看到预测掩膜叠在原图上是什么样。适合两类人刚接触医学图像分割、想找一个结构清晰能改的代码底子以及做过检测分类、第一次碰像素级任务、需要快速建立手感的人。下面按我实际复现这类工程时的顺序讲参数和坑都给到能直接抄的程度。2. 把 Unet 心脏分割的工程骨架拆开看2.1 Unet 的编码器-解码器结构在心脏分割里到底干了什么Unet 的形状是个 U左边编码器一路下采样特征图尺寸减半、通道翻倍把空间信息压缩成语义信息右边解码器一路上采样把语义信息还原回原图分辨率中间用跳跃连接把编码器同层的特征直接拼到解码器对应层。心脏分割里这个设计特别关键因为左心室和右心室在短轴切片上就是两个相邻的亮环边界只差几个像素纯靠深层语义恢复不出这种细节跳跃连接把浅层的高分辨率边缘信息送回去边界才不至于糊成一团。常见的心脏分割 Unet 输入是单通道灰度切片输出通道数等于类别数。二分类任务前景心肌 背景输出 1 通道配 Sigmoid多分类背景 左心室 右心室 心肌输出 4 通道配 Softmax。编码器每层通常是「两次 3×3 卷积 BN ReLU再 2×2 最大池化」解码器是「上采样 拼接 两次 3×3 卷积」。这个配置不是玄学是无数复现验证过的稳定组合第一次跑不要动它。2.2 数据准备心脏 MRI 切片的读取与归一化心脏分割公开数据常见的是短轴 cine MRI格式多为 NIfTI 或 DICOM。源码工程一般会约定一个数据目录结构图像和掩膜按病例分文件夹。读取时最容易翻车的是灰度范围MRI 没有 CT 那种固定 HU 值不同扫描仪、不同序列的强度分布差异很大直接送进网络会让训练极不稳定。import numpy as np import nibabel as nib def load_slice(img_path, mask_path): # 读取 NIfTI取三维体数据 img nib.load(img_path).get_fdata() # 形状 (H, W, D) mask nib.load(mask_path).get_fdata() # 取中间层心脏短轴通常中间层结构最完整 d img.shape[-1] // 2 img_slice img[:, :, d].astype(np.float32) mask_slice mask[:, :, d].astype(np.uint8) # 按切片做 z-score 归一化而不是按整个体数据 mean, std img_slice.mean(), img_slice.std() 1e-8 img_slice (img_slice - mean) / std return img_slice, mask_slice逻辑说明按切片归一化而不是按整个体数据是因为不同层的亮度差异本身就携带信息全局归一化会把层间对比抹平。1e-8是防止 std 为 0 时除零。掩膜转uint8是为了后面做 one-hot 和计算损失时省内存。参数上如果你用的是多分类标签掩膜里前景类别值通常是 1、2、3背景是 0送进损失函数前要确认这一点标签错位是新手最常见的静默失败。2.3 损失函数与评价指标Dice 和交叉熵怎么配心脏分割里前景像素占比远小于背景纯交叉熵会让网络倾向于全预测背景Dice 系数看着还行但掩膜是空的。常见做法是 Dice Loss 和交叉熵按权重相加Dice 管区域重叠交叉熵管逐像素分类稳定。import torch import torch.nn as nn class DiceBCELoss(nn.Module): def __init__(self, weight0.5): super().__init__() self.weight weight self.bce nn.BCEWithLogitsLoss() def forward(self, logits, targets): # logits: (B,1,H,W) 未过 sigmoidtargets: (B,1,H,W) 0/1 bce self.bce(logits, targets) probs torch.sigmoid(logits) probs probs.view(probs.size(0), -1) targets targets.view(targets.size(0), -1) inter (probs * targets).sum(dim1) dice 1 - (2 * inter 1e-6) / (probs.sum(dim1) targets.sum(dim1) 1e-6) return bce self.weight * dice.mean()逻辑说明BCEWithLogitsLoss内部做了 sigmoid所以前向输出不要提前过激活否则等于做了两次。Dice 里加1e-6是平滑项防止某张图前景为空时除零。weight控制 Dice 占比我一般从 0.5 起如果发现边界还是糊调到 1.0 让 Dice 主导。评价指标单独算 Dice 和 IoU不要直接拿损失值当指标看损失下降不代表 Dice 上升这两个经常不同步。3. 从零跑通训练环境、数据加载和训练循环3.1 环境搭建与依赖版本这类工程对版本不算敏感但 PyTorch 和 CUDA 的匹配要确认。先装 Python再按显卡驱动选对应 CUDA 版本的 PyTorch。CPU 也能跑只是心脏分割的 Unet 在 CPU 上训一轮要很久建议至少有一张 8G 显存的卡。# 创建独立环境避免和系统包打架 conda create -n heart_unet python3.9 -y conda activate heart_unet # 按你的 CUDA 版本选这里以 cu118 为例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install nibabel numpy opencv-python matplotlib tqdm逻辑说明用 conda 建独立环境是血泪经验医学图像工程经常依赖老版本 numpy和系统里其他项目的包冲突后报错很难查。nibabel读 NIfTIopencv-python做 resize 和可视化tqdm看进度。装完先跑一句python -c import torch; print(torch.cuda.is_available())返回 True 再往下走返回 False 说明 CUDA 没配好别急着训。3.2 Dataset 和 DataLoader 的写法自定义 Dataset 要返回图像张量和掩膜张量形状统一成(C, H, W)。心脏切片尺寸不一通常 resize 到 256×256 或 512×512。resize 掩膜必须用最近邻插值用双线性会把标签值插成小数one-hot 直接崩。from torch.utils.data import Dataset, DataLoader import cv2 class HeartDataset(Dataset): def __init__(self, samples, size256): self.samples samples # [(img_path, mask_path), ...] self.size size def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, mask_path self.samples[idx] img, mask load_slice(img_path, mask_path) img cv2.resize(img, (self.size, self.size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.size, self.size), interpolationcv2.INTER_NEAREST) img torch.from_numpy(img).unsqueeze(0).float() # (1,H,W) mask torch.from_numpy(mask).unsqueeze(0).float() # (1,H,W) return img, mask loader DataLoader(HeartDataset(samples), batch_size8, shuffleTrue, num_workers4)逻辑说明图像用INTER_LINEAR掩膜用INTER_NEAREST这是硬规矩。unsqueeze(0)把(H,W)变成(1,H,W)因为 Unet 期望有通道维。num_workers在 Windows 上设 0 更稳Linux 上设 4 到 8。batch_size受显存限制256×256 输入下 8G 卡跑 batch 8 一般没问题爆显存就减半。3.3 训练循环与学习率调度训练循环本身不复杂关键是保存验证集 Dice 最高的那个权重而不是最后一个 epoch 的权重。心脏分割很容易过拟合最后一个 epoch 往往已经退化了。import torch.optim as optim from torch.optim.lr_scheduler import ReduceLROnPlateau device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_ch1, out_ch1).to(device) criterion DiceBCELoss(weight0.5) optimizer optim.Adam(model.parameters(), lr1e-3) scheduler ReduceLROnPlateau(optimizer, modemax, factor0.5, patience5) best_dice 0.0 for epoch in range(100): model.train() for img, mask in loader: img, mask img.to(device), mask.to(device) optimizer.zero_grad() logits model(img) loss criterion(logits, mask) loss.backward() optimizer.step() val_dice evaluate(model, val_loader, device) scheduler.step(val_dice) if val_dice best_dice: best_dice val_dice torch.save(model.state_dict(), best_heart_unet.pth)逻辑说明ReduceLROnPlateau监控验证 Dice连续 5 个 epoch 不涨就把学习率砍半比固定学习率稳。modemax因为 Dice 越大越好。保存state_dict而不是整个模型加载时先实例化同结构再load_state_dict这样换设备不会出问题。学习率 1e-3 是 Adam 的常用起点如果 loss 一开始就震荡降到 1e-4。4. 推理、可视化和模型导出4.1 单张切片推理与掩膜叠加训练完要能直观看到效果把预测掩膜以半透明色叠在原图上边界对不对一眼就看出来。import matplotlib.pyplot as plt def predict_and_show(model, img_tensor, device, threshold0.5): model.eval() with torch.no_grad(): logits model(img_tensor.unsqueeze(0).to(device)) prob torch.sigmoid(logits).squeeze().cpu().numpy() pred (prob threshold).astype(np.uint8) img img_tensor.squeeze().numpy() plt.imshow(img, cmapgray) plt.imshow(pred, cmapjet, alpha0.4) # 半透明叠加 plt.axis(off) plt.show() return pred逻辑说明threshold0.5是默认阈值但心脏分割里这个值可以调。如果发现预测掩膜偏小、边界往里缩把阈值降到 0.3 到 0.4如果掩膜外溢、把周围组织也包进来提到 0.6。这个阈值本质是在查全率和查准率之间挪没有绝对最优看你的下游任务更怕漏还是更怕多。alpha0.4保证底图还能看清。4.2 导出 ONNX 供部署如果要把模型接到别的推理框架导出 ONNX 是常见做法。导出时固定输入尺寸动态轴按需开。dummy torch.randn(1, 1, 256, 256).to(device) torch.onnx.export( model, dummy, heart_unet.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version11 )逻辑说明dynamic_axes把 batch 维设成动态这样部署时 batch 可以变。opset_version11兼容性好别盲目追高版本有些推理引擎对高 opset 支持不全。导出后拿onnxruntime跑一遍和 PyTorch 输出对比误差在 1e-4 以内算正常差太多说明有算子没对齐。5. 心脏分割训练里那些避坑记录5.1 掩膜全黑但 Dice 还有 0.9现象训练几个 epoch 后可视化发现预测掩膜几乎全黑但打印 Dice 接近 0.9。原因心脏切片里背景占比极高如果验证集里混进了没有前景的层Dice 的分母被背景撑大空预测也能拿高分。解决算 Dice 时只统计有前景的切片或者改用前景 IoU 作为主指标别被虚高的 Dice 骗了。5.2 损失变 NaN现象训练到一半 loss 突然变成 nan之后再也降不下来。原因多半是学习率太大导致梯度爆炸或者归一化时 std 为 0 除了零。解决先把学习率降到 1e-4归一化里加1e-8平滑项再开梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。这三个一起上基本能压住。5.3 验证集 Dice 高但换一批数据就崩现象在自己划分的验证集上 Dice 0.9换一家医院的数据掉到 0.6。原因不同扫描仪的强度分布和分辨率差异大模型学到了设备相关的伪特征。解决训练时加强度增强随机调对比度和亮度resize 前先做统一的物理间距重采样。跨中心泛化是心脏分割的老大难别指望一个数据集训完就通用。5.4 显存够但训练特别慢现象显存占用不高但每个 epoch 要跑很久。原因num_workers设成 0数据加载在主进程里串行GPU 一直在等数据。解决Linux 下把num_workers提到 4 以上pin_memoryTrue让数据预取和 GPU 计算重叠。这一步经常能让训练速度翻倍。5.5 加载模型报 key 不匹配现象load_state_dict报一堆 missing keys 和 unexpected keys。原因保存时用了DataParallel包了一层state_dict 的 key 前面多了module.前缀。解决加载前把 key 前缀去掉或者保存时用model.module.state_dict()。这个坑在多卡训练转单卡推理时必踩。6. 把 Unet 心脏分割做扎实的两个进阶习惯第一个习惯是永远留一个「肉眼验证」环节。指标再好看也要定期把预测掩膜叠到原图上翻几张尤其是边界区域和乳头肌附近。我见过太多次 Dice 0.92 但掩膜在心肌薄的地方直接断开的案例指标根本反映不出来。具体做法是每个 epoch 抽 3 张验证图存成 png训练完翻一遍比盯 loss 曲线有用得多。第二个习惯是数据增强别只做旋转平移。心脏短轴切片的解剖结构有强先验左右心室的位置关系基本固定你把它水平翻转等于造了一个解剖上不存在的样本模型学到的可能是错的。安全的增强是弹性形变、小幅旋转±15 度以内、强度扰动和随机裁剪。弹性形变对心肌这种薄壁结构特别有效能让模型对边界形变更鲁棒。再补一个验证方法拿同一病人的收缩末期和舒张末期切片分别推理看分割结果随心动周期的变化是否连续。如果两个时相预测出来的心室面积跳变特别大说明模型对形变不鲁棒这时候回去加弹性形变增强比调网络结构见效快。参数上弹性形变的alpha从 30 起sigma从 5 起太大会把解剖结构扭得不像话。最后说个我自己的教训早期做心脏分割时我总想换更复杂的网络从 Unet 换到各种注意力变体结果发现瓶颈根本不在网络而在数据清洗和归一化。把标签错位的几张图修掉、把归一化改成按切片做之后同一个 Unet 的 Dice 直接涨了 5 个点。所以拿到这套源码先把数据管线走通、把可视化做出来再谈改模型。希望帮到你。本文还有配套的精品资源点击获取