48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
from gallery_builder.align import align_face_5pts
|
|
from gallery_builder.detector import OnnxFaceDetector, load_det_outputs_config
|
|
|
|
# ============ 只改这里 ============
|
|
INPUT_IMAGE = "003.jpg"
|
|
# =================================
|
|
|
|
DET_MODEL = "./models/RetinaFace_mobile320.onnx"
|
|
DET_CONFIG = "./models/retinaface_mobile320_config.json"
|
|
RECOG_MODEL = "./models/mobilefacenet_arcface_prenorm.onnx"
|
|
OUT_DIR = "./debug_output"
|
|
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
base_name = os.path.splitext(os.path.basename(INPUT_IMAGE))[0]
|
|
|
|
det_cfg = load_det_outputs_config(DET_CONFIG)
|
|
detector = OnnxFaceDetector(DET_MODEL, det_cfg, score_thresh=0.5, pick_face="largest")
|
|
recog_session = ort.InferenceSession(RECOG_MODEL, providers=["CPUExecutionProvider"])
|
|
|
|
img_bgr = cv2.imread(INPUT_IMAGE)
|
|
det = detector.detect_one(img_bgr)
|
|
print(f"landmarks5:\n{det.landmarks5}")
|
|
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
aligned = align_face_5pts(img_rgb, det.landmarks5)
|
|
|
|
# 保存对齐图
|
|
cv2.imwrite(f"{OUT_DIR}/{base_name}_aligned.png", cv2.cvtColor(aligned, cv2.COLOR_RGB2BGR))
|
|
np.save(f"{OUT_DIR}/{base_name}_aligned_rgb.npy", aligned)
|
|
|
|
# 提取 embedding
|
|
x = (aligned.astype(np.float32) - 127.5) / 128.0
|
|
x = x.transpose(2, 0, 1)[np.newaxis, ...]
|
|
out = recog_session.run(None, {recog_session.get_inputs()[0].name: x})[0]
|
|
emb = out.flatten()
|
|
emb = emb / np.linalg.norm(emb)
|
|
|
|
np.save(f"{OUT_DIR}/{base_name}_emb.npy", emb)
|
|
|
|
print(f"\n保存到 {OUT_DIR}/")
|
|
print(f" {base_name}_aligned.png")
|
|
print(f" {base_name}_aligned_rgb.npy")
|
|
print(f" {base_name}_emb.npy")
|
|
print(f"\nembedding[0:8]: {emb[:8]}") |