#!/usr/bin/env python3 # -*- coding: utf-8 -*- """自动准备 LoRA 训练数据 把 train_images/ 里的原始人像图处理成 kohya LoRA 训练格式： 原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签 用法: cd /home/xsl/change_hair/project/hair_service_sd /home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/prepare_train_data.py \ --input /home/xsl/change_hair/train_images \ --hair-id new_hairstyle_001 \ --gender boy """ import os import sys import ssl import uuid import argparse # 禁用 SSL 验证（解决 resnet18 权重缓存的 SSL 证书问题） ssl._create_default_https_context = ssl._create_unverified_context # 必须在 hair_service_sd 目录运行（依赖其模块） 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 # Monkey-patch torch.load: 部分第三方代码用旧式 map_location lambda 触发 legacy_load, # 但权重文件是新版 zip 格式，导致 UnpicklingError。统一用 map_location='cpu'。 _orig_torch_load = torch.load def _patched_torch_load(*args, **kwargs): if 'map_location' in kwargs and callable(kwargs['map_location']) and not isinstance(kwargs['map_location'], str): kwargs['map_location'] = 'cpu' return _orig_torch_load(*args, **kwargs) torch.load = _patched_torch_load from models.detector import RetinaFaceDetector from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K from utils.landmark_processor import get_max_rect # 用 core/process_modules 的 Generator_Matte（从 weights/ 加载真实权重， # 而非 hair_matting/ 目录的 LFS 指针文件） from core.process_modules import Generator_Matte # kohya 标准训练标签（白底版，和已有发型完全一致） CAPTION = "titor hairstyle, easyphoto, faceless, no human, white background, simple background, " # 训练图分辨率列表（和 prepare_single_hairstyle_v2 一致） RESOLUTIONS = [512, 768, 1024, 1280, 1536] def process_one(img_path, detector, aligner, matte_gen, out_dir, idx): """处理单张原始图：检测→抠图→白底合成→多分辨率保存""" img = cv2.imread(img_path) if img is None: print(f" [{idx}] ✗ 读取失败: {img_path}") return 0 h, w = img.shape[:2] # 1. 人脸检测 dets, landms = detector.forward(img) if len(dets) == 0: print(f" [{idx}] ✗ 未检测到人脸: {os.path.basename(img_path)}") return 0 max_idx = get_max_rect(dets) det = dets[max_idx] # 2. 1k 关键点 landmarks_1k = aligner.stable_forward(img.copy(), [det]) if landmarks_1k is None or len(landmarks_1k) == 0: print(f" [{idx}] ✗ 关键点检测失败: {os.path.basename(img_path)}") return 0 pt1k = landmarks_1k[0] # 3. 头发抠图（matte_inference 返回 pred_fg, pred, image_resize） with torch.no_grad(): pred_fg, pred, image_resize = matte_gen.matte_inference(img, pt1k) # pred 是 alpha mask，已 resize 回原图尺寸 alpha = pred h, w = img.shape[:2] if alpha.shape != (h, w): alpha = cv2.resize(alpha, (w, h)) # 4. 白底合成（公式同 hairstyle_model.py:3666-3669） alpha_f = alpha.astype(np.float32) / 255.0 white_bg = np.ones_like(img, dtype=np.float32) * 255.0 img_f = img.astype(np.float32) composite = (img_f * alpha_f[:, :, None] + white_bg * (1 - alpha_f[:, :, None])) composite = np.clip(composite, 0, 255).astype(np.uint8) # 5. 居中裁剪到正方形（以人脸为中心） pts = pt1k[:312] # 前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) # 调整确保正方形 actual_side = min(x2 - x1, y2 - y1) x2, y2 = x1 + actual_side, y1 + actual_side composite_square = composite[y1:y2, x1:x2] if composite_square.shape[0] < 100: # fallback: 直接用整图短边 side = min(h, w) composite_square = composite[:side, :side] # 6. 多分辨率保存 + 同名标签 count = 0 for res in RESOLUTIONS: if composite_square.shape[0] >= res // 2: # 只生成 >= res/2 的（避免过度放大） if composite_square.shape[0] >= res: resized = cv2.resize(composite_square, (res, res), interpolation=cv2.INTER_AREA) else: resized = cv2.resize(composite_square, (res, res), interpolation=cv2.INTER_LANCZOS4) file_uuid = str(uuid.uuid4()) png_path = os.path.join(out_dir, f"{file_uuid}.png") txt_path = os.path.join(out_dir, f"{file_uuid}.txt") cv2.imwrite(png_path, resized) with open(txt_path, "w") as f: f.write(CAPTION) count += 1 print(f" [{idx}] ✓ {os.path.basename(img_path)} → {count}张训练图") return count def main(): parser = argparse.ArgumentParser(description="准备 LoRA 训练数据") parser.add_argument("--input", required=True, help="原始图目录") parser.add_argument("--hair-id", required=True, help="新发型ID") parser.add_argument("--gender", default="boy", choices=["boy", "girl"], help="性别") args = parser.parse_args() from common.logger import config train_dir = config.get('default', 'train_dir') out_base = os.path.join(train_dir, args.hair_id, "images", "1_hairstyle") os.makedirs(out_base, exist_ok=True) print(f"=== 准备训练数据 ===") print(f"输入: {args.input}") print(f"发型ID: {args.hair_id}, 性别: {args.gender}") print(f"输出: {out_base}") # 加载模型 print("加载模型（RetinaFace + 1k关键点 + Generator_Matte）...") gpu_id = 0 detector = RetinaFaceDetector(gpu_id=gpu_id) aligner = MomocvFaceAlignment1K(gpu_id=gpu_id) matte_gen = Generator_Matte(gpu=True, device_id=gpu_id) print("模型加载完成") # 处理每张图 import glob imgs = sorted(glob.glob(os.path.join(args.input, "*.jpg")) + glob.glob(os.path.join(args.input, "*.png")) + glob.glob(os.path.join(args.input, "*.jpeg"))) print(f"共 {len(imgs)} 张原始图") total = 0 for i, img_path in enumerate(imgs): total += process_one(img_path, detector, aligner, matte_gen, out_base, i + 1) print(f"\n=== 完成: {len(imgs)}张原始图 → {total}张训练图 ===") print(f"训练数据目录: {out_base}") print(f"\n下一步训练命令:") print(f"curl -X POST http://127.0.0.1:32678/api/hair/train \\") print(f' -H "Content-Type: application/json" \\') print(f' -d \'{{"task_id":"train_{args.hair_id}","hair_id":"{args.hair_id}",' f'"hair_material_dir":"{os.path.join(train_dir, args.hair_id)}",' f'"is_tj":"1","device_id":"0","webui_addr":"http://0.0.0.0:32678/"}}\'') if __name__ == "__main__": main()