UNet迁移实战：如何用Labelme标注自己的数据，并快速替换官方数据集进行训练

UNet迁移实战：从Labelme标注到自定义数据集训练全流程指南

当你在GitHub上成功运行了UNet的官方Demo后，下一步自然是想让这个强大的语义分割模型为你自己的项目服务——无论是分析医学影像中的病变区域，还是识别卫星图片中的特定地物。本文将手把手带你完成从原始图像标注、格式转换到模型训练的全过程，重点解决两个核心问题：如何用Labelme高效标注自己的数据，以及如何无缝替换官方数据集。

1. 环境准备与工具选择

在开始标注前，需要确保你的开发环境已经就绪。不同于简单的Demo运行，真实项目往往需要处理更大的数据量和更复杂的场景。

1.1 基础环境配置

推荐使用Anaconda创建独立Python环境，避免依赖冲突：

conda create -n unet_labelme python=3.8 conda activate unet_labelme

安装UNet训练所需的核心库：

pip install torch torchvision pip install opencv-python pillow numpy matplotlib

1.2 标注工具选型对比

提示：对于个人研究者和小型项目，Labelme以其轻量化和灵活性成为首选。它生成的JSON格式也能方便地转换为各种深度学习框架所需的数据格式。

2. Labelme标注实战技巧

2.1 安装与基础标注

通过pip安装Labelme：

pip install labelme

启动标注界面：

labelme

高效标注工作流：

1. 使用Open Dir加载图像文件夹
2. 点击Create Polygons开始绘制多边形
3. 右键完成当前多边形绘制
4. 为每个区域指定类别标签
5. 使用Ctrl+S保存当前标注

2.2 高级标注策略

对于医学影像等专业领域，这些技巧能显著提升标注质量：

• 边缘精确控制：放大图像(鼠标滚轮)进行像素级调整
• 快捷键加速：
  • Ctrl+Z撤销上一步操作
  • Del删除选中多边形
  • Ctrl+J复制选中形状
• 批量处理：通过Next Image快速切换未标注图像

标注完成后，每个图像会生成对应的JSON文件，包含所有多边形的坐标和类别信息。

3. 数据格式转换：从JSON到UNet掩码

UNet训练需要的是二值化的掩码图像(mask)，而Labelme生成的是矢量标注。我们需要编写转换脚本实现这一关键步骤。

3.1 基础转换脚本

创建labelme2mask.py文件：

import json import os import numpy as np import cv2 from glob import glob def json_to_mask(json_path, output_dir, class_mapping): with open(json_path) as f: data = json.load(f) img_shape = (data['imageHeight'], data['imageWidth']) mask = np.zeros(img_shape, dtype=np.uint8) for shape in data['shapes']: label = shape['label'] points = np.array(shape['points'], dtype=np.int32) cv2.fillPoly(mask, [points], color=class_mapping[label]) base_name = os.path.basename(json_path).replace('.json', '.png') cv2.imwrite(os.path.join(output_dir, base_name), mask) # 示例使用 class_mapping = {'background': 0, 'tumor': 1, 'organ': 2} # 根据实际类别修改 json_files = glob('path/to/labelme_json/*.json') os.makedirs('masks', exist_ok=True) for json_file in json_files: json_to_mask(json_file, 'masks', class_mapping)

3.2 处理多类别场景

对于多类别分割，需要特别注意：

1. 类别映射表：确保class_mapping字典包含所有可能的标签
2. 边缘处理：重叠区域的处理策略（后标注覆盖 or 取最大值）
3. 可视化验证：生成检查图像确认转换正确性

def visualize_mask(image_path, mask_path): image = cv2.imread(image_path) mask = cv2.imread(mask_path, 0) # 为不同类别赋予不同颜色 colored_mask = np.zeros_like(image) colored_mask[mask == 1] = [0, 0, 255] # 红色表示类别1 colored_mask[mask == 2] = [0, 255, 0] # 绿色表示类别2 overlay = cv2.addWeighted(image, 0.7, colored_mask, 0.3, 0) cv2.imshow('Validation', overlay) cv2.waitKey(0)

4. 数据集集成与UNet适配

4.1 文件结构规范

UNet通常期望特定的数据集结构：

my_dataset/ ├── images/ │ ├── train/ │ │ ├── case1.png │ │ └── case2.png │ └── val/ │ ├── case3.png │ └── case4.png └── masks/ ├── train/ │ ├── case1.png │ └── case2.png └── val/ ├── case3.png └── case4.png

4.2 关键代码修改点

在UNet训练脚本中，通常需要调整以下参数：

1. 数据加载器修改：

# 原代码可能类似这样 train_dataset = Dataset( img_dir="original_images_dir", mask_dir="original_masks_dir", transform=transforms ) # 修改为你的路径 train_dataset = Dataset( img_dir="my_dataset/images/train", mask_dir="my_dataset/masks/train", transform=transforms )

2. 类别数量调整：

# 修改模型输出通道数 model = UNet(n_channels=3, n_classes=len(class_mapping)) # 原可能是n_classes=1

3. 损失函数适配：

# 二分类常用BCEWithLogitsLoss # 多分类则需要CrossEntropyLoss criterion = nn.CrossEntropyLoss() if len(class_mapping) > 2 else nn.BCEWithLogitsLoss()

4.3 数据增强策略

针对不同领域数据的特性，需要定制化的增强策略：

医学影像增强示例：

from albumentations import ( Compose, Rotate, RandomBrightnessContrast, ElasticTransform, GridDistortion, OpticalDistortion ) transform = Compose([ Rotate(limit=15, p=0.5), RandomBrightnessContrast(p=0.3), ElasticTransform(p=0.2, alpha=120, sigma=6), GridDistortion(p=0.1) ])

卫星图像增强示例：

transform = Compose([ RandomRotate90(p=0.5), Flip(p=0.5), Transpose(p=0.5), RandomResizedCrop(height=256, width=256, p=0.3) ])

5. 训练优化与调试技巧

5.1 学习率策略对比

5.2 常见问题排查

问题1：损失值不下降

• 检查数据路径是否正确
• 验证掩码是否与图像对齐
• 尝试减小学习率

问题2：预测结果全黑/全白

• 检查类别权重是否平衡
• 验证损失函数是否适合多分类
• 检查最后一层激活函数是否正确

问题3：GPU内存不足

• 减小batch size
• 使用梯度累积：

for i, (images, masks) in enumerate(train_loader): outputs = model(images) loss = criterion(outputs, masks) loss = loss / accumulation_steps # 梯度累积 loss.backward() if (i+1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

在实际项目中，最耗时的往往不是模型训练本身，而是数据准备和调试过程。使用小样本(10-20张)进行快速验证可以节省大量时间——先确保在小样本上能过拟合(训练损失趋近于0)，再扩展到全量数据。