# -*- coding: utf-8 -*- """换发型（发际线带重绘实验版） 与 hair_swap_debug.py 工作流完全一致，唯一区别在步骤③： - debug版：重绘区域 = 原头发mask ∪ 新发型mask（整个头发区域） - 本版： 重绘区域 = 原mask边界带 ∪ 新mask边界带（只重绘发际线交界带，主体保留原图） 工作流：输入 → 粗推理 → 发际线带mask → warpAffine → SD推理 → 贴回 → 消除接缝 """ import os import json import glob import time import shutil import base64 import cv2 import numpy as np from gen_super_image import webui_img2img from common.logger import config from utils.landmark_processor import high_quality_warpAffine from utils.landmark_processor import pts_1k_to_137 def hair_swap_hairline(origin_img, hair_id, hairstyle_process, landmark_processor, task_id, # ===== 流程开关（真正生效）===== method="mediapipe", # 发际线mask方案: boundary_band/mediapipe/landmark_1k/deeplab strict_mask=False, # 步骤⑥ 严格按mask贴回 seamless_blend=True, # 步骤⑦ 泊松融合消除接缝 # ===== 尺寸/对齐（真正生效）===== is_hr=True, # 高清(1152x1536) vs 标清(576x768) dilate_kernel=(6, 18), # 步骤④ mask膨胀核 (x,y) # ===== SD 推理（仅 denoising 可调，其余webui端固定）===== denoising_strength=0.6, # 重绘强度（唯一能透传到webui的SD参数） # ===== 贴回/融合（真正生效）===== blend_dilate=(5, 5), # 步骤⑥ strict mask贴回时mask膨胀核 seamless_dilate=(9, 9), # 步骤⑦ 泊松融合mask膨胀核 # ===== 发际线带专属参数 ===== band_width=15, # 边界带宽(像素)，形态学梯度核大小 preview_only=False, # ★ 只验证mask（跳过SD推理，几秒出结果，默认False） # ===== landmark_1k 方案形状参数 ===== height_ratio=0.432, # 高度缩放比（越小越窄） width_ratio=0.144, # 左右各扩展比（相对眉宽，越大越宽） corner_ratio=0.25, # 圆角半径比（相对短边） vertical_offset=0.0, # 上下平移比（相对额头高度，正值下移） # ===== webui SD 参数 ===== refiner_switch_at=0.5, # refiner切换点（0~1，仅高清模式生效，0=不用refiner） ): """换发型 + 全参数可视化。 返回: result_img: 最终结果图 BGR steps: [{title, desc, images:[{label, b64}]}] params_used: 实际使用的参数（回显） """ start_all = time.time() steps = [] hairstyle_dir = config.get('default', 'hairstyleDir') user_dir = config.get('default', 'userDir') train_dir = config.get('default', 'train_dir') userInfo_dir = config.get('default', 'userInfo_dir') hair_material_dir = os.path.join(train_dir, hair_id) material_save_path = os.path.join(hairstyle_dir, hair_id) ref_img_path = os.path.join(material_save_path, "ref_rgb_8uc3_768.png") if not os.path.exists(ref_img_path): raise FileNotFoundError(f"发型材质不存在: {ref_img_path}") user_img_name = f"swapdbg_{task_id}.jpg" new_user_img_path = os.path.join(user_dir, user_img_name) if os.path.exists(new_user_img_path): os.remove(new_user_img_path) cv2.imwrite(new_user_img_path, origin_img) origin_img = cv2.imread(new_user_img_path) ref_img = cv2.imread(ref_img_path) # === 步骤1: 输入 === steps.append({ "title": "① 输入", "desc": "用户人像 + 目标发型参考图。参考图来自训练材质 ref_rgb_8uc3_768.png", "images": [ {"label": "用户原图", "b64": _enc(origin_img)}, {"label": f"发型参考图({hair_id})", "b64": _enc(cv2.resize(ref_img, (origin_img.shape[1], origin_img.shape[0])))}, ] }) # === 步骤2: 粗推理 infer_hairstyle_diy_jy === start = time.time() work_dir = os.path.join(userInfo_dir, task_id) os.makedirs(work_dir, exist_ok=True) import torch with torch.no_grad(): img_res, status, _, landmarks_1k, isEyeOccluded = hairstyle_process.infer_hairstyle_diy_jy( origin_img, ref_img, os.path.join(work_dir, task_id), f"{task_id}.png") if status != 0: raise RuntimeError("发型粗推理失败") t_coarse = time.time() - start user_material_dir = os.path.join(work_dir, task_id) hair_matting_path = os.path.join(user_material_dir, "hair_mask_2.png") user_orig_mask_path = os.path.join(user_material_dir, "user_orig_mask.png") origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) steps.append({ "title": "② 粗推理 (infer_hairstyle_diy_jy)", "desc": f"用 SPADE 风格迁移模型把目标发型粗略合成到用户脸上。这步只产生粗略效果，细节由后续 SD 推理完善。耗时 {t_coarse:.1f}s", "images": [ {"label": "粗推理结果", "b64": _enc(img_res)}, ] }) # === 步骤3: 发际线带 mask（核心：用所选方案自动识别发际线区域）=== start = time.time() # 调用 hairline_mask 统一入口 import hairline_mask as hm matting_merge, method_info, debug_imgs = hm.detect_hairline( origin_matting, landmarks_1k, img_res, band_width=band_width, method=method, height_ratio=height_ratio, width_ratio=width_ratio, corner_ratio=corner_ratio, vertical_offset=vertical_offset) if matting_merge is None: # 方案失败：回退到边界带法 matting_merge, method_info, debug_imgs = hm.detect_hairline( origin_matting, landmarks_1k, img_res, band_width=band_width, method="boundary_band") method_info = f"[回退到边界带] " + method_info # 收集步骤③的展示图 step3_images = [] step3_images.append({"label": "原头发mask(参考)", "b64": _enc_mask(origin_matting)}) for label, arr in debug_imgs.items(): if arr.ndim == 2: step3_images.append({"label": label, "b64": _enc_mask(arr)}) else: step3_images.append({"label": label, "b64": _enc(arr)}) step3_images.append({"label": "发际线重绘区", "b64": _enc_mask(matting_merge)}) # ★ mask半透明叠加到粗推理图（目标发型图） overlay_res = hm.make_overlay(img_res, matting_merge, alpha=0.45) step3_images.append({"label": "★ 叠加目标发型图", "b64": _enc(overlay_res)}) # ★ mask半透明叠加到用户原图（直观看重绘区域在用户原图上的位置） overlay_orig = hm.make_overlay(origin_img, matting_merge, alpha=0.45) step3_images.append({"label": "★ 叠加用户原图", "b64": _enc(overlay_orig)}) steps.append({ "title": f"③ 发际线mask [方案: {method}]", "desc": f"{method_info}。其余头发主体保留原图。可调：band_width={band_width}px。看「★ 叠加用户原图」直观判断重绘区位置。", "images": step3_images }) # ★ preview_only 模式：只验证 mask，跳过后续 SD 推理（几秒出结果） if preview_only: params_used = { "method": method, "band_width": band_width, "preview_only": True, "height_ratio": height_ratio, "width_ratio": width_ratio, "corner_ratio": corner_ratio, "vertical_offset": vertical_offset, "is_hr": is_hr, "total_time": round(time.time() - start_all, 1), } print(f"[swap_hairline] preview_only 完成，耗时 {params_used['total_time']}s") return overlay_orig, steps, params_used # === 步骤4: warpAffine 变换 === dst_size = (1152, 1536) if is_hr else (576, 768) box_info = hairstyle_process.get_body_info(img_res) box_w = box_info[2] - box_info[0] box_h = box_info[3] - box_info[1] scale = min(dst_size[1] / max(box_h, 1), dst_size[0] / max(box_w, 1)) rotate_center = [(box_info[2] + box_info[0]) * 0.5, (box_info[3] + box_info[1]) * 0.5] M = cv2.getRotationMatrix2D(rotate_center, 0, scale) M[:, 2] += np.float32([dst_size[0] * 0.5, dst_size[1] * 0.5]) - np.float32(rotate_center) crop_result = high_quality_warpAffine(img_res, M, dst_size) crop_matting = cv2.warpAffine(matting_merge, M, dst_size) mask = (crop_matting > 10).astype(np.float32) dk = tuple(max(1, int(x)) for x in dilate_kernel) mask_dilate = cv2.dilate(mask, np.ones(dk, np.uint8)) mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8) final_img = crop_result steps.append({ "title": "④ warpAffine 对齐裁剪", "desc": f"检测身体框，计算缩放/平移矩阵，把图和【发际线带mask】对齐裁剪到 {dst_size[0]}x{dst_size[1]} 送入 webui。可调：is_hr={is_hr}(决定尺寸), dilate_kernel={dk}", "images": [ {"label": f"裁剪后图({dst_size[0]}x{dst_size[1]})", "b64": _enc(_shrink(crop_result))}, {"label": f"发际线带mask膨胀后(dilate={dk})", "b64": _enc_mask(_shrink(mask_dilate, nearest=True))}, ] }) # === 步骤5: 读取 config/prompt === config_json_path = os.path.join(material_save_path, "config.json") with open(config_json_path, "r") as f: in_gender = json.load(f)["gender"] images_dir = os.path.join(hair_material_dir, "images") txt_dir = os.path.join(images_dir, os.listdir(images_dir)[0]) txt_path = glob.glob(txt_dir + '/*.txt')[0] with open(txt_path, 'r') as f: p_tag = f.readline() if "titor hairstyle, faceless, no human, gray background, simple background" in p_tag: p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):] else: p_tag = "" # === 步骤6: webui 推理（走 photo_service + LoRA）=== start = time.time() sd_result = webui_img2img( img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id, hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr, denoising_strength=denoising_strength, inference_port="57860", refiner_switch_at=refiner_switch_at) t_sd = time.time() - start steps.append({ "title": "⑤ SD 推理 (webui img2img inpainting + LoRA)", "desc": f"加载发型LoRA，在mask区域用SD重绘新发型。这是生成发丝细节的关键步骤。耗时 {t_sd:.1f}s。可调：denoising={denoising_strength}（其余 cfg=7/steps=20/sampler=DPM++ 2M Karras/mask_blur=11/seed=123456789 在webui端固定）", "images": [ {"label": "webui输出(SD重绘结果)", "b64": _enc(_shrink(sd_result))}, ] }) # === 步骤7: warpAffine 贴回原图 === M_inv = cv2.invertAffineTransform(M) result_full = origin_img.copy() cv2.warpAffine(sd_result, M_inv, (origin_img.shape[1], origin_img.shape[0]), dst=result_full, borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4) result_strict = origin_img.copy() sd_result_back = np.zeros_like(origin_img) cv2.warpAffine(sd_result, M_inv, (origin_img.shape[1], origin_img.shape[0]), dst=sd_result_back, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0), flags=cv2.INTER_LANCZOS4) mask_back = np.zeros(origin_img.shape[:2], dtype=np.uint8) cv2.warpAffine(mask_dilate, M_inv, (origin_img.shape[1], origin_img.shape[0]), dst=mask_back, borderMode=cv2.BORDER_CONSTANT, borderValue=0, flags=cv2.INTER_NEAREST) bd = tuple(max(1, int(x)) for x in blend_dilate) mask_back = cv2.dilate(mask_back, np.ones(bd, np.uint8)) mask_blend = (mask_back.astype(np.float32) / 255)[..., None] result_strict = (sd_result_back.astype(np.float32) * mask_blend + origin_img.astype(np.float32) * (1 - mask_blend)) result_strict = np.clip(result_strict, 0, 255).astype(np.uint8) result_strict[mask_back == 0] = origin_img[mask_back == 0] result = result_strict if strict_mask else result_full steps.append({ "title": "⑥ 贴回原图", "desc": f"把SD结果逆warpAffine贴回用户原图。可调：strict_mask={strict_mask}, blend_dilate={bd}。整框覆盖=整个裁剪框覆盖原图；严格mask=只在mask区域融合，mask外保留原图", "images": [ {"label": "用户原图", "b64": _enc(origin_img)}, {"label": "整框覆盖", "b64": _enc(result_full)}, {"label": "严格mask贴回", "b64": _enc(result_strict)}, {"label": f"当前选用({'严格mask' if strict_mask else '整框覆盖'})", "b64": _enc(result)}, ] }) # === 步骤8: 消除接缝（泊松无缝融合）=== if strict_mask and seamless_blend: start_enhance = time.time() try: sd_back = np.zeros_like(origin_img) cv2.warpAffine(sd_result, M_inv, (origin_img.shape[1], origin_img.shape[0]), dst=sd_back, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0), flags=cv2.INTER_LANCZOS4) mask_back2 = np.zeros(origin_img.shape[:2], dtype=np.uint8) cv2.warpAffine(mask_dilate, M_inv, (origin_img.shape[1], origin_img.shape[0]), dst=mask_back2, borderMode=cv2.BORDER_CONSTANT, borderValue=0, flags=cv2.INTER_NEAREST) sd2 = tuple(max(1, int(x)) for x in seamless_dilate) mask_back2 = cv2.dilate(mask_back2, np.ones(sd2, np.uint8)) ys2, xs2 = np.where(mask_back2 > 10) if len(ys2) > 10: cx2 = int((xs2.min() + xs2.max()) / 2) cy2 = int((ys2.min() + ys2.max()) / 2) result = cv2.seamlessClone(sd_back, origin_img, mask_back2, (cx2, cy2), cv2.NORMAL_CLONE) t_seamless = time.time() - start_enhance steps.append({ "title": "⑦ 泊松融合消除接缝 (seamlessClone)", "desc": f"严格mask模式下边缘有接缝/色差。泊松融合保持SD结果内部梯度，把边缘梯度过渡到原图，消除突变。耗时 {t_seamless:.1f}s。可调：seamless_dilate={sd2}", "images": [ {"label": "融合前(有接缝)", "b64": _enc(result_strict)}, {"label": "融合后(最终)", "b64": _enc(result)}, ] }) except Exception as e: steps.append({ "title": "⑦ 泊松融合（失败）", "desc": f"泊松融合失败({e})，回退到严格mask结果", "images": [{"label": "最终结果", "b64": _enc(result)}] }) # 清理临时文件 try: shutil.rmtree(user_material_dir) os.remove(new_user_img_path) except Exception: pass params_used = { "method": method, "strict_mask": strict_mask, "seamless_blend": seamless_blend, "is_hr": is_hr, "dilate_kernel": list(dilate_kernel), "denoising_strength": denoising_strength, "blend_dilate": list(blend_dilate), "seamless_dilate": list(seamless_dilate), "band_width": band_width, "height_ratio": height_ratio, "width_ratio": width_ratio, "corner_ratio": corner_ratio, "vertical_offset": vertical_offset, "refiner_switch_at": refiner_switch_at, "gender": in_gender, "tag": p_tag, "dst_size": list(dst_size), "total_time": round(time.time() - start_all, 1), } print(f"[swap_hairline] 总耗时: {params_used['total_time']}s") return result, steps, params_used def _boundary_band(mask, kernel): """提取 mask 的边界带（形态学梯度：dilate - erode）。 输入软mask(0-255)，输出边界带(0或255)，带宽≈2*kernel_size。 """ binary = (mask > 10).astype(np.uint8) * 255 dilated = cv2.dilate(binary, kernel) eroded = cv2.erode(binary, kernel) band = cv2.subtract(dilated, eroded) return band def _shrink(img, max_side=1024, nearest=False): """缩小图片便于前端展示""" h, w = img.shape[:2] if max(h, w) <= max_side: return img sc = max_side / max(h, w) interp = cv2.INTER_NEAREST if nearest else cv2.INTER_AREA return cv2.resize(img, (int(w * sc), int(h * sc)), interpolation=interp) def _enc(img): _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 80]) return base64.b64encode(buf).decode() def _enc_mask(mask): if mask.ndim == 2: colored = cv2.applyColorMap(mask, cv2.COLORMAP_JET) else: colored = mask _, buf = cv2.imencode(".jpg", colored, [cv2.IMWRITE_JPEG_QUALITY, 80]) return base64.b64encode(buf).decode()