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>
@@ -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"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -504,17 +504,22 @@ async def hair_grow(
|
|||||||
return err(1002, "人像分辨率过低")
|
return err(1002, "人像分辨率过低")
|
||||||
|
|
||||||
try:
|
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
|
# 预览 + 生发(ComfyUI) 都是阻塞且较慢,放线程池避免卡住事件循环
|
||||||
if previews is None:
|
items = await run_in_threadpool(generate_grow_results, image, gender)
|
||||||
|
if items is None:
|
||||||
return err(1001, "无法识别人像")
|
return err(1001, "无法识别人像")
|
||||||
|
|
||||||
results = []
|
results = []
|
||||||
for p in previews:
|
for p in items:
|
||||||
ok_enc, png = cv2.imencode(".png", p["image_bgr"])
|
_ok, png = cv2.imencode(".png", p["image_bgr"])
|
||||||
|
grown_b64 = (base64.b64encode(p["grown_png"]).decode()
|
||||||
|
if p["grown_png"] else None)
|
||||||
results.append({
|
results.append({
|
||||||
"image_base64": base64.b64encode(png.tobytes()).decode(),
|
"image_base64": base64.b64encode(png.tobytes()).decode(),
|
||||||
|
"grown_image_base64": grown_b64,
|
||||||
"hairline_type": p["hairline_type"],
|
"hairline_type": p["hairline_type"],
|
||||||
"order": p["order"],
|
"order": p["order"],
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -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)")
|
||||||
@@ -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())}")
|
||||||
@@ -78,19 +78,30 @@ def render_hairline_overlay(photo_bgr: np.ndarray,
|
|||||||
points502_norm: (502, 3) 归一化坐标(x,y ∈ [0,1]),**MP 顺序**(extract_hairline 输出)。
|
points502_norm: (502, 3) 归一化坐标(x,y ∈ [0,1]),**MP 顺序**(extract_hairline 输出)。
|
||||||
"""
|
"""
|
||||||
H, W = photo_bgr.shape[:2]
|
H, W = photo_bgr.shape[:2]
|
||||||
TH, TW = texture_rgba.shape[:2]
|
overlay = build_overlay_layer(H, W, points502_norm, ext_faces, uv502, texture_rgba)
|
||||||
pts_obj = mp_order_to_obj_order(points502_norm) # MP序 → OBJ序
|
# alpha 合成(RGBA→BGR:贴图 RGB 顺序需反成 BGR)
|
||||||
img_xy = pts_obj[:, :2] * np.array([W, H], dtype=np.float32) # (502,2)
|
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)
|
tex = texture_rgba.astype(np.float32)
|
||||||
for (i, j, k) in ext_faces:
|
for (i, j, k) in ext_faces:
|
||||||
dst = img_xy[[i, j, k]].astype(np.float32)
|
dst = img_xy[[i, j, k]].astype(np.float32)
|
||||||
# UV → 贴图像素;flipY:贴图 y = (1 - v_raw) * TH(与 head3d Three.js flipY=true 一致)
|
# 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)],
|
src = np.array([[uv502[v][0] * TW, (1.0 - uv502[v][1]) * TH] for v in (i, j, k)],
|
||||||
dtype=np.float32)
|
dtype=np.float32)
|
||||||
# 退化三角形(投影到一条线)跳过,避免 getAffineTransform 奇异
|
if cv2.contourArea(dst.astype(np.int32)) < 1.0: # 退化三角形跳过
|
||||||
if cv2.contourArea(dst.astype(np.int32)) < 1.0:
|
|
||||||
continue
|
continue
|
||||||
M = cv2.getAffineTransform(src, dst)
|
M = cv2.getAffineTransform(src, dst)
|
||||||
warped = cv2.warpAffine(tex, M, (W, H), flags=cv2.INTER_LINEAR,
|
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)
|
cv2.fillConvexPoly(tri_mask, dst.astype(np.int32), 255)
|
||||||
sel = tri_mask > 0
|
sel = tri_mask > 0
|
||||||
overlay[sel] = warped[sel]
|
overlay[sel] = warped[sel]
|
||||||
|
return overlay
|
||||||
# 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)
|
|
||||||
|
|||||||
@@ -11,13 +11,22 @@ import cv2
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from . import constants as C
|
from . import constants as C
|
||||||
|
from . import comfyui
|
||||||
from .face_landmarks import FaceLandmarker
|
from .face_landmarks import FaceLandmarker
|
||||||
from .face_parsing import FaceParser
|
from .face_parsing import FaceParser
|
||||||
from .hairline_2d import sample_hairline, smooth_hairline
|
from .hairline_2d import sample_hairline, smooth_hairline
|
||||||
from .lift_3d import lift_hairline_to_3d, build_middle_row, assemble_full
|
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 .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 跑算子会报
|
# ⚠️ 本 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。
|
# "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):
|
def extract_502(image_bgr: np.ndarray):
|
||||||
"""照片(BGR) → (points502 MP序, valid17)。无人脸返回 (None, None)。"""
|
"""照片(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)
|
rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
|
||||||
landmarks = get_landmarker().detect(rgb)
|
landmarks = get_landmarker().detect(rgb)
|
||||||
if landmarks is None:
|
if landmarks is None:
|
||||||
return None, None
|
return None
|
||||||
parse_map = get_parser().parse(rgb)
|
parse_map = get_parser().parse(rgb)
|
||||||
hairline_2d, valid = sample_hairline(landmarks, parse_map)
|
hairline_2d, valid = sample_hairline(landmarks, parse_map)
|
||||||
hairline_2d = smooth_hairline(hairline_2d, valid)
|
hairline_2d = smooth_hairline(hairline_2d, valid)
|
||||||
hairline_3d = lift_hairline_to_3d(landmarks, hairline_2d)
|
hairline_3d = lift_hairline_to_3d(landmarks, hairline_2d)
|
||||||
middle_3d = build_middle_row(landmarks, hairline_3d)
|
middle_3d = build_middle_row(landmarks, hairline_3d)
|
||||||
points = assemble_full(landmarks, middle_3d, 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):
|
def generate_previews(image_bgr: np.ndarray, gender: str):
|
||||||
"""生成该性别全部发际线预览图。
|
"""生成该性别全部发际线预览图(仅预览,不生发)。
|
||||||
|
|
||||||
Returns: list[dict],每项 {"hairline_type": key, "image_bgr": ndarray, "order": 1..N}。
|
Returns: list[dict] {"hairline_type", "image_bgr", "order"};无人脸返回 None。
|
||||||
无人脸返回 None。gender 必须是 male/female。
|
|
||||||
"""
|
"""
|
||||||
if gender not in ("male", "female"):
|
if gender not in ("male", "female"):
|
||||||
raise ValueError(f"gender 必须是 male/female,收到 {gender!r}")
|
raise ValueError(f"gender 必须是 male/female,收到 {gender!r}")
|
||||||
points, _valid = extract_502(image_bgr)
|
ctx = extract_context(image_bgr)
|
||||||
if points is None:
|
if ctx is None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
uv, ext_faces = load_ext_mesh()
|
uv, ext_faces = load_ext_mesh()
|
||||||
results = []
|
results = []
|
||||||
for order, (key, path) in enumerate(get_texture_map()[gender], start=1):
|
for order, (key, path) in enumerate(get_texture_map()[gender], start=1):
|
||||||
tex = load_texture_rgba(path)
|
preview = render_hairline_overlay(image_bgr, ctx["points"], ext_faces, uv,
|
||||||
preview = render_hairline_overlay(image_bgr, points, ext_faces, uv, tex)
|
load_texture_rgba(path))
|
||||||
results.append({"hairline_type": key, "image_bgr": preview, "order": order})
|
results.append({"hairline_type": key, "image_bgr": preview, "order": order})
|
||||||
return results
|
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__":
|
if __name__ == "__main__":
|
||||||
import sys
|
import sys
|
||||||
g = sys.argv[2] if len(sys.argv) > 2 else "female"
|
g = sys.argv[2] if len(sys.argv) > 2 else "female"
|
||||||
|
|||||||
|
Before Width: | Height: | Size: 5.5 KiB After Width: | Height: | Size: 5.8 KiB |
|
After Width: | Height: | Size: 3.8 KiB |
|
After Width: | Height: | Size: 4.2 KiB |
|
After Width: | Height: | Size: 3.9 KiB |
|
After Width: | Height: | Size: 3.8 KiB |
|
After Width: | Height: | Size: 4.1 KiB |
|
After Width: | Height: | Size: 4.3 KiB |
|
After Width: | Height: | Size: 3.8 KiB |
|
After Width: | Height: | Size: 4.0 KiB |
|
After Width: | Height: | Size: 3.4 KiB |
@@ -81,7 +81,16 @@ def test_grow_missing_gender_1004(client):
|
|||||||
assert r.json()["code"] == 1004
|
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")}
|
files = {"image_file": ("frontal.jpg", open(fixture("frontal.jpg"), "rb"), "application/octet-stream")}
|
||||||
r = client.post(GROW, headers=H, files=files, data={"gender": "female"})
|
r = client.post(GROW, headers=H, files=files, data={"gender": "female"})
|
||||||
body = r.json()
|
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["hairline_type"] for x in results] == ["ellipse", "flower", "heart", "straight", "wave"]
|
||||||
assert [x["order"] for x in results] == [1, 2, 3, 4, 5]
|
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]["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]
|
assert "image_url" not in results[0]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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(不透明)
|
||||||