xsl
9b7324cd26
将所有绝对路径从 /home/ubuntu 改回 /home/xsl，覆盖启动脚本、 configure.ini 及训练/测试脚本。纯路径替换，无功能性改动。 本分支(xsl5090)用于本机(RTX 5090)运行；master 维持 /home/ubuntu 供远程 ubuntu 机器使用。
248 lines
11 KiB
Python
248 lines
11 KiB
Python
#!/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()
|