Files
hair/hairline/service.py
T
xsl cd8985e93d feat(接口5): 新增 generate_grow_image 参数控制是否生成生发效果图
- app.py: 接口5 路由加表单参数 generate_grow_image(bool, 默认 True)并透传
- hairline/service.py: generate_hairline_pngs 加同名参数, False 时跳过 ComfyUI 生发、grown_png 恒 None
- 默认行为不变(向后兼容); false 时仅返回三档发际线叠图与中心点, 大幅降低耗时
- 网关字节级透传 multipart, 新参数自动到达 worker, 无需改网关
2026-07-24 00:52:04 +08:00

441 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""接口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, build_overlay_layer
from .mask import build_inpaint_mask, compose_comfy_rgba, mask_from_curve
from .marker_detect import detect_marker_hairline, path_to_curve_mask
import base64
import io
import logging
logger = logging.getLogger("hair.worker")
# 接口2 女性发型 key → change_hair hair_idchang_*)映射:换发型+Flux-2 整帧重绘用。
# 与接口12 final 的 5 型一一对应。
_FEMALE_KEY_TO_CHANG = {
"ellipse": "chang_tuoyuan", # 椭圆
"flower": "chang_huaban", # 花瓣
"heart": "chang_xinxing", # 心形
"straight": "chang_zhixian", # 直线
"wave": "chang_bolang", # 波浪
}
_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")
def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0):
"""直接调 ComfyUI 重绘(替代原 local_test HTTP 服务)。
传 final 图 + 纯红遮罩 PNG,返回重绘后的 PNG bytes。
失败抛异常(调用方负责 try/except 跳过)。
"""
from .redraw import run_redraw
# front=True:接口2 时延敏感,插到 ComfyUI 队列最前,避免排在接口3/5 的批量任务后面
return run_redraw(image_png_bytes, mask_png_bytes, timeout=timeout, front=True)
# 发际线贴图档位:middle=默认(hairline_texture/)high/low 各自独立文件夹。
_TEXTURE_DIRS = {
"middle": _TEXTURE_DIR,
"high": os.path.join(_REPO, "hairline_texture_high"),
"low": os.path.join(_REPO, "hairline_texture_low"),
}
# torch 2.7.1+cu128 已支持 RTX 5090 (sm_120)SegFormer 走 GPU~0.05s/张)
_SEG_DEVICE = os.getenv("SEG_DEVICE", "cuda")
_landmarker = None
_parser = None
_texture_maps: dict = {} # {level: {gender: [(key, path)]}},按档位缓存
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(level: str = "middle") -> dict:
"""扫描指定档位贴图目录建 {gender: [(key, path)]},按 key 排序、按档位缓存。
levelmiddle(默认) / high / low,分别对应 hairline_texture[/_high|/_low]。
文件名规范化去空格(如 `man_ inverse_arc.png` → key `inverse_arc`)。
"""
if level not in _TEXTURE_DIRS:
raise ValueError(f"hairline_level 必须是 middle/high/low,收到 {level!r}")
cached = _texture_maps.get(level)
if cached is not None:
return cached
mapping: dict[str, list] = {"female": [], "male": []}
for path in sorted(glob.glob(os.path.join(_TEXTURE_DIRS[level], "*.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_maps[level] = mapping
return mapping
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, use_mask: bool = True,
prompt: str = None, hair_styles: list[int] | None = None,
workflow_path: str | None = None):
"""指定发际线类型:发际线透明叠图(白线 RGBA) + 生发图(ComfyUI)。
hair_styles1-indexed 列表):指定生成哪几张发际线(按贴图排序)。female: 1..5male: 1..4。
为 None 时生成全部(兼容旧调用)。
use_mask(默认 True):是否启用 inpaint 遮罩,用于测试对比(同接口3)。
False 时用**干净原图 + 空遮罩**送 ComfyUI(不烧黑色模板线)。
prompt(默认 None):ComfyUI 提示词,非 None 时替换工作流节点60文本。
workflow_path(默认 None):ComfyUI 工作流 JSON 路径,None 用默认 add_hair.json。
Returns: list[dict] {"hairline_type","order","overlay"((H,W,4) RGBA 透明层),
"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()
textures = get_texture_map()[gender] # [(key, path), ...] 已排序
if hair_styles is not None:
items = [(s, textures[s - 1]) for s in hair_styles]
else:
items = list(enumerate(textures, start=1))
# 禁用遮罩:干净原图 + 空遮罩,与模板无关 → 只跑一次 ComfyUI,下面 N 项复用
shared_grown = None
if not use_mask:
try:
h, w = image_bgr.shape[:2]
buf = io.BytesIO()
compose_comfy_rgba(image_bgr, np.zeros((h, w), np.uint8)).save(buf, format="PNG")
# front=True:接口2 时延敏感,插到 ComfyUI 队列最前
shared_grown = comfyui.run(buf.getvalue(), prompt=prompt,
workflow_path=workflow_path, front=True)
except Exception as e: # noqa: BLE001
logger.warning("接口2 生发图失败(无遮罩):%s", e)
results = []
h, w = image_bgr.shape[:2]
for order, (key, white_path) in items:
white = load_texture_rgba(white_path)
overlay = build_overlay_layer(h, w, ctx["points"], ext_faces, uv, white)
if not use_mask:
grown_png = shared_grown
else:
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")
# front=True:接口2 时延敏感,插到 ComfyUI 队列最前
grown_png = comfyui.run(buf.getvalue(), prompt=prompt,
workflow_path=workflow_path, front=True)
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
logger.warning("接口2 生发图失败 type=%s%s", key, e)
results.append({"hairline_type": key, "order": order,
"overlay": overlay, "grown_png": grown_png})
return results
def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | None,
redraw_defaults: dict):
"""接口2 女性专用:发际线透明叠图(同 generate_grow_results+ 换发型重绘图。
grown 图来源(新流程):对每个选中发型把 female key 映射到 change_hair 的 chang_* hair_id
调 face_analysis.hairline_grow.generate_hairline_redraw= 接口12 final 管线,参数用
redraw_defaults)拿到 ④ final(接缝融合基底)+ ⑤-② 纯红遮罩 PNG,再**后端直接调
ComfyUI**0716add-hair-api.json 工作流)完成发际线带重绘,重绘结果作为生发图。
overlay 仍是发际线曲线透明层(与 generate_grow_results 完全一致)。
Returns: list[dict] {"hairline_type","order","overlay","grown_png"(jpg bytes 或 None)}
无人脸返回 None。单个发型换发型/重绘失败时 grown_png=None,不抛异常。
"""
from face_analysis.hairline_grow import generate_hairline_redraw, NoFaceError
from face_analysis.head_mask import SEGFORMER_HAIR
ctx = extract_context(image_bgr)
if ctx is None:
return None
uv, ext_faces = load_ext_mesh()
# 复用 extract_context 已算好的 SegFormer parse_map,避免 generate_hairline_redraw 内部重复分割
hair_mask_reuse = (ctx["parse_map"] == SEGFORMER_HAIR)
textures = get_texture_map()["female"] # [(key, path), ...] 已排序
if hair_styles is not None:
items = [(s, textures[s - 1]) for s in hair_styles]
else:
items = list(enumerate(textures, start=1))
results = []
h, w = image_bgr.shape[:2]
for order, (key, white_path) in items:
white = load_texture_rgba(white_path)
overlay = build_overlay_layer(h, w, ctx["points"], ext_faces, uv, white)
grown_png = None
chang_id = _FEMALE_KEY_TO_CHANG.get(key)
if chang_id is None:
logger.warning("接口2 换发型:female key=%s 无对应 chang_id,跳过生发图", key)
else:
try:
data = generate_hairline_redraw(image_bgr, chang_id, hair_mask=hair_mask_reuse, **redraw_defaults)
steps = data.get("steps") or {}
# ④ final(接缝融合基底)+ ⑤-② 纯红遮罩 PNG
final_b64 = steps.get("final_base64") or ""
mask_b64 = steps.get("redraw_band_mask_base64") or ""
if not final_b64 or not mask_b64:
logger.warning("接口2 换发型:type=%s final/遮罩缺失(final=%d mask=%d",
key, len(final_b64), len(mask_b64))
else:
# 去掉 data URI 前缀
if final_b64.startswith("data:"):
final_b64 = final_b64.split(",", 1)[1]
if mask_b64.startswith("data:"):
mask_b64 = mask_b64.split(",", 1)[1]
final_bytes = base64.b64decode(final_b64)
mask_bytes = base64.b64decode(mask_b64)
# 后端直接调 ComfyUI 重绘,返回重绘后的 PNG
grown_png = _call_local_redraw(final_bytes, mask_bytes)
if grown_png is None:
logger.warning("接口2 换发型:type=%s 重绘结果为空", key)
except NoFaceError:
logger.warning("接口2 换发型:type=%s 未检出人脸", key)
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
logger.warning("接口2 换发型图失败 type=%s%s", key, e)
results.append({"hairline_type": key, "order": order,
"overlay": overlay, "grown_png": grown_png})
return results
def _grow_from_texture(image_bgr: np.ndarray, ctx: dict, white_path: str | None,
use_mask: bool, prompt: str | None):
"""对单个发际线做生发(ComfyUI)。黑模板固定取 hairline_texture_black/middle),
与 hairline_level 无关(high/low 贴图与 middle 同名,basename 映射即落回 middle 黑模板)。
use_mask=False 时用干净原图+空遮罩(与贴图无关,white_path 可为 None)。
失败返回 None,不抛异常。
"""
try:
if use_mask:
black = load_texture_rgba(_black_texture_path(white_path))
marked, mask = build_inpaint_mask(
image_bgr, ctx["landmarks"], ctx["parse_map"], ctx["points"], black)
else:
h, w = image_bgr.shape[:2]
marked, mask = image_bgr, np.zeros((h, w), np.uint8)
buf = io.BytesIO()
compose_comfy_rgba(marked, mask).save(buf, format="PNG")
return comfyui.run(buf.getvalue(), prompt=prompt)
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
logger.warning("接口5 生发图失败:%s", e)
return None
def generate_hairline_pngs(image_bgr: np.ndarray, gender: str,
hair_styles: list[int], use_mask: bool = True,
prompt: str | None = None,
generate_grow_image: bool = True):
"""接口5:对选中发型返回 middle/high/low 三档发际线透明叠图 + 生发图(同接口2)。
入参同接口2:先选 gender,再多选 hair_styles(必填,1-indexed 按贴图排序)。
每个选中发型返回三档叠图(middle/high/low,RGBA 透明层只含发际线曲线)与一张生发图;
三档贴图同名,生发黑模板固定取自 hairline_texture_black/middle),故生发目标固定 middle 档。
use_mask/prompt:同接口2 的生发参数。
generate_grow_image(默认 True):是否生成生发图(ComfyUI,最耗时)。False 时跳过生发,
各发型 grown_png 恒为 None,可大幅降低耗时(仅留三档发际线叠图与中心点)。
Returns: {"images":[{hairline_type,order,overlays:{middle,high,low}((H,W,4) RGBA 透明层),grown_png}],
"best_centers":{"middle":(x,y),"high":(x,y),"low":(x,y)}};无人脸 None。
best_centers 取首个选中发型三档各自的发际线中点。
"""
if gender not in ("male", "female"):
raise ValueError(f"gender 必须是 male/female,收到 {gender!r}")
if not hair_styles:
raise ValueError("hair_styles 必填且不能为空")
ctx = extract_context(image_bgr)
if ctx is None:
return None
h, w = image_bgr.shape[:2]
uv, ext_faces = load_ext_mesh()
lm = ctx["landmarks"]
# 面部中轴 x = 眉心(9/151 中点)
face_cx = float((lm[9, 0] + lm[151, 0]) / 2 * w)
# 三档贴图表(同性别、同 key 顺序,因三个文件夹同名)
tex_by_level = {lv: get_texture_map(lv)[gender] for lv in _TEXTURE_DIRS}
# use_mask=False:干净原图+空遮罩与贴图无关,只跑一次 ComfyUI,选中项复用
# generate_grow_image=False:完全跳过生发(最耗时),grown_png 恒为 None
shared_grown = None
if generate_grow_image and not use_mask:
shared_grown = _grow_from_texture(image_bgr, ctx, None, use_mask=False, prompt=prompt)
def _center_of(overlay):
"""从某档发际线透明叠图取面部中轴处的发际线中点 (x,y),无像素返回 None。"""
ys, xs = np.where(overlay[:, :, 3] > 40)
if not xs.size:
return None
near = np.abs(xs - face_cx) <= max(2, int(w * 0.02))
col_ys = ys[near] if near.any() else ys[np.argsort(np.abs(xs - face_cx))[:20]]
return (int(round(face_cx)), int(round(float(col_ys.mean()))))
images, best_centers = [], None
for s in hair_styles: # s = 1-indexed 发型序号
key, mid_path = tex_by_level["middle"][s - 1]
overlays = {}
for lv in _TEXTURE_DIRS:
white = load_texture_rgba(tex_by_level[lv][s - 1][1])
overlays[lv] = build_overlay_layer(h, w, ctx["points"], ext_faces, uv, white)
# 生发:固定 middle 黑模板(generate_grow_image=False 时跳过,恒 None
if not generate_grow_image:
grown_png = None
elif not use_mask:
grown_png = shared_grown
else:
grown_png = _grow_from_texture(image_bgr, ctx, mid_path, use_mask=True, prompt=prompt)
images.append({"hairline_type": key, "order": s,
"overlays": overlays, "grown_png": grown_png})
# best_centers:首个选中发型三档(middle/high/low)发际线中点
if best_centers is None:
best_centers = {lv: _center_of(overlays[lv]) for lv in _TEXTURE_DIRS}
return {"images": images, "best_centers": best_centers}
def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True, prompt: str = None):
"""接口3:检测医生手绘发际线 → 遮罩 → 送 ComfyUI 生发(仅需划线图一张)。
检测路径只用来**建遮罩**ComfyUI 输入图用 **marked 原图**(含医生手绘线,
工作流提示词会清除黑线再生发)。
use_mask(默认 True):是否启用自动检测的遮罩,用于测试对比。
- True:检测手绘线 → 建遮罩 → alpha=255−mask(透明区=重绘区,节点44 画黄色参考区)。
- False:跳过检测,直接送划线图,alpha 全 255(空遮罩,节点26 mask 为空),
模型仅凭医生黑线参考生发。无需改工作流,唯一变量是遮罩。
Returns: {"grown_png": bytes 或 None, "status": "ok"|"no_face"|"no_line"}。
"""
h, w = marked_bgr.shape[:2]
if use_mask:
rgb = cv2.cvtColor(marked_bgr, cv2.COLOR_BGR2RGB)
landmarks = get_landmarker().detect(rgb)
if landmarks is None:
return {"grown_png": None, "status": "no_face"}
parse_map = get_parser().parse(rgb)
path = detect_marker_hairline(marked_bgr, landmarks, parse_map)
if path is None:
return {"grown_png": None, "status": "no_line"}
line_w = max(2, int(w * 0.006))
curve_mask = path_to_curve_mask(path, h, w, thickness=max(3, line_w))
mask = mask_from_curve(curve_mask, landmarks, parse_map)
else:
mask = np.zeros((h, w), np.uint8) # 空遮罩:alpha 全 255,跳过检测
buf = io.BytesIO()
compose_comfy_rgba(marked_bgr, mask).save(buf, format="PNG") # marked 原图 + 遮罩
grown_png = comfyui.run(buf.getvalue(), prompt=prompt)
return {"grown_png": grown_png, "status": "ok"}
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}")