xsl/change_hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

初始化：换发型/换发色/训练发型服务

Browse Source 
包含：
- 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密钥已脱敏为环境变量，原文件备份在本地
This commit is contained in:
xsl
2026-07-07 13:53:52 +08:00
commit 443cfa298f
312 changed files with 67065 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train_hairstyle_full.py
+247
View File
@@ -0,0 +1,247 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""全自动训练新发型
完整流程：准备训练数据 → 训练LoRA → 生成ref材质 → 生成模板材质 → 生成预览图
串接 prepare_train_data / gen_template_material / photo_service训练 / 回调 / 预览图
用法:
cd /home/xsl/change_hair/project/hair_service_sd
python /home/xsl/change_hair/train_hairstyle_full.py --hair-id huaban1 --input /home/xsl/change_hair/huaban1 --gender girl --template-img /home/xsl/change_hair/huaban1/huaban1.jpg
"""
import os
import sys
import ssl
import json
import time
import uuid
import pickle
import base64
import argparse
import requests as req
ssl._create_default_https_context = ssl._create_unverified_context
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1")
import cv2
import numpy as np
import torch
_orig_torch_load = torch.load
def _patched_torch_load(*a, **kw):
if 'map_location' in kw and callable(kw['map_location']) and not isinstance(kw['map_location'], str):
kw['map_location'] = 'cpu'
return _orig_torch_load(*a, **kw)
torch.load = _patched_torch_load
from models.detector import RetinaFaceDetector
from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K
from utils.landmark_processor import get_max_rect, pts_1k_to_137
from core.process_modules import Generator_Matte
from common.logger import config
CAPTION = "titor hairstyle, easyphoto, faceless, no human, white background, simple background, "
RESOLUTIONS = [512, 768, 1024, 1280, 1536]
PHOTO_TRAIN = "http://127.0.0.1:32678/api/hair/train"
HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack"
SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1"
GIRL_IMG = "/home/xsl/change_hair/images/girl.png"
BOY_IMG = "/home/xsl/change_hair/images/boy.png"
PREVIEW_DIR = "/home/xsl/change_hair/hair_grow_service/static/previews"
_detector = _aligner = _matte = None
def get_models():
global _detector, _aligner, _matte
if _detector is None:
print(" 加载模型...")
_detector = RetinaFaceDetector(gpu_id=0)
_aligner = MomocvFaceAlignment1K(gpu_id=0)
_matte = Generator_Matte(gpu=True, device_id=0)
print(" 模型加载完成")
return _detector, _aligner, _matte
def step1_prepare_data(hair_id, input_dir, gender):
"""准备训练数据（白底抠图+标签）"""
print(f"\n{'='*50}\n[步骤1] 准备训练数据: {hair_id}\n{'='*50}")
train_dir = config.get('default', 'train_dir')
out_dir = os.path.join(train_dir, hair_id, "images", "1_hairstyle")
os.makedirs(out_dir, exist_ok=True)
os.makedirs(os.path.join(train_dir, hair_id, "model"), exist_ok=True)
import glob
imgs = sorted(glob.glob(os.path.join(input_dir, "*.jpg")) + glob.glob(os.path.join(input_dir, "*.png")))
det, aligner, matte = get_models()
total = 0
for i, img_path in enumerate(imgs):
img = cv2.imread(img_path)
if img is None: continue
h, w = img.shape[:2]
dets, _ = det.forward(img)
if len(dets) == 0:
print(f" [{i+1}] ✗ 未检测到人脸: {os.path.basename(img_path)}"); continue
det_box = dets[get_max_rect(dets)]
lm = aligner.stable_forward(img.copy(), [det_box])
if lm is None or len(lm) == 0: continue
pt1k = lm[0]
with torch.no_grad():
_, alpha, _ = matte.matte_inference(img, pt1k)
if alpha.shape != (h, w): alpha = cv2.resize(alpha, (w, h))
alpha_f = alpha.astype(np.float32) / 255.0
white = np.ones_like(img, np.float32) * 255
composite = np.clip(img.astype(np.float32) * alpha_f[:,:,None] + white * (1-alpha_f[:,:,None]), 0, 255).astype(np.uint8)
# 居中正方形
pts = pt1k[:312]; cx, cy = int(pts[:,0].mean()), int(pts[:,1].mean())
side = min(h, w); x1=max(0,cx-side//2); y1=max(0,cy-side//2)
x2=min(w,x1+side); y2=min(h,y1+side); asd=min(x2-x1,y2-y1); x2,y2=x1+asd,y1+asd
sq = composite[y1:y2, x1:x2]
if sq.shape[0] < 100: side=min(h,w); sq=composite[:side,:side]
# 数据增广：原图 + 水平翻转 + 轻微缩放(0.9/1.1)，生成多角度样本
variants = [("orig", sq)]
variants.append(("flip", cv2.flip(sq, 1))) # 水平翻转
h2, w2 = sq.shape[:2]
for sc_name, sc_val in [("s09",0.9),("s11",1.1)]:
scaled = cv2.resize(sq, (int(w2*sc_val), int(h2*sc_val)), interpolation=cv2.INTER_AREA if sc_val<1 else cv2.INTER_LANCZOS4)
# resize后仍是正方形
variants.append((sc_name, scaled))
for vname, vimg in variants:
for res in RESOLUTIONS:
if vimg.shape[0] >= res//2:
r = cv2.resize(vimg,(res,res), interpolation=cv2.INTER_AREA if vimg.shape[0]>=res else cv2.INTER_LANCZOS4)
uid=str(uuid.uuid4()); cv2.imwrite(os.path.join(out_dir,f"{uid}.png"),r)
open(os.path.join(out_dir,f"{uid}.txt"),"w").write(CAPTION); total+=1
print(f" [{i+1}] ✓ {os.path.basename(img_path)} → {len(variants)*len(RESOLUTIONS)}张(含增广)")
print(f" 训练数据: {len(imgs)}张原始图 → {total}张训练图")
return total > 0
def step2_train(hair_id):
"""触发LoRA训练"""
print(f"\n{'='*50}\n[步骤2] 训练LoRA: {hair_id}\n{'='*50}")
train_dir = config.get('default', 'train_dir')
payload = {"task_id":f"train_{hair_id}","hair_id":hair_id,
"hair_material_dir":os.path.join(train_dir,hair_id),
"is_tj":"0","device_id":"0","webui_addr":"http://0.0.0.0:32678/"}
r = req.post(PHOTO_TRAIN, json=payload, timeout=30)
d = r.json()
if d.get("state") != 0:
print(f" ✗ 训练启动失败: {d}"); return False
print(f" 训练已启动，等待完成（约30-40分钟）...")
return True
def step3_wait_and_callback(hair_id, template_img):
"""等待训练完成（检测LoRA文件）+ 触发回调生成材质"""
print(f"\n{'='*50}\n[步骤3] 等待训练完成+生成材质: {hair_id}\n{'='*50}")
train_dir = config.get('default', 'train_dir')
lora_path = os.path.join(train_dir, hair_id, "model", "hairstyle_hd_lora.safetensors")
# 先放模板图到upload_train_imgs（回调需要）
upload_dir = config.get('default', 'upload_train_dir')
save_dir = os.path.join(upload_dir, hair_id)
os.makedirs(save_dir, exist_ok=True)
dst = os.path.join(save_dir, f"first##{hair_id}.jpg")
if template_img and os.path.exists(template_img):
import shutil; shutil.copy(template_img, dst)
# 等待LoRA生成
print(" 等待LoRA权重生成...")
waited = 0
while waited < 3000:
if os.path.exists(lora_path) and os.path.getsize(lora_path) > 100000000:
print(f" ✓ LoRA训练完成! ({os.path.getsize(lora_path)//1024//1024}MB, 等待{waited}s)")
break
time.sleep(15); waited += 15
if waited % 60 == 0: print(f" ...已等待{waited}s")
else:
print(" ✗ 训练超时"); return False
time.sleep(5) # 等训练进程写完
# 触发回调生成ref材质
print(" 触发回调生成ref材质...")
r = req.post(HAIR_CALLBACK, json={"task_id":f"train_{hair_id}","hair_id":hair_id,
"state":0,"msg":"头发lora训练成功","is_tj":"1"}, timeout=180)
d = r.json()
ref_dir = os.path.join(config.get('default','hairstyleDir'), hair_id)
ok = d.get("state")==0 and os.path.exists(os.path.join(ref_dir,"ref_rgb_8uc3_768.png"))
print(f" 回调{'成功' if ok else '失败'}: {d.get('msg')}")
return ok
def step4_template_material(hair_id, template_img):
"""生成 hair_template_material（换发型功能4需要）"""
print(f"\n{'='*50}\n[步骤4] 生成模板材质: {hair_id}\n{'='*50}")
template_dir = config.get('default', 'hair_template_material_dir')
out_dir = os.path.join(template_dir, hair_id)
os.makedirs(out_dir, exist_ok=True)
img = cv2.imread(template_img)
if img is None: print(" ✗ 模板图读取失败"); return False
det, aligner, matte = get_models()
dets, _ = det.forward(img)
if len(dets)==0: print(" ✗ 未检测到人脸"); return False
det_box = dets[get_max_rect(dets)]
lm = aligner.stable_forward(img.copy(), [det_box])
pt1k = lm[0]
with torch.no_grad():
_, alpha, _ = matte.matte_inference(img, pt1k)
if alpha.shape != img.shape[:2]: alpha = cv2.resize(alpha,(img.shape[1],img.shape[0]))
uid = f"first##{uuid.uuid4()}"
cv2.imwrite(os.path.join(out_dir,f"{uid}.png"), img)
cv2.imwrite(os.path.join(out_dir,f"{uid}_matting.png"), alpha)
with open(os.path.join(out_dir,f"{uid}.pkl"),"wb") as f:
pickle.dump({"human_pt1k": np.asarray(pt1k, dtype=np.float32)}, f)
print(f" ✓ 模板材质生成完成")
return True
def step5_preview(hair_id, gender):
"""生成预览图"""
print(f"\n{'='*50}\n[步骤5] 生成预览图: {hair_id}\n{'='*50}")
os.makedirs(PREVIEW_DIR, exist_ok=True)
base_img = GIRL_IMG if gender=="girl" else BOY_IMG
with open(base_img,"rb") as f:
img_b64 = "data:image/png;base64," + base64.b64encode(f.read()).decode()
r = req.post(SWAP_API, json={"hair_id":hair_id,"task_id":f"preview_{hair_id}",
"is_hr":"false","user_img_path":img_b64,"output_format":"base64"}, timeout=180)
d = r.json()
if d.get("state")==0 and d.get("data"):
res = cv2.imdecode(np.frombuffer(base64.b64decode(d["data"]),np.uint8), cv2.IMREAD_COLOR)
h,w=res.shape[:2]; sc=min(300/w,400/h)
res=cv2.resize(res,(int(w*sc),int(h*sc)),interpolation=cv2.INTER_AREA)
cv2.imwrite(os.path.join(PREVIEW_DIR,f"{hair_id}.jpg"), res, [cv2.IMWRITE_JPEG_QUALITY,85])
print(f" ✓ 预览图生成完成")
return True
print(f" ✗ 预览图失败: {d.get('msg')}")
return False
def main():
parser = argparse.ArgumentParser(description="全自动训练新发型")
parser.add_argument("--hair-id", required=True)
parser.add_argument("--input", required=True, help="训练图目录")
parser.add_argument("--gender", required=True, choices=["boy","girl"])
parser.add_argument("--template-img", required=True, help="模板图(发型图,用于材质和预览)")
args = parser.parse_args()
t0 = time.time()
print(f"\n{'#'*60}")
print(f"# 全自动训练发型: {args.hair_id} ({args.gender})")
print(f"{'#'*60}")
ok = True
ok = step1_prepare_data(args.hair_id, args.input, args.gender) and ok
if not ok: print("数据准备失败，终止"); sys.exit(1)
ok = step2_train(args.hair_id) and ok
if not ok: print("训练启动失败，终止"); sys.exit(1)
ok = step3_wait_and_callback(args.hair_id, args.template_img) and ok
ok = step4_template_material(args.hair_id, args.template_img) and ok
ok = step5_preview(args.hair_id, args.gender) and ok
print(f"\n{'#'*60}")
print(f"# 训练完成: {args.hair_id} | {'✓全部成功' if ok else '⚠️部分失败'} | 总耗时{time.time()-t0:.0f}s")
print(f"{'#'*60}")
if __name__ == "__main__":
main()