医学图像分割实战：基于PyTorch的U-Net完整实现与调优

简介：基于Python与深度学习实现的医学图像分割系统，面向计算机、人工智能、自动化等专业的学生、教师及从业者，适用于毕业设计、课程设计、项目进阶练习。项目以U-Net为分割核心，覆盖数据预处理、模型训练、评估与预测等完整流程，代码经调试可正常运行，模块化设计便于按需修改。压缩包内的Word手册对目录结构和运行方式做了梳理，能帮助降低上手门槛。压缩包共136个文件，包括6个Python源码、120张PNG图像（数据集样本及可视化结果）、6个XML标注文件、1份Word手册和工程配置文件，整体约13.75MB，目录结构清晰。目前已有316人学习下载；作为答辩98分的高分毕设项目，既能帮助新手快速理解分割任务实现思路，也适合在此基础上做功能扩展，完成课程大作业或进一步研究。

1. 医学图像分割：一套“源码+数据集”高分项目拆出的完整链路

医学图像分割在公开数据集上刷高 Dice，拼的往往不是更深的网络，而是“数据预处理 + 损失函数 + 推理后处理”这组配套工程。标题里的“源码+数据集”之所以被反复标记为高分项目，是因为它覆盖了从 Python 环境到模型训练再到评估可视化的完整生命周期。这套系统的本质是：用深度学习网络对每个像素做类别判定，在 CT、MRI 或病理切片上把器官边界、病灶区域与背景精确分开。适合想做医疗影像方向的算法工程师、准备竞赛或毕设的学生，以及需要快速在自采数据上复现分割效果的团队。读完前四章，你应该能拿到一套直接运行的 U-Net 代码；最后一章专门处理从 2D 到 3D 时才会暴露的工程差异。

2. Python与深度学习环境准备：把数据集从下载变成可训练样本

2.1 Python虚拟环境与PyTorch安装的推荐组合

医学图像分割项目对环境依赖比较敏感，nibabel、opencv、torch 之间出现版本冲突时，排查成本远高于装环境的成本。常见做法是先用 conda 建一个独立虚拟环境，再装 PyTorch，不要直接在全局 Python 里混装。我的默认选择是 Python 3.9 或 3.10，对 PyTorch 的算子支持最稳定，遇到旧项目源码也能少碰兼容性问题。

conda create -n medseg python=3.9 -y conda activate medseg pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install numpy opencv-python pillow scikit-learn tqdm matplotlib nibabel

依赖说明：torchvision用于部分数据增强和预训练权重加载，nibabel负责读取 nii.gz 格式的医学影像，opencv-python处理切片的缩放与形态学后处理。安装完成后先用下面这条命令确认 CUDA 是否对当前 PyTorch 可见，True才代表训练能真正用上 GPU。

python -c "import torch; print(torch.cuda.is_available(), torch.__version__)"

2.2 医学图像数据集的读取与预处理要点

公开的医学分割数据集通常有三种形态：ISIC 2018 皮肤镜数据是 PNG 原图加 PNG 掩膜，DRIVE 眼底血管是 TIF 格式，Synapse 多器官 CT 则是 nii.gz 序列。nii.gz 最特殊，它包含体素间距和方向信息，直接当普通图片读会把 z 轴间距忽略掉，导致 3D 造标签时器官形状被拉长。我一般会先用 nibabel 读取并统一方向到 RAS，再按切片提取 2D 训练样本。

import nibabel as nib import numpy as np def load_nii_as_slices(path, axis=-1): img = nib.load(path) data = img.get_fdata() data = np.asarray(data, dtype=np.float32) slices = np.moveaxis(data, axis, 0) return slices # 形状 [num_slices, H, W]

逻辑说明：get_fdata()返回的是浮点体数据，moveaxis把指定轴移到第一个维度，方便循环切片。CT 影像在输入网络前必须做窗宽窗位，常见做法是把取值裁剪到 [-200, 400] Hu 区间再归一化到 [0,1]，这样可以抑制骨骼和空气对软组织分割的干扰。2D 自然影像则只需做全局均值方差归一化。掩膜数据要确认背景是否为 0，前景是否为 1 或 255，某些数据集标注是 255 的 PNG，不转成 0/1 会导致损失函数数值异常。

2.3 数据集划分与目录规范

训练集、验证集、测试集的划分直接在文件层面完成，比在代码里写随机索引要直观得多。我习惯把数据集整理成固定目录结构，后续 DataLoader 只需要拼接路径。

import os, random, shutil random.seed(42) img_dir = "raw/images" mask_dir = "raw/masks" out_train_img, out_val_img = "dataset/train/images", "dataset/val/images" out_train_mask, out_val_mask = "dataset/train/masks", "dataset/val/masks" for d in [out_train_img, out_val_img, out_train_mask, out_val_mask]: os.makedirs(d, exist_ok=True) names = [f for f in os.listdir(img_dir) if f.endswith(".png")] random.shuffle(names) val_n = int(len(names) * 0.15) for i, name in enumerate(names): if i < val_n: os.rename(f"{img_dir}/{name}", f"{out_val_img}/{name}") os.rename(f"{mask_dir}/{name}", f"{out_val_mask}/{name}") else: os.rename(f"{img_dir}/{name}", f"{out_train_img}/{name}") os.rename(f"{mask_dir}/{name}", f"{out_train_mask}/{name}")

参数说明：val_n取总样本的 15% 作为验证集，这个比例在数据量几千张时已经足够稳定估计 Dice。掩膜和原图保持同名，是为了在Dataset里只维护一份文件名列表，避免标签错位。random.seed(42)保证每次执行划分结果一致，对复现实验很重要。

3. 医学图像分割源码拆解：用PyTorch实现U-Net

3.1 U-Net为什么是医学分割的“超经典”基线

医学数据集的标注成本高，样本量通常只有几百到几千张，大规模预训练模型在这种尺度下容易过拟合。U-Net 的结构决定了它对小样本友好：编码器逐步下采样压缩语义信息，解码器逐步恢复空间分辨率，跳跃连接把编码器的高分辨率特征直接拼到解码器对应层，让网络同时保留“在哪里”和“是什么”的信息。医学图像大多是灰度、边界弱、前景占比小的结构，U-Net 这种对称编解码设计天然适合。3D U-Net 只是把卷积核和池化换成三维版本，整体思想完全一致。

3.2 一份可直接运行的最小U-Net源码

下面这个实现去掉了一切花哨模块，只保留 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, num_classes=1): super().__init__() self.enc1 = DoubleConv(in_channels, 64) self.enc2 = DoubleConv(64, 128) self.enc3 = DoubleConv(128, 256) self.enc4 = DoubleConv(256, 512) self.pool = nn.MaxPool2d(2) self.bottleneck = DoubleConv(512, 1024) self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) self.dec4 = DoubleConv(1024, 512) self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.dec3 = DoubleConv(512, 256) self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.dec2 = DoubleConv(256, 128) self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.dec1 = DoubleConv(128, 64) self.out = nn.Conv2d(64, num_classes, kernel_size=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], dim=1)) d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1)) d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) return self.out(d1)

参数说明：in_channels=1对应灰度医学图像，num_classes=1表示二分割网络只输出一张 logit 图，Sigmoid 放到损失函数内部处理。所有卷积都用padding=1保持特征图尺寸，MaxPool2d(2)让特征图降采样为原来一半，转置卷积则恢复到两倍。torch.cat沿通道维度拼接编码器特征，是跳跃连接的具体实现。

3.3 配套数据加载器、Dice损失与评估指标的配套实现

模型只是骨架，训练还需要一个从目录读图的 Dataset 和针对医学分割的损失函数。Dice 损失对前景占比极小的病灶更友好，BCE 则提供稳定的梯度，实际项目中更常用两者组合。

import cv2 from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, size=(256, 256)): self.img_dir = img_dir self.mask_dir = mask_dir self.size = size self.names = [f for f in os.listdir(img_dir) if f.endswith(".png")] def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = cv2.imread(f"{self.img_dir}/{name}", cv2.IMREAD_GRAYSCALE) mask = cv2.imread(f"{self.mask_dir}/{name}", cv2.IMREAD_GRAYSCALE) img = cv2.resize(img, self.size) mask = cv2.resize(mask, self.size, interpolation=cv2.INTER_NEAREST) img = img.astype("float32") / 255.0 mask = (mask.astype("float32") / 255.0).round() return (torch.from_numpy(img).unsqueeze(0), torch.from_numpy(mask).unsqueeze(0)) def dice_loss(pred, target, eps=1.0): pred = torch.sigmoid(pred) inter = 2 * (pred * target).sum() union = pred.sum() + target.sum() return 1.0 - (inter + eps) / (union + eps)

逻辑说明：从目录读图并 resize 到统一尺寸，掩膜用INTER_NEAREST最近邻插值防止类别标签被拉出小数。dice_loss里先过 Sigmoid 把 logit 压到 0~1，再按元素乘计算交集。eps的存在是为了避免前景像素为 0 时除零，在空洞或者小病灶图像上尤其重要。

4. 训练策略与参数调优：epoch、batch_size与过拟合控制

4.1 训练主循环

训练医学分割模型时，一个很常见的坑是只保存最后一个 epoch，而验证集 Dice 往往在训练中段达到峰值。用验证集挑选权重文件，比盲目增加 epoch 更有效。下面这个主循环把训练和验证拆开，每次验证都计算 Dice，只在比历史最优更高时保存模型。

def train_one_epoch(model, loader, opt, criterion): model.train() total_loss = 0.0 for img, mask in loader: img, mask = img.cuda(), mask.cuda() pred = model(img) loss = criterion(pred, mask) opt.zero_grad() loss.backward() opt.step() total_loss += loss.item() return total_loss / len(loader) @torch.no_grad() def validate(model, loader): model.eval() dice_total = 0.0 for img, mask in loader: img, mask = img.cuda(), mask.cuda() pred = torch.sigmoid(model(img)) > 0.5 inter = (pred & mask.bool()).sum().float() union = pred.sum() + mask.sum() dice_total += (2 * inter / (union + 1e-8)).item() return dice_total / len(loader)

参数说明：train_one_epoch中损失函数使用 Dice 与 BCE 加权组合，这个组合针对目标占比极小的分割任务比单一 Dice 更稳。validate使用@torch.no_grad()关闭梯度计算，推理阈值取 0.5。注意pred是布尔张量，mask.bool()要先把目标转成同一类型，否则&运算会报类型错误。

4.2 三个必调参数与推荐区间

参数之间的关系在显存受限时很关键：输入 512×512 加 batch_size 8 大约需要 11GB 显存，如果换到 256×256 加 batch_size 16，显存占用反而更低且吞吐量更高。分割任务的感受野不是越大越好，关键结构较小的时候，256 输入足够覆盖绝大多数 CT 切片。

4.3 过拟合、数据增强与样本不均衡

医学图像数据集普遍只有几百上千张，模型很容易把背景学得特别好、前景学得特别差。对抗过拟合最有效的手段不是换大网络，而是两个组合拳：在线数据增强 + 早停。以下增强方法适用于任何 2D 医学分割源码。

import random, math def augment(img, mask): if random.random() < 0.5: img = cv2.flip(img, 1) mask = cv2.flip(mask, 1) if random.random() < 0.3: angle = random.uniform(-15, 15) M = cv2.getRotationMatrix2D((img.shape[1]//2, img.shape[0]//2), angle, 1.0) img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]), flags=cv2.INTER_LINEAR) mask = cv2.warpAffine(mask, M, (mask.shape[1], mask.shape[0]), flags=cv2.INTER_NEAREST) return img, mask

逻辑说明：翻转与旋转能显著提升分割模型的旋转鲁棒性，但旋转掩膜的插值必须用INTER_NEAREST，否则类的边界会出现“模糊”的非整数标签。另一个有效做法是固定编码器预训练权重，把 BatchNorm 也设置为冻结状态，先只训练解码器，保存解码器收敛后的权重再整体微调。样本不均衡则可以在 DiceLoss 中给前景类加权重，权重数值一般是背景的 2 到 5 倍，具体取决于前景像素占比。

5. 推理、评估可视化：把模型权重应用到单张切片上

5.1 单张切片的推理与后处理

训练时的预处理步骤必须原封不动地应用到推理输入上，否则模型会因数据分布偏移而输出异常。推理不仅要输出概率图，还要把概率图转成干净的二值掩膜。

def predict_mask(model, img_path, size=(256, 256), thr=0.5): img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) img = cv2.resize(img, size).astype("float32") / 255.0 x = torch.from_numpy(img).unsqueeze(0).unsqueeze(0) x = x.cuda() with torch.no_grad(): prob = torch.sigmoid(model(x)).squeeze().cpu().numpy() mask = (prob > thr).astype("uint8") kernel = np.ones((3, 3), np.uint8) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) return mask

参数说明：thr=0.5是最常用的二值化阈值。后处理里的开运算MORPH_OPEN可以消除预测结果中的散点噪声，对血管、病灶这类连通的区域特别有用，但对细小结构要谨慎使用，核大小超过 5 会把细血管腐蚀掉。squeeze()必须同时去掉 batch 维和通道维，得到的数组形状才能和原图对应。

5.2 结果叠加可视化的实现

分割结果单独看掩膜很难判断边界质量，叠加到原图上才能一眼看出过分割和欠分割的位置。

def overlay(img_path, mask, save_path, alpha=0.4): img = cv2.imread(img_path) img_resized = cv2.resize(img, mask.shape[:2][::-1]) color = img_resized.copy() color[mask == 1] = (0, 0, 255) # BGR 红色标注病灶 out = cv2.addWeighted(img_resized, 1 - alpha, color, alpha, 0) cv2.imwrite(save_path, out)

参数说明：alpha=0.4是透明度推荐的起点值，过大会掩盖原始解剖结构，过小则看不清预测区域。直接在拷贝的图上给 mask 区域染色，再用addWeighted融合，比直接修改原图像素值更平滑。建议同时保存掩膜 PNG 和叠加图，前者用于指标计算，后者用于人工目检。

5.3 评估Dice与可视化之间容易被忽略的坑

计算 Dice 时最常见的三个问题：一是输入预测图没有做阈值处理直接与标签比较，导致指标虚低；二是对多类别分割用argmax后忘记排除背景通道 0，把背景也当成类别参与平均；三是对前景极少的病例，空掩膜与全背景预测会得到一个本不该存在的恒高值。正确做法是先检查标签是否真的存在前景，再计算指标。验证集 Dice 稳定在 0.75 以上时，叠加可视化才能看到合理的区域边界，低于 0.6 则说明模型学到的主要是背景信息。

6. 进阶：3D U-Net、滑窗推理与工程化注意点

6.1 从2D切片合成3D预测的两种方案

真实 CT 和 MRI 是三维体数据，直接切 2D 训练会丢失 z 轴上下文。常见做法有两种：第一种是保持 2D 模型，推理时逐层切片预测再堆叠成 3D 掩膜，这种方法处理非标准切面时会出现“阶梯感”。第二种是切换到 3D U-Net，把输入改成[batch, 1, D, H, W]，编码器每层卷积替换成Conv3d，池化替换成MaxPool3d。3D 模型显存开销随 D 线性增长，需要在 z 方向按块切分训练。我的建议是先用 2D 跑通基线，再决定是否换 3D，3D 在器官分割上通常能提升 2~5 个 Dice 点，但在小样本数据集上极易过拟合。

6.2 滑窗推理与内存控制的快捷实现

完整 3D 体数据一般超过显存容量，滑窗推理是标准解法。

def sliding_window_3d(model, volume, window=(64, 256, 256), stride=(64, 192, 192)): D, H, W = volume.shape pred = np.zeros((D, H, W), dtype=np.float32) count = np.zeros((D, H, W), dtype=np.float32) for z in range(0, D - window[0] + 1, stride[0]): for y in range(0, H - window[1] + 1, stride[1]): for x in range(0, W - window[2] + 1, stride[2]): patch = volume[z:z+window[0], y:y+window[1], x:x+window[2]] p = model(torch.from_numpy(patch).unsqueeze(0).unsqueeze(0).cuda()) pred[z:z+window[0], y:y+window[1], x:x+window[2]] += \ torch.sigmoid(p).squeeze().cpu().numpy() count[z:z+window[0], y:y+window[1], x:x+window[2]] += 1 return pred / np.maximum(count, 1e-8)

参数说明：stride小于window时窗口之间有重叠，重叠区域取平均能有效抑制边缘拼接伪影。z 方向 stride 与 window 相等时速度最快但切面间不够平滑，推荐 z 方向重叠 8~16 个切片。count矩阵记录每个体素被覆盖次数，最后做归一化。

6.3 部署阶段必须对齐的输入输出规范

工程化落地时最容易出的问题不是模型精度，而是输入输出格式没对齐。训练时的 resize 尺寸、归一化均值和方差、通道顺序都必须在推理服务中原样复现。导出 ONNX 时固定输入尺寸可以去掉动态维度的额外延展逻辑，二分割模型只需保存一个 logit 输出通道，后处理阈值放在部署端调整即可。整套医学图像分割系统从数据到源码的闭环里，真正决定交付质量的往往是这些看起来不起眼的转换函数和阈值参数。

本文还有配套的精品资源，点击获取