logo

开发者热搜

从零掌握UNet图像分割：数据集制作、训练与推理全流程详解

作者：php是最好的2025.09.18 17:15浏览量：77

简介：本文详细介绍如何使用UNet模型进行图像语义分割，涵盖从自定义数据集制作、模型训练到推理测试的全流程，适合有一定编程基础的深度学习爱好者。

从零掌握UNet图像分割：数据集制作、训练与推理全流程详解

一、引言：为什么选择UNet进行图像语义分割？

UNet（U-shaped Network）是2015年由Olaf Ronneberger等人提出的经典卷积神经网络架构，专为医学图像分割任务设计，但因其独特的编码器-解码器结构和跳跃连接（skip connections），被广泛应用于各类图像语义分割场景。相比其他模型，UNet具有以下优势：

  • 轻量化：参数量相对较少，适合小规模数据集训练
  • 多尺度特征融合：通过跳跃连接保留低级特征，提升分割精度
  • 端到端训练：可直接输出与输入图像尺寸相同的分割结果

本文将详细介绍如何使用UNet模型，从零开始制作自定义数据集，完成模型训练，并进行推理测试。整个流程基于PyTorch框架实现，适合有一定Python和深度学习基础的读者。

二、数据集制作：从原始图像到标注文件

1. 数据集准备

首先需要准备原始图像和对应的标注掩码（mask）。标注掩码是一张与原始图像尺寸相同的单通道灰度图，其中每个像素值代表对应的类别。例如，在二分类任务中，背景像素值为0，目标像素值为1。

建议：

  • 图像和标注文件应一一对应，文件名建议使用相同前缀（如img_001.jpg对应mask_001.png）
  • 图像尺寸建议统一为256×256或512×512，便于模型处理
  • 划分训练集、验证集和测试集，比例建议为7:1:2

2. 标注工具推荐

  • Labelme：开源标注工具，支持多边形、矩形等多种标注方式，可导出JSON格式标注文件，需转换为掩码图
  • CVAT：专业视频/图像标注平台，支持团队协作
  • Photoshop：手动绘制掩码，适合简单场景

示例：使用Labelme标注后，通过以下Python代码将JSON文件转换为掩码图：

  1. import json
  2. import numpy as np
  3. from PIL import Image
  4. import os
  6. def json_to_mask(json_path, output_path):
  7.     with open(json_path) as f:
  8.         data = json.load(f)
  10.     height = data['imageHeight']
  11.     width = data['imageWidth']
  12.     mask = np.zeros((height, width), dtype=np.uint8)
  14.     for shape in data['shapes']:
  15.         points = shape['points']
  16.         label = shape['label']
  17.         # 简单将所有标注区域设为1（实际应根据label区分类别）
  18.         # 更复杂的实现可使用多边形填充算法
  19.         # 这里仅作示例
  20.         cv2.fillPoly(mask, [np.array(points, dtype=np.int32)], 1)
  22.     Image.fromarray(mask * 255).save(output_path)
  24. # 使用示例
  25. json_to_mask('labelme_output.json', 'mask.png')

3. 数据增强

为提升模型泛化能力，建议对训练数据进行增强。常用方法包括：

  • 随机水平/垂直翻转
  • 随机旋转（±15度）
  • 随机缩放（0.9-1.1倍）
  • 随机亮度/对比度调整
  • 添加高斯噪声

PyTorch实现示例：

  1. import torchvision.transforms as transforms
  2. from torchvision.transforms import functional as F
  3. import random
  5. class CustomAugmentation:
  6.     def __init__(self):
  7.         self.augmentations = [
  8.             self.random_flip,
  9.             self.random_rotation,
  10.             self.random_scale
  11.         ]
  13.     def __call__(self, img, mask):
  14.         for augment in self.augmentations:
  15.             img, mask = augment(img, mask)
  16.         return img, mask
  18.     def random_flip(self, img, mask):
  19.         if random.random() > 0.5:
  20.             img = F.hflip(img)
  21.             mask = F.hflip(mask)
  22.         if random.random() > 0.5:
  23.             img = F.vflip(img)
  24.             mask = F.vflip(mask)
  25.         return img, mask
  27.     def random_rotation(self, img, mask):
  28.         angle = random.uniform(-15, 15)
  29.         img = F.rotate(img, angle)
  30.         mask = F.rotate(mask, angle, fill=(0,))
  31.         return img, mask
  33.     def random_scale(self, img, mask):
  34.         scale = random.uniform(0.9, 1.1)
  35.         h, w = img.size[1], img.size[0]
  36.         new_h, new_w = int(h * scale), int(w * scale)
  37.         img = F.resize(img, (new_h, new_w))
  38.         mask = F.resize(mask, (new_h, new_w), Image.NEAREST)
  39.         # 随机裁剪回原尺寸
  40.         i, j = random.randint(0, new_h - h), random.randint(0, new_w - w)
  41.         img = F.crop(img, i, j, h, w)
  42.         mask = F.crop(mask, i, j, h, w)
  43.         return img, mask

三、UNet模型实现：从架构设计到代码实现

1. UNet架构解析

UNet由编码器（下采样路径）和解码器（上采样路径）组成，通过跳跃连接融合多尺度特征：

  • 编码器：4次下采样，每次通道数翻倍
  • 解码器：4次上采样，每次通道数减半
  • 跳跃连接：将编码器的特征图与解码器的上采样特征图拼接

2. PyTorch实现代码

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  5. class DoubleConv(nn.Module):
  6.     """(convolution => [BN] => ReLU) * 2"""
  7.     def __init__(self, in_channels, out_channels):
  8.         super().__init__()
  9.         self.double_conv = nn.Sequential(
  10.             nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
  11.             nn.BatchNorm2d(out_channels),
  12.             nn.ReLU(inplace=True),
  13.             nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
  14.             nn.BatchNorm2d(out_channels),
  15.             nn.ReLU(inplace=True)
  16.         )
  18.     def forward(self, x):
  19.         return self.double_conv(x)
  21. class Down(nn.Module):
  22.     """Downscaling with maxpool then double conv"""
  23.     def __init__(self, in_channels, out_channels):
  24.         super().__init__()
  25.         self.maxpool_conv = nn.Sequential(
  26.             nn.MaxPool2d(2),
  27.             DoubleConv(in_channels, out_channels)
  28.         )
  30.     def forward(self, x):
  31.         return self.maxpool_conv(x)
  33. class Up(nn.Module):
  34.     """Upscaling then double conv"""
  35.     def __init__(self, in_channels, out_channels, bilinear=True):
  36.         super().__init__()
  37.         self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) if bilinear else nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)
  38.         self.conv = DoubleConv(in_channels, out_channels)
  40.     def forward(self, x1, x2):
  41.         x1 = self.up(x1)
  42.         # 输入是CHW
  43.         diffY = x2.size()[2] - x1.size()[2]
  44.         diffX = x2.size()[3] - x1.size()[3]
  46.         x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
  47.                         diffY // 2, diffY - diffY // 2])
  48.         x = torch.cat([x2, x1], dim=1)
  49.         return self.conv(x)
  51. class OutConv(nn.Module):
  52.     def __init__(self, in_channels, out_channels):
  53.         super(OutConv, self).__init__()
  54.         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
  56.     def forward(self, x):
  57.         return self.conv(x)
  59. class UNet(nn.Module):
  60.     def __init__(self, n_channels, n_classes, bilinear=True):
  61.         super(UNet, self).__init__()
  62.         self.n_channels = n_channels
  63.         self.n_classes = n_classes
  64.         self.bilinear = bilinear
  66.         self.inc = DoubleConv(n_channels, 64)
  67.         self.down1 = Down(64, 128)
  68.         self.down2 = Down(128, 256)
  69.         self.down3 = Down(256, 512)
  70.         self.down4 = Down(512, 1024)
  71.         self.up1 = Up(1024, 512, bilinear)
  72.         self.up2 = Up(512, 256, bilinear)
  73.         self.up3 = Up(256, 128, bilinear)
  74.         self.up4 = Up(128, 64, bilinear)
  75.         self.outc = OutConv(64, n_classes)
  77.     def forward(self, x):
  78.         x1 = self.inc(x)
  79.         x2 = self.down1(x1)
  80.         x3 = self.down2(x2)
  81.         x4 = self.down3(x3)
  82.         x5 = self.down4(x4)
  83.         x = self.up1(x5, x4)
  84.         x = self.up2(x, x3)
  85.         x = self.up3(x, x2)
  86.         x = self.up4(x, x1)
  87.         logits = self.outc(x)
  88.         return logits

四、模型训练：从数据加载到优化策略

1. 数据加载器实现

  1. from torch.utils.data import Dataset, DataLoader
  2. from PIL import Image
  3. import torch
  4. import os
  6. class SegmentationDataset(Dataset):
  7.     def __init__(self, img_dir, mask_dir, transform=None):
  8.         self.img_dir = img_dir
  9.         self.mask_dir = mask_dir
  10.         self.transform = transform
  11.         self.images = os.listdir(img_dir)
  13.     def __len__(self):
  14.         return len(self.images)
  16.     def __getitem__(self, idx):
  17.         img_path = os.path.join(self.img_dir, self.images[idx])
  18.         mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '.png'))
  20.         image = Image.open(img_path).convert('RGB')
  21.         mask = Image.open(mask_path).convert('L')  # 转换为灰度图
  23.         if self.transform:
  24.             image, mask = self.transform(image, mask)
  26.         # 归一化并添加通道维度
  27.         image = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
  28.         mask = torch.from_numpy(np.array(mask)).long()  # 分类任务使用long类型
  30.         return image, mask
  32. # 使用示例
  33. train_transform = CustomAugmentation()  # 前文定义的数据增强
  34. train_dataset = SegmentationDataset('data/train/images', 'data/train/masks', transform=train_transform)
  35. train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4)

2. 训练脚本实现

  1. import torch.optim as optim
  2. from tqdm import tqdm
  3. import numpy as np
  5. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  7. def train_model(model, train_loader, val_loader, epochs=50):
  8.     model = model.to(device)
  9.     criterion = nn.CrossEntropyLoss()  # 多分类任务
  10.     optimizer = optim.Adam(model.parameters(), lr=1e-4)
  11.     scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3)
  13.     best_val_loss = float('inf')
  15.     for epoch in range(epochs):
  16.         model.train()
  17.         train_loss = 0
  19.         progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}')
  20.         for images, masks in progress_bar:
  21.             images, masks = images.to(device), masks.to(device)
  23.             optimizer.zero_grad()
  24.             outputs = model(images)
  25.             loss = criterion(outputs, masks)
  26.             loss.backward()
  27.             optimizer.step()
  29.             train_loss += loss.item()
  30.             progress_bar.set_postfix(loss=loss.item())
  32.         train_loss /= len(train_loader)
  34.         # 验证阶段
  35.         val_loss = validate(model, val_loader, criterion)
  36.         scheduler.step(val_loss)
  38.         print(f'Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')
  40.         if val_loss < best_val_loss:
  41.             best_val_loss = val_loss
  42.             torch.save(model.state_dict(), 'best_model.pth')
  44. def validate(model, val_loader, criterion):
  45.     model.eval()
  46.     val_loss = 0
  47.     with torch.no_grad():
  48.         for images, masks in val_loader:
  49.             images, masks = images.to(device), masks.to(device)
  50.             outputs = model(images)
  51.             loss = criterion(outputs, masks)
  52.             val_loss += loss.item()
  54.     return val_loss / len(val_loader)
  56. # 初始化模型
  57. model = UNet(n_channels=3, n_classes=2)  # 假设二分类任务
  58. # 开始训练
  59. train_model(model, train_loader, val_loader)

五、推理测试：从模型加载到结果可视化

1. 模型加载与预处理

  1. def load_model(model_path, n_channels=3, n_classes=2):
  2.     model = UNet(n_channels, n_classes)
  3.     model.load_state_dict(torch.load(model_path))
  4.     model.eval()
  5.     return model
  7. def preprocess_image(image_path, target_size=(256, 256)):
  8.     image = Image.open(image_path).convert('RGB')
  9.     if image.size != target_size:
  10.         image = image.resize(target_size, Image.BILINEAR)
  11.     transform = transforms.Compose([
  12.         transforms.ToTensor(),
  13.         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  14.     ])
  15.     return transform(image).unsqueeze(0)  # 添加batch维度

2. 推理与后处理

  1. import matplotlib.pyplot as plt
  2. import matplotlib.colors as mcolors
  4. def predict(model, image_tensor):
  5.     with torch.no_grad():
  6.         output = model(image_tensor.to(device))
  7.         pred = torch.argmax(output.squeeze(), dim=0).cpu().numpy()
  8.     return pred
  10. def visualize_prediction(image, pred_mask, original_mask=None):
  11.     plt.figure(figsize=(12, 6))
  13.     plt.subplot(1, 3, 1)
  14.     plt.imshow(image)
  15.     plt.title('Original Image')
  16.     plt.axis('off')
  18.     plt.subplot(1, 3, 2)
  19.     # 创建颜色映射（0:黑色，1:红色）
  20.     cmap = mcolors.ListedColormap(['black', 'red'])
  21.     plt.imshow(pred_mask, cmap=cmap)
  22.     plt.title('Predicted Mask')
  23.     plt.axis('off')
  25.     if original_mask is not None:
  26.         plt.subplot(1, 3, 3)
  27.         plt.imshow(original_mask, cmap=cmap)
  28.         plt.title('Original Mask')
  29.         plt.axis('off')
  31.     plt.tight_layout()
  32.     plt.show()
  34. # 使用示例
  35. model = load_model('best_model.pth')
  36. image_tensor = preprocess_image('test_image.jpg')
  37. pred_mask = predict(model, image_tensor)
  39. # 显示结果（假设有原始图像和标注）
  40. original_image = Image.open('test_image.jpg')
  41. original_mask = Image.open('test_mask.png').convert('L')
  42. visualize_prediction(original_image, pred_mask, np.array(original_mask))

六、进阶优化建议

  1. 损失函数选择：

    • 二分类任务：BCEWithLogitsLoss
    • 多分类任务：CrossEntropyLoss
    • 类别不平衡时：加权CrossEntropy或Dice Loss

  2. 模型改进：

    • 使用ResNet或EfficientNet作为编码器 backbone
    • 添加注意力机制（如CBAM）
    • 使用DeepLabv3+的空洞卷积结构

  3. 训练技巧：

    • 使用学习率预热（warmup）
    • 实现早停（Early Stopping）
    • 使用混合精度训练（AMP）

  4. 部署优化：

    • 模型量化（INT8）
    • TensorRT加速
    • ONNX格式导出

七、总结与展望

本文详细介绍了使用UNet进行图像语义分割的完整流程，包括自定义数据集制作、模型实现、训练优化和推理测试。通过实践，读者可以掌握：

  • 如何准备符合要求的语义分割数据集
  • 如何实现经典的UNet架构
  • 如何设计有效的训练策略
  • 如何进行模型推理和结果可视化

未来工作可以探索：

  • 3D医学图像分割（如使用3D UNet）
  • 实时语义分割（如使用轻量化模型）
  • 弱监督/半监督学习方法

希望本文能为深度学习初学者提供实用的入门指南，也为研究人员提供有价值的实践参考。

相关文章推荐

发表评论

最热文章

关于作者

  • 被阅读数
  • 被赞数
  • 被收藏数
活动
咨询