基于U-Net的细胞图像分割源码Demo实战：从训练到预测全流程解析

简介这份资源是面向深度学习初学者与图像处理爱好者的入门级实践项目聚焦细胞图像分割这一医疗图像分析中的典型任务帮助零基础读者理解并跑通从数据到模型的完整流程。压缩包共625个文件、约34.32MB其中516张jpg与90张png构成细胞图像数据集为训练与验证提供充足样本8个py源码包含UNet_train.py、UNet_predict.py等核心训练与预测逻辑另有transformer.py及utils工具模块6个xml保存训练参数与状态1个ipynb交互式笔记用于实验展示并附txt说明、LICENSE、gitignore、iml等工程配置文件。目前已有376人学习。读者可借此掌握数据预处理、UNet模型构建、训练评估与结果预测的完整链路理解深度学习框架在细胞分割中的落地方式并基于清晰目录结构快速搭建自己的实验环境适合作为课程设计或入门练手参考。1. 从一堆细胞图到可复现的分割 Demo这套源码到底能跑出什么如果你手头有一批显微镜下拍的细胞图像想用深度学习把细胞区域一个个抠出来又不想从零搭训练框架这套基于 Python 的图像细胞分割源码 Demo 值得先跑一遍再决定要不要改。它包含 626 个文件其中 516 张 jpg、90 张 png 作为训练与验证数据8 个 py 源码文件里UNet_train.py和UNet_predict.py是核心另有一个gan.ipynb交互笔记和transformer.py预处理脚本。说白了这是一个入门级但结构完整的细胞分割工程数据、训练、预测、配置四块都齐了。适合刚接触深度学习图像分割的从业者也适合想快速验证 U-Net 在自己数据上表现的熟手。下面按「它是什么 → 怎么跑 → 坑在哪 → 怎么改」的顺序拆开讲。2. U-Net 为什么是细胞分割的默认起点结构、损失与数据流2.1 细胞分割任务与 U-Net 的匹配逻辑细胞分割本质是像素级二分类每个像素属于细胞或背景。医学图像里细胞边界模糊、样本少、目标密集这三点决定了模型必须能同时利用浅层纹理和深层语义。U-Net 的编码器-解码器加跳跃连接正好干这个编码器逐层下采样提取语义解码器逐层上采样恢复分辨率跳跃连接把编码器的高分辨率特征直接拼到解码器对应层让边界细节不丢。相比 FCNU-Net 在小数据集上收敛更快相比 DeepLab 系列它不需要大规模预训练权重也能出可用结果。这套 Demo 选 U-Net 作为主模型是入门场景下最稳的选择。2.2 数据目录与预处理脚本的对应关系项目里 516 张 jpg 和 90 张 png 不是随便堆的。常见做法是 jpg 放原始显微图像png 放对应的二值掩膜mask因为 png 支持无损和透明通道适合存标注。transformer.py负责把图像和掩膜对齐、归一化、可能还做了尺寸裁剪。跑之前先确认图像和掩膜文件名能一一对应否则训练时读不到标签会直接报错。# transformer.py 常见预处理逻辑示意 import os import cv2 import numpy as np def load_pair(img_path, mask_path, size(256, 256)): img cv2.imread(img_path) # 读原始细胞图 img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # OpenCV 默认 BGR转 RGB mask cv2.imread(mask_path, 0) # 灰度读掩膜0 表示单通道 img cv2.resize(img, size) # 统一尺寸U-Net 要求输入固定 mask cv2.resize(mask, size, interpolationcv2.INTER_NEAREST) mask (mask 127).astype(np.float32) # 二值化阈值 127 可按数据调 img img / 255.0 # 归一化到 0-1 return img, mask逻辑说明cv2.imread第二参数为 0 时读灰度图掩膜必须用最近邻插值INTER_NEAREST否则边缘会产生中间灰度值二值化后边界错位。参数size要和训练脚本里的输入尺寸一致改一处不改另一处是新手最常翻的车。2.3 训练脚本的关键参数与启动方式UNet_train.py是入口。常见结构是 argparse 收参数、DataLoader 加载、定义 U-Net、选损失函数、循环训练。细胞分割里正负样本极不平衡背景远多于细胞所以损失函数一般用 Dice Loss 或 BCEDice 组合而不是纯交叉熵。# 启动训练参数按自己数据改 python UNet_train.py \ --data_dir ./data \ --epochs 50 \ --batch_size 4 \ --lr 1e-4 \ --img_size 256 \ --save_dir ./checkpoints参数说明batch_size受显存限制256×256 输入下 4 到 8 比较稳lr用 1e-4 是 U-Net 的常见起点太大震荡、太小收敛慢epochs50 对几百张图通常够看验证集 Dice 不再涨就可以停。如果脚本里没有 argparse 而是硬编码路径先把路径改成自己的再跑。2.4 预测脚本与结果输出UNet_predict.py加载训练好的权重对单张或整个目录做推理输出分割掩膜。常见做法是输出概率图后取阈值 0.5 二值化再叠加到原图上看效果。# UNet_predict.py 推理核心示意 import torch import cv2 import numpy as np from model import UNet # 假设模型定义在 model.py device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_channels3, out_channels1).to(device) model.load_state_dict(torch.load(./checkpoints/best.pth, map_locationdevice)) model.eval() img cv2.imread(image_03082.jpg) img cv2.resize(img, (256, 256)) / 255.0 tensor torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0).to(device) with torch.no_grad(): pred torch.sigmoid(model(tensor)) # 输出概率 pred (pred 0.5).cpu().numpy()[0, 0] # 阈值二值化 cv2.imwrite(pred_mask.png, (pred * 255).astype(np.uint8))逻辑说明permute(2,0,1)把 HWC 转 CHWunsqueeze(0)加 batch 维度这两步漏一个就会维度报错。map_location保证在无 GPU 机器上也能加载 GPU 训的权重。阈值 0.5 是默认值如果细胞偏小可以降到 0.4 试试。3. 把 Demo 跑起来环境、数据与训练全流程3.1 环境依赖与安装顺序这套源码基于 Python深度学习框架从文件名和结构看大概率是 PyTorchgan.ipynb和 U-Net 组合在 PyTorch 生态里最常见。先装 Python 3.8 以上再按顺序装依赖顺序错了容易出玄学问题。# 建议在虚拟环境里操作避免污染全局 python -m venv cellseg source cellseg/bin/activate # Windows 用 cellseg\Scripts\activate pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python numpy matplotlib pillow pip install jupyter # 如果要跑 gan.ipynb参数说明cu118对应 CUDA 11.8没有 GPU 就装 CPU 版去掉 index-url。opencv-python用于图像读写和预处理numpy是数组运算基础。装完用python -c import torch; print(torch.cuda.is_available())确认 GPU 是否可用。3.2 数据组织与路径校验项目里图像文件是平铺的但训练脚本通常期望data/train/images和data/train/masks这种结构。先按脚本要求整理或者改脚本里的路径。整理完写个校验脚本确认图像和掩膜数量一致、文件名能配对。import os img_dir ./data/train/images mask_dir ./data/train/masks imgs sorted(os.listdir(img_dir)) masks sorted(os.listdir(mask_dir)) print(f图像数: {len(imgs)}, 掩膜数: {len(masks)}) assert len(imgs) len(masks), 图像和掩膜数量不一致检查数据 # 检查文件名主干是否一一对应 for i, m in zip(imgs, masks): assert os.path.splitext(i)[0] os.path.splitext(m)[0], f不匹配: {i} vs {m} print(数据配对校验通过)逻辑说明sorted保证顺序一致os.path.splitext去掉扩展名比对主干。如果掩膜是 png 而图像是 jpg主干相同就能配对。这一步不做训练时读不到标签会报KeyError或直接崩。3.3 训练过程监控与中断恢复训练时重点看两个指标训练 loss 和验证集 Dice。loss 下降但 Dice 不涨说明过拟合或阈值不对loss 震荡多半是学习率太大或 batch 太小。建议每几个 epoch 存一次权重方便中断后恢复。# 训练循环里加 checkpoint 保存 best_dice 0.0 for epoch in range(epochs): model.train() # ... 训练一个 epoch ... model.eval() dice evaluate(model, val_loader) # 自定义评估函数 if dice best_dice: best_dice dice torch.save(model.state_dict(), f{save_dir}/best.pth) print(fEpoch {epoch}: 新最优 Dice{dice:.4f}) torch.save(model.state_dict(), f{save_dir}/last.pth) # 每轮都存防丢参数说明best.pth存验证集最优last.pth存最后一轮。中断后从last.pth恢复把 epoch 起点改成保存时的轮次。Dice 系数是分割任务最直观的指标0.8 以上算可用0.9 以上算好。3.4 用 gan.ipynb 做实验性训练gan.ipynb是交互式笔记可能用于 GAN 做数据增强或分割后处理。Jupyter 的好处是能逐块跑、即时看中间结果。打开后先确认 kernel 选的是刚建的虚拟环境否则 import 会失败。jupyter notebook gan.ipynb在 notebook 里按顺序执行 cell重点看生成器/判别器的 loss 曲线和生成的样本图。如果只是想做分割这个 notebook 可以先跳过不影响主流程。4. 避坑与排查跑这套 Demo 最容易翻车的五个地方4.1 现象训练 loss 一直是 nan原因学习率太大、输入没归一化、或者损失函数里出现了 log(0)。细胞分割用 Dice Loss 时如果预测和标签都为空分母为 0 也会 nan。解决先把学习率降到 1e-5 试确认img / 255.0这步没漏Dice Loss 分母加个极小值1e-6。4.2 现象验证集 Dice 很高但预测图全黑原因类别极度不平衡模型学会全预测背景也能拿到高准确率但 Dice 对这种情况敏感度不够或者阈值设太高。解决换用 Dice BCE 组合损失预测阈值从 0.5 降到 0.3 到 0.4 之间试检查掩膜是不是反了细胞标成 0、背景标成 1。4.3 现象CUDA out of memory原因batch_size 太大、图像尺寸太大、或者没加torch.no_grad()导致预测时还在建计算图。解决batch_size 降到 2 或 1img_size 从 512 降到 256预测代码里确认有with torch.no_grad():。4.4 现象图像和掩膜对不上训练报错原因文件名主干不一致、排序方式不同、或者有的图没有对应掩膜。解决跑 3.2 的校验脚本统一用sorted排序删掉没有配对的孤立文件。4.5 现象预测结果边界锯齿严重原因掩膜 resize 用了双线性插值、或者上采样后没做后处理。解决掩膜 resize 必须用INTER_NEAREST预测后可以用cv2.morphologyEx做开闭运算平滑边界或者训练时加边界加权损失。5. 从能跑到好用改模型、换损失与结果验证的进阶手法跑通 Demo 只是起点真正要用到自己的数据上得会改。第一刀改输入通道如果自己的细胞图是灰度单通道把 U-Net 的in_channels从 3 改成 1同时预处理里去掉COLOR_BGR2RGB。第二刀改损失细胞密集且小的时候Dice 对小块不敏感可以换成 Tversky Loss通过调 alpha/beta 控制查全和查准的权衡。# Tversky Loss适合小目标细胞分割 import torch import torch.nn as nn class TverskyLoss(nn.Module): def __init__(self, alpha0.3, beta0.7): super().__init__() self.alpha alpha # 控制假阳性权重 self.beta beta # 控制假阴性权重beta 大偏向查全 def forward(self, pred, target): pred torch.sigmoid(pred) tp (pred * target).sum() fp (pred * (1 - target)).sum() fn ((1 - pred) * target).sum() return 1 - (tp 1e-6) / (tp self.alpha * fp self.beta * fn 1e-6)参数说明alpha和beta之和通常为 1细胞漏检代价高就把beta调大比如 0.7误检代价高就调大alpha。这个损失在细胞边界模糊、目标小的情况下比纯 Dice 稳。第三刀改评估别只看 Dice加上 IoU 和 Hausdorff 距离。IoU 反映重叠度Hausdorff 反映边界最大偏差后者对分割边缘质量更敏感。验证时把预测掩膜和真实掩膜叠加到原图上肉眼过一遍比只看数字靠谱。# 叠加可视化快速判断分割质量 overlay img.copy() overlay[pred 1] [255, 0, 0] # 预测区域标红 overlay[target 1] [0, 255, 0] # 真实区域标绿 # 红绿重叠处偏黄说明分割准纯红是误检纯绿是漏检 cv2.imwrite(overlay.png, overlay)从那以后我每次改完模型或损失都强制先跑 10 个 epoch 看 overlay 图数字好看但图不对的情况见得太多了。希望这套拆解帮你少走点弯路把这份 Demo 真正用起来。本文还有配套的精品资源点击获取