feat(接口2): 新增生发后图片(ComfyUI/Flux inpaint)
在发际线预览基础上，每种发际线再出一张「植发3个月」生发图： - hairline/mask.py: headmark 5步法遮罩(额头上部区域∩SegFormer头部=ROI， 取发际线曲线以上闭合区域)；用 hairline_texture_black 渲染黑线替代手绘检测； compose_comfy_rgba 合成 RGBA(alpha=255-mask, 透明=重绘区, 对齐 ComfyUI mask=1-alpha) - hairline/comfyui.py: ComfyUI 客户端(默认8182)，/upload/image+/prompt(改节点26+随机seed) +轮询/history+/view 取回生发图 - hairline/render.py: 抽出 build_overlay_layer 供遮罩取曲线像素 - hairline/service.py: extract_context 一次出 landmarks/parse_map/502点； generate_grow_results 每种=预览+生发图(同步串行N张，单张ComfyUI失败则grown置空不拖垮整请求) - app.py: /hair/grow 返回 results[].grown_image_base64；重活放线程池避免卡事件循环 - add_hair.json 工作流 + hairline_texture_black/ 黑贴图入库 - 测试: test_mask.py(遮罩几何) + test_api mock ComfyUI 验 grown 字段，35 全绿 实测(5090): female 5张生发图同步约18s；预览/生发图人物五官服饰背景保持、黑线已清除。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
+64
-11
@@ -11,13 +11,22 @@ 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
|
||||
|
||||
_TEXTURE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "hairline_texture")
|
||||
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。
|
||||
@@ -73,40 +82,84 @@ def get_texture_map() -> dict:
|
||||
|
||||
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, 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 points, valid
|
||||
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": key, "image_bgr": ndarray, "order": 1..N}。
|
||||
无人脸返回 None。gender 必须是 male/female。
|
||||
Returns: list[dict] {"hairline_type", "image_bgr", "order"}；无人脸返回 None。
|
||||
"""
|
||||
if gender not in ("male", "female"):
|
||||
raise ValueError(f"gender 必须是 male/female，收到 {gender!r}")
|
||||
points, _valid = extract_502(image_bgr)
|
||||
if points is None:
|
||||
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):
|
||||
tex = load_texture_rgba(path)
|
||||
preview = render_hairline_overlay(image_bgr, points, ext_faces, uv, tex)
|
||||
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"
|
||||
|
||||
Reference in New Issue
Block a user