头发分割实战：基于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 语义分割模型横向对比

训练任务定下来之后，最纠结的就是选哪个模型。当时市面上主流方案大概有这么几类：

2000多张数据集属于典型的小规模数据，在这个量级下，复杂Transformer结构很容易欠拟合或过拟合——不是训练不动，就是验证集涨到一定程度再也不动。DeepLabV3+的ASPP模块理论上能捕获多尺度上下文，但在数据量不够时优势不容易体现。

我最终选了 UNet + ResNet34 编码器。理由有三点：一是 UNet 的跳连接结构在浅层特征和深层语义之间做了显式桥接，对边缘细节的保留在小数据集上明显优于纯encoder-decoder结构；二是训练稳定，对学习率不敏感，适合快速迭代；三是ONNX导出和部署生态成熟，后续换TensorRT、TFLite都有现成方案。如果你不想自己搭，用现成的 segmentation-models-pytorch 库也可以，它有UNet、MAnet、Linknet等一堆现成模型和预训练权重，我的基础版本就是在它基础上改的。

3.2 实际训练配置参考

训练细节可以参考我跑稳定的一组配置：

• 输入分辨率：384x384。最初试了512x512，边缘确实更精细，但训练速度慢了一半以上，而且显存占用从6G涨到11G；256x256又会在发梢处丢失太多细节。384是一个不错的折中。
• 归一化：使用ImageNet的mean和std，因为用了在ImageNet上预训练的ResNet34权重，这个必须对齐。
• 优化器：AdamW，初始学习率3e-4，weight_decay设1e-4，使用CosineAnnealingLR把学习率在80个epoch内降到接近0。
• Batch size：单张A6000上可以到16，8G显存的卡建议8。batch size太小会明显放大batch normalization的噪声，导致验证集曲线抖得厉害。
• Epoch：80轮，配合早停（patience=15，监控验证集Dice）。我的经验是大约40轮基本稳定，但后面20轮的微调往往能把发梢细节再往上推一两个点。
• 数据增强：随机水平翻转、随机裁剪（先随机crop到384x384再缩放）、旋转±10度、色彩抖动（亮度±20、对比度±15、饱和度±15）、少量随机遮挡（对小区域用灰色矩形块mask掉，模拟被手或衣领挡住的情况）。

3.3 损失函数设计：BCE+Dice为什么够了

这个任务我最终用了BCE + Dice Loss，权重比是0.5比0.5。

BCE（Binary Cross Entropy）对每个像素独立计算损失，它的优势是给模型提供稠密的梯度信号，让每个像素都参与学习。但问题也很典型：背景像素数量通常是头发像素的3到5倍，模型只要把大面积背景判断正确，BCE就能降到比较低，但这会掩盖头发区域的错误。

Dice Loss是根据预测mask和真实mask的区域重叠率来算的，天然对类别不均衡不敏感。它更关注“头发区域是否被完整找到”，缺点是在目标区域很小时梯度会不稳定。两者混合之后，BCE负责让每个像素有明确的收敛方向，Dice负责把整体注意力聚焦到头发区域上。

如果你遇到训练刚开始loss下降缓慢的情况，建议把Dice Loss加一个smooth=1e-5的平滑项，避免除零问题。另外强烈建议在验证阶段同时跑Dice和mIoU两个指标，因为二者侧重点不同——Dice更贴近“这个mask用户能不能直接用”，mIoU对误检更敏感。业务导向的项目可以以Dice为主。

4. 训练高频问题排查记录：边缘毛刺、暗色头发与类别失衡

4.1 边缘毛刺与二值化后处理

第一个明显的问题是：模型输出的预测图是概率值，但实际使用必须转成二值mask。直接mask = (prob > 0.5)会得到边缘充满锯齿的硬mask，尤其是发梢位置，锯齿感会严重到“假发片”效果。

排查之后我发现问题出在两个层面。一是标注阶段发梢处本身就很难标得精细，模型拟合的其实是“大致的、平滑的发梢区域”；二是训练用的下采样会把细碎发丝的标注平滑掉，恢复时也没法还原。这个问题的核心认知是：不可能也不应该追求每个发丝像素都和标注完全一致。

解决方案是按“区域完整、边缘平滑”的思路处理后处理流程。我最终的pipeline是：模型输出概率图 → 阈值0.45转二值 → 形态学闭运算（kernel 5x5，先膨胀后腐蚀，填补内部小黑洞）→ 高斯模糊（sigma=3的边缘羽化）。前两步是为了区域完整性，第三步是为了边缘过渡自然。具体在发色替换场景里，羽化边缘会让换色后和肤色、背景的过渡柔和很多，而不是一条硬边。

4.2 黑色头发漏检：颜色空间的局限

第二个坑是黑色头发频繁漏检。主要原因是很多训练图里，黑色头发和黑色背景、深色衣服在RGB空间的像素值非常接近，模型很难靠颜色区分，只能靠上下文推测。有时模型会直接把深色背景也预测成头发，有时又把黑色头发预测成背景。

这里我试验了两种有效缓解手段。一是在数据增强里加入强烈的亮度扰动和直方图均衡化，打乱颜色绝对值的稳定性，逼模型去学形状和纹理，而不是死记颜色。二是把YCrCb色彩空间作为一个额外输入通道喂给模型，因为YCrCb把亮度和色度解耦，在没有明显颜色区分时，模型可以重点利用Cr、Cb通道里肤色和头发的色度差异。实测把这两条加进去之后，深色头发场景的召回率提升了约4个百分点，代价是训练时间略微增加。

4.3 类别不均衡导致预测倾向背景

分类别占比统计了一下，头发平均占图像总像素的18%到25%。虽然不如很多医学图像里病灶只占1%那么极端，但已经足以让模型对头发区域“不敢下手”，典型的症状是预测mask面积明显比真实mask小，尤其在头发区域较小的人像里。

Dice Loss已经缓解了一部分，但我发现还不够。又加了两个手段：一是给BCE部分做类别加权，把头发类的权重设为背景的1.5倍；二是在训练时对头发区域占比低于20%的样本做了过采样，因为这类样本是最容易被模型忽略的困难样本。调完这版之后，小面积头发场景的误判率明显下降。

4.4 一轮调优实验记录

把我比较典型的一轮实验记录放出来，数值是GPU训练日志里的真实记录，大家做同样实验时可以参考变化趋势：

可以看到，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_shift=None): """通过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, ksize=3) 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_version=12, 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做图像编辑的条件约束，这些都是同一个底座上长出来的新应用。

本文还有配套的精品资源，点击获取