"""方案 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")