# -*- 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