简介这份资源面向计算机视觉入门与进阶开发者围绕UNet语义分割的人物抠图任务展开帮助读者理解二分类分割的完整实现路径。内容涵盖单通道输出、Sigmoid归一化、交叉熵损失计算与反向传播等关键环节适合想掌握图像分割基础、动手复现抠图效果的读者。压缩包共3430个文件以3404张png图像数据为主辅以11个py训练与推理脚本、7个pyc缓存文件以及Dockerfile、sh、md、txt等环境配置与说明文档整体约987.41MB目录结构便于按数据、代码、配置分块查阅。目前已有2816人学习下载热度较高。读者可从中获得可直接运行的分割工程、配套数据集与训练脚本快速搭建UNet人物抠图实验环境并对照代码理解二分类分割从数据加载到模型输出的完整流程适合作为课程设计、毕业项目或算法练手的参考素材。1. 从一张人像图到二值掩膜UNet 抠图资源包拆解手里这份UNet语义分割实战使用UNet实现对人物的抠图.zip解压后其实很朴素一个Dockerfile、一份README.md、LICENSE、.gitignore外加四张示例图00315.png、00458.png、00916.png、00109.png。没有花哨的目录树也没有预训练权重它更像一个「最小可跑骨架」——把 UNet 做人物二分类分割这条链路的关键文件摆在你面前让你自己把训练、推理、后处理串起来。它解决的是「人像抠图」这个具体问题输入一张 RGB 人像输出一张同尺寸的单通道掩膜前景人为 1背景为 0。适合谁适合已经跑通过分类网络、想第一次把语义分割落到真实图片上的工程师也适合做证件照换底、直播虚拟背景、电商主图去背这类需求想自己掌控模型而不是调第三方 API 的人。下面按「资源是什么 → 怎么用 → 坑在哪」的顺序把这份包拆到能复现的程度。2. 二分类分割的两种输出形态为什么这份资源走单通道 sigmoid2.1 单通道 sigmoid 与双通道 softmax 的取舍语义分割做二分类业界常见两条路。第一条是网络输出[batch_size, 1, height, width]数值任意训练时先过sigmoid压到[0,1]再和只有 0/1 的target算交叉熵。第二条是输出[batch_size, 2, height, width]过softmax得到两类概率取argmax得到掩膜。这份资源走的是第一条。原因很实际人物抠图本质是「前景 vs 背景」的二元判断单通道省一半输出计算量显存占用更低小批量也能训。更重要的是推理阶段你只需要一个阈值常见 0.5就能二值化后处理链路短。双通道 softmax 在类别极不平衡时反而容易把背景概率压得过高需要额外调 class weight。代价也要说清单通道 sigmoid 对类别不平衡更敏感。如果训练集里人占画面比例很小损失会被大量背景像素主导模型倾向全预测背景。常见做法是加pos_weight或者改用 Dice Loss、BCEDice 组合。这份包没带训练脚本所以这部分要你自己补后面第 4 章会给可抄的损失配置。2.2 从输出张量到可保存的掩膜完整数据流把链路拆开看一次前向到出图要经过这些步骤读图转 RGB缩放到网络输入尺寸UNet 常用 256×256 或 512×512需被 16 整除。归一化常见mean[0.485,0.456,0.406]、std[0.229,0.224,0.225]或简单除以 255。前向得到output形状[1,1,H,W]。torch.sigmoid(output)得到概率图。按阈值二值化得到 0/1 掩膜。把掩膜缩回原图尺寸与原图做逐像素乘得到抠出的前景。第 6 步是很多人翻车的地方先缩放再二值化还是先二值化再缩放正确顺序是先二值化再最近邻缩放否则插值会在边缘产生 0 到 1 之间的灰边抠图边缘发虚。这个细节资源包里没写但它是抠图能不能用的分水岭。2.3 资源包文件与用途对照文件类型用途Dockerfile构建脚本固定 Python、PyTorch、CUDA 版本避免环境漂移README.md文档项目说明与运行入口LICENSE授权使用前确认许可范围.gitignore/.keep版本控制忽略产物、保留空目录00315.png等四张示例图验证推理链路是否跑通的输入样本四张图不是数据集是「冒烟测试」样本。别拿它们当训练集否则模型只会记住这四张。真正训练要自己准备人像数据集第 3 章讲怎么造。3. 用 Dockerfile 固定环境并跑通推理链路3.1 先看 Dockerfile 再动手别急着 pip install拿到包第一件事不是pip install是读Dockerfile。它决定了基础镜像、CUDA 版本、依赖安装顺序。典型结构长这样FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime WORKDIR /workspace COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt COPY . . CMD [python, infer.py]逻辑说明基础镜像直接带 PyTorch 和 CUDA省去编译时间requirements.txt单独 COPY 再安装利用镜像层缓存改代码不会重装依赖最后 COPY 全部代码。参数上cuda12.1要和宿主机驱动匹配驱动太旧就换cuda11.8的 tag。构建命令docker build -t unet-matting:latest . docker run --gpus all -v $(pwd)/data:/workspace/data unet-matting:latest--gpus all暴露 GPU-v把宿主机数据目录挂进去避免每次改图都重建镜像。如果没 GPU去掉--gpus all镜像换成 CPU 版 tag。3.2 补一个最小推理脚本资源包没给推理代码按这个场景补一份能直接跑的import torch import torch.nn as nn import numpy as np from PIL import Image from torchvision import transforms class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch3, out_ch1): super().__init__() # 编码器逐层下采样通道翻倍 self.d1 DoubleConv(in_ch, 64) self.d2 DoubleConv(64, 128) self.d3 DoubleConv(128, 256) self.d4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) # 解码器上采样后与编码器特征拼接 self.up4 nn.ConvTranspose2d(512, 256, 2, stride2) self.u4 DoubleConv(512, 256) self.up3 nn.ConvTranspose2d(256, 128, 2, stride2) self.u3 DoubleConv(256, 128) self.up2 nn.ConvTranspose2d(128, 64, 2, stride2) self.u2 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): c1 self.d1(x) c2 self.d2(self.pool(c1)) c3 self.d3(self.pool(c2)) c4 self.d4(self.pool(c3)) x self.u4(torch.cat([self.up4(c4), c3], dim1)) x self.u3(torch.cat([self.up3(x), c2], dim1)) x self.u2(torch.cat([self.up2(x), c1], dim1)) return self.out(x) def infer(img_path, ckpt_path, size256, thresh0.5): device cuda if torch.cuda.is_available() else cpu model UNet().to(device) model.load_state_dict(torch.load(ckpt_path, map_locationdevice)) model.eval() img Image.open(img_path).convert(RGB) orig_size img.size tf transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225]), ]) x tf(img).unsqueeze(0).to(device) with torch.no_grad(): prob torch.sigmoid(model(x)) # 概率图 [1,1,H,W] mask (prob thresh).float() # 先二值化 mask torch.nn.functional.interpolate( mask, size(orig_size[1], orig_size[0]), modenearest ) # 再最近邻放大回原尺寸 mask_np mask.squeeze().cpu().numpy() rgba np.array(img.convert(RGBA)) rgba[:, :, 3] (mask_np * 255).astype(np.uint8) # 掩膜写进 alpha 通道 Image.fromarray(rgba).save(out.png) if __name__ __main__: infer(00315.png, unet.pth)逻辑说明DoubleConv是 UNet 的基本块两次卷积加 BN 和 ReLU编码器四次下采样解码器三次上采样并和对应层特征cat拼接这是 UNet 能恢复边缘细节的关键。参数上size必须能被 16 整除否则拼接时尺寸对不上thresh默认 0.5边缘发虚就调到 0.6漏前景就降到 0.4。3.3 用四张示例图做冒烟测试没有权重时先随机初始化跑一遍确认形状和显存没问题python -c import torch from infer import UNet m UNet() x torch.randn(1,3,256,256) print(m(x).shape) # 期望 torch.Size([1, 1, 256, 256]) 输出[1,1,256,256]说明网络结构对。这一步能提前暴露通道数写错、拼接维度不匹配这类低级错误比训到一半才发现省事得多。四张示例图分别跑一遍检查输出out.png的 alpha 通道是否只有 0 和 255 两个值有中间值说明二值化或缩放顺序写反了。4. 训练自己的数据集标签制作、损失配置与显存控制4.1 人像掩膜标签怎么造训练 UNet 抠图标签质量决定上限。常见做法有三条手工标注用 labelme、CVAT 画多边形导出二值 PNG。精度高但人像边缘头发丝极费时间。半自动先用现成模型粗抠再人工修边。适合批量生产。合成把前景人像贴到不同背景上自动生成掩膜。适合扩充数据但要注意边缘融合自然。标签统一存成单通道 PNG前景 255、背景 0文件名和原图一一对应。目录结构建议dataset/ images/ 00315.png ... masks/ 00315.png ...注意掩膜不要用 JPG 存JPG 有损压缩会在边缘产生 0 到 255 之间的灰值训练时被当成 0.5 附近的软标签边缘学不准。4.2 损失函数与类别不平衡处理单通道 sigmoid 配 BCE 是基线但人像占比小时要加权import torch import torch.nn as nn class BCEDiceLoss(nn.Module): def __init__(self, pos_weight2.0): super().__init__() self.bce nn.BCEWithLogitsLoss(pos_weighttorch.tensor([pos_weight])) def forward(self, logits, target): bce_loss self.bce(logits, target) prob torch.sigmoid(logits) inter (prob * target).sum() dice_loss 1 - (2 * inter 1e-6) / (prob.sum() target.sum() 1e-6) return bce_loss dice_loss逻辑说明BCEWithLogitsLoss内部做 sigmoid数值更稳不要在外面再套一层 sigmoidpos_weight给前景加权人越小调越大常见 2 到 5。Dice 项直接优化重叠度对不平衡更鲁棒。参数上1e-6防止除零别省。训练循环里记得把target转成float并归一化到[0,1]否则 BCE 会报维度或数值错误。4.3 显存不够时的三个调节旋钮UNet 在 512×512 输入下显存吃紧按优先级调降batch_size到 2 或 1配合梯度累积模拟大 batch。降输入尺寸到 256×256边缘精度会掉但能先跑通。把基础通道从 64 降到 32参数量约减到四分之一。梯度累积写法accum 4 for i, (img, mask) in enumerate(loader): loss criterion(model(img), mask) / accum loss.backward() if (i 1) % accum 0: optimizer.step() optimizer.zero_grad()accum4表示每 4 个 batch 更新一次等效 batch 放大 4 倍显存不变。注意 loss 要除以accum否则梯度被放大。5. 抠图链路的避坑与排查五个真实翻车点5.1 边缘出现灰边或半透明现象抠出的图边缘一圈发灰头发丝糊成一片。 原因先缩放掩膜再二值化双线性插值在 0 和 1 之间造出中间值。 解决严格按「sigmoid → 阈值二值化 → 最近邻缩放」顺序缩放模式用nearest别用bilinear。5.2 模型全预测背景现象训练 loss 很快降到很低但推理出来全是黑图。 原因前景像素占比太小BCE 被背景主导模型学到「全 0」就是局部最优。 解决加pos_weight或换 BCEDice 组合同时检查标签是否真的包含前景别把全黑掩膜喂进去。5.3 拼接时尺寸不匹配报错现象torch.cat报Sizes of tensors must match。 原因输入尺寸不能被 16 整除下采样再上采样后尺寸和编码器特征差 1 像素。 解决输入统一 resize 到 256 或 512若必须用任意尺寸上采样后做一次F.interpolate对齐到编码器特征尺寸再拼接。5.4 推理结果和训练时差很多现象训练时指标不错推理图一塌糊涂。 原因推理忘了model.eval()BN 还在用 batch 统计量或归一化参数和训练不一致。 解决推理前model.eval()加torch.no_grad()把训练用的 mean/std 固化成常量训练推理共用同一份。5.5 Docker 里 GPU 不可用现象容器内torch.cuda.is_available()返回 False。 原因宿主机没装 nvidia-container-toolkit或基础镜像 CUDA 版本和驱动不匹配。 解决宿主机装好 toolkitdocker run加--gpus all镜像 tag 的 CUDA 版本不高于驱动支持的上限。6. 把掩膜做成可用抠图阈值扫描与 alpha 羽化技巧跑通推理只是起点真正交付的抠图要边缘自然。我一般会做两件事阈值扫描和 alpha 羽化。阈值扫描是找最佳thresh。固定一张验证图遍历 0.3 到 0.7算预测掩膜和真值的 IoU取峰值import numpy as np import torch def scan_threshold(prob, target): best_t, best_iou 0.5, 0.0 for t in np.arange(0.3, 0.71, 0.05): pred (prob t).astype(np.uint8) inter (pred target).sum() union (pred | target).sum() iou inter / (union 1e-6) if iou best_iou: best_iou, best_t iou, t return best_t, best_iou逻辑说明prob是 sigmoid 后的概率图target是 0/1 真值。遍历步长 0.05 够用太细收益递减。参数上1e-6防除零。得到的最佳阈值写进推理脚本别硬编码 0.5。alpha 羽化解决硬边。二值掩膜直接当 alpha边缘是锯齿。做法是对概率图做一次小半径高斯模糊再当 alphaimport cv2 alpha (prob * 255).astype(np.uint8) alpha cv2.GaussianBlur(alpha, (5, 5), 0) # 半径 5边缘过渡 2 到 3 像素 rgba[:, :, 3] alpha半径别太大5×5 足够超过 9×9 会把前景吃掉一圈。这一步对头发丝效果明显但会牺牲一点边缘锐度按场景取舍。还有一个容易忽略的点UNet 对细长结构发丝、手指缝天生偏弱因为下采样丢细节。如果业务对发丝要求高常见做法是在 UNet 输出后再接一个轻量 refine 网络或者用高分辨率输入加多尺度推理。这份资源包没带这部分属于进阶改造先把基线跑稳再动。从那以后我每次交付抠图模型都强制走一遍「阈值扫描 羽化半径对比」把两组参数固化进配置文件而不是凭感觉写 0.5。希望帮到你。本文还有配套的精品资源点击获取