|
|
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 |
|
|
|
|
|
|
|
|
faces = self.app.get(frame) |
|
|
|
|
|
|
|
|
for face in faces: |
|
|
|
|
|
detected_embedding = face.normed_embedding |
|
|
|
|
|
|
|
|
scores = np.dot(detected_embedding, np.array(self.embeddings).T) |
|
|
scores = np.clip(scores, 0., 1.) |
|
|
|
|
|
|
|
|
idx = np.argmax(scores) |
|
|
max_score = scores[idx] |
|
|
|
|
|
|
|
|
threshold = 0.7 |
|
|
if max_score >= threshold: |
|
|
recognized_name = self.names[idx] |
|
|
else: |
|
|
recognized_name = "Unknown" |
|
|
|
|
|
|
|
|
bbox = face.bbox.astype(int) |
|
|
cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2) |
|
|
|
|
|
cv2.putText(frame, recognized_name, (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) |
|
|
|
|
|
|
|
|
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 |