ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

头发分割实战:基于UNet的小样本语义分割全流程解析

头发分割实战:基于UNet的小样本语义分割全流程解析 简介语义分割是计算机视觉中的核心任务之一其目标是对图像中的每个像素进行分类从而实现精细的区域划分。与目标检测的矩形框和图像分类的粗粒度标签不同语义分割能够输出像素级的mask在美颜、虚拟试戴、人像编辑等场景中具有重要价值。头发分割便是一个典型的像素级应用但由于头发与背景边界模糊、发丝细节复杂传统方法难以取得理想效果。以UNet为代表的编码器-解码器结构通过跳连接融合浅层与深层特征在小样本数据集上表现出色。本文从2000余张头发分割数据集的构建出发涵盖标注规范、模型选型、损失函数设计、训练调优以及部署落地等关键环节结合真实业务场景中的边缘处理与发色替换技巧为从事人像分割和语义分割的工程师提供一套可复用的工程参考。 做头发分割这个项目最初的动机其实很朴素美发App要做一个试发色功能用户上传一张照片后得先把头发区域精确抠出来才能改色。我一开始想走捷径直接靠人脸关键点去猜头发区域试了几轮效果都很勉强——发际线、刘海、披肩发这些不规则边界靠几个关键点根本兜不住。最后老老实实回到语义分割这条路把头发区域逐像素标注出来这才真正解决问题。这篇文章会从数据集构建、标注规范、模型选型、训练调优到部署落地完整复盘一遍重点包括一份2000多张的头发语义分割数据集和对应标签是怎么整理出来的UNet系列模型怎么训练才能在实际场景中稳定输出干净mask以及这个过程中踩过的坑和沉下来的经验。适合两类人一类是原本做目标检测或图像分类、想转向语义分割的新手另一类是正在做人像相关分割项目的工程师尤其是想找“小数据集也能跑通”的完整参考案例的人。1. 为什么单独做头发分割应用需求与任务定义1.1 真实业务场景里的头发分割需求头发分割在计算机视觉体系里不算大任务但实际需求密度很高。除了美发App试色还有美颜相机的“蓬松发际线”“发缝修复”电商平台假发与发饰的虚拟试戴直播软件里的虚拟形象换装甚至还有影视后期里给演员做头发的局部合成。这类功能听起来都只是“把头发抠出来”但落到工程上每一步都有完全不同的技术选型。拿试发色来说用户要的不是一个粗糙的选框而是每一根发丝边缘都干净的像素级区域。换色时如果不小心把额头或耳朵边缘带进去效果立刻就是“假发片既视感”。同样美颜里的“发际线调整”要移动的是发际线内侧和外侧的边界区域如果mask边缘不够平滑调整后会出现明显的锯齿和断裂感。这些需求都导向同一个技术方向像素级语义分割而不是目标检测那种矩形框或者图像分类那种粗粒度结果。1.2 任务边界语义分割不是检测也不是抠图把任务限定清楚后面很多选择才不会跑偏。头发语义分割的定义是给定一张输入图像对每个像素做二分类判断——属于头发还是不属于头发输出一张与输入图同尺寸的 mask习惯上用 0 表示背景、255 表示头发区域保存为单通道 8-bit PNG比直接用 0/1 更通用很多标注工具和可视化代码默认都认 255。这个任务看起来简单但有个容易被忽视的点头发和背景之间没有清晰的物理边界。人的衣服有轮廓、车有硬边但头发是细丝状材质发梢处大量像素是半透明的和背后的墙、天空、衣服混在一起。所以头发分割天然是个“边缘难度大于区域难度”的任务。这一点会贯穿后续所有选择为什么标注规则要反复强调边缘判断为什么训练阶段要在边缘和发梢上加额外注意力为什么后处理阶段要花大力气做羽化和形态学修正。另外要区分一下语义分割和实例分割本任务只需要所有头发像素归为一类不需要区分“这缕是额前发、那缕是后侧发”。如果以后要做“分层染发”“挑染设计”那需要的是实例或更细粒度的类别划分属于任务升级不在这次范围里。2. 2000多张头发分割数据集的构建全流程2.1 数据来源与筛选做分割任务数据质量直接决定模型上限。我整理的这套数据集最终是2080张可用图像加对应标签比标题里“2000多张”略多一点。这2000多张图不是一次性下载下来就能用的中间经过了好几轮筛选和清洗。来源主要有三块。第一块是公开人脸数据集中带分割标注的部分比如CelebAMask-HQ这类它本身有人脸部件标注可以把hair那一类的mask单独抽出来用第二块是从FFHQ等高质量人像数据里挑出合适图片自己标注第三块是少量自采和合作方提供的合规授权人像图用来补充一些特殊场景——比如戴帽子、湿发、卷发、白发这类公开数据集里比例较低的样本。整个筛选过程的重点是去重和去低质量图公开数据里同一个人的多张不同表情图训练时如果不过滤会让验证集失灵模糊图、极端侧脸、遮挡过半的图要么删掉要么在标注时把被遮挡的头发区域明确标为背景。这里有个很关键的经验数据集规模不是唯一指标分布覆盖比数量重要得多。2000张图如果全是白墙背景、黑长直发训练出来一到真实场景立刻效果崩坏。我后面在调优阶段发现往数据集里补了约150张复杂背景图之后测试集Dice比单纯增加500张同类图提升更明显。2.2 标注工具与人工修正流程标注工具这块我试过 Labelme、EISeg 和 Photoshop 三条路子最后实际量产采用的是“EISeg自动预标注 Labelme人工修正”的组合。EISeg 是百度开源的交互式分割标注工具背后是RITM等交互式分割模型。它可以先用点提示快速生成一个初始 mask把头发主体大致圈出来然后导出到 Labelme 里做精细修正。这个流程比完全从零用多边形一点点描要快很多尤其是在大面积平整的头发区域预标注能省掉80%的精力人工只需要专注修正边缘发丝和误检区域。Labelme 里我们统一用多边形polygon标注头发外轮廓导出为 JSON 格式然后再用脚本把多边形渲染成 PNG mask。渲染时我踩过一个坑Labelme 导出的 JSON 里坐标是小数直接cv2.fillPoly会因为在图像边界外的坐标出错必须先做取整和越界裁剪。这个细节在数据量小的时候看不出来但批量处理几千张时就是硬伤。双人交叉抽检也值得做。把标注好的 mask 叠到原图上以半透明红色显示头发区域抽检人专门看三类问题耳朵和发际线附近是否误标、发梢是否被粗暴截断、头顶的高光区是否被当成背景漏掉。抽检比例建议20%如果一个人的误标率超过5%他标的数据全部退回重做。之前我为了赶进度跳过这步结果训练出来模型在发际线处的假阳性明显偏高后来排查到是一批标签里额头皮肤被大面积标成了头发返工成本远高于一开始就抽检。2.3 标签格式与目录组织数据结构建议按下面的方式组织这也是目前比较通用的语义分割数据集布局hair_dataset/ ├── images/ │ ├── train/ │ │ ├── 000001.jpg │ │ └── ... │ ├── val/ │ └── test/ ├── masks/ │ ├── train/ │ │ ├── 000001.png │ │ └── ... │ ├── val/ │ └── test/ └── vis/ ├── train/ │ ├── 000001.jpg │ └── ...masks 目录下每张 PNG 是和对应 JPG 同名的单通道 mask像素值只允许出现 0 或 255。vis目录用于可视化检查把原图和 mask 按半透明方式叠加方便快速翻看标注质量。划分比例上最终1080张做训练、300张做验证、700张留作测试。这个划分不是随机切而是按人划分同一人的所有图片只能出现在一个集合里否则模型会通过记忆人脸特征“作弊”验证指标会虚高得离谱。测试集单独从合作方和自采数据里留出和训练集分布有略微差异这样测出来的数才有参考价值。2.4 质量验收哪些图被淘汰数据处理阶段我淘汰了大约300张图标准可以给读者参考图像分辨率低于800x600的放大后边缘标注难度大直接不要头发区域小于整张图10%的比如远距离全身照分割价值低且类别不均衡问题会被放大头发被大面积遮挡、无法判断真实边界的标了也是带偏模型多人同框、头发互相重叠的语义边界难以定义暂时排除JPEG压缩痕迹过重的边缘会出现明显的块状伪影。淘汰之后保留2080张覆盖了黑、棕、金、红、白/灰等常见发色直发、卷发、波浪发、盘发、短发、光头作为全背景负样本等发型以及室内、街景、纯色背景、复杂背景等环境。负样本无头发图片我特意放了60张这对抑制背景误检很有效。3. 模型选型与训练配置选UNet系是理性的选择3.1 语义分割模型横向对比训练任务定下来之后最纠结的就是选哪个模型。当时市面上主流方案大概有这么几类模型参数量在小数据集表现部署难度典型场景UNet / UNet30M-50M很好低医学图像、小样本分割DeepLabV340M-60M好中通用语义分割SegFormer-B04M-15M一般中轻量级分割BiSeNetV23M-10M好低实时人像分割YOLOv8-seg10M-50M中低实例分割2000多张数据集属于典型的小规模数据在这个量级下复杂Transformer结构很容易欠拟合或过拟合——不是训练不动就是验证集涨到一定程度再也不动。DeepLabV3的ASPP模块理论上能捕获多尺度上下文但在数据量不够时优势不容易体现。我最终选了 UNet ResNet34 编码器。理由有三点一是 UNet 的跳连接结构在浅层特征和深层语义之间做了显式桥接对边缘细节的保留在小数据集上明显优于纯encoder-decoder结构二是训练稳定对学习率不敏感适合快速迭代三是ONNX导出和部署生态成熟后续换TensorRT、TFLite都有现成方案。如果你不想自己搭用现成的 segmentation-models-pytorch 库也可以它有UNet、MAnet、Linknet等一堆现成模型和预训练权重我的基础版本就是在它基础上改的。3.2 实际训练配置参考训练细节可以参考我跑稳定的一组配置输入分辨率384x384。最初试了512x512边缘确实更精细但训练速度慢了一半以上而且显存占用从6G涨到11G256x256又会在发梢处丢失太多细节。384是一个不错的折中。归一化使用ImageNet的mean和std因为用了在ImageNet上预训练的ResNet34权重这个必须对齐。优化器AdamW初始学习率3e-4weight_decay设1e-4使用CosineAnnealingLR把学习率在80个epoch内降到接近0。Batch size单张A6000上可以到168G显存的卡建议8。batch size太小会明显放大batch normalization的噪声导致验证集曲线抖得厉害。Epoch80轮配合早停patience15监控验证集Dice。我的经验是大约40轮基本稳定但后面20轮的微调往往能把发梢细节再往上推一两个点。数据增强随机水平翻转、随机裁剪先随机crop到384x384再缩放、旋转±10度、色彩抖动亮度±20、对比度±15、饱和度±15、少量随机遮挡对小区域用灰色矩形块mask掉模拟被手或衣领挡住的情况。3.3 损失函数设计BCEDice为什么够了这个任务我最终用了BCE Dice Loss权重比是0.5比0.5。BCEBinary Cross Entropy对每个像素独立计算损失它的优势是给模型提供稠密的梯度信号让每个像素都参与学习。但问题也很典型背景像素数量通常是头发像素的3到5倍模型只要把大面积背景判断正确BCE就能降到比较低但这会掩盖头发区域的错误。Dice Loss是根据预测mask和真实mask的区域重叠率来算的天然对类别不均衡不敏感。它更关注“头发区域是否被完整找到”缺点是在目标区域很小时梯度会不稳定。两者混合之后BCE负责让每个像素有明确的收敛方向Dice负责把整体注意力聚焦到头发区域上。如果你遇到训练刚开始loss下降缓慢的情况建议把Dice Loss加一个smooth1e-5的平滑项避免除零问题。另外强烈建议在验证阶段同时跑Dice和mIoU两个指标因为二者侧重点不同——Dice更贴近“这个mask用户能不能直接用”mIoU对误检更敏感。业务导向的项目可以以Dice为主。4. 训练高频问题排查记录边缘毛刺、暗色头发与类别失衡4.1 边缘毛刺与二值化后处理第一个明显的问题是模型输出的预测图是概率值但实际使用必须转成二值mask。直接mask (prob 0.5)会得到边缘充满锯齿的硬mask尤其是发梢位置锯齿感会严重到“假发片”效果。排查之后我发现问题出在两个层面。一是标注阶段发梢处本身就很难标得精细模型拟合的其实是“大致的、平滑的发梢区域”二是训练用的下采样会把细碎发丝的标注平滑掉恢复时也没法还原。这个问题的核心认知是不可能也不应该追求每个发丝像素都和标注完全一致。解决方案是按“区域完整、边缘平滑”的思路处理后处理流程。我最终的pipeline是模型输出概率图 → 阈值0.45转二值 → 形态学闭运算kernel 5x5先膨胀后腐蚀填补内部小黑洞→ 高斯模糊sigma3的边缘羽化。前两步是为了区域完整性第三步是为了边缘过渡自然。具体在发色替换场景里羽化边缘会让换色后和肤色、背景的过渡柔和很多而不是一条硬边。4.2 黑色头发漏检颜色空间的局限第二个坑是黑色头发频繁漏检。主要原因是很多训练图里黑色头发和黑色背景、深色衣服在RGB空间的像素值非常接近模型很难靠颜色区分只能靠上下文推测。有时模型会直接把深色背景也预测成头发有时又把黑色头发预测成背景。这里我试验了两种有效缓解手段。一是在数据增强里加入强烈的亮度扰动和直方图均衡化打乱颜色绝对值的稳定性逼模型去学形状和纹理而不是死记颜色。二是把YCrCb色彩空间作为一个额外输入通道喂给模型因为YCrCb把亮度和色度解耦在没有明显颜色区分时模型可以重点利用Cr、Cb通道里肤色和头发的色度差异。实测把这两条加进去之后深色头发场景的召回率提升了约4个百分点代价是训练时间略微增加。4.3 类别不均衡导致预测倾向背景分类别占比统计了一下头发平均占图像总像素的18%到25%。虽然不如很多医学图像里病灶只占1%那么极端但已经足以让模型对头发区域“不敢下手”典型的症状是预测mask面积明显比真实mask小尤其在头发区域较小的人像里。Dice Loss已经缓解了一部分但我发现还不够。又加了两个手段一是给BCE部分做类别加权把头发类的权重设为背景的1.5倍二是在训练时对头发区域占比低于20%的样本做了过采样因为这类样本是最容易被模型忽略的困难样本。调完这版之后小面积头发场景的误判率明显下降。4.4 一轮调优实验记录把我比较典型的一轮实验记录放出来数值是GPU训练日志里的真实记录大家做同样实验时可以参考变化趋势版本关键改动验证集Dice测试集Dice备注v1 baselineUNetBCE0.9380.895边缘毛糙v2BCEDice0.9550.919整体召回提升v3v2边界加权0.9620.930边缘明显改善v4v3颜色空间增强0.9680.941深色头发场景提升v5v4复杂背景补充0.9710.947测试集涨幅小于v4但业务场景观感提升很大可以看到v4到v5测试集数值涨幅不大但实际效果差异最明显这就是“指标没涨但体验涨了”的典型情况。做业务项目千万别只盯一个测试集指标定期在真实样本上人工看mask效果比刷0.1个点的mIoU有价值得多。5. 让mask真正可用发色替换与边缘羽化的实现细节5.1 发色替换的核心算法分割模型给了mask接下来就是产品功能逻辑了。拿最核心的“试发色”功能举例网上很多教程直接改RGB通道效果很生硬。我最终的方案是把原图转到HSV色彩空间只修改头发区域的Hue通道色调保留Saturation饱和度和Value明度这样发丝的高光、阴影和纹理都能保留看起来更像真实染发效果。代码逻辑大概是这样的import cv2 import numpy as np def recolor_hair(image_bgr, mask_gray, target_bgr, hue_shiftNone): 通过HSV空间进行发色替换 image_bgr: 原图BGR mask_gray: 单通道mask, 像素值0或255 target_bgr: 目标颜色BGR, 例如灰发(180, 180, 180) hue_shift: 若给定额外偏移量则优先使用偏移量 hsv cv2.cvtColor(image_bgr, cv2.COLOR_BGR2HSV) H, S, V cv2.split(hsv) target_hsv cv2.cvtColor(np.uint8([[target_bgr]]), cv2.COLOR_BGR2HSV)[0][0] H_new H.copy() if hue_shift is None: H_new np.full_like(H, target_hsv[0]) S_new np.clip(S * 1.0 (target_hsv[1] - S) * 0.75, 0, 255) else: H_new H.copy() H_new (H_new hue_shift) % 180 S_new S hsv_new cv2.merge([H_new, S_new, V]) result cv2.cvtColor(hsv_new, cv2.COLOR_HSV2BGR) # 只把mask区域替换mask边缘用羽化过渡 mask_float cv2.GaussianBlur(mask_gray.astype(np.float32), (7, 7), 3) / 255.0 mask_float mask_float[:, :, None] output image_bgr * (1 - mask_float) result * mask_float return np.uint8(output)这个函数里的mask_float是对mask做了高斯模糊后的渐变权值这样替换区域不是硬切而是有自然过渡。核心思路就一句话只换颜色相貌不换明暗光影。5.2 保留高光与发丝质感的小技巧实际做了几版demo后又发现了两个观感差异很大的细节。第一是发丝质感。直接把H通道统一替换为目标色相发丝之间原本的色彩层次几乎全没了看起来像块平整的色板。解决办法是只替换H通道目标色相S通道保留原图值V通道也就是明暗度根据原图不变再利用原图在灰度上的纹理作为alpha叠加一个细密的噪声或发丝层。这个纹理层可以用一个小的拉普拉斯滤波提取def hair_texture_detail(image_gray): laplacian cv2.Laplacian(image_gray, cv2.CV_32F, ksize3) texture cv2.GaussianBlur(laplacian, (3, 3), 1) return texture然后在替换后的颜色上叠加小幅度的纹理变化发丝的层次感就回来了。第二是发际线处的过渡。mask在额头的边界如果太硬换色后会像戴了顶帽子。这里我对mask的边界区域做了一次特殊的膨胀处理把mask的二值图先腐蚀再取差值作为“边缘带”边缘带内部做梯度过渡效果是发际线处的肤色和发色能自然融合。这个在美颜类赛道非常重要用户对发际线处的“假”特别敏感。5.3 生成式扩展从分割到换发型mask还能做更多事情。拿到精确的头发mask之后可以进一步计算头发的轮廓、面积、大致长度和发量占比再加上一个简单的分类器就能得到“长发/短发/中发”“直发/卷发/波浪发”这类发型的结构化描述。这些特征组合起来可以做“智能发型推荐”——用户上传照片后系统根据脸型、发量、头发长度推荐合适的发型而不是随机展示一堆结果。再往上走可以接生成模型用原图和头发mask做引导的换发型。这个时候语义分割的mask就不再是最终产品而是作为生成模型的输入条件用来锁定待修改区域让生成器只改头发附近的像素而不动五官和背景。这个方向我后来试过一阵比直接全图生成的效果好很多本质原因就是mask大幅缩小了生成的不确定性范围。6. 部署到端侧的两条路线与工程建议6.1 ONNX导出与CPU推理速度验证训练完的PyTorch模型要落地我的习惯是导出成ONNX后面无论是CPU、GPU、TensorRT还是移动端都从这个中间格式出发。import torch import onnxruntime as ort model.eval() dummy_input torch.randn(1, 3, 384, 384) torch.onnx.export( model, dummy_input, hair_seg.onnx, input_names[input], output_names[output], opset_version12, dynamic_axes{input: {0: batch}, output: {0: batch}}, )走ONNX的主要原因一是脱离PyTorch生态部署端不用装重量级的运行时二是onnxruntime的CPU推理速度比PyTorch原生的CPU推理快不少在桌面端的i5-12400上384x384的单帧推理大约在120ms到180ms之间这个速度对“用户上传照片→等待结果”的异步流程完全够用。如果还想更快可以尝试把输入分辨率降到320x320速度能进入100ms以内代价是发梢细节轻微损失。6.2 轻量级方案从UNet到BiSeNetV2/PP-LiteSeg如果目标平台是手机或者需要摄像头实时预览UNet的参数量和计算量就不合适了。这时候我在工程上会更倾向于下面两条路线。一是把UNet模型做蒸馏用上一节训练好的UNet作为teacher模型让一个小型模型比如BiSeNetV2去学它的输出而不是直接学真值标注。因为teacher模型的输出比人工标注更平滑小模型学起来难度更低分割效果能保持在teacher的95%以上。这里可以顺手用unet作为教师模型蒸馏出约1/8参数量的学生模型。二是直接用PP-LiteSeg这类专为移动端设计的模型它用到的Decoder结构比较简单模型只有几MB在手机端CPU上跑320x320分辨率大约在40ms到60ms。我用它重新训练了一版测试集Dice比UNet低了大概2个百分点但换来的是实时可用。如果业务对精度要求不高这个方案开发成本更低不用管蒸馏流程如果精度敏感蒸馏更好。6.3 部署时容易忽略的三个问题最后记录几个部署阶段常踩的坑。第一个是输入预处理必须和训练时完全一致。ImageNet的mean/std、归一化顺序、BGR还是RGB、缩放用的插值方式我训练时用cv2.INTER_LINEAR这些只要有一处不一致部署后精度会掉得莫名其妙。有一个经典坑训练用PIL读图是RGB部署用OpenCV读图是BGR忘了转过来效果立刻崩掉。第二个是ONNX导出的Dynamic axes。动态维度虽然灵活但部分推理引擎尤其是一些移动端框架对动态shape支持不完善可能莫名其妙地慢或者报错。如果业务里输入尺寸固定就干脆把shape固定死换来的是更多的优化机会和更稳的推理行为。第三个是推理结果的坐标系mask要和原图尺寸一致。如果模型输入是384x384而原图是1200x1600那么推理出来的mask要么先缩放到原图尺寸再做形态学处理要么在模型前处理时就同时缩放原图和mask。这个乍看不是问题但我在某个项目里见过同事把mask直接resize回原图后边缘出现了一格一格的马赛克因为他在模型输出之后才resize中间漏了一步双线性插值的平滑处理。我在实际部署中养成的习惯是任何一步预处理都写成和训练推理完全一致的函数并且在部署前用一张固定图片对比PyTorch输出和ONNX输出的差异数值差异应该低于0.01一旦超过就优先检查归一化和通道顺序。这个习惯帮我省下了大量线上排查时间。如果以后要继续扩展大家可以考虑把头发mask和头部关键点结合起来做发际线微调或者用mask做图像编辑的条件约束这些都是同一个底座上长出来的新应用。本文还有配套的精品资源点击获取
返回列表