# -*- coding: utf-8 -*- """换发型（带可视化中间产物） 和换发型 /api/swapHair/v1 完全一致的工作流，但保存每一步的中间产物 用于前端可视化展示换发型过程。 mask 处理与换发型一致：origin_matting - 刘海圆 ∪ new_matting （与生发不同：不减刘海、不加手绘mask） """ 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_viz(origin_img, hair_id, hairstyle_process, landmark_processor, task_id, is_hr=True, strict_mask=False): """换发型 + 保存中间产物。 返回: result_img: 最终结果图 BGR steps: [{title, desc, images:[{label, b64}]}] 每步的中间产物 """ 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"swapviz_{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": "用户人像 + 目标发型参考图", "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("发型粗推理失败") print(f"[swap_viz] 粗推理: {time.time()-start:.1f}s") 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": "② 粗推理", "desc": f"把目标发型粗略合成到用户脸上（耗时{time.time()-start:.1f}s）", "images": [ {"label": "粗推理结果", "b64": _enc(img_res)}, ] }) # === 步骤3: mask 合并（换发型：减刘海 ∪ 新发型）=== start = time.time() # 刘海圆（与换发型功能8一致） no_bang_result = origin_matting.copy() try: landmarks_137 = pts_1k_to_137(np.asarray(landmarks_1k)) # 眼睛中点（137点第7、14点） eye_centers = [] for i in [7, 14]: if i < len(landmarks_137): eye_centers.append(landmarks_137[i]) if len(eye_centers) == 2: cx = int((eye_centers[0][0] + eye_centers[1][0]) / 2) cy = int((eye_centers[0][1] + eye_centers[1][1]) / 2) r = int(np.linalg.norm(np.array(eye_centers[0]) - np.array(eye_centers[1]))) cv2.circle(no_bang_result, (cx, cy), r, 0, -1) except Exception as e: print(f"[swap_viz] 刘海圆计算失败({e})，不减刘海") matting_merge = np.max(np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8) steps.append({ "title": "③ 生成重绘mask", "desc": "用户原头发mask(减刘海圆) ∪ 新发型mask = 需要重绘的区域", "images": [ {"label": "用户原头发mask", "b64": _enc_mask(origin_matting)}, {"label": "新发型mask", "b64": _enc_mask(new_matting)}, {"label": "合并mask(重绘区)", "b64": _enc_mask(matting_merge)}, ] }) # === 步骤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) dilate_kernel = (6, 18) if is_hr else (3, 9) mask_dilate = cv2.dilate(mask, np.ones(dilate_kernel, 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", "images": [ {"label": "变换后的图", "b64": _enc(cv2.resize(crop_result, (crop_result.shape[1]//2, crop_result.shape[0]//2)))}, {"label": "变换后的mask", "b64": _enc_mask(cv2.resize(mask_dilate, (mask_dilate.shape[1]//2, mask_dilate.shape[0]//2), interpolation=cv2.INTER_NEAREST))}, ] }) # === 步骤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=0.6, inference_port="57860") print(f"[swap_viz] webui推理: {time.time()-start:.1f}s") steps.append({ "title": "⑤ SD推理(LoRA)", "desc": f"webui img2img inpainting，加载发型LoRA，在mask区重绘新发型（耗时{time.time()-start:.1f}s）", "images": [ {"label": "webui输出", "b64": _enc(cv2.resize(sd_result, (sd_result.shape[1]//2, sd_result.shape[0]//2)))}, ] }) # === 步骤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) # 严格按mask贴回：只覆盖mask_dilate逆变换后的区域 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_dilate逆变换到原图尺寸 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) # 轻微膨胀+高斯模糊让边缘自然 mask_back = cv2.dilate(mask_back, np.ones((5, 5), 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) # mask_back=0 的区域（裁剪框外）恢复原图 result_strict[mask_back == 0] = origin_img[mask_back == 0] # 根据strict_mask选择返回哪个 result = result_strict if strict_mask else result_full # === 步骤7.5: 消除接缝（泊松无缝融合）=== # strict_mask模式下mask边缘有接缝，用seamlessClone消除 # 原理：保持sd_result内部梯度，把边缘梯度过渡到原图，消除色差/亮度突变 if strict_mask: start_enhance = time.time() try: # sd_result逆变换到原图尺寸 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逆变换 + 膨胀（扩大融合区域） 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) mask_back2 = cv2.dilate(mask_back2, np.ones((9, 9), np.uint8)) # 泊松融合需要mask外接矩形中心 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) # NORMAL_CLONE: 在mask区域用sd_back的颜色，边缘梯度自然过渡 result = cv2.seamlessClone(sd_back, origin_img, mask_back2, (cx2, cy2), cv2.NORMAL_CLONE) print(f"[swap_viz] 泊松融合消除接缝: {time.time()-start_enhance:.1f}s") except Exception as e: print(f"[swap_viz] 泊松融合失败({e})，回退到严格mask结果") steps.append({ "title": "⑥ 贴回+增强重绘" + ("（严格按mask+消除接缝）" if strict_mask else "（整框覆盖+增强）"), "desc": "先贴回原图，再用低denoise(0.35)全图增强重绘消除mask接缝（耗时见日志）", "images": [ {"label": "用户原图", "b64": _enc(origin_img)}, {"label": "整框覆盖(增强前)", "b64": _enc(result_full)}, {"label": "严格mask(增强前)", "b64": _enc(result_strict)}, {"label": "最终(增强后)", "b64": _enc(result)}, ] }) # 清理临时文件 try: shutil.rmtree(user_material_dir) os.remove(new_user_img_path) except Exception: pass print(f"[swap_viz] 总耗时: {time.time()-start_all:.1f}s") return result, steps def _enc(img): """ndarray -> base64 jpeg""" _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 80]) return base64.b64encode(buf).decode() def _enc_mask(mask): """单通道mask -> base64 jpeg（伪彩色化便于查看）""" 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()