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]}")