ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Keras U-Net医学图像分割实战:从项目结构到Dice Loss调优

Keras U-Net医学图像分割实战:从项目结构到Dice Loss调优 简介面向医学影像分析初学者与课程设计人群这份基于KerasUnet的医学图像分割项目提供了从数据准备到模型输出的完整代码链路解决Unet分割实验难以快速复现的痛点。压缩包共394个文件以383个nii.gz医学影像数据为核心配合7个Python脚本、1个环境依赖说明及模型权重文件包体约990MB。数据已按训练集、剩余测试集和课程要求预测集分类并预留10%训练数据作为验证集便于直接开展训练与调参。目前已有1032人学习下载。项目按data.py、unet_model.py、train.py、predict.py等脚本串联操作说明覆盖环境安装、逐模块执行顺序及结果保存位置预测阶段生成多次结果目录还提供.nii文件查看脚本方便对比分割效果。对于想深入理解Unet网络结构、数据增强策略与医学影像分割评估方式的读者这套代码加数据包的组合能显著降低上手门槛也可作为毕业设计或医学图像分割竞赛的基础框架。1. KerasU-Net 医学图像分割先搞清楚这个 zip 里到底有什么很多人在网上下载到「基于KerasUnet实现医学图像分割源码数据集项目运行操作说明.zip」这类压缩包时第一反应是解压、找train.py、直接python train.py然后等着看 loss 曲线。但实际上这类项目包的结构远比一个训练脚本复杂它至少包含四块内容U-Net 模型的 Keras 实现、预处理与数据增强工具、医学影像数据集通常是 PNG 格式的切片或原始 DICOM 的导出结果以及一份运行说明文档。多数人卡住的地方不在模型代码本身而在数据路径、标签格式和训练参数这三处。另一个反直觉的事实是Keras 跑 U-Net 做医学图像分割真正决定成果上限的往往不是网络结构而是损失函数的选择和预处理时对标签的处理方式。binary_crossentropy在类别不平衡的医学数据集上经常让模型退化成“全预测背景”而dice_loss这类基于区域重叠的损失函数才是更常见的选择。这篇文章会按“拿到项目包后怎么快速跑通”的顺序来讲先拆解项目结构再做数据预处理然后给出可直接改用的 Keras U-Net 实现最后落在一组实际训练参数和推理可视化技巧上。2. 项目结构解剖Keras 版 U-Net 项目的目录与关键文件识别拿到 zip 后第一件事不是运行而是建立对项目文件的全局认知。医学图像分割项目包的目录组织方式虽然各有差异但大体遵循同一个逻辑数据目录、源码目录、输出目录、文档。先把这个摸清楚后面每一步才不会瞎猜。2.1 典型的目录树与各文件职责常见做法是解压后你会看到类似下面的结构medical_segmentation_keras/ ├── data/ │ ├── images/ # 原始图像通常是 PNG 或 JPG 灰度图 │ ├── masks/ # 标签图二值图或三值图 │ └── train_val_split.py ├── src/ │ ├── model.py # U-Net 网络定义 │ ├── data_generator.py # 数据加载与增强 │ ├── loss.py # 损失函数定义 │ ├── train.py # 训练入口 │ └── predict.py # 推理脚本 ├── models/ # 训练好的 .h5 权重文件 ├── results/ # 预测结果输出 └── README.md # 运行操作说明model.py是整个项目的心脏它定义 U-Net 的编码器-解码器结构data_generator.py负责把磁盘上的图像和掩码读入内存并在每个 epoch 做在线数据增强loss.py里定义损失函数通常包含 Dice 系数或 Tversky 损失的实现。README 里的运行说明往往是按 Linux 环境写的如果你在 Windows 上跑路径分隔符和绝对路径问题会是头两个坑。2.2 识别数据集的真实内容与标注格式医学图像分割数据集最常见的格式是“原图 同名掩码”的配对结构。打开data/masks目录看几张图确认掩码是单通道二值图0 和 255还是多类别标签0、1、2这决定了模型输出层应该用sigmoid还是softmax。部分数据集提供的是三通道的“伪彩色掩码”这种情况下需要把彩色掩码映射回类别索引否则直接读入会在损失计算阶段报错。2.2.1 用脚本快速统计数据分布我一般会写一个简短的脚本来核对数据的基本信息这一步能避免后续八成以上的运行错误import os import cv2 import numpy as np img_dir data/images mask_dir data/masks imgs sorted(os.listdir(img_dir)) masks sorted(os.listdir(mask_dir)) print(f图像数量: {len(imgs)}, 掩码数量: {len(masks)}) # 检查图像尺寸和掩码像素值分布 img cv2.imread(os.path.join(img_dir, imgs[0]), cv2.IMREAD_GRAYSCALE) mask cv2.imread(os.path.join(mask_dir, masks[0]), cv2.IMREAD_GRAYSCALE) print(f图像尺寸: {img.shape}, 掩码尺寸: {mask.shape}) unique_vals np.unique(mask) print(f掩码像素值: {unique_vals}) # 常见输出 [0, 255] 或 [0, 1, 2]这段代码用 OpenCV 读取灰度图避免了三通道读取时cv2.imread默认返回 BGR 三通道的问题。np.unique统计掩码中的像素类别如果输出是[0, 255]训练时需将 255 归一化为 1如果是[0, 1, 2]就是多类别分割问题需要改模型输出通道数。2.2.2 数据配对的常见问题数据文件名不匹配是高频报错点。原图叫case001.png掩码叫case001_mask.png如果加载代码用filename.replace(.png, _mask.png)拼接没问题但如果原图命名带有_img后缀而掩码不匹配就会导致读到空数组。跑训练前用一个断言检查所有文件名是否一一对应比训练到一半报ValueError再回头查要快得多for f in imgs: mask_name f.replace(.png, _mask.png) assert mask_name in masks, f缺少掩码: {mask_name}2.3 运行说明文档的“无效信息”识别项目包里的 README 或操作说明文档质量参差不齐有的写清楚了 Python 版本和依赖列表有的只是复制了标准 U-Net 论文的摘要。你需要做的是提取其中的有效信息Python 版本Keras 2.x 还是 3.x 决定了 API 写法、依赖库版本tensorflow2.4还是tensorflow2.0、训练命令和推理命令。如果文档里写的命令是python train.py --data_path ./data那就照着做如果文档缺失直接看train.py里的argparse参数定义即可这往往比 README 更准确。3. 数据预处理与增强医学图像分割的数据集加载与归一化细节医学图像和自然图像在预处理上有本质差异。自然图像的 ImageNet 均值归一化在医学灰度图上没有太大意义因为 CT、MRI、超声等影像的像素值分布由成像设备决定需要的是截断、归一化到[0,1]区间并保证图像和掩码在空间维度上的对齐。3.1 灰度图的读入、截断与归一化流程CT 影像的原始像素值通常以 Hounsfield UnitHU表示范围在 -1000 到 3000 之间直接归一化到[0,1]会让软组织对比度极低模型学不到有效特征。虽然多数公开数据集已经导出为 PNG 格式像素值被映射到 0~255但保不准你拿到的数据集是 DICOM 格式或未截断的 NIfTI 文件。这个时候预处理就十分关键。对于 PNG 格式的数据常规预处理代码如下def load_and_preprocess(img_path, mask_path, target_size(256, 256)): # 读取灰度图 img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 归一化到 [0, 1] 区间 img img.astype(np.float32) / 255.0 mask (mask 127).astype(np.float32) # 二值化处理 # 统一缩放到目标尺寸插值方式要区分对待 img cv2.resize(img, target_size, interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, target_size, interpolationcv2.INTER_NEAREST) # 扩展通道维度: (H, W) - (H, W, 1) img np.expand_dims(img, axis-1) mask np.expand_dims(mask, axis-1) return img, mask注意这里的两个关键点。第一图像用INTER_LINEAR插值掩码必须用INTER_NEAREST否则掩码边缘会出现灰色过渡值导致损失函数计算时把背景误判为前景。第二如果原数据像素范围已经是 0~255除以 255 即可如果是 0~1 的浮点图则不应再做除法。处理好这两点训练才能正常收敛。3.2 在线数据增强不只是防过拟合医学图像数据集普遍偏小几百张图的规模是常态。数据增强在这类项目中承担的责任比在自然图像任务中更重它不仅增加样本多样性还能缓解标注不足带来的类别不平衡问题。Keras 的ImageDataGenerator虽然方便但在医学分割场景下有几个致命限制它无法自动对掩码做同步变换而且它对灰度图的增强策略不一定符合医学常识。因此多数从业者选择自定义数据生成器。import numpy as np import tensorflow as tf from tensorflow.keras.utils import Sequence class SegmentationDataGenerator(Sequence): def __init__(self, img_paths, mask_paths, batch_size4, dim(256, 256), augmentTrue): self.img_paths img_paths self.mask_paths mask_paths self.batch_size batch_size self.dim dim self.augment augment self.indices np.arange(len(self.img_paths)) def __len__(self): return int(np.ceil(len(self.img_paths) / self.batch_size)) def __getitem__(self, idx): batch_indices self.indices[idx * self.batch_size:(idx 1) * self.batch_size] X, y [], [] for i in batch_indices: img, mask self._load_pair(i) X.append(img) y.append(mask) return np.array(X), np.array(y) def _load_pair(self, idx): img cv2.imread(self.img_paths[idx], cv2.IMREAD_GRAYSCALE) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) img cv2.resize(img, self.dim, interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, self.dim, interpolationcv2.INTER_NEAREST) img np.expand_dims(img.astype(np.float32) / 255.0, axis-1) mask np.expand_dims((mask 127).astype(np.float32), axis-1) if self.augment: # 随机左右翻转图像和掩码同步翻转 if np.random.rand() 0.5: img img[:, ::-1, :] mask mask[:, ::-1, :] # 随机旋转 ±10 度 angle np.random.uniform(-10, 10) M cv2.getRotationMatrix2D((self.dim[0] // 2, self.dim[1] // 2), angle, 1.0) img cv2.warpAffine(img, M, self.dim, flagscv2.INTER_LINEAR) mask cv2.warpAffine(mask, M, self.dim, flagscv2.INTER_NEAREST) return img, mask这个生成器继承tf.keras.utils.Sequence使用多线程预取数据不会在 epoch 之间重新打乱顺序导致训练中断。增强操作只用简单翻转和旋转是因为对于医学图像弹性形变虽然更有效但需要额外处理边界值和掩码的同步插值复杂度较高。如果你的数据集是血管、器官这类变形大的结构可以考虑后续加入imgaug库的ElasticTransformation但在项目跑通之前先用简单的翻转旋转足够验证模型能否正常收敛。3.3 训练集、验证集划分的两种策略医学分割项目里验证集的划分粒度是很多人容易忽略的问题。如果数据来自多个病人不同病人的切片可能高度相似随机按切片划分会导致验证集和训练集之间存在数据泄露评估结果虚高。正确的做法是按病人 ID 划分确保同一个病人的所有切片只在训练集或验证集中出现。# 假设文件名格式为 patient001_slice_001.png patient_ids sorted(set([f.split(_)[0] for f in img_files])) train_patients patient_ids[:int(len(patient_ids) * 0.8)] val_patients patient_ids[int(len(patient_ids) * 0.8):] train_files [f for f in img_files if f.split(_)[0] in train_patients] val_files [f for f in img_files if f.split(_)[0] in val_patients]如果数据集本身就按目录分好了训练和测试那么直接使用即可但依然要检查train/和val/目录是否存在病人重叠。这一步做对了后续评估的 Dice 分数才有参考价值。4. 模型构建与训练参数Keras 实现 U-Net 的关键模块与损失函数调优U-Net 的结构本身并不复杂编码器下采样捕获语义信息解码器上采样恢复空间分辨率跳跃连接把浅层细节传递给深层。但在 Keras 中实现时有几个细节决定了训练能否稳定进行。4.1 用 Keras 函数式 API 搭建 U-Net 的完整代码import tensorflow as tf from tensorflow.keras import layers, models def conv_block(input_tensor, num_filters): x layers.Conv2D(num_filters, 3, paddingsame)(input_tensor) x layers.BatchNormalization()(x) x layers.ReLU()(x) x layers.Conv2D(num_filters, 3, paddingsame)(x) x layers.BatchNormalization()(x) x layers.ReLU()(x) return x def encoder_block(input_tensor, num_filters): x conv_block(input_tensor, num_filters) p layers.MaxPooling2D((2, 2))(x) return x, p def decoder_block(input_tensor, skip_features, num_filters): x layers.Conv2DTranspose(num_filters, (2, 2), strides(2, 2), paddingsame)(input_tensor) x layers.Concatenate()([x, skip_features]) x conv_block(x, num_filters) return x def build_unet(input_shape(256, 256, 1), num_classes1): inputs layers.Input(input_shape) # 编码器 s1, p1 encoder_block(inputs, 64) s2, p2 encoder_block(p1, 128) s3, p3 encoder_block(p2, 256) s4, p4 encoder_block(p3, 512) # 瓶颈层 b1 conv_block(p4, 1024) # 解码器 d1 decoder_block(b1, s4, 512) d2 decoder_block(d1, s3, 256) d3 decoder_block(d2, s2, 128) d4 decoder_block(d3, s1, 64) # 输出层 if num_classes 1: outputs layers.Conv2D(1, 1, activationsigmoid)(d4) else: outputs layers.Conv2D(num_classes, 1, activationsoftmax)(d4) model models.Model(inputs, outputs) return model model build_unet(input_shape(256, 256, 1), num_classes1) model.summary()编码器中每个卷积块包含两个Conv2D BatchNormalization ReLU这是与原始 U-Net 最大的不同点原论文没有 BatchNorm但在现代框架下加入 BatchNorm 对梯度传播和收敛速度有显著帮助。Conv2DTranspose在解码器中负责上采样代码中用Concatenate实现跳跃连接。输出层sigmoid对应二分类softmax对应多分类。参数量约 3100 万在 GTX 1080Ti 级别的显卡上以batch_size4训练无明显压力。4.2 损失函数为什么选 Dice Loss 而不是交叉熵在医学图像分割中前景区域往往只占整张图像的很小比例。如果直接用binary_crossentropy模型会把所有像素预测为背景因为这样损失的期望值反而更低。Dice 系数衡量预测掩码与真实掩码的重叠度天然对类别不平衡不敏感。def dice_loss(y_true, y_pred, smooth1e-6): y_true_f tf.reshape(y_true, [-1]) y_pred_f tf.reshape(y_pred, [-1]) intersection tf.reduce_sum(y_true_f * y_pred_f) return 1 - (2.0 * intersection smooth) / (tf.reduce_sum(y_true_f) tf.reduce_sum(y_pred_f) smooth)代码中把y_true和y_pred展平为一维向量intersection是逐元素相乘后求和相当于预测正确的前景像素数。smooth参数防止分母为零。训练时采用 Dice Loss 和二元交叉熵的组合是推荐做法因为 Dice Loss 的梯度在重叠度极低时不稳定交叉熵能提供更平滑的梯度信号model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-4), losslambda y_true, y_pred: binary_crossentropy(y_true, y_pred) dice_loss(y_true, y_pred), metrics[accuracy, dice_coef] )4.3 训练参数的默认建议值与调整方向下表是训练这类项目的推荐初始参数按常见数据集规模给出建议值参数建议值说明batch_size4超过 8 容易 OOM低于 2 时 BatchNorm 统计不稳定learning_rate1e-4Adam 优化器下从 1e-4 起步20 个 epoch 不降再调小epochs100配合EarlyStopping实际训练约 40~60 轮收敛image_size256×256原始图像过大时可先用该尺寸跑通validation_split0.2按病人 ID 划分而非随机切片划分steps_per_epochceil(num_train/batch_size)使用自定义生成器时必填训练入口的代码可以这样写from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau callbacks [ ModelCheckpoint(models/unet_best.h5, save_best_onlyTrue, monitorval_loss), EarlyStopping(patience15, restore_best_weightsTrue), ReduceLROnPlateau(factor0.5, patience5, min_lr1e-7) ] history model.fit( train_gen, validation_dataval_gen, epochs100, callbackscallbacks, verbose1 )ModelCheckpoint保存验证损失最优的权重EarlyStopping在连续 15 轮无改善时停止训练ReduceLROnPlateau在损失平台期自动降低学习率。这三个回调组合几乎是医学图像分割项目的标配缺一不可。训练时观察val_loss和val_dice_coef两条曲线如果val_loss前期不降反升先停掉训练检查数据预处理和标签是否对齐而不是盲目加大训练轮数。5. 推理与结果可视化预测脚本、阈值分割与指标计算技巧模型训练完成后预测阶段通常有两个目标出结果图用于报告以及计算量化指标用于评估。这两件事可以合并到一个脚本里完成。5.1 单张图像的预测与后处理流程def predict_single_image(model, img_path, threshold0.5): img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) original_shape img.shape img_resized cv2.resize(img, (256, 256), interpolationcv2.INTER_LINEAR) img_norm np.expand_dims(img_resized.astype(np.float32) / 255.0, axis0) img_norm np.expand_dims(img_norm, axis-1) mask_prob model.predict(img_norm)[0, :, :, 0] mask_binary (mask_prob threshold).astype(np.uint8) * 255 # 恢复到原始尺寸 mask_full_size cv2.resize(mask_binary, (original_shape[1], original_shape[0]), interpolationcv2.INTER_NEAREST) return mask_full_size, mask_prob预测时的输入形状要与训练完全一致(1, 256, 256, 1)四维张量中第一位是 batch 维度。需要注意cv2.resize的坐标参数顺序是(宽, 高)也就是(original_shape[1], original_shape[0])写反了会得到旋转 90 度的结果这种错误在可视化阶段很容易被忽略。对于概率图threshold0.5只是一个默认值实际项目里往往要根据验证集上的最优结果调整。可以扫描阈值区间[0.3, 0.4, 0.5, 0.6, 0.7]画一条 Dice 随阈值变化的曲线取最大值对应的阈值作为最终推理阈值。这比固定用 0.5 能带来 1~2 个百分点的 Dice 提升。5.2 Dice 系数与 IoU 的计算方法评估指标的计算要避免用 Keras 训练时出现的metrics[accuracy]像素级准确率在前景占比小的情况下没有参考意义。正确的做法是直接用掩码矩阵计算def calculate_metrics(y_true, y_pred): y_true_f y_true.flatten().astype(np.bool_) y_pred_f y_pred.flatten().astype(np.bool_) intersection np.logical_and(y_true_f, y_pred_f).sum() union np.logical_or(y_true_f, y_pred_f).sum() dice (2.0 * intersection) / (y_true_f.sum() y_pred_f.sum() 1e-6) iou intersection / (union 1e-6) return dice, iouy_true需要从 PNG 掩码中读取并二值化y_pred是模型预测的输出图两者尺寸必须一致。dice和iou的计算都加了1e-6防止除零代码里的np.bool_保证了logical_and按位运算的正确性。5.3 批量预测到结果目录的保存策略最后把全部验证集结果画在一起便于人工检查包括原图、真实掩码、预测掩码三列import matplotlib.pyplot as plt def visualize_results(img_path, mask_path, pred_mask, save_path): fig, axes plt.subplots(1, 3, figsize(15, 5)) img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) axes[0].imshow(img, cmapgray) axes[0].set_title(原图) axes[1].imshow(mask, cmapgray) axes[1].set_title(真实掩码) axes[2].imshow(pred_mask, cmapgray) axes[2].set_title(预测掩码) for ax in axes: ax.axis(off) plt.tight_layout() plt.savefig(save_path, dpi150, bbox_inchestight) plt.close()plt.close()一定要加否则循环生成大量图像时内存会持续增长最终导致RuntimeError: main thread is not in main loop。可视化结果不仅用于汇报更是检验模型错误模式的直接手段如果预测掩码出现大量孤立小点说明后处理缺少形态学开运算如果边界出现连续的锯齿状误差就要考虑增强策略或上采样方式的问题。这一步做到位基于 Keras 与 U-Net 的医学图像分割项目从数据到指标的全链路才算真正跑通。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进