345 lines
16 KiB
Python
345 lines
16 KiB
Python
# -*- 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)
|