xsl
443cfa298f
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
318 lines
14 KiB
Python
318 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""区域生发模块
|
||
|
||
在用户指定的 mask 区域内，通过 webui inpainting 生成头发（生发）。
|
||
不依赖换发型的 photo_service/LoRA 链路，直接调 webui 的 /sdapi/v1/img2img。
|
||
|
||
核心函数:
|
||
hair_grow(img, mask, strength) -> BGR ndarray
|
||
"""
|
||
import os
|
||
import cv2
|
||
import time
|
||
import base64
|
||
import numpy as np
|
||
import requests
|
||
|
||
# webui 地址：优先用环境变量，否则默认本机 57860
|
||
# （不依赖 common.logger，使本模块可被独立服务直接 import）
|
||
WEBUI_URL = os.environ.get("WEBUI_URL", "http://0.0.0.0:57860/")
|
||
|
||
# 固定随机种子，保证同输入同输出（可复现）
|
||
SEED = 123456789
|
||
|
||
|
||
def _encode_numpy_to_base64(img):
|
||
"""ndarray -> 裸 base64 字符串（webui sdapi 约定，不带 data: 前缀）"""
|
||
retval, bytes = cv2.imencode('.png', img)
|
||
return base64.b64encode(bytes).decode('utf-8')
|
||
|
||
|
||
def _webui_inpaint(img, mask, prompt, negative_prompt,
|
||
denoising_strength, mask_blur, steps=30, cfg_scale=7.0):
|
||
"""直连 webui img2img inpainting。
|
||
|
||
参数:
|
||
img: BGR ndarray (H,W,3) 原图
|
||
mask: 灰度 ndarray (H,W)，白色(255)=重绘区
|
||
prompt / negative_prompt: 文本提示
|
||
denoising_strength: 重绘强度 0~1
|
||
mask_blur: mask 边缘羽化像素数
|
||
返回:
|
||
BGR ndarray 生发结果
|
||
"""
|
||
url = f"{WEBUI_URL}sdapi/v1/img2img"
|
||
request_dict = {
|
||
"prompt": prompt,
|
||
"negative_prompt": negative_prompt,
|
||
"sampler_name": "DPM++ 2M Karras",
|
||
"batch_size": 1,
|
||
"steps": steps,
|
||
"width": img.shape[1],
|
||
"height": img.shape[0],
|
||
"cfg_scale": cfg_scale,
|
||
"seed": SEED,
|
||
"mask_blur": mask_blur,
|
||
"init_images": [_encode_numpy_to_base64(img)],
|
||
"inpaint_full_res": False, # 全图重绘模式，mask 像素精确对应原图
|
||
"inpainting_fill": 1, # mask 区填噪声后去噪（即 inpaint）
|
||
"inpainting_mask_invert": 0, # 重绘 mask 白色区域
|
||
"mask": _encode_numpy_to_base64(mask),
|
||
"denoising_strength": denoising_strength,
|
||
"alwayson_scripts": {}
|
||
}
|
||
|
||
start = time.time()
|
||
response = requests.post(url=url, json=request_dict, timeout=300)
|
||
ret_json = response.json()
|
||
|
||
if 'images' not in ret_json or not ret_json['images']:
|
||
err = ret_json.get('detail') or ret_json.get('error') or str(ret_json)
|
||
raise RuntimeError(f"webui inpainting 返回异常: {err}")
|
||
|
||
result_b64 = ret_json['images'][0]
|
||
# 去掉可能的 data:image/png;base64, 前缀
|
||
if "," in result_b64 and result_b64.startswith("data:"):
|
||
result_b64 = result_b64.split(",", 1)[1]
|
||
result_img = cv2.imdecode(
|
||
np.frombuffer(base64.b64decode(result_b64), np.uint8), cv2.IMREAD_COLOR)
|
||
print(f"[hair_grow] webui inpainting 完成，耗时 {time.time()-start:.1f}s, "
|
||
f"denoising={denoising_strength}, mask_blur={mask_blur}")
|
||
return result_img
|
||
|
||
|
||
def _estimate_hair_color(img, mask):
|
||
"""从原图估计头发的颜色统计（LAB空间均值/方差）。
|
||
|
||
策略：在 mask 周边（外扩区域）取"深色像素"作为已有头发样本。
|
||
用户涂抹的生发区周边必然紧邻真实头发，比从图像顶部取样更可靠。
|
||
返回 (mean_l, mean_a, mean_b), (std_l, std_a, std_b) in LAB。
|
||
"""
|
||
h, w = img.shape[:2]
|
||
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
mb = mask > 127
|
||
|
||
# 1. mask 外扩一圈（25px），作为"周边采样区"
|
||
dilated = cv2.dilate(mb.astype(np.uint8) * 255, np.ones((51, 51), np.uint8))
|
||
surround = (dilated > 0) & (~mb)
|
||
|
||
# 2. 取周边全部像素作为参考（含头发+皮肤过渡区，色调更平衡）
|
||
# 单独筛"深色"会偏向阴影色导致发黑，用全采样更稳健
|
||
if surround.sum() > 200:
|
||
samples = lab[surround]
|
||
mean = samples.mean(axis=0)
|
||
# 方差用原图全局方差（保证生发区有合理色彩范围，不发灰）
|
||
global_std = lab.reshape(-1, 3).std(axis=0)
|
||
std = np.maximum(global_std * 0.7, [15, 8, 8])
|
||
print(f"[hair_grow] 头发色估计(mask周边全采样) LAB: 均值={mean.astype(int)} "
|
||
f"方差={std.astype(int)} 样本{surround.sum()}px")
|
||
return mean, std
|
||
|
||
# 3. fallback: 图像顶部深色像素
|
||
top_gray = gray[:max(1, int(h * 0.25))]
|
||
top_lab = lab[:max(1, int(h * 0.25))]
|
||
thresh = np.percentile(top_gray, 30)
|
||
hair_mask = (top_gray <= thresh) & (top_gray >= 5)
|
||
samples = top_lab[hair_mask] if hair_mask.sum() > 200 else top_lab.reshape(-1, 3)
|
||
mean = samples.mean(axis=0)
|
||
std = np.maximum(samples.std(axis=0), [15, 8, 8])
|
||
print(f"[hair_grow] 头发色估计(顶部fallback) LAB: 均值={mean.astype(int)} "
|
||
f"方差={std.astype(int)} 样本{len(samples)}px")
|
||
return mean, std
|
||
|
||
|
||
def _color_transfer_to_hair(result_img, orig_img, mask):
|
||
"""让生发区域颜色与原图自然融合。
|
||
|
||
放弃强制染色（会导致偏绿/偏蓝），改用：
|
||
1. 保留 webui 生成的原始颜色和纹理（自然发色）
|
||
2. 轻度亮度/对比度匹配，让生发区与周边亮度一致
|
||
3. 轻微降低饱和度（防止异常色相凸显）
|
||
"""
|
||
mb = mask > 127
|
||
if not mb.any():
|
||
return result_img
|
||
|
||
# 1. 亮度匹配：让生发区平均亮度对齐周边
|
||
dilated = cv2.dilate(mb.astype(np.uint8) * 255, np.ones((31, 31), np.uint8))
|
||
surround = (dilated > 0) & (~mb)
|
||
if surround.sum() > 100:
|
||
orig_gray = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY)
|
||
ref_lum = orig_gray[surround].mean()
|
||
else:
|
||
ref_lum = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY).mean()
|
||
|
||
res_gray = cv2.cvtColor(result_img, cv2.COLOR_BGR2GRAY)
|
||
grow_lum = res_gray[mb].mean() if mb.sum() > 0 else ref_lum
|
||
if grow_lum > 5:
|
||
lum_ratio = np.clip(ref_lum / grow_lum, 0.7, 1.3)
|
||
else:
|
||
lum_ratio = 1.0
|
||
|
||
# 2. 轻度饱和度降低 + 去除异常色相（绿/蓝/青）
|
||
# SD1.5 生成头发时常出现偏绿/偏蓝，需要把异常色相拉回头发自然色(棕/灰)
|
||
res_hsv = cv2.cvtColor(result_img, cv2.COLOR_BGR2HSV).astype(np.float32)
|
||
# 2a. 饱和度降低15%
|
||
res_hsv[mb, 1] *= 0.85
|
||
# 2b. 亮度匹配
|
||
res_hsv[mb, 2] *= lum_ratio
|
||
# 2c. 去除绿色/青色/蓝色色相：
|
||
# OpenCV HSV: H=0-180。绿35-85, 青85-105, 蓝105-130
|
||
# 把这些异常色相改成棕色(H≈15-25，头发的自然色)
|
||
h_flat = res_hsv[:, :, 0][mb].copy()
|
||
s_flat = res_hsv[:, :, 1][mb].copy()
|
||
abnormal = (h_flat >= 35) & (h_flat <= 130) # 绿~蓝全范围
|
||
# 异常色相 → 棕色(H=20)，饱和度降到40%（保留一些色调但不刺眼）
|
||
h_flat[abnormal] = 20.0
|
||
s_flat[abnormal] *= 0.4
|
||
res_hsv[:, :, 0][mb] = h_flat
|
||
res_hsv[:, :, 1][mb] = s_flat
|
||
|
||
res_hsv = np.clip(res_hsv, 0, 255).astype(np.uint8)
|
||
result = cv2.cvtColor(res_hsv, cv2.COLOR_HSV2BGR)
|
||
print(f"[hair_grow] 去色完成: 异常色相(绿/青/蓝)像素 {abnormal.mean()*100:.0f}% 已转为棕色")
|
||
return result
|
||
|
||
|
||
def _seamless_blend(result_img, orig_img, mask):
|
||
"""把生发区域贴回原图，消除拼接痕迹。
|
||
|
||
关键原则：mask 内部全部用生发结果（不让皮肤透出形成白边），
|
||
只在 mask 外侧做轻微羽化过渡。
|
||
"""
|
||
mb = mask > 127
|
||
if mb.sum() < 10:
|
||
return result_img
|
||
|
||
# 1. alpha: mask 内 = 1.0（全覆盖，不透出原图皮肤），mask 外 = 0
|
||
alpha = mb.astype(np.float32)
|
||
|
||
# 2. mask 边缘内缩 2px（避免最边缘像素因生成质量差透出），
|
||
# 但不渐变到0——用腐蚀保证内边缘是实心的
|
||
eroded = cv2.erode(mb.astype(np.uint8) * 255, np.ones((5, 5), np.uint8))
|
||
alpha = np.maximum(alpha, eroded.astype(np.float32) / 255)
|
||
|
||
# 3. mask 外侧轻微羽化（让生发边缘自然延伸 3-4px 到皮肤，柔化过渡）
|
||
# 用很小的 GaussianBlur，不会让 mask 内部透出皮肤
|
||
alpha = cv2.GaussianBlur(alpha, (7, 7), sigmaX=1.5)
|
||
alpha = np.clip(alpha, 0, 1.0)
|
||
|
||
# 4. 亮度微调：生发区与周边的亮度适度靠拢（80%生发 + 20%周边）
|
||
# 避免头发过暗形成明显色块，但又保留头发的深色调
|
||
dil_outer = cv2.dilate(mb.astype(np.uint8) * 255, np.ones((31, 31), np.uint8))
|
||
ring_outer = (dil_outer > 0) & (~mb)
|
||
if ring_outer.sum() > 100:
|
||
orig_gray = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY)
|
||
ref_lum = orig_gray[ring_outer].mean()
|
||
res_gray = cv2.cvtColor(result_img, cv2.COLOR_BGR2GRAY)
|
||
grow_lum = res_gray[mb].mean() if mb.sum() > 0 else ref_lum
|
||
if grow_lum > 5:
|
||
# 生发区亮度向周边靠拢20%（不至于让头发太亮，但减少色块感）
|
||
lum_ratio = np.clip(ref_lum / grow_lum, 0.8, 1.2)
|
||
# 只对较暗的生发区提亮一点
|
||
result_adj = result_img.astype(np.float32).copy()
|
||
dark_pixels = res_gray[mb] < ref_lum * 0.85
|
||
mb_idx = np.where(mb)
|
||
# 对暗像素适度提亮
|
||
res_gray_full = res_gray.astype(np.float32)
|
||
lift = np.clip((ref_lum * 0.9 - res_gray_full) / np.maximum(res_gray_full, 1), 0, 0.3)
|
||
result_adj = np.clip(result_adj * (1 + lift[..., None] * 0.5), 0, 255).astype(np.uint8)
|
||
else:
|
||
result_adj = result_img
|
||
else:
|
||
result_adj = result_img
|
||
|
||
# 5. alpha 混合
|
||
alpha3 = alpha[..., None]
|
||
blended = (result_adj.astype(np.float32) * alpha3 +
|
||
orig_img.astype(np.float32) * (1 - alpha3))
|
||
return np.clip(blended, 0, 255).astype(np.uint8)
|
||
|
||
|
||
def hair_grow(img, mask, strength=0.5):
|
||
"""在 mask 区域内生发。
|
||
|
||
参数:
|
||
img: BGR ndarray (H,W,3) 人头像原图
|
||
mask: 灰度 ndarray (H,W)，与 img 同分辨率，白色(255)=生发区
|
||
strength: 生发强度 0.1~1.0，控制重绘强度与边缘处理
|
||
- 低强度(0.1~0.3): 轻微生发，严格保持 mask 边界
|
||
- 中强度(0.4~0.6): 明显生发
|
||
- 高强度(0.7~1.0): 浓密生发，mask 轻微膨胀+高羽化让边缘自然过渡
|
||
返回:
|
||
BGR ndarray 生发结果图（与原图同尺寸）
|
||
"""
|
||
# ---- 参数校验 ----
|
||
if img is None or mask is None:
|
||
raise ValueError("img 和 mask 不能为空")
|
||
if img.ndim != 3 or mask.ndim != 2:
|
||
raise ValueError(f"img 须为3通道BGR(H,W,3)，mask 须为单通道灰度(H,W)；"
|
||
f"got img.ndim={img.ndim}, mask.ndim={mask.ndim}")
|
||
if img.shape[:2] != mask.shape[:2]:
|
||
raise ValueError(f"img 与 mask 分辨率不一致: img={img.shape[:2]}, mask={mask.shape[:2]}")
|
||
|
||
# 二值化 mask（防止用户传入带灰度的 mask）
|
||
_, mask_bin = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
|
||
if cv2.countNonZero(mask_bin) == 0:
|
||
raise ValueError("mask 全黑，没有指定生发区域")
|
||
|
||
# ---- 强度参数映射 ----
|
||
# strength -> denoising_strength (0.7~0.85)
|
||
# 经实测：生发需要 denoise≥0.7 才能生成全新头发纹理（低于此值只会磨皮）
|
||
# denoise 0.85 是甜点（纹理量最高），超过 0.9 反而细节下降
|
||
denoising_strength = 0.7 + float(strength) * 0.15
|
||
# strength -> mask_blur (4~8)，边缘羽化
|
||
mask_blur = int(4 + float(strength) * 4)
|
||
|
||
# ---- mask 边缘处理（强度控制）----
|
||
# 高强度时 mask 轻微 dilate，让生发边缘与原有头发自然衔接
|
||
work_mask = mask_bin
|
||
if strength >= 0.7:
|
||
work_mask = cv2.dilate(work_mask, np.ones((3, 3), np.uint8), iterations=1)
|
||
print(f"[hair_grow] 高强度({strength})，mask 已 dilate 3px")
|
||
|
||
# ---- 生发 prompt ----
|
||
# 经实测：SD1.5 inpainting 对极简短词最敏感，冗长描述反而效果差（纹理量降30%+）
|
||
prompt = "hair, thick hair"
|
||
negative_prompt = ("(bald:1.5), (skin:1.4), smooth skin, scalp, "
|
||
"(green:1.5), (blue:1.5), (cyan:1.4), (green hair:1.5), "
|
||
"(colored hair:1.3), low quality, blurry")
|
||
|
||
# ---- 尺寸对齐到8的倍数(SD要求) ----
|
||
# webui 内部会把图 resize 到 8 的倍数，这里提前对齐，保证结果与原图同尺寸
|
||
orig_h, orig_w = img.shape[:2]
|
||
align_w = orig_w - (orig_w % 8)
|
||
align_h = orig_h - (orig_h % 8)
|
||
if (align_w, align_h) != (orig_w, orig_h):
|
||
img_in = cv2.resize(img, (align_w, align_h), interpolation=cv2.INTER_AREA)
|
||
mask_in = cv2.resize(work_mask, (align_w, align_h), interpolation=cv2.INTER_NEAREST)
|
||
else:
|
||
img_in, mask_in = img, work_mask
|
||
|
||
# ---- 调 webui inpainting ----
|
||
result_img = _webui_inpaint(
|
||
img_in, mask_in, prompt, negative_prompt,
|
||
denoising_strength=denoising_strength,
|
||
mask_blur=mask_blur,
|
||
steps=40, cfg_scale=7.0 # 40步(高于30)减少镂空，细节更连续
|
||
)
|
||
|
||
# ---- 结果 resize 回原图尺寸 ----
|
||
if result_img.shape[:2] != (orig_h, orig_w):
|
||
result_img = cv2.resize(result_img, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
|
||
# mask 也 resize 回原图尺寸，后续后处理要用
|
||
mask_full = cv2.resize(mask_in, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
|
||
|
||
# ---- 后处理（顺序很重要）----
|
||
# 1. 先泊松融合：消除生发区与原图的拼接痕迹（结构层面无缝）
|
||
result_img = _seamless_blend(result_img, img, mask_full)
|
||
print("[hair_grow] 泊松融合完成（消除拼接痕迹）")
|
||
|
||
# 2. 再颜色迁移：泊松融合会把颜色拉向周边皮肤，需重新染回头发色
|
||
# （放最后，确保最终颜色与原图头发一致）
|
||
result_img = _color_transfer_to_hair(result_img, img, mask_full)
|
||
print("[hair_grow] 颜色迁移完成（生发区已对齐原图发色）")
|
||
|
||
return result_img
|