接口5:新增 hairline_level 档位(middle/high/low)贴图 + 集成接口2生发能力
- hairline_level: 可选 middle(默认)/high/low,选用不同高度档位发际线贴图; 新增 hairline_texture_high/low 两套贴图,get_texture_map 改为按档位缓存 - hair_style: 可选逗号分隔序号,对选中发际线类型同步生发(ComfyUI), 结果合并进 hairline_images[].grown_image_base64(未选中为 null); 生发黑模板固定取 middle(hairline_texture_black/),与档位无关 - 新增 use_mask/prompt 生发控制参数(同接口2) - 测试页/接口文档同步更新 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+66
-14
@@ -29,13 +29,20 @@ _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")
|
||||
|
||||
# 发际线贴图档位: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"),
|
||||
}
|
||||
|
||||
# ⚠️ 本 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
|
||||
_texture_maps: dict = {} # {level: {gender: [(key, path)]}},按档位缓存
|
||||
|
||||
|
||||
def get_landmarker() -> FaceLandmarker:
|
||||
@@ -61,24 +68,27 @@ def _gender_key(stem: str):
|
||||
return None, None
|
||||
|
||||
|
||||
def get_texture_map() -> dict:
|
||||
"""扫描 hairline_texture/ 建 {gender: [(key, path)]},按 key 排序、缓存。
|
||||
def get_texture_map(level: str = "middle") -> dict:
|
||||
"""扫描指定档位贴图目录建 {gender: [(key, path)]},按 key 排序、按档位缓存。
|
||||
|
||||
level:middle(默认) / high / low,分别对应 hairline_texture[/_high|/_low]。
|
||||
文件名规范化去空格(如 `man_ inverse_arc.png` → key `inverse_arc`)。
|
||||
"""
|
||||
global _texture_map
|
||||
if _texture_map is not None:
|
||||
return _texture_map
|
||||
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_DIR, "*.png"))):
|
||||
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_map = mapping
|
||||
return _texture_map
|
||||
_texture_maps[level] = mapping
|
||||
return mapping
|
||||
|
||||
|
||||
def extract_502(image_bgr: np.ndarray):
|
||||
@@ -190,10 +200,40 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
|
||||
return results
|
||||
|
||||
|
||||
def generate_hairline_pngs(image_bgr: np.ndarray, gender: str):
|
||||
"""接口5:该性别全部发际线叠图(同接口2预览) + 最佳(order1)发际线曲线的面部中间点。
|
||||
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
|
||||
|
||||
Returns: {"images":[{hairline_type,order,image_bgr}], "best_center":(x,y)};无人脸 None。
|
||||
|
||||
def generate_hairline_pngs(image_bgr: np.ndarray, gender: str, hairline_level: str = "middle",
|
||||
hair_styles: list[int] | None = None, use_mask: bool = True,
|
||||
prompt: str | None = None):
|
||||
"""接口5:该性别全部发际线叠图(同接口2预览) + 最佳(order1)发际线曲线的面部中间点,
|
||||
并对 hair_styles 选中的类型生发(同接口2)。
|
||||
|
||||
hairline_level:middle(默认)/high/low,选用不同档位的发际线贴图(仅影响叠图)。
|
||||
hair_styles(1-indexed 列表,按贴图排序):指定对哪些类型生发;None/[] 时不生发。
|
||||
生发黑模板固定取自 hairline_texture_black/(middle),与 hairline_level 无关。
|
||||
use_mask/prompt:同接口2 的生发参数。
|
||||
Returns: {"images":[{hairline_type,order,image_bgr,grown_png}], "best_center":(x,y)};无人脸 None。
|
||||
"""
|
||||
if gender not in ("male", "female"):
|
||||
raise ValueError(f"gender 必须是 male/female,收到 {gender!r}")
|
||||
@@ -206,12 +246,24 @@ def generate_hairline_pngs(image_bgr: np.ndarray, gender: str):
|
||||
# 面部中轴 x = 眉心(9/151 中点)
|
||||
face_cx = float((lm[9, 0] + lm[151, 0]) / 2 * w)
|
||||
|
||||
textures = get_texture_map()[gender]
|
||||
textures = get_texture_map(hairline_level)[gender]
|
||||
grow_set = set(hair_styles or [])
|
||||
|
||||
# use_mask=False:干净原图+空遮罩与贴图无关,只跑一次 ComfyUI,选中项复用
|
||||
shared_grown = None
|
||||
if grow_set and not use_mask:
|
||||
shared_grown = _grow_from_texture(image_bgr, ctx, None, use_mask=False, prompt=prompt)
|
||||
|
||||
images, best_center = [], None
|
||||
for order, (key, path) in enumerate(textures, start=1):
|
||||
white = load_texture_rgba(path)
|
||||
preview = render_hairline_overlay(image_bgr, ctx["points"], ext_faces, uv, white)
|
||||
images.append({"hairline_type": key, "order": order, "image_bgr": preview})
|
||||
grown_png = None
|
||||
if order in grow_set:
|
||||
grown_png = shared_grown if not use_mask else \
|
||||
_grow_from_texture(image_bgr, ctx, path, use_mask=True, prompt=prompt)
|
||||
images.append({"hairline_type": key, "order": order,
|
||||
"image_bgr": preview, "grown_png": grown_png})
|
||||
if order == 1: # 最佳发际线曲线的中点(面部中轴处的发际线 y)
|
||||
overlay = build_overlay_layer(h, w, ctx["points"], ext_faces, uv, white)
|
||||
ys, xs = np.where(overlay[:, :, 3] > 40)
|
||||
|
||||
Reference in New Issue
Block a user