#!/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()