100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""为新发型生成 hair_template_material（first##_matting.png + .pkl）
|
||
|
||
换发型功能4 需要：
|
||
hair_template_material/<hair_id>/first##<name>.png 原图
|
||
hair_template_material/<hair_id>/first##<name>_matting.png 头发抠图mask
|
||
hair_template_material/<hair_id>/first##<name>.pkl 1k关键点
|
||
|
||
用法:
|
||
cd /home/ubuntu/change_hair/project/hair_service_sd
|
||
python gen_template_material.py --hair-id new_test_001 --img /path/to/template.jpg
|
||
"""
|
||
import os
|
||
import sys
|
||
import ssl
|
||
import uuid
|
||
import pickle
|
||
import argparse
|
||
|
||
ssl._create_default_https_context = ssl._create_unverified_context
|
||
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
|
||
|
||
_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
|
||
from core.process_modules import Generator_Matte
|
||
from common.logger import config
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--hair-id", required=True)
|
||
parser.add_argument("--img", required=True, help="模板图(发型图)")
|
||
args = parser.parse_args()
|
||
|
||
template_dir = config.get('default', 'hair_template_material_dir')
|
||
out_dir = os.path.join(template_dir, args.hair_id)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
img = cv2.imread(args.img)
|
||
if img is None:
|
||
print(f"✗ 读取失败: {args.img}"); sys.exit(1)
|
||
print(f"模板图: {img.shape}")
|
||
|
||
# 1. 人脸检测 + 1k关键点
|
||
print("加载模型...")
|
||
detector = RetinaFaceDetector(gpu_id=0)
|
||
aligner = MomocvFaceAlignment1K(gpu_id=0)
|
||
matte_gen = Generator_Matte(gpu=True, device_id=0)
|
||
print("模型加载完成")
|
||
|
||
dets, landms = detector.forward(img)
|
||
if len(dets) == 0:
|
||
print("✗ 未检测到人脸"); sys.exit(1)
|
||
det = dets[get_max_rect(dets)]
|
||
landmarks_1k = aligner.stable_forward(img.copy(), [det])
|
||
pt1k = landmarks_1k[0]
|
||
|
||
# 2. 头发抠图
|
||
with torch.no_grad():
|
||
_, alpha, _ = matte_gen.matte_inference(img, pt1k)
|
||
if alpha.shape != img.shape[:2]:
|
||
alpha = cv2.resize(alpha, (img.shape[1], img.shape[0]))
|
||
|
||
# 3. 保存三个文件（first##前缀）
|
||
file_uuid = f"first##{uuid.uuid4()}"
|
||
png_path = os.path.join(out_dir, f"{file_uuid}.png")
|
||
matting_path = os.path.join(out_dir, f"{file_uuid}_matting.png")
|
||
pkl_path = os.path.join(out_dir, f"{file_uuid}.pkl")
|
||
|
||
cv2.imwrite(png_path, img)
|
||
cv2.imwrite(matting_path, alpha)
|
||
with open(pkl_path, 'wb') as f:
|
||
# 存 numpy array（不能用 tolist()，代码用 landmarks[index, :] 索引需要 ndarray）
|
||
pickle.dump({'human_pt1k': np.asarray(pt1k, dtype=np.float32)}, f)
|
||
|
||
print(f"✓ 生成完成:")
|
||
print(f" {png_path}")
|
||
print(f" {matting_path}")
|
||
print(f" {pkl_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|