简介:一套面向深度学习道路提取任务的nnU-Net格式训练与测试数据集,源自马萨诸塞道路遥感数据集(Massachusetts Roads Dataset),适合使用nnU-Net框架进行遥感图像分割的研究者与工程师。数据包共59个文件,含58张PNG图像和1份dataset.json配置,压缩后约142.02MB;目录严格按nnU-Net规范组织,imagesTr/imagesTs存放训练与测试影像,labelsTr/labelsTs存放道路标注掩膜,训练与测试集划分清晰,默认命名符合nnU-Net预处理要求,可直接被读取使用。影像覆盖城市、郊区和农村等多种地形,道路标签由OpenStreetMap道路中心线栅格化生成,适合检验模型在树木遮挡、阴影等复杂场景下的提取能力。目前已有1772人浏览学习,拿到后无需格式转换即可用于道路分割实验,也可作为遥感语义分割基准数据集对比不同算法。
1. nnunet训练测试数据集:为什么这个框架能“开箱即用”
“nnunet训练测试数据集”这几个词放在一起,容易让人误以为这是某个现成项目的操作手册。实际经历过的人都知道,nnU-Net最反直觉的地方在于:它不靠发明新网络结构,而是靠一套自动配置流程,把预处理、训练策略和推理参数按数据集特征给你定好。你只需要按它的规矩把数据集摆放整齐、写好dataset.json,剩下的plan、train、predict三步命令就能跑起来。对做医学影像分割、又不想在参数调优上耗太多时间的团队来说,这套流程省掉的不是几分钟,是几周。我按自己惯用的落地路径来写:数据集怎么组织、预处理怎么跑、训练怎么训、测试推理怎么出结果,以及最容易返工的那几个坑。
2. 把数据集整理成nnU-Net要的样子:任务目录与JSON配置
2.1 从原始标注到nnU-Net的数据组织:labelsTr、imagesTr与dataset.json
nnU-Net v2的数据组织是按“任务(Task)”管理的。每个任务是一个独立的文件夹,放在nnUNet_raw_data目录底下,命名规则是Task序号_任务名称,比如Task001_liver或Task002_brain。序号就是后面plan和train命令里用的任务编号,名称是给你自己看的,两个部分用下划线连接。文件夹内部固定是imagesTr、labelsTr、imagesTs三个目录,分别放训练图像、训练标签和测试图像,命名中的Tr是training的缩写,Ts是testing的缩写。nnU-Net只吃NIfTI格式(.nii.gz),其他格式需要先转,这个没得商量。
文件名要按“病例名_模态标识符_病例名”来组织。比如一个病例叫case_001的CT图像,在imagesTr里就叫case_001_0000.nii.gz,其中_0000代表模态序号0。标签文件在labelsTr里就叫case_001.nii.gz。这里最容易翻车的是:图像文件和标签文件名必须完全对应,但标签不能带模态后缀。你有一张case_001_0000.nii.gz,就必须有一张case_001.nii.gz在labelsTr里,否则训练时nnU-Net报“找不到配对”的错误。另一个容易忽略的规则是case_001这个名称本身不能再包含下划线,因为nnU-Net解析文件名时靠下划线切分字段,case名字里再带下划线会让解析逻辑直接混乱。
下面的小脚本可以快速把散落的图像和标签整理成这个结构:
import shutil from pathlib import Path raw_root = Path("raw_data") # 原始图像目录 label_root = Path("raw_label") # 原始标签目录 target = Path("nnUNet_raw/nnUNet_raw_data/Task001_demo") (target / "imagesTr").mkdir(parents=True, exist_ok=True) (target / "labelsTr").mkdir(parents=True, exist_ok=True) (target / "imagesTs").mkdir(parents=True, exist_ok=True) for i, img_path in enumerate(sorted(raw_root.glob("*.nii.gz"))): case_id = img_path.stem.split(".")[0] # 取文件名当作病例名 tgt_img = target / f"imagesTr/{case_id}_0000.nii.gz" shutil.copy(img_path, tgt_img) # 找到同名标签 lab_path = label_root / f"{case_id}.nii.gz" if lab_path.exists(): shutil.copy(lab_path, target / f"labelsTr/{case_id}.nii.gz") else: print(f"缺少标签: {case_id}")这段脚本做的事情:遍历原始图像目录,把每个.nii.gz文件复制到imagesTr,并自动加上_0000后缀;同时检查同名标签是否存在,存在的复制到labelsTr,不存在的单独打印出来提醒你补数据。我用它处理过一次几百例的肝脏分割数据,唯一要注意的是case_id不能包含下划线以外的特殊字符,否则训练阶段解析文件名时会异常。
2.2 写dataset.json的字段:channel_names、labels与file_ending
每个Task目录下都要有一个dataset.json,nnU-Net从plan阶段开始就读它。最核心的三个字段是channel_names、labels和file_ending。channel_names描述每个模态通道是什么,单模态CT就是["CT"],多模态比如MRI的T1和T2就写["T1","T2"],顺序要和文件名里的_0000、_0001一一对应。labels是字典,从0开始,0必须是background,后面的1、2就是你要分割的解剖结构。file_ending固定写.nii.gz。
{ "channel_names": { "0": "CT" }, "labels": { "0": "background", "1": "liver", "2": "tumor" }, "numTraining": 120, "file_ending": ".nii.gz" }这里的numTraining要和imagesTr里实际图像数量一致,不一致时plan阶段会警告。labels的键必须是字符串,值是类别名,建议用英文,方便在日志里识别。多模态场景常见做法是把MRI的多个序列按固定顺序排列,比如T1加权、T2加权、FLAIR分别作为_0000、_0001、_0002,同一病例的所有模态都必须存在,缺失一个模态时nnU-Net不会替你补全,训练会以报错告终。如果你处理的是带子区域的分割任务,比如肝脏和肝脏肿瘤分别标注,还可以在dataset.json里额外写regions字段来描述子区域归属,但这个字段是可选的,新手阶段先不用管。
2.3 检查图像与标签形状:跑plan之前先做自检
数据放好、json写完之后,我不会立刻跑plan_and_preprocess,而是先做一个自检:把每个病例的图像和标签都load进来,比对shape和spacing。因为nnU-Net在训练阶段会自动把图像裁剪到非零区域,但如果你图像的尺寸和标签尺寸不一致,等到预处理重采样之后才发现,再回去改数据就晚了。下面这个脚本遍历一遍:
import nibabel as nib from pathlib import Path img_dir = Path("nnUNet_raw/nnUNet_raw_data/Task001_demo/imagesTr") lab_dir = Path("nnUNet_raw/nnUNet_raw_data/Task001_demo/labelsTr") for img_path in sorted(img_dir.glob("*.nii.gz")): case_id = img_path.name.replace("_0000.nii.gz", "") lab_path = lab_dir / f"{case_id}.nii.gz" if not lab_path.exists(): print(f"标签缺失: {case_id}") continue img = nib.load(img_path) lab = nib.load(lab_path) if img.shape != lab.shape: print(f"shape不一致: {case_id} img={img.shape} lab={lab.shape}") else: print(f"OK: {case_id} {img.shape}")这一步花五分钟,能避免后面plan阶段报一堆形状不匹配的错。shape不一致通常是因为原始标注软件把标签裁剪过或重采样过,常见做法是把标签重新resample到图像的空间,建议用SimpleITK的Resample,插值方式选最近邻,避免引入新的标签值。多模态数据还要把每个模态的shape也互相核对一遍,我遇到过T1和T2序列采集范围不同导致shape不一致的情况,这种病例要么裁到共同区域,要么干脆剔除。
3. 跑通预处理与训练:从nnUNet_plan_and_preprocess到nnUNet_train
3.1 安装与环境变量:nnUNet_raw、nnUNet_preprocessed、nnUNet_results三件套
nnU-Net v2用pip安装,命令是pip install nnunetv2,它的命令行工具前缀是nnUNet_。要跑起来,必须先设置三个环境变量。nnUNet_raw对应上面放原始数据的目录;nnUNet_preprocessed是预处理输出目录,plan阶段生成的plans.json和预处理后的数组都落在那里;nnUNet_results是训练结果目录,权重和日志都在里面。这三个目录建议都放到同一块剩余空间足够的盘上,因为预处理后的数据往往比原始数据大好几倍。
export nnUNet_raw="/data/nnUNet/nnUNet_raw" export nnUNet_preprocessed="/data/nnUNet/nnUNet_preprocessed" export nnUNet_results="/data/nnUNet/nnUNet_results"这三个环境变量在每个终端都要export一遍,如果忘记设置,命令会直接提示找不到路径。我一般把它们写进~/.bashrc,省得每次重开终端都要敲一遍。另外要注意nnunetv1和v2的命令不兼容,网上一搜一大把旧教程里nnUNet_plan_and_preprocess的用法是v1风格,v2里已经统一成nnUNet_plan_and_preprocess -t 任务ID这种写法,不需要再手动指定数据集名称。
3.2 nnUNet_plan_and_preprocess:数据指纹与自动配置
预处理阶段做的事情是:扫描所有数据的统计信息(形状、spacing、强度分布、前景占比),然后在3d_fullres、3d_lowres、2d这些配置里选出适合当前数据的方案,同时完成重采样、裁剪、z-score归一化。这一步的输出是预处理后的数据、plans.json和dataset.json副本。命令很简单:
nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity-t 1对应Task001那个序号,--verify_dataset_integrity是额外做一次数据完整性校验,会检查每个病例的图像和标签是否一一对应、尺寸是否一致,比2.3节那个自检脚本更严格。跑完之后看nnUNet_preprocessed/Task001_demo目录,里面应该出现一个plans.json和一堆预处理后的.npz、.npy文件。预处理时间取决于数据量和分辨率,几十例CT大概十几分钟,几百例高分辨率MRI可能要跑几个小时,这个速度是正常的,别看到进度条不动就去Ctrl+C。
plan阶段不是所有配置都会跑。如果数据量和前景区域都不满足3d_lowres的条件,nnU-Net会自动跳过低分辨率配置,只跑2d和3d_fullres。这一步看起来像个黑匣子,但它其实把传统调参过程里“图像应该重采样成什么大小”“loss怎么配”“要不要用级联”这些问题都自动回答了。对新手来说,不需要手动干预;对想干预的熟手,plans.json里的preprocessor_name、resampling_scheme这些字段都可以改,但改完要能说得清为什么,否则不如不改。
3.3 nnUNet_train开始训练:配置、Trainer与fold参数
预处理完成后就能训练了。nnUNet_train命令格式是:
nnUNet_train 3d_fullres nnUNetTrainerV2 1 0第一个参数是配置名,3d_fullres是完整分辨率的三维模型,数据少或显存小时用2d;第二个参数是Trainer类,一般就写nnUNetTrainerV2,想换loss或改学习率就继承这个类写自己的Trainer;第三个参数是任务ID;第四个是fold。fold从0到4,代表5折交叉验证的哪一折,如果只想快速验证流程通不通,可以用0这一个fold训完,后续要出稳定结果再训全5个fold。
模型权重默认写在nnUNet_results/Task001_demo/nnUNetTrainerV2__3d_fullres/fold_0下面,文件名是model_best和model_final_checkpoint。训练日志在同一目录下的training_log_*.txt里,每个epoch会打印loss、Dice等指标。显存不够时,可以在命令里加--disable_deep_supervision减少显存占用,但会牺牲一点精度;也可以手动把batch size调小,训练脚本默认会按数据大小自动选batch size,手动改的话一般从2开始,逐步往下试,同时观察GPU利用率别掉到30%以下。
3.4 用--val跑验证集,先拿到训练阶段的客观指标
训练完fold_0之后,我想先看一眼这个fold在验证集上的表现,再决定要不要把5个fold都训练完。做法是重新执行一次nnUNet_train,但加上--val参数:
nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --val这个命令不会重新训练,而是用刚才训练好的权重在fold_0的验证集(就是训练时留出来的那20%)上做推理,然后输出一个validation文件夹,里面是预测结果,同时打印每个类别的Dice系数。这个验证指标的意义在于:它是nnU-Net自己按数据划分算出来的,和后面nnUNet_predict阶段你自己拿测试集去推的结果,在数据分布上是一致的,可以用来判断训练是否收敛、需不需要加大epochs或者换Trainer。如果你的数据类别极度不均衡,这里还要关注一下每个类别的Dice而不是只看平均Dice,某个类别Dice特别低通常是该类别样本太少,需要回去检查标注质量或考虑用更大的学习率。
4. 用nnUNet_predict做测试推理:输入输出与后处理
4.1 nnUNet_predict的标准命令与关键参数
测试阶段的核心命令是nnUNet_predict。它读取一个输入文件夹里所有图像,输出对应的分割结果。一个完整的调用是:
nnUNet_predict -i /data/test_images -o /data/test_pred \ -t 1 -m 3d_fullres -f 0 \ --save_npz参数含义:-i是测试图像目录,里面放和训练集同名的.nii.gz文件,模态后缀_0000也是必需的;-o是输出目录;-t是任务ID;-m是配置名;-f是指定用哪几个fold的模型,多个fold时写-f 0 1 2 3 4,nnU-Net会把多个fold的结果取平均再argmax,这是它比单模型更稳的重要原因,5折ensemble的预测结果通常比单折高1到2个点Dice;--save_npz会把softmax概率保存下来,占空间但后续做集成或分析方便。
如果训练时用的自定义Trainer,还要加-tr nnUNetTrainerV2指定;如果plan阶段生成了多个配置,要选对你想用的那个。这些参数不一致时最常见的报错是“checkpoint not found”——因为你指定的fold没有训练过,或者Trainer名对不上。排查时先看nnUNet_results/Task001_demo目录下有哪些配置和fold,再回头对照命令里的参数。
4.2 单张与批量推理:输入目录结构决定成败
nnUNet_predict按目录批量处理,输入目录里的每个文件都会被当成一个病例。所以如果你只测一张图,也得放在一个目录里,不能直接给文件路径。输出目录不需要提前创建,程序会自动建。推理完成之后,输出目录里是case名加.nii.gz后缀的分割结果,标签值对应dataset.json里定义的1、2等类别。
结果读进Python里很简单:
import nibabel as nib import numpy as np pred = nib.load("/data/test_pred/case_001.nii.gz") data = pred.get_fdata().astype(np.int16) print(np.unique(data)) # 应该出现0、1、(2)这些标签这里一个容易踩的坑是输出的是label还是probabilities。默认输出的是标签图,只有加了--save_npz才会额外生成npz概率文件。如果你拿到结果发现值全是0和1,而训练时有两个目标类别,多半是label id写错或者模型本身没训好。批量推理时建议分批拷贝输入文件,比如一次放50例,跑完一批再放下批,这样万一某个文件有问题,不会整个任务失败重来。
4.3 TTA开不开、step_size怎么设
nnU-Net推理时默认会启用test time augmentation(TTA),即对输入做几种空间变换后分别预测再取平均,这个机制常被称为测试时训练。TTA通常能涨0.5到1个点Dice,代价是推理时间翻倍。显存够、时间不急时我不关;但如果是做线上服务或处理几百例的大批量数据,加--disable_tta能省一半时间。这里我一般先开着TTA跑一遍测试集,把结果存好,再关掉TTA跑一遍,对比指标差异,如果差异不到0.3个点,后续就关掉TTA换速度。
step_size是滑窗推理的步长比例,默认0.5,表示窗口每次移动一半。对特别大的图像,step_size调到0.8能加快推理,但边界区域的分割会略差;调小到0.3会变慢但更稳。这个参数是推理阶段少数值得手动调的,因为它直接控制滑窗重叠程度,对分割稳定性影响明显。调的时候观察一下输出mask的连通性,如果目标区域出现很多细碎的碎片,多半是step_size太大导致窗口间预测不一致。
4.4 模型checkpoint的“后悔药”:中断续训与权重选择
训练中途断了,不用从头跑。nnUNet_train会记录断点,重新执行同样的训练命令会从最近的checkpoint继续。这里要分清两个文件:model_best是验证集上表现最好的权重,model_final_checkpoint是训练结束时的权重,两个文件在同名目录下。如果训练到一半被kill掉,只有model_final_checkpoint存在,可以用它继续;如果要推理,优先用model_best还是model_final_checkpoint可以看验证指标,常见做法是哪个在验证集上高用哪个。
我在实际项目中习惯每个fold训完后,把model_best单独复制一份到另一个目录,避免后续误操作把结果覆盖。权重文件是几百MB级别,不占空间,但这个习惯救过我两次,一次是重跑训练时把fold_0的结果覆盖了,另一次是清理目录时误删了训练输出。nnU-Net的自动化程度再高,也架不住手动操作失误,备份永远是成本最低的后悔药。
5. nnunet训练测试数据集的5个常见问题:现象、原因与解决
5.1 标签是RGB索引而不是单通道标签:loss怎么都不降
现象:训练几十个epoch,loss在0.6左右不动,Dice一直是0或极低。
原因:很多人从标注工具(比如Labelme或某些Web标注平台)导出的标签是RGB彩色图,比如背景是(0,0,0),目标是(255,0,0),直接存成PNG转成NIfTI后,每个像素是一个三维向量,不是单一类别ID。nnU-Net读进去会把每个向量当成多通道标签来处理,模型无法学习这种离散映射,loss自然降不下去。
解决:把RGB标签映射成索引标签。常见做法是按颜色构建一个查找表,把每个像素的RGB向量映射成整数类别:
import numpy as np from PIL import Image # 颜色到类别ID的映射,按你的实际标注颜色填写 color_map = { (0, 0, 0): 0, # background (255, 0, 0): 1, # liver (0, 0, 255): 2, # tumor } label_img = np.array(Image.open("label.png")) # 形状 (H, W, 3) out = np.zeros((label_img.shape[0], label_img.shape[1]), dtype=np.int16) for rgb, cls_id in color_map.items(): mask = (label_img[..., 0] == rgb[0]) & (label_img[..., 1] == rgb[1]) & (label_img[..., 2] == rgb[2]) out[mask] = cls_id nib.save(nib.Nifti1Image(out, affine), "label_mapped.nii.gz")关键是映射后要检查:np.unique(out)只能出现0、1、2这些连续整数,中间不能有缺口,否则模型训练时类别数量和dataset.json定义对不上,又产生新的报错。我遇到过标注工具把两种颜色混在一起的情况,比如两个类别都含有(255,0,0)分量,这时候要回到标注软件里重新修正,靠脚本补不干净。
5.2 plan_and_preprocess在计算数据指纹时中断
现象:nnUNet_plan_and_preprocess跑到某个病例时直接报错或Killed,日志里没有明确报错信息。
原因:最常见的是图像非零区域dense,比如整幅图都没有全零切片,导致裁剪时申请的内存过大;或者图像spacing异常(某个维度spacing为0),导致重采样计算溢出;再就是内存不够被系统OOM kill。
解决:先看plans.json是否已经生成,如果生成了一部分说明是在统计阶段之后就挂的。然后单跑一遍数据完整性校验:
nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity如果这步能过,再把内存上限放开,或者分批预处理。nnU-Net支持--mode参数,把plan和preprocess拆开执行,先plan再preprocess,这样能定位到底是哪一步挂的。数据本身有问题的,用2.3节的自检脚本把报错case找出来,单独修掉再重跑。这一步最忌讳反复直接重跑完整命令,日志一大片根本看不出哪个case有问题。
5.3 重采样后目标器官变形严重:spacing不一致先做归一化
现象:预处理后的图像和标签里,目标器官的形状看起来和原始图像不一致,细小的结构扭曲、断裂。
原因:nnU-Net的plan阶段会计算所有病例的median spacing,把所有图像重采样到目标spacing。如果一部分病例是1x1x1mm体素,另一部分是5x5x5mm体素,重采样到中间值后,后者被放大的倍数过大,细节就糊了。
解决:在数据准备阶段统一spacing。常见做法是先用SimpleITK把所有图像重采样到同一个spacing范围,比如CT统一到1x1x1mm,然后再交给nnU-Net。
注意:标签重采样必须用最近邻插值,用线性或三次样条会在边缘产生中间值,破坏标签语义。
重采样后还要再跑一次自检,确认图像和标签的shape仍然一致。另外不同设备采集的同一器官数据,spacing差异太大时不要硬塞进同一个数据集训练,先按spacing分桶处理,看各自的分割效果再决定是否合并。
5.4 训练完找不到model_final_checkpoint:fold和命令没对应
现象:训练日志正常打印,但nnUNet_results对应目录里是空的,或者只有partial_checkpoint。
原因:训练进程被提前终止,或者store_checkpoint的路径没写对。nnU-Net训练过程中会周期性保存partial_checkpoint,只有正常跑完一个fold才会落model_best和model_final_checkpoint。如果你只是快速验证模型能跑通,训练被Ctrl+C中断,自然找不到完整权重。
解决:确认训练日志里出现了类似“Epoch done”的完整周期输出。如果只是验证流程,建议把epochs调小、data少放点,但最好还是让它完整跑完一个fold,不然测试阶段nnUNet_predict找不到checkpoint,白忙一场。也可以训练过程中手动把partial_checkpoint复制出来续训,但我不建议新手这么做,续训的路径和参数比较容易搞错,我见过有人把partial_checkpoint当成final去推理,结果输出全黑。
5.5 推理结果全黑或全背景:label IDs和JSON定义对不上
现象:测试集推理完成了,结果图打开全黑,np.unique输出只有0。
原因:测试图和训练图强度分布差太远,比如训练数据是增强CT,测试数据是平扫CT,归一化之后目标区域的响应完全不同;更常见的是labels定义和标注数据不一致,比如标注文件里类别ID是1到5,dataset.json里只定义了0和1两个类别,模型从未见过那些标签。
解决:先跑验证集看看是否正常。如果验证集正常、测试集全黑,优先怀疑数据分布差异,把测试图像的强度分布和训练集对比一下,看是不是范围差太多。如果验证集也全黑,检查标注数据里到底有哪些值,把labels字典补全。不要先怀疑模型,这一步排查顺序能省很多时间。还要顺带检查一下是不是测试图像命名里的模态序号和训练集不一致,_0000写成了_0001,模型拿到的通道根本不是它认识的那个。
6. 一个让nnU-Net结果更稳的习惯:先读预测图再谈指标
我见过不少同事拿到Dice 0.9就急着交付,结果临床医生打开一看,分割结果比原始图像小了一圈,边缘完全对不上。Dice是整体重叠率,对边界偏差不敏感,但对医学分割来说,边界才是最能看出问题的。所以我现在有一个习惯:无论训练还是测试,推理出来的预测图一定先叠加到原图上肉眼看一遍,再去看指标。
用SimpleITK或nibabel都能做这个叠加,最快速的方法是先把预测转成numpy,然后做伪彩色覆盖:
import numpy as np import nibabel as nib import matplotlib.pyplot as plt img = nib.load("case_001_0000.nii.gz").get_fdata() pred = nib.load("case_001_pred.nii.gz").get_fdata() plt.figure(figsize=(12, 6)) plt.subplot(1, 2, 1) plt.imshow(img[:, :, img.shape[2] // 2], cmap="gray") plt.subplot(1, 2, 2) plt.imshow(img[:, :, img.shape[2] // 2], cmap="gray") plt.imshow(np.ma.masked_where(pred[:, :, img.shape[2] // 2] == 0, pred[:, :, img.shape[2] // 2]), alpha=0.5, cmap="autumn") plt.savefig("overlay_check.png")这个脚本每次推理后我都会跑一遍,挑几个典型切片存成图。如果要系统验证一个测试集,我会从中随机抽10个病例做叠加图,肉眼检查有没有:分割区域跑到图像外面去、目标器官之间的边界粘连、小结构被漏掉。这三个问题,指标上不一定看得出来,但叠加图一眼就能发现。
做nnU-Net这两年,我最大的教训就是:不要迷信框架的自动化,自动化的plan和train只是把常规参数调好了,数据本身的问题它管不了。每一次拿到新数据集,先把数据质量关卡住,再谈训练和推理。希望这个习惯对你也有用。
本文还有配套的精品资源,点击获取