hair/hairline/mask.py

"""接口2 第二步：inpaint 遮罩 + 黑色发际线划线合成（参考 headmark 5步法）。

算法（用 hairline_texture_black 渲染黑线替代 headmark 的手绘检测）：
  ① 额头上部区域：MediaPipe 额头边界关键点连线，向上+两侧补到图像边缘填充
  ② 头部轮廓：SegFormer 头部类（hair∪skin∪…，排除 bg/neck/cloth）
  ③ ROI = ① ∩ ②
  ④ 渲染黑色发际线 → 烧进照片(marked) + 得到曲线像素
  ⑤ mask = ROI 中"发际线曲线以上"，闭运算去洞 + 最大连通域 + 轻羽化
合成 RGBA：RGB=marked，alpha=255×(1−mask)（透明=重绘区，对齐 ComfyUI mask=1−alpha）。
"""
from __future__ import annotations

import cv2
import numpy as np
from PIL import Image

from .render import load_ext_mesh, load_texture_rgba, render_hairline_overlay, build_overlay_layer

# headmark 额头边界关键点（MediaPipe canonical 索引，左→右沿上额）
FOREHEAD_LANDMARKS = [21, 68, 104, 69, 108, 151, 337, 299, 333, 298, 251]
# SegFormer 头部类（含 skin..hat；排除 bg=0 / ear_r=15 / neck_l=16 / neck=17 / cloth=18）
_HEAD_CLASSES = list(range(1, 15))


def forehead_upper_region(landmarks_mp: np.ndarray, w: int, h: int) -> np.ndarray:
    """headmark step1：额头边界关键点以上的"上部区域"填充 mask（uint8 0/255）。"""
    pts = [(int(landmarks_mp[i, 0] * w), int(landmarks_mp[i, 1] * h)) for i in FOREHEAD_LANDMARKS]
    left_ext = (0, pts[0][1])
    right_ext = (w - 1, pts[-1][1])
    polygon = np.array([left_ext] + pts + [right_ext, (w - 1, 0), (0, 0)], dtype=np.int32)
    m = np.zeros((h, w), np.uint8)
    cv2.fillPoly(m, [polygon], 255)
    return m


def head_silhouette(parse_map: np.ndarray) -> np.ndarray:
    """headmark step2：SegFormer 头部轮廓 mask（uint8 0/255）。"""
    return (np.isin(parse_map, _HEAD_CLASSES).astype(np.uint8) * 255)


def _curve_bottom_per_column(curve_mask: np.ndarray):
    """每列发际线曲线的**最低**像素 y（线下沿），返回 (xs, ys) 仅含有曲线的列。"""
    ys_idx, xs_idx = np.where(curve_mask > 0)
    if xs_idx.size == 0:
        return None, None
    w = curve_mask.shape[1]
    bottom = np.full(w, -1, np.int32)
    np.maximum.at(bottom, xs_idx, ys_idx)
    cols = np.where(bottom >= 0)[0]
    return cols, bottom[cols]


def _above_curve_region(curve_mask: np.ndarray, h: int, w: int) -> np.ndarray:
    """由发际线曲线得到"曲线以上"区域（uint8 0/255）。

    曲线 x 跨度内逐列插值出下沿 y_line(x)，两侧按端点 y 水平延伸；
    above = 所有 y ≤ y_line(x)。曲线缺失（极端）则返回全 1（交给 ROI 兜底）。
    """
    cols, ybot = _curve_bottom_per_column(curve_mask)
    if cols is None:
        return np.full((h, w), 255, np.uint8)
    x0, x1 = int(cols.min()), int(cols.max())
    # 全列插值 y_line：[x0,x1] 内线性插值，两侧水平延伸
    yline = np.interp(np.arange(w), cols, ybot,
                      left=float(ybot[0]), right=float(ybot[-1])).astype(np.int32)
    yy = np.arange(h)[:, None]                  # (h,1)
    above = (yy <= yline[None, :]).astype(np.uint8) * 255   # (h,w)
    return above


def _clean_mask(mask: np.ndarray, w: int) -> np.ndarray:
    """闭运算去洞 + 取最大连通域填充 + 轻羽化。"""
    k = max(3, (int(w * 0.015) | 1))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
    closed = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    cnts, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    out = np.zeros_like(mask)
    if cnts:
        largest = max(cnts, key=cv2.contourArea)
        cv2.drawContours(out, [largest], -1, 255, -1)
    # 轻羽化（柔化边缘，利于扩散衔接）
    out = cv2.GaussianBlur(out, (0, 0), sigmaX=max(1.0, w * 0.004))
    return out


def build_inpaint_mask(photo_bgr: np.ndarray, landmarks_mp: np.ndarray,
                       parse_map: np.ndarray, points502: np.ndarray,
                       black_texture_rgba: np.ndarray):
    """返回 (marked_bgr 划线图, mask uint8 0..255 重绘区)。"""
    h, w = photo_bgr.shape[:2]
    uv, ext_faces = load_ext_mesh()

    # ④ 渲染黑线：marked = 烧进照片；curve_mask = 曲线像素
    marked = render_hairline_overlay(photo_bgr, points502, ext_faces, uv, black_texture_rgba)
    overlay = build_overlay_layer(h, w, points502, ext_faces, uv, black_texture_rgba)
    curve_mask = (overlay[:, :, 3] > 40).astype(np.uint8) * 255

    # ①②③ ROI
    upper = forehead_upper_region(landmarks_mp, w, h)
    head = head_silhouette(parse_map)
    roi = cv2.bitwise_and(upper, head)

    # ⑤ ROI ∩ 曲线以上 → 清理
    above = _above_curve_region(curve_mask, h, w)
    mask = cv2.bitwise_and(roi, above)
    mask = _clean_mask(mask, w)
    return marked, mask


def compose_comfy_rgba(marked_bgr: np.ndarray, mask: np.ndarray) -> Image.Image:
    """合成 ComfyUI LoadImage 用的 RGBA：RGB=划线图，alpha=255−mask（透明=重绘区）。"""
    rgb = cv2.cvtColor(marked_bgr, cv2.COLOR_BGR2RGB)
    alpha = (255 - mask).astype(np.uint8)
    rgba = np.dstack([rgb, alpha])
    return Image.fromarray(rgba, mode="RGBA")


if __name__ == "__main__":
    import sys, os
    from .service import get_landmarker, get_parser

    path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
    tex_name = sys.argv[2] if len(sys.argv) > 2 else "girl_straight"
    img = cv2.imread(path)
    h, w = img.shape[:2]
    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    from .hairline_2d import sample_hairline, smooth_hairline
    from .lift_3d import lift_hairline_to_3d, build_middle_row, assemble_full
    lm = get_landmarker().detect(rgb)
    parse_map = get_parser().parse(rgb)
    h2d, valid = sample_hairline(lm, parse_map); h2d = smooth_hairline(h2d, valid)
    h3d = lift_hairline_to_3d(lm, h2d); mid = build_middle_row(lm, h3d)
    pts = assemble_full(lm, mid, h3d)

    black = load_texture_rgba(f"hairline_texture_black/{tex_name}.png")
    marked, mask = build_inpaint_mask(img, lm, parse_map, pts, black)
    os.makedirs("tests/output", exist_ok=True)
    cv2.imwrite("tests/output/mask_marked.png", marked)
    cv2.imwrite("tests/output/mask_binary.png", mask)
    # 三联可视化：划线图 / ROI / mask 叠加
    upper = forehead_upper_region(lm, w, h); head = head_silhouette(parse_map)
    roi = cv2.bitwise_and(upper, head)
    vis = marked.copy()
    vis[roi > 0] = (vis[roi > 0] * 0.6 + np.array([0, 40, 0])).clip(0, 255).astype(np.uint8)
    vis[mask > 128] = (vis[mask > 128] * 0.4 + np.array([0, 0, 150])).clip(0, 255).astype(np.uint8)
    cv2.imwrite("tests/output/mask_vis.png", vis)
    compose_comfy_rgba(marked, mask).save("tests/output/comfy_input.png")
    print(f"saved mask_marked/mask_binary/mask_vis/comfy_input；mask 像素 {int((mask>128).sum())}")