hair/hairline/service.py

"""接口2 服务层：模型单例 + 性别贴图映射 + 「照片→N 张发际线预览图」管线。

把 head3d 的 extract_hairline 步骤包成单例复用（避免每请求重建模型），再按性别
对每张贴图调 render.render_hairline_overlay 生成预览图。
"""
from __future__ import annotations
import glob
import os

import cv2
import numpy as np

from . import constants as C
from . import comfyui
from .face_landmarks import FaceLandmarker
from .face_parsing import FaceParser
from .hairline_2d import sample_hairline, smooth_hairline
from .lift_3d import lift_hairline_to_3d, build_middle_row, assemble_full
from .render import load_ext_mesh, load_texture_rgba, render_hairline_overlay
from .mask import build_inpaint_mask, compose_comfy_rgba

import io
import logging

logger = logging.getLogger("hair.worker")

_REPO = os.path.dirname(os.path.dirname(__file__))
_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture")
_BLACK_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture_black")

# ⚠️ 本 worker 是 RTX 5090(sm_120)，torch 2.2.2(cu121) 只编到 sm_90，CUDA 跑算子会报
#    "no kernel image"。SegFormer 默认走 CPU（~2.5s/张）。换 torch cu128 后可设 SEG_DEVICE=cuda。
_SEG_DEVICE = os.getenv("SEG_DEVICE", "cpu")

_landmarker = None
_parser = None
_texture_map = None


def get_landmarker() -> FaceLandmarker:
    global _landmarker
    if _landmarker is None:
        _landmarker = FaceLandmarker(static_image_mode=True)
    return _landmarker


def get_parser() -> FaceParser:
    global _parser
    if _parser is None:
        _parser = FaceParser(device=_SEG_DEVICE)
    return _parser


def _gender_key(stem: str):
    """文件名 stem → (gender, key)；非 girl_/man_ 前缀返回 (None, None)。"""
    if stem.startswith("girl_"):
        return "female", stem[5:].replace(" ", "").strip()
    if stem.startswith("man_"):
        return "male", stem[4:].replace(" ", "").strip()
    return None, None


def get_texture_map() -> dict:
    """扫描 hairline_texture/ 建 {gender: [(key, path)]}，按 key 排序、缓存。

    文件名规范化去空格（如 `man_ inverse_arc.png` → key `inverse_arc`）。
    """
    global _texture_map
    if _texture_map is not None:
        return _texture_map
    mapping: dict[str, list] = {"female": [], "male": []}
    for path in sorted(glob.glob(os.path.join(_TEXTURE_DIR, "*.png"))):
        stem = os.path.splitext(os.path.basename(path))[0]
        gender, key = _gender_key(stem)
        if gender:
            mapping[gender].append((key, path))
    for g in mapping:
        mapping[g].sort(key=lambda kp: kp[0])
    _texture_map = mapping
    return _texture_map


def extract_502(image_bgr: np.ndarray):
    """照片(BGR) → (points502 MP序, valid17)。无人脸返回 (None, None)。"""
    ctx = extract_context(image_bgr)
    if ctx is None:
        return None, None
    return ctx["points"], ctx["valid"]


def extract_context(image_bgr: np.ndarray):
    """照片(BGR) → {landmarks, parse_map, points, valid}。无人脸返回 None。"""
    rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    landmarks = get_landmarker().detect(rgb)
    if landmarks is None:
        return None
    parse_map = get_parser().parse(rgb)
    hairline_2d, valid = sample_hairline(landmarks, parse_map)
    hairline_2d = smooth_hairline(hairline_2d, valid)
    hairline_3d = lift_hairline_to_3d(landmarks, hairline_2d)
    middle_3d = build_middle_row(landmarks, hairline_3d)
    points = assemble_full(landmarks, middle_3d, hairline_3d)
    return {"landmarks": landmarks, "parse_map": parse_map, "points": points, "valid": valid}


def _black_texture_path(white_path: str) -> str:
    """白贴图路径 → 同名黑贴图路径（hairline_texture_black/）。"""
    return os.path.join(_BLACK_TEXTURE_DIR, os.path.basename(white_path))


def generate_previews(image_bgr: np.ndarray, gender: str):
    """生成该性别全部发际线预览图（仅预览，不生发）。

    Returns: list[dict] {"hairline_type", "image_bgr", "order"}；无人脸返回 None。
    """
    if gender not in ("male", "female"):
        raise ValueError(f"gender 必须是 male/female，收到 {gender!r}")
    ctx = extract_context(image_bgr)
    if ctx is None:
        return None
    uv, ext_faces = load_ext_mesh()
    results = []
    for order, (key, path) in enumerate(get_texture_map()[gender], start=1):
        preview = render_hairline_overlay(image_bgr, ctx["points"], ext_faces, uv,
                                          load_texture_rgba(path))
        results.append({"hairline_type": key, "image_bgr": preview, "order": order})
    return results


def generate_grow_results(image_bgr: np.ndarray, gender: str):
    """该性别全部发际线：预览图(白线) + 生发图(ComfyUI)。同步、串行。

    Returns: list[dict] {"hairline_type","order","image_bgr"(预览), "grown_png"(bytes 或 None)}。
    无人脸返回 None。某张 ComfyUI 失败时该项 grown_png=None，不影响其余。
    """
    if gender not in ("male", "female"):
        raise ValueError(f"gender 必须是 male/female，收到 {gender!r}")
    ctx = extract_context(image_bgr)
    if ctx is None:
        return None
    uv, ext_faces = load_ext_mesh()
    results = []
    for order, (key, white_path) in enumerate(get_texture_map()[gender], start=1):
        white = load_texture_rgba(white_path)
        preview = render_hairline_overlay(image_bgr, ctx["points"], ext_faces, uv, white)

        grown_png = None
        try:
            black = load_texture_rgba(_black_texture_path(white_path))
            marked, mask = build_inpaint_mask(
                image_bgr, ctx["landmarks"], ctx["parse_map"], ctx["points"], black)
            buf = io.BytesIO()
            compose_comfy_rgba(marked, mask).save(buf, format="PNG")
            grown_png = comfyui.run(buf.getvalue())
        except Exception as e:  # noqa: BLE001  单张失败不拖垮整请求
            logger.warning("接口2 生发图失败 type=%s：%s", key, e)

        results.append({"hairline_type": key, "order": order,
                        "image_bgr": preview, "grown_png": grown_png})
    return results


if __name__ == "__main__":
    import sys
    g = sys.argv[2] if len(sys.argv) > 2 else "female"
    img = cv2.imread(sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg")
    os.makedirs("tests/output", exist_ok=True)
    print("texture map:", {k: [kp[0] for kp in v] for k, v in get_texture_map().items()})
    res = generate_previews(img, g)
    if res is None:
        print("无人脸")
        sys.exit(1)
    for r in res:
        out = f"tests/output/preview_{g}_{r['hairline_type']}.png"
        cv2.imwrite(out, r["image_bgr"])
        print(f"  order={r['order']} type={r['hairline_type']} -> {out}")