diff --git a/add_hair.json b/add_hair.json new file mode 100644 index 0000000..78e4b68 --- /dev/null +++ b/add_hair.json @@ -0,0 +1,450 @@ +{ + "1": { + "inputs": { + "scheduler": "simple", + "steps": 6, + "denoise": 1, + "model": [ + "2", + 0 + ] + }, + "class_type": "BasicScheduler", + "_meta": { + "title": "基本调度器" + } + }, + "2": { + "inputs": { + "max_shift": 1.15, + "base_shift": 0.5, + "width": [ + "14", + 0 + ], + "height": [ + "14", + 1 + ], + "model": [ + "16", + 0 + ] + }, + "class_type": "ModelSamplingFlux", + "_meta": { + "title": "采样算法（Flux）" + } + }, + "3": { + "inputs": { + "vae_name": "flux2-vae.safetensors" + }, + "class_type": "VAELoader", + "_meta": { + "title": "加载VAE" + } + }, + "5": { + "inputs": { + "conditioning": [ + "19", + 0 + ], + "latent": [ + "13", + 0 + ] + }, + "class_type": "ReferenceLatent", + "_meta": { + "title": "参考Latent" + } + }, + "6": { + "inputs": { + "noise_seed": 690815303279000 + }, + "class_type": "RandomNoise", + "_meta": { + "title": "随机噪波" + } + }, + "7": { + "inputs": { + "width": [ + "14", + 0 + ], + "height": [ + "14", + 1 + ], + "batch_size": 1 + }, + "class_type": "EmptySD3LatentImage", + "_meta": { + "title": "空Latent图像（SD3）" + } + }, + "8": { + "inputs": { + "sampler_name": "euler" + }, + "class_type": "KSamplerSelect", + "_meta": { + "title": "K采样器选择" + } + }, + "9": { + "inputs": { + "noise": [ + "6", + 0 + ], + "guider": [ + "20", + 0 + ], + "sampler": [ + "8", + 0 + ], + "sigmas": [ + "1", + 0 + ], + "latent_image": [ + "7", + 0 + ] + }, + "class_type": "SamplerCustomAdvanced", + "_meta": { + "title": "自定义采样器（高级）" + } + }, + "10": { + "inputs": { + "samples": [ + "9", + 0 + ], + "vae": [ + "3", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE解码" + } + }, + "13": { + "inputs": { + "pixels": [ + "44", + 0 + ], + "vae": [ + "3", + 0 + ] + }, + "class_type": "VAEEncode", + "_meta": { + "title": "VAE编码" + } + }, + "14": { + "inputs": { + "image": [ + "44", + 0 + ] + }, + "class_type": "GetImageSize+", + "_meta": { + "title": "🔧 Get Image Size" + } + }, + "16": { + "inputs": { + "unet_name": "flux2.0/flux-2-klein-9b-fp8.safetensors", + "weight_dtype": "fp8_e4m3fn" + }, + "class_type": "UNETLoader", + "_meta": { + "title": "UNet加载器" + } + }, + "17": { + "inputs": { + "filename_prefix": "ComfyUI", + "images": [ + "62", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "保存图像" + } + }, + "19": { + "inputs": { + "guidance": 1, + "conditioning": [ + "22", + 0 + ] + }, + "class_type": "FluxGuidance", + "_meta": { + "title": "Flux引导" + } + }, + "20": { + "inputs": { + "model": [ + "2", + 0 + ], + "conditioning": [ + "5", + 0 + ] + }, + "class_type": "BasicGuider", + "_meta": { + "title": "基本引导器" + } + }, + "22": { + "inputs": { + "text": [ + "60", + 0 + ], + "clip": [ + "61", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP文本编码" + } + }, + "26": { + "inputs": { + "image": "clipspace/clipspace-painted-masked-1780841075435.png [input]" + }, + "class_type": "LoadImage", + "_meta": { + "title": "加载图像" + } + }, + "31": { + "inputs": { + "image": [ + "26", + 0 + ] + }, + "class_type": "easy imageSize", + "_meta": { + "title": "图像尺寸" + } + }, + "32": { + "inputs": { + "aspect_ratio": "custom", + "proportional_width": [ + "31", + 0 + ], + "proportional_height": [ + "31", + 1 + ], + "fit": "letterbox", + "method": "lanczos", + "round_to_multiple": "8", + "scale_to_side": "None", + "scale_to_length": 1024, + "background_color": "#000000", + "image": [ + "26", + 0 + ], + "mask": [ + "37", + 0 + ] + }, + "class_type": "LayerUtility: ImageScaleByAspectRatio V2", + "_meta": { + "title": "LayerUtility: ImageScaleByAspectRatio V2" + } + }, + "33": { + "inputs": { + "masks": [ + "26", + 1 + ] + }, + "class_type": "Mask Fill Holes", + "_meta": { + "title": "遮罩填充漏洞" + } + }, + "36": { + "inputs": { + "masks": [ + "33", + 0 + ] + }, + "class_type": "Convert Masks to Images", + "_meta": { + "title": "遮罩到图像" + } + }, + "37": { + "inputs": { + "method": "intensity", + "image": [ + "39", + 0 + ] + }, + "class_type": "Image To Mask", + "_meta": { + "title": "图像到遮罩" + } + }, + "39": { + "inputs": { + "upscale_method": "nearest-exact", + "width": [ + "31", + 0 + ], + "height": [ + "31", + 1 + ], + "crop": "disabled", + "image": [ + "36", + 0 + ] + }, + "class_type": "ImageScale", + "_meta": { + "title": "缩放图像" + } + }, + "44": { + "inputs": { + "mask_opacity": 1, + "mask_color": "FFFF00", + "pass_through": true, + "image": [ + "32", + 0 + ], + "mask": [ + "32", + 1 + ] + }, + "class_type": "ImageAndMaskPreview", + "_meta": { + "title": "图像与遮罩预览" + } + }, + "45": { + "inputs": { + "images": [ + "44", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "预览图像" + } + }, + "53": { + "inputs": { + "rgthree_comparer": { + "images": [ + { + "name": "A", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_nljpo_00001_.png&type=temp&subfolder=&rand=0.7003773747423834" + }, + { + "name": "B", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_nljpo_00002_.png&type=temp&subfolder=&rand=0.10574778041280719" + } + ] + }, + "image_a": [ + "62", + 0 + ], + "image_b": [ + "26", + 0 + ] + }, + "class_type": "Image Comparer (rgthree)", + "_meta": { + "title": "图像对比" + } + }, + "60": { + "inputs": { + "text": "严格保留原图人物五官、面部肤色、神态、服饰、背景、原图光影与整体色调，画面除划线区域外全部细节保持原样不变；先清除画面内所有黑色标注划线，仅在原划线划定范围内生成植发术后 3 个月长发效果，新生头发生长边界刚好止于原划线位置，新发沿原划线轮廓边缘和外部原生发丝交错融合、自然无缝衔接，无生硬分界线；发色和自身原生黑发保持统一，发丝细腻写实、带有真实毛发纹理与自然光泽，发量均匀，符合植发 3 个月刚长出的轻微稀疏生长状态，毛发走向遵从头部原生头发生长规律，头发光影明暗和原图环境统一，整体写实自然，无贴片假发质感。" + }, + "class_type": "JjkText", + "_meta": { + "title": "Text" + } + }, + "61": { + "inputs": { + "clip_name": "qwen_3_8b_fp8mixed.safetensors", + "type": "flux2", + "device": "default" + }, + "class_type": "CLIPLoader", + "_meta": { + "title": "加载CLIP" + } + }, + "62": { + "inputs": { + "method": "mkl", + "strength": 1, + "multithread": true, + "image_ref": [ + "26", + 0 + ], + "image_target": [ + "10", + 0 + ] + }, + "class_type": "ColorMatch", + "_meta": { + "title": "Color Match" + } + } +} \ No newline at end of file diff --git a/app.py b/app.py index 1aab43b..6539a78 100644 --- a/app.py +++ b/app.py @@ -504,17 +504,22 @@ async def hair_grow( return err(1002, "人像分辨率过低") try: - from hairline.service import generate_previews + from fastapi.concurrency import run_in_threadpool + from hairline.service import generate_grow_results - previews = generate_previews(image, gender) # 无人脸 → None - if previews is None: + # 预览 + 生发(ComfyUI) 都是阻塞且较慢，放线程池避免卡住事件循环 + items = await run_in_threadpool(generate_grow_results, image, gender) + if items is None: return err(1001, "无法识别人像") results = [] - for p in previews: - ok_enc, png = cv2.imencode(".png", p["image_bgr"]) + for p in items: + _ok, png = cv2.imencode(".png", p["image_bgr"]) + grown_b64 = (base64.b64encode(p["grown_png"]).decode() + if p["grown_png"] else None) results.append({ "image_base64": base64.b64encode(png.tobytes()).decode(), + "grown_image_base64": grown_b64, "hairline_type": p["hairline_type"], "order": p["order"], }) diff --git a/hairline/comfyui.py b/hairline/comfyui.py new file mode 100644 index 0000000..e9422aa --- /dev/null +++ b/hairline/comfyui.py @@ -0,0 +1,110 @@ +"""ComfyUI 客户端：用 add_hair.json 工作流跑生发图（Flux-2 inpaint）。 + +worker 不跑 Flux，只把「划线图 + 遮罩」的 RGBA 上传到本机 ComfyUI(默认 8182)， +替换工作流节点 26 的输入图、随机 seed，提交 /prompt，轮询 /history，取回 /view 输出。 +""" +from __future__ import annotations + +import copy +import json +import os +import random +import time +import uuid + +import httpx + +COMFYUI_URL = os.getenv("COMFYUI_URL", "http://127.0.0.1:8182").rstrip("/") +WORKFLOW_PATH = os.getenv( + "ADD_HAIR_WORKFLOW", + os.path.join(os.path.dirname(os.path.dirname(__file__)), "add_hair.json"), +) +COMFY_TIMEOUT = float(os.getenv("COMFYUI_TIMEOUT", "600")) # 单张出图最长等待(秒) + +_INPUT_NODE = "26" # LoadImage：外部输入图(含 alpha 遮罩) +_SEED_NODE = "6" # RandomNoise +_OUTPUT_NODE = "17" # SaveImage + +_workflow = None + + +def _load_workflow() -> dict: + global _workflow + if _workflow is None: + with open(WORKFLOW_PATH, encoding="utf-8") as f: + _workflow = json.load(f) + return _workflow + + +def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT) -> bytes: + """提交一次生发任务，返回输出 PNG 字节。失败抛异常。""" + client_id = uuid.uuid4().hex + with httpx.Client(base_url=COMFYUI_URL, timeout=30.0) as cli: + # 1. 上传输入图（含 alpha 遮罩）到 ComfyUI input 目录 + fname = f"hair_{client_id}.png" + r = cli.post("/upload/image", files={"image": (fname, rgba_png_bytes, "image/png")}, + data={"overwrite": "true", "type": "input"}) + r.raise_for_status() + up = r.json() + name = (up.get("subfolder") + "/" if up.get("subfolder") else "") + up["name"] + + # 2. 改工作流：节点26 输入图 + 随机 seed + wf = copy.deepcopy(_load_workflow()) + wf[_INPUT_NODE]["inputs"]["image"] = name + wf[_SEED_NODE]["inputs"]["noise_seed"] = random.randint(0, 2**63 - 1) + + # 3. 提交 + r = cli.post("/prompt", json={"prompt": wf, "client_id": client_id}) + r.raise_for_status() + prompt_id = r.json()["prompt_id"] + + # 4. 轮询 /history + deadline = time.time() + timeout + outputs = None + while time.time() < deadline: + hr = cli.get(f"/history/{prompt_id}") + hr.raise_for_status() + hist = hr.json() + if prompt_id in hist: + entry = hist[prompt_id] + status = entry.get("status", {}) + if status.get("status_str") == "error": + raise RuntimeError(f"ComfyUI 执行报错: {status}") + outputs = entry.get("outputs") + if outputs and _OUTPUT_NODE in outputs: + break + time.sleep(1.0) + if not outputs or _OUTPUT_NODE not in outputs: + raise TimeoutError(f"ComfyUI 出图超时({timeout}s) prompt_id={prompt_id}") + + # 5. 取回输出图 + imgs = outputs[_OUTPUT_NODE].get("images") or [] + if not imgs: + raise RuntimeError("ComfyUI 输出无图像") + info = imgs[0] + vr = cli.get("/view", params={"filename": info["filename"], + "subfolder": info.get("subfolder", ""), + "type": info.get("type", "output")}) + vr.raise_for_status() + return vr.content + + +def ping() -> bool: + """探测 ComfyUI 是否在线（/system_stats）。""" + try: + with httpx.Client(base_url=COMFYUI_URL, timeout=3.0) as cli: + return cli.get("/system_stats").status_code == 200 + except Exception: # noqa: BLE001 + return False + + +if __name__ == "__main__": + import sys + inp = sys.argv[1] if len(sys.argv) > 1 else "tests/output/comfy_input.png" + print("ComfyUI:", COMFYUI_URL, "online:", ping()) + with open(inp, "rb") as f: + png = run(f.read()) + out = "tests/output/grown.png" + with open(out, "wb") as f: + f.write(png) + print(f"生发图已存 {out}（{len(png)} bytes）") diff --git a/hairline/mask.py b/hairline/mask.py new file mode 100644 index 0000000..f726437 --- /dev/null +++ b/hairline/mask.py @@ -0,0 +1,149 @@ +"""接口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())}") diff --git a/hairline/render.py b/hairline/render.py index 5ac8d7a..2939029 100644 --- a/hairline/render.py +++ b/hairline/render.py @@ -78,19 +78,30 @@ def render_hairline_overlay(photo_bgr: np.ndarray, points502_norm: (502, 3) 归一化坐标（x,y ∈ [0,1]），**MP 顺序**（extract_hairline 输出）。 """ H, W = photo_bgr.shape[:2] - TH, TW = texture_rgba.shape[:2] - pts_obj = mp_order_to_obj_order(points502_norm) # MP序 → OBJ序 - img_xy = pts_obj[:, :2] * np.array([W, H], dtype=np.float32) # (502,2) + overlay = build_overlay_layer(H, W, points502_norm, ext_faces, uv502, texture_rgba) + # alpha 合成（RGBA→BGR：贴图 RGB 顺序需反成 BGR） + a = overlay[:, :, 3:4] / 255.0 + rgb = overlay[:, :, :3][..., ::-1] # RGB→BGR + out = photo_bgr.astype(np.float32) * (1.0 - a) + rgb * a + return np.clip(out, 0, 255).astype(np.uint8) - overlay = np.zeros((H, W, 4), np.float32) # 累积曲线层 RGBA + +def build_overlay_layer(H, W, points502_norm, ext_faces, uv502, texture_rgba) -> np.ndarray: + """渲染发际线曲线层，返回 (H, W, 4) float32 RGBA（未合成到照片）。 + + 供渲染合成（render_hairline_overlay）与遮罩（mask.py 取 alpha=曲线像素）共用。 + """ + TH, TW = texture_rgba.shape[:2] + pts_obj = mp_order_to_obj_order(points502_norm) + img_xy = pts_obj[:, :2] * np.array([W, H], dtype=np.float32) + overlay = np.zeros((H, W, 4), np.float32) tex = texture_rgba.astype(np.float32) for (i, j, k) in ext_faces: dst = img_xy[[i, j, k]].astype(np.float32) # UV → 贴图像素；flipY：贴图 y = (1 - v_raw) * TH（与 head3d Three.js flipY=true 一致） src = np.array([[uv502[v][0] * TW, (1.0 - uv502[v][1]) * TH] for v in (i, j, k)], dtype=np.float32) - # 退化三角形（投影到一条线）跳过，避免 getAffineTransform 奇异 - if cv2.contourArea(dst.astype(np.int32)) < 1.0: + if cv2.contourArea(dst.astype(np.int32)) < 1.0: # 退化三角形跳过 continue M = cv2.getAffineTransform(src, dst) warped = cv2.warpAffine(tex, M, (W, H), flags=cv2.INTER_LINEAR, @@ -99,9 +110,4 @@ def render_hairline_overlay(photo_bgr: np.ndarray, cv2.fillConvexPoly(tri_mask, dst.astype(np.int32), 255) sel = tri_mask > 0 overlay[sel] = warped[sel] - - # alpha 合成（RGBA→BGR：贴图 RGB 顺序需反成 BGR） - a = overlay[:, :, 3:4] / 255.0 - rgb = overlay[:, :, :3][..., ::-1] # RGB→BGR - out = photo_bgr.astype(np.float32) * (1.0 - a) + rgb * a - return np.clip(out, 0, 255).astype(np.uint8) + return overlay diff --git a/hairline/service.py b/hairline/service.py index b66c6a5..1130749 100644 --- a/hairline/service.py +++ b/hairline/service.py @@ -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" diff --git a/hairline_texture/man_ inverse_arc.png b/hairline_texture/man_ inverse_arc.png index 70ce13a..1212535 100644 Binary files a/hairline_texture/man_ inverse_arc.png and b/hairline_texture/man_ inverse_arc.png differ diff --git a/hairline_texture_black/girl_ellipse.png b/hairline_texture_black/girl_ellipse.png new file mode 100644 index 0000000..bad192e Binary files /dev/null and b/hairline_texture_black/girl_ellipse.png differ diff --git a/hairline_texture_black/girl_flower.png b/hairline_texture_black/girl_flower.png new file mode 100644 index 0000000..debae8f Binary files /dev/null and b/hairline_texture_black/girl_flower.png differ diff --git a/hairline_texture_black/girl_heart.png b/hairline_texture_black/girl_heart.png new file mode 100644 index 0000000..cd48126 Binary files /dev/null and b/hairline_texture_black/girl_heart.png differ diff --git a/hairline_texture_black/girl_straight.png b/hairline_texture_black/girl_straight.png new file mode 100644 index 0000000..023d366 Binary files /dev/null and b/hairline_texture_black/girl_straight.png differ diff --git a/hairline_texture_black/girl_wave.png b/hairline_texture_black/girl_wave.png new file mode 100644 index 0000000..4c2762f Binary files /dev/null and b/hairline_texture_black/girl_wave.png differ diff --git a/hairline_texture_black/man_ inverse_arc.png b/hairline_texture_black/man_ inverse_arc.png new file mode 100644 index 0000000..6742d5e Binary files /dev/null and b/hairline_texture_black/man_ inverse_arc.png differ diff --git a/hairline_texture_black/man_ellipse.png b/hairline_texture_black/man_ellipse.png new file mode 100644 index 0000000..367acc9 Binary files /dev/null and b/hairline_texture_black/man_ellipse.png differ diff --git a/hairline_texture_black/man_m.png b/hairline_texture_black/man_m.png new file mode 100644 index 0000000..d87b691 Binary files /dev/null and b/hairline_texture_black/man_m.png differ diff --git a/hairline_texture_black/man_straight.png b/hairline_texture_black/man_straight.png new file mode 100644 index 0000000..632174e Binary files /dev/null and b/hairline_texture_black/man_straight.png differ diff --git a/tests/test_api.py b/tests/test_api.py index f5638cb..75b340e 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -81,7 +81,16 @@ def test_grow_missing_gender_1004(client): assert r.json()["code"] == 1004 -def test_grow_female_returns_5(client): +_PNG_1x1 = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) + + +def test_grow_female_returns_5(client, monkeypatch): + # mock ComfyUI：不依赖 8182、不跑 Flux，只验证管线接线 + grown 字段 + import hairline.comfyui as comfy + monkeypatch.setattr(comfy, "run", lambda *a, **k: _PNG_1x1) + files = {"image_file": ("frontal.jpg", open(fixture("frontal.jpg"), "rb"), "application/octet-stream")} r = client.post(GROW, headers=H, files=files, data={"gender": "female"}) body = r.json() @@ -90,6 +99,7 @@ def test_grow_female_returns_5(client): assert [x["hairline_type"] for x in results] == ["ellipse", "flower", "heart", "straight", "wave"] assert [x["order"] for x in results] == [1, 2, 3, 4, 5] assert base64.b64decode(results[0]["image_base64"])[:8] == b"\x89PNG\r\n\x1a\n" + assert base64.b64decode(results[0]["grown_image_base64"])[:8] == b"\x89PNG\r\n\x1a\n" assert "image_url" not in results[0] diff --git a/tests/test_mask.py b/tests/test_mask.py new file mode 100644 index 0000000..42d472f --- /dev/null +++ b/tests/test_mask.py @@ -0,0 +1,47 @@ +"""接口2 第二步 遮罩几何单元测试（纯 numpy/cv2，无需 SegFormer/torch）。""" +import numpy as np + +from hairline.mask import ( + forehead_upper_region, head_silhouette, _above_curve_region, + compose_comfy_rgba, FOREHEAD_LANDMARKS, +) + + +def test_forehead_upper_region(): + w = h = 200 + lm = np.zeros((468, 3), np.float32) + for i in FOREHEAD_LANDMARKS: # 额头边界点都放在 y=0.4 + lm[i, 0] = 0.5 + lm[i, 1] = 0.4 + m = forehead_upper_region(lm, w, h) + assert m[10, 100] == 255 # 上方在区域内 + assert m[150, 100] == 0 # 下方(0.75h)在区域外 + + +def test_head_silhouette(): + pm = np.zeros((50, 50), np.int32) + pm[:25] = 13 # hair + pm[25:40] = 1 # skin + pm[40:] = 0 # bg + m = head_silhouette(pm) + assert m[10, 10] == 255 and m[30, 10] == 255 # hair/skin → 头部 + assert m[45, 10] == 0 # bg → 非头部 + + +def test_above_curve_region(): + h = w = 100 + curve = np.zeros((h, w), np.uint8) + curve[60, 20:80] = 255 # 水平发际线在 row=60 + above = _above_curve_region(curve, h, w) + assert above[30, 50] == 255 # 线以上 + assert above[80, 50] == 0 # 线以下 + + +def test_compose_comfy_rgba_alpha(): + marked = np.zeros((10, 10, 3), np.uint8) + mask = np.zeros((10, 10), np.uint8) + mask[2:5, 2:5] = 255 # 重绘区 + rgba = np.array(compose_comfy_rgba(marked, mask)) + assert rgba.shape == (10, 10, 4) + assert rgba[3, 3, 3] == 0 # 重绘区 alpha=0(透明) + assert rgba[8, 8, 3] == 255 # 非重绘区 alpha=255(不透明)