简介面向医学图像处理与CT结节分割任务的Pytorch 3D图像分割工程以Luna16公开数据集为案例完整演示了UNet3d与VNet3d两种CNN结构实现。压缩包共92个文件以49个Python源码脚本为核心覆盖从数据重采样、掩码与bbox标注生成到模型训练、验证、推理、评估、可视化和后处理的全流程另有16个pyc编译文件、7个npy数据文件、5张PNG曲线或结果图、3个NII医学图像样本及CSV标注文件等整体约61.66MB。代码目录按数据预处理、数据集加载、模型定义、推理、后处理和结果展示等模块清晰划分并提供多个训练主脚本适应多类别、patch采样等不同任务变体可帮助读者直接对照代码理解3D分割的完整落地思路。资源已有542人学习适合具备一定Pytorch基础、希望上手医学影像3D分割项目的开发者参考。1. 3D 图像分割没做对数据准备后面全是白忙先把这句话放前面很多人一上来就搜 Pytorch 3D 图像分割的模型代码把 nnUNet、VNet、3D U-Net 的权重加载跑通就觉得自己会了。但真正落到自己的数据集上最先卡住的根本不是网络结构而是数据准备医学影像的 NIfTI 格式怎么读、体素间距不一致要不要重采样、标签和原图空间位置对不对得上、显存不够怎么切 patch、归一化按什么统计量算。这些事没理清楚模型训练得再久出来的分割结果也是一团噪声。这篇笔记就围绕「数据准备过程和代码思路」展开把从原始影像到 Pytorch Dataset 的整条链路拆开讲适合正在做 CT/MRI 分割、想把 nnUNet 之外的自定义流程跑通的工程师。我会按自己常用的方案给出可复现代码并把参数和边界说明白便于照着改。2. 3D 分割的数据长什么样NIfTI、体素间距和标签的坐标系2.1 为什么 3D 图像分割绕不开 NIfTI 和 DICOM 的转换3D 医学图像分割最常见的原始格式是 DICOM 和 NIfTI。DICOM 是医院影像设备直接输出的格式一个 3D 体积往往拆成几百张二维切片还带着大量患者信息和扫描参数NIfTI 是科研和开源数据集里更常用的单文件格式一个 .nii.gz 就能装下整个三维体积并且自带仿射矩阵affine能把体素坐标映射到物理空间。做 Pytorch 训练时我一般不会直接拿 DICOM 喂给模型而是先用 dcm2niix 这类工具把 DICOM 转成 NIfTI理由很简单NIfTI 读取快、文件数少、标签配准容易而且大多数开源分割模型和数据增强库都以 NIfTI 为默认输入。转换时要注意一个常见认知DICOM 里的像素值可能是无符号整数且带有 rescale slope 和 rescale intercept转成 NIfTI 后需要确认是否保留了原始 CT 值。用 dcm2niix 转换时默认会应用 rescale但不同厂商的影像可能会有例外。拿到 NIfTI 后第一件事不是急着看像素而是打印 shape、spacing、orientation 和 affine确认数据没有翻转或方向错乱。2.2 读取 NIfTI 的代码思路用 nibabel 还是 SimpleITK在 Pytorch 的数据准备阶段我常用 nibabel 或 SimpleITK 读取 NIfTI。nibabel 更轻量适合快速读取和数组操作SimpleITK 内置了重采样、方向调整等函数在预处理阶段更省事。两者都行但代码思路有区别nibabel 读出来的是 nib.Nifti1Image需要通过 get_fdata() 拿到 numpy 数组SimpleITK 读出来的是 sitk.Image可以直接获取 spacing 和 direction。下面是我推荐的读取和检查代码import nibabel as nib import numpy as np def load_nifti(path): img nib.load(path) data img.get_fdata() # 得到 numpy 数组注意可能返回 float64 affine img.affine # 4x4 仿射矩阵记录体素到物理空间的映射 header img.header spacing header.get_zooms() # 体素间距比如 (0.7, 0.7, 1.0) print(shape:, data.shape, spacing:, spacing) print(affine:\n, affine) return data, affine, header这段代码的逻辑是先通过 nibabel 加载 NIfTI 文件再用 get_fdata() 获取数值数组。这里有个坑——get_fdata() 默认返回 float64如果不注意会把 CPU 内存撑爆尤其是 512x512x500 这种体积float64 比原始 int16 多占 4 倍内存。所以拿到数据后尽快转成 float32 或直接按需归一化不要一直抱着 float64 不放。spacing 的获取有两种方式header.get_zooms() 返回每个维度的体素间距img.affine 也可以推算出间距但直接读取 zooms 更直观。2.3 体素间距不一致为什么不直接训练重采样与 z-score 归一化不同扫描设备、不同医院的图像体素间距通常不一样。有的 CT 是 0.7x0.7x1.0有的是 1.0x1.0x5.0如果直接丢进 3D 模型训练模型会把间距差异当成解剖结构差异分割精度会明显下降。常用的办法是把所有数据重采样到统一的 target_spacing比如 (1.0, 1.0, 1.0)这样模型看到的每个体素都对应相同的物理尺寸。重采样的代码思路如下使用 SimpleITK 会比较简洁import SimpleITK as sitk def resample_to_spacing(sitk_img, target_spacing(1.0, 1.0, 1.0), is_labelFalse): original_spacing sitk_img.GetSpacing() original_size sitk_img.GetSize() target_size [ int(round(original_size[i] * original_spacing[i] / target_spacing[i])) for i in range(3) ] resampler sitk.ResampleImageFilter() resampler.SetOutputSpacing(target_spacing) resampler.SetSize(target_size) resampler.SetOutputDirection(sitk_img.GetDirection()) resampler.SetOutputOrigin(sitk_img.GetOrigin()) if is_label: resampler.SetInterpolator(sitk.sitkNearestNeighbor) # 标签用最近邻避免引入新类别 else: resampler.SetInterpolator(sitk.sitkLinear) # 图像用线性插值 return resampler.Execute(sitk_img)这里的参数说明target_spacing 越小重采样后的体素数越多训练显存消耗越大标签必须用最近邻插值否则 0/1/2 这些标签值会被插值成 0.4、1.7 之类的中间值模型就懵了。图像数据用线性插值就够不要用三次样条虽然更平滑但会改变原始强度分布。重采样后最好再检查一次 shape 和 spacing因为取整操作会导致实际 size 和理论值有偏差。z-score 归一化也是 3D 分割数据准备的核心步骤。常见做法是先计算整个训练集的均值和标准差再用 (data - mean) / std 做归一化。但医学图像里不同扫描协议会产生不同强度范围直接对每个样本分别做 z-score 往往更稳妥尤其是 CT 图像还可以先做窗宽窗位截断再归一化。我一般会按器官和任务决定如果是 MRI 的 T1/T2直接全图 z-score如果是 CT先用 [-200, 400] 这类窗位裁剪肝脏或肾脏区域再计算该区域内的均值方差做归一化效果比全图归一化好很多。3. Pytorch Dataset 的完整代码思路从 NIfTI 文件到训练张量3.1 自定义 Dataset 类的基本框架读取、缓存、返回字典Pytorch 训练时数据加载的核心是 torch.utils.data.Dataset。对 3D 图像分割任务我通常会让 Dataset 返回一个字典包含 image、label、filename 等信息。文件名建议返回这样推理时能回溯到原始样本也方便定位坏样本。下面是一个实用的基础框架import torch from torch.utils.data import Dataset import numpy as np import nibabel as nib class SegmentationDataset3D(Dataset): def __init__(self, case_list, transformNone): self.case_list case_list # 每个元素是 dict含 image_path 和 label_path self.transform transform def __len__(self): return len(self.case_list) def __getitem__(self, idx): case self.case_list[idx] image nib.load(case[image_path]).get_fdata().astype(np.float32) label nib.load(case[label_path]).get_fdata().astype(np.int16) # 保证 image 和 label 的 shape 一致 assert image.shape label.shape, fshape mismatch: {image.shape} vs {label.shape} # 统一转成 (C, D, H, W) 的格式C1 image np.expand_dims(image, axis0) label np.expand_dims(label, axis0) sample { image: image, label: label, image_path: case[image_path], label_path: case[label_path], } if self.transform: sample self.transform(sample) return sample这个类直接把 image 和 label 读进内存在小数据集上是够用的。注意 label 用 int16 而不是 int64因为 3D 标签体积很大int64 会浪费内存。如果数据集太大就需要在init里做预处理把重采样和归一化的结果缓存成 .npy 文件避免每次 epoch 都重新读原始 NIfTI。3.2 训练集验证集怎么划分按病人划分而不是按图像划分3D 图像分割数据准备最容易犯的错误是把同一个病人的多次扫描或同一病例的多个序列同时分进训练集和验证集。这会导致验证集的评估结果虚高因为模型已经见过同一个人的解剖结构。正确的做法是把数据按病人 ID 分组再按组划分。比如前列腺 MRI 中一个病人可能有 T2 和 ADC 两个序列这两个序列必须放在同一个 fold 里。代码思路如下import random from collections import defaultdict def split_by_patient(case_list, patient_keypatient_id, val_ratio0.2, seed42): patient_dict defaultdict(list) for i, case in enumerate(case_list): patient_dict[case[patient_key]].append(i) patients list(patient_dict.keys()) random.seed(seed) random.shuffle(patients) val_count max(1, int(len(patients) * val_ratio)) val_patients set(patients[:val_count]) train_idx [] val_idx [] for i, case in enumerate(case_list): if case[patient_key] in val_patients: val_idx.append(i) else: train_idx.append(i) return train_idx, val_idx这里的逻辑是先按 patient_id 把样本索引分组再对病人列表做随机打乱按比例取验证集病人。这样同一个病人的所有序列只会出现在一边。实际项目中我还会进一步做分层采样保证训练集和验证集中不同类别器官的体积分布大致接近但至少按病人划分是最基本的底线。3.3 显存不够怎么办patch 切块与随机采样的实现3D 图像通常无法整图放进 GPU最常见的是随机采样固定大小的 patch比如 128x128x64。训练时从图像和标签的对应位置裁出相同区域验证时用滑窗把整张图切成多个 patch 分别推理最后再拼回去。patch 大小的选择取决于显存和模型结构3D U-Net 的编码器下采样 4 次patch 每个维度至少要被 16 整除否则特征图尺寸对不上。下面是一个随机裁剪的代码片段import random def random_crop_3d(image, label, crop_size(128, 128, 64)): d, h, w image.shape[1:] # 注意这里 image 是 (C, D, H, W) cd, ch, cw crop_size if d cd or h ch or w cw: raise ValueError(fimage size {image.shape} smaller than crop size {crop_size}) start_d random.randint(0, d - cd) start_h random.randint(0, h - ch) start_w random.randint(0, w - cw) image_crop image[:, start_d:start_dcd, start_h:start_hch, start_w:start_wcw] label_crop label[:, start_d:start_dcd, start_h:start_hch, start_w:start_wcw] return image_crop, label_crop随机裁剪的问题是可能裁到大量背景区域导致类别极度不均衡。改进方案是让裁剪中心偏向前景体素做法是先计算 label 中前景的连通域或所有前景坐标再随机选一个前景体素作为裁剪中心如果越界就裁掉边界。这个简单策略能显著提高训练效率。另一个常见做法是在每个 epoch 前生成一组裁剪位置保证每个位置在整个训练过程中至少被覆盖到一次避免某些区域永远看不到。4. 数据增强在 3D 分割里怎么做别把空间变换和强度变换搞混4.1 空间增强旋转、翻转、弹性形变的参数边界3D 数据增强比 2D 谨慎得多因为医学图像有明确的解剖方向。轴向翻转上下翻转对大多数器官是可以接受的但左右翻转要分任务肝脏、肾脏左右对称性没有强到可以随便翻脑部分割也要考虑左右半球的镜像差异。旋转最好限制在小角度范围比如 ±10 度到 ±20 度之间超过 30 度会破坏解剖结构。弹性形变常用来模拟器官受呼吸或肌肉牵拉产生的形变但幅度要保守否则会让标签边界失真。如果用 torchio 库做增强代码思路如下import torchio as tio transform tio.Compose([ tio.RandomFlip(axes(0, 1), flip_probability0.5), tio.RandomAffine( scales(0.9, 1.1), degrees10, translation5, p0.5 ), tio.RandomElasticDeformation( num_control_points7, max_displacement5, p0.3 ), ])这里的参数说明RandomAffine 中的 scales 是各向同性的缩放范围0.9~1.1 对应体积缩小/放大 10%degrees10 表示绕每个轴最多旋转 10 度translation5 是最大平移体素数。torchio 对标签和图像会使用相同的变换并自动用最近邻处理标签所以不需要额外写对齐逻辑。注意弹性形变的 max_displacement 单位受 spacing 影响在预处理阶段统一重采样到相同 spacing 后这个参数才有稳定的物理意义。4.2 强度增强对比度、噪声、Gamma 变换对 MRI 和 CT 的区别MRI 图像的强度分布易受场不均匀性影响同一组织在不同扫描中的信号强度差异较大所以强度增强对 MRI 分割很有帮助。CT 图像则有明确的物理单位Hounsfield UnitHU一般不建议做大幅度强度变换最多做加性高斯噪声或小范围的 Gamma 变换否则会破坏 CT 值的一致性。一个实用的强度增强方案如下class RandomIntensityAugment: def __call__(self, sample): image sample[image] # 以 0.5 概率对图像做对比度变换 if np.random.rand() 0.5: gamma np.random.uniform(0.8, 1.2) image np.power(np.clip(image, 0, 1), gamma) # 以 0.3 概率添加高斯噪声 if np.random.rand() 0.3: noise_std np.random.uniform(0.01, 0.03) image image np.random.normal(0, noise_std, sizeimage.shape) sample[image] image.astype(np.float32) return sample这里的逻辑是假设图像已经做了 min-max 或 z-score 归一化。如果 z-score 后数据均值接近 0 方差接近 1Gamma 变换不能直接套在负值上所以上面的代码里先 clip 到了 0~1。实际操作中我通常会对 z-score 后的数据跳过 Gamma 变换只做加性噪声和轻微模糊。MRI 多模态T1、T2、ADC场景下各模态可以共享同一组 Gamma 参数也可以独立采样但共享参数能保持不同模态间的对比度关系。4.3 增强顺序有讲究先重采样再增强还是先增强再重采样数据预处理和增强的顺序是个容易被忽略的细节。我的固定顺序是读 NIfTI → 重采样到统一 spacing → 裁剪或全图归一化 → 空间增强 → 强度增强 → 随机裁剪 patch。原因很简单重采样会改变图像和标签的空间对应关系如果先做旋转再重采样插值会引入额外误差而空间增强最好在 patch 裁剪之前做这样增强后的边界不会因为裁剪而丢失。如果先裁 patch 再做旋转旋转后的图像角落会出现空值虽然可以 fill 成 0但会在 label 边界引入伪影。实际项目里我见过不少人把增强放在 Dataset 外单独写一个函数每次调用时顺序不一致导致同一套代码在本地和服务器上效果不一样。建议把所有增强封装成一个 tio.Compose 或自定义 transform 类固定顺序统一入口。5. 避坑合集3D 分割数据准备的 5 个典型翻车现场5.1 标签和原图错位方向、origin、affine 三个检查点现象训练时 loss 一直不降可视化 patch 发现分割结果整体偏移了几个体素但肉眼看起来形状相似只是位置不对。原因最常见是 NIfTI 的方向或 origin 不一致。nibabel 读取时把数据按某个固定方向输出如果原始数据的 affine 有问题或者在重采样后没有保留 origin标签和原图就会错位。另一个场景是从 DICOM 转换时标签是在不同软件里勾画的导出的 NIfTI 和原图虽然 shape 一致但方向正交变换不同。解决检查三件事——img.affine 是否一致、header.get_qform() 和 sform 是否包含正确的空间信息、重采样后是否用 SetOutputOrigin 保留原始坐标。最直接的办法是把原图和标签叠加用切片可视化逐层看边界是否对齐。我在代码里会做一个快速检查def check_alignment(image, label, slice_idxNone): import matplotlib.pyplot as plt slice_idx slice_idx or image.shape[2] // 2 plt.figure(figsize(12, 5)) plt.subplot(1, 2, 1) plt.imshow(image[:, :, slice_idx], cmapgray) plt.subplot(1, 2, 2) plt.imshow(label[:, :, slice_idx], cmaphot, alpha0.5) plt.show()如果发现 label 出现了镜像翻转或者 90 度旋转优先检查 affine 的 rotation 部分而不是手动把数组翻转。5.2 显存爆炸不是模型的错patch 大小和 batch size 的匹配关系现象能跑通 2D 分割的服务器换成 3D U-Net 后 batch size 设为 2 直接 OOM把 crop size 从 128 减到 96 还是不行。原因3D 卷积的特征图占用显存随着 patch 尺寸呈立方增长。128x128x64 的输入3D U-Net 第一层下采样后就是 64x64x32逐层卷积的中间结果非常多。很多人只看输入大小忽略了下采样后的特征图数量。解决先固定 patch 为 96x96x64batch size 设为 1然后逐步加大 patch 和 batch。我的经验公式是显存 24GB 的卡3D U-Net 基础版本crop 128x128x64 batch 2 基本是临界点想加大 patch就得把 batch 降到 1 或用梯度累积。梯度累积的实现很简单每 n 个 batch 更新一次参数等价于有效 batch size 扩大 n 倍但要注意 BatchNorm 在这种场景下会有问题建议用 GroupNorm 或 InstanceNorm 替代。5.3 归一化统计量泄漏验证集看到了训练集的均值和方差现象训练损失正常下降验证精度很高但换了一台设备的新数据后效果暴跌。原因预处理时把整个数据集的均值和方差一起算完再划分训练集验证集导致验证集的归一化参数包含了自身信息评估结果偏乐观。这在 k-fold 交叉验证时尤其严重。解决归一化统计量必须在训练集划分完成后再计算验证集和测试集使用训练集的统计量。如果对每个样本独立做 z-score则不存在泄漏问题但要注意 CT 的窗宽窗位裁剪参数也要只从训练集统计。更好的做法是把归一化过程封装成 fit_transform 和 transform 两个阶段像 sklearn 的 StandardScaler 一样。5.4 标签类别不平衡到了离谱的程度背景占 99.5%前景器官占 0.5%现象模型预测结果几乎全为背景Dice 分数在 0.1 以下波动训练曲线的 loss 下降得很平滑但就是分割不出器官。原因3D 体数据中器官体积占比往往极小比如胰腺在 CT 中通常只占全部体素的 0.5% 左右。交叉熵损失对每个体素等权模型很容易把全部体素预测为背景来获得低损失。解决除了使用 soft Dice loss 或 focal loss数据准备阶段就要用随机采样补偿类别不平衡。做法是在随机裁剪 patch 时以一定概率保证 patch 内包含前景体素。我在这类任务里通常会设置 70% 的 patch 从前景体素附近采样30% 的 patch 随机采样这样既能保证均衡又能保留背景上下文。另外计算 Dice loss 时对前景类加入平滑项也能缓解零梯度问题。5.5 重采样太慢每次迭代都读一次 NIfTI现象数据量只有 100 例但一个 epoch 要跑 10 分钟其中 8 分钟都在等数据加载。原因Dataset 的getitem里每次读取原始 NIfTI 并做重采样和归一化CPU 成了瓶颈。虽然设置了 num_workers8但因为每个 worker 都在重复读取和重采样整体速度还是上不去。解决把重采样和归一化结果预处理成 .npy 文件缓存到磁盘Dataset 直接读取 .npy。文件数量可以是原始数据的 2 倍image 和 label 各一个也可以为了减少小文件 IO 把多个 patch 拼成单个 .npy。另一个有效手段是用内存映射文件 np.load(path, mmap_moder)这样只把需要的部分读入内存3D patch 裁剪时效率极高。但要注意 mmap 模式下返回的数组不能直接做增强需要先 copy 出来再操作。6. 把数据准备链路串起来一个可以直接改的模板工程结构6.1 文件组织方式我在实际项目中习惯把所有数据准备步骤拆成三个脚本而不是把所有逻辑堆在一个 Dataset 类里。第一个脚本是 1_convert_dicom.py负责把 DICOM 转成 NIfTI第二个脚本是 2_preprocess.py负责重采样、归一化、缓存 npy第三个脚本是 3_dataset.py定义 Dataset 和 transform。这样做的好处是每个阶段都可以单独验证比如重采样后可以单独跑一轮检查确认 spacing 和 shape 都正确再进入缓存步骤。目录结构大致如下project/ ├── data/ │ ├── raw/ # 原始 NIfTI │ ├── preprocessed/ # 重采样和归一化后的 .npy │ └── splits/ # 训练验证集划分的 json ├── src/ │ ├── dataset.py │ ├── transform.py │ └── utils.py └── config.yaml6.2 一个完整的训练前数据流水线代码示例下面是我简化后的预处理主流程适合直接当模板改import os import numpy as np import SimpleITK as sitk from tqdm import tqdm def preprocess_case(image_path, label_path, output_dir, target_spacing(1.0, 1.0, 1.0)): image_sitk sitk.ReadImage(image_path) label_sitk sitk.ReadImage(label_path) image_resampled resample_to_spacing(image_sitk, target_spacing, is_labelFalse) label_resampled resample_to_spacing(label_sitk, target_spacing, is_labelTrue) image_np sitk.GetArrayFromImage(image_resampled).transpose(2, 1, 0).astype(np.float32) label_np sitk.GetArrayFromImage(label_resampled).transpose(2, 1, 0).astype(np.int16) # CT 窗宽窗位裁剪然后 z-score image_np np.clip(image_np, -200, 400) mean image_np[label_np 0].mean() if (label_np 0).sum() 0 else image_np.mean() std image_np[label_np 0].std() if (label_np 0).sum() 0 else image_np.std() image_np (image_np - mean) / (std 1e-8) np.save(os.path.join(output_dir, image_path.split(/)[-1].replace(.nii.gz, _img.npy)), image_np) np.save(os.path.join(output_dir, label_path.split(/)[-1].replace(.nii.gz, _label.npy)), label_np)这个示例里用 SimpleITK 的 GetArrayFromImage 得到的数组顺序是 (z, y, x)和 NIfTI 常见的 (x, y, z) 不一致所以我用了 transpose(2,1,0) 把顺序调整成深度、高度、宽度。很多人忽略这一点导致标签和图像在某个维度翻转。如果你用的是 nibabel它的 get_fdata() 返回的维度顺序就是 NIfTI 定义的顺序和 SimpleITK 不同这个差异非常容易踩坑。6.3 验证数据准备结果可视化切片加统计指标数据准备完在训练之前花十分钟做三层验证第一层看 shape 和 spacing 统计表确认所有病例的 shape 一致第二层随机抽 5 个样本可视化原图和标签叠加第三层检查标签类别数是否符合预期比如预期只有 0、1、2结果出现了 4就说明标签存在异常。这个验证脚本通常不长但收益很高。我有一次在一批公开数据集上做预处理重采样后把第 10 个病例的标签统计出来发现类别 ID 从 1 变成了连续的血管编号原来数据集里标签是实例级而不是语义级。如果不是提前检查类别分布训练出来的模型就毫无意义。3D 分割的数据准备就是这样的活每一步都琐碎但漏掉一步后面都是连锁反应。希望这套思路帮你在自己的数据上少走一段弯路。本文还有配套的精品资源点击获取
