facefuse / webcam_app.py
N3tron's picture
Update webcam_app.py
e070794 verified
import os
import cv2
import numpy as np
import streamlit as st
import insightface
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
from glob import glob
from tqdm import tqdm
from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
import shutil
import zipfile
import image_app
import utils
from utils import app
class FaceRecognitionTransformer(VideoTransformerBase):
def __init__(self):
self.app = FaceAnalysis(name='buffalo_l')
self.app.prepare(ctx_id=0, det_size=(640, 640))
self.names = None
self.embeddings = None
def _recognize_faces(self, frame):
if self.names is None or self.embeddings is None:
return frame
# Perform face analysis on the frame
faces = self.app.get(frame)
# Process each detected face separately
for face in faces:
# Retrieve the embedding for the detected face
detected_embedding = face.normed_embedding
# Calculate similarity scores with known embeddings
scores = np.dot(detected_embedding, np.array(self.embeddings).T)
scores = np.clip(scores, 0., 1.)
# Find the index with the highest score
idx = np.argmax(scores)
max_score = scores[idx]
# Check if the maximum score is above a certain threshold (adjust as needed)
threshold = 0.7
if max_score >= threshold:
recognized_name = self.names[idx]
else:
recognized_name = "Unknown"
# Draw bounding box around the detected face
bbox = face.bbox.astype(int)
cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
# Write recognized name within the bounding box
cv2.putText(frame, recognized_name, (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Debug print
print("Detected face:", recognized_name, "with confidence:", max_score)
return frame
def transform(self, frame):
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = self._recognize_faces(frame)
return frame