208 lines
8.7 KiB
Python
208 lines
8.7 KiB
Python
"""方案 B：BiSeNet 头发/耳朵分割 + 发际线/头顶定位。
|
||
|
||
加载 face-parsing BiSeNet（CelebAMask-HQ 19 类，hair=17、l_ear=7、r_ear=8），对整图
|
||
做像素级语义分割：得到头发 mask 用于沿面部中轴线扫描真实发际线与头顶，并得到耳朵
|
||
mask 供标注图取人头最左/最右竖线（耳朵外缘）。GPU 可用时走 CUDA，否则 CPU。
|
||
单例加载权重，避免每请求重载。详见技术方案 §1.4 / §4.0。
|
||
"""
|
||
import os
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
# ⚠️ torch / torchvision / BiSeNet 仅在 HairSegmenter.__init__ 内惰性导入，
|
||
# 使本模块的纯 numpy 函数 locate_hairline_by_segmentation 可在无 torch/GPU
|
||
# 的环境（如 Tier-1 合成几何测试、方案 A only 降级版）被安全导入。
|
||
|
||
_WEIGHTS = os.path.join(os.path.dirname(__file__), "weights", "79999_iter.pth")
|
||
HAIR_CLASS = 17 # CelebAMask-HQ 19 类中 hair 的索引
|
||
EAR_CLASSES = (7, 8) # 7=l_ear / 8=r_ear（类名以人为参照，图像左右另行判定，不依赖类名）
|
||
N_CLASSES = 19
|
||
_INPUT_SIZE = 512 # BiSeNet 推理输入边长
|
||
|
||
|
||
def _select_device(torch):
|
||
"""选择推理设备：优先 CUDA，但实测一次小算子确认当前 GPU 架构被本 torch 支持。
|
||
|
||
场景：本机为 RTX 5090（sm_120/Blackwell），而 torch 2.2.2+cu121 仅编译到 sm_90，
|
||
.cuda() 会在执行时抛 "no kernel image is available"。此处用一次小 matmul 探测，
|
||
失败则回退 CPU（BiSeNet CPU 推理 ~0.3–1s/张，方案B 仍可用）。
|
||
换装支持 sm_120 的 torch（cu128）后会自动改用 GPU，无需改代码。
|
||
可用环境变量 FORCE_CPU=1 强制 CPU。
|
||
"""
|
||
import os as _os
|
||
if _os.getenv("FORCE_CPU") == "1" or not torch.cuda.is_available():
|
||
return torch.device("cpu")
|
||
try:
|
||
_ = (torch.zeros(8, 8, device="cuda") @ torch.zeros(8, 8, device="cuda")).cpu()
|
||
return torch.device("cuda")
|
||
except Exception: # noqa: BLE001 GPU 架构不被支持 → 回退 CPU
|
||
return torch.device("cpu")
|
||
|
||
|
||
class HairSegmenter:
|
||
"""BiSeNet 头发分割封装。建议经 get_segmenter() 取单例。"""
|
||
|
||
def __init__(self, weights_path=_WEIGHTS):
|
||
import torch
|
||
import torchvision.transforms as transforms
|
||
|
||
from face_analysis.bisenet_model import BiSeNet
|
||
|
||
self._torch = torch
|
||
self.device = _select_device(torch)
|
||
self.net = BiSeNet(n_classes=N_CLASSES)
|
||
state = torch.load(weights_path, map_location="cpu", weights_only=False)
|
||
self.net.load_state_dict(state)
|
||
self.net.to(self.device)
|
||
self.net.eval()
|
||
self._to_tensor = transforms.Compose([
|
||
transforms.ToTensor(),
|
||
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
|
||
])
|
||
|
||
def _parse(self, image_bgr):
|
||
"""整图语义分割，返回原图尺寸的类别图（H×W int，值为 0–18 类别号）。"""
|
||
torch = self._torch
|
||
h, w = image_bgr.shape[:2]
|
||
rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
|
||
resized = cv2.resize(rgb, (_INPUT_SIZE, _INPUT_SIZE),
|
||
interpolation=cv2.INTER_LINEAR)
|
||
inp = self._to_tensor(resized).unsqueeze(0).to(self.device)
|
||
with torch.no_grad():
|
||
out = self.net(inp)[0] # 主输出 (1, C, 512, 512)
|
||
parsing = out.squeeze(0).argmax(0).cpu().numpy() # (512, 512) 类别图
|
||
# 还原到原图尺寸（最近邻保持类别边界）
|
||
return cv2.resize(parsing.astype(np.int32), (w, h),
|
||
interpolation=cv2.INTER_NEAREST)
|
||
|
||
def segment_hair(self, image_bgr):
|
||
"""返回 hair_mask（H×W bool，True=头发），尺寸同输入原图。"""
|
||
return self._parse(image_bgr) == HAIR_CLASS
|
||
|
||
def _parse_face_cropped(self, image_bgr, face_box):
|
||
"""按人脸框裁剪后再分割，结果映射回原图尺寸（裁剪外填背景 0）。
|
||
|
||
BiSeNet 在 CelebAMask-HQ「紧裁对齐人脸」上训练，整张大场景图（全身/街拍，
|
||
脸只占一小块、背景复杂）会严重欠分割、丢耳朵。先按人脸放大裁剪，让脸接近
|
||
训练分布，耳朵/头发分割明显更稳。裁剪含足够上/侧边距以纳入发顶与双耳。
|
||
"""
|
||
h, w = image_bgr.shape[:2]
|
||
x0, y0, x1, y1 = face_box
|
||
fw, fh = max(1.0, x1 - x0), max(1.0, y1 - y0)
|
||
cx0 = int(max(0, x0 - fw * 0.8)); cx1 = int(min(w, x1 + fw * 0.8))
|
||
cy0 = int(max(0, y0 - fh * 1.0)); cy1 = int(min(h, y1 + fh * 0.5))
|
||
if cx1 - cx0 < 2 or cy1 - cy0 < 2:
|
||
return self._parse(image_bgr)
|
||
full = np.zeros((h, w), dtype=np.int32)
|
||
full[cy0:cy1, cx0:cx1] = self._parse(image_bgr[cy0:cy1, cx0:cx1])
|
||
return full
|
||
|
||
def segment_hair_and_ears(self, image_bgr, face_box=None):
|
||
"""单次推理返回 (hair_mask, ear_mask)，均为 H×W bool，尺寸同原图。
|
||
|
||
ear_mask = 左耳(7) ∪ 右耳(8)；耳朵被头发/侧脸遮挡时对应区域天然为空，
|
||
正好用于「看不到耳朵就不画线」的判定。两类合并、左右按图像位置另判，
|
||
不依赖以人为参照的类名（详见 EAR_CLASSES 注释）。
|
||
|
||
face_box=(x0,y0,x1,y1)（人脸关键点包围盒像素坐标）给定时先按人脸裁剪再
|
||
分割（见 _parse_face_cropped），整张大场景图也能稳定分出耳朵；不给则整图分割。
|
||
"""
|
||
if face_box is None:
|
||
parsing = self._parse(image_bgr)
|
||
else:
|
||
parsing = self._parse_face_cropped(image_bgr, face_box)
|
||
hair_mask = parsing == HAIR_CLASS
|
||
ear_mask = np.isin(parsing, EAR_CLASSES)
|
||
return hair_mask, ear_mask
|
||
|
||
|
||
_segmenter = None
|
||
|
||
|
||
def get_segmenter():
|
||
"""惰性单例：首次调用时加载权重（并占用显存），后续复用。"""
|
||
global _segmenter
|
||
if _segmenter is None:
|
||
_segmenter = HairSegmenter()
|
||
return _segmenter
|
||
|
||
|
||
def locate_hairline_by_segmentation(hair_mask, brow_center_x, image_height):
|
||
"""从头发 mask 定位发际线与头顶。
|
||
|
||
Args:
|
||
hair_mask: H×W bool/uint8，True=头发。
|
||
brow_center_x: 面部中轴线 x（像素）。
|
||
image_height: 图高（保留参数，便于后续边界判断）。
|
||
Returns:
|
||
(hairline_y, hair_top_y) 像素坐标；失败返回 None（交给方案 A 兜底）。
|
||
"""
|
||
if hair_mask is None:
|
||
return None
|
||
mask = np.asarray(hair_mask).astype(bool)
|
||
if mask.sum() == 0:
|
||
return None
|
||
|
||
w = mask.shape[1]
|
||
cx = int(round(brow_center_x))
|
||
cx = max(0, min(cx, w - 1))
|
||
# 中轴线附近窄列带（±3px）求稳，避免单列噪声
|
||
band = mask[:, max(0, cx - 3): min(w, cx + 4)]
|
||
col = band.any(axis=1)
|
||
hair_rows = np.where(col)[0]
|
||
if hair_rows.size == 0:
|
||
return None
|
||
|
||
# 发际线：中轴线列带上头发区域最靠下的行（头发→皮肤交界，y 向下为正）
|
||
hairline_y = int(hair_rows.max())
|
||
# 头顶：整张头发 mask 的最高点（最小 y），用全图更鲁棒
|
||
top_rows = np.where(mask.any(axis=1))[0]
|
||
hair_top_y = int(top_rows.min())
|
||
|
||
# 合理性校验：头顶必须严格在发际线上方
|
||
if hair_top_y >= hairline_y:
|
||
return None
|
||
return hairline_y, hair_top_y
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
from face_analysis.detector import detector
|
||
from face_analysis.calibration import normalized_to_pixel
|
||
from face_analysis.face_mesh_landmarks import GLABELLA_9, GLABELLA_151
|
||
|
||
path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
|
||
img = cv2.imread(path)
|
||
if img is None:
|
||
print(f"无法读取图片: {path}")
|
||
sys.exit(1)
|
||
h, w = img.shape[:2]
|
||
segmenter = get_segmenter()
|
||
print("device:", segmenter.device)
|
||
mask = segmenter.segment_hair(img)
|
||
print("hair pixels:", int(mask.sum()))
|
||
|
||
lms = detector.detect(img)
|
||
if lms is None:
|
||
print("未检出人脸，跳过定位")
|
||
sys.exit(0)
|
||
lm = lms.landmark
|
||
bx = (normalized_to_pixel(lm[GLABELLA_9], w, h)[0]
|
||
+ normalized_to_pixel(lm[GLABELLA_151], w, h)[0]) / 2
|
||
by = (normalized_to_pixel(lm[GLABELLA_9], w, h)[1]
|
||
+ normalized_to_pixel(lm[GLABELLA_151], w, h)[1]) / 2
|
||
res = locate_hairline_by_segmentation(mask, bx, h)
|
||
if res is None:
|
||
print("定位失败（将回退方案 A）")
|
||
else:
|
||
hairline_y, hair_top_y = res
|
||
print(f"brow_y={by:.1f} hairline_y={hairline_y} hair_top_y={hair_top_y}")
|
||
print("自洽校验 hair_top_y < hairline_y < brow_y:",
|
||
hair_top_y < hairline_y < by)
|
||
|
||
# dump mask 预览
|
||
os.makedirs("tests/output", exist_ok=True)
|
||
cv2.imwrite("tests/output/hair_mask.png", (mask.astype(np.uint8) * 255))
|
||
print("mask 预览已存 tests/output/hair_mask.png")
|