Update app.py

app.py

@@ -5,6 +5,7 @@ import streamlit as st
 from insightface.app import FaceAnalysis
 from glob import glob
 from tqdm import tqdm
+from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
 import shutil
 import zipfile
 
@@ -13,38 +14,53 @@ def extract_zip(zip_file_path, extract_dir):
     with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
         zip_ref.extractall(extract_dir)
 
-# Function to recognize faces
-def recognize_faces(frame, names, embeddings, app):
-    # Perform face analysis on the frame
-    faces = app.get(frame)
 
-    # Process each detected face separately
-    for face in faces:
-        # Retrieve the embedding for the detected face
-        detected_embedding = face.normed_embedding
+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
 
-        # Calculate similarity scores with known embeddings
-        scores = np.dot(detected_embedding, np.array(embeddings).T)
-        scores = np.clip(scores, 0., 1.)
+    def _recognize_faces(self, frame):
+        if self.names is None or self.embeddings is None:
+            return frame
 
-        # Find the index with the highest score
-        idx = np.argmax(scores)
-        max_score = scores[idx]
+        # Perform face analysis on the frame
+        faces = self.app.get(frame)
 
-        # Check if the maximum score is above a certain threshold (adjust as needed)
-        threshold = 0.7
-        if max_score >= threshold:
-            recognized_name = names[idx]
-        else:
-            recognized_name = "Unknown"
+        # Process each detected face separately
+        for face in faces:
+            # Retrieve the embedding for the detected face
+            detected_embedding = face.normed_embedding
 
-        # 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)
+            # Calculate similarity scores with known embeddings
+            scores = np.dot(detected_embedding, np.array(self.embeddings).T)
+            scores = np.clip(scores, 0., 1.)
 
-    return frame
+            # 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)
+
+        return frame
+
+    def transform(self, frame):
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame = self._recognize_faces(frame)
+        return frame
 
 # Function to get embeddings
 def get_embeddings(db_dir):
@@ -124,39 +140,21 @@ def main():
         uploaded_embeddings = st.file_uploader("Upload embeddings.npy", type="npy")
 
         if uploaded_names and uploaded_embeddings:
-            # Load names and embeddings
             names = np.load(uploaded_names)
             embeddings = np.load(uploaded_embeddings)
 
-            # Initialize FaceAnalysis app
-            app = FaceAnalysis(name='buffalo_l')
-            app.prepare(ctx_id=0, det_size=(640, 640))
-
-            # Display a button to start webcam
-            if st.button("Start Webcam"):
-                # Start capturing video from webcam
-                cap = cv2.VideoCapture(0)
-
-                # Process each frame in real-time
-                while True:
-                    # Capture frame-by-frame
-                    ret, frame = cap.read()
-                    if not ret:
-                        break
-
-                    # Perform face recognition
-                    frame = recognize_faces(frame, names, embeddings, app)
-
-                    # Display the resulting frame
-                    st.image(frame, channels="BGR", use_column_width=True)
+            # Initialize transformer with names and embeddings
+            transformer = FaceRecognitionTransformer()
+            transformer.names = names
+            transformer.embeddings = embeddings
 
-                    # Break the loop if 'q' is pressed or the user closes the Streamlit app
-                    if st.button("Stop"):
-                        break
+            # Create WebRTC streamer
+            webrtc_ctx = webrtc_streamer(
+                key="example",
+                video_transformer_factory=FaceRecognitionTransformer,
+                async_transform=True,
+            )
 
-                # Release the capture
-                cap.release()
-                cv2.destroyAllWindows()