Files
change_hair/project/hair_service_sd/hair_swap_manual.py
T
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密钥已脱敏为环境变量,原文件备份在本地
2026-07-07 13:53:52 +08:00

426 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# -*- coding: utf-8 -*-
"""换发型(手绘mask重绘版)
与 hair_swap_debug.py 工作流完全一致,唯一区别在步骤③:
- debug版:重绘区域 = 原头发mask ∪ 新发型mask(自动计算)
- 本版: 重绘区域 = 用户手绘mask(白色=重绘,黑色=保留原图)
用户在页面上画哪里,SD 就重绘哪里,完全由用户控制重绘范围。
工作流:输入 → 粗推理 → 手绘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_manual(origin_img, hand_mask, hair_id, hairstyle_process, landmark_processor,
task_id,
# ===== 流程开关(真正生效)=====
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膨胀核
feather_px=0, # ★ strict贴回边缘羽化像素(0=无羽化硬边缘,>0=高斯模糊边缘)
seamless_dilate=(9, 9), # 步骤⑦ 泊松融合mask膨胀核
# ===== enhance 二次增强重绘 =====
enhance=False, # 步骤⑧ 是否对结果再做一次低强度SD重绘(让发丝更清晰)
enhance_denoising=0.35, # enhance 重绘强度(默认0.35,越低越保留原图)
):
"""手绘mask版换发型。
参数:
hand_mask: 用户手绘mask,灰度ndarray,白色(255)=重绘区,与 origin_img 同尺寸
其余参数同 hair_swap_debug
返回:
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
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()
# 手绘mask归一化为 0/255 灰度,保证尺寸与原图一致
if hand_mask.shape[:2] != origin_img.shape[:2]:
hand_mask = cv2.resize(hand_mask, (origin_img.shape[1], origin_img.shape[0]),
interpolation=cv2.INTER_NEAREST)
if hand_mask.ndim == 3:
hand_mask = cv2.cvtColor(hand_mask, cv2.COLOR_BGR2GRAY)
# 二值化(前端传过来已是黑白,这里再保底阈值化)
_, hand_mask_bin = cv2.threshold(hand_mask, 30, 255, cv2.THRESH_BINARY)
matting_merge = hand_mask_bin
# 统计手绘区域占比,便于调试
mask_area_ratio = float((hand_mask_bin > 0).sum()) / hand_mask_bin.size * 100
# 同时展示自动计算的mask作对比(仅展示,不参与重绘)
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)
auto_merge = None
if origin_matting is not None and new_matting is not None:
auto_merge = np.max(np.stack([origin_matting, new_matting], axis=2), axis=2).astype(np.uint8)
step3_images = [
{"label": "用户手绘mask(重绘区)", "b64": _enc_mask(hand_mask_bin)},
]
if origin_matting is not None:
step3_images.append({"label": "自动:原头发mask(仅对比)", "b64": _enc_mask(origin_matting)})
if new_matting is not None:
step3_images.append({"label": "自动:新发型mask(仅对比)", "b64": _enc_mask(new_matting)})
if auto_merge is not None:
step3_images.append({"label": "自动合并(普通换发型会用)", "b64": _enc_mask(auto_merge)})
steps.append({
"title": "③ 手绘mask(核心:用户控制重绘区)",
"desc": f"重绘区域完全由用户手绘决定,白色区域={mask_area_ratio:.1f}%。SD 只在这些区域重绘,其余保留原图。下方对比图展示自动计算的mask(仅供参考,不参与本次重绘)。",
"images": step3_images
})
# === 步骤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")
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_back做高斯模糊,让融合边缘平滑过渡(feather_px=模糊核半径)
feather_info = ""
# 先算一份硬边缘版本(feather=0),用于对比展示
hard_blend = (mask_back.astype(np.float32) / 255)[..., None]
result_strict_hard = (sd_result_back.astype(np.float32) * hard_blend +
origin_img.astype(np.float32) * (1 - hard_blend))
result_strict_hard = np.clip(result_strict_hard, 0, 255).astype(np.uint8)
result_strict_hard[mask_back == 0] = origin_img[mask_back == 0]
if feather_px and feather_px > 0:
k = max(1, int(feather_px)) * 2 + 1 # 高斯核必须是奇数
mask_back_blur = cv2.GaussianBlur(mask_back, (k, k), 0)
mask_blend = (mask_back_blur.astype(np.float32) / 255)[..., None]
feather_info = f",边缘羽化{feather_px}px"
else:
mask_blend = hard_blend
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]
# ★ 羽化对比图:左=硬边缘(feather=0) 右=当前羽化,裁剪mask边界附近放大
feather_compare = None
if feather_px and feather_px > 0:
ys, xs = np.where(mask_back > 10)
if len(ys) > 0:
x1 = max(0, xs.min()-20); x2 = min(origin_img.shape[1], xs.max()+20)
y1 = max(0, ys.min()-20); y2 = min(origin_img.shape[0], ys.max()+20)
hard_crop = result_strict_hard[y1:y2, x1:x2]
blur_crop = result_strict[y1:y2, x1:x2]
sep = np.full((hard_crop.shape[0], 3, 3), 128, dtype=np.uint8)
feather_compare = np.hstack([hard_crop, sep, blur_crop])
result = result_strict if strict_mask else result_full
step6_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)},
]
if feather_compare is not None:
step6_images.append({"label": f"★ 羽化对比(左=硬边 右=羽化{feather_px}px)", "b64": _enc(feather_compare)})
steps.append({
"title": "⑥ 贴回原图",
"desc": f"把SD结果逆warpAffine贴回用户原图。可调:strict_mask={strict_mask}, blend_dilate={bd}{feather_info}。整框覆盖=整个裁剪框覆盖原图;严格mask=只在mask区域融合,mask外保留原图",
"images": step6_images
})
# === 步骤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)}]
})
# === 步骤⑧: enhance 二次增强重绘(可选,让发丝更清晰)===
# ★ enhance 固定基于「严格mask贴回」的图(result_strict),不受 strict_mask 开关影响
if enhance:
start_enh = time.time()
try:
from utils import enhance_hair
# enhance 的输入图固定用严格mask贴回的结果(mask外保留原图)
enh_input = result_strict
# ★ enhance 重绘区域 = 原图头发区域 ∪ 手绘mask(并集)
# origin_matting 是用户原头发maskhand_mask_bin 是用户手绘mask
if origin_matting is not None:
enh_mask = np.max(np.stack([origin_matting, hand_mask_bin], axis=2), axis=2).astype(np.uint8)
else:
enh_mask = hand_mask_bin
# ★ mask 再往外扩大 10%(相对mask外接矩形短边,避免边缘漏掉)
ys, xs = np.where(enh_mask > 10)
if len(ys) > 0:
short_side = min(xs.max() - xs.min(), ys.max() - ys.min())
dilate_px = max(1, int(short_side * 0.10)) # 短边的10%
enh_mask = cv2.dilate(enh_mask, np.ones((dilate_px, dilate_px), np.uint8))
# ★ 用瞳距换算 3cm = 多少像素(瞳距≈6.3cm)
# 1k关键点 [691]=左眼中心, [792]=右眼中心
cm3_px = 0 # 3cm对应的像素数
try:
lm = np.asarray(landmarks_1k)
left_eye = lm[691]; right_eye = lm[792]
pupil_dist_px = np.linalg.norm(np.array(left_eye) - np.array(right_eye))
cm3_px = max(1, int(pupil_dist_px / 6.3 * 3)) # 瞳距/6.3cm × 3cm
except Exception as e:
print(f"[enhance] 瞳距换算失败({e})3cm按40px")
# ★ 减掉发际线带:从头发下边界向上 3cm 的带状区域
# 头发下边界 = enh_mask 每列最下面的头发像素;向上3cm内的区域擦除
enh_mask_before = enh_mask.copy() # 保存减发际线前的mask(可视化用)
if cm3_px > 0:
# 找每列头发mask的最下像素行
cols = np.where(enh_mask.max(axis=0) > 10)[0]
for c in cols:
col_ys = np.where(enh_mask[:, c] > 10)[0]
if len(col_ys) == 0:
continue
bottom_y = col_ys.max() # 该列头发最下像素
# 擦除 [bottom_y - cm3_px, bottom_y] 这一段(发际线带)
top_erase = max(0, bottom_y - cm3_px)
enh_mask[top_erase:bottom_y+1, c] = 0
prompt = "high quality, detailed hair strands, realistic hair texture, sharp focus"
enhanced = enhance_hair.webui_img2img(enh_input, enh_mask, prompt=prompt,
denoising_strength=enhance_denoising)
t_enh = time.time() - start_enh
# enhance 重绘区域可视化
try:
from hairline_mask import make_overlay
enh_overlay = make_overlay(enh_input, enh_mask, alpha=0.45)
enh_overlay_before = make_overlay(enh_input, enh_mask_before, alpha=0.45)
except Exception:
enh_overlay = enh_input; enh_overlay_before = enh_input
steps.append({
"title": "⑧ enhance 二次增强重绘",
"desc": f"基于「严格mask贴回」图做低强度SD重绘(denoising={enhance_denoising})。重绘区=原图头发∪手绘mask外扩10%({dilate_px}px),再从发际线向上减掉3cm({cm3_px}px)。耗时 {t_enh:.1f}s",
"images": [
{"label": "严格mask贴回(enhance输入)", "b64": _enc(enh_input)},
{"label": "① 重绘区(外扩后,未减发际线)", "b64": _enc_mask(enh_mask_before)},
{"label": "② 减发际线带后(最终mask)", "b64": _enc_mask(enh_mask)},
{"label": "③ 减之前叠加图", "b64": _enc(enh_overlay_before)},
{"label": "④ 最终mask叠加图", "b64": _enc(enh_overlay)},
{"label": "⑤ 增强后(最终)", "b64": _enc(enhanced)},
]
})
result = enhanced
except Exception as e:
steps.append({
"title": "⑧ enhance(失败)",
"desc": f"enhance增强失败({e}),使用未增强结果",
"images": [{"label": "最终结果", "b64": _enc(result)}]
})
# 清理临时文件
try:
shutil.rmtree(user_material_dir)
os.remove(new_user_img_path)
except Exception:
pass
params_used = {
"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), "feather_px": feather_px,
"seamless_dilate": list(seamless_dilate),
"enhance": enhance, "enhance_denoising": enhance_denoising,
"mask_area_ratio_pct": round(mask_area_ratio, 1),
"gender": in_gender, "tag": p_tag, "dst_size": list(dst_size),
"total_time": round(time.time() - start_all, 1),
}
print(f"[swap_manual] 总耗时: {params_used['total_time']}s")
return result, steps, params_used
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()