xsl
6f34e8876c
硬件迁移：从 RTX 5090 (32GB) 迁移到 RTX 3090 (24GB) 主要改动： - 所有启动脚本和配置文件中的 /home/xsl/ 路径替换为 /home/ubuntu/ - 适配新的 24GB VRAM 环境
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""为测试图自动生成发际线/额头区域的 mask（模拟用户手画遮罩）。
|
|
用 RetinaFace + 1k关键点定位额头，生成该区域的白色 mask。
|
|
"""
|
|
import os
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
|
|
HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
|
|
sys.path.insert(0, HAIR_SERVICE_DIR)
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
|
|
import torch
|
|
from utils.landmark_processor import get_max_rect
|
|
|
|
|
|
def gen_hairline_mask(img_path, mask_path):
|
|
"""用关键点定位额头，生成发际线区域 mask"""
|
|
from models.detector import RetinaFaceDetector
|
|
from utils.MomocvFaceAlignment1K import MomocvFaceAlignment1K
|
|
|
|
img = cv2.imread(img_path)
|
|
if img is None:
|
|
print(f"❌ 读取失败: {img_path}")
|
|
return False
|
|
|
|
print(f"原图: {img_path} shape={img.shape}")
|
|
gpu_id = 0
|
|
|
|
# 1. RetinaFace 检测
|
|
detector = RetinaFaceDetector(gpu_id=gpu_id)
|
|
dets, landms = detector.forward(img)
|
|
if len(dets) == 0:
|
|
print("❌ 未检测到人脸")
|
|
return False
|
|
max_idx = get_max_rect(dets)
|
|
det = dets[max_idx]
|
|
print(f"检测到人脸: {det}")
|
|
|
|
# 2. 1k 关键点
|
|
aligner = MomocvFaceAlignment1K(gpu_id=gpu_id)
|
|
landmarks_1k = aligner.stable_forward(img.copy(), [det])
|
|
if landmarks_1k is None or len(landmarks_1k) == 0:
|
|
print("❌ 关键点检测失败")
|
|
return False
|
|
pts = landmarks_1k[0] # (1000, 2)
|
|
print(f"1k关键点 shape: {pts.shape}")
|
|
|
|
# 3. 构造发际线/额头 mask
|
|
# 1k点的前312点是整脸外轮廓（含发际线）
|
|
# 额头区：用眉骨以上的外轮廓点 + 发际线点构造凸包
|
|
h, w = img.shape[:2]
|
|
mask = np.zeros((h, w), dtype=np.uint8)
|
|
|
|
# 取脸部上半部分轮廓点（发际线到太阳穴）构造额头区域
|
|
# 索引 0-311 是脸部外轮廓，取上半部分（额头/太阳穴）
|
|
forehead_pts = []
|
|
# 发际线区域：取轮廓的上半段 + 稍微往下扩展
|
|
for i in range(0, 312, 3): # 隔点采样，减少密度
|
|
pt = pts[i]
|
|
# 只取上半脸（y < 图像中线偏上）
|
|
if pt[1] < h * 0.55:
|
|
forehead_pts.append(pt)
|
|
# 加入额头中心区域一些点，确保覆盖额头
|
|
forehead_pts = np.array(forehead_pts, dtype=np.int32)
|
|
|
|
if len(forehead_pts) < 3:
|
|
print("❌ 额头点不足")
|
|
return False
|
|
|
|
# 凸包填充
|
|
hull = cv2.convexHull(forehead_pts)
|
|
cv2.fillConvexPoly(mask, hull, 255)
|
|
|
|
# 膨胀让区域稍大一点（模拟用户涂抹）
|
|
mask = cv2.dilate(mask, np.ones((5, 5), np.uint8), iterations=2)
|
|
|
|
# 统计 mask 区域
|
|
nonzero = cv2.countNonZero(mask)
|
|
print(f"mask 生成完成: 白色区域={nonzero}px ({nonzero/(h*w)*100:.1f}%)")
|
|
cv2.imwrite(mask_path, mask)
|
|
print(f"✅ mask 已保存: {mask_path}")
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
img_path = sys.argv[1] if len(sys.argv) > 1 else \
|
|
"/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
|
|
mask_path = sys.argv[2] if len(sys.argv) > 2 else \
|
|
"/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
|
|
gen_hairline_mask(img_path, mask_path)
|