UNet人物抠图实战:从数据集制作到模型训练与推理全流程

📅 发布时间:2026/9/28 7:29:33
UNet人物抠图实战:从数据集制作到模型训练与推理全流程
简介这份资源面向计算机视觉入门与进阶开发者围绕UNet语义分割的二分类任务提供一套可直接运行的实战工程用于实现人物抠图。内容聚焦单通道输出方案网络输出经Sigmoid归一化到[0,1]后与0/1标签图计算交叉熵损失反向传播更新权重最终使预测逼近目标掩码帮助读者理解二分类分割的训练与推理流程。压缩包共3430个文件以3404个png图像数据为主另含11个py训练与推理脚本、7个pyc缓存、Dockerfile、sh启动脚本、README.md说明文档及LICENSE等整体约987.41MB目录结构便于按数据、代码、配置分层查阅。目前已有2816人学习下载适合希望掌握UNet抠图全流程、复用训练脚本与数据组织方式的读者参考实践。1. UNet人物抠图为什么语义分割比传统抠图更值得投入用 GIMP 手动抠一张人像边缘发丝要放大到 800% 一点点描半小时起步换一张背景复杂的时间直接翻倍。而 UNet 语义分割做人物抠图训练完之后推理一张 512×512 的图只要几十毫秒批量处理几百张也不用重新调参数。这就是为什么越来越多的从业者从「钢笔工具 通道抠图」转向「UNet 训练自己的数据集」这条路。UNet 最初是为医学影像分割设计的编码器-解码器加跳跃连接的结构让它在小样本、边缘精细的任务上表现很稳。人物抠图本质上是一个二分类语义分割问题每个像素判断是「人」还是「背景」。和实例分割不同语义分割不区分「张三」和「李四」只关心「是不是人」这恰好匹配抠图的需求。这篇文章会从数据集制作、UNet 代码搭建、训练调参、推理导出一路讲到踩坑排查目标是让你看完能自己跑通一套人物抠图流程而不是停留在「跑一个 UNet 网络」的 Demo 层面。适合谁看有 Python 和 PyTorch 基础、想做语义分割落地但没完整走过一遍的工程师被传统抠图效率折磨、想用模型批量处理人像的从业者以及想理解「语义分割数据集制作」和「UNet 模型改进」之间关系的人。下面按「先立住原理、再动手复现、最后避坑」的顺序展开。2. UNet 结构与人物抠图任务的匹配逻辑2.1 编码器-解码器为什么适合边缘精细的二分类UNet 的核心结构分三部分下采样编码器、上采样解码器、跳跃连接。编码器每经过一次下采样特征图尺寸减半、通道数翻倍逐步提取从边缘纹理到语义区域的抽象特征。解码器做相反操作逐步恢复空间分辨率。关键在于跳跃连接——它把编码器同层的高分辨率特征直接拼接到解码器对应层让浅层的边缘细节不会在深层抽象中丢失。人物抠图最难的恰恰是边缘头发丝、手指缝隙、半透明衣物。纯编码器-解码器结构比如 SegNet在下采样过程中会丢掉这些细节解码器再怎么上采样也补不回来。跳跃连接相当于给解码器开了一条「后悔药」通道浅层特征直接送过来边缘精度能提升一个档次。另一个匹配点是 UNet 对小数据集友好。医学影像标注成本极高UNet 设计之初就考虑了少量样本下的泛化能力。人物抠图的数据集标注同样费时UNet 这种结构在几百到几千张图上就能收敛出可用结果不像 Transformer 类分割模型动辄需要上万张标注。2.2 从输入到 mask一次前向传播里发生了什么假设输入是一张 3×256×256 的 RGB 人像。经过编码器四次下采样后特征图变成 512×16×16此时每个像素的感受野覆盖了原图大部分区域语义信息足够判断「这块区域是不是人」。解码器从 16×16 逐步上采样回 256×256每次上采样后与编码器对应层特征拼接最后经过一个 1×1 卷积输出 1×256×256 的 logits再经 Sigmoid 得到每个像素属于「人」的概率。训练时的损失函数通常用 BCEWithLogitsLoss 或 Dice Loss。BCE 对每个像素独立计算二分类交叉熵简单直接Dice Loss 直接优化预测 mask 和真实 mask 的重叠度对前景背景不平衡更鲁棒。人物抠图里背景像素通常远多于人物像素我一般用 BCE Dice 加权组合权重 0.5:0.5 起步根据验证集 IoU 再调。推理阶段Sigmoid 输出大于 0.5 的像素判为人得到二值 mask。把 mask 与原图做逐像素乘法背景置黑或置透明就完成了抠图。如果要输出 PNG 带透明通道把 mask 作为 alpha 通道写入即可。2.3 和实例分割、传统抠图的边界在哪语义分割输出的是「类别 mask」所有人像素都是同一个标签。实例分割比如 YOLO 系列里的 mask 分支会区分不同个体输出「人 A」「人 B」各自的 mask。抠图只需要「人 vs 背景」语义分割足够。但如果场景里有多人且需要分别抠出就得上实例分割。传统抠图工具GIMP、PS依赖人工交互或颜色通道先验单张精度可以很高但无法批量、无法自动化。UNet 语义分割的优势在于训练一次之后可以批量推理代价是需要标注数据。实际项目中常见的混合方案是用 UNet 做粗抠边缘再用传统算法如导向滤波、Matting做精修兼顾效率和精度。3. 人物抠图数据集制作从原始图片到可训练 mask3.1 数据采集与标注工具选择数据集质量直接决定模型上限。人物抠图的数据来源常见有几类公开人像数据集如 Supervisely Person、P3M-10k、自己拍摄或爬取的人像、以及业务场景里的真实图片。不管来源如何标注格式最终要统一成「原图 二值 mask」的配对。标注工具推荐 labelme 或 CVAT。labelme 轻量、Python 生态好适合小规模标注CVAT 支持多人协作和视频标注适合团队。标注时用多边形勾勒人物轮廓导出为 JSON再转成二值 PNG mask。注意 mask 的像素值人物区域为 255背景为 0不要用 1 和 0否则可视化时几乎全黑排查问题很不方便。提示标注边缘时尽量贴合人物实际轮廓发丝区域可以适当放宽但不要大面积包含背景。标注质量差的数据比没有数据更糟糕模型会学到错误的边缘模式。3.2 用脚本把 labelme 标注转成训练用 masklabelme 导出的 JSON 里包含多边形点坐标和标签。下面脚本把 JSON 转成二值 mask并统一尺寸。import json import numpy as np import cv2 import os from glob import glob def labelme_to_mask(json_path, output_dir, target_size(256, 256)): with open(json_path, r, encodingutf-8) as f: data json.load(f) img_h data[imageHeight] img_w data[imageWidth] mask np.zeros((img_h, img_w), dtypenp.uint8) for shape in data[shapes]: # 只处理标签为 person 的多边形 if shape[label] ! person: continue points np.array(shape[points], dtypenp.int32) cv2.fillPoly(mask, [points], 255) # 统一尺寸最近邻插值保持二值特性 mask_resized cv2.resize(mask, target_size, interpolationcv2.INTER_NEAREST) base_name os.path.splitext(os.path.basename(json_path))[0] out_path os.path.join(output_dir, base_name .png) cv2.imwrite(out_path, mask_resized) return out_path if __name__ __main__: json_files glob(annotations/*.json) os.makedirs(masks, exist_okTrue) for jf in json_files: labelme_to_mask(jf, masks) print(f转换完成共 {len(json_files)} 张)逻辑说明fillPoly把多边形内部填充为 255多个多边形会叠加。cv2.resize用INTER_NEAREST而不是默认的双线性插值因为 mask 是二值的双线性会产生 0-255 之间的灰度值破坏二值特性。target_size根据你的训练输入尺寸设定常见 256×256 或 512×512。参数说明shape[label]要和标注时用的标签名一致如果标注用了中文标签这里要对应修改。output_dir需要提前创建。如果一张图有多个人物所有 person 多边形都会填充到同一张 mask符合语义分割的定义。3.3 数据增强与训练集/验证集划分人物抠图数据增强要小心水平翻转、随机裁剪、亮度对比度调整是安全的垂直翻转要谨慎因为人物通常不会倒立翻转后可能引入不合理姿态旋转角度不要太大超过 15 度人物姿态会失真。颜色抖动对抠图任务帮助有限因为模型学的是形状和边缘不是颜色分布。import random import cv2 import numpy as np def augment_pair(image, mask): # 水平翻转 if random.random() 0.5: image cv2.flip(image, 1) mask cv2.flip(mask, 1) # 随机裁剪并 resize h, w image.shape[:2] scale random.uniform(0.8, 1.2) new_h, new_w int(h * scale), int(w * scale) image cv2.resize(image, (new_w, new_h)) mask cv2.resize(mask, (new_w, new_h), interpolationcv2.INTER_NEAREST) # 裁剪回原尺寸 if new_h h: top random.randint(0, new_h - h) image image[top:toph, :w] mask mask[top:toph, :w] else: pad_h h - new_h image cv2.copyMakeBorder(image, 0, pad_h, 0, 0, cv2.BORDER_CONSTANT, value0) mask cv2.copyMakeBorder(mask, 0, pad_h, 0, 0, cv2.BORDER_CONSTANT, value0) return image, mask逻辑说明图像和 mask 必须同步做相同的几何变换否则标签就错了。裁剪时如果缩放后尺寸大于原图随机裁一块小于原图则补零。mask 的 resize 始终用最近邻。划分比例一般 8:1:1 或 7:2:1。验证集用来监控过拟合测试集只在最后评估用。如果数据量少于 500 张建议用 5 折交叉验证而不是固定划分否则验证集波动会很大。4. UNet 模型搭建与训练代码逐段拆解4.1 双卷积块与下采样/上采样模块实现UNet 的基本单元是「双卷积块」两次 3×3 卷积 BN ReLU。下采样用 2×2 最大池化上采样用转置卷积或双线性插值 卷积。下面是一个可直接用的 PyTorch 实现。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool_conv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.pool_conv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch, bilinearTrue): super().__init__() if bilinear: self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv DoubleConv(in_ch, out_ch) else: self.up nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 self.up(x1) # 处理尺寸不匹配奇数尺寸输入时可能出现 diff_h x2.size(2) - x1.size(2) diff_w x2.size(3) - x1.size(3) x1 F.pad(x1, [diff_w // 2, diff_w - diff_w // 2, diff_h // 2, diff_h - diff_h // 2]) x torch.cat([x2, x1], dim1) return self.conv(x)逻辑说明DoubleConv里biasFalse是因为后面接了 BNBN 会减去均值bias 冗余。Up里bilinearTrue时用双线性插值上采样比转置卷积更不容易产生棋盘伪影这是我踩过坑之后的默认选择。F.pad处理编码器和解码器特征图尺寸不一致的情况输入尺寸不是 2 的整数次幂时会出现。参数说明in_ch和out_ch是通道数第一层通常 3→64之后 64→128→256→512。bilinear建议保持 True除非你有明确理由用转置卷积。4.2 完整 UNet 组装与输出层设计class UNet(nn.Module): def __init__(self, n_channels3, n_classes1, bilinearTrue): super().__init__() self.n_channels n_channels self.n_classes n_classes self.bilinear bilinear self.inc DoubleConv(n_channels, 64) self.down1 Down(64, 128) self.down2 Down(128, 256) self.down3 Down(256, 512) factor 2 if bilinear else 1 self.down4 Down(512, 1024 // factor) self.up1 Up(1024, 512 // factor, bilinear) self.up2 Up(512, 256 // factor, bilinear) self.up3 Up(256, 128 // factor, bilinear) self.up4 Up(128, 64, bilinear) self.outc nn.Conv2d(64, n_classes, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.outc(x) # 输出 logits不加 Sigmoid逻辑说明输出层用 1×1 卷积把 64 通道压成n_classes。人物抠图是二分类n_classes1配合BCEWithLogitsLoss使用。不要在模型里加 Sigmoid因为BCEWithLogitsLoss内部做了 log-sum-exp 技巧数值更稳定。推理时再手动加 Sigmoid。参数说明n_channels3对应 RGB 输入。如果输入是 RGBA 或灰度改这个值。bilinear控制上采样方式和Up模块保持一致。4.3 训练循环、损失函数与学习率策略import torch.optim as optim from torch.utils.data import DataLoader, Dataset import cv2 import os class PersonDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size256, augmentFalse): self.img_dir img_dir self.mask_dir mask_dir self.img_size img_size self.augment augment self.names [f for f in os.listdir(img_dir) if f.endswith((.jpg, .png))] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img cv2.imread(os.path.join(self.img_dir, name)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask_name os.path.splitext(name)[0] .png mask cv2.imread(os.path.join(self.mask_dir, mask_name), 0) img cv2.resize(img, (self.img_size, self.img_size)) mask cv2.resize(mask, (self.img_size, self.img_size), interpolationcv2.INTER_NEAREST) if self.augment: img, mask augment_pair(img, mask) img img.astype(np.float32) / 255.0 img np.transpose(img, (2, 0, 1)) mask (mask 127).astype(np.float32) mask np.expand_dims(mask, axis0) return torch.from_numpy(img), torch.from_numpy(mask) def dice_loss(pred, target, smooth1.0): pred torch.sigmoid(pred) pred pred.view(-1) target target.view(-1) intersection (pred * target).sum() return 1 - (2. * intersection smooth) / (pred.sum() target.sum() smooth) # 训练主循环 device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(n_channels3, n_classes1).to(device) optimizer optim.Adam(model.parameters(), lr1e-3) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, patience5, factor0.5) bce nn.BCEWithLogitsLoss() train_dataset PersonDataset(data/train/images, data/train/masks, augmentTrue) train_loader DataLoader(train_dataset, batch_size8, shuffleTrue, num_workers4) for epoch in range(100): model.train() epoch_loss 0 for imgs, masks in train_loader: imgs, masks imgs.to(device), masks.to(device) optimizer.zero_grad() outputs model(imgs) loss 0.5 * bce(outputs, masks) 0.5 * dice_loss(outputs, masks) loss.backward() optimizer.step() epoch_loss loss.item() avg_loss epoch_loss / len(train_loader) scheduler.step(avg_loss) print(fEpoch {epoch1}, Loss: {avg_loss:.4f}, LR: {optimizer.param_groups[0][lr]:.6f})逻辑说明损失用 BCE 和 Dice 各占 0.5。BCE 提供稳定的逐像素梯度Dice 直接优化重叠度两者互补。ReduceLROnPlateau在验证 loss 不下降时自动降学习率patience5 表示连续 5 个 epoch 没改善就降。dice_loss里先对 pred 做 Sigmoid因为模型输出是 logits。参数说明batch_size8在 8GB 显存下跑 256×256 输入比较稳显存不够降到 4。lr1e-3是 Adam 的常用起点如果 loss 震荡明显降到 5e-4。num_workers根据 CPU 核数设一般 4 或 8。4.4 验证指标与模型保存策略训练过程中要监控验证集的 IoU 和 Dice 系数。IoU 交集 / 并集Dice 2×交集 / (预测和 真实和)。人物抠图里 IoU 到 0.85 以上基本可用0.9 以上算不错。def evaluate(model, loader, device): model.eval() iou_sum, dice_sum, count 0, 0, 0 with torch.no_grad(): for imgs, masks in loader: imgs, masks imgs.to(device), masks.to(device) outputs torch.sigmoid(model(imgs)) preds (outputs 0.5).float() intersection (preds * masks).sum(dim(1,2,3)) union preds.sum(dim(1,2,3)) masks.sum(dim(1,2,3)) - intersection iou (intersection 1e-6) / (union 1e-6) dice (2 * intersection 1e-6) / (preds.sum(dim(1,2,3)) masks.sum(dim(1,2,3)) 1e-6) iou_sum iou.sum().item() dice_sum dice.sum().item() count imgs.size(0) return iou_sum / count, dice_sum / count保存策略不要只保存最后一个 epoch。我一般保存验证 IoU 最高的模型权重同时每 10 个 epoch 存一个 checkpoint 方便回溯。如果训练 loss 持续下降但验证 IoU 停滞或下降说明过拟合加数据增强或加 Dropout。5. 推理、后处理与踩坑排查5.1 单张图片推理与透明 PNG 导出def inference_single(model, img_path, device, img_size256): model.eval() img cv2.imread(img_path) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w img_rgb.shape[:2] img_resized cv2.resize(img_rgb, (img_size, img_size)) img_tensor img_resized.astype(np.float32) / 255.0 img_tensor np.transpose(img_tensor, (2, 0, 1)) img_tensor torch.from_numpy(img_tensor).unsqueeze(0).to(device) with torch.no_grad(): output torch.sigmoid(model(img_tensor)) mask output.squeeze().cpu().numpy() mask cv2.resize(mask, (w, h)) mask_binary (mask 0.5).astype(np.uint8) * 255 # 导出带透明通道的 PNG rgba cv2.cvtColor(img, cv2.COLOR_BGR2BGRA) rgba[:, :, 3] mask_binary cv2.imwrite(output_cutout.png, rgba) return mask_binary逻辑说明推理时把原图 resize 到训练尺寸预测完再把 mask resize 回原图尺寸。mask 0.5得到二值 mask作为 alpha 通道写入 BGRA 图像。如果边缘有锯齿可以对 mask 做一次高斯模糊再二值化或者用导向滤波精修。参数说明img_size必须和训练时一致。0.5是阈值如果发现前景偏少可以降到 0.4偏多升到 0.6根据验证集表现调。5.2 避坑人物抠图训练中最容易翻车的 5 个点现象一训练 loss 正常下降但推理出来全是黑或全是白。原因通常是 mask 像素值不是 0/255 而是 0/1或者标签反了人标成 0背景标成 1。解决用np.unique(mask)检查 mask 取值确保人物区域是 255。如果反了在 Dataset 里做mask 255 - mask。现象二边缘发丝区域糊成一团IoU 卡在 0.8 上不去。原因是下采样丢细节或者训练尺寸太小。解决把输入从 256 提到 512或者在损失里给边缘区域加权。我一般用边界带权重对 mask 做形态学膨胀和腐蚀差值就是边缘区域在 loss 里给这部分乘 23 倍权重。现象三验证集 IoU 比训练集低很多过拟合明显。数据量太少或增强不够。解决加水平翻转、随机裁剪、亮度调整如果还不行在编码器加 Dropout2dp0.20.3或者用预训练编码器如 ResNet34替换 UNet 原生编码器。现象四推理时显存溢出。输入尺寸太大或 batch 太大。解决推理时 batch_size 设 1用torch.no_grad()包住必要时把图分块推理再拼接。另外检查是否有中间变量没释放比如在循环里累积了 tensor。现象五换一批新数据推理效果骤降。域偏移问题。训练集和推理数据的背景、光照、人物风格差异大。解决在新数据上做少量微调fine-tune只训练解码器部分学习率设小一点1e-4几十张图就能拉回来。5.3 模型导出与批量处理脚本训练完的模型可以导出为 ONNX方便部署到不同环境。dummy_input torch.randn(1, 3, 256, 256).to(device) torch.onnx.export( model, dummy_input, unet_person.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version11 )批量处理时把推理函数套一个for循环遍历文件夹即可。注意每张图推理完及时释放中间变量避免显存累积。如果图片数量大可以用DataLoader做批量推理比单张循环快 35 倍。6. 把 UNet 抠图推到可用的几个进阶技巧训练出一个能跑的 UNet 只是起点真正到业务可用还有一段路。分享几个我实际项目里验证过的技巧。用预训练编码器替换原生编码器。UNet 原生编码器从零训练小数据集上收敛慢且精度有限。把编码器换成 ResNet34 或 EfficientNet-B0 的预训练权重解码器保持随机初始化训练时编码器用更小的学习率1e-4解码器用 1e-3。这样 IoU 通常能涨 35 个点收敛 epoch 数减少一半。代价是模型参数量增加推理速度略降。边缘精修用 Matting 而不是继续堆 UNet。UNet 输出的 mask 在发丝区域始终偏硬。与其把 UNet 改得越来越复杂不如在 UNet 粗 mask 基础上跑一个轻量 Matting 网络如 Guided Filter 或 FBA Matting输入是原图 粗 mask输出是精细 alpha。这个组合在多个项目里把边缘质量从「能用」推到「可交付」。多尺度推理提升小目标。人物在图中占比较小时固定尺寸推理会丢细节。做法是原图、1.5 倍放大、0.75 倍缩小各推理一次把三个 mask 平均后再二值化。代价是推理时间 ×3但小人物抠图质量提升明显。阈值不要固定 0.5。根据业务需求调需要保留更多前景如后续还要精修就降到 0.4需要干净背景如直接合成就升到 0.6。在验证集上画 Precision-Recall 曲线选 F1 最高的阈值。技巧适用场景预期收益代价预训练编码器数据少于 2000 张IoU 3~5参数量增加Matting 精修发丝/半透明边缘边缘质量显著提升多一个模型多尺度推理小人物占比小目标 IoU 5~8推理 ×3动态阈值不同业务需求匹配业务指标需验证集调参最后说一个习惯每次训练完我一定会在验证集里挑 10 张最差的样本可视化出来看。IoU 数字会骗人但眼睛不会。很多时候模型在某个特定背景比如树林、栏杆上集体翻车看数字发现不了看图一眼就清楚。这个习惯帮我省了很多盲目调参的时间。希望帮到你。本文还有配套的精品资源点击获取