File size: 2,409 Bytes
74a6c37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | 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()
|