# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Face Matching: OpenVINO Runtime sample.
# See README.md for prerequisites and inputs. Run: python3 openvino_sample.py

import cv2
import numpy as np
import openvino as ov

DETECTION_MODEL = "intel/face-detection-adas-0001/FP16/face-detection-adas-0001.xml"
REID_MODEL = "intel/face-reidentification-retail-0095/FP16/face-reidentification-retail-0095.xml"
REFERENCE_IMAGE = "face_a.jpg"  # reference: left subject captured from the video
SCENE_VIDEO = "test_video.mp4"  # scene containing the reference subject plus others
SCENE_FRAME = 480               # frame index of the two-person scene (~40s)
CONF_THRESHOLD = 0.5
MATCH_THRESHOLD = 0.5

core = ov.Core()

# Change device to "GPU" or "NPU" to run on integrated GPU or NPU.
det_compiled = core.compile_model(core.read_model(DETECTION_MODEL), "CPU")
reid_compiled = core.compile_model(core.read_model(REID_MODEL), "CPU")

det_input = det_compiled.input(0)
det_h, det_w = det_input.shape[2], det_input.shape[3]
reid_input = reid_compiled.input(0)
reid_h, reid_w = reid_input.shape[2], reid_input.shape[3]


def detect_faces(img):
    h0, w0 = img.shape[:2]
    blob = cv2.resize(img, (det_w, det_h))
    blob = blob.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
    dets = det_compiled([blob])[det_compiled.output(0)][0][0]
    faces = []
    for d in dets:
        if float(d[2]) < CONF_THRESHOLD:
            continue
        x1 = max(0, int(d[3] * w0))
        y1 = max(0, int(d[4] * h0))
        x2 = min(w0, int(d[5] * w0))
        y2 = min(h0, int(d[6] * h0))
        if x2 > x1 and y2 > y1:
            faces.append((x1, y1, x2, y2))
    return faces


def get_embedding(img, bbox):
    x1, y1, x2, y2 = bbox
    crop = img[y1:y2, x1:x2]
    blob = cv2.resize(crop, (reid_w, reid_h))
    blob = blob.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
    emb = reid_compiled([blob])[reid_compiled.output(0)].flatten()
    return emb / np.linalg.norm(emb)


# 1. Embed the reference face.
reference = cv2.imread(REFERENCE_IMAGE)
ref_faces = detect_faces(reference)
if not ref_faces:
    raise SystemExit("No face detected in the reference image")
ref_bbox = max(ref_faces, key=lambda b: (b[2] - b[0]) * (b[3] - b[1]))
ref_emb = get_embedding(reference, ref_bbox)

# 2. Read a scene frame that contains several people.
cap = cv2.VideoCapture(SCENE_VIDEO)
cap.set(cv2.CAP_PROP_POS_FRAMES, SCENE_FRAME)
ok, frame = cap.read()
cap.release()
if not ok:
    raise SystemExit(f"Could not read frame {SCENE_FRAME} from {SCENE_VIDEO}")

# 3. Compare every face in the scene to the reference; keep the best match.
best_bbox = None
best_sim = 0.0
for bbox in detect_faces(frame):
    sim = float(np.dot(get_embedding(frame, bbox), ref_emb))
    if sim > best_sim:
        best_sim = sim
        best_bbox = bbox

if best_bbox is not None and best_sim >= MATCH_THRESHOLD:
    x1, y1, x2, y2 = best_bbox
    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(frame, f"MATCH {best_sim:.2f}", (x1, max(15, y1 - 8)),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    print(f"Matched reference identity, similarity {best_sim:.4f}")
else:
    print(f"No matching face found (best similarity {best_sim:.4f})")

cv2.imwrite("output_openvino.jpg", frame)
print("Saved: output_openvino.jpg")