import gradio as gr
import cv2
import os
import insightface
import numpy as np
from collections import deque

# Load InsightFace model
model = insightface.app.FaceAnalysis(allowed_modules=['detection', 'recognition'])
model.prepare(ctx_id=0, det_size=(640, 640))  # CPU: ctx_id=-1

# Normalize function
def normalize(v):
    return v / np.linalg.norm(v)

# Cosine similarity function
def cosine_similarity(a, b):
    return np.dot(a, b)

# Load known embeddings
known_embs = []
names = []
for fname in os.listdir("images"):
    if fname.lower().endswith(('.jpg', '.png')):
        img = cv2.imread(os.path.join("images", fname))
        faces = model.get(img)
        if faces:
            emb = normalize(faces[0].embedding)
            known_embs.append(emb)
            names.append(os.path.splitext(fname)[0])
            print(f"Loaded {fname}")
        else:
            print(f"No face in {fname}")

# Recognition function for an uploaded image
def recognize(image):
    face_buffers = {}
    frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    faces = model.get(frame)
    current_buffers = {}

    for face in faces:
        x1, y1, x2, y2 = face.bbox.astype(int)
        emb = normalize(face.embedding)
        face_id = f"{x1}-{y1}-{x2}-{y2}"

        if face_id not in face_buffers:
            face_buffers[face_id] = deque(maxlen=5)

        face_buffers[face_id].append(emb)
        current_buffers[face_id] = face_buffers[face_id]

        avg_emb = normalize(np.mean(face_buffers[face_id], axis=0))
        sims = [cosine_similarity(avg_emb, known) for known in known_embs]
        max_idx = np.argmax(sims)
        name = "Unknown"
        if sims[max_idx] > 0.5:
            name = names[max_idx]

        # Draw on the frame
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

# Gradio interface
iface = gr.Interface(
    fn=recognize,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Recognized Faces"),
    title="Face Recognition with InsightFace",
    description="Upload an image, and the system will identify known faces from the 'images/' folder."
)

if __name__ == "__main__":
    iface.launch()