1. 从零开始为什么你需要亲手构建一个分割数据集如果你对图像分割感兴趣尤其是医学影像分析、自动驾驶感知或者卫星图像解译那么UNet这个名字你一定不陌生。它就像一个“看图说话”的专家能把一张图片里的不同物体比如肺部CT里的病灶、街景照片里的行人车辆用不同颜色精准地“圈”出来。但很多朋友卡在了第一步网上公开的数据集要么类别不对要么数据量不够想用自己的图片训练一个专属的UNet模型却不知道从何下手。我刚开始做项目时也踩过不少坑。网上教程要么太理论只讲模型原理要么太零碎代码跑不通。最头疼的就是数据准备环节图片怎么标注标注完的JSON文件怎么变成模型能“吃”的格式目录结构怎么摆参数怎么改这些问题不解决再厉害的模型也是空中楼阁。所以这篇指南就是来解决这些实际问题的。我会带你走一遍我从头构建一个多分类语义分割数据集的完整流程从安装标注工具开始到最终训练出能用的模型。整个过程就像搭积木每一步我都会解释清楚“为什么这么做”并提供可以直接复制粘贴的代码。我们的目标很明确让你能用自己的图片训练一个能实际跑起来的多分类UNet模型。无论你是学生、研究员还是开发者只要你有Python和PyTorch的基础跟着做就能搞定。2. 实战第一步用Labelme搞定图像标注万事开头难数据标注就是这“开头”。对于语义分割我们需要告诉模型图片上每一个像素点属于哪一类。手动给每个像素涂色那太疯狂了。我们需要一个高效的工具——Labelme。它是一个图形化的标注工具特别适合多边形标注简单易用。2.1 安装与配置你的标注环境首先确保你的电脑上安装了Python。我强烈建议使用Anaconda来管理环境它能避免各种包版本冲突的“玄学”问题。打开你的终端Windows用Anaconda PromptMac/Linux用Terminal我们创建一个专属的标注环境conda create -n labelme_env python3.7 conda activate labelme_env环境建好后安装Labelme就一行命令pip install labelme安装完成后在命令行输入labelme并回车一个图形界面就会弹出来。恭喜你的“数字画笔”已经就位。注意在开始标注前检查一下你的图片。Labelme对图片格式有要求最好是常见的.jpg或.png格式并且是24位真彩色。你可以用任何图片查看器检查属性或者用PIL库简单验证。如果图片格式不对可以用画图工具另存一下或者写个简单的Python脚本批量转换。2.2 标注技巧与文件管理点击Labelme界面的“Open Dir”按钮打开存放原始图片的文件夹。选择一张图片点击左侧的“Create Polygons”按钮就可以开始沿着目标物体的边缘点点连线了。比如我们要分割街景中的“汽车”、“行人”、“道路”那么每看到一辆车就沿着它的轮廓点一圈形成一个闭合多边形。画完一个多边形后会弹出一个对话框让你输入标签名。这里非常关键你需要预先规划好所有类别。例如你可以设定_background_(背景这是必须的)carpersonroadbuilding请保持标签名称的一致性不要一会儿用“car”一会儿用“auto”。标注完一张图片后保存会生成一个同名的.json文件。这个JSON文件里记录了图片路径、所有多边形的顶点坐标以及对应的标签名。我的经验是建立一个清晰的项目文件夹结构从最开始就养成好习惯。我会这样组织my_segmentation_project/ ├── raw_images/ # 存放所有的原始.jpg图片 ├── annotations/ # 存放Labelme生成的.json标注文件 └── labels.txt # 记录所有类别名称的文本文件这个labels.txt文件内容如下第一行固定__ignore__ _background_ car person road building踩坑提醒标注是个细致活多边形的点要尽量贴合物体边缘特别是拐角处。但也不必过分追求像素级完美对于训练来说大体准确且一致的标注比少数几张完美标注更重要。建议先标注50-100张图片作为启动数据集。3. 核心转换将JSON标注变成模型能懂的Mask图Labelme生成的JSON文件对人类友好但UNet模型可不认识。模型训练需要的是标签掩码图即一张和原图尺寸相同的图片其中每个像素点的值代表了其类别ID比如背景是0汽车是1行人是2。这一步转换是构建数据集的核心。3.1 使用Labelme2VOC脚本进行批量转换我们需要一个“翻译官”脚本把JSON“翻译”成Mask。Labelme社区提供了一个非常实用的脚本通常叫labelme2voc.py。你可以直接在网上找到它或者使用我下面提供的这个稳定版本。创建一个新的Python文件比如就叫labelme2voc.py把以下代码复制进去#!/usr/bin/env python from __future__ import print_function import argparse import glob import json import os import os.path as osp import sys import imgviz import numpy as np import PIL.Image import labelme def main(): parser argparse.ArgumentParser( formatter_classargparse.ArgumentDefaultsHelpFormatter ) parser.add_argument(input_dir, helpinput annotated directory) parser.add_argument(output_dir, helpoutput dataset directory) parser.add_argument(--labels, helplabels file, requiredTrue) parser.add_argument( --noviz, helpno visualization, actionstore_true ) args parser.parse_args() if osp.exists(args.output_dir): print(Output directory already exists:, args.output_dir) sys.exit(1) os.makedirs(args.output_dir) os.makedirs(osp.join(args.output_dir, JPEGImages)) os.makedirs(osp.join(args.output_dir, SegmentationClass)) os.makedirs(osp.join(args.output_dir, SegmentationClassPNG)) if not args.noviz: os.makedirs( osp.join(args.output_dir, SegmentationClassVisualization) ) print(Creating dataset:, args.output_dir) class_names [] class_name_to_id {} for i, line in enumerate(open(args.labels).readlines()): class_id i - 1 # starts with -1 class_name line.strip() class_name_to_id[class_name] class_id if class_id -1: assert class_name __ignore__ continue elif class_id 0: assert class_name _background_ class_names.append(class_name) class_names tuple(class_names) print(class_names:, class_names) out_class_names_file osp.join(args.output_dir, class_names.txt) with open(out_class_names_file, w) as f: f.writelines(\n.join(class_names)) print(Saved class_names:, out_class_names_file) for label_file in glob.glob(osp.join(args.input_dir, *.json)): print(Generating dataset from:, label_file) with open(label_file) as f: base osp.splitext(osp.basename(label_file))[0] out_img_file osp.join( args.output_dir, JPEGImages, base .jpg) out_lbl_file osp.join( args.output_dir, SegmentationClass, base .npy) out_png_file osp.join( args.output_dir, SegmentationClassPNG, base .png) if not args.noviz: out_viz_file osp.join( args.output_dir, SegmentationClassVisualization, base .jpg, ) data json.load(f) img_file osp.join(osp.dirname(label_file), data[imagePath]) img np.asarray(PIL.Image.open(img_file)) PIL.Image.fromarray(img).save(out_img_file) lbl labelme.utils.shapes_to_label( img_shapeimg.shape, shapesdata[shapes], label_name_to_valueclass_name_to_id, ) labelme.utils.lblsave(out_png_file, lbl) np.save(out_lbl_file, lbl) if not args.noviz: viz imgviz.label2rgb( labellbl, imgimgviz.rgb2gray(img), font_size15, label_namesclass_names, locrb, ) imgviz.io.imsave(out_viz_file, viz) if __name__ __main__: main()这个脚本做了几件重要的事它读取每个JSON文件根据多边形信息生成一张灰度图Mask其中像素值就是类别ID并将这张Mask保存为PNG格式。同时它还会把原始图片复制一份到指定位置。3.2 执行转换命令并理解输出现在让我们来运行这个脚本。假设你的文件结构如下project/ ├── data_annotated/ # 这个文件夹里同时放了原始.jpg图片和对应的.json文件 ├── labels.txt # 刚才定义的类别文件 └── labelme2voc.py # 转换脚本打开终端激活你的labelme环境导航到project目录下运行这条命令python labelme2voc.py data_annotated data_dataset_voc --labels labels.txt解释一下参数data_annotated: 输入目录存放图片.jpg和图片.json。data_dataset_voc: 输出目录脚本会自动创建。--labels labels.txt: 指定你的类别定义文件。运行成功后你会得到一个崭新的data_dataset_voc文件夹里面结构如下data_dataset_voc/ ├── class_names.txt # 复制的类别文件 ├── JPEGImages/ # 存放所有原始图片的副本 ├── SegmentationClass/ # 存放.npy格式的标签NumPy数组可用于后续处理 ├── SegmentationClassPNG/ # **最重要**存放.png格式的掩码图训练就靠它 └── SegmentationClassVisualization/ # 可视化效果图方便你检查标注是否正确重点检查打开SegmentationClassVisualization文件夹看看生成的效果图。彩色覆盖层是否准确贴合了你标注的物体这是验证转换是否成功最直观的方法。如果发现标注有误回到Labelme修改JSON文件重新运行转换命令即可。4. 构建UNet专属的数据集目录上一步我们得到了核心的图片和Mask但UNet的PyTorch实现通常期望一种类似PASCAL VOC格式的目录结构。我们需要手动组装一下。4.1 创建标准化的数据集文件夹在你的项目根目录下新建一个文件夹比如叫MyDataset。在里面创建三个子文件夹MyDataset/ ├── ImageSets/ │ └── Segmentation/ # 这个路径是很多代码默认寻找的 ├── JPEGImages/ └── SegmentationClass/现在进行文件搬运将data_dataset_voc/JPEGImages/下的所有.jpg文件复制到MyDataset/JPEGImages/。将data_dataset_voc/SegmentationClassPNG/下的所有.png掩码文件复制到MyDataset/SegmentationClass/。关键点确保JPEGImages里的每张图片在SegmentationClass里都有一个同名的.png掩码文件。这是数据配对的基础。4.2 划分训练集与验证集我们不能把所有数据都用来训练需要留出一部分做验证以监控模型是否过拟合。这就是ImageSets/Segmentation/文件夹的作用。你需要创建两个文本文件train.txt和val.txt。这两个文件里只写图片的文件名不带后缀每行一个。例如你有200张图片可以按8:2的比例划分train.txt里写160个名字img_001,img_002, ...,img_160。val.txt里写40个名字img_161, ...,img_200。如何生成这两个列表文件手动写太累。我常用一个简单的Python脚本完成import os import random # 设置路径 jpeg_dir MyDataset/JPEGImages all_images [f.split(.)[0] for f in os.listdir(jpeg_dir) if f.endswith(.jpg)] random.shuffle(all_images) # 打乱顺序 split_ratio 0.8 # 80%训练 split_idx int(len(all_images) * split_ratio) train_list all_images[:split_idx] val_list all_images[split_idx:] with open(MyDataset/ImageSets/Segmentation/train.txt, w) as f: f.write(\n.join(train_list)) with open(MyDataset/ImageSets/Segmentation/val.txt, w) as f: f.write(\n.join(val_list)) print(f训练集数量: {len(train_list)}) print(f验证集数量: {len(val_list)})运行这个脚本train.txt和val.txt就自动生成了。记得检查一下确保这些名字在JPEGImages和SegmentationClass文件夹里都存在。5. 修改UNet代码以适配你的数据集现在我们有了标准格式的数据集接下来就是让UNet代码“认识”它。网上有很多UNet的PyTorch实现我们以一份结构清晰的代码为例进行修改。你需要找到以下几个关键文件。5.1 配置数据集路径和类别数首先找到mypath.py或类似的定义路径的文件。你需要修改数据集根目录的路径指向你刚创建的MyDataset文件夹。# mypath.py class Path: staticmethod def db_root_dir(database): if database pascal: return /path/to/your/MyDataset/ # 修改为你的绝对路径 else: print(Database {} not available..format(database)) raise NotImplementedError接着找到定义数据集的文件通常是dataloaders/datasets/pascal.py因为很多代码沿用VOC命名。在这里你需要修改最重要的一个参数类别数量。# dataloaders/datasets/pascal.py class VOCSegmentation(data.Dataset): NUM_CLASSES 21 # 默认是PASCAL VOC的21类将NUM_CLASSES 21修改为你的实际类别数。记住这个数量要包含背景。如果你定义了_background_,car,person,road,building那么类别数就是5。5.2 修改模型定义与训练脚本然后找到模型定义文件。在UNet的初始化函数里同样需要修改n_classes参数。# modeling/unet.py 或 train.py 中模型定义部分 model UNet(n_channels3, n_classes5) # 3通道输入RGB图5个输出类别最后修改训练脚本train.py。主要修改两个地方模型初始化确保和上面一致。GPU设置如果你只有一张显卡保持默认即可如果有多张并想使用需要修改--gpu-ids参数。# train.py 中的相关部分 import argparse parser argparse.ArgumentParser(descriptionPyTorch UNet Training) parser.add_argument(--gpu-ids, typestr, default0, # 单卡用0双卡用0,1 helpuse which gpu to train) # ... 其他参数 ... model UNet(n_channels3, n_classes5) # 修改类别数5.3 一个常见的“坑”标签颜色映射在dataloaders/utils.py文件中经常会找到一个decode_segmap函数它负责将预测的类别ID0,1,2...映射回可视化的RGB颜色。PASCAL VOC有固定的颜色表但我们的自定义类别没有。如果你不关心可视化可以不管它。但如果想看到彩色的预测结果你需要修改这个颜色映射字典为你的每个类别分配一个独特的RGB颜色。# dataloaders/utils.py 片段 def get_cityscapes_labels(): return np.array([ [0, 0, 0], # 背景 - 黑色 [128, 0, 0], # 类别1 - 深红色 [0, 128, 0], # 类别2 - 深绿色 [128, 128, 0], # 类别3 - 橄榄色 [0, 0, 128], # 类别4 - 深蓝色 # ... 为你的每个类别添加RGB值 ])6. 启动训练与模型测试所有配置修改完毕最激动人心的时刻来了——开始训练。6.1 运行训练命令与监控在终端中进入代码目录运行类似下面的命令python train.py --dataset pascal --batch-size 4 --epochs 50参数解释--dataset pascal因为我们沿用了VOC的数据集格式和加载器。--batch-size根据你的显卡内存调整。如果出现“CUDA out of memory”错误就调小这个值比如从4调到2。--epochs训练轮数50是一个常见的起点。训练开始后控制台会打印每个epoch的训练损失和验证集上的精度如mIoU。一定要关注验证集精度的变化趋势。理想情况下训练损失和验证精度都应该随着epoch增加而改善损失下降精度上升。如果训练损失持续下降但验证精度停滞不前甚至下降很可能出现了过拟合这时候就需要考虑提前停止训练、增加数据量或使用数据增强。6.2 使用训练好的模型进行预测训练完成后通常在checkpoints或saved_models目录下会保存性能最好的模型文件如model_best.pth.tar。现在我们可以写一个简单的预测脚本或者修改代码库中自带的demo.py来测试新图片。你需要修改demo.py中的几个关键参数模型加载时的类别数 (n_classes)。测试图片的输入路径 (--in-path)。预测结果的保存路径。一个简化版的预测核心代码如下# 加载模型 model UNet(n_channels3, n_classes5) # 类别数要对齐 checkpoint torch.load(path/to/your/model_best.pth.tar, map_locationcpu) model.load_state_dict(checkpoint[state_dict]) model.eval() # 预处理图片 transform transforms.Compose([ transforms.Resize((512, 512)), # 调整到训练时的大小 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(your_test_image.jpg).convert(RGB) input_tensor transform(image).unsqueeze(0) # 增加batch维度 # 预测 with torch.no_grad(): output model(input_tensor) prediction torch.argmax(output, dim1).squeeze().cpu().numpy() # prediction现在是一个二维数组值就是类别ID(0,1,2,3,4) # 你可以将其转换为彩色图片保存运行修改后的demo脚本就能在你的测试图片上看到分割效果了。第一次看到模型正确识别出你自己标注的物体时那种成就感是非常棒的。7. 避坑指南与效果优化走完全流程你可能已经成功了但也可能遇到一些“坑”。这里分享几个我实战中总结的经验。数据层面的坑类别不平衡如果“背景”像素远多于“汽车”像素模型会倾向于把所有东西都预测为背景。解决方法是在损失函数中使用class_weight给像素少的类别更高的权重。PyTorch的CrossEntropyLoss可以传入weight参数。数据量太少深度学习是数据饥渴型的。如果只有几百张图片很容易过拟合。除了收集更多数据数据增强是救命稻草。在dataloaders/transforms中可以增加随机翻转、旋转、亮度对比度调整等操作能极大提升模型的泛化能力。标注噪声人工标注难免有误。如果发现模型在某个类别上始终学不好可以回头检查这个类别的标注是否一致、准确。训练技巧学习率不要死守默认值。如果训练初期损失不下降可以尝试调大学习率如果后期验证精度剧烈波动可以调小学习率。使用学习率调度器如ReduceLROnPlateau是个好习惯。批量大小在显存允许的前提下适当调大batch_size有时能使训练更稳定。早停耐心观察验证集指标。如果连续10个epoch验证精度不再提升就可以考虑停止训练避免过拟合。模型调试如果预测结果全是一个类别首先检查数据加载是否正确图片和Mask是否配对然后检查n_classes参数是否在代码的每一处都修改一致了。预测时记得将图片用和训练时相同的均值和标准差进行归一化。构建自己的数据集并完成训练是一个从理论到实践的完整闭环。这个过程可能繁琐但每一步都加深了你对模型和数据关系的理解。当你亲手调教的模型开始准确地识别出你关心的目标时所有的努力都是值得的。记住第一次训练结果不理想是常态根据问题回头检查数据、调整参数、迭代优化才是深度学习项目的日常。