图像分割实战：UNet 语义分割入门到训练

图像分割实战：UNet 语义分割入门到训练

分类告诉你"这是什么"，检测告诉你"它在哪"，分割告诉你"哪些像素属于它"。医学影像、遥感、抠图、缺陷检测都靠分割。这篇用 UNet 从数据到训练到评估完整走一遍，代码真实可跑。

前言

做完分类和检测项目后，客户又提了个新需求：把产品图片里的"缺陷区域"精确圈出来——检测框不够用，要像素级。这就是语义分割：给图片里每个像素一个类别标签。

语义分割的经典模型是 UNet（医学影像分割的标配，也是很多分割任务的最佳起点）。这篇从数据格式到训练评估一次讲透，用 segmentation-models-pytorch 库把代码量压到最少。

环境准备（先装依赖）：

pip install torch torchvision pillow numpy segmentation-models-pytorch

一、语义分割是什么

任务	输出	例子
图像分类	一个标签	"这是一只猫"
目标检测	若干框	"猫在这里（坐标框）"
语义分割	每个像素一个标签	"这些像素是猫，那些是背景"

应用场景：医学影像（肿瘤/器官分割）、遥感（建筑/道路提取）、抠图、缺陷检测、自动驾驶（道路/行人分割）。

二、UNet 架构（一句话记住）

编码器（不断下采样，学特征）
   ↓
解码器（不断上采样，恢复尺寸）
   ↓
跳跃连接（编码器和解码器对应层相连，保留细节）

形状像字母 U，所以叫 UNet。三个特点：

1. 跳跃连接保留高分辨率细节（分割边界的细节全靠它）；
2. 在小数据集上表现好（医学图像数据集通常很小）；
3. 配合预训练编码器（ResNet 等）效果更稳。

三、数据格式（分割和分类最大的不同）

分割数据集是图片 + 掩码（mask）成对：

dataset/
├── images/xxx.jpg        # 原图
└── masks/xxx.png         # 掩码：和原图同尺寸，像素值=类别号（0 背景 / 1 前景）

• 掩码是单通道灰度图，不是彩色图；
• 二分类任务：背景=0，目标=1；
• 公开数据集：VOC2012、Cityscapes（街景）、医学分割数据集（如心脏/细胞分割）。

四、完整训练代码（真实可跑）

pip install segmentation-models-pytorch

import torch
import segmentation_models_pytorch as smp
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np

# 1. 数据集类：读图 + 读掩码
class SegDataset(Dataset):
    def __init__(self, img_dir, mask_dir, size=256):
        self.images = sorted(img_dir.iterdir())
        self.masks = sorted(mask_dir.iterdir())
        self.size = size
    def __len__(self):
        return len(self.images)
    def __getitem__(self, i):
        img = Image.open(self.images[i]).convert("RGB").resize((self.size, self.size))
        msk = Image.open(self.masks[i]).convert("L").resize((self.size, self.size))
        img = transforms.ToTensor()(img)
        msk = torch.from_numpy(np.array(msk, dtype=np.float32) / 255.0).unsqueeze(0)
        return img, msk

# 2. 模型：UNet + 预训练 ResNet34 编码器
model = smp.Unet(
    encoder_name="resnet34",
    encoder_weights="imagenet",
    in_channels=3,
    classes=1,                  # 二分类（前景/背景）
)
model = model.to("cuda" if torch.cuda.is_available() else "cpu")

# 3. 损失 + 优化器（DiceLoss 适合分割，类别不平衡时比 BCE 稳）
loss_fn = smp.losses.DiceLoss(mode="binary")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 4. 训练（标准循环）
train_ds = SegDataset(img_dir=..., mask_dir=...)
loader = DataLoader(train_ds, batch_size=8, shuffle=True)
for epoch in range(30):
    running = 0.0
    for images, masks in loader:
        images, masks = images.to("cuda"), masks.to("cuda")
        optimizer.zero_grad()
        pred = model(images)
        loss = loss_fn(pred, masks)
        loss.backward()
        optimizer.step()
        running += loss.item()
    print(f"Epoch {epoch+1}/30  loss: {running/len(loader):.3f}")

五、评估指标：IoU（分割的"准确率"）

分割不用准确率（背景像素太多，虚高），用 IoU（交并比）：预测和真实区域的交集 ÷ 并集，0~1，越高越好。

def iou(pred, mask, threshold=0.5):
    pred = (torch.sigmoid(pred) > threshold).float()
    inter = (pred * mask).sum()
    union = (pred + mask).gt(0).sum()
    return (inter / (union + 1e-6)).item()

经验值：二分类分割 IoU 0.7 以上算不错，0.85+ 很优秀。

运行验证：运行训练代码，期望看到每个 epoch 的 loss 持续下降（0.6 → 0.3 → 0.2 → …），训练完用 iou() 函数在验证集上计算，二分类 IoU 达到 0.7 以上。如果 loss 不降或 IoU 很低，先查掩码数据对不对（单通道、值 0/1），再查数据集路径。

六、部署

训练完 torch.save(model.state_dict(), "unet.pth")，导出 ONNX 后用 《PyTorch 模型导出 ONNX》的 FastAPI 套路部署成"传图返回掩码"的接口：

torch.onnx.export(model.eval(), dummy, "unet.onnx",
                  input_names=["image"], output_names=["mask"])

七、常见坑

坑	解决
掩码读成彩色图	掩码用 convert("L") 单通道读；标签值 0/1 别缩放错
类别不平衡（目标很小）	用 DiceLoss/加权 BCE；或对目标区域做增强
IoU 一直很低	检查掩码和原图是否对得上（数据错位最常见）
显存不足	batch 降到 2~4、尺寸 256 起步
边缘锯齿	后处理加条件随机场（CRF）或换更大模型

八、总结

• 一句话：UNet + 预训练编码器 + DiceLoss，就是分割任务的黄金起手式；
• 接单角度：医学图像分割、遥感分割、工业缺陷分割都是付费场景，这套代码换数据集就能交付；
• 组合拳：《PyTorch 模型导出 ONNX》+ 《Docker 容器化实战》把它包装成服务，就是完整交付物。

觉得有用点个赞收藏，有问题评论区见，看到都会回。