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

MODEL_PATH = "model/version-slim-320.onnx"   # 240×320 模型
IMG_PATH   = "img/friends.jpg"

# ---------- 1. 加载模型 ----------
sess = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
in_name = sess.get_inputs()[0].name   # [1,3,240,320]
out_score = sess.get_outputs()[0].name  # [1,N,2]
out_box   = sess.get_outputs()[1].name  # [1,N,4]  x1,y1,x2,y2 归一化

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

# ---------- 3. 预处理（exact 240×320 + FP32） ----------
blob = cv2.resize(img0, (320, 240))          # 宽 320，高 240
blob = blob.astype(np.float32) / 255.0
blob = blob.transpose(2, 0, 1)[None]         # NCHW

# ---------- 4. 推理 ----------
scores, boxes = sess.run([out_score, out_box], {in_name: blob})
scores = np.squeeze(scores)[:, 1]            # 人脸置信度 [N]
boxes  = np.squeeze(boxes)                   # 归一化框 [N,4]

# ---------- 5. NMS 去重 + 自适应框 ----------
conf_thres = 0.7
nms_iou    = 0.3
rel_size   = 0.12                            # 框边长 = 0.12×图像宽

# 先过滤低置信
keep = scores > conf_thres
boxes, scores = boxes[keep], scores[keep]

# OpenCV NMSBoxes 需要 [x,y,w,h] + 分数
nms_boxes = []
for b in boxes:
    x1, y1, x2, y2 = b
    w_box = int((x2 - x1) * w0)              # 真实像素宽
    h_box = int((y2 - y1) * h0)              # 真实像素高
    nms_boxes.append([int(x1 * w0), int(y1 * h0), w_box, h_box])

if nms_boxes:
    indices = cv2.dnn.NMSBoxes(nms_boxes, scores.tolist(), score_threshold=conf_thres, nms_threshold=nms_iou)
    for idx in indices.flatten():
        x, y, w, h = nms_boxes[idx]
        # 自适应框：以检测框为中心，用相对尺寸画框
        cx = x + w // 2
        cy = y + h // 2
        side = int(rel_size * w0)              # 统一边长
        x1 = max(0, cx - side // 2)
        y1 = max(0, cy - side // 2)
        x2 = min(w0, cx + side // 2)
        y2 = min(h0, cy + side // 2)

        cv2.rectangle(img0, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(img0, f"{scores[idx]:.2f}", (x1, max(5, y1 - 5)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

# ---------- 6. 弹窗 ----------
cv2.imshow("Face-Detection-ONNX", img0)
cv2.waitKey(0)
cv2.destroyAllWindows()
