细胞图像分割实战：UNet与UNet++选型及Python实现

简介这份资源面向计算机相关专业正在做毕业设计、课程设计或期末大作业的学生以及需要医学图像分割实战练习的学习者提供基于UNet与UNet两种经典网络对细胞图像进行分割的完整Python实现。压缩包共48个文件以44个py源码为主另含requirements.txt依赖清单、Dockerfile容器配置、readme.md说明文档及.gitignore等辅助文件整体约95KB体量轻便易于本地部署与二次开发。代码按功能拆分为unet模型定义、sahi切片推理、数据加载与dice_score评估、train.py训练入口、predict.py与slicePredict.py预测脚本等模块结构清晰便于对照理解网络结构与训练流程。目前已有181人学习下载适合作为分割任务的入门参考与项目实战模板帮助读者快速跑通训练与推理链路掌握医学细胞图像分割的关键实现思路。1. 细胞图像分割为什么总在边界翻车UNet 与 UNet 的选型起点细胞图像分割是医学图像分析里最基础也最磨人的一环。你拿到的往往是一张相差显微镜或荧光染色图细胞核挤在一起边界模糊背景还有杂质。用阈值分割能跑出个大概但一遇到粘连细胞就崩。语义分割模型里UNet 是绕不开的基线UNet 则是在它基础上针对多尺度融合做的改进。这个标题指向的是一套 Python 源码核心就是用这两个网络把细胞图像里的细胞区域逐像素分出来。适合谁做病理图像分析、细胞计数、药物筛选的算法工程师以及想拿医学图像练手语义分割的学生。它解决的不是“有没有模型”而是“怎么把边界抠准、把粘连分开、把训练跑通”。2. UNet 与 UNet 的结构差异为什么细胞边界更吃多尺度融合2.1 从编码器-解码器到嵌套跳跃连接UNet 的结构很直白左边编码器不断下采样提特征右边解码器不断上采样恢复分辨率中间用跳跃连接把编码器的浅层特征直接拼到解码器对应层。这个设计对医学图像很友好因为细胞边界这种细节信息在浅层特征里保留得最多。但 UNet 的跳跃连接是“同层直连”编码器第一层只跟解码器最后一层拼中间隔了太多下采样浅层细节传到后面已经被稀释了。UNet 改的就是这一点。它在编码器和解码器之间加了一堆嵌套的密集跳跃连接每一层解码器都能拿到前面所有层编码器的特征。具体说UNet 的节点用 $X^{i,j}$ 表示$i$ 是下采样层数$j$ 是沿跳跃路径的卷积层索引。每个节点的输入来自同一层前一个节点和下一层对应节点的上采样结果。这样浅层细节可以走多条路径到达深层解码器边界信息保留得更好。对细胞图像来说这个差异很关键。细胞核边界往往只有几个像素宽UNet 在深层解码时容易把边界平滑掉UNet 因为多路径融合边界响应更强。但代价是参数量和显存占用上去了训练时间也更长。2.2 参数量与显存的实际对比我一般会在选型前先跑一个 batch 的显存测试。以输入 256×256 的 RGB 细胞图为例UNet 在 batch size 8 时显存大约 3.2GBUNet 同样 batch size 会到 4.5GB 左右。如果显卡只有 6GBUNet 可能得把 batch size 降到 4 或者用混合精度。对比项UNetUNet跳跃连接方式同层直连嵌套密集连接参数量约7.8M9.2M256×256 batch8 显存约 3.2GB约 4.5GB边界分割 Dice基线通常高 2-5 个点训练收敛速度较快稍慢这个表不是绝对的具体跟你的数据难度有关。如果细胞边界本身很清晰UNet 就够了如果粘连严重、边界模糊UNet 的收益才明显。2.3 损失函数选型Dice 还是 BCE细胞图像有个典型问题前景像素占比远小于背景。一张 256×256 的图细胞区域可能只占 10%-20%。这时候纯 BCE 会被背景主导模型倾向于全预测背景。常见做法是 Dice Loss 和 BCE 按权重加起来比如loss 0.5 * BCE 0.5 * Dice。Dice 直接优化重叠度对类别不平衡更鲁棒。代码里一般这样写import torch import torch.nn as nn class DiceBCELoss(nn.Module): def __init__(self, weight_bce0.5): super().__init__() self.weight_bce weight_bce self.bce nn.BCEWithLogitsLoss() def forward(self, pred, target): # pred: (B, 1, H, W) logits # target: (B, 1, H, W) 0/1 bce_loss self.bce(pred, target) pred_sigmoid torch.sigmoid(pred) intersection (pred_sigmoid * target).sum() dice_loss 1 - (2. * intersection 1e-6) / (pred_sigmoid.sum() target.sum() 1e-6) return self.weight_bce * bce_loss (1 - self.weight_bce) * dice_loss逻辑说明BCE 提供稳定的梯度Dice 处理不平衡。1e-6是平滑项防止分母为零。weight_bce一般设 0.5如果前景特别少可以降到 0.3让 Dice 主导。参数说明pred是模型原始输出不要先 sigmoidtarget要转成 float 类型。如果多分类把 BCE 换成 CrossEntropyDice 按类别算。3. 用 Python 跑通细胞图像分割的最小闭环3.1 数据准备与目录结构细胞图像数据集常见格式是图像和掩码分开存放文件名一一对应。我一般整理成这样的结构data/ train/ images/ cell_001.png cell_002.png masks/ cell_001.png cell_002.png val/ images/ masks/掩码是单通道二值图前景 255背景 0。如果原始掩码是彩色或者多类先转成灰度再二值化。注意有些数据集掩码是 0/1有些是 0/255读进来后统一除以 255 转成 0/1。import os import cv2 import numpy as np from torch.utils.data import Dataset class CellDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.names sorted(os.listdir(img_dir)) def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img cv2.imread(os.path.join(self.img_dir, name)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) mask (mask 127).astype(np.float32) # 二值化并转 0/1 if self.transform: augmented self.transform(imageimg, maskmask) img augmented[image] mask augmented[mask] img img.astype(np.float32) / 255.0 img np.transpose(img, (2, 0, 1)) # HWC - CHW mask np.expand_dims(mask, axis0) # HW - 1HW return torch.from_numpy(img), torch.from_numpy(mask)逻辑说明cv2.imread读进来是 BGR转 RGB 是为了跟预训练权重对齐。掩码用127二值化避免插值产生的中间值。transform用 albumentations 做同步增强图像和掩码要一起变。参数说明img_dir和mask_dir路径要对齐文件名必须一致。如果数据量小增强可以加随机旋转、翻转、弹性变形弹性变形对细胞形状模拟很好。3.2 模型定义与训练循环UNet 和 UNet 可以直接用现成实现也可以自己搭。自己搭的好处是能改结构。下面是一个简化版 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.conv 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.conv(x) class UNet(nn.Module): def __init__(self, in_ch3, out_ch1): super().__init__() self.down1 DoubleConv(in_ch, 64) self.down2 DoubleConv(64, 128) self.down3 DoubleConv(128, 256) self.pool nn.MaxPool2d(2) self.up1 nn.ConvTranspose2d(256, 128, 2, stride2) self.conv1 DoubleConv(256, 128) self.up2 nn.ConvTranspose2d(128, 64, 2, stride2) self.conv2 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): d1 self.down1(x) d2 self.down2(self.pool(d1)) d3 self.down3(self.pool(d2)) u1 self.up1(d3) u1 torch.cat([u1, d2], dim1) u1 self.conv1(u1) u2 self.up2(u1) u2 torch.cat([u2, d1], dim1) u2 self.conv2(u2) return self.out(u2)逻辑说明编码器每层 DoubleConv 后池化解码器用转置卷积上采样再跟对应编码器特征拼接。最后 1×1 卷积输出单通道 logits。参数说明in_ch根据输入图像通道数改灰度图设 1RGB 设 3。out_ch二分类设 1多分类设类别数。如果显存不够可以把 64 改成 32。训练循环import torch from torch.utils.data import DataLoader from torch.optim import Adam device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_ch3, out_ch1).to(device) criterion DiceBCELoss(weight_bce0.5) optimizer Adam(model.parameters(), lr1e-3) train_loader DataLoader(CellDataset(data/train/images, data/train/masks), batch_size8, shuffleTrue) for epoch in range(50): model.train() epoch_loss 0 for img, mask in train_loader: img, mask img.to(device), mask.to(device) optimizer.zero_grad() pred model(img) loss criterion(pred, mask) loss.backward() optimizer.step() epoch_loss loss.item() print(fEpoch {epoch1}, Loss: {epoch_loss/len(train_loader):.4f})逻辑说明每个 epoch 遍历训练集前向传播算 loss反向传播更新参数。optimizer.zero_grad()清梯度防止累加。参数说明lr1e-3是 Adam 的常用起点如果 loss 震荡可以降到 1e-4。batch_size8根据显存调显存小就降到 4 或 2。epoch 数看收敛情况一般 50-100 够用。3.3 验证与指标计算训练完要在验证集上算 Dice 和 IoU。Dice 衡量重叠度IoU 更严格。def compute_metrics(pred, target, threshold0.5): pred_bin (torch.sigmoid(pred) threshold).float() intersection (pred_bin * target).sum() dice (2. * intersection 1e-6) / (pred_bin.sum() target.sum() 1e-6) iou (intersection 1e-6) / (pred_bin.sum() target.sum() - intersection 1e-6) return dice.item(), iou.item()逻辑说明先 sigmoid 再阈值化得到二值预测然后算交集。1e-6防止除零。参数说明threshold0.5是默认值如果模型偏向保守可以降到 0.4偏向激进就升到 0.6。验证时记得model.eval()和torch.no_grad()。4. 细胞图像分割的避坑与排查清单4.1 掩码值域不统一导致 loss 不降现象训练几个 epoch 后 loss 一直在 0.7 左右不动预测全是背景。原因掩码读进来是 0/255没归一化BCE 的 target 期望 0/1导致梯度方向错了。解决在 Dataset 里统一mask (mask 127).astype(np.float32)确保 target 在 0 到 1 之间。4.2 数据增强把掩码插值出灰边现象训练集 Dice 很高验证集一塌糊涂边界处预测出很多毛刺。原因用了带插值的几何变换如旋转、缩放掩码被双线性插值出 0 到 1 之间的值模型学到模糊边界。解决掩码增强用最近邻插值albumentations 里设interpolationcv2.INTER_NEAREST。或者增强后重新二值化。4.3 显存溢出但 batch size 已经降到 1现象报 CUDA out of memory但 batch size 已经是 1。原因输入图像分辨率太高或者 UNet 的嵌套连接中间特征图没释放。解决把输入 resize 到 256×256 或 512×512用混合精度训练torch.cuda.amp或者把模型通道数减半。4.4 验证集 Dice 波动大现象每个 epoch 验证 Dice 忽高忽低没有稳定上升趋势。原因验证集太小或者 batch size 太小导致 BatchNorm 统计量不稳定。解决增大验证集比例验证时用model.eval()固定 BN或者把 BN 换成 GroupNorm。4.5 粘连细胞分不开现象两个挨着的细胞被预测成一坨边界处没有分开。原因UNet 的跳跃连接不够密浅层边界信息传到深层被稀释。解决换 UNet或者在 loss 里加边界加权用距离变换图给边界像素更高权重。5. 从 UNet 到 UNet 的进阶调优与验证习惯如果你已经跑通了 UNet 基线想进一步压边界误差我一般会做三件事。第一把 UNet 的跳跃连接改成 UNet 的嵌套结构不用全改先在最深层加两条密集连接看验证 Dice 有没有提升。第二在损失函数里加边界权重用 OpenCV 的distanceTransform算每个前景像素到背景的距离距离小的像素权重高这样模型会更关注边界。第三用 Test Time Augmentation验证时把图像翻转、旋转几次预测结果平均通常能涨 1-2 个点。验证习惯上我坚持每个实验只改一个变量。比如这轮只换模型结构下轮只调 loss 权重不然出了问题不知道是哪个改动导致的。另外细胞图像分割的评估不能只看整体 Dice要单独看边界区域的 Dice因为整体 Dice 会被大片背景拉高边界才是真正难的地方。# 边界加权 loss 示例 def boundary_weighted_loss(pred, target, dist_map, weight2.0): bce F.binary_cross_entropy_with_logits(pred, target, reductionnone) # dist_map 越小越靠近边界权重越高 boundary_weight 1.0 weight * (1.0 - dist_map) weighted_bce (bce * boundary_weight).mean() return weighted_bce逻辑说明dist_map是归一化后的距离图边界处接近 0内部接近 1。boundary_weight让边界像素的 loss 贡献更大。参数说明weight2.0是边界权重系数太大可能导致内部像素欠拟合一般 1.5 到 3 之间调。最后说个血泪经验细胞图像分割的玄学在于数据质量比模型结构重要。如果原始图像模糊、染色不均再好的 UNet 也救不回来。我一般会先花时间做图像预处理比如 CLAHE 增强对比度、中值滤波去噪再喂给模型。模型选型上UNet 是保底UNet 是上限但前提是你的数据值得这个上限。希望帮到你。本文还有配套的精品资源点击获取