Update app.py

app.py

@@ -5,13 +5,10 @@ import time
 import os
 from keras.models import load_model
 from PIL import Image
-from mtcnn import MTCNN  # MTCNN for better face detection
-from keras.preprocessing import image
-from tensorflow.keras.applications.inception_v3 import preprocess_input
 import pymongo
 from datetime import datetime
 import tempfile
-from facenet_pytorch import MTCNN, InceptionResnetV1  # FaceNet model for face recognition
+from facenet_pytorch import MTCNN
 
 # MongoDB Atlas Connection String
 MONGO_URI = "mongodb+srv://test:[email protected]/?retryWrites=true&w=majority"
@@ -42,57 +39,67 @@ print("Time taken to load model: ", time.time() - start)
 # Emotion labels
 emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
 
-# Load FaceNet Model for Face Recognition
-@st.cache_resource
-def load_facenet_model():
-    # Load FaceNet model for face recognition
-    facenet_model = InceptionResnetV1(pretrained='vggface2').eval()
-    return facenet_model
+# Initialize MTCNN for face detection
+mtcnn = MTCNN()
+
+# Load known faces and names
+known_faces = []
+known_names = []
+face_recognizer = cv2.face.LBPHFaceRecognizer_create()
+
+def load_known_faces():
+    folder_path = "known_faces"  # Folder containing known face images
+    for image_name in os.listdir(folder_path):
+        if image_name.endswith(('.jpg', '.jpeg', '.png')): 
+            image_path = os.path.join(folder_path, image_name)
+            image = cv2.imread(image_path)
+            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            # Detect face in the image
+            faces = mtcnn.detect(image)[0]  # Use the correct detect method
 
-facenet_model = load_facenet_model()
+            for face in faces:
+                x, y, w, h = face[0], face[1], face[2], face[3]
+                roi_gray = gray[y:y+h, x:x+w]
+                known_faces.append(roi_gray)
+                known_names.append(image_name.split('.')[0])  # Assuming file name is the person's name
 
-# MTCNN for face detection
-detector = MTCNN()
+    # Train the recognizer with the known faces
+    face_recognizer.train(known_faces, np.array([i for i in range(len(known_faces))]))
+
+load_known_faces()
 
 # Process a single frame
 def process_frame(frame):
-    # Detect faces
-    faces = detector.detect_faces(frame)
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    faces = mtcnn.detect(frame)[0]  # Use the correct detect method
 
     result_text = ""  # Initialize result text
 
-    if len(faces) > 0:
+    if faces is not None and len(faces) > 0:
         for face in faces:
-            x, y, w, h = face['box']
+            x, y, w, h = face[0], face[1], face[2], face[3]
             roi_color = frame[y:y+h, x:x+w]
+            roi_gray = gray[y:y+h, x:x+w]
 
             # Apply histogram equalization for better feature extraction
-            roi_gray = cv2.cvtColor(roi_color, cv2.COLOR_BGR2GRAY)
             roi_gray = cv2.equalizeHist(roi_gray)
 
-            # Emotion detection
             face_roi = cv2.resize(roi_color, (48, 48))
             face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)
             face_roi = np.expand_dims(face_roi, axis=0) / 255.0  # Normalize
 
+            # Emotion detection
             predictions = model.predict(face_roi)
             emotion = emotion_labels[np.argmax(predictions[0])]
 
-            # Extract embedding for face recognition using FaceNet
-            face_embedding = facenet_model.predict(preprocess_input(np.expand_dims(face['keypoints'], axis=0)))
-
-            # Compare face embeddings with known faces (replace with your own database)
-            known_faces = []  # Load known face embeddings from database
-            known_names = []  # Corresponding names
-
+            # Face recognition
             name = "Unknown"
-            min_distance = float('inf')
+            label, confidence = face_recognizer.predict(roi_gray)
+            if confidence < 100:
+                name = known_names[label]
 
-            for i, known_face in enumerate(known_faces):
-                distance = np.linalg.norm(face_embedding - known_face)
-                if distance < min_distance:
-                    min_distance = distance
-                    name = known_names[i]
+            # Format result text
+            result_text = f"{name} is feeling {emotion}"
 
             # Save data to MongoDB if face is recognized (name != Unknown)
             if name != "Unknown":
@@ -106,8 +113,6 @@ def process_frame(frame):
                 collection.insert_one(document)
                 print(f"Data inserted into MongoDB: {document}")
 
-            result_text = f"{name} is feeling {emotion}"
-
             # Draw bounding box and label
             cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
             cv2.putText(frame, result_text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)