ARTICLE · INTELLIGENCE

战地情报 · 详情页

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

基于Keras与U-Net的医学图像分割实战:从数据预处理到模型训练

基于Keras与U-Net的医学图像分割实战:从数据预处理到模型训练 简介一套基于Keras与Unet的医学图像分割实战资源面向想要入门深度学习图像分割的学生、科研人员或课程设计开发者可用于训练、验证并推断MRI等医学影像的像素级分割结果。压缩包内包含394个文件以383个nii.gz医学影像数据文件为主配套7个Python源码脚本另有预训练模型、环境依赖txt及Markdown说明文档整体大小约990MB数据与代码划分清晰。其中train目录用于模型训练并自动拆分10%作为验证集rest与test目录提供测试图片和课程设计要求预测文件同时源码覆盖数据准备、Unet模型搭建、训练、预测及nii文件可视化等完整流程。已有1032人学习下载适合学习基于深度学习的医学图像分割或直接用于相关实验课设与论文预研。1. 用 KerasUnet 做医学图像分割真正难的不是模型用 KerasUnet 做医学图像分割源码包和数据集通常都是齐的但真正把模型跑通的人往往会卡在三个地方TensorFlow/Keras 版本不一致、数据维度对不上、训练完了看不见有效的分割结果。常见做法是先固定环境版本再把数据读取和预处理独立成管线最后才碰网络结构。下面按这个思路把基于 KerasUnet 的医学图像分割从网络设计、数据准备、训练到验证拆开讲代码可以直接改成自己的路径跑。适合要在私有 CT、MRI 或病理切片数据上快速出分割结果的算法工程师以及准备用 Unet 训练自己数据集的学生。2. U-Net架构与Keras选型源码包里的网络结构图怎么读2.1 医学图像分割为什么默认选Unet医学图像分割的标签往往是小目标、边界模糊、样本量少。Unet 的 U 形结构通过编码器逐级下采样提取语义特征再通过解码器逐级上采样恢复空间分辨率跳跃连接把底层细节和高层语义拼接起来所以即使在小数据集上也能收敛得不错。跟 DeepLab V3 或 Swin-UNet 相比Unet 结构简单显存占用可控在 Keras 里几十行就能搭出来这也是大量源码包默认选 Unet 的原因。实际使用时要区分两类 Unet一类是纯粹下采样四次再上采样四次每层两个卷积加 ReLU另一类把 BatchNorm、Dropout、注意力模块嵌进去。源码包的网络结构图画得再清楚不打印层名也容易看走眼尤其是输入尺寸在哪些层发生变化。建议先把输入固定成(H, W, 1)或(H, W, 3)每一层构建完成后用model.summary()打印形状比读结构图直观得多。2.2 从源码包提取Unet结构的关键点拿到任何 Unet 源码先找三样东西卷积核数量是否按 64/128/256/512 翻倍跳跃连接有没有用Concatenate上采样用的是反卷积还是UpSampling2D。常见 Keras 实现有两种全用Conv2DTranspose或者先UpSampling2D再Conv2D。前者参数多能学上采样但容易出现棋盘效应后者参数少训练更稳边界细节稍弱。医学分割里我一般用UpSampling2D加两个卷积因为过拟合风险低。还要确认所有卷积都用了paddingsame。只要有一层没有特征图尺寸就会在跳跃连接concat时报维度错误这是源码包跑不起来最常见的坑。下面这个代码骨架是从常见项目里改出来的保留name参数排查shape时会方便很多import tensorflow as tf from tensorflow.keras import layers, models def conv_block(x, filters, name): x layers.Conv2D(filters, 3, paddingsame, namename _conv1)(x) x layers.BatchNormalization(namename _bn1)(x) x layers.ReLU(namename _relu1)(x) x layers.Conv2D(filters, 3, paddingsame, namename _conv2)(x) x layers.BatchNormalization(namename _bn2)(x) x layers.ReLU(namename _relu2)(x) return x def unet_encoder(input_shape(256, 256, 1)): inputs layers.Input(shapeinput_shape, nameinput_image) c1 conv_block(inputs, 64, enc1) p1 layers.MaxPool2D(2, namepool1)(c1) c2 conv_block(p1, 128, enc2) p2 layers.MaxPool2D(2, namepool2)(c2) c3 conv_block(p2, 256, enc3) p3 layers.MaxPool2D(2, namepool3)(c3) c4 conv_block(p3, 512, enc4) p4 layers.MaxPool2D(2, namepool4)(c4) c5 conv_block(p4, 1024, bottleneck) return inputs, c1, c2, c3, c4, c5这段代码里编码器每一层做两次 3x3 卷积激活函数放在 BatchNorm 之后是 Unet 最常见的写法。shape(256,256,1)表示灰度医学影像如果输入是彩色或多模态叠加图把最后一个维度改成 3 或更多。调用model.summary()时能看到enc1_relu1这样的层名报错时通过层名定位维度比看数字快捷很多。Unet 实现方式对最终训练效果有直接影响下表列出三种上采样方案的取舍实现方式参数规模上采样质量源码包常见度Conv2DTranspose大可学习易棋盘效应高UpSampling2D Conv2D小平滑细节略弱中双线性插值 Conv2D小不稳定低2.3 Keras与TensorFlow版本选型Keras 3.x 与 TensorFlow 2.x 的 API 差异会影响源码包能不能直接跑。很多老源码用的是from keras.layers import ...而新环境默认 TensorFlow 2.16 之后只有tf.keras才是主流如果项目里还用了keras.backend.clear_session()或旧式merge会直接报AttributeError。建议直接pip install tensorflow2.13.1配合keras2.13.1这个组合在医学分割代码库里最常被验证过。提示运行任何 Keras 医学分割项目前先执行tf.config.list_physical_devices(GPU)确认 GPU 可见再跑一个最小训练循环避免一上来就调网络。如果源码包自带requirements.txt严格按版本装不要顺手把包升级成最新版。医学图像分割依赖的nibabel、SimpleITK、opencv-python对版本也有隐性要求后面数据读取阶段再说明。3. 医学图像数据集预处理把原始切片变成能训练Unet的输入3.1 数据集目录结构与读取方式源码包里的数据集通常分为images和masks两个目录文件名一一对应。少数项目提供 NIfTI 格式的.nii.gz比如 MIMIC 或公开 CT 数据集下载下来的文件这时需要nibabel或SimpleITK读取三维体数据再逐层提取 2D 切片。这里最典型的坑是 NIfTI 中的体数据可能是(D, H, W)或(H, W, D)直接做切片会让 Unet 学到错误的空间关系。一般做法是先统一到(H, W)把空白层筛掉再做 resize 和归一化。图像格式上PNG 和 JPG 要注意位深。医学灰度图经常是 16 位 PNG直接cv2.imread默认按 8 位读会把灰度范围压爆。正确做法是先用cv2.IMREAD_UNCHANGED读成原始位深再手动线性映射到[0,1]。掩码文件一般单通道但有些源码包会用彩色掩码读进来是(H, W, 3)必须二值化转成(H, W, 1)。3.2 归一化、resize与掩码二值化的完整代码下面这段预处理代码是常见“源码数据集”项目里数据管线的核心按路径批量处理import os import cv2 import numpy as np IMG_SIZE 256 num_channels 1 def read_image(path): if isinstance(path, bytes): path path.decode(utf-8) img cv2.imread(path, cv2.IMREAD_UNCHANGED) if img.ndim 3: img cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img cv2.resize(img, (IMG_SIZE, IMG_SIZE), interpolationcv2.INTER_AREA) img img.astype(np.float32) # 线性归一化到 0~1去掉了最大最小拉伸的常数 min_v img.min() max_v img.max() if max_v min_v: img (img - min_v) / (max_v - min_v) if num_channels 1: img img[..., np.newaxis] else: img np.stack([img] * num_channels, axis-1) return img def read_mask(path): if isinstance(path, bytes): path path.decode(utf-8) mask cv2.imread(path, cv2.IMREAD_UNCHANGED) if mask.ndim 3: mask cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY) mask cv2.resize(mask, (IMG_SIZE, IMG_SIZE), interpolationcv2.INTER_NEAREST) mask (mask 127).astype(np.float32) return mask[..., np.newaxis]这段代码里有两个关键参数IMG_SIZE和num_channels。IMG_SIZE必须与 Unet 输入尺寸一致源码包若写死 256这里就不能改。interpolation也值得推敲原图缩小用INTER_AREA抑制混叠掩码缩放用INTER_NEAREST保持边缘不会把 0 和 1 变成灰色。接收路径时先判断是不是bytes是因为tf.numpy_function传入的路径是字节串不 decode 会让cv2.imread读不出文件。3.3 数据增强与训练验证集划分医学图像样本量少数据增强是必须的。常见做法是随机水平翻转、小角度旋转、亮度对比度调节但不要随便做弹性形变医学图像的结构对应关系太强。Keras 的ImageDataGenerator能写但控制变量不如手动用tf.image或 albumentations。源码包里如果是旧式生成器迁移到tf.data时会遇到输出 shape 对不上原因通常是yield的数组没有带 batch 维。建议把增强写成tf.data.Dataset的 map 函数以下是可直接用的版本import tensorflow as tf def aug_func(image, mask): seed tf.random.uniform([2], maxval1000, dtypetf.int32) image tf.image.random_flip_left_right(image, seedseed) mask tf.image.random_flip_left_right(mask, seedseed) image tf.image.random_brightness(image, max_delta0.05) return image, mask def build_dataset(img_paths, mask_paths, batch_size, trainingTrue): ds tf.data.Dataset.from_tensor_slices((img_paths, mask_paths)) ds ds.map( lambda p, m: tf.numpy_function( lambda a, b: (read_image(a), read_mask(b)), (p, m), (tf.float32, tf.float32) ), num_parallel_callstf.data.AUTOTUNE ) if training: ds ds.map(aug_func, num_parallel_callstf.data.AUTOTUNE) ds ds.batch(batch_size).prefetch(tf.data.AUTOTUNE) return ds注意tf.numpy_function会把输入路径变成字节串所以read_image和read_mask里都要先做decode(utf-8)这一步源码包经常漏掉。另外aug_func使用同一个seed翻转 image 和 mask保证配对关系。batch_size根据显存决定12GB 显存跑 256x256 输入建议 816再高容易 OOM。参数推荐值影响IMG_SIZE256 或 512越大保留细节越多显存占用越高batch_size8 / 16过小收敛慢过大直接 OOMnum_channels1灰度图默认 1彩色或多模态改为 3旋转范围10° 以内过大引入非解剖结构形变4. 基于Keras搭建Unet损失函数、训练参数与模型保存4.1 从骨干到完整Unet的拼接第二章只编了编码器完整 Unet 还需要解码器和跳跃连接。拿到c1到c5之后先对c5上采样再与c4concat。因为编码器里每一层都用 same 卷积concat 前不需要裁剪比很多旧版实现简单。完整的解码器代码如下def build_unet(input_shape(256, 256, 1)): inputs layers.Input(shapeinput_shape) c1 conv_block(inputs, 64, enc1) p1 layers.MaxPool2D(2)(c1) c2 conv_block(p1, 128, enc2) p2 layers.MaxPool2D(2)(c2) c3 conv_block(p2, 256, enc3) p3 layers.MaxPool2D(2)(c3) c4 conv_block(p3, 512, enc4) p4 layers.MaxPool2D(2)(c4) c5 conv_block(p4, 1024, bottleneck) d4 layers.UpSampling2D((2, 2))(c5) # 32x32 d4 layers.Concatenate()([d4, c4]) # 跳跃连接 d4 conv_block(d4, 512, dec4) d3 layers.UpSampling2D((2, 2))(d4) # 64x64 d3 layers.Concatenate()([d3, c3]) d3 conv_block(d3, 256, dec3) d2 layers.UpSampling2D((2, 2))(d3) # 128x128 d2 layers.Concatenate()([d2, c2]) d2 conv_block(d2, 128, dec2) d1 layers.UpSampling2D((2, 2))(d2) # 256x256 d1 layers.Concatenate()([d1, c1]) d1 conv_block(d1, 64, dec1) outputs layers.Conv2D(1, 1, activationsigmoid, nameoutput_mask)(d1) model models.Model(inputs, outputs) return model这段代码把 Unet 的跳跃连接和上采样完整串了起来。最后Conv2D(1, 1)输出通道数对应二分类多器官分割把 1 改成num_classes激活函数改成softmax。读源码包时先搜Conv2D(1还是Conv2D(num_classes就能判断项目是二分类还是多标签分割也决定了损失函数怎么写。4.2 损失函数与评估指标的选择医学图像分割中背景像素远多于前景直接binary_crossentropy会让模型偏向预测背景。常见做法是 Dice Loss 或 Dice Loss 与交叉熵组合。Keras 没有内置 Dice Loss需要自己写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) union tf.reduce_sum(y_true_f) tf.reduce_sum(y_pred_f) return 1 - (2.0 * intersection smooth) / (union smooth) def combined_loss(y_true, y_pred): return tf.keras.losses.binary_crossentropy(y_true, y_pred) dice_loss(y_true, y_pred)y_pred已经过 sigmoid所以dice_loss里不需要再激活。smooth防止分母为 0取1e-6或1e-5都行。这里把真实值和预测值展平到一维计算全局 Dice不对每个 batch 单独算因为 batch 小的时候单独算会抖动剧烈。评估指标一般用tf.keras.metrics.IoU再加一个dice_coef就是上面dice_loss去掉1 -的部分。4.3 训练参数、回调与模型保存源码包运行说明里的训练命令通常长这样python train.py --data_dir ./dataset --epochs 100 --batch_size 8 --lr 1e-4 --save_dir ./checkpoints这套参数不是万能的。医学图像训练建议 epoch 至少 50配合早停设 30 轮学习率 1e-4 对 Adam 来说适中。下面是比较稳的回调配置from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau callbacks [ ModelCheckpoint(unet_best.h5, monitorval_loss, save_best_onlyTrue, save_weights_onlyFalse), EarlyStopping(monitorval_loss, patience30, restore_best_weightsTrue), ReduceLROnPlateau(monitorval_loss, factor0.5, patience10, min_lr1e-6), ] model.compile(optimizertf.keras.optimizers.Adam(learning_rate1e-4), losscombined_loss, metrics[dice_coef, tf.keras.metrics.IoU(num_classes2, target_class_ids[1])]) model.fit(train_ds, validation_dataval_ds, epochs100, callbackscallbacks)ModelCheckpoint的save_weights_onlyFalse会连优化器状态一起保存恢复训练时损失不跳变。EarlyStopping的restore_best_weightsTrue让训练结束后模型回滚到验证集最优位置而不是停在最后一轮。ReduceLROnPlateau的factor0.5表示验证损失卡住时学习率减半min_lr1e-6防止学习率过低。参数推荐值说明learning_rate1e-4Adam 下比 1e-3 稳定batch_size8~16由显存决定epochs100~200配合早停实际跑五六十轮Dice threshold0.5预测概率转掩码阈值5. 模型评估、推理与常见坑分割结果怎么验证5.1 推理预测与后处理代码训练好的unet_best.h5不能直接对整图做预测因为数据类型和尺寸可能与训练时不一致。常见做法是先 normal 到IMG_SIZE预测后再 resize 回原尺寸。对单张图片推理的代码model tf.keras.models.load_model( unet_best.h5, custom_objects{combined_loss: combined_loss, dice_coef: dice_coef} ) img read_image(test.png) # (256, 256, 1) pred model.predict(img[np.newaxis, ...])[0] # 添加batch维再去除 pred_mask (pred[..., 0] 0.5).astype(np.uint8) * 255 cv2.imwrite(pred.png, pred_mask)load_model时必须传入自定义的 loss 和 metric否则会报Unknown loss function: combined_loss。pred[..., 0]只取第一个通道对应单类分割。0.5是二分类的常见阈值如果前景占比极小可以调低到 0.3这个阈值属于 Unet 使用时的注意事项里高频出现的一项。后处理最好再加两步去掉小连通域或对掩码做开运算去噪。医学分割里少量孤立误检会影响 Dice 评分用cv2.connectedComponents过滤小于min_area的区域这一步能挽回 23 个百分点的 Dice但源码包往往不写。5.2 常见报错与解决对照下面这张表总结了跑 KerasUnet 医学分割源码包时最常碰到的报错报错信息原因解决方法Allocation of xxx exceeds 10% of free memory显存不足降 batch_size 或 IMG_SIZEDimensions must be equalconv 没有 sameconcat 对不上检查所有 Conv2D 的 paddingUnknown loss function自定义 loss 没传custom_objects 中添加Input 0 of layer conv2d_1 incompatible输入通道数不一致核对 num_channelsdescriptor decode requires a bytes objecttf.numpy_function 路径未解码函数里先 decode(utf-8)第一行是 TensorFlow 的 OOM 变形提示很多人误以为系统内存不够实际上是 CUDA 上下文占满显存。解决优先级是先降 batch_size再降 IMG_SIZE最后才换 GPU。第三行问题在源码包和数据集一起打包时很常见运行说明里只写了load_model没提示要传custom_objects。5.3 验证分割效果除了看Dice还看什么训练完成后只看验证集最后 loss 是不够的需要把预测掩码叠加到原图上检查边界。医学图像分割最怕的不是整体灰度相似而是“预测面积对但边缘偏了一个像素”。几个数值指标一起看Dice 系数、IoU、Hausdorff 距离。Hausdorff 对边界敏感但源码包里一般不写只有 Dice 的话可以在自己推理时补上。更好的做法是只打印验证集平均指标不如把每个样本的 Dice 从低到高排序挑最低的 10 张截图保存。这个操作能暴露出模型对哪些器官、哪种对比度失效如果总是把暗区误判成病灶再去逆向检查输入图像的归一化方式往往是用img / 255.0而不是分位数截断归一化所致。6. 进阶Unet模型改进与验证指标的边界6.1 轻量化和注意力机制改进基础 Unet 在医学分割项目里够用但如果直接上超大尺寸图像显存会先撑不住。常见改进是在瓶颈层后加深度可分离卷积或者把编码器中大通道数的普通卷积换成SeparableConv2D。这类改动对 Keras 来说只需要替换层名但注意参数量下降后收敛速度变慢学习率继续沿用 1e-4 就好。上注意力机制时优先在跳跃连接上加入 AttentionGate比在解码器末端加 SE 模块更贴近 Unet 的落地场景。另一种流行做法是把 Unet 的下采样改为残差连接降低深层梯度消失风险。对源码包做这种修改最容易犯的错是网络结构变了却加载旧权重导致shape mismatch。建议每次结构变化后重新训练至少 20 个 epoch对比验证集 Dice不要沿用历史权重。6.2 批量预测与导出模型推理阶段要把整个验证集跑完并保存概率图不要只在终端打印单张掩码。批量预测用model.predict(ds, verbose1)配合tf.data输出全部结果之后导出成 SavedModel 方便部署tf.keras.models.save_model(model, unet_savedmodel, save_formattf)导出后用saved_model_cli show --dir unet_savedmodel --all检查输入输出签名特别是输出节点的形状避免上线时维度又对不上。这一点在源码包里几乎不会写但却是“项目运行操作说明”与真实生产环境之间最大的鸿沟。6.3 验证改进是否有效的具体做法改进网络后要同时对比三组数值验证集 Dice 均值、Dice 中位数、最低 5% 样本的 Dice 均值。平均值容易被高分样本拉高医学图像分割的数据分布又极不均衡只看均值很容易被那点 0.001 的提升误导。写一段小代码把 Dice 列表算出来取中位数和 P5 值能直观看出改进是否真的救了难样本。对所有调整都记录在train.log里确保事后能复现。Unet 项目的落地能力说到底取决于有没有把改进压到难样本上用同样的 epoch 跑基础模型和改进模型如果 P5 没有提升说明改进只帮助了易分样本这种改进在医学场景里不值得上线。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

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