import cv2
import insightface
import numpy as np

# Load InsightFace model
app = insightface.app.FaceAnalysis(
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)

# Larger detection size for better accuracy
app.prepare(ctx_id=0, det_size=(640, 640))

known_face_embeddings = []
known_face_names = []

# -------------------------
# Add Known Person
# -------------------------
def add_person(image_path, person_name):
    image = cv2.imread(image_path)

    faces = app.get(image)

    if len(faces) > 0:
        embedding = faces[0].embedding
        known_face_embeddings.append(embedding)
        known_face_names.append(person_name)
        print(f"{person_name} added successfully")
    else:
        print(f"No face found in {image_path}")

# Add your saved faces
add_person(r"D:\face\Lavi.png", "Lavi")
add_person(r"D:\face\gaurav image.jpg", "Gaurav")

# Webcam
cap = cv2.VideoCapture(1)

cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

while True:
    ret, frame = cap.read()

    if not ret:
        break

    faces = app.get(frame)

    for face in faces:
        bbox = face.bbox.astype(int)
        x1, y1, x2, y2 = bbox

        current_embedding = face.embedding

        name = "Unknown"
        best_score = -1

        for i, known_embedding in enumerate(known_face_embeddings):
            similarity = np.dot(current_embedding, known_embedding) / (
                np.linalg.norm(current_embedding) *
                np.linalg.norm(known_embedding)
            )

            if similarity > best_score:
                best_score = similarity
                name = known_face_names[i]

        # Threshold
        if best_score < 0.45:
            name = "Unknown"

        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

        label = f"{name} {best_score:.2f}"
        cv2.putText(
            frame,
            label,
            (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.8,
            (0, 255, 0),
            2
        )

    cv2.imshow("InsightFace Recognition", frame)

    if cv2.waitKey(1) == 27:
        break

cap.release()
cv2.destroyAllWindows()