ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

3D图像分割数据准备全指南:从体数据到PyTorch Dataset

3D图像分割数据准备全指南:从体数据到PyTorch Dataset 简介面向医学图像处理与深度学习入门者及科研人员的PyTorch 3D图像分割工程包以Luna16肺结节CT数据集为真实病例完整覆盖从原始CT数据预处理、各向同性重采样、掩膜标签生成、结节边界框标注到UNet3d/VNet3d模型搭建、训练、验证、推理、指标评估与后处理的端到端流程。压缩包共92个文件以49个Python脚本为核心按数据准备、模型定义、训练主程序、损失函数、推理预测、结果评估等模块清晰组织7个npy文件可直接观察中间特征或标签png训练曲线便于查看loss与Dice变化xml、nii/gz、csv等文件提供配置、样本数据与标注信息整体仅61.66MB。已有543人学习下载工程内还保留多版本训练脚本、测试脚本与日志方便对比不同训练策略和代码演进适合结合系列文章边读边跑通完整3D分割基线。1. 3D 图像分割数据准备为什么才是真正的拦路虎基于 PyTorch 的 3D 图像分割任务很多人一上来就扎进模型结构真正卡住大家的其实是数据准备这一关。原因很直白3D 分割要处理的不再是“一张图”而是一整卷体数据——CT、MRI、显微镜堆栈文件动辄几百 MB还带像素间距、方向矩阵、窗宽窗位这些 2D 任务里很少碰的物理信息。标签要重采样对齐、体数据要按显存极限切成 patch、前景往往只占体积的 1%任何一个环节处理得不严谨模型训练出来就是废的。别急着背 U-Net 结构或调参技巧。本文把“从原始体数据到 PyTorch Dataset”这条链路的每个环节拆开讲透格式解析、重采样、归一化、滑窗切块、采样策略、数据增强以及那些只有在 3D 里才会爆发的翻车现场。每一段都带可直接改的代码和参数你在自己机器上跟着走一遍就能把数据管线架起来而不是又攒了一堆收藏没动过手。2. 从原始体数据到标准化采样格式解析、重采样与归一化的代码路线2.1 先摸清体数据的三件事size、spacing 与轴顺序拿到一个 NIfTI 或 DICOM 序列后我习惯先把元数据完整打印出来而不是急着看像素值。三个东西必须确认size 决定后面 patch 怎么切spacing 决定要不要重采样direction 决定三个坐标轴在物理空间里的朝向。同一台机器的 CT 可能是 0.5×0.5×1.5mmMRI 可能是 1×1×1mm如果不管 spacing 直接送到网络里卷积核在三个方向上的物理感受野就是歪的学出来的特征没有可比性。import SimpleITK as sitk import numpy as np def inspect_volume(path): itk_img sitk.ReadImage(path) img sitk.GetArrayFromImage(itk_img) print(itk size (x, y, z):, itk_img.GetSize()) print(numpy shape (z, y, x):, img.shape) print(spacing:, itk_img.GetSpacing()) print(origin:, itk_img.GetOrigin()) print(direction:, itk_img.GetDirection()) print(dtype:, img.dtype, min/max:, img.min(), img.max()) return itk_img, img itk_img, img inspect_volume(case_001_ct.nii.gz)这段代码没有做任何数据处理但它能帮你避掉 3D 分割里最隐蔽的坑轴顺序。SimpleITK 的GetSize()返回的是物理坐标下的(x, y, z)而GetArrayFromImage()转出来的 numpy 数组 shape 是(z, y, x)。两个顺序永远相反这是设计如此不是 bug。你后续所有切片操作都要以 numpy 的(z, y, x)为准代码里注释写清楚免得过两周自己看都犯迷糊。还有一个实际经验统一用 SimpleITK 读文件不要和 nibabel 混用。nibabel 对 NIfTI 的处理方式不同它给你的数组在某些情况下和 SimpleITK 读出来是转置过的。团队协作时更要定死一个读库——否则张三预处理的数据和李四的增强代码在轴上悄悄对不上训练前期看不出来中后期才开始掉点排查起来非常痛苦。2.2 重采样到各向同性一张表说清插值方式怎么选spacing 不统一3D 卷积就学不到一致的物理结构。常见做法是把所有训练数据重采样到各向同性CT 类数据默认 1.0×1.0×1.0mm。比如原始层厚 1.5mm、层内 0.5mm重采样到 1mm 立方体体素数会变化z 方向相当于上采样xy 方向相当于下采样。这个操作里最容易翻车的不是重采样本身而是插值方式。数据类型插值方式原因默认填充值CT 图像sitkLinear保留 HU 值的连续梯度-1000空气MRI 图像sitkLinear保留强度渐变0标签掩膜sitkNearestNeighbor保持类别为离散整数0背景标签掩膜用线性插值是大忌。0/1 的标签经过线性插值会变成 0.3、0.7 这种中间值模型学出来的边界是模糊的Dice 永远上不去。图像用线性插值、标签用最近邻插值这套组合是 3D 分割数据准备的基本盘。def resample_to_isotropic(itk_img, target_spacing(1.0, 1.0, 1.0), is_labelFalse): original_spacing itk_img.GetSpacing() original_size itk_img.GetSize() new_size [ int(round(orig_sz * orig_spc / target_spc)) for orig_sz, orig_spc, target_spc in zip(original_size, original_spacing, target_spacing) ] resampler sitk.ResampleImageFilter() resampler.SetSize(new_size) resampler.SetOutputSpacing(target_spacing) resampler.SetOutputOrigin(itk_img.GetOrigin()) resampler.SetOutputDirection(itk_img.GetDirection()) if is_label: resampler.SetInterpolator(sitk.sitkNearestNeighbor) resampler.SetDefaultPixelValue(0) else: resampler.SetInterpolator(sitk.sitkLinear) resampler.SetDefaultPixelValue(-1000) return resampler.Execute(itk_img) ct_resampled resample_to_isotropic(itk_img, target_spacing(1.0, 1.0, 1.0)) label_resampled resample_to_isotropic(itk_label, target_spacing(1.0, 1.0, 1.0), is_labelTrue)注意new_size的计算用round而不是int()直接截断。如果直接截断重采样后的物理范围会比原始范围小一两个体素图像和标签就会在边界上错位。另外SetDefaultPixelValue对 CT 图像要填 -1000这是空气的 HU 值如果填 0重采样边界会出现一圈假的高密度环模型会把这圈环当成真实结构去学。2.3 归一化的临床与工程视角窗宽窗位还是 z-scoreCT 的原始 HU 值范围很大从 -1024 到 3000直接做 min-max 归一化会被金属伪影或极高密度结构毁掉。临床常规是先做窗宽窗位裁剪把关注的密度范围映射出来。做肝脏分割我一般用窗宽 1500、窗位 -600做肺结节分割用窗宽 1500、窗位 -600 也很常见。裁剪后再线性映射到 [0,1] 或 z-score。def hu_windowing_norm(img, window_center-600, window_width1500): lower window_center - window_width / 2.0 upper window_center window_width / 2.0 img np.clip(img, lower, upper) img (img - lower) / (upper - lower) return img.astype(np.float32)参数含义window_center是窗位决定你关注哪个 HU 密度段window_width是窗宽决定这个密度段的跨度。裁剪掉的部分不是噪声而是你不想看的密度范围。MRI 没有 CT 这种绝对物理单位不能照搬窗宽窗位最稳的做法是每个 volume 单独做 z-score 归一化或者按 percentile 裁剪再标准化。多中心数据要特别注意归一化参数的记录。不同 CT 设备的重建算法、辐射剂量会影响 HU 分布如果每个中心都用自己的窗宽窗位模型会把扫描仪差异当成解剖差异。我一般会把每批数据的均值和标准差存到 JSON 里推理阶段用训练集的参数对测试数据做同样的变换保证前后一致。3. 把体数据切成 patch滑窗索引、前景优先与显存这笔账3.1 显存账怎么算patch 尺寸与 batch 的取舍表3D U-Net 这类模型不能像 2D 那样整图输入。一张 512×512×300 的 CT单通道 float32 就要 300MB经过四层下采样后中间特征图爆炸式增长显存根本扛不住。所以要切 patch比如 128×128×128 甚至 96×96×96。显存占用不是线性涨的patch 从 96 提到 128体积变成 2.37 倍显存需求也跟着翻两倍多复杂度是三次方关系。patch 尺寸单样本估算显存3D U-Net 单通道batch2 的建议显存64³2-3 GB8 GB 可跑96³6-10 GB16 GB 起步128³16-20 GB24 GB 或以上160³30 GB多卡或需特殊优化这张表不是绝对值特征图通道数、网络深度不同都会变化但量级关系是稳定的。调显存优先动 patch 而不是 batch。batch 降到 1 省不了多少patch 降一档立竿见影。先拿 96³ 把流程跑通再根据显存余量逐步往上升这一条能省掉很多深夜 OOM 的烦恼。3.2 滑窗索引表把块切让给离线计算去做patch 怎么切是有讲究的。直接随机裁切片大多数 patch 落在背景区如果固定裁中心又会错过边缘的病灶。常见做法是把滑窗得到的所有候选位置离线算好存成索引表训练时从表里抽。stride 一般取 patch_size 的一半这样相邻 patch 有 50% 重叠保证分割目标跨 patch 时不至于信息断裂。def build_patch_index(volume_shape, patch_size, stride): z, y, x volume_shape pz, py, px patch_size z_steps list(range(0, z - pz 1, stride[0])) y_steps list(range(0, y - py 1, stride[1])) x_steps list(range(0, x - px 1, stride[2])) # 最后一步如果覆盖不到边界手动补一个贴边的位置 if z_steps[-1] ! z - pz: z_steps.append(z - pz) if y_steps[-1] ! y - py: y_steps.append(y - py) if x_steps[-1] ! x - px: x_steps.append(x - px) indices [] for zi in z_steps: for yi in y_steps: for xi in x_steps: indices.append((zi, yi, xi)) return indicesstride传三元组(sz, sy, sx)三个方向可以不同步。z 轴分辨率较低的体数据z 方向 stride 可以调大一些。补边的逻辑是为了让滑窗覆盖到体积的每个角落否则靠近边缘的体素永远没有对应的完整 patch训练时没问题推理时就会在边缘出现空洞。索引表建好后存成 pickle 或 npy 文件Dataset 初始化时直接加载。不要在__getitem__里每次现算滑窗位置那是重复劳动还会带来随机性的不可控。预处理阶段把索引表固化下来训练时才能保证每个 epoch 采样的一致性排查问题也容易——同样的索引表同样的 seed结果应该完全可复现。3.3 前景与背景不平衡按比例采样比各种 loss 更直接器官分割里前景占到体积的 5% 都算多的肺结节这种小目标可能只有成千上万个体素占整体的千分之一。如果不处理采样策略模型看到的 patch 几乎全是背景训练到后面 loss 不怎么降Dice 也上不去。与其依赖 loss 函数去硬扛不如在采样阶段直接动手把候选 patch 分成“含前景”和“纯背景”两桶按比例抽。def split_fg_bg_indices(label, patch_size, stride): all_indices build_patch_index(label.shape, patch_size, stride) fg_indices [] bg_indices [] for (zi, yi, xi) in all_indices: patch_label label[ zi:zi patch_size[0], yi:yi patch_size[1], xi:xi patch_size[2] ] if patch_label.sum() 0: fg_indices.append((zi, yi, xi)) else: bg_indices.append((zi, yi, xi)) return fg_indices, bg_indices fg_idx, bg_idx split_fg_bg_indices(label_resampled_np, patch_size(96, 96, 96), stride(48, 48, 48)) print(fforeground patches: {len(fg_idx)}, background patches: {len(bg_idx)})foreground_ratio一般设 0.5 到 0.7意思是每 10 个 patch 里有 5 到 7 个是含前景的。这个比例不是越高越好前景 patch 之间重叠度大采样相关性太高模型容易对特定位置过拟合。如果前景桶里的 patch 数量很少比如只有几十个那每次 epoch 都在反复看同样的位置这时候要把数据增强的幅度加一些或者把 stride 调小增加前景桶的样本量。背景桶也不要直接扔它负责让模型学会什么是“正常组织”。完全不要背景 patch模型会把所有非前景区域都误判成病灶这个方向同样不对。4. 撑起训练循环的 Dataset 与 DataLoader索引表、缓存与多进程的取舍4.1 Dataset 只做单一职责切 patch、上变换、返回张量像 mmdetection 这类成熟的 2D 工具包把数据准备和训练循环解耦得很干净3D 任务虽然没有现成框架可以照搬但思路可以借用。Dataset 里只做三件事按索引取 patch、做数据增强、返回张量。读文件、重采样、归一化这些重活全部放到预处理阶段不要在 Dataset 里碰 NIfTI。import torch import torch.utils.data import random class PatchDataset(torch.utils.data.Dataset): def __init__(self, volume, label, fg_idx, bg_idx, patch_size, foreground_ratio0.6, transformNone): self.volume volume self.label label self.fg_idx fg_idx self.bg_idx bg_idx self.patch_size patch_size self.foreground_ratio foreground_ratio self.transform transform def __len__(self): # 固定长度训练时用 iteration 数控制每个 epoch 的步数 return 100000 def __getitem__(self, _): if self.fg_idx and random.random() self.foreground_ratio: zi, yi, xi random.choice(self.fg_idx) else: zi, yi, xi random.choice(self.bg_idx) pz, py, px self.patch_size img_patch self.volume[zi:zi pz, yi:yi py, xi:xi px] lbl_patch self.label[zi:zi pz, yi:yi py, xi:xi px] if self.transform: img_patch, lbl_patch self.transform(img_patch, lbl_patch) # 输入加 channel 维输出与输入保持相同空间尺寸 return torch.from_numpy(img_patch).float().unsqueeze(0), \ torch.from_numpy(lbl_patch).long().unsqueeze(0)__len__返回 100000 这个固定大数和实际 patch 数量无关。这不算作弊因为采样是加权随机抽样每次__getitem__都可能落在不同位置epoch 的概念被弱化了训练步数完全由你自己控制。foreground_ratio起的作用是调节两类 patch 的比例而不是精确计数。如果想更严格可以写成len(fg_idx) len(bg_idx)但那样 shuffle 后每轮恰好见一轮不如固定长度灵活。输入张量加unsqueeze(0)是因为 3D 卷积期望的输入是(B, C, D, H, W)单样本也要有 channel 维。标签用long()是为了配合交叉熵和 Dice loss如果标签是 0/1 掩膜float()也行但很多框架里混合 loss 对输入类型有要求提前统一成 long 能少踩一个错。4.2 DataLoader 的性能参数num_workers、pin_memory 与缓存3D patch 的数据增强计算量不小如果 CPU 来不及就需要调 DataLoader 的并行参数。我的默认配置是num_workers8、pin_memoryTrue、persistent_workersTrue。num_workers看机器核心数8 到 12 都是合理区间开太多反而会因为进程切换和内存竞争变慢。pin_memory对 GPU 训练是必要的它把 CPU 内存锁页减少显存拷贝时间。prefetch_factor4可以让每个 worker 预取 4 个 batch掩盖数据加载延迟。train_loader torch.utils.data.DataLoader( train_ds, batch_size2, shuffleFalse, # 采样的随机性在 Dataset 内部完成 num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue, )shuffleFalse是因为索引抽样的随机性已经在__getitem__里做了外部 shuffle 反而打散不了什么。如果未来改成确定性迭代再开shuffleTrue也不迟。persistent_workersTrue让 worker 进程在 epoch 之间不销毁重建省掉一大笔进程创建开销代价是内存占用会高一些训练完记得释放。数据缓存策略上首选把预处理完的 volume 直接留在内存里。一个 case 重采样到 1mm 后是 256×256×256 的 float32大约 64MB几十个 case 也就几个 GB一次全载入没问题。只有数据总量大到内存放不下时才考虑 npy 文件按需读取但那样要注意磁盘 IO 成为瓶颈——机械盘跑 3D 数据准备就是灾难至少换 SSD。在执行数据管线优化之前先确认 PyTorch 和 CUDA 版本匹配3D 卷积对 cuDNN 依赖更强环境不对的话 DataLoader 再快训练也白搭。4.3 2D 堆叠不是 3D别用切片栈骗过自己的数据链路很多刚接触体数据的人以为把 2D 切片一张张叠起来就是 3D 数据。严格说那是“3D 数组”不是“3D 影像”。关键区别是 spacing如果层内分辨率 0.5mm、层厚 5mm直接堆出来的体数据在 z 轴上是严重欠采样的3D 卷积在 z 方向的感受野换来的全是无效信息。slice_paths sorted(glob.glob(case_dicom_slices/*.dcm)) # 方案一SimpleITK 直接读整个 DICOM 序列 itk_vol sitk.ReadImage(case_dicom_slices/) # 方案二如果只有逐层 PNG需要手动拼接并设置物理间距 itk_img_list [sitk.ReadImage(p) for p in slice_paths] itk_vol sitk.JoinSeries(itk_img_list) itk_vol.SetSpacing((x_spacing, y_spacing, z_spacing))z_spacing必须来自两层之间的物理距离不是 DICOM 标签里的 slice thickness。有的扫描序列层与层之间有间隙gapslice thickness 是扫描层本身的厚度两者不是一回事。取 z 坐标列表做差分才是准确的。这个细节是个容易随意填黑匣子的地方填错了数据链路照样跑模型性能差却找不出原因。脑不懒预处理时多打印几个数值对账后续能省几个失眠夜。5. 3D 数据准备的高频翻车现场5 个典型错误与排查顺序5.1 标签被插值污染了验证集 Dice 为 0 的隐性原因现象训练 loss 正常下降验证集 Dice 始终在 0 附近或者涨到 0.5 就再也上不去。原因重采样标签时用了线性插值。0/1 的掩膜被插值成 0.3、0.7 这样的中间值模型输出的预测和这种模糊标签算 Dice边界区域永远对不齐。更隐蔽的是有些数据用 nearest 重采样后没有检查标签是否还是纯 0/1中间值已经悄悄混进去了。解决标签重采样只允许sitkNearestNeighbor并且预处理脚本里加一行断言unique_values np.unique(label_after_resample) assert set(unique_values).issubset({0, 1}), flabel has unexpected values: {unique_values}把这条断言写进预处理流程里跑完一批数据就验证一批不要等到训练结束才回头查。5.2 SimpleITK 与 nibabel 混用轴的错位在训练中后期才爆发现象训练前几个 epoch 指标正常到中后期 loss 出现周期性尖峰或者验证指标忽高忽低。原因一部分数据用了 SimpleITK 读另一部分用了 nibabel 读。两个库对 NIfTI 的轴顺序定义不一致模型在不同样本上看到的是翻转或转置后的结构权重被来回拉扯。解决团队统一读库所有预处理脚本只允许从一个库取数。我在项目里固定用 SimpleITK因为它的GetArrayFromImage返回值轴顺序有明确文档而 nibabel 的轴方向处理更灵活灵活到容易用错。一旦定了读库预处理产物统一存成 npy 或 npz后续所有环节都不要碰原始文件。5.3 滑窗推理接缝重叠融合参数不一致现象推理输出出现了棋盘格一样的接缝尤其是分割目标跨 patch 边界时预测结果在边界处有明显的跳变。原因训练时 stride 是 patch 的一半推理时为了省时间把 stride 调大了或者重叠区用了硬切的融合方式。训练和推理的采样分布不一致模型在重叠区的预测置信度自然不连续。解决推理的 stride 和训练保持一致重叠区的融合用高斯权重或线性加权而不是简单地取最大值。预测概率图在重叠区按权重叠加再取 argmax边缘会平滑很多。5.4 显存 OOM 不只 batch 的问题patch 是三次方增长现象8GB 显卡跑 batch2 直接 OOM报错经常出现在 backward 阶段而不是 forward。原因显存峰值在反向传播时达到最大和 forward 相比多了一倍的激活值存储。patch 从 96³ 提到 128³显存需求变成 2.37 倍batch 反而只占了个零头。解决先降 patch 尺寸再降 batch。如果 96³ 也放不下就开混合精度AMP或者用梯度累积模拟大 batch。不要一上来就为省显存牺牲感受野——patch 太小导致分割目标超出感受野范围这个损失比 OOM 更难发现。5.5 数据增强反而拉低了分数强度扰动的边界现象开了数据增强之后loss 比不开还高Dice 也在下降。原因强度扰动超出了合理值域。比如归一化后的图像是 [0,1]高斯噪声加到了 0.5 的标准差把器官边界彻底淹没或者弹性形变网格过猛器官结构扭曲到解剖学上不存在。解决增强组合按强度递增逐级加。先把随机翻转和 ±10% 的强度扰动跑通确认 loss 正常后再加任意角度旋转、弹性形变。每加一种增强跑 20 个 iteration 对比 loss 变化。增强不是为了炫技是让模型见过更多合理的变体强度越过边界就是人为制造噪声。6. 建模前最后的验证用 10 分钟检查你的数据链路是否真的通了6.1 smoke test8 个样本暴露 80% 的数据问题数据准备脚本写完之后不要直接开始训练先跑一个 smoke test。它的目标不是验证模型而是验证数据链路本身。8 个样本足够暴露绝大多数问题轴错了、标签有中间值、前景为空、归一化范围不对。def smoke_test(dataset, n8): for i in range(n): x, y dataset[i] assert x.shape y.shape, fshape mismatch: {x.shape} vs {y.shape} assert x.dtype torch.float32, finput dtype: {x.dtype} assert set(torch.unique(y).tolist()).issubset({0, 1}), label leak fg_ratio (y 0).float().mean().item() assert fg_ratio 0, fsample {i} has empty foreground print(fsample {i}: shape{tuple(x.shape)}, fg_ratio{fg_ratio:.4f}) print(smoke test passed)fg_ratio 0是这里最容易被忽略的检查。如果加权采样逻辑写错或者前景桶索引为空有些样本会吐出全背景 patch模型在训练时学到的全是背景评估时自然全错。dtype 检查也很关键3D 卷积对 float32 的精度更敏感float16 做推理可以训练阶段不要省这点内存。6.2 不调模型先看 loss过拟合最小配置是数据链路的试金石smoke test 过了下一步是拿 2 到 3 个样本跑一个过拟合测试。固定 seed、固定模型为小型 3D U-Net、固定优化器训练 30 个 iterationloss 应该显著下降。如果 loss 纹丝不动或者下降后迅速反弹问题大概率还在数据侧图像和标签没有对齐数据增强用了不同的变换参数或者归一化把目标区域压成了接近 0 的数值。我踩过最深的坑是在重采样标签时图省事用了 linear 插值验证集 Dice 卡了整整两天最后是 smoke test 的 unique 值检查救了命。从那以后每个 3D 项目我都先把数据管线验证跑完再碰模型。你也可以把 smoke test 脚本写进预处理流程的固定环节每次数据更新后跑一遍确认链路是真的通再进入调模型的黑匣子。这十分钟的检查价值超过很多天的盲目调参希望帮到你。本文还有配套的精品资源点击获取
返回列表