简介:基于PyTorch与Unet架构的MRI肝脏图像分割毕业设计项目,定位为医学影像方向的高分参考实现,面向需要完成相关课题的本科生、研究生,以及希望动手掌握图像分割技术的初学者。压缩包内含1070个文件,主体为1065张png格式的MRI肝脏图像,包含原始影像与数据增强后的样本,可直接用于模型训练和评估;4个py文件覆盖Unet网络定义、数据读取、训练与预测流程,另有1个md文档说明环境配置与运行步骤,整体大小21.82MB,部署门槛低。代码中配有详细注释,新手可以对照理清网络结构、数据预处理和前向传播逻辑,配合现成数据集能快速跑通实验,直观观察肝脏分割效果。文档部分还提供了项目设计思路与实现说明,便于扩展修改,适合在毕业设计、期末大作业或课程设计中作为技术方案参考。目前已有472人浏览学习。
1. 用PyTorch+Unet做MRI肝脏分割,难的不是模型而是数据
MRI肝脏图像分割并没有一个所谓的“官方最优解法”,但PyTorch加Unet是目前学术论文和毕业设计里出现频率最高的组合,原因很直白:医学影像数据集普遍不大,样本类别也少,Unet这种带跳跃连接的对称编码解码结构,能让深层语义信息和浅层边缘细节同时参与最终预测,恰好适配肝脏和周围软组织边界模糊、背景占比高的场景。而PyTorch的模块化接口让每一步张量形状变化都暴露在明处,训练到一半想改网络里的某一层,不需要重新建图,这对毕设阶段的迭代尤其友好。真正容易翻车的地方往往不在网络结构,而在数据读取、标签取值和损失函数的选择上:标签读出来是255而不是1,训练时类别不平衡导致Loss和Dice一起震荡,测试阶段忘掉训练时的归一化参数,这些问题几乎每届都能遇到。这篇博文按完整闭环来组织,从环境搭建、数据集组织、Unet模型实现、训练与调参,到最后的验证和文档收尾,给出一套可直接复现的源码骨架。
2. 环境与数据准备:PyTorch配置、数据集布局和U-Net网络结构
2.1 conda + pip搭建可复现的PyTorch运行环境
开始写代码之前,先建一个独立的conda环境,避免和日常开发环境互相污染。大多数做毕业设计的机器是单卡NVIDIA GPU,下面这套命令按CUDA 11.8准备:
conda create -n liver-seg python=3.9 -y conda activate liver-seg pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118 pip install monai nibabel opencv-python-headless scikit-image matplotlib tensorboard第一行创建名为liver-seg的虚拟环境,Python版本选3.9是为了兼容后续可能用到的一些医学影像库;第三行的--index-url指到PyTorch官方whl目录,cu118表示CUDA 11.8,这一步务必和显卡驱动匹配。装完以后可以用python -c "import torch;print(torch.cuda.is_available())"验证,输出True再继续。如果机器没有NVIDIA GPU,也可以临时把whl路径换成cpu版,小规模跑通流程后再换GPU机器训练,毕设阶段的调试不需要一开始就追求完整算力。
MONAI在这里是可选项,它的主要价值是提供现成的LoadImaged、ScaleIntensityRange等医学影像transform链;如果不想引入额外依赖,只装nibabel和opencv也完全够用。实际项目中依赖越少越好,因为提交源码时评审老师会在另一台机器上复现,依赖复杂等于给自己埋坑。
2.2 数据集目录与nii.gz标签常见的坑
MRI数据通常以nii或nii.gz格式存储,建议目录结构按“训练/验证/测试”三块分开,标签放在独立路径下:
lits/ ├── train/ │ ├── images/ │ │ ├── liver_001.nii.gz │ │ └── liver_002.nii.gz │ └── labels/ │ ├── liver_001.nii.gz │ └── liver_002.nii.gz ├── val/ │ ├── images/ │ └── labels/ └── test/ └── images/公开肝脏分割数据集里比较常见的是LiTS和CHAOS,拿来做毕设实验没有问题,但要注意遵守各自的使用协议。读取nii.gz用nibabel就能搞定:
import nibabel as nib import numpy as np img = nib.load("lits/train/images/liver_001.nii.gz") volume = img.get_fdata().astype(np.float32) label = nib.load("lits/train/labels/liver_001.nii.gz").get_fdata() print(volume.shape) # (512, 512, slice数) print(np.unique(label)) # 期望看到 [0. 1.],肝肿瘤数据集可能会多一个 2第一步检查np.unique(label)的输出,这是每次换数据集必做的事。有些数据集的标注像素值是0和255,如果不做处理直接作为标签训练,损失函数会一直无法收敛。出现255时用label = (label > 0).astype(np.float32)转成二值标签即可。另外要注意nii读取出来的维度顺序是(H, W, D),和CT图像常见的轴位切片方向一致,但有些公开数据在读取时会带旋转矩阵,打印形状确认一下再进网络。
| 检查项 | 判断标准 | 处理方式 |
|---|---|---|
| 标签取值 | 应包含0和1,背景为0 | 出现255时按阈值重映射 |
| 尺寸一致性 | 同一批样本的H、W、D应一致 | 不一致就统一裁剪或重采样 |
| 背景占比 | 肝脏区域通常占全图比例较低 | 配合Dice损失或ROI裁剪 |
2.3 U-Net网络结构最该看懂的三处设计
Unet的核心是三条线。第一条是编码路径,通过多次Conv - ReLU - MaxPool把分辨率一层层降低,通道数翻倍,网络在这里提炼全局语义;第二条是解码路径,用上采样和卷积把分辨率逐层恢复回来,结合高层语义生成逐像素预测;第三条是跳跃连接,把编码器中同分辨率的特征图在通道维度直接拼接到解码器上,这一步把浅层的边缘、纹理信息补充到深层特征里。正是第三条让Unet在医学任务里表现稳定,因为器官边界往往是低层特征比高层语义更清晰。
网络里的卷积普遍使用kernel_size=3, padding=1,目的是让特征图的高宽不随卷积变化,这样跳跃连接两边才能直接拼接。大部分实现会在每层卷积后接BatchNorm,它对小batch训练时的稳定性帮助很大,尤其是医学图像中样本间灰度差异大的情况。理解了这三处,再看任何Unet变体都不会觉得陌生,注意力Unet和ResUnet改动的基本都是这三个位置的细节。
3. 在PyTorch里实现Unet模型并写和MRI切片匹配的Dataset加载器
3.1 一个带BatchNorm的可复现Unet代码
网络上流行的Unet源码有很多版本,毕设阶段我用得比较多的是下面这套结构,它把双卷积块单独抽出来,后续加注意力机制或者残差连接时只需要改DoubleConv内部:
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=1, out_channels=1): super().__init__() self.inc = DoubleConv(in_channels, 64) self.down1 = nn.MaxPool2d(2) self.conv1 = DoubleConv(64, 128) self.down2 = nn.MaxPool2d(2) self.conv2 = DoubleConv(128, 256) self.down3 = nn.MaxPool2d(2) self.conv3 = DoubleConv(256, 512) self.down4 = nn.MaxPool2d(2) self.conv4 = DoubleConv(512, 1024) self.up1 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) self.conv5 = DoubleConv(1024, 512) self.up2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.conv6 = DoubleConv(512, 256) self.up3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.conv7 = DoubleConv(256, 128) self.up4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.conv8 = DoubleConv(128, 64) self.outc = nn.Conv2d(64, out_channels, kernel_size=1) def forward(self, x): x1 = self.inc(x) x2 = self.conv1(self.down1(x1)) x3 = self.conv2(self.down2(x2)) x4 = self.conv3(self.down3(x3)) x5 = self.conv4(self.down4(x4)) x = self.up1(x5) x = torch.cat([x, x4], dim=1) x = self.conv5(x) x = self.up2(x) x = torch.cat([x, x3], dim=1) x = self.conv6(x) x = self.up3(x) x = torch.cat([x, x2], dim=1) x = self.conv7(x) x = self.up4(x) x = torch.cat([x, x1], dim=1) x = self.conv8(x) return self.outc(x)forward里的cat拼接是整段代码的关键,dim=1指通道维,因为PyTorch默认张量布局是(Batch, Channel, Height, Width)。输入输出都是单通道灰度图,所以在类初始化里设置in_channels=1, out_channels=1;如果未来想迁移到肝脏肿瘤分割,只需要把out_channels改成2,再在损失函数里切换成多分类版本。
3.2 给肝脏切片设计的Dataset类
医学影像训练要先想清楚一个取舍:是把整卷3D数据喂给3D Unet,还是切成2D切片用2D Unet处理。毕业设计最常使用的是2D方案,因为数据加载快、显存占用小、调试直观。数据类里两种做法都要考虑到:
import torch import cv2 import numpy as np import nibabel as nib from glob import glob from torch.utils.data import Dataset import torch.nn.functional as F class LiverSliceDataset(Dataset): def __init__(self, image_dir, label_dir, target_size=(256, 256), use_middle_slice_only=True): self.image_paths = sorted(glob(f"{image_dir}/*.nii.gz")) self.label_paths = sorted(glob(f"{label_dir}/*.nii.gz")) self.target_size = target_size self.use_middle_slice_only = use_middle_slice_only def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img_vol = nib.load(self.image_paths[idx]).get_fdata().astype(np.float32) lab_vol = nib.load(self.label_paths[idx]).get_fdata().astype(np.uint8) # 只分割肝脏,把肿瘤标签并入肝脏区域 lab_vol[lab_vol > 1] = 1 if self.use_middle_slice_only: z = img_vol.shape[2] // 2 img_2d = img_vol[:, :, z] lab_2d = lab_vol[:, :, z] else: # 全部切片采样时会覆盖更多组织,但数据量会大很多 z = np.random.randint(0, img_vol.shape[2]) img_2d = img_vol[:, :, z] lab_2d = lab_vol[:, :, z] # 缩放到统一尺寸,标签用最近邻避免产生新类别 img_2d = cv2.resize(img_2d, self.target_size, interpolation=cv2.INTER_LINEAR) lab_2d = cv2.resize(lab_2d, self.target_size, interpolation=cv2.INTER_NEAREST) # z-score归一化,统计值只来自当前样本 img_2d = (img_2d - img_2d.mean()) / (img_2d.std() + 1e-8) # 转成网络输入格式 img_tensor = torch.FloatTensor(img_2d).unsqueeze(0) lab_tensor = torch.LongTensor(lab_2d) return img_tensor, lab_tensor代码里cv2.resize的标签部分必须用INTER_NEAREST,用线性插值会产生介于0和1之间的小数标签,这就是很多复现项目训练时Loss异常的直接原因。归一化这里按每个切片独立做z-score,比用全数据集的全局均值和方差要稳,因为不同卷的MRI灰度范围差异较大。use_middle_slice_only是为了先跑通流程,确认模型能收敛后再打开全切片采样。
3.3 预处理参数参考表
| 参数 | 建议值 | 理由 |
|---|---|---|
| 输入尺寸 | 256×256 | 兼顾显存和分辨率,512分辨率时显存占用约翻4倍 |
| 归一化 | 单样本z-score | 适配不同扫描仪导致的灰度差异 |
| 标签插值 | INTER_NEAREST | 避免生成非整数标签 |
| 切片策略 | 先中间层后全层 | 先用稳定数据调模型,再扩展数据量提升泛化 |
| 数据增强 | 水平翻转+旋转±15 | 肝脏左右部软组织分布天然对称,翻转不会产生错误先验 |
预处理顺序的先后有讲究:先裁剪或缩放,再做归一化,最后转tensor。如果先归一化再缩放,缩放后边界上的统计值会被二次插值污染,这类问题上不容易看出来但会体现在Dice系数的细微波动上。
4. 训练循环与参数调优:把U-Net在肝脏数据上跑稳
4.1 为什么损失函数要混合使用Dice和交叉熵
肝脏分割里背景像素数量远大于肝脏像素,如果只用交叉熵,模型容易学到“全部预测为背景”这种又懒又稳的解。Dice Loss直接优化预测区域和真实区域的交并比,天然不依赖像素比例,但它在前期的梯度又容易不稳定。常用做法是把两者按权重混合,用一个可复现的实现:
import torch import torch.nn.functional as F def dice_loss(y_prob, y_true): smooth = 1.0 y_prob = y_prob.view(y_prob.size(0), -1) y_true = y_true.view(y_true.size(0), -1) intersection = (y_prob * y_true).sum(dim=1) union = y_prob.sum(dim=1) + y_true.sum(dim=1) dice = (2 * intersection + smooth) / (union + smooth) return 1 - dice.mean() def mixed_loss(logits, y_true): bce = F.binary_cross_entropy_with_logits(logits.view(logits.size(0), -1), y_true.view(y_true.size(0), -1)) prob = torch.sigmoid(logits) dice = dice_loss(prob, y_true) return 0.5 * bce + 0.5 * dice上面代码里smooth放在分子分母一起加,作用是防止当预测和标签都为全零时产生除零错误。binary_cross_entropy_with_logits函数内部已经做了sigmoid计算,所以输入logits即可。权重0.5/0.5是较稳妥的比例,当发现预测区域偏小时可以调成0.3/0.7,偏大时反向调。
4.2 最小可行训练主循环
训练循环按下面的骨架写就足够支撑实验,不需要再堆额外功能:
from torch.utils.data import DataLoader device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = UNet(in_channels=1, out_channels=1).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100) train_data = LiverSliceDataset("lits/train/images", "lits/train/labels") train_loader = DataLoader(train_data, batch_size=8, shuffle=True, num_workers=4) for epoch in range(100): model.train() running_loss = 0.0 for images, labels in train_loader: images = images.to(device) labels = labels.to(device).float() logits = model(images) loss = mixed_loss(logits, labels) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=12.0) optimizer.step() running_loss += loss.item() scheduler.step() if (epoch + 1) % 10 == 0: print(f"Epoch {epoch + 1:3d}, Loss: {running_loss / len(train_loader):.4f}") if (epoch + 1) % 20 == 0: torch.save(model.state_dict(), f"weights/unet_liver_{epoch + 1}.pth")训练循环里labels经过.float()后形状是(Batch, H, W),在mixed_loss内部通过view压成一维向量再和预测比较,避免了维度对不上的问题。梯度裁剪max_norm=12.0这个值在Unet上效果稳定,它防止偶尔的异常损失把参数直接打飞。保存checkpoint时同时记录epoch和loss,避免最后只留了一个过拟合的最终版本。
4.3 训练参数参考与故障排查
| 超参数 | 建议范围 | 备注 |
|---|---|---|
| 学习率 | 1e-4到5e-4 | Adam下超过1e-3极易崩 |
| batch_size | 8到16 | 显存不够时优先降batch,再用梯度累积 |
| epochs | 100到200 | 医学数据量小,过早停会欠拟合 |
| 数据增强 | 翻转+旋转 | 不要用强度太高的增强,肝脏结构变形过度会伤害边缘预测 |
| Checkpoint保存 | 每20轮一次 | 崩了不用重跑 |
训练时最常遇到的三种情况分别是Loss不下降、只预测背景、验证集指标震荡。Loss不下降先去查标签值,确认不是0和255混用;只预测背景是类别不平衡压过了模型,把混合损失的Dice权重调大;验证集震荡多半是学习率过大,切到余弦退火或把初始学习率降到3e-5再观察两轮。
5. 验证与改进:Dice评估、逐切片可视化与文档收尾
5.1 一个能输出Dice和IoU的验证函数
训练结束后要在验证集上跑一次彻底评估,指标只保留Dice和IoU两个就够,肝脏分割任务不需要强上AUC或PA这些跟区域重叠关系不大的指标:
def evaluate(model, val_loader, device): model.eval() dices, ious = [], [] with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) logits = model(images) pred = torch.sigmoid(logits) > 0.5 # 默认threshold=0.5 pred_np = pred.cpu().numpy().astype(np.uint8) label_np = labels.cpu().numpy().astype(np.uint8) for p, g in zip(pred_np, label_np): inter = np.logical_and(p, g).sum() union = np.logical_or(p, g).sum() dice = 2 * inter / (p.sum() + g.sum() + 1e-8) iou = inter / (union + 1e-8) dices.append(dice) ious.append(iou) return np.mean(dices), np.mean(ious)预测阈值0.5是默认值,当肝脏在切片上占比极小且目标区域呈细长形状时,可以把阈值提到0.6左右重新评估一次,最终文档里比较两种阈值的差别会显得工作更完整。
5.2 把预测切片导出成三视图PNG
数值指标之外,还需要视觉结果支撑论文里的定性分析。保存预测切片的代码通常和评估循环放在一起:
import matplotlib.pyplot as plt def save_visualization(model, vol_path, out_path, device, z_slice): model.eval() vol = nib.load(vol_path).get_fdata().astype(np.float32) slice_img = vol[:, :, z_slice] slice_img_norm = (slice_img - slice_img.mean()) / (slice_img.std() + 1e-8) input_tensor = torch.FloatTensor(slice_img_norm).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): logits = model(input_tensor) pred = (torch.sigmoid(logits) > 0.5).squeeze().cpu().numpy() fig, axes = plt.subplots(1, 3, figsize=(12, 4)) axes[0].imshow(slice_img, cmap="gray") axes[0].set_title("Input MRI slice") axes[1].imshow(pred, cmap="gray") axes[1].set_title("Unet prediction") axes[2].imshow(slice_img, cmap="gray") axes[2].imshow(pred, cmap="jet", alpha=0.4) axes[2].set_title("Overlay") for ax in axes: ax.axis("off") plt.savefig(out_path, bbox_inches="tight", dpi=150)输出图建议统一采用“原图、预测二值图、原图叠加半透明预测”三列布局。答辩时老师最容易关注的是预测边界是否贴合真实器官轮廓,而不是单看Dice数字。选可视化切片时要避开那些肝脏很小甚至没有肝脏进入的端点切片,否则整页图几乎没有目标区域,视觉冲击力不够。
5.3 把实验记录收进毕设文档
文档整理是标题里“源码+数据集+文档”三件套的最后一块。不需要长篇大论,但至少要包含三个固定内容:网络结构图和对应的参数表、训练与验证指标曲线、三组以上不同设置的对比结果。对比实验建议是Unet原版、去掉跳跃连接的版本、以及加一个简单注意力模块的版本,这样能证明每个改动都有来源。最后单独加一节“数据预处理细节”,写明归一化方式、裁剪尺寸、标签二值化规则,这些是复现实验的关键信息。最后把最优模型的checkpoint文件和对应epoch、learning rate写在同一份README里,确保任何一个人拿到源码和数据后都能按文档全程跑通。
本文还有配套的精品资源,点击获取