Update app.py

app.py

@@ -1,111 +1,103 @@
 import streamlit as st
 import cv2
 import numpy as np
-import time
-import os
+from datetime import datetime
+import torch
+from facenet_pytorch import MTCNN, InceptionResnetV1
 from keras.models import load_model
 from PIL import Image
 import pymongo
-from datetime import datetime
+import os
 import tempfile
-from facenet_pytorch import MTCNN
 
 # MongoDB Atlas Connection String
 MONGO_URI = "mongodb+srv://test:[email protected]/?retryWrites=true&w=majority"
-
-# Connect to MongoDB
 client = pymongo.MongoClient(MONGO_URI)
 db = client.get_database("emotion_detection")
 collection = db.get_collection("face_data")
 
-# Larger title
+# Page title
 st.markdown("<h1 style='text-align: center;'>Emotion Detection with Face Recognition</h1>", unsafe_allow_html=True)
-
-# Smaller subtitle
 st.markdown("<h3 style='text-align: center;'>angry, fear, happy, neutral, sad, surprise</h3>", unsafe_allow_html=True)
 
-# Start time for measuring performance
-start = time.time()
-
 # Load the emotion detection model
 @st.cache_resource
 def load_emotion_model():
-    model = load_model('CNN_Model_acc_75.h5')  # Ensure this file is in your Space
+    model = load_model('CNN_Model_acc_75.h5')  # Ensure this file is present
     return model
 
-model = load_emotion_model()
-print("Time taken to load model: ", time.time() - start)
-
-# Emotion labels
+emotion_model = load_emotion_model()
 emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
 
-# Initialize MTCNN for face detection
+# Initialize FaceNet model and MTCNN
+facenet = InceptionResnetV1(pretrained='vggface2').eval()
 mtcnn = MTCNN()
 
-# Load known faces and names
+# Known faces and embeddings
 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')): 
+        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 using mtcnn
-            faces, _ = mtcnn.detect(image)  # Use the correct method detect()
-
-            if faces is not None:
-                for face in faces:
-                    # Convert bounding box coordinates to integers
-                    x, y, w, h = map(int, face)  # Ensure the coordinates are integers
-                    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
-
-    # Train the recognizer with the known faces
-    if known_faces:  # Only train if we have faces
-        face_recognizer.train(known_faces, np.array([i for i in range(len(known_faces))]))
+            image = Image.open(image_path).convert("RGB")
+            face, _ = mtcnn.detect(image)
+
+            if face is not None:
+                face_box = face[0].astype(int)
+                cropped_face = image.crop((face_box[0], face_box[1], face_box[2], face_box[3]))
+                cropped_face = cropped_face.resize((160, 160))
+                face_tensor = np.array(cropped_face).transpose(2, 0, 1) / 255.0
+                face_tensor = torch.tensor(face_tensor, dtype=torch.float32).unsqueeze(0)
+                
+                with torch.no_grad():
+                    embedding = facenet(face_tensor).numpy()
+                
+                known_faces.append(embedding)
+                known_names.append(image_name.split('.')[0])
 
 load_known_faces()
 
-# Process a single frame
-def process_frame(frame):
-    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-    faces, _ = mtcnn.detect(frame)  # Use the correct detect method
-
-    result_text = ""  # Initialize result text
-
-    if faces is not None and len(faces) > 0:
-        for face in faces:
-            # Convert bounding box coordinates to integers
-            x, y, w, h = map(int, face)  # Ensure the coordinates are integers
-            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.equalizeHist(roi_gray)
+def recognize_face(embedding):
+    min_distance = float('inf')
+    name = "Unknown"
+    for idx, known_embedding in enumerate(known_faces):
+        distance = np.linalg.norm(known_embedding - embedding)
+        if distance < min_distance and distance < 0.6:  # Threshold for similarity
+            min_distance = distance
+            name = known_names[idx]
+    return name
 
-            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)
+def process_frame(frame):
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    faces, _ = mtcnn.detect(frame_rgb)
+    result_text = ""
+
+    if faces is not None:
+        for face_box in faces:
+            x1, y1, x2, y2 = map(int, face_box)
+            cropped_face = frame_rgb[y1:y2, x1:x2]
+            resized_face = cv2.resize(cropped_face, (48, 48))
+            face_normalized = resized_face / 255.0
+            face_array = np.expand_dims(face_normalized, axis=0)
+
+            # Emotion prediction
+            predictions = emotion_model.predict(face_array)
             emotion = emotion_labels[np.argmax(predictions[0])]
 
             # Face recognition
-            name = "Unknown"
-            label, confidence = face_recognizer.predict(roi_gray)
-            if confidence < 100:
-                name = known_names[label]
+            cropped_face_for_recognition = cv2.resize(cropped_face, (160, 160))
+            face_tensor = np.array(cropped_face_for_recognition).transpose(2, 0, 1) / 255.0
+            face_tensor = torch.tensor(face_tensor, dtype=torch.float32).unsqueeze(0)
 
-            # Format result text
-            result_text = f"{name} is feeling {emotion}"
+            with torch.no_grad():
+                face_embedding = facenet(face_tensor).numpy()
+
+            name = recognize_face(face_embedding)
 
-            # Save data to MongoDB if face is recognized (name != Unknown)
+            # Save record in MongoDB
             if name != "Unknown":
                 timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                 document = {
@@ -113,45 +105,35 @@ def process_frame(frame):
                     "emotion": emotion,
                     "timestamp": timestamp
                 }
-                
-                # Check if the same record already exists (to prevent duplicates)
-                existing_record = collection.find_one({"name": name, "timestamp": timestamp})
-                if not existing_record:
-                    # Insert the data into MongoDB
+                if not collection.find_one({"name": name, "timestamp": timestamp}):
                     collection.insert_one(document)
-                    print(f"Data inserted into MongoDB: {document}")
-
-            # 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)
 
+            # Display result
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            result_text = f"{name} is feeling {emotion}"
+            cv2.putText(frame, result_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
     else:
         result_text = "No face detected!"
 
     return frame, result_text
 
-# Video feed display
 def video_feed(video_source):
-    frame_placeholder = st.empty()  # Placeholder for displaying video frames
-    text_placeholder = st.empty()  # Placeholder for displaying result text
+    frame_placeholder = st.empty()
+    text_placeholder = st.empty()
 
     while True:
         ret, frame = video_source.read()
         if not ret:
             break
-
         frame, result_text = process_frame(frame)
-
-        # Display frame and result text
         frame_placeholder.image(frame, channels="BGR", use_column_width=True)
         text_placeholder.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
 
-# Sidebar for user input source selection
+# Sidebar for input source
 upload_choice = st.sidebar.radio("Choose Input Source", ["Upload Image", "Upload Video", "Camera"])
 
 if upload_choice == "Camera":
     image = st.camera_input("Take a picture")
-
     if image:
         frame = np.array(Image.open(image))
         frame, result_text = process_frame(frame)
@@ -159,8 +141,7 @@ if upload_choice == "Camera":
         st.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
 
 elif upload_choice == "Upload Image":
-    uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "gif"])
-
+    uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
     if uploaded_image:
         image = Image.open(uploaded_image)
         frame = np.array(image)
@@ -169,29 +150,18 @@ elif upload_choice == "Upload Image":
         st.markdown(f"<h3 style='text-align: center;'>{result_text}</h3>", unsafe_allow_html=True)
 
 elif upload_choice == "Upload Video":
-    uploaded_video = st.file_uploader("Upload Video", type=["mp4", "mov", "avi", "mkv", "webm"])
-
+    uploaded_video = st.file_uploader("Upload Video", type=["mp4", "mov", "avi"])
     if uploaded_video:
         with tempfile.NamedTemporaryFile(delete=False) as tfile:
             tfile.write(uploaded_video.read())
             video_source = cv2.VideoCapture(tfile.name)
             video_feed(video_source)
 
-# Display the latest 5 records stored in MongoDB with timestamp and proper alignment
-st.markdown("### MongoDB Records")
-records = collection.find().sort("timestamp", -1).limit(5)  # Limit to 5 records
-
+# Display recent MongoDB records
+st.markdown("### Recent Records")
+records = collection.find().sort("timestamp", -1).limit(5)
 for record in records:
-    col1, col2, col3 = st.columns([3, 3, 2])
-    with col1:
-        st.write(f"**Name**: {record['name']}")
-    with col2:
-        st.write(f"**Emotion**: {record['emotion']}")
-    with col3:
-        st.write(f"**Timestamp**: {record['timestamp']}")
-    
-    # Delete record
-    delete_button = st.button("Delete", key=f"delete_{record['_id']}")
-    if delete_button:
-        collection.delete_one({"_id": record["_id"]})
-        st.success(f"Record with ID {record['_id']} has been deleted.")
+    col1, col2, col3 = st.columns(3)
+    col1.write(f"**Name**: {record['name']}")
+    col2.write(f"**Emotion**: {record['emotion']}")
+    col3.write(f"**Timestamp**: {record['timestamp']}")