简介这份压缩包面向深度学习与图像分割初学者及进阶开发者提供基于segmentation_models.pytorch框架完成人物抠图二分类语义分割的完整实现方案。内容围绕框架调用、语义分割数据集加载、交叉熵与Dice Loss组合、WandB可视化等关键环节展开并完整给出训练与预测两阶段的可运行工程代码。压缩包共2000个文件包括大量png训练图、Python源码py/pyc、Dockerfile、说明文档等整体约990MB代码注释清晰、目录结构规整适合边读博文边实践也可直接下载后复现算法节省环境与数据准备时间。通过该资源可重点掌握segmentation_models.pytorch的模型构建、二分类损失设计、训练监控及推理流程也能借鉴其中的数据读取、验证集划分和预测结果保存等工程细节。目前已有1764人学习下载是快速入门语义分割、落地人物抠图任务的实用参考。1. segmentation_models.pytorch 为什么成了人物抠图的首选框架把人物从背景里抠出来听起来是图形软件的活儿但遇到大批量图片、复杂发丝、逆光轮廓时PS 套索和 GIMP 的魔术棒根本不抗造。segmentation_models.pytorch简称 SMP这类图像分割框架把人物抠图变成像素级分类问题模型输出一张每个像素属于前景还是背景的概率图再转成 alpha 通道就能得到透明 PNG。SMP 把 ResNet、EfficientNet 等预训练编码器和 UNet、DeepLabV3 这类解码器封装成一行代码让抠图的落地成本压到一台 GPU 加几百张标注图并且可以直接拿来做批量出图和二次开发。本文按我实际落地的路径走一遍环境搭建、数据准备、模型训练、常见坑以及最后交付透明图。新手能照做熟手能拿去调参。2. 环境搭建与数据准备人物抠图的地基打在哪2.1 安装 segmentation_models.pytorch版本搭配与验证SMP 本身是一个高层封装库底层依赖 PyTorch、torchvision 和 timm。常见的翻车点不是 SMP 装不上去而是 PyTorch 版本和 CUDA 版本对不上或者 import 的时候报找不到timm。我一般会先创建一个干净的 conda 环境把 PyTorch 装好再装 SMP顺序反了容易出现包依赖解析问题。conda create -n smp-matting python3.9 -y conda activate smp-matting # 先装 PyTorch以 CUDA 11.8 为例具体按你的显卡驱动选版本 pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu118 # 再装 SMP pip install segmentation-models-pytorch逻辑说明--index-url后面的地址是 PyTorch 官方预编译 wheel 源按cu118、cu121这样的后缀区分 CUDA 版本。如果你的机器是纯 CPU去掉--index-url直接装 CPU 版即可。装完 SMP 后先跑一个最小验证避免后面训练时才发现环境有问题。import torch import segmentation_models_pytorch as smp model smp.Unet( encoder_nameresnet50, encoder_weightsimagenet, in_channels3, classes1, ) x torch.randn(2, 3, 256, 256) y model(x) print(y.shape) # 期望输出 torch.Size([2, 1, 256, 256])参数说明encoder_name是编码器骨干网络encoder_weightsimagenet表示加载在 ImageNet 上预训练好的权重in_channels3对应 RGB 输入classes1表示二分类任务即前景和背景。模型输出的不是概率而是 logits后面要接torch.sigmoid才能得到 0 到 1 的置信度图。如果第一次跑这个验证网络会自动下载预训练权重大约几百 MB建议提前准备好网络环境。2.2 数据准备开源数据集、标注工具还是用 PS/GIMP 已有结果人物抠图本质上是做一个二值分割任务所以数据要是“原图 对应 mask”的成对形式。mask 是黑白图白色代表人物黑色代表背景。准备数据有三条路我建议按你的实际场景选。第一条路是直接用开源人物分割数据集比如 COCO 数据集里的 person 类标注或者 Supervisely Person 这类专门的人物分割数据集。好处是标注质量有保障缺点是开源数据集的场景和你要处理的业务图可能差异很大后面会有领域偏移的问题。第二条路是用标注工具自己标比如 labelme、AnyLabeling 这类工具导出的多边形 JSON 再转成 mask适合只有几十张图的小场景。第三条路是把你用 PS 或 GIMP 已经抠好的透明 PNG 反推成 mask这个做法速度最快而且标注质量完全是你自己控制的。有了数据后就要写 Dataset 类。这里有一个非常容易踩的坑mask 在做 resize 时必须用最近邻插值不能用线性插值否则会生成 0 和 255 之间的灰阶相当于给模型灌了错误的模糊标签。import glob import cv2 import numpy as np import torch from torch.utils.data import Dataset class PersonMattingDataset(Dataset): def __init__(self, image_dir, mask_dir, size256): self.image_paths sorted(glob.glob(f{image_dir}/*.jpg)) self.mask_paths sorted(glob.glob(f{mask_dir}/*.png)) self.size size def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image cv2.imread(self.image_paths[idx]) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) image cv2.resize(image, (self.size, self.size), interpolationcv2.INTER_LINEAR) # 关键mask 必须用最近邻插值否则边缘会被混出新灰度值 mask cv2.resize(mask, (self.size, self.size), interpolationcv2.INTER_NEAREST) image torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(mask).float() / 255.0 # 基础版人物抠图用二值 mask先阈值化边缘细节留给后处理 mask (mask 0.5).float().unsqueeze(0) return image, mask参数说明mask 0.5会把所有大于半灰度的像素变成 1小于的变成 0。如果你的标注是带抗锯齿的软边直接阈值化会丢掉发丝和半透明边缘的信息。一个常见的“后悔药”是训练时保留连续灰度值也就是不0.5直接把mask / 255.0当软标签损失函数里要用适合连续标签的版本。但这样做模型输出的边缘会更细腻训练收敛也更慢。我一般先跑二值版本把整个流程打通之后再根据交付要求决定要不要上软标签。2.3 数据增强与归一化这些参数不能乱填数据增强对人物抠图的影响非常大尤其是边缘类别的样本。常用的增强库是 albumentations它最方便的地方是能同时处理 image 和 mask并且 mask 的插值方式被固定为最近邻不会像你自己写那样子破坏标签。import albumentations as A from albumentations.pytorch import ToTensorV2 train_transform A.Compose([ A.RandomResizedCrop(256, 256, scale(0.6, 1.0), p0.8), A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(brightness_limit0.2, contrast_limit0.2, p0.3), A.HueSaturationValue(hue_shift_limit5, sat_shift_limit20, val_shift_limit20, p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) valid_transform A.Compose([ A.Resize(256, 256), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ])训练时会这样调用transformed transform(imageimage, maskmask)然后从transformed[image]和transformed[mask]里取数据。参数说明RandomResizedCrop的scale(0.6, 1.0)意思是随机裁剪原图的 60% 到 100% 再缩放到 256相当于一种随机缩放能提升模型对不同人物占比的适应力。HueSaturationValue的偏移量不要给太大给太大容易把肤色调成怪异颜色模型反而学到了错误的颜色关联。归一化用 ImageNet 的均值和标准差因为 SMP 加载的预训练权重就是基于 ImageNet 统计量训练的跳过了这一层前面的预训练权重就白加载了收敛速度会明显变慢。3. 用 SMP 构建人物分割模型编码器选型与损失函数3.1 编码器选型ResNet50 稳EfficientNet 准MobileNet 快SMP 支持几十种编码器大部分来自 timm 库。人物抠图任务对编码器的要求并不苛刻关键看你的部署场景。我常用的三档对比看下表编码器参数量推理速度人物边缘细节适用场景resnet50约 67M中等一般万能起步兼容性好efficientnet-b3约 12M快较细精度优先显存有限时mobilenet_v3_large约 5M最快中等CPU 或移动端部署如果只是要把流程跑通直接选resnet50是最省心的如果发现边缘细节不够好或者显存只有 6G、8G换成efficientnet-b3通常能在差不多的精度下省下一半显存。考虑实时场景比如做 OBS 直播抠像插件这类应用就得换mobilenet_v3_large同时把输入分辨率降到 256 甚至 192。编码器对最终效果的影响往往没有输入分辨率和数据质量影响大。3.2 UNet 和 DeepLabV3两个都能“一行代码”建起来SMP 最舒服的一点是换模型结构几乎不写新代码。同一份训练脚本里把模型对象换一行就行。import segmentation_models_pytorch as smp # UNet ResNet50先作为基线跑通 model smp.Unet( encoder_nameresnet50, encoder_weightsimagenet, in_channels3, classes1, ) # DeepLabV3 EfficientNet边缘细节通常更好 model smp.DeepLabV3Plus( encoder_nameefficientnet-b3, encoder_weightsimagenet, in_channels3, classes1, decoder_aspp_sep_kernel_size1, )参数说明decoder_aspp_sep_kernel_size控制 ASPP 模块的空洞卷积核大小默认值通常不必改。对于人物抠图UNet 的上采样路径能让分割结果更完整不容易出现人物身体断成两截的情况DeepLabV3 对边缘轮廓保持更好但在小目标上偶尔会漏掉躯干。我一般先跑 UNet 验证数据没问题再切 DeepLabV3 对比效果。encoder_weightsimagenet会联网下载权重如果下载失败训练会直接报错可以提前手动把权重文件下载好放到缓存目录里避免在机器上干等。3.3 损失函数为什么 BCE 和 Dice 要搭配用人物抠图的标签里背景面积通常远大于人物面积直接训练二分类交叉熵模型会倾向于把不确定的像素都判成背景因为这样能让整体 loss 更低。Dice loss 关注的是预测区域和真实区域的重叠度正好能抑制这种倾向。import torch import torch.nn as nn class BCEWithLogitsDiceLoss(nn.Module): def __init__(self, bce_weight0.5, dice_weight0.5, smooth1e-6): super().__init__() self.bce nn.BCEWithLogitsLoss() self.bce_weight bce_weight self.dice_weight dice_weight self.smooth smooth def forward(self, logits, targets): prob torch.sigmoid(logits) bce_loss self.bce(logits, targets) intersection (prob * targets).sum() dice_loss 1 - (2 * intersection self.smooth) / ( prob.sum() targets.sum() self.smooth ) return self.bce_weight * bce_loss self.dice_weight * dice_loss参数说明smooth是平滑系数防止分母为 0也起到一个轻微正则作用。bce_weight0.5、dice_weight0.5是最常见的配比但如果你发现模型抠出来的区域总是比真实区域大一圈可以把dice_weight调高到 0.7让模型更看重区域重叠如果边缘总是抠不干净可以反过来把bce_weight调高到 0.7让每个像素都受到更直接的监督。另一个常见变体是 Focal Loss处理背景占比极端的情况我放在后面避坑章节里展开。3.4 评估指标IoU 不是损失函数别混着看训练时打印的 loss 只是优化目标不能直接反映人物抠得干不干净。分割任务最常用的指标是 IoU也就是预测 mask 和真实 mask 的交集除以并集。数值越接近 1 越好。def compute_iou(logits, masks, threshold0.5): prob torch.sigmoid(logits) pred (prob threshold).int().view(-1) target masks.int().view(-1) intersection (pred target).sum().float() union (pred | target).sum().float() return (intersection / union).item()pred target是按位与pred | target是按位或在 PyTorch 里做二值集合运算非常直观。IoU 对大面积重叠很敏感但人物边缘那一两圈像素在整体 IoU 里占比很小所以我会额外看一眼“边界 IoU”也就是只计算真实 mask 边缘带状区域内的 IoU。这个指标能直接反映发丝、衣物轮廓这些最容易被吐槽的地方比只看整体 IoU 有用得多。4. 训练与调参人物边缘抠得干不干净全看这里的设置4.1 训练循环AMP、AdamW、学习率调度三项一个都不能少训练循环本身不难但有几个细节决定训练能不能稳定收敛。第一用混合精度训练显存占用能降一半速度提升明显第二用 AdamW 而不是普通 Adam因为 AdamW 把权重衰减和梯度更新解耦对预训练模型微调更友好第三给梯度加裁剪防止个别批次数据异常导致梯度爆炸训练过程一下子崩掉。import torch import torch.nn as nn epochs 40 device torch.device(cuda if torch.cuda.is_available() else cpu) optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs) criterion BCEWithLogitsDiceLoss(bce_weight0.5, dice_weight0.5) scaler torch.cuda.amp.GradScaler() for epoch in range(epochs): model.train() total_loss 0.0 total_iou 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): logits model(images) loss criterion(logits, masks) scaler.scale(loss).backward() # 混合精度下梯度可能异常裁剪一下更稳 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) scaler.step(optimizer) scaler.update() total_loss loss.item() total_iou compute_iou(logits, masks) scheduler.step() print(fepoch {epoch:02d} loss {total_loss/len(train_loader):.4f} fiou {total_iou/len(train_loader):.4f})这里有一个版本兼容问题老版本 PyTorch 用torch.cuda.amp.autocast()和torch.cuda.amp.GradScaler()新版本更推荐torch.amp.autocast(cuda)和torch.amp.GradScaler(cuda)。两个写法在我的代码里都能跑通但我建议你按自己安装的 PyTorch 版本来选不然会看到一条 DeprecationWarning不影响训练但看着闹心。max_norm5.0是梯度裁剪的阈值一般取 1 到 10 之间人物分割任务我习惯取 5。4.2 关键超参图像尺寸、batch size、epochs怎么定人物抠图的超参没有绝对最优但有一张参考表能减少大量试错时间参数默认值调整方向说明图像尺寸256384 / 512尺寸越大边缘细节越丰富显存占用成倍增加batch size84 / 16受显存限制batch 越大收敛越稳epochs4030 / 60更看验证集表现不是越多越好初始学习率1e-35e-4 / 2e-3换小编码器可适当调低图像尺寸是人物抠图里最值得投资的超参。256 适合快速迭代和验证流程512 的边缘细节明显比 256 好但显存占用大约是 4 倍。如果显卡只有 8G我建议用 384 作为折中方案。batch size 反过来调如果显存紧张可以保持 batch size 为 4用梯度累积模拟更大的 batch。比如每 4 个 batch 做一次优化器更新效果接近 batch size 16代码上就是在累加 loss 后延迟调用optimizer.step()和optimizer.zero_grad()。4.3 检查点保存与训练可视化别让训练白跑跑人物分割训练动辄几个小时一旦中断就得从头再来很难受。所以保存检查点是必做项。我习惯每轮都保存一套完整的 checkpoint而不是只保存最后的权重。torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), config: { encoder: resnet50, input_size: 256, batch_size: 8, }, }, fcheckpoints/smp_epoch{epoch:02d}.pth)加载的时候这样操作checkpoint torch.load(checkpoints/smp_epoch30.pth, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) start_epoch checkpoint[epoch] 1说明map_locationcpu是为了防止服务器和本地 CUDA 版本不一致导致加载报错。保存optimizer_state_dict的意义在于中断后能恢复学习率状态不然余弦退火的进度就断了。日志方面除了打印 loss 和 IoU我还会每轮随机拿几张验证图片出来跑一次推理把预测 mask 存下来肉眼观察。loss 曲线只能说明模型在收敛不能说明人物边缘真的干净可视化是最直接的反馈。5. 人物抠图落地常见坑与排查现象、原因、对策5.1 边缘发虚像蒙了一层毛玻璃现象模型输出的透明 PNG人物边缘有一圈半透明残留放到深色背景上尤其明显像没擦干净。原因有三个第一训练集 mask 在 resize 时用了线性插值导致标签边缘出现灰阶第二训练标签本身就是抗锯齿边缘模型学出来的是一个渐变 alpha而你用的是二值化第三推理时直接按 0.5 阈值切边缘置信度本身就低切出来的边缘参差不齐。解决方法是先检查训练数据确认 mask 是纯黑白的然后在推理后处理里加一圈形态学操作。# 对概率图先取阈值再做形态学清理 mask (prob 0.5).astype(np.uint8) * 255 mask cv2.medianBlur(mask, 5) # 去除边缘上的孤立噪点 mask cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)) # 填住边缘小洞medianBlur的核大小取 3 或 5取太大人物发丝会被抹平。MORPH_CLOSE是先膨胀后腐蚀能填平 mask 上细小的孔洞。这个操作不会让边缘变锋利但能显著减少半透明残留。如果你的业务要求发丝级别的抠图二值分割路线本身就满足不了得用专门的 matting 模型这个后处理只能救边缘救不了头发丝。5.2 人物和背景撞色时背景被一起抠了进来现象深色衣服配深色背景模型预测出的 mask 比人物真实轮廓大一圈背景里颜色相近的区域也被归成前景。原因训练数据里这类撞色样本太少模型本质上还是在靠颜色先验判断而不是真正理解人物轮廓另一个叠加原因是 BCE 损失对大面积误分类的惩罚不够模型把一大块背景误判成人物整体 loss 也不会飙升太多。解决思路是用 Focal Loss 替代普通 BCE把难分的像素权重抬高。class FocalLoss(nn.Module): def __init__(self, alpha0.25, gamma2.0): super().__init__() self.alpha alpha self.gamma gamma self.bce nn.BCEWithLogitsLoss(reductionnone) def forward(self, logits, targets): bce_loss self.bce(logits, targets) prob torch.sigmoid(logits) # 标准 focal 权重预测越自信权重越小 focal_weight torch.where( targets 0, 1 - prob, prob ).pow(self.gamma) return (self.alpha * focal_weight * bce_loss).mean()gamma2.0是常用值模型已经预测对的简单像素权重会被压得非常低难分像素主导梯度更新alpha0.25用来调正负样本的整体比例。但要注意Focal Loss 也会让训练收敛变慢建议先拿普通 BCEDice 跑 10 个 epoch确认数据没问题再加 Focal。另外在数据增强里把亮度对比度扰动加大一点也能缓解颜色撞车问题。5.3 训练正常推理时显存突然爆掉现象训练 8 个 batch 都没问题推理单张图却 OOM。原因推理脚本没有包torch.no_grad()图还在建反向图显存翻了至少一倍或者输入图尺寸比训练尺寸大很多模型的 feature map 显存占用成倍膨胀。解决方法是推理时强制不追踪梯度并改成单张输入。model.eval() with torch.no_grad(): for images, _ in valid_loader: images images.to(device) logits model(images) prob torch.sigmoid(logits).cpu().numpy()注意model.eval()只切换 BN 和 Dropout 状态不关闭梯度追踪必须配合torch.no_grad()。如果单张图还是 OOM看一下原图分辨率超过训练尺寸太多就先cv2.resize到训练尺寸。你要让推理图和训练图的分辨率保持一致模型的分割效果才会稳定。5.4 换到电商图、网页图之后效果突然跳水现象在开源人像数据集上 IoU 能到 0.96拿到商品详情页的模特图上一测边缘多了一大圈背景或者人物腿上缺了一块。原因这是典型的领域偏移开源数据集大多是摄影作品光照自然、背景简单而电商图大量使用白底、浅灰底人物姿态和裁切方式也更极端。解决方式不是盲目加数据而是收集 50 到 100 张目标场景图做微调。一个和效率有关的习惯做法是先用 PS 或 GIMP 手工抠出这批图的精确 mask再拿来微调模型。几十张图就足以让模型在新场景上恢复可用水平比从零训练划算得多。这里也解释了一个现象网上有些“AI 一键抠图”的网页版工具换到专业图上不好用本质就是训练数据和你的业务图不在同一个域。OBS 直播的 AI 抠像插件在直播画面里表现好是因为直播背景相对固定把它拿去抠海报图也会翻车。所以不要指望一个通用模型解决所有场景小规模微调几乎是必做的一步。5.5 视频人物抠像边缘像噪声一样在“呼吸”现象对视频逐帧调用模型单帧静止看都还行连起来播放时人物边缘一直在抖动发丝边缘一会粗一会细。原因模型对每一帧独立推理相邻帧的像素置信度波动被直接输出成了 mask 边界的抖动这不叫算法坏了而是时间一致性没人管。解决手段是加一个时间维度的平滑最常见的是指数移动平均。# 简单的时间平滑alpha 控制当前帧权重 prev_mask None for frame in video_frames: curr_mask predict(frame) # 得到当前帧概率图 if prev_mask is None: fused_mask curr_mask else: fused_mask alpha * curr_mask (1 - alpha) * prev_mask prev_mask fused_mask output_mask (fused_mask 0.5).astype(np.uint8)alpha我一般取 0.7 到 0.8场景快速切换时调低到 0.5否则人物快速移动会拖影。更复杂一点可以用光流把前一帧的 mask 对齐到当前帧再融合但对多数人物抠像需求指数平滑已经能把边缘抖动压到肉眼可接受的程度。6. 让模型真正变成“一键抠图”后处理与交付6.1 从概率图到透明 PNG 的完整后处理模型输出的 logits 不能直接用要经过 sigmoid 转概率、阈值化转 mask、形态学清理、缩放到原图分辨率。我习惯把所有步骤封装成一个函数方便在批量脚本里反复调用。def predict_to_mask(model, image_bgr, size256, threshold0.5): image cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) h, w image.shape[:2] image cv2.resize(image, (size, size), interpolationcv2.INTER_LINEAR) image torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 image normalize(image).unsqueeze(0).to(device) with torch.no_grad(): prob torch.sigmoid(model(image)).squeeze().cpu().numpy() mask (prob threshold).astype(np.uint8) * 255 mask cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)) mask cv2.resize(mask, (w, h), interpolationcv2.INTER_LINEAR) return mask输出 mask 后用cv2.merge把人物原图和 mask 合成 RGBA 图再写 PNG 就是带透明通道的成品。threshold可以做成可调参数边缘偏粗就调高到 0.6偏窄就调到 0.4。注意最后一步cv2.resize回到原图时用线性插值这样 mask 边缘会有一点点柔化视觉上反而自然。6.2 和 PS/GIMP 的流程衔接训练好的模型更适合做批量初筛而不是取代精修。我通常的交付流程是先跑模型把所有图的透明底拿出来再用 GIMP 的“从 Alpha 选择”功能加载 mask对需要精修的大图手动调整对不需要的图直接保存。这样模型承担 80% 的重复劳动PS 框选这类手工操作只留给真正复杂的边缘区域。如果交付的是头像、商品主图这类固定尺寸素材模型跑完直接导出透明 PNG 就够用了不需要再过一遍图形软件。6.3 验证模型到底把人物抠干净了没有除了 IoU 数值我还会额外看三个维度一是边缘 2 到 3 像素内的准确率边缘干净程度比整体面积重叠更重要二是对高光、透明材质、深色背景三类困难样本的单独表现三是单张推理耗时控制在应用可接受范围内。验证集我固定留 30 到 50 张不参与训练的图片每次训练完都肉眼逐张翻一遍预测结果。我现在换数据集之后的第一件事是打印三张最失败的预测图而不是盯着 loss 曲线。loss 曲线只能告诉你模型在拟合训练数据不能告诉你它在人物发丝边缘上输得有多难看。希望帮到你也让你的分割模型少走几趟弯路。本文还有配套的精品资源点击获取