
简介这是一份基于3D-UNet和TensorFlow实现的人类大脑图像分割算法项目面向医学影像分析、深度学习及三维分割领域的研究者与入门学者可用于脑部病灶定位、解剖结构划分等场景。压缩包内共18个文件以13个Python脚本为主并含2张结果示意图、1份Markdown说明文档等脚本覆盖数据格式转换与预处理、3D-UNet编码器-解码器及注意力模块构建、模型训练、评估指标如Dice、Hausdorff距离计算等完整流程压缩包整体仅201KB轻量易复现。目前已有203人学习浏览。配套详细流程教程从环境配置到推理预测均有讲解便于逐步理解跳跃连接、多尺度特征提取等核心原理源码模块划分清晰研究者可在此基础上调整网络结构、更换损失函数或优化策略进一步提升分割精度。整体而言这是一份兼具教学与实用价值的优质开源实现。1. 为什么大脑分割必须上 3D-UNet一维信息丢失带来的边界塌陷先看一个反直觉的事实很多在 2D 医学图像分割上表现优秀的模型直接迁移到大脑 MRI 分割时会发现皮层沟回区域的边界预测变成一团糊。原因不是模型容量不够而是大脑 MRI 本质上是三维体数据单张切片只包含约 1mm 层厚的信息相邻切片间的空间连续性被 2D 卷积天然丢弃了——这是 2D 分割对 3D 医学影像的老大难问题。3D-UNet 把卷积核从 2D 换成了 3D 形式用滑动窗口直接吃三维体素块让模型能同时感知冠状面、矢状面和横断面上的空间关联。这个项目的价值在于它把从 NIfTI 数据预处理、TFRecord 生成、3D 网络构建到 Dice 评估和 Hausdorff 距离验证的完整流程全部开源适合想用 TensorFlow 做 3D 分割但不想从零搭管线的工程师和研究生。2. NIfTI 到 TFRecord生成 3D patches 与数据增强的预处理管线2.1 为什么不直接喂原始 NIfTI 文件项目里generate_tfrecord.py的存在说明作者选择了先把数据转成 TFRecord 再训练而不是像普通分类任务那样直接在 dataset 里读原始文件。这个选择背后有几个现实约束需要先说清楚。原始 NIfTI 文件.nii/.nii.gz里存的是带 affine 变换矩阵的体素数组尺寸通常在160×192×160到256×256×256之间。如果直接把整个 3D 体积读进内存并做 one-hot 编码单样本就会占用几百 MB 显存batch size 根本提不起来。另一个问题是 NIfTI 的 header 信息如 voxel spacing、slope/intercept在训练时并不都需要每次解析都会浪费 I/O。TFRecord 把这些体积切成小的 patches以序列化字符串形式存储读取时用tf.io.parse_single_example反序列化既减少了磁盘 I/O 次数也方便 TensorFlow 的原生 pipeline 做预取和乱序。2.2 Patch 提取的窗口设计与参数实际项目中patch 大小直接决定模型的感受野和显存占用。常见做法是取64×64×64或128×128×64的 overlap sliding window。考虑到大脑解剖结构的对称性patch 的中心点尽量覆盖灰质、白质和脑脊液交界区域这样模型才不会在推理时只看到单一组织类型。def extract_patches(volume, label, patch_size(64, 64, 64), stride(32, 32, 32)): patches [] labels [] D, H, W volume.shape pd, ph, pw patch_size sd, sh, sw stride for d in range(0, D - pd 1, sd): for h in range(0, H - ph 1, sh): for w in range(0, W - pw 1, sw): vol_patch volume[d:d pd, h:h ph, w:w pw] lbl_patch label[d:d pd, h:h ph, w:w pw] patches.append(vol_patch) labels.append(lbl_patch) return np.stack(patches), np.stack(labels)这段代码是generate_tfrecord.py里最核心的窗口逻辑。stride参数除了控制 patch 数量之外还间接充当数据增强——stride 越小patch 之间的重叠区域越多相同样本能被模型看到的次数越多。在训练阶段常用stride32来扩增数据推理时则改成stride16甚至更小的重叠度配合最后一步的 probability map 平均来消除拼接伪影。注意volume的 shape不同 MRI 扫描仪的图像方向约定不同有的按[D, H, W]存储有的按[H, W, D]写代码前要先用 nibabel 读一个样本确认维度顺序否则后面所有切片操作都会错位。2.3 标准化与类别不平衡的预处理大脑 MRI 的像素强度是相对值不同扫描仪、不同受试者之间的灰度范围差异很大。项目里对 patch 做了 z-score 标准化这是 3D 医学分割任务的标配。预处理顺序不能乱先对全图算均值和方差再切成 patches而不是先切再算统计量。先切会导致位于低灰度区域如脑室附近的 patch 被过度拉伸。def normalize_volume(volume): NIfTI 原始体积标准化返回 float32 类型的 z-score 数组 mean np.mean(volume) std np.std(volume) normalized (volume - mean) / (std 1e-8) return normalized.astype(np.float32)类别不平衡问题在这里比对自然图像分割更严重。脑组织中灰质占比较大而海马体、杏仁核等小结构可能只有几百个 voxel。generate_tfrecord.py里虽然没有集成采样器但在实际复现时我会在 patch 提取阶段做 label-aware 采样计算每个 patch 中非背景类别的 voxel 比例并以 30% 的概率随机跳过全背景 patch。这个策略能显著提升小结构的分割召回率且不会破坏背景区域的空间上下文。2.4 从 patch 到 TFRecord 的写盘细节TFRecord 的每条 record 用tf.train.Example协议格式封装。3D 数组不能直接存需要通过tf.io.serialize_tensor转成字符串再写入bytes_list字段。def write_tfrecord(patches, labels, record_file): with tf.io.TFRecordWriter(record_file) as writer: for vol, lbl in zip(patches, labels): feature { volume: tf.train.Feature( bytes_listtf.train.BytesList( value[tf.io.serialize_tensor(vol).numpy()])), label: tf.train.Feature( bytes_listtf.train.BytesList( value[tf.io.serialize_tensor(lbl).numpy()])), } example tf.train.Example( featurestf.train.Features(featurefeature)) writer.write(example.SerializeToString())input_fn.py里读取时用tf.io.parse_single_example解析再通过tf.io.parse_tensor恢复为张量。需要提醒的是parse_tensor必须显式指定输出 dtype否则解析出的张量 shape 会丢失。这里常见的一个坑是写 TFRecord 时如果忘了.numpy()转换序列化的字符串在 eager 模式下会被当成 EagerTensor 而报TypeError通常加上tf.py_function包装就能绕过但性能会打折扣。项目generate_tfrecord.py里用的是序列化后再写入的方式这也是推荐做法——保持tf.io.serialize_tensor在 eager 上下文外执行读写速度差距在 3D 数据上非常明显。3. 3D-UNet 结构拆解编码器、跳跃连接与注意力门控的实现路径3.1 网络文件的分工network.py / model.py / attention.py项目里network.py和model.py的职责划分值得先捋清楚。network.py倾向于定义通用的网络组件如3D Conv Block、Downsample Block、Upsample Block而model.py负责把编码器、解码器、跳跃连接串成一个可实例化的模型类。attention.py则是可插拔的注意力门控模块。这种分法的好处是可以单独测试某一段网络比如只调试编码器的下采样倍数不影响解码器路径。3.2 编码器路径下采样与多尺度特征3D-UNet 的编码器和 2D UNet 逻辑一致每一层两个 3D 卷积 ReLU group normalization然后接一个 stride2 的 3D 卷积做下采样特征图空间尺寸减半、通道数翻倍。区别在于 kernel size 和 stride 的设置方式。def encoder_block(inputs, filters, name): 单个编码器块两个 3D 卷积 GroupNorm ReLU conv1 tf.keras.layers.Conv3D( filters, kernel_size3, paddingsame, namef{name}_conv1)(inputs) norm1 tf.keras.layers.GroupNormalization( groups8, namef{name}_gn1)(conv1) act1 tf.keras.layers.ReLU(namef{name}_relu1)(norm1) conv2 tf.keras.layers.Conv3D( filters, kernel_size3, paddingsame, namef{name}_conv2)(act1) norm2 tf.keras.layers.GroupNormalization( groups8, namef{name}_gn2)(conv2) act2 tf.keras.layers.ReLU(namef{name}_relu2)(norm2) return act2 def downsample_block(inputs, filters, name): stride2 的 3D 卷积实现下采样同时改变通道数 return tf.keras.layers.Conv3D( filters, kernel_size3, strides2, paddingsame, namef{name}_down)(inputs)这里用了GroupNormalization而不是BatchNorm原因很直接3D 医学图像训练时 batch size 通常只有 2 到 8BatchNorm 的 batch 统计量极不稳定。GroupNorm 按通道分组计算归一化统计量不依赖 batch 大小因此在小 batch 场景下收敛速度反而更快。groups8是常用配置——如果通道数是 16分成 8 组就是每组 2 个通道。在显存允许的情况下把kernel_size保持为 3 不变增加层数比增加单层通道数带来的收益更明显因为下采样次数多了才能捕获更长距离的空间依赖。3.3 解码器路径3D 转置卷积与跳跃连接本文还有配套的精品资源点击获取