尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

UNet人物抠图实战:从语义分割到透明底PNG全流程

UNet人物抠图实战:从语义分割到透明底PNG全流程 简介这份资源面向计算机视觉入门与进阶开发者聚焦UNet在图像二分类分割中的实战应用核心场景是人物抠图。内容围绕单通道输出方案展开网络输出形状为[batch_size,1,height,width]经Sigmoid归一化到[0,1]后与仅含0、1的标签图计算交叉熵损失反向传播更新权重使预测逐步逼近真实掩码帮助读者理解二分类分割的训练逻辑与实现细节。资源包共3430个文件以3404张png图像数据为主体辅以11个py训练与推理脚本、7个pyc缓存文件以及Dockerfile、sh、md、txt等环境配置与说明文档压缩包约987.41MB目录结构便于按数据、代码、配置分层查阅。目前已有2816人学习下载适合希望掌握UNet抠图流程、复现二分类分割实验并积累工程排错经验的读者参考。1. UNet人物抠图从语义分割到像素级掩膜的落地路径人像抠图这件事很多团队第一反应是上商业API或者PS批处理但一旦遇到批量、定制化、私有化部署的需求这些方案要么贵要么不可控。UNet语义分割做人物抠图本质是把“人”当成一个二分类目标让网络输出每个像素属于前景的概率最终得到一张二值掩膜。它不需要绿幕不依赖深度相机一张普通RGB图就能跑。适合谁适合手里有几百到几千张标注图、想自己训一个抠图模型、又不想从头设计网络结构的工程师。热搜里“unet训练自己的数据集”“语义分割数据集制作”这些词恰恰说明大家卡在的不是网络本身而是数据怎么标、怎么喂、怎么把输出变成能用的透明底PNG。这篇就把这条链路拆开从数据准备到推理后处理每一步都给可复现的命令和参数。2. 数据准备与标注把人物抠图变成二分类语义分割任务2.1 为什么人物抠图可以当成二分类语义分割语义分割的输出是每个像素的类别标签。人物抠图只需要区分“人”和“背景”两类所以输出通道数设为2背景人或者用1通道加Sigmoid做二分类。常见做法是后者最后一层用1x1卷积把特征图压到1通道接Sigmoid输出值在0到1之间大于0.5判为人小于0.5判为背景。这样做的好处是显存占用小训练时用BCEWithLogitsLoss或DiceLoss都方便。热搜里“yolo26 中实例分割与语义分割的区别”其实点到了一个关键实例分割要区分不同的人语义分割只关心“是不是人”。抠图场景下如果画面里只有一个人语义分割足够如果多人且需要分别抠出那得走实例分割路线。本文锁在UNet语义分割所以默认单人或多人的前景合并成一张掩膜。2.2 标注工具选型与标注规范标注工具用Labelme或CVAT都行。Labelme安装简单pip install labelme就能跑。标注时用多边形沿人物轮廓点一圈类别名统一写person。注意几个细节头发丝、半透明衣物、手指缝隙这些区域多边形不要偷懒画成直线否则训练出来的模型边缘会像狗啃。我一般要求标注员在发丝区域至少每5个像素点一个顶点。标注完成后每张图对应一个JSON文件里面存了多边形坐标。接下来要转成训练用的掩膜图。# 安装labelme pip install labelme # 启动标注 labelme --autosave --nodata--autosave让标注自动保存--nodata不把图片base64塞进JSON减小文件体积。标注完的目录结构一般是images/放原图jsons/放JSON。2.3 从Labelme JSON生成二值掩膜Labelme自带的labelme_json_to_dataset只能生成可视化结果不能直接生成训练用的PNG掩膜。我一般写个脚本批量转。import json import numpy as np import cv2 import os def json_to_mask(json_path, output_path): with open(json_path, r, encodingutf-8) as f: data json.load(f) h, w data[imageHeight], data[imageWidth] mask np.zeros((h, w), dtypenp.uint8) for shape in data[shapes]: if shape[label] ! person: continue points np.array(shape[points], dtypenp.int32) cv2.fillPoly(mask, [points], 1) cv2.imwrite(output_path, mask * 255) # 批量处理 json_dir jsons mask_dir masks os.makedirs(mask_dir, exist_okTrue) for name in os.listdir(json_dir): if name.endswith(.json): json_to_mask( os.path.join(json_dir, name), os.path.join(mask_dir, name.replace(.json, .png)) )逻辑说明fillPoly把多边形内部填成1最后乘255存成灰度图。参数上dtypenp.uint8保证掩膜是8位训练时读进来除以255就变成0/1。注意shape[label]要和你标注时写的类别名完全一致大小写敏感。如果一张图有多个人多个多边形会依次填充到同一张掩膜上这正是语义分割要的合并效果。2.4 数据集划分与增强策略划分比例一般7:2:1或8:1:1。人物抠图的数据增强要小心水平翻转没问题但垂直翻转和大幅旋转会让人物倒立不符合真实分布。我常用的增强是随机缩放0.8到1.2、随机裁剪至少保留人物主体、亮度对比度微调。颜色抖动幅度不要太大否则肤色偏移会让模型把某些背景误判成人。用Albumentations库比较顺手。import albumentations as A train_transform A.Compose([ A.HorizontalFlip(p0.5), A.RandomScale(scale_limit0.2, p0.5), A.RandomCrop(height512, width512, p1.0), A.ColorJitter(brightness0.1, contrast0.1, saturation0.1, p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ])RandomCrop固定到512x512是因为UNet下采样5次尺寸最好是32的倍数。Normalize用的ImageNet均值方差如果你从零训也可以换成自己数据集的统计值但用预训练权重时必须保持一致。3. UNet网络搭建与训练从零写一个能跑通的版本3.1 UNet结构拆解与关键参数UNet分编码器、瓶颈、解码器三部分。编码器每层两个3x3卷积加ReLU然后2x2最大池化下采样。解码器每层先上采样转置卷积或双线性插值再和编码器对应层特征拼接然后两个3x3卷积。最后1x1卷积输出1通道。关键参数初始特征数base_channels一般设32或64设大了显存爆炸设小了欠拟合。深度方面输入512x512时下采样5次到16x16再深就失去空间信息了。热搜里“unet模型改进”很多人换注意力模块我的建议是先把基础版跑通再考虑在跳跃连接上加CBAM或SE否则连基线都没有改了什么效果都说不清。import torch import torch.nn as nn 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), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch3, out_ch1, base64): super().__init__() self.down1 DoubleConv(in_ch, base) self.down2 DoubleConv(base, base*2) self.down3 DoubleConv(base*2, base*4) self.down4 DoubleConv(base*4, base*8) self.down5 DoubleConv(base*8, base*16) self.pool nn.MaxPool2d(2) self.up4 nn.ConvTranspose2d(base*16, base*8, 2, stride2) self.conv4 DoubleConv(base*16, base*8) self.up3 nn.ConvTranspose2d(base*8, base*4, 2, stride2) self.conv3 DoubleConv(base*8, base*4) self.up2 nn.ConvTranspose2d(base*4, base*2, 2, stride2) self.conv2 DoubleConv(base*4, base*2) self.up1 nn.ConvTranspose2d(base*2, base, 2, stride2) self.conv1 DoubleConv(base*2, base) self.out nn.Conv2d(base, out_ch, 1) def forward(self, x): d1 self.down1(x) d2 self.down2(self.pool(d1)) d3 self.down3(self.pool(d2)) d4 self.down4(self.pool(d3)) d5 self.down5(self.pool(d4)) u4 self.up4(d5) u4 torch.cat([u4, d4], dim1) u4 self.conv4(u4) u3 self.up3(u4) u3 torch.cat([u3, d3], dim1) u3 self.conv3(u3) u2 self.up2(u3) u2 torch.cat([u2, d2], dim1) u2 self.conv2(u2) u1 self.up1(u2) u1 torch.cat([u1, d1], dim1) u1 self.conv1(u1) return self.out(u1)base64时参数量约31M显存占用在512x512输入下约6GBbatch4。如果显存不够把base降到32或者用混合精度训练。ConvTranspose2d的stride2保证上采样倍数和池化对应。拼接时注意通道数up4输出base*8d4也是base*8拼完base*16正好喂给conv4。3.2 损失函数选择BCE、Dice还是组合人物抠图有个典型问题前景像素占比可能只有20%到40%背景占大头。纯BCE会让模型偏向预测背景边缘糊。DiceLoss直接优化重叠度对小目标更友好。我一般用BCEWithLogitsLoss加DiceLoss权重各0.5。注意BCEWithLogitsLoss内部带Sigmoid所以网络输出不要加Sigmoid。class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, logits, targets): probs torch.sigmoid(logits) probs probs.view(-1) targets targets.view(-1) intersection (probs * targets).sum() dice (2. * intersection self.smooth) / (probs.sum() targets.sum() self.smooth) return 1 - dice criterion lambda pred, target: 0.5 * nn.BCEWithLogitsLoss()(pred, target) 0.5 * DiceLoss()(pred, target)smooth1.0防止分母为零。如果训练时loss震荡把Dice权重降到0.3试试。另外注意target要是float类型且值在0到1之间。3.3 训练循环与学习率调度优化器用AdamW初始学习率1e-3权重衰减1e-4。调度器用CosineAnnealingLRT_max设总epoch数。Batch size在显存允许下尽量大8或16。每个epoch跑完在验证集上算IoU和Dice保存IoU最高的权重。from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR model UNet(base64).cuda() optimizer AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max100) scaler torch.cuda.amp.GradScaler() for epoch in range(100): model.train() for img, mask in train_loader: img, mask img.cuda(), mask.cuda().float() optimizer.zero_grad() with torch.cuda.amp.autocast(): pred model(img) loss criterion(pred, mask) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step() # 验证并保存最佳权重混合精度训练用GradScaler能省30%左右显存。autocast下模型输出是float16但损失计算会自动转回float32。验证时记得model.eval()和torch.no_grad()。如果IoU在某个epoch后不升反降大概率是过拟合加早停或者增大权重衰减。4. 推理与后处理把概率图变成可用的透明底PNG4.1 滑窗推理与大图拼接训练时裁到512x512推理时原图可能更大。直接resize会丢细节滑窗推理更稳。窗口大小512步长25650%重叠每个窗口预测后取中心区域拼回原图。重叠区域取平均边缘过渡更自然。def sliding_window_inference(model, image, window512, stride256): model.eval() h, w image.shape[:2] prob_map np.zeros((h, w), dtypenp.float32) count_map np.zeros((h, w), dtypenp.float32) for y in range(0, h, stride): for x in range(0, w, stride): y1, y2 y, min(y window, h) x1, x2 x, min(x window, w) patch image[y1:y2, x1:x2] patch cv2.resize(patch, (window, window)) tensor torch.from_numpy(patch).permute(2,0,1).unsqueeze(0).float().cuda() / 255.0 with torch.no_grad(): pred torch.sigmoid(model(tensor)).cpu().numpy()[0,0] pred cv2.resize(pred, (x2-x1, y2-y1)) prob_map[y1:y2, x1:x2] pred count_map[y1:y2, x1:x2] 1 return prob_map / np.maximum(count_map, 1)stride256保证每个像素至少被覆盖一次。count_map记录覆盖次数最后取平均。如果显存够窗口可以设768步长384减少拼接痕迹。4.2 阈值选取与边缘优化概率图出来后阈值0.5是默认但人物抠图边缘往往需要更精细。我一般用0.4到0.6之间扫一遍看验证集Dice。另外可以对概率图做条件随机场CRF或者引导滤波让边缘贴合原图纹理。CRF用pydensecrf库引导滤波用OpenCV的ximgproc.guidedFilter。import cv2 import numpy as np def refine_mask(prob_map, image, radius8, eps1e-3): guide image.astype(np.float32) / 255.0 mask prob_map.astype(np.float32) refined cv2.ximgproc.guidedFilter(guide, mask, radius, eps) return refinedradius8控制滤波窗口eps越小边缘越锐利。引导滤波比CRF快很多效果对人物抠图够用。如果头发丝区域还是糊可以再叠一个Matting方法但那就超出UNet语义分割的范围了。4.3 生成透明底PNG与批量处理最终掩膜二值化后把原图RGB和掩膜alpha通道合并存成RGBA PNG。def save_transparent(image, mask, output_path): b, g, r cv2.split(image) rgba cv2.merge([b, g, r, mask]) cv2.imwrite(output_path, rgba) # 批量 for name in os.listdir(test_images): img cv2.imread(ftest_images/{name}) prob sliding_window_inference(model, img) prob refine_mask(prob, img) mask (prob 0.5).astype(np.uint8) * 255 save_transparent(img, mask, foutput/{name.replace(.jpg, .png)})注意OpenCV读进来是BGR存RGBA时通道顺序要对应。如果下游要用在网页或设计软件里PNG的alpha通道就是透明底直接拖进PS或GIMP就能用。热搜里“gimp怎么抠图详细步骤”和“网页版ps图艺图怎么抠图”其实都是在找替代方案UNet批量出透明底PNG正好补上这个缺口。5. 避坑与排查人物抠图训练中常见的五个翻车现场5.1 掩膜全黑或全白现象训练几个epoch后验证集预测全是背景或全是人。原因学习率太大导致模型崩溃或者损失函数权重失衡。解决把学习率降到1e-4检查DiceLoss权重是否过高。另外确认掩膜读取时有没有除以255如果掩膜值是0和255不归一化的话BCE会爆炸。5.2 边缘出现锯齿或光晕现象人物轮廓像楼梯或者周围有一圈半透明灰边。原因下采样丢失高频信息或者上采样用了转置卷积导致棋盘效应。解决把转置卷积换成双线性插值加卷积或者在跳跃连接前对编码器特征做一次3x3卷积平滑。后处理加引导滤波也能缓解。5.3 多人场景只抠出一个人现象画面里两个人掩膜只覆盖了其中一个。原因标注时只标了一个人或者损失函数对大面积前景过拟合。解决检查标注JSON是否漏标训练时用DiceLoss对每个连通域加权。如果多人重叠严重语义分割本身就有局限考虑换实例分割。5.4 验证集IoU很高但实际抠图效果差现象IoU 0.9但透明底PNG边缘一塌糊涂。原因验证集和训练集同分布且IoU对边缘不敏感。解决单独算边界区域的IoU或者用BFScore。另外把验证集换成不同背景、不同光照的图别用训练集里截出来的。5.5 显存溢出OOM现象训练到一半报CUDA out of memory。原因batch size太大或者输入尺寸没固定。解决用混合精度把base_channels从64降到32或者用梯度累积模拟大batch。推理时用滑窗别一次性把整张大图塞进去。6. 进阶技巧用TTA和模型集成把Dice再提两个点训练完一个UNet后别急着上线。先做测试时增强TTA对同一张图做水平翻转、多尺度缩放分别推理后取平均。这个操作不需要重新训练推理时间翻几倍但Dice通常能涨1到2个点。具体做法是原图、水平翻转图、0.8倍缩放图、1.2倍缩放图各跑一次滑窗推理把概率图对齐后平均。def tta_inference(model, image): probs [] # 原图 probs.append(sliding_window_inference(model, image)) # 水平翻转 flipped cv2.flip(image, 1) p_flip sliding_window_inference(model, flipped) probs.append(cv2.flip(p_flip, 1)) # 多尺度 for scale in [0.8, 1.2]: h, w image.shape[:2] resized cv2.resize(image, (int(w*scale), int(h*scale))) p sliding_window_inference(model, resized) p cv2.resize(p, (w, h)) probs.append(p) return np.mean(probs, axis0)如果还想再进一步可以训两个不同base_channels的UNet比如32和64或者一个用BCE为主、一个用Dice为主推理时概率图加权平均。集成两个模型通常比单模型涨2到3个点但推理成本翻倍。线上服务如果QPS要求高可以只对边缘区域做TTA背景区域用单次推理。最后说个血泪教训别在训练集上调后处理参数。我曾经在训练集上把阈值调到0.35Dice刷到0.93结果测试集一跑掉到0.81。后来固定用0.5只在验证集上选引导滤波的radius和eps泛化才稳。另外模型保存时把优化器状态和epoch一起存不然想接着训还得重头来。希望帮到你。本文还有配套的精品资源点击获取
返回列表