初始化：换发型/换发色/训练发型服务
包含： - 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
@@ -0,0 +1,182 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user