1134 lines
57 KiB
Python
1134 lines
57 KiB
Python
"""接口11：发际线生发。
|
||
|
||
输入一张发际线较高 / 头发稀少的正脸图 + 发际线类型 ID（= change_hair 的 hair_id，
|
||
如 chang_tuoyuan/chang_bolang/...），输出同一个人、同一发型、按该发际线类型压低发际线
|
||
后的图片。管线（见 docs/发际线增强算法.md）：
|
||
|
||
1. 用接口9 的算法算出头发遮罩（含额头闭合区域，外缘内缩 erode_cm）。
|
||
—— seg_model 选 bisenet/segformer，mask_type 选 eroded(内缩)/closed(未内缩闭合区域)。
|
||
2. 调 change_hair 换发型服务（/api/swapHair/v1）生成该发际线类型的图。返回的结果图已被
|
||
change_hair 用 M_inv 贴回、与输入原图**同分辨率同对齐**，可直接按遮罩合成。
|
||
两种取图模式（swap_mode）：
|
||
- ext_mask：把步骤1 的遮罩作为 ext_mask 传给 swapHair，让 webui 精确重绘该区域
|
||
（忠于算法文档「换发型的遮罩用接口9 遮罩」）。
|
||
- as_is：不改 change_hair，swapHair 用它自己的内部遮罩，贴回时再裁到接口9 遮罩。
|
||
3. 严格按接口9 遮罩把生成图贴回原图（遮罩外=原图，纹丝不动）。
|
||
4. 融合接缝：blend_method 选 feather(高斯羽化) / alpha_gradient(距离变换内渐变) /
|
||
seamless(泊松无缝克隆) / multiband(多频段金字塔融合)；feather_px、edge_erode_px
|
||
控制过渡细节（feather_px 仅 feather/alpha_gradient 用；multiband 用 mb_levels 控制金字塔层数）。
|
||
可选 color_match=True 先在遮罩区做 Reinhard 颜色统计迁移，消除生成图与原图
|
||
的整体色差（对 feather/alpha_gradient/multiband 有效；seamless 自带色彩调和，自动跳过）。
|
||
|
||
对外返回每一步可视化（base64，data URI），供测试页逐步展示。经网关时 *_base64 字段会被
|
||
落盘改写为 *_url。
|
||
"""
|
||
import base64
|
||
import logging
|
||
import os
|
||
import time
|
||
from uuid import uuid4
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
from face_analysis.detector import detector
|
||
from face_analysis.calibration import estimate_scale_factor
|
||
from face_analysis.head_mask import (
|
||
NoFaceError,
|
||
BASELINE_IDX,
|
||
_baseline_points,
|
||
_upper_region_mask,
|
||
_bisenet_hair_mask,
|
||
_segformer_hair_mask,
|
||
_fill_to_baseline,
|
||
_erode,
|
||
_largest_cc,
|
||
_overlay,
|
||
_draw_baseline,
|
||
)
|
||
|
||
# 调试日志：写 <仓库根>/log/hairline_grow.log，每个步骤详细记录（可用 HAIR_LOG_DIR 覆盖）
|
||
_LOG_DIR = os.getenv(
|
||
"HAIR_LOG_DIR",
|
||
os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "log"),
|
||
)
|
||
os.makedirs(_LOG_DIR, exist_ok=True)
|
||
logger = logging.getLogger("hairline_grow")
|
||
_log_fh = logging.FileHandler(os.path.join(_LOG_DIR, "hairline_grow.log"), encoding="utf-8")
|
||
_log_fh.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
|
||
logger.addHandler(_log_fh)
|
||
logger.setLevel(logging.DEBUG)
|
||
|
||
# change_hair 服务地址（可用环境变量覆盖）
|
||
SWAP_URL = os.getenv("SWAP_HAIR_URL", "http://127.0.0.1:8801/api/swapHair/v1")
|
||
HAIRGROW_URL = os.getenv("HAIR_GROW_URL", "http://127.0.0.1:8801/api/hairGrow/v1")
|
||
SWAP_TIMEOUT = float(os.getenv("SWAP_HAIR_TIMEOUT", "300"))
|
||
|
||
# 多频段融合最细层羽化：羽化最细 FEATHER_LAYERS 层（每层核尺寸按尺度放大）。
|
||
# 只羽最细1层效果极弱（其拉普拉斯系数幅度小），羽化 3 层才能明显软化发丝边缘锯齿。
|
||
FEATHER_LAYERS = 3
|
||
|
||
DEFAULTS = {
|
||
"gen_backend": "swaphair", # swaphair(换发型LoRA) | hairgrow(区域生发inpaint)
|
||
"is_hr": False,
|
||
"seg_model": "segformer", # bisenet | segformer
|
||
"hairline_push_cm": 1.0, # 发际线径向外推距离（厘米）
|
||
"hairline_edge": "column", # column（逐列下沿）
|
||
"swap_mode": "ext_mask", # ext_mask | as_is（仅 swaphair）
|
||
"denoising_strength": 0.6, # 仅 swaphair
|
||
"hairgrow_strength": 0.75, # 仅 hairgrow
|
||
"edge_erode_px": 3,
|
||
"mb_levels": 5, # multiband 金字塔层数（2~6，越大色差抹得越宽）
|
||
"erode_cm": 0.6, # 接口12 固定值（pushed 模式下仅用于 baseline 截断参考，影响很小）
|
||
}
|
||
|
||
|
||
class SwapError(Exception):
|
||
"""调用 change_hair 换发型服务失败。"""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 编码
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _jpg_b64(bgr):
|
||
ok, buf = cv2.imencode(".jpg", bgr, [cv2.IMWRITE_JPEG_QUALITY, 92])
|
||
return "data:image/jpeg;base64," + base64.b64encode(buf.tobytes()).decode()
|
||
|
||
|
||
def _png_b64(bgr_or_gray):
|
||
ok, buf = cv2.imencode(".png", bgr_or_gray)
|
||
return "data:image/png;base64," + base64.b64encode(buf.tobytes()).decode()
|
||
|
||
|
||
def _gray_b64(gray_float):
|
||
"""0~1 的浮点图 → 灰度 PNG data URI。"""
|
||
g = np.clip(gray_float * 255.0, 0, 255).astype(np.uint8)
|
||
return _png_b64(g)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 步骤1：接口9 头发遮罩（复用 head_mask 构件）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _close_cyclic(mask, gap):
|
||
"""对 1D 布尔环形序列做闭运算，填掉短于 gap 的 False 缺口（消除内侧判定的小抖动）。"""
|
||
if gap <= 0 or mask.all() or not mask.any():
|
||
return mask
|
||
n = len(mask)
|
||
tripled = np.concatenate([mask, mask, mask]).astype(np.uint8).reshape(1, -1)
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2 * gap + 1, 1))
|
||
closed = cv2.morphologyEx(tripled, cv2.MORPH_CLOSE, kernel).reshape(-1)
|
||
return closed[n:2 * n].astype(bool)
|
||
|
||
|
||
def _longest_true_run_cyclic(mask):
|
||
"""返回环形布尔序列中最长连续 True 段的索引（按顺序）。全 True 返回全体索引。"""
|
||
n = len(mask)
|
||
if mask.all():
|
||
return np.arange(n)
|
||
if not mask.any():
|
||
return np.empty(0, dtype=np.int64)
|
||
# 从一个 False 处断开成线性序列，避免最长段跨越首尾
|
||
start = int(np.where(~mask)[0][0])
|
||
order = np.roll(np.arange(n), -start)
|
||
rmask = mask[order]
|
||
best_len = best_s = cur = cur_s = 0
|
||
for i, v in enumerate(rmask):
|
||
if v:
|
||
if cur == 0:
|
||
cur_s = i
|
||
cur += 1
|
||
if cur > best_len:
|
||
best_len, best_s = cur, cur_s
|
||
else:
|
||
cur = 0
|
||
return order[best_s:best_s + best_len]
|
||
|
||
|
||
def _draw_polyline(image, pts, color, thickness=3):
|
||
"""把有序点列 (Nx2) 画成折线（用于内轮廓/外推线可视化）。点<2 原样返回。"""
|
||
out = image.copy()
|
||
if pts is not None and len(pts) >= 2:
|
||
poly = np.ascontiguousarray(np.asarray(pts).reshape(-1, 1, 2), dtype=np.int32)
|
||
cv2.polylines(out, [poly], isClosed=False, color=color,
|
||
thickness=thickness, lineType=cv2.LINE_AA)
|
||
return out
|
||
|
||
|
||
def _extract_hairline(hair_mask, center, chin_y=None, rid="", sample_px=12):
|
||
"""提取头发区域朝脸一侧的内轮廓线（发际线）。
|
||
|
||
结果是一条有序折线：额头弧线 + 左右两侧鬓角/脸颊边界，一直向下到下颌 chin_y。
|
||
做法（轮廓 + 内侧判定，见 docs/发际线生发遮罩算法_pushed模式.md ①-f）：
|
||
1. 取头发 mask 最大连通域的外轮廓（CHAIN_APPROX_NONE，逐像素稠密点，保序）。
|
||
2. 逐点判定「内侧」：从该轮廓点朝脸中心 center 采样 sample_px 像素，若落点是
|
||
非头发像素 → 该点朝向脸（头发/皮肤交界的内轮廓）；否则是朝背景的外侧剪影，丢弃。
|
||
3. 内轮廓点在闭合轮廓上本是一段连续弧，先环形闭运算填小缝，再取最长连续段并保序。
|
||
4. 下颌截断：丢掉 y > chin_y 的点（把两侧末端截到下颌一带），得到红真值那样的环脸弧。
|
||
|
||
center: 脸中心 (cx, cy)，取 151 点（眉心）。内侧判定与 ①-g 径向外推共用此圆心。
|
||
chin_y: 下颌截断 y（一般取下巴关键点 152 的 y）；None 则不截断。
|
||
sample_px: 内侧判定的采样距离（像素），一般 ≈ 0.4cm。
|
||
返回按顺序排列的内轮廓点 np.ndarray(Nx2, int32)；不足 2 点返回空数组 (0,2)。
|
||
"""
|
||
lg = lambda msg: logger.info("[%s] %s", rid, msg) if rid else None
|
||
h, w = hair_mask.shape
|
||
empty = np.empty((0, 2), dtype=np.int32)
|
||
|
||
m = _largest_cc(hair_mask).astype(np.uint8)
|
||
if m.sum() == 0:
|
||
lg("_extract_hairline: 头发 mask 为空")
|
||
return empty
|
||
cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
||
if not cnts:
|
||
lg("_extract_hairline: 无轮廓")
|
||
return empty
|
||
big = max(cnts, key=cv2.contourArea)
|
||
pts = big[:, 0].astype(np.float64) # N×2 (x, y)，沿边界有序
|
||
n = len(pts)
|
||
lg(f"_extract_hairline: 最大轮廓点数={n} 面积={cv2.contourArea(big):.0f} center={center} chin_y={chin_y}")
|
||
if n < 3 or center is None:
|
||
return empty
|
||
|
||
cx, cy = float(center[0]), float(center[1])
|
||
# 内侧判定：每个轮廓点朝脸中心方向采样 sample_px，落在非头发上 → 内轮廓点
|
||
d = np.stack([cx - pts[:, 0], cy - pts[:, 1]], axis=1)
|
||
dist = np.hypot(d[:, 0], d[:, 1])
|
||
dist[dist < 1e-3] = 1.0
|
||
step = d / dist[:, None] * float(sample_px)
|
||
sx = np.clip(np.round(pts[:, 0] + step[:, 0]).astype(int), 0, w - 1)
|
||
sy = np.clip(np.round(pts[:, 1] + step[:, 1]).astype(int), 0, h - 1)
|
||
inner = m[sy, sx] == 0
|
||
lg(f" 内侧点数={int(inner.sum())}/{n}")
|
||
if inner.sum() < 3:
|
||
lg(" 内侧点<3，退化为空内轮廓")
|
||
return empty
|
||
|
||
# 环形闭运算填掉短缝隙，再取最长连续内侧段（= 朝脸的整段内轮廓）
|
||
inner = _close_cyclic(inner, gap=max(3, n // 100))
|
||
idx = _longest_true_run_cyclic(inner)
|
||
arc = pts[idx].astype(np.int32)
|
||
lg(f" 最长内侧段点数={len(arc)}")
|
||
|
||
# 下颌截断：只保留 y <= chin_y 的点（把两侧末端截到下颌）
|
||
if chin_y is not None and len(arc):
|
||
keep = arc[:, 1] <= int(chin_y)
|
||
arc = arc[keep]
|
||
lg(f" 下颌截断(chin_y={int(chin_y)})后点数={len(arc)}")
|
||
if len(arc) < 2:
|
||
return empty
|
||
|
||
# 稠密点降采样，减小后续多边形/外推开销（保序）
|
||
if len(arc) > 600:
|
||
arc = arc[:: len(arc) // 600 + 1]
|
||
return arc
|
||
|
||
|
||
def _pushed_mask(hair_mask, upper, baseline_pts, push_px, rid="",
|
||
center=None, chin_y=None, sample_px=12):
|
||
"""发际线外推遮罩：先提取头发内轮廓线（①-f），以眉心为圆心把内轮廓逐点径向
|
||
「向外」（远离脸中心、推进现有头发）外推 push_px 得到外推线（①-g 黄线）；
|
||
再取外推线与 baseline（①-a 分割线）组成的闭合区域作为最终遮罩。
|
||
|
||
即逐列从外推线 pushed_y[x] 填充到 baseline_y[x]（仅 pushed_y 在 baseline 以上的列），
|
||
与旧逻辑一致：遮罩顶界=外推发际线（覆盖现有头发 push_cm），底界=baseline。
|
||
区别只是现在 pushed_y 来自修正后的整条内轮廓（额头弧已延伸到两侧鬓角），
|
||
额头遮罩宽度不再被截短。
|
||
|
||
center: 圆心 (cx, cy)，取 151 点（眉心）；内侧判定与径向外推共用。
|
||
chin_y: 下颌截断 y，透传给 _extract_hairline。
|
||
sample_px: 内侧判定采样距离，透传给 _extract_hairline。
|
||
返回 (mask_bool, inner_pts, outer_pts)：
|
||
inner_pts —— 头发内轮廓有序点列 (Nx2)。
|
||
outer_pts —— 内轮廓径向外推 push_px 后的有序点列 (Nx2)。
|
||
"""
|
||
lg = lambda msg: logger.info("[%s] %s", rid, msg) if rid else None
|
||
h, w = hair_mask.shape
|
||
empty = np.empty((0, 2), dtype=np.int32)
|
||
|
||
inner_pts = _extract_hairline(hair_mask, center=center, chin_y=chin_y,
|
||
rid=rid, sample_px=sample_px)
|
||
lg(f"_pushed_mask: push_px={push_px} 内轮廓点数={len(inner_pts)} 圆心={center}")
|
||
if len(inner_pts) < 2 or center is None:
|
||
return np.zeros((h, w), dtype=bool), inner_pts, empty
|
||
|
||
cx, cy = float(center[0]), float(center[1])
|
||
# 逐点径向外推：沿「从圆心指向该点」方向（远离脸中心 = 推进现有头发）外推 push_px
|
||
d = inner_pts.astype(np.float64) - np.array([cx, cy])
|
||
dist = np.hypot(d[:, 0], d[:, 1])
|
||
dist[dist < 1e-3] = 1.0
|
||
u = d / dist[:, None]
|
||
outer = inner_pts.astype(np.float64) + u * float(push_px)
|
||
outer[:, 0] = np.clip(outer[:, 0], 0, w - 1)
|
||
outer[:, 1] = np.clip(outer[:, 1], 0, h - 1)
|
||
outer_pts = np.round(outer).astype(np.int32)
|
||
|
||
# baseline 每列的 y（①-a 分割线，含左右水平延长）
|
||
x0, y0 = baseline_pts[0]
|
||
x1, y1 = baseline_pts[-1]
|
||
chain_x = np.array([0] + [p[0] for p in baseline_pts] + [w - 1])
|
||
chain_y = np.array([y0] + [p[1] for p in baseline_pts] + [y1])
|
||
baseline_y = np.interp(np.arange(w), chain_x, chain_y)
|
||
|
||
# 外推线逐列归并：只取 baseline 以上的外推点，逐列取最靠上 y 作为遮罩顶界 pushed_y。
|
||
# 两侧鬓角向下的段落到 baseline 以下，自然被排除（与旧逻辑一致，遮罩=额头闭合区域）。
|
||
ox = np.clip(outer_pts[:, 0], 0, w - 1)
|
||
oy = outer_pts[:, 1].astype(np.float64)
|
||
above = oy < baseline_y[ox]
|
||
pushed_y = np.full(w, np.nan)
|
||
for xi, yi in zip(ox[above], oy[above]):
|
||
if np.isnan(pushed_y[xi]) or yi < pushed_y[xi]:
|
||
pushed_y[xi] = yi
|
||
cols_valid = np.where(~np.isnan(pushed_y))[0]
|
||
if len(cols_valid) >= 2:
|
||
lo, hi = int(cols_valid.min()), int(cols_valid.max())
|
||
pushed_y[lo:hi + 1] = np.interp(np.arange(lo, hi + 1), cols_valid,
|
||
pushed_y[cols_valid])
|
||
lg(f" 外推线额带[{lo},{hi}] 推后y范围[{int(np.nanmin(pushed_y))},{int(np.nanmax(pushed_y))}]")
|
||
|
||
# 逐列从 pushed_y 填充到 baseline_y（闭合区域），仅 pushed_y<baseline_y 的列
|
||
baseline_yi = baseline_y.astype(np.int32)
|
||
mask = np.zeros((h, w), dtype=bool)
|
||
fill_valid = ~np.isnan(pushed_y)
|
||
pushed_int = np.nan_to_num(pushed_y, nan=float(h)).astype(np.int32)
|
||
cols = np.where(fill_valid & (pushed_int < baseline_yi))[0]
|
||
lg(f" 有效填充列数={len(cols)} (pushed_y<baseline_y)")
|
||
for x in cols:
|
||
mask[int(pushed_int[x]):baseline_yi[x] + 1, x] = True
|
||
mask = _largest_cc(mask & upper)
|
||
lg(f" 闭合区域 mask 像素={int(mask.sum())}")
|
||
return mask, inner_pts, outer_pts
|
||
|
||
|
||
|
||
def _redraw_band_mask(inner_pts, outer_pts, h, w, rid="", upper=None,
|
||
lo_mult=0.5, hi_mult=1.5):
|
||
"""重绘带遮罩：由发际线（①-f 内轮廓 inner_pts）沿径向外推方向，取
|
||
`lo_mult × push` 与 `hi_mult × push` 两条外推线之间的带状区域作为重绘 mask。
|
||
|
||
inner_pts / outer_pts 是一一对应的有序点列（outer = inner 径向外推 1.0×push_px），
|
||
故每点的 1.0× 位移向量 = outer - inner；下界线 = inner + lo_mult×位移，
|
||
上界线 = inner + hi_mult×位移。闭合环 = 下界线（正向）+ 上界线（反向）首尾相接。
|
||
|
||
lo_mult / hi_mult：外推倍率（相对 hairline_push_cm）。默认 0.5 / 1.5，即带位于
|
||
0.5×push ~ 1.5×push 之间（以内轮廓为 0×、原外推线为 1.0×）。
|
||
|
||
upper：①-a baseline 以上区域布尔掩码。传入时把重绘带与它求交集，只保留 baseline
|
||
以上的部分（两侧鬓角落到 baseline 以下的段会被截掉）。
|
||
|
||
返回 band_bool。
|
||
"""
|
||
lg = lambda msg: logger.info("[%s] %s", rid, msg) if rid else None
|
||
if len(inner_pts) < 2 or len(outer_pts) < 2:
|
||
return np.zeros((h, w), dtype=bool)
|
||
inner_f = np.asarray(inner_pts, dtype=np.float32)
|
||
outer_f = np.asarray(outer_pts, dtype=np.float32)
|
||
disp = outer_f - inner_f # 每点 1.0×push 的径向位移向量
|
||
lo_line = inner_f + float(lo_mult) * disp # 下界外推线（lo_mult×push）
|
||
hi_line = inner_f + float(hi_mult) * disp # 上界外推线（hi_mult×push）
|
||
# 闭合多边形：下界线正向 + 上界线反向，端点自然相连
|
||
ring = np.vstack([lo_line.astype(np.int32), hi_line[::-1].astype(np.int32)])
|
||
band_u8 = np.zeros((h, w), dtype=np.uint8)
|
||
cv2.fillPoly(band_u8, [ring], 255)
|
||
band = band_u8 > 0
|
||
raw_px = int(band.sum())
|
||
# ①-a baseline 截断：只保留 baseline 以上的重绘带
|
||
if upper is not None:
|
||
band = band & upper
|
||
lg(f"_redraw_band_mask: 内轮廓点={len(inner_pts)} lo_mult={lo_mult} hi_mult={hi_mult} "
|
||
f"band像素(截断前)={raw_px} band像素(截断后)={int(band.sum())}")
|
||
return band
|
||
|
||
|
||
def compute_mask(image_bgr, landmarks, seg_model, mask_type, erode_cm, px_per_cm,
|
||
hairline_push_cm=0.0, hairline_edge="column", rid=""):
|
||
"""算出布尔遮罩 + 可视化。
|
||
|
||
seg_model: bisenet | segformer。
|
||
mask_type: eroded(外缘内缩) | closed(闭合区域未内缩) | pushed(发际线外推)。
|
||
hairline_push_cm: 仅 pushed 模式——发际线往头发方向外推的厘米数（进入现有头发）。
|
||
hairline_edge: 仅 pushed 模式——发际线提取方式 column(逐列最低点) | contour(形态学轮廓)。
|
||
rid: 调用方的 request id，用于日志关联。
|
||
返回 (mask_bool, viz_dict)。
|
||
"""
|
||
lg = lambda msg: logger.info("[%s] %s", rid, msg) if rid else None
|
||
lg(f"compute_mask 入参: mask_type={mask_type!r} erode_cm={erode_cm} "
|
||
f"px_per_cm={px_per_cm:.3f} hairline_push_cm={hairline_push_cm} hairline_edge={hairline_edge!r}")
|
||
|
||
h, w = image_bgr.shape[:2]
|
||
r = int(round(max(0.0, erode_cm) * px_per_cm))
|
||
lg(f"图像尺寸 {w}x{h}, erode_px={r}")
|
||
baseline_pts = _baseline_points(landmarks, w, h)
|
||
upper = _upper_region_mask(baseline_pts, w, h)
|
||
lg(f"baseline 第一点={baseline_pts[0]} 末点={baseline_pts[-1]} upper像素={int(upper.sum())}")
|
||
|
||
if seg_model == "bisenet":
|
||
hair_mask = _bisenet_hair_mask(image_bgr, landmarks, w, h)
|
||
elif seg_model == "segformer":
|
||
hair_mask = _segformer_hair_mask(image_bgr)
|
||
else:
|
||
raise ValueError(f"未知 seg_model: {seg_model}")
|
||
lg(f"头发分割完成 seg_model={seg_model} hair_pixels={int(hair_mask.sum())}")
|
||
|
||
top_fill = _fill_to_baseline(hair_mask, upper) # 含额头，延伸到图底
|
||
closed = _largest_cc(top_fill & upper) # 闭合区域：头发+额头，底=基线
|
||
eroded = _largest_cc(_erode(top_fill, r) & upper) # 外缘内缩 r、底线不动
|
||
lg(f"旧流程: top_fill像素={int(top_fill.sum())} closed像素={int(closed.sum())} eroded像素={int(eroded.sum())}")
|
||
|
||
# pushed 模式：发际线外推遮罩（额外保留 hairline_y/pushed_y/额带边界 供可视化）
|
||
pushed_info = None
|
||
if mask_type == "pushed":
|
||
push_px = int(round(max(0.0, hairline_push_cm) * px_per_cm))
|
||
# 圆心 = 151 点（眉心）完整坐标，内侧判定与径向外推共用
|
||
# 按值查 151 在 BASELINE_IDX 中的位置，避免列表变动后索引错位（曾硬编码 [5]）
|
||
_idx151 = BASELINE_IDX.index(151) if 151 in BASELINE_IDX else -1
|
||
center = baseline_pts[_idx151] if _idx151 >= 0 else None
|
||
# 下颌截断线：下巴关键点 152 的 y（内轮廓两侧向下画到这里为止）
|
||
try:
|
||
chin_y = int(round(landmarks.landmark[152].y * h))
|
||
except Exception: # noqa: BLE001
|
||
chin_y = None
|
||
# 内侧判定采样距离 ≈ 0.4cm
|
||
sample_px = max(6, int(round(0.4 * px_per_cm)))
|
||
lg(f"进入 PUSHED 分支: push_px={push_px} 圆心(151)={center} chin_y={chin_y} "
|
||
f"sample_px={sample_px} edge={hairline_edge}")
|
||
mask_bool, inner_pts, outer_pts = _pushed_mask(
|
||
hair_mask, upper, baseline_pts, push_px, rid=rid,
|
||
center=center, chin_y=chin_y, sample_px=sample_px)
|
||
pushed_info = (inner_pts, outer_pts, push_px)
|
||
lg(f"PUSHED 结果: 内轮廓点数={len(inner_pts)} mask_pixels={int(mask_bool.sum())}")
|
||
elif mask_type == "eroded":
|
||
mask_bool = eroded
|
||
lg(f"进入 ERODED 分支: 用 eroded 遮罩 pixels={int(eroded.sum())}")
|
||
else:
|
||
mask_bool = closed
|
||
lg(f"进入 CLOSED 分支: 用 closed 遮罩 pixels={int(closed.sum())}")
|
||
|
||
lg(f"最终遮罩 mask_type={mask_type} mask_pixels={int(mask_bool.sum())}")
|
||
|
||
# 遮罩计算过程可视化：
|
||
# eroded/closed 走 top_fill→closed/eroded 流程；
|
||
# pushed 走 baseline→头发分割→发际线→外推 流程，与 top_fill/closed 无关，故置空。
|
||
viz = {
|
||
"erode_px": r,
|
||
"hair_pixels": int(hair_mask.sum()),
|
||
"closed_pixels": int(closed.sum()),
|
||
"mask_pixels": int(mask_bool.sum()),
|
||
# 1. 发际线分割线（baseline）：151 中心点标红，其余点标绿，黄线含左右延长线
|
||
"baseline_overlay_base64": _jpg_b64(_draw_baseline(image_bgr, baseline_pts, w)),
|
||
# 2. 分割线以上区域（upper 半区）：青色叠加
|
||
"upper_overlay_base64": _jpg_b64(_overlay(image_bgr, upper, (0, 255, 255))),
|
||
# 3. 头发分割原始结果（hair_mask）：绿色叠加在原图上
|
||
"hair_seg_overlay_base64": _jpg_b64(_overlay(image_bgr, hair_mask, (0, 255, 0))),
|
||
# 4. top_fill / closed —— 仅 eroded/closed 流程用；pushed 流程无关，留空
|
||
"top_fill_overlay_base64": "" if mask_type == "pushed"
|
||
else _jpg_b64(_overlay(image_bgr, top_fill, (255, 0, 0))),
|
||
"closed_overlay_base64": "" if mask_type == "pushed"
|
||
else _jpg_b64(_overlay(image_bgr, closed, (255, 0, 255))),
|
||
# 5. pushed 模式专有（发际线提取/外推）—— 非 pushed 留空
|
||
"hairline_overlay_base64": "",
|
||
"pushed_overlay_base64": "",
|
||
# —— 最终遮罩 ——
|
||
"mask_overlay_base64": _jpg_b64(_overlay(image_bgr, mask_bool, (0, 0, 255))),
|
||
"mask_base64": _png_b64((mask_bool.astype(np.uint8)) * 255),
|
||
}
|
||
# pushed 模式：补充内轮廓提取 + 外推线可视化
|
||
if pushed_info is not None:
|
||
inner_pts, outer_pts, push_px = pushed_info
|
||
# ①-f 提取内轮廓：绿=头发内轮廓线（额头弧+两侧到下颌），黄=baseline 折线
|
||
hl_img = _draw_baseline(image_bgr, baseline_pts, w) # 画 baseline（黄线+关键点）
|
||
hl_img = _draw_polyline(hl_img, inner_pts, (0, 255, 0), 3)
|
||
viz["hairline_overlay_base64"] = _jpg_b64(hl_img)
|
||
# ①-g 外推：圆心红点(151) + 内轮廓（绿）+ 外推线（青）+ 遮罩（红半透明）
|
||
ps_img = _draw_polyline(image_bgr.copy(), inner_pts, (0, 255, 0), 2)
|
||
ps_img = _draw_polyline(ps_img, outer_pts, (0, 255, 255), 3)
|
||
# 画圆心（151 点）红点，标示径向外推的中心（_idx151 上方已按值查到）
|
||
if center is not None:
|
||
cx151, cy151 = center
|
||
cv2.circle(ps_img, (cx151, cy151), 6, (0, 0, 255), -1, cv2.LINE_AA)
|
||
cv2.putText(ps_img, "151", (cx151 + 8, cy151 - 8),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1, cv2.LINE_AA)
|
||
ps_img = _overlay(ps_img, mask_bool, (0, 0, 255), 0.3)
|
||
viz["pushed_overlay_base64"] = _jpg_b64(ps_img)
|
||
viz["push_px"] = push_px
|
||
# 重绘带用原始数据：内轮廓点 + 外推点（供 _redraw_band_mask 连端点成带）
|
||
viz["_inner_pts"] = inner_pts
|
||
viz["_outer_pts"] = outer_pts
|
||
# baseline 以上区域，供重绘带按 ①-a baseline 截断（只留上面）
|
||
viz["_upper_mask"] = upper
|
||
# 记录 viz 各字段是否非空（长度），便于排查前端取不到图的问题
|
||
viz_summary = {k: (len(v) if isinstance(v, str) and v else 0)
|
||
for k, v in viz.items() if k.endswith("_base64")}
|
||
lg(f"viz 生成完毕，各图字节长度: {viz_summary}")
|
||
return mask_bool, viz
|
||
|
||
|
||
def _segment_hair(image_bgr, seg_model, landmarks, w, h):
|
||
"""对任意图（如 hard_paste 重绘结果）重跑头发分割，返回 bool 掩码。
|
||
|
||
与 compute_mask 内部用的同一个 seg_model 逻辑（bisenet 需 landmarks，
|
||
segformer 不需要），保证第1步（原图头发）与第2步（重绘后头发）分割口径一致。
|
||
"""
|
||
if seg_model == "bisenet":
|
||
return _bisenet_hair_mask(image_bgr, landmarks, w, h)
|
||
elif seg_model == "segformer":
|
||
return _segformer_hair_mask(image_bgr)
|
||
else:
|
||
raise ValueError(f"未知 seg_model: {seg_model}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 步骤2：调 change_hair 换发型
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _call_swap(image_bgr, hairline_id, is_hr, ext_mask_bool, denoising_strength,
|
||
inpainting_fill=1, mask_blur=11, mask_dilate_scale=1.0):
|
||
"""调 change_hair /api/swapHair/v1，返回与输入同分辨率同对齐的换发型结果（BGR）。
|
||
|
||
ext_mask_bool 非 None 时作为 ext_mask 传入（swap_mode=ext_mask）。
|
||
denoising_strength：webui img2img 重绘强度（越大生发越激进），透传给换发型。
|
||
inpainting_fill / mask_blur / mask_dilate_scale：服务端重绘参数（透传给 change_hair，
|
||
默认值=服务端原始硬编码值，未传时行为不变）。详见 change_hair 文档。
|
||
"""
|
||
import requests
|
||
|
||
ok, ibuf = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
payload = {
|
||
"hair_id": hairline_id,
|
||
"task_id": "if11-" + uuid4().hex[:12],
|
||
"is_hr": "true" if is_hr else "false",
|
||
"user_img_path": "data:image/jpeg;base64," + base64.b64encode(ibuf.tobytes()).decode(),
|
||
"output_format": "base64",
|
||
"denoising_strength": float(denoising_strength),
|
||
"inpainting_fill": int(inpainting_fill),
|
||
"mask_blur": int(mask_blur),
|
||
"mask_dilate_scale": float(mask_dilate_scale),
|
||
}
|
||
if ext_mask_bool is not None:
|
||
mbuf = cv2.imencode(".png", (ext_mask_bool.astype(np.uint8)) * 255)[1]
|
||
payload["ext_mask"] = "data:image/png;base64," + base64.b64encode(mbuf.tobytes()).decode()
|
||
|
||
try:
|
||
resp = requests.post(SWAP_URL, json=payload, timeout=SWAP_TIMEOUT)
|
||
except Exception as ex: # noqa: BLE001
|
||
raise SwapError(f"换发型服务不可达（{SWAP_URL}）：{ex}")
|
||
|
||
try:
|
||
j = resp.json()
|
||
except Exception: # noqa: BLE001
|
||
raise SwapError(f"换发型服务返回非 JSON（HTTP {resp.status_code}）：{resp.text[:200]}")
|
||
|
||
if j.get("state") != 0 or not j.get("data"):
|
||
raise SwapError(f"换发型失败：{j.get('msg', j)}")
|
||
|
||
b64 = j["data"]
|
||
if "," in b64 and b64.startswith("data:"):
|
||
b64 = b64.split(",", 1)[1]
|
||
result = cv2.imdecode(np.frombuffer(base64.b64decode(b64), np.uint8), cv2.IMREAD_COLOR)
|
||
if result is None:
|
||
raise SwapError("换发型结果解码失败")
|
||
|
||
# 保险：与原图对齐（change_hair 已贴回原尺寸，若极端情况尺寸不符则拉回）
|
||
if result.shape[:2] != image_bgr.shape[:2]:
|
||
result = cv2.resize(result, (image_bgr.shape[1], image_bgr.shape[0]),
|
||
interpolation=cv2.INTER_LANCZOS4)
|
||
return result
|
||
|
||
|
||
def _call_hairgrow(image_bgr, mask_bool, strength):
|
||
"""调 change_hair /api/hairGrow/v1（区域生发 inpaint），在遮罩区域长出头发。
|
||
|
||
返回与输入同分辨率的结果（BGR）。hairGrow 内部已做贴回与颜色迁移，
|
||
这里再套接口11 的遮罩羽化贴回以保证遮罩外严格不动。
|
||
"""
|
||
import requests
|
||
|
||
ok, ibuf = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
mbuf = cv2.imencode(".png", (mask_bool.astype(np.uint8)) * 255)[1]
|
||
payload = {
|
||
"img": "data:image/jpeg;base64," + base64.b64encode(ibuf.tobytes()).decode(),
|
||
"mask": "data:image/png;base64," + base64.b64encode(mbuf.tobytes()).decode(),
|
||
"strength": float(strength),
|
||
"output_format": "base64",
|
||
}
|
||
try:
|
||
resp = requests.post(HAIRGROW_URL, json=payload, timeout=SWAP_TIMEOUT)
|
||
except Exception as ex: # noqa: BLE001
|
||
raise SwapError(f"区域生发服务不可达（{HAIRGROW_URL}）：{ex}")
|
||
try:
|
||
j = resp.json()
|
||
except Exception: # noqa: BLE001
|
||
raise SwapError(f"区域生发返回非 JSON（HTTP {resp.status_code}）：{resp.text[:200]}")
|
||
if j.get("state") != 0 or not j.get("result"):
|
||
raise SwapError(f"区域生发失败：{j.get('msg', j)}")
|
||
b64 = j["result"]
|
||
if "," in b64 and b64.startswith("data:"):
|
||
b64 = b64.split(",", 1)[1]
|
||
result = cv2.imdecode(np.frombuffer(base64.b64decode(b64), np.uint8), cv2.IMREAD_COLOR)
|
||
if result is None:
|
||
raise SwapError("区域生发结果解码失败")
|
||
if result.shape[:2] != image_bgr.shape[:2]:
|
||
result = cv2.resize(result, (image_bgr.shape[1], image_bgr.shape[0]),
|
||
interpolation=cv2.INTER_LANCZOS4)
|
||
return result
|
||
|
||
|
||
_REPAINT_WORKFLOW = os.path.join(os.path.dirname(os.path.dirname(__file__)), "hair_repaint.json")
|
||
|
||
|
||
def _call_comfyui(image_bgr, mask_bool, prompt=None):
|
||
"""调本机 ComfyUI 的 Flux-2 inpaint 工作流（hair_repaint.json），返回与输入同分辨率的 BGR。
|
||
|
||
与 swapHair 的区别：ComfyUI 把「原图 VAE 编码作 reference latent + ColorMatch」双重保色，
|
||
天生不易染色；提示词自由可调（中文）。mask 经 RGBA alpha 通道传入（透明=重绘区）。
|
||
ComfyUI 不在线时抛 SwapError（由调用方捕获降级）。prompt=None 用工作流内置默认提示词。
|
||
"""
|
||
import io
|
||
from hairline.mask import compose_comfy_rgba
|
||
from hairline.comfyui import run as comfyui_run, ping
|
||
|
||
if not ping():
|
||
raise SwapError("ComfyUI 不可达（http://127.0.0.1:8188），redraw Flux-2 路跳过")
|
||
mask_u8 = (mask_bool.astype(np.uint8)) * 255
|
||
rgba_img = compose_comfy_rgba(image_bgr, mask_u8) # alpha=255-mask：透明=重绘区
|
||
buf = io.BytesIO()
|
||
rgba_img.save(buf, format="PNG")
|
||
png_bytes = comfyui_run(buf.getvalue(), prompt=prompt, workflow_path=_REPAINT_WORKFLOW)
|
||
result = cv2.imdecode(np.frombuffer(png_bytes, np.uint8), cv2.IMREAD_COLOR)
|
||
if result is None:
|
||
raise SwapError("ComfyUI 结果解码失败")
|
||
if result.shape[:2] != image_bgr.shape[:2]:
|
||
result = cv2.resize(result, (image_bgr.shape[1], image_bgr.shape[0]),
|
||
interpolation=cv2.INTER_LANCZOS4)
|
||
return result
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 步骤3+4：按遮罩贴回 + 接缝融合
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _color_match_to_orig(swap_result, orig, mask_bool, strength=1.0):
|
||
"""在 mask_bool 区域内做 Reinhard 颜色迁移：逐通道把 swap_result 的均值/方差对齐 orig。
|
||
|
||
strength 控制迁移强度：1.0=完全对齐到 orig（原行为），<1.0 只迁移部分，
|
||
防止 Reinhard 在某些图上过度改色（如把生成发色整体拉向皮肤色）。
|
||
遮罩外保持 swap_result 原样（不会越界污染）。返回 uint8 BGR。
|
||
"""
|
||
m = mask_bool.astype(bool)
|
||
src_f = swap_result.astype(np.float32)
|
||
out = src_f.copy()
|
||
if m.sum() < 30:
|
||
return swap_result.copy()
|
||
strength = float(min(max(strength, 0.0), 1.0))
|
||
for c in range(3):
|
||
src_pix = swap_result[..., c][m].astype(np.float32)
|
||
dst_pix = orig[..., c][m].astype(np.float32)
|
||
s_mean, s_std = src_pix.mean(), src_pix.std() + 1e-6
|
||
d_mean, d_std = dst_pix.mean(), dst_pix.std() + 1e-6
|
||
aligned = (out[..., c] - s_mean) * (d_std / s_std) + d_mean
|
||
out[..., c] = src_f[..., c] * (1.0 - strength) + aligned * strength
|
||
return np.clip(out, 0, 255).astype(np.uint8)
|
||
|
||
|
||
def _feather_alpha(mask_bool, blend_method, feather_px, edge_erode_px):
|
||
"""由布尔遮罩生成 0~1 的 alpha（贴图权重）。遮罩外恒为 0（原图纹丝不动）。"""
|
||
m = mask_bool.astype(np.uint8)
|
||
if edge_erode_px > 0:
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * edge_erode_px + 1,) * 2)
|
||
m = cv2.erode(m, k)
|
||
|
||
fp = max(1, int(feather_px))
|
||
if blend_method == "alpha_gradient":
|
||
# 距离变换：过渡只发生在遮罩内侧（边界 0 → 内部 feather_px 处 1），遮罩外严格为 0
|
||
dist = cv2.distanceTransform(m, cv2.DIST_L2, 3)
|
||
alpha = np.clip(dist / fp, 0.0, 1.0)
|
||
else: # feather（高斯羽化，默认）
|
||
ksz = fp * 2 + 1
|
||
alpha = cv2.GaussianBlur(m.astype(np.float32), (ksz, ksz), sigmaX=fp / 2.0)
|
||
alpha = np.clip(alpha, 0.0, 1.0)
|
||
return alpha
|
||
|
||
|
||
def _multiband_alpha(mask_bool, edge_erode_px):
|
||
"""多频段融合用的二值掩码：先内缩、保证最小边距，否则最小一层金字塔会塌缩。
|
||
|
||
返回 uint8 二值 {0,255}（拉普拉斯金字塔融合要求起始掩码为二值，否则粗层会把整图混色）。
|
||
"""
|
||
m = mask_bool.astype(np.uint8) * 255
|
||
if edge_erode_px > 0:
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * edge_erode_px + 1,) * 2)
|
||
m = cv2.erode(m, k)
|
||
return m
|
||
|
||
|
||
def _multiband_blend(orig, swap_result, mask_bool, levels, edge_erode_px,
|
||
feather_px=1, transition_band_px=-1):
|
||
"""多频段（拉普拉斯金字塔）融合：低频用宽窗抹色差，高频用窄窗保发丝。
|
||
|
||
levels：金字塔层数（2~6），越大则低频色差在越宽范围被抹平。
|
||
feather_px：最细若干层掩码轻羽化像素（0=不羽化，保持硬二值）。羽化最细 FEATHER_LAYERS
|
||
层（核尺寸按层尺度放大），消除发丝边缘 1px 硬切锯齿；粗层仍保持二值（否则粗层会
|
||
把整图混色）。注意：这是消除锯齿的微调，幅度有限（边界 Δ 约 1~3/255），
|
||
不要指望它做大范围过渡——那是 mb_levels/transition_band_px 的事。
|
||
transition_band_px：keep-region 外缘边距。-1=自动按层数 2**n（旧行为）；
|
||
>=0 则用绝对像素，使过渡带宽度与金字塔层数解耦。
|
||
返回 uint8 BGR。
|
||
"""
|
||
m = _multiband_alpha(mask_bool, edge_erode_px)
|
||
if m.sum() < 255:
|
||
return orig.copy()
|
||
|
||
# 层数受分辨率上限约束：每层尺寸减半，最小一层至少 4px，否则金字塔塌缩
|
||
min_dim = min(orig.shape[:2])
|
||
max_by_res = int(np.floor(np.log2(min_dim / 4))) if min_dim >= 16 else 1
|
||
n = int(max(1, min(levels, max_by_res)))
|
||
if n < 2:
|
||
# 极小图退化：直接按内缩遮罩硬贴，避免单层金字塔无意义
|
||
out = orig.copy()
|
||
m_bool = _multiband_alpha(mask_bool, edge_erode_px) > 127
|
||
out[m_bool] = swap_result[m_bool]
|
||
return out
|
||
|
||
def lap_pyr(img, n):
|
||
pyr = [img.astype(np.float32)]
|
||
cur = img.astype(np.float32)
|
||
for _ in range(n):
|
||
cur = cv2.pyrDown(cur)
|
||
pyr.append(cur)
|
||
laps = [pyr[-1]]
|
||
for i in range(n, 0, -1):
|
||
size = (pyr[i - 1].shape[1], pyr[i - 1].shape[0])
|
||
up = cv2.pyrUp(pyr[i], dstsize=size)
|
||
laps.append(pyr[i - 1] - up)
|
||
return laps # [最粗层, 细节层L1, ..., 最细层Ln]
|
||
|
||
def mask_pyr(mask_u8, n):
|
||
# 起始必须二值；逐层 pyrDown 后自动变软（金字塔天然多频段软掩码）。
|
||
# 返回顺序与 lap_pyr 一致：粗 → 细。
|
||
pyr = [mask_u8.astype(np.float32) / 255.0]
|
||
cur = mask_u8.astype(np.float32) / 255.0
|
||
for _ in range(n):
|
||
cur = cv2.pyrDown(cur)
|
||
pyr.append(cur)
|
||
return list(reversed(pyr)) # 与 lap_pyr 同尺度（最粗层在前）
|
||
|
||
la = lap_pyr(orig, n)
|
||
lb = lap_pyr(swap_result, n)
|
||
ma = mask_pyr(m, n)
|
||
|
||
# 最细层（reversed 后末元素 = 全分辨率原始二值掩码）及其下若干层轻羽化，
|
||
# 消除发丝边缘 1px 硬切锯齿。注意：多频段融合中各层都贡献边界过渡，但最细层的
|
||
# 拉普拉斯系数幅度最小，只羽化它效果很弱（实测边界 Δ 仅 ~0.25/255）。因此对最细
|
||
# FEATHER_LAYERS 层都做按尺度放大的羽化（越细的层核越大），才能明显软化边缘。
|
||
# 粗层（低频）仍保持二值，否则会把整图混色，违反多频段融合的二值掩码前提。
|
||
fp = int(max(0, feather_px))
|
||
if fp > 0:
|
||
for li in range(1, FEATHER_LAYERS + 1):
|
||
idx = -li
|
||
if abs(idx) > len(ma):
|
||
break
|
||
scale = 2 ** (li - 1)
|
||
ksz = fp * 2 * scale + 1
|
||
ma[idx] = cv2.GaussianBlur(ma[idx], (ksz, ksz), sigmaX=fp * scale / 2.0)
|
||
|
||
merged = []
|
||
for a, b, mk in zip(la, lb, ma):
|
||
m3 = mk[:, :, None]
|
||
merged.append(a * (1 - m3) + b * m3)
|
||
|
||
out = merged[0]
|
||
for i in range(1, len(merged)):
|
||
size = (merged[i].shape[1], merged[i].shape[0])
|
||
out = cv2.pyrUp(out, dstsize=size)
|
||
out = out + merged[i]
|
||
out = np.clip(out, 0, 255).astype(np.uint8)
|
||
|
||
# 契约：遮罩远区纹丝不动，但保留多频段的过渡带。多频段融合的意义就在于低频层
|
||
# （粗层）的掩码在 pyrDown/pyrUp 后向外扩散变软，形成一条随层数变宽的过渡带——
|
||
# 这条带正是 mb_levels 要控制的东西。若像旧实现那样用原始硬二值遮罩钳回，
|
||
# 过渡带会被整条抹掉（实测 levels 2↔6 边界差恒为 0），mb_levels 形同虚设。
|
||
# 故按层数膨胀出一个外缘 keep 区：keep 内允许过渡，keep 外才强制还原原图。
|
||
if transition_band_px is not None and transition_band_px >= 0:
|
||
margin = int(transition_band_px) # 与金字塔层数解耦，用绝对像素
|
||
else:
|
||
margin = 2 ** n # n=2→4px … n=6→64px，与粗层掩码的自然扩散宽度匹配
|
||
margin = max(0, margin)
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * margin + 1, 2 * margin + 1))
|
||
keep = cv2.dilate(mask_bool.astype(np.uint8), k).astype(bool)
|
||
out[~keep] = orig[~keep]
|
||
return out
|
||
|
||
|
||
def _seamless_clone(orig, swap_result, mask_bool, edge_erode_px):
|
||
"""泊松无缝克隆（cv2.seamlessClone NORMAL_CLONE）：梯度域调和整体色调。
|
||
|
||
返回调色后的整帧 uint8 BGR；掩码过小（<10px）时返回原图。
|
||
供 seamless 分支与 two_stage 两段式融合的第一段复用。
|
||
"""
|
||
m = mask_bool.astype(np.uint8)
|
||
if edge_erode_px > 0:
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * edge_erode_px + 1,) * 2)
|
||
m = cv2.erode(m, k)
|
||
if m.sum() < 10:
|
||
return orig.copy()
|
||
ys, xs = np.where(m > 0)
|
||
center = (int((xs.min() + xs.max()) / 2), int((ys.min() + ys.max()) / 2))
|
||
return cv2.seamlessClone(swap_result, orig, m * 255, center, cv2.NORMAL_CLONE)
|
||
|
||
|
||
def _composite(orig, swap_result, mask_bool, blend_method, feather_px, edge_erode_px,
|
||
color_match=False, mb_levels=5, color_match_strength=1.0,
|
||
mb_feather_px=1, transition_band_px=-1):
|
||
"""把 swap_result 按遮罩贴回 orig，返回 (final_bgr, alpha_float or None)。
|
||
|
||
blend_method:
|
||
- multiband : 多频段金字塔融合（默认）
|
||
- seamless : 泊松无缝克隆（梯度域调色，自带色彩调和，故跳过 color_match）
|
||
- two_stage : 先 seamless 统一整体色调，再 multiband 贴发丝细节（大色差场景）
|
||
- feather/alpha_gradient : 单层 alpha 过渡
|
||
"""
|
||
# seamless / two_stage 自带梯度域色彩调和，不叠 Reinhard 颜色迁移
|
||
if blend_method == "seamless":
|
||
final = _seamless_clone(orig, swap_result, mask_bool, edge_erode_px)
|
||
return final, None
|
||
|
||
if blend_method == "two_stage":
|
||
# 第一段：seamless 把整体色调拉平（生成图色调对齐到原图）
|
||
harmonized = _seamless_clone(orig, swap_result, mask_bool, edge_erode_px)
|
||
# 第二段：对调色后的结果再做 multiband 贴发丝细节（不加 color_match，避免重复改色）
|
||
final = _multiband_blend(orig, harmonized, mask_bool, mb_levels, edge_erode_px,
|
||
feather_px=mb_feather_px,
|
||
transition_band_px=transition_band_px)
|
||
alpha = (_multiband_alpha(mask_bool, edge_erode_px).astype(np.float32)) / 255.0
|
||
return final, alpha
|
||
|
||
# multiband / feather / alpha_gradient：先做 Reinhard 颜色迁移消除整体色差
|
||
src = (_color_match_to_orig(swap_result, orig, mask_bool, color_match_strength)
|
||
if color_match else swap_result)
|
||
|
||
if blend_method == "multiband":
|
||
final = _multiband_blend(orig, src, mask_bool, mb_levels, edge_erode_px,
|
||
feather_px=mb_feather_px,
|
||
transition_band_px=transition_band_px)
|
||
# 可视化用：用多频段的二值掩码做一层 alpha 标记（展示实际合成区）
|
||
alpha = (_multiband_alpha(mask_bool, edge_erode_px).astype(np.float32)) / 255.0
|
||
return final, alpha
|
||
|
||
alpha = _feather_alpha(mask_bool, blend_method, feather_px, edge_erode_px)
|
||
a3 = alpha[:, :, None]
|
||
final = (orig.astype(np.float32) * (1 - a3) + src.astype(np.float32) * a3)
|
||
return np.clip(final, 0, 255).astype(np.uint8), alpha
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 主入口
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _grow_core(image_bgr, hairline_id, *, is_hr, seg_model, erode_cm, swap_mode,
|
||
edge_erode_px, denoising_strength, gen_backend, hairgrow_strength,
|
||
mb_levels, hairline_push_cm, hairline_edge, blend_method, color_match,
|
||
color_match_strength, mb_feather_px, transition_band_px,
|
||
inpainting_fill, mask_blur, mask_dilate_scale, rid):
|
||
"""接口11 共享核心：遮罩(pushed)→生成→硬贴回→接缝融合，产出 ④ final。
|
||
|
||
不做任何重绘。返回中间产物 dict（供接口11 构造响应、接口12 取 final+重绘带用）：
|
||
final / swap_result / hard_paste / alpha / mask_bool / mask_viz /
|
||
px_per_cm / t_mask / t_swap / t_blend / h / w
|
||
未检出人脸抛 NoFaceError。
|
||
"""
|
||
mask_type = "pushed" # 固定：只支持 pushed 遮罩算法
|
||
logger.info("[%s] _grow_core 参数(固定 mask=pushed): erode_cm=%s hairline_push_cm=%s "
|
||
"hairline_edge=%r mb_levels=%s seg=%s gen_backend=%s swap_mode=%s blend=%s "
|
||
"color_match=%s cm_strength=%s mb_feather_px=%s transition_band_px=%s "
|
||
"inpainting_fill=%s mask_blur=%s mask_dilate_scale=%s",
|
||
rid, erode_cm, hairline_push_cm, hairline_edge, mb_levels,
|
||
seg_model, gen_backend, swap_mode, blend_method, color_match,
|
||
color_match_strength, mb_feather_px, transition_band_px,
|
||
inpainting_fill, mask_blur, mask_dilate_scale)
|
||
h, w = image_bgr.shape[:2]
|
||
landmarks = detector.detect(image_bgr)
|
||
if landmarks is None:
|
||
logger.warning("[%s] 未检出人脸", rid)
|
||
raise NoFaceError()
|
||
px_per_cm = estimate_scale_factor(landmarks, w, h)
|
||
logger.info("[%s] 人脸检出 px_per_cm=%.3f 图尺寸=%dx%d", rid, px_per_cm, w, h)
|
||
|
||
# 步骤1：接口9 遮罩（固定 pushed）
|
||
t0 = time.time()
|
||
mask_bool, mask_viz = compute_mask(
|
||
image_bgr, landmarks, seg_model, mask_type, erode_cm, px_per_cm,
|
||
hairline_push_cm=hairline_push_cm, hairline_edge=hairline_edge, rid=rid)
|
||
t_mask = time.time() - t0
|
||
logger.info("[%s] 步骤1 遮罩完成 耗时=%dms mask_pixels=%d", rid, int(t_mask*1000), int(mask_bool.sum()))
|
||
|
||
# 步骤2：生成（按后端）
|
||
t0 = time.time()
|
||
if gen_backend == "hairgrow":
|
||
swap_result = _call_hairgrow(image_bgr, mask_bool, hairgrow_strength)
|
||
else:
|
||
ext_mask = mask_bool if swap_mode == "ext_mask" else None
|
||
swap_result = _call_swap(image_bgr, hairline_id, is_hr, ext_mask, denoising_strength,
|
||
inpainting_fill=inpainting_fill, mask_blur=mask_blur,
|
||
mask_dilate_scale=mask_dilate_scale)
|
||
t_swap = time.time() - t0
|
||
|
||
# 步骤3：严格按遮罩硬贴回（无融合，用于对比）
|
||
hard_paste = image_bgr.copy()
|
||
hard_paste[mask_bool] = swap_result[mask_bool]
|
||
|
||
# 步骤4：接缝融合（默认 multiband）→ ④ final
|
||
t0 = time.time()
|
||
final, alpha = _composite(
|
||
image_bgr, swap_result, mask_bool, blend_method, 0, edge_erode_px,
|
||
color_match=color_match, mb_levels=mb_levels,
|
||
color_match_strength=color_match_strength,
|
||
mb_feather_px=mb_feather_px, transition_band_px=transition_band_px)
|
||
t_blend = time.time() - t0
|
||
|
||
return {
|
||
"final": final, "swap_result": swap_result, "hard_paste": hard_paste,
|
||
"alpha": alpha, "mask_bool": mask_bool, "mask_viz": mask_viz,
|
||
"px_per_cm": px_per_cm, "t_mask": t_mask, "t_swap": t_swap, "t_blend": t_blend,
|
||
"h": h, "w": w,
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 主入口
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def generate_hairline_grow(image_bgr, hairline_id, is_hr=False, seg_model="segformer",
|
||
erode_cm=0.6, swap_mode="ext_mask",
|
||
edge_erode_px=3,
|
||
denoising_strength=0.6, gen_backend="swaphair",
|
||
hairgrow_strength=0.75, mb_levels=5,
|
||
hairline_push_cm=1.0, hairline_edge="column",
|
||
blend_method="multiband", color_match=True,
|
||
color_match_strength=1.0, mb_feather_px=1,
|
||
transition_band_px=-1,
|
||
inpainting_fill=1, mask_blur=11, mask_dilate_scale=1.0,
|
||
rid=None):
|
||
"""接口11 完整管线（**不含重绘**，重绘见接口12 generate_hairline_redraw）。
|
||
返回可直接进 ok() 的 data dict。未检出人脸抛 NoFaceError。
|
||
|
||
遮罩算法固定为 pushed（发际线外推）。
|
||
融合算法 blend_method 默认 multiband（多频段金字塔），可选 seamless(泊松)/
|
||
two_stage(泊松→多频段两段式)/feather(羽化)/alpha_gradient(距离变换)。
|
||
color_match 默认开启 Reinhard 颜色迁移消除整体色差（对 multiband/feather 有效）。
|
||
inpainting_fill/mask_blur/mask_dilate_scale：透传 change_hair 服务端换发型重绘参数。
|
||
rid: 调用方的 request id，用于日志关联。为 None 时自动生成。
|
||
"""
|
||
if rid is None:
|
||
rid = uuid4().hex[:8]
|
||
logger.info("[%s] ===== generate_hairline_grow 开始 =====", rid)
|
||
core = _grow_core(
|
||
image_bgr, hairline_id, is_hr=is_hr, seg_model=seg_model, erode_cm=erode_cm,
|
||
swap_mode=swap_mode, edge_erode_px=edge_erode_px, denoising_strength=denoising_strength,
|
||
gen_backend=gen_backend, hairgrow_strength=hairgrow_strength, mb_levels=mb_levels,
|
||
hairline_push_cm=hairline_push_cm, hairline_edge=hairline_edge, blend_method=blend_method,
|
||
color_match=color_match, color_match_strength=color_match_strength,
|
||
mb_feather_px=mb_feather_px, transition_band_px=transition_band_px,
|
||
inpainting_fill=inpainting_fill, mask_blur=mask_blur,
|
||
mask_dilate_scale=mask_dilate_scale, rid=rid)
|
||
mask_viz = core["mask_viz"]
|
||
alpha = core["alpha"]
|
||
w, h = core["w"], core["h"]
|
||
|
||
data = {
|
||
"hairline_id": hairline_id,
|
||
"gen_backend": gen_backend,
|
||
"hairgrow_strength": round(float(hairgrow_strength), 3),
|
||
"is_hr": is_hr,
|
||
"seg_model": seg_model,
|
||
"mask_type": "pushed",
|
||
"erode_cm": round(float(erode_cm), 2),
|
||
"swap_mode": swap_mode,
|
||
"blend_method": blend_method,
|
||
"edge_erode_px": int(edge_erode_px),
|
||
"mb_levels": int(mb_levels),
|
||
"hairline_push_cm": round(float(hairline_push_cm), 2),
|
||
"hairline_edge": hairline_edge,
|
||
"denoising_strength": round(float(denoising_strength), 3),
|
||
"color_match": bool(color_match),
|
||
"color_match_strength": round(float(color_match_strength), 3),
|
||
"mb_feather_px": int(mb_feather_px),
|
||
"transition_band_px": int(transition_band_px),
|
||
"inpainting_fill": int(inpainting_fill),
|
||
"mask_blur": int(mask_blur),
|
||
"mask_dilate_scale": round(float(mask_dilate_scale), 3),
|
||
"px_per_cm": round(float(core["px_per_cm"]), 4),
|
||
"erode_px": mask_viz["erode_px"],
|
||
"hair_pixels": mask_viz["hair_pixels"],
|
||
"closed_pixels": mask_viz["closed_pixels"],
|
||
"mask_pixels": mask_viz["mask_pixels"],
|
||
"image_size": {"width": w, "height": h},
|
||
"timings_ms": {
|
||
"mask": int(core["t_mask"] * 1000),
|
||
"swap": int(core["t_swap"] * 1000),
|
||
"blend": int(core["t_blend"] * 1000),
|
||
},
|
||
"steps": {
|
||
"input_base64": _jpg_b64(image_bgr),
|
||
# 遮罩计算全过程（接口9 子步骤）
|
||
"baseline_overlay_base64": mask_viz["baseline_overlay_base64"],
|
||
"upper_overlay_base64": mask_viz["upper_overlay_base64"],
|
||
"hair_seg_overlay_base64": mask_viz["hair_seg_overlay_base64"],
|
||
"top_fill_overlay_base64": mask_viz["top_fill_overlay_base64"],
|
||
"closed_overlay_base64": mask_viz["closed_overlay_base64"],
|
||
# pushed 模式专有（非 pushed 时为空串）
|
||
"hairline_overlay_base64": mask_viz["hairline_overlay_base64"],
|
||
"pushed_overlay_base64": mask_viz["pushed_overlay_base64"],
|
||
# 最终遮罩
|
||
"mask_overlay_base64": mask_viz["mask_overlay_base64"],
|
||
"mask_base64": mask_viz["mask_base64"],
|
||
"swap_raw_base64": _jpg_b64(core["swap_result"]),
|
||
"hard_paste_base64": _jpg_b64(core["hard_paste"]),
|
||
"alpha_base64": _gray_b64(alpha) if alpha is not None else mask_viz["mask_base64"],
|
||
"final_base64": _jpg_b64(core["final"]),
|
||
},
|
||
"_rid": rid,
|
||
}
|
||
steps_summary = {k: (len(v) if isinstance(v, str) and v else 0)
|
||
for k, v in data["steps"].items() if k.endswith("_base64")}
|
||
logger.info("[%s] 返回 steps 字段长度: %s", rid, steps_summary)
|
||
logger.info("[%s] ===== generate_hairline_grow 完成 =====", rid)
|
||
return data
|
||
|
||
|
||
def generate_hairline_redraw(image_bgr, hairline_id, is_hr=False, seg_model="segformer",
|
||
erode_cm=0.6, swap_mode="ext_mask",
|
||
edge_erode_px=3,
|
||
denoising_strength=0.6, gen_backend="swaphair",
|
||
hairgrow_strength=0.75, mb_levels=5,
|
||
hairline_push_cm=1.0, hairline_edge="column",
|
||
blend_method="multiband", color_match=True,
|
||
color_match_strength=1.0, mb_feather_px=1,
|
||
transition_band_px=-1,
|
||
inpainting_fill=1, mask_blur=11, mask_dilate_scale=1.0,
|
||
comfyui_prompt=None, beauty_alpha=0.6,
|
||
band_lo_mult=0.5, band_hi_mult=1.5, rid=None):
|
||
"""接口12 发际线带重绘。内部先跑接口11 核心拿到 ④ final，再取 ⑤-① 发际线重绘带
|
||
（外推↔内推之间、经 baseline 截断只留上部）作遮罩，用 Flux-2（ComfyUI，hair_repaint.json）
|
||
重绘（band 经 alpha 送进 ComfyUI 决定加发位置，ComfyUI 输出为整帧重绘+美颜图）。
|
||
|
||
**同时产出两版结果供对比**：
|
||
- `redraw_full`：ComfyUI 整帧输出（全脸美颜 + 全脸重绘），与手动跑 ComfyUI 一致。
|
||
- `redraw_band`：加发只在发际线带、美颜保留全脸。band 内完全用 ComfyUI 重绘（加发），
|
||
band 外用 `final` 结构 + 按 `beauty_alpha` 融入 ComfyUI 的全脸美颜。
|
||
|
||
返回可直接进 ok() 的 data dict。未检出人脸抛 NoFaceError。
|
||
|
||
comfyui_prompt：Flux-2 提示词，None 用默认「补充遮罩区域的头发，加一点美颜」。
|
||
beauty_alpha：redraw_band 版 band 外的全脸美颜融入强度（0=完全保留 final 无美颜，
|
||
1=band 外也完全用 ComfyUI 输出≈redraw_full），默认 0.6。
|
||
band_lo_mult / band_hi_mult：重绘带外推倍率（相对 hairline_push_cm），带位于
|
||
lo×push ~ hi×push 之间（内轮廓=0×、原外推线=1.0×），默认 0.5 / 1.5。
|
||
其余参数含义与接口11 相同（用于内部生成 final 与重绘带）。
|
||
"""
|
||
if rid is None:
|
||
rid = uuid4().hex[:8]
|
||
logger.info("[%s] ===== generate_hairline_redraw 开始 =====", rid)
|
||
core = _grow_core(
|
||
image_bgr, hairline_id, is_hr=is_hr, seg_model=seg_model, erode_cm=erode_cm,
|
||
swap_mode=swap_mode, edge_erode_px=edge_erode_px, denoising_strength=denoising_strength,
|
||
gen_backend=gen_backend, hairgrow_strength=hairgrow_strength, mb_levels=mb_levels,
|
||
hairline_push_cm=hairline_push_cm, hairline_edge=hairline_edge, blend_method=blend_method,
|
||
color_match=color_match, color_match_strength=color_match_strength,
|
||
mb_feather_px=mb_feather_px, transition_band_px=transition_band_px,
|
||
inpainting_fill=inpainting_fill, mask_blur=mask_blur,
|
||
mask_dilate_scale=mask_dilate_scale, rid=rid)
|
||
final = core["final"]
|
||
mask_viz = core["mask_viz"]
|
||
w, h = core["w"], core["h"]
|
||
px_per_cm = core["px_per_cm"]
|
||
|
||
# ① 算重绘带（⑤-①）：发际线(内轮廓)↔外推发际线成带，经 baseline 截断只留上部
|
||
t0 = time.time()
|
||
inner_pts = mask_viz.get("_inner_pts")
|
||
outer_pts = mask_viz.get("_outer_pts")
|
||
upper_mask = mask_viz.get("_upper_mask")
|
||
push_px = int(round(max(0.0, hairline_push_cm) * px_per_cm))
|
||
redraw_band_overlay_b64 = ""
|
||
redraw_full_b64 = "" # A：ComfyUI 整帧（全脸美颜+全脸重绘）
|
||
redraw_band_b64 = "" # B：加发只在发际线带、美颜保留全脸
|
||
redraw_info = {"enabled": False}
|
||
band_mask = None
|
||
try:
|
||
band_mask = _redraw_band_mask(inner_pts, outer_pts, h, w, rid=rid, upper=upper_mask,
|
||
lo_mult=band_lo_mult, hi_mult=band_hi_mult)
|
||
if band_mask.sum() < 30:
|
||
raise RuntimeError("重绘带像素过少，可能内轮廓/外推线缺失")
|
||
logger.info("[%s] 重绘带 push_px=%d lo_mult=%s hi_mult=%s band_pixels=%d",
|
||
rid, push_px, band_lo_mult, band_hi_mult, int(band_mask.sum()))
|
||
redraw_band_overlay_b64 = _jpg_b64(_overlay(final, band_mask, (255, 0, 255)))
|
||
redraw_info = {"enabled": True, "band_pixels": int(band_mask.sum()), "push_px": push_px,
|
||
"band_lo_mult": float(band_lo_mult), "band_hi_mult": float(band_hi_mult)}
|
||
except Exception as ex: # noqa: BLE001
|
||
logger.exception("[%s] 重绘带计算失败，整个重绘跳过", rid)
|
||
redraw_info = {"enabled": False, "error": f"band: {ex}"}
|
||
|
||
# ② Flux-2 路：final + band 调 ComfyUI（hair_repaint.json 整图重绘 + reference latent
|
||
# 保色 + ColorMatch + 美颜）。band 经 alpha 送进 ComfyUI 决定加发位置。
|
||
beauty_alpha = float(min(max(beauty_alpha, 0.0), 1.0))
|
||
if redraw_info.get("enabled"):
|
||
try:
|
||
prompt = comfyui_prompt if comfyui_prompt else "补充遮罩区域的头发，加一点美颜"
|
||
redraw_c_raw = _call_comfyui(final, band_mask, prompt=prompt)
|
||
# A：整帧输出（与手动跑 ComfyUI 一致）
|
||
redraw_full_b64 = _jpg_b64(redraw_c_raw)
|
||
# B：加发只在 band、美颜保留全脸。
|
||
# alpha = band 内 1（羽化边缘）；band 外 = beauty_alpha。
|
||
# band 内完全用 ComfyUI（加发）；band 外用 final 结构 + beauty_alpha 融入全脸美颜。
|
||
feather_px = max(8, int(round(push_px * 0.6)))
|
||
band_a = _feather_alpha(band_mask, "feather", feather_px, 0) # 0..1，band 外=0
|
||
a = band_a + (1.0 - band_a) * beauty_alpha
|
||
a3 = a[:, :, None]
|
||
band_mix = (final.astype(np.float32) * (1.0 - a3)
|
||
+ redraw_c_raw.astype(np.float32) * a3)
|
||
redraw_band_b64 = _jpg_b64(np.clip(band_mix, 0, 255).astype(np.uint8))
|
||
redraw_info["beauty_alpha"] = beauty_alpha
|
||
logger.info("[%s] Flux-2 重绘完成：redraw_full(整帧) + redraw_band(局部加发+全脸美颜 beauty_alpha=%.2f)",
|
||
rid, beauty_alpha)
|
||
except Exception as ex: # noqa: BLE001
|
||
logger.warning("[%s] Flux-2 路重绘失败，跳过: %s", rid, ex)
|
||
redraw_info["c_error"] = str(ex)
|
||
t_redraw = time.time() - t0
|
||
|
||
data = {
|
||
"hairline_id": hairline_id,
|
||
"blend_method": blend_method,
|
||
"hairline_push_cm": round(float(hairline_push_cm), 2),
|
||
"comfyui_prompt": comfyui_prompt or "补充遮罩区域的头发，加一点美颜",
|
||
"beauty_alpha": beauty_alpha,
|
||
"px_per_cm": round(float(px_per_cm), 4),
|
||
"mask_pixels": mask_viz["mask_pixels"],
|
||
"image_size": {"width": w, "height": h},
|
||
"timings_ms": {
|
||
"mask": int(core["t_mask"] * 1000),
|
||
"swap": int(core["t_swap"] * 1000),
|
||
"blend": int(core["t_blend"] * 1000),
|
||
"redraw": int(t_redraw * 1000),
|
||
},
|
||
"steps": {
|
||
"input_base64": _jpg_b64(image_bgr),
|
||
# 接口11 的 ④ final —— 作为本接口的重绘输入基底
|
||
"final_base64": _jpg_b64(final),
|
||
# ⑤-① 发际线重绘带（紫，已按 baseline 截断只留上部）
|
||
"redraw_band_overlay_base64": redraw_band_overlay_b64,
|
||
# A：ComfyUI 整帧重绘+美颜（与手动跑 ComfyUI 一致）
|
||
"redraw_full_base64": redraw_full_b64,
|
||
# B：加发只在发际线带、美颜保留全脸
|
||
"redraw_band_base64": redraw_band_b64,
|
||
# 兼容旧字段：指向 A（整帧版）
|
||
"redraw_c_base64": redraw_full_b64,
|
||
},
|
||
"redraw": redraw_info,
|
||
"_rid": rid,
|
||
}
|
||
steps_summary = {k: (len(v) if isinstance(v, str) and v else 0)
|
||
for k, v in data["steps"].items() if k.endswith("_base64")}
|
||
logger.info("[%s] 返回 steps 字段长度: %s", rid, steps_summary)
|
||
logger.info("[%s] ===== generate_hairline_redraw 完成 =====", rid)
|
||
return data
|