# -*- coding: utf-8 -*- """发际线区域 mask 自动识别（4 种方案） 需求约束： - 只取额头部分的发际线（前面，不要侧面/鬓角） - 下面不超过眉毛（用眉毛峰值点作下界） - 左右不超过耳朵/太阳穴（用太阳穴点作左右界） - 上面不超过已有头发（用头发mask的上边界作上界） 4 种方案： 1. boundary_band: 头发mask边界带（现有方案，沿头发外轮廓的带） 2. mediapipe: MediaPipe 478关键点 → 额头多边形 3. landmark_1k: 现有1k模型关键点 → 几何裁剪额头区 4. deeplab: 复用3分类DeepLab的skin/hair边界 所有方案输入输出统一： 输入: origin_matting(头发mask 0-255), landmarks_1k(np.array), img(BGR), band_width(int) 输出: (mask_8u, info_str, debug_images_dict) mask_8u: 发际线mask, 0或255, 与img同尺寸 info_str: 方案说明文字 debug_images_dict: {label: ndarray} 用于前端展示中间产物 """ import os import cv2 import numpy as np # MediaPipe 模型路径 MP_MODEL = "/home/ubuntu/change_hair/project/hair_service_sd/weights/mediapipe/face_landmarker.task" _mp_detector = None def detect_hairline(origin_matting, landmarks_1k, img, band_width=15, method="mediapipe", height_ratio=0.432, width_ratio=0.144, corner_ratio=0.25, vertical_offset=0.0): """统一入口：按 method 调用对应方案。 height_ratio/width_ratio/corner_ratio/vertical_offset 仅 landmark_1k 方案使用。 vertical_offset: 重绘区上下平移比例（相对额头高度，正值下移，负值上移） 返回 (mask_8u, info_str, debug_images) """ if method == "boundary_band": return _method_boundary_band(origin_matting, landmarks_1k, band_width) elif method == "mediapipe": return _method_mediapipe(img, band_width) elif method == "landmark_1k": return _method_landmark_1k(origin_matting, landmarks_1k, band_width, height_ratio, width_ratio, corner_ratio, vertical_offset) elif method == "deeplab": return _method_deeplab(img, origin_matting, landmarks_1k, band_width) else: raise ValueError(f"未知方法: {method}，可选: boundary_band/mediapipe/landmark_1k/deeplab") # ============================================================ # 方案1: 边界带法（现有方案） # ============================================================ def _method_boundary_band(origin_matting, landmarks_1k, band_width): """头发mask的形态学边界带（沿整个头发外轮廓的带，不限于额头）。 注意：这是最宽的方案，会包含鬓角/发尾的边界。 """ bw = max(1, int(band_width)) kernel = np.ones((bw, bw), np.uint8) mask = _boundary_band(origin_matting, kernel) info = f"边界带法(band_width={bw})：沿整个头发外轮廓的带，带宽≈{2*bw}px。注意：会包含鬓角/侧面边界" debug = {"原头发mask": origin_matting, "边界带(重绘区)": mask} return mask, info, debug # ============================================================ # 方案2: MediaPipe 478关键点 → 额头多边形 # ============================================================ def _method_mediapipe(img, band_width): """用 MediaPipe Face Mesh 的额头/太阳穴/眉毛关键点构造额头多边形。 额头多边形顶点（MediaPipe索引）： 发际线侧：10(顶), 151(上沿), 104(右太阳穴上), 333(左太阳穴上), 346(右发际侧) 眉毛峰值（下界）：70(左眉峰), 300(右眉峰) 眉间（下界中）：8(眉间上) 多边形 = [左太阳穴上 → 顶 → 右太阳穴上 → 右发际侧 → 右眉峰 → 眉间 → 左眉峰 → 左太阳穴下] 然后膨胀成带（band_width）。 """ global _mp_detector try: if _mp_detector is None: import mediapipe as mp from mediapipe.tasks import python from mediapipe.tasks.python import vision os.environ.setdefault("GLOG_minloglevel", "3") # 抑制 mediapipe 日志 base = python.BaseOptions(model_asset_path=MP_MODEL) opts = vision.FaceLandmarkerOptions(base_options=base, num_faces=1) _mp_detector = vision.FaceLandmarker.create_from_options(opts) import mediapipe as mp h, w = img.shape[:2] mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) res = _mp_detector.detect(mp_img) if not res.face_landmarks: return _fallback_empty("MediaPipe未检测到人脸") lm = res.face_landmarks[0] # 额头多边形关键点索引（顺时针） # 左太阳穴上(104) → 顶(10) → 右太阳穴上(333) → 右发际侧(346) # → 右眉峰(300) → 眉间(8) → 左眉峰(70) → 回到左太阳穴 poly_indices = [104, 103, 67, 109, 10, 337, 332, 333, 299, 346, 300, 293, 8, 63, 105, 107, 70, 63] poly_indices = [104, 109, 10, 338, 333, 346, 300, 8, 70] # 精简版 pts = [] for idx in poly_indices: if idx < len(lm): p = lm[idx] pts.append((int(p.x * w), int(p.y * h))) if len(pts) < 3: return _fallback_empty("MediaPipe关键点不足") # 画多边形 mask = np.zeros((h, w), dtype=np.uint8) pts_arr = np.array(pts, dtype=np.int32) cv2.fillPoly(mask, [pts_arr], 255) # 关键点可视化（用于debug图） vis = img.copy() for i, idx in enumerate(poly_indices): if idx < len(lm): p = lm[idx] x, y = int(p.x * w), int(p.y * h) cv2.circle(vis, (x, y), 4, (0, 255, 255), -1) cv2.putText(vis, str(idx), (x+4, y-4), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1) cv2.polylines(vis, [pts_arr], True, (0, 255, 0), 2) # 膨胀成带（让边缘羽化） bw = max(1, int(band_width)) mask_band = cv2.dilate(mask, np.ones((bw, bw), np.uint8)) info = f"MediaPipe法(478点)：额头多边形(顶点10/104/333等)+眉毛峰70/300作下界，膨胀带={bw}px。约束：下不超眉毛，左右不超太阳穴" debug = {"MediaPipe关键点+多边形": vis, "额头多边形": mask, "膨胀后(重绘区)": mask_band} return mask_band, info, debug except Exception as e: return _fallback_empty(f"MediaPipe失败: {e}") # ============================================================ # 方案3: 现有1k关键点 → 圆角矩形额头区 # ============================================================ def _method_landmark_1k(origin_matting, landmarks_1k, band_width, height_ratio=0.432, width_ratio=0.144, corner_ratio=0.25, vertical_offset=0.0): """用现有1k模型的关键点（眉毛/眼睛/脸轮廓）构造【圆角矩形】额头区。 形状参数（可调）： - height_ratio: 高度缩放比（相对原矩形高度，越小越窄） - width_ratio: 左右各扩展比例（相对眉宽，越大越宽） - corner_ratio: 圆角半径比例（相对短边） - vertical_offset: 上下平移比例（相对额头高度，正值下移） 下界：眉毛最高点（137索引 121:129 右眉, 129:137 左眉） 左右界：脸轮廓太阳穴（向左右扩展） 上界：头发mask的上边界 """ from utils.landmark_processor import pts_1k_to_137 try: lm137 = pts_1k_to_137(np.asarray(landmarks_1k)) h, w = origin_matting.shape[:2] # 下界：眉毛峰值（137索引 121:129 右眉, 129:137 左眉） brow_y = [] for seg in [lm137[121:129], lm137[129:137]]: if len(seg) > 0: brow_y.append(int(seg[:, 1].min())) # 眉毛最高点(y最小) if not brow_y: return _fallback_empty("1k关键点无眉毛点") brow_top = min(brow_y) # 取最高的眉毛点 # 左右界：眉尾向外扩展（左右长一点） left_brow_x = int(lm137[129:137][:, 0].min()) # 左眉最左 right_brow_x = int(lm137[121:129][:, 0].max()) # 右眉最右 brow_width = right_brow_x - left_brow_x # 左右各扩展 width_ratio（向两侧扩，默认0.144） extend = int(brow_width * width_ratio) left = max(0, left_brow_x - extend) right = min(w, right_brow_x + extend) # 上界：头发mask在左右界范围内每列的最上像素 sub_hair = origin_matting[:, left:right+1] if right > left else origin_matting hair_rows = np.where(sub_hair.max(axis=1) > 30)[0] hair_top = int(hair_rows.min()) if len(hair_rows) > 0 else 0 # 基础矩形 raw_top = max(0, hair_top) raw_bottom = min(h, brow_top) raw_left = max(0, left) raw_right = min(w, right) if raw_bottom <= raw_top or raw_right <= raw_left: return _fallback_empty("额头矩形无效") # ★ 上下短一点：以矩形中心为基准，高度缩到 height_ratio cy = (raw_top + raw_bottom) / 2.0 half_h = (raw_bottom - raw_top) / 2.0 * height_ratio # ★ vertical_offset 上下平移（正值下移，相对额头高度） forehead_h = raw_bottom - raw_top offset_px = forehead_h * vertical_offset top = int(max(0, cy - half_h + offset_px)) bottom = int(min(h, cy + half_h + offset_px)) # ★ 圆角矩形 corner_radius = int(min(bottom - top, raw_right - raw_left) * corner_ratio) mask = _rounded_rect((h, w), raw_left, top, raw_right, bottom, corner_radius) # 可视化关键点 + 圆角矩形描边 vis = np.zeros((h, w, 3), dtype=np.uint8) vis[:, :, 1] = origin_matting # 头发mask用绿色显示 cv2.rectangle(vis, (raw_left, raw_top), (raw_right, raw_bottom), (60, 60, 60), 1) # 原矩形(灰) # 描圆角矩形边 contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) cv2.drawContours(vis, contours, -1, (0, 0, 255), 2) # 圆角矩形(红) for seg, color in [(lm137[121:129], (255, 255, 0)), (lm137[129:137], (255, 255, 0))]: for p in seg: cv2.circle(vis, (int(p[0]), int(p[1])), 3, color, -1) # 膨胀 bw = max(1, int(band_width)) mask_band = cv2.dilate(mask, np.ones((bw, bw), np.uint8)) info = (f"1k关键点法(圆角矩形)：原矩形[{raw_left},{raw_top},{raw_right},{raw_bottom}] → " f"高度×{height_ratio}(上下短) 宽度+{width_ratio*2*100:.0f}%(左右长) 圆角r={corner_radius} " f"下移{offset_px:.0f}px(offset={vertical_offset})。" f"下界=眉毛峰 左右界=眉尾扩展 上界=头发顶") debug = {"1k关键点+圆角矩形(灰=原矩形/红=优化后)": vis, "圆角矩形mask": mask, "膨胀后(重绘区)": mask_band} return mask_band, info, debug except Exception as e: return _fallback_empty(f"1k关键点失败: {e}") def _rounded_rect(size, x1, y1, x2, y2, r): """画圆角矩形 mask（白色填充）。r=圆角半径。""" h, w = size mask = np.zeros((h, w), dtype=np.uint8) r = max(1, min(r, (x2 - x1) // 2, (y2 - y1) // 2)) # 中间矩形（去掉四角的圆角区） cv2.rectangle(mask, (x1 + r, y1), (x2 - r, y2), 255, -1) cv2.rectangle(mask, (x1, y1 + r), (x2, y2 - r), 255, -1) # 四个角的扇形 cv2.ellipse(mask, (x1 + r, y1 + r), (r, r), 180, 0, 90, 255, -1) # 左上 cv2.ellipse(mask, (x2 - r, y1 + r), (r, r), 270, 0, 90, 255, -1) # 右上 cv2.ellipse(mask, (x1 + r, y2 - r), (r, r), 90, 0, 90, 255, -1) # 左下 cv2.ellipse(mask, (x2 - r, y2 - r), (r, r), 0, 0, 90, 255, -1) # 右下 return mask # ============================================================ # 方案4: 复用3分类DeepLab（skin/hair边界） # ============================================================ def _method_deeplab(img, origin_matting, landmarks_1k, band_width): """用项目现有的3分类DeepLab（bg/skin/hair）找skin与hair的交界线。 发际线 = skin区域中紧邻hair的像素带。 用 Evaluator(gpu, nclass=3).eval(image, landmark1k) 得到分割图。 """ _deeplab_model = getattr(_method_deeplab, "_model", None) try: import torch from core.seg.hairseg_single_model import Evaluator if _deeplab_model is None: model_root = os.path.join(os.path.dirname(__file__), "weights") model_path = os.path.join(model_root, "deeplabv3_hair512_360_0520_wl.pth") _deeplab_model = Evaluator(gpu_id=0, output_img_size=512, nclass=3, seg_model_path=model_path) _method_deeplab._model = _deeplab_model # 缓存 with torch.no_grad(): seg = _deeplab_model.eval(img, np.asarray(landmarks_1k)) # 0=bg, 128=skin, 255=hair if seg is None: return _fallback_empty("DeepLab分割失败") skin = ((seg > 60) & (seg < 200)).astype(np.uint8) * 255 # skin hair = (seg > 200).astype(np.uint8) * 255 # hair # 发际线带 = skin侧紧邻hair的带：skin膨胀后 ∩ hair的边界带 bw = max(2, int(band_width)) kernel = np.ones((bw, bw), np.uint8) hair_band = _boundary_band(hair, kernel) # hair边界带 skin_dilate = cv2.dilate(skin, kernel) # skin膨胀 hairline = cv2.bitwise_and(hair_band, skin_dilate) # 二者交集≈发际线 # 用眉毛点约束下界（不跑到眉毛以下） from utils.landmark_processor import pts_1k_to_137 lm137 = pts_1k_to_137(np.asarray(landmarks_1k)) brow_y = [] for s in [lm137[121:129], lm137[129:137]]: if len(s) > 0: brow_y.append(int(s[:, 1].min())) if brow_y: brow_top = min(brow_y) hairline[brow_top:, :] = 0 # 可视化分割结果（伪彩色） seg_vis = cv2.applyColorMap(seg, cv2.COLORMAP_JET) info = f"DeepLab法(3分类)：skin/hair交界线带(band={bw}px)，下界=眉毛峰。复用现有triseg模型，无需新依赖" debug = {"3分类(蓝=bg/绿=skin/红=hair)": seg_vis, "skin/hair交界带(发际线)": hairline} return hairline, info, debug except ImportError as e: return _fallback_empty(f"DeepLab模块未找到: {e}") except Exception as e: return _fallback_empty(f"DeepLab失败: {e}") # ============================================================ # 工具函数 # ============================================================ def _boundary_band(mask, kernel): """形态学梯度边界带：dilate - erode""" binary = (mask > 10).astype(np.uint8) * 255 dilated = cv2.dilate(binary, kernel) eroded = cv2.erode(binary, kernel) return cv2.subtract(dilated, eroded) def _fallback_empty(reason): """方案失败时的空mask回退""" return None, f"⚠️ {reason}", {} def make_overlay(img, mask, alpha=0.45, color=(0, 0, 255)): """把 mask 半透明叠加到 img 上（红色高亮重绘区域）。 参数: img: BGR ndarray mask: 灰度mask (0/255 或 0-255) alpha: 透明度 0~1（mask区域的颜色强度） color: 高亮颜色 BGR，默认红色 返回: 叠加后的 BGR ndarray """ if mask is None: return img.copy() m = (mask > 10).astype(np.float32) if m.ndim == 3: m = m[:, :, 0] overlay = img.copy().astype(np.float32) # mask区域混入颜色 for c in range(3): overlay[:, :, c] = overlay[:, :, c] * (1 - m * alpha) + color[c] * (m * alpha) # mask边界描边（更清晰） m_u8 = (m * 255).astype(np.uint8) contours, _ = cv2.findContours(m_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) cv2.drawContours(overlay.astype(np.uint8), contours, -1, color, 1) return np.clip(overlay, 0, 255).astype(np.uint8)