change_hair/prepare_train_data.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""自动准备 LoRA 训练数据

把 train_images/ 里的原始人像图处理成 kohya LoRA 训练格式：
原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签

用法:
    cd /home/ubuntu/change_hair/project/hair_service_sd
    /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/prepare_train_data.py \
        --input /home/ubuntu/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/ubuntu/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()