【PyTorch实战：图像分割】19、端到端实战：从Labelme标注到UNet部署

内容分享2天前发布

0 0 0

🌟 项目全景图

一、项目背景：三大场景实战（卫星/医学/细胞）

🔥 为什么选择图像分割？

卫星遥感：Massachusetts Buildings数据集，识别建筑物轮廓（本文主案例）医学影像：细胞分割（ISBI挑战赛同款），肿瘤边界标注工业检测：PCB板缺陷定位（附标注技巧）

📊 效果预览（卫星图像分割）

【PyTorch实战：图像分割】19、端到端实战：从Labelme标注到UNet部署
左：原始卫星图 | 中：人工标注 | 右：UNet预测

二、数据标注：从工具到格式转换（含避坑指南）

🔧 5大标注工具对比（2025最新）

工具	支持类型	效率	特色功能	适用场景
Labelme	多边形/语义分割	★★★★☆	JSON转VOC/COCO，支持插件	通用场景（本文主用）
EISeg	交互式分割	★★★★★	集成PaddleSeg预训练模型	医疗/遥感快速标注
CVAT	视频标注	★★★☆☆	在线协作，支持跟踪	视频序列标注
RITM	点击式分割	★★★☆☆	小样本快速修正	精细物体标注
自研工具	定制化需求	★★★☆☆	集成业务逻辑	企业级流水线

📝 Labelme实战流程（卫星数据为例）

安装启动：


pip install labelme
labelme data/train/images --labels labels.txt --nodata  # 关键参数：--nodata跳过原图保存

标注技巧：

快捷键：W画多边形，D删除，Ctrl+Z撤销分层标注：建筑物用红色，道路用蓝色（通过labels.txt定义颜色）批量处理：使用labelme_json_to_dataset批量转PNG掩码

格式转换代码（关键）：


# label2voc.py（节选）
import labelme
import numpy as np
from PIL import Image

def json_to_mask(json_path, output_dir, label_names):
    data = labelme.LabelFile.load_json_file(json_path)
    shapes = data['shapes']
    img_h, img_w = data['imageHeight'], data['imageWidth']
    
    mask = np.zeros((img_h, img_w), dtype=np.uint8)
    for i, shape in enumerate(shapes):
        label = shape['label']
        points = np.array(shape['points'], dtype=np.int32)
        cv2.fillPoly(mask, [points], label_names.index(label)+1)  # 背景为0
    
    Image.fromarray(mask).save(f"{output_dir}/{os.path.basename(json_path)[:-5]}.png")

🚦 避坑指南：

标注精度：多边形顶点间隔建议≤5像素（卫星图），医学图需≤2像素类别平衡：建筑物占比＜10%时，使用OHEM（在线难例挖掘）掩码检查：用plt.subplot同时显示原图+掩码，确保边界对齐

三、数据管道：从Loader到增强（附性能优化）

🛠️ 通用Dataset类（支持3大场景）


class SegmentationDataset(Dataset):
    def __init__(self, img_dir, mask_dir, transform=None, mode='satellite'):
        self.imgs = sorted(os.listdir(img_dir))
        self.masks = sorted(os.listdir(mask_dir))
        self.transform = transform
        self.mode = mode  # 区分卫星/医学/细胞的预处理差异

    def __getitem__(self, idx):
        img = Image.open(os.path.join(img_dir, self.imgs[idx])).convert('RGB')
        mask = Image.open(os.path.join(mask_dir, self.masks[idx])).convert('L')
        
        # 特殊预处理（医学图去黑边）
        if self.mode == 'medical':
            img = np.array(img)
            mask = np.array(mask)
            _, img = cv2.threshold(img, 5, 255, cv2.THRESH_TOZERO)  # 去除低亮度噪声
        
        if self.transform:
            img = self.transform['img'](img)
            mask = self.transform['mask'](mask)
        
        return img/255., np.array(mask)/255.  # 归一化到[0,1]

🌪️ 高级数据增强（卫星图专用）


import albumentations as A
from albumentations.pytorch import ToTensorV2

train_transform = A.Compose([
    A.Resize(256, 256),
    A.HorizontalFlip(p=0.7),
    A.RandomRotate90(p=0.5),
    A.GaussNoise(var_limit=(10, 50), p=0.3),  # 模拟卫星图噪声
    A.OneOf([
        A.RandomGamma(gamma_limit=(80, 120)),
        A.RandomBrightnessContrast(),
    ], p=0.6),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.5),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
], additional_targets={'mask': 'mask'})

📊 性能优化技巧：

多线程加载：DataLoader(num_workers=8, pin_memory=True)（需关闭杀毒软件）缓存机制：对医学大图（＞1024×1024），预处理后保存为.npy混合精度：AMP加速训练（卫星图batch=16时，显存占用降低40%）

四、UNet架构：从基础版到魔改（含可视化）

🧱 核心Block设计（附激活函数对比）


class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch, act='relu'):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True) if act=='relu' else nn.SiLU(),  # 医学图建议SiLU
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True) if act=='relu' else nn.SiLU()
        )
    
    def forward(self, x):
        return self.conv(x)

🌐 完整UNet结构（含跳跃连接可视化）

🚀 魔改技巧（针对不同场景）：

场景	改进点	效果提升
卫星图	加入CBAM注意力	IoU+2.3%
医学图	替换为3D UNet	Dice+1.8%
细胞分割	浅层网络（减少下采样次数）	速度×2

五、训练策略：从Loss到调度（附曲线分析）

🔍 复合损失函数（Dice+BCE+Lovász）


class MultiLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.bce = nn.BCEWithLogitsLoss()
        self.dice = DiceLoss()
        self.lovasz = LovaszLoss()  # 自定义边界敏感损失
    
    def forward(self, pred, mask):
        return 0.5*self.bce(pred, mask) + 0.3*self.dice(pred, mask) + 0.2*self.lovasz(pred, mask)

📈 训练曲线解读（卫星图25epoch）

【PyTorch实战：图像分割】19、端到端实战：从Labelme标注到UNet部署

蓝色：训练IoU（从0.65→0.89）橙色：验证IoU（稳定在0.85±0.02）阴影区：第10epoch学习率下降（1e-4→5e-5）

⚙️ 超参数调优表（最佳组合）

参数	卫星图	医学图	细胞图
学习率	1e-4（AdamW）	5e-5（RAdam）	2e-4（SGD）
Batch Size	16（8GB显存）	8（32GB显存）	32（小图）
Epochs	50	100	30
权重初始化	HeUniform	Kaiming	Xavier

六、评估与预测：从指标到后处理

📌 5大评估指标（附代码）


def calculate_metrics(pred, mask, num_classes=1):
    pred = (pred > 0.5).float()
    metrics = {}
    
    # IoU
    intersection = (pred * mask).sum()
    union = (pred + mask - pred*mask).sum()
    metrics['iou'] = (intersection + 1e-6) / (union + 1e-6)
    
    # Dice
    metrics['dice'] = 2 * intersection / (pred.sum() + mask.sum() + 1e-6)
    
    # 边界误差（医学图专用）
    if num_classes == 1:
        contours_pred, _ = cv2.findContours(pred.numpy().astype(np.uint8), 1, 2)
        contours_mask, _ = cv2.findContours(mask.numpy().astype(np.uint8), 1, 2)
        if len(contours_pred) and len(contours_mask):
            dist = cv2.pointPolygonTest(contours_mask[0], contours_pred[0][0][0], True)
            metrics['border_err'] = abs(dist)
    
    return metrics

🛠️ 后处理三板斧：

形态学操作：cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel=3x3)（消除小孔）连通域过滤：保留面积＞100像素的区域（卫星图建筑物最小尺寸）边界优化：skimage.morphology.binary_border_mask（医学图专用）

🚀 推理加速技巧：

TensorRT转换：FP16精度推理速度提升3倍（RTX3090）滑动窗口推理：对大图（＞1024×1024）分块预测，重叠率20%缓存机制：保存中间特征图，加速连续帧处理（视频分割）

七、可视化：从训练到部署（附交互式工具）

📊 训练过程可视化（TensorBoard）


# 日志记录
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir='runs/satellite_seg')

# 每epoch记录
writer.add_scalar('Train/Loss', loss, epoch)
writer.add_images('Visualization', torch.cat([img, mask, pred], dim=2), epoch)

🎨 交互式预测工具（Streamlit）


import streamlit as st
from PIL import Image

st.title("卫星图像建筑物分割")
uploaded_file = st.file_uploader("选择图片", type=["jpg", "png", "tif"])

if uploaded_file:
    img = Image.open(uploaded_file)
    st.image(img, caption="原始图像", width=400)
    
    # 预测
    pred_mask = predict(img)
    st.image(pred_mask, caption="预测掩码", width=400)
    
    # 对比图
    st.image(np.hstack([np.array(img), pred_mask*255]), caption="原图vs预测", width=800)

八、部署实战：从模型到API（3种方案）

🔧 轻量化部署对比

方案	卫星图模型大小	推理时间（RTX3060）	特色
PyTorch原生	128MB	12ms/张	支持动态形状
ONNX	65MB	8ms/张	跨平台部署
TensorRT	32MB	4ms/张	工业级高性能推理
Flask API	–	50ms/请求（含IO）	快速原型开发

🚀 Flask API代码（核心）


from flask import Flask, request, jsonify
import torch
import cv2

app = Flask(__name__)
model = torch.load('best_model.pth').eval().cuda()

@app.route('/predict', methods=['POST'])
def predict():
    file = request.files['image']
    img = cv2.imdecode(np.fromstring(file.read(), np.uint8), cv2.IMREAD_COLOR)
    img = preprocess(img)  # 预处理函数
    
    with torch.no_grad():
        pred = model(img.cuda()).sigmoid()
    mask = (pred > 0.5).cpu().numpy()
    
    # 返回掩码PNG
    _, buffer = cv2.imencode('.png', mask[0]*255)
    return buffer.tobytes(), 200, {'Content-Type': 'image/png'}

九、进阶优化：从数据到模型（附研究方向）

🧠 前沿技术融合：

CLIP引导分割：用CLIP文本提示（如”building”）优化小样本标注（CVPR 2025最新）扩散模型：生成高质量合成数据（解决医学数据稀缺问题）动态UNet：根据图像复杂度自动调整下采样层数（ECCV 2024）

🔍 常见问题解决方案：

问题现象	可能原因	解决方案
训练IoU震荡	数据增强过强	降低旋转角度至10°以内
边界模糊	缺少边界监督	加入Canny边缘损失
小物体漏检	下采样次数过多	改用Strided Conv替代MaxPool
显存不足	Batch Size过大	梯度累积（accumulate_grad_batches=2）

十、总结：从入门到精通的学习路径

📚 推荐学习资源：

论文：《U-Net: Convolutional Networks for Biomedical Image Segmentation》数据集：Massachusetts Buildings（卫星）、ISBI 2012（细胞）、BTCV（医学）工具：Labelme官方文档、albumentations增强指南