#!/usr/bin/env python3
import cv2
import numpy as np
import onnxruntime as ort
from pathlib import Path

MODEL_PATH = "model/RetinaFace_mobile320.onnx"  # INT8 模型
IMG_PATH   = "img/friends.jpg"

# ---------- 1. 加载 INT8 模型 ----------
sess = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
in_name = sess.get_inputs()[0].name   # [1,3,320,320]

# ---------- 2. 读图 ----------
img0 = cv2.imread(IMG_PATH)
if img0 is None: raise FileNotFoundError(IMG_PATH)
h0, w0 = img0.shape[:2]

# ---------- 3. 预处理（exact 320×320 + FP32） ----------
blob = cv2.resize(img0, (320, 320))
blob = blob.astype(np.float32) / 255.0   # 输入保持 FP32，权重 INT8
blob = blob.transpose(2, 0, 1)[None]     # NCHW

# 4. 推理 & 维度整理
outs = sess.run(None, {in_name: blob})
bboxes, scores, lands = map(np.squeeze, outs)   # [N,4] [N,?] [N,10]
# 确保一维置信度
if scores.ndim > 1 and scores.shape[1] == 2:
    scores = scores[:, 1]          # 二分类取“人脸”列
else:
    scores = scores.flatten()

print("scores shape:", scores.shape)   # 应有 (N,) 或 (N,1)
print("scores前10个:", scores[:10])   # 看是否全 < 0.7
# 5. 画框 + landmarks
conf_thres = 0.7   # 置信度阈值
for i, conf in enumerate(scores):
    if conf < conf_thres: continue
    x1, y1, x2, y2 = bboxes[i]
    x1 = int(x1 * w0); y1 = int(y1 * h0)
    x2 = int(x2 * w0); y2 = int(y2 * h0)
    cv2.rectangle(img0, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(img0, f"{conf:.2f}", (x1, max(5, y1 - 5)),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    # 5 点 landmarks
    for k in range(5):
        px = int(lands[i, k*2]   * w0)
        py = int(lands[i, k*2+1] * h0)
        cv2.circle(img0, (px, py), 2, (0, 0, 255), -1)

# ---------- 6. 弹窗 ----------
cv2.imshow("RetinaFace-INT8-320", img0)
cv2.waitKey(0)
cv2.destroyAllWindows()
