Update app.py

app.py

@@ -1,177 +1,74 @@
-import streamlit as st
+import os
 import cv2
 import numpy as np
-import time
-import os
-import requests
-from keras.models import load_model
-from PIL import Image
 import face_recognition
-import pymongo
-from datetime import datetime
-import tempfile
-
-# 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
-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
-    return model
-
-model = load_emotion_model()
-st.write("Time taken to load model: ", time.time() - start)
-
-# Emotion labels
-emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
-
-# Load known faces and names
-known_faces = []
-known_names = []
-
-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 = face_recognition.load_image_file(image_path)
-            encoding = face_recognition.face_encodings(image)
-            
-            if encoding:  # Ensure encoding was found
-                known_faces.append(encoding[0])  # Store the first encoding (assumes 1 face per image)
-                known_names.append(image_name.split('.')[0])  # Assuming file name is the person's name
-
-load_known_faces()
-
-# Process a single frame
-def process_frame(frame):
-    # Convert the image to RGB for face_recognition
-    rgb_frame = frame[:, :, ::-1]
-    
-    # Detect faces
-    face_locations = face_recognition.face_locations(rgb_frame)
-    face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)
-    
-    result_text = ""  # Initialize result text
-
-    if face_encodings:
-        for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
-            # Check if the detected face matches any known faces
-            matches = face_recognition.compare_faces(known_faces, face_encoding)
-            name = "Unknown"
-            
-            # Find the name of the recognized face
-            if True in matches:
-                first_match_index = matches.index(True)
-                name = known_names[first_match_index]
-            
-            # Emotion detection
-            face_roi = frame[top:bottom, left:right]
-            face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)
-            face_roi = cv2.resize(face_roi, (48, 48))
-            face_roi = np.expand_dims(face_roi, axis=0) / 255.0  # Normalize
-
-            predictions = model.predict(face_roi)
-            emotion = emotion_labels[np.argmax(predictions[0])]
-
-            # Format result text
-            result_text = f"{name} is feeling {emotion}"
-
-            # Save data to MongoDB if face is recognized (name != Unknown)
-            if name != "Unknown":
-                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                document = {
-                    "name": name,
-                    "emotion": emotion,
-                    "timestamp": timestamp
-                }
-                # Insert the data into MongoDB
-                collection.insert_one(document)
-                print(f"Data inserted into MongoDB: {document}")
-
-            # Draw bounding box and label
-            cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
-            cv2.putText(frame, result_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-
-    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
-
-    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
-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)
-        st.image(frame, caption='Processed Image', use_column_width=True)
-        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"])
-
-    if uploaded_image:
-        image = Image.open(uploaded_image)
-        frame = np.array(image)
-        frame, result_text = process_frame(frame)
-        st.image(frame, caption='Processed Image', use_column_width=True)
-        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"])
-
-    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 records stored in MongoDB with latest records on top
-st.markdown("### MongoDB Records")
-records = collection.find().sort("timestamp", -1)  # Sort records by timestamp in descending order
+import streamlit as st
+from PIL import Image
 
-for record in records:
-    col1, col2, col3 = st.columns([3, 3, 1])
-    with col1:
-        st.write(f"**Name**: {record['name']}")
-    with col2:
-        st.write(f"**Emotion**: {record['emotion']}")
-    with col3:
-        st.write(f"**Timestamp**: {record['timestamp']}")
+# List for storing images and class names
+Images = []
+classnames = []
+
+# Directory where known faces are stored
+directory = "known_faces"  # Directory containing images of known faces
+
+# Load images and classnames from the directory
+myList = os.listdir(directory)
+
+for cls in myList:
+    if os.path.splitext(cls)[1] in [".jpg", ".jpeg"]:
+        img_path = os.path.join(directory, cls)
+        curImg = cv2.imread(img_path)
+        Images.append(curImg)
+        classnames.append(os.path.splitext(cls)[0])
+
+# Function to find face encodings
+def findEncodings(Images):
+    encodeList = []
+    for img in Images:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        encode = face_recognition.face_encodings(img)[0]
+        encodeList.append(encode)
+    return encodeList
+
+# Find encodings for known faces
+encodeListknown = findEncodings(Images)
+
+# Streamlit UI for capturing image using the camera
+img_file_buffer = st.camera_input("Take a picture")
+
+if img_file_buffer is not None:
+    # Open the captured image using PIL
+    test_image = Image.open(img_file_buffer)
+    image = np.asarray(test_image)
+
+    # Resize and convert image for face recognition
+    imgS = cv2.resize(image, (0, 0), None, 0.25, 0.25)  # Resize for faster processing
+    imgS = cv2.cvtColor(imgS, cv2.COLOR_BGR2RGB)
     
-    # Delete record button
-    delete_button = st.button(f"Delete {record['_id']}", key=record['_id'])
-    if delete_button:
-        collection.delete_one({"_id": record["_id"]})
-        st.success(f"Record with ID {record['_id']} has been deleted.")
+    # Find faces in the current frame
+    facesCurFrame = face_recognition.face_locations(imgS)
+    encodesCurFrame = face_recognition.face_encodings(imgS, facesCurFrame)
+
+    name = "Unknown"  # Default name for unknown faces
+
+    if len(encodesCurFrame) > 0:
+        for encodeFace, faceLoc in zip(encodesCurFrame, facesCurFrame):
+            # Compare the face with known faces
+            matches = face_recognition.compare_faces(encodeListknown, encodeFace)
+            faceDis = face_recognition.face_distance(encodeListknown, encodeFace)
+            matchIndex = np.argmin(faceDis)  # Find the closest match
+
+            if matches[matchIndex]:
+                name = classnames[matchIndex].upper()  # Assign name if match is found
+
+            # Draw a rectangle around the face and label it with the name
+            y1, x2, y2, x1 = faceLoc
+            y1, x2, y2, x1 = y1 * 4, x2 * 4, y2 * 4, x1 * 4  # Scale coordinates back
+            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.rectangle(image, (x1, y2 - 35), (x2, y2), (0, 255, 0), cv2.FILLED)
+            cv2.putText(image, name, (x1 + 6, y2 - 6), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 2)
+
+    # Display the image with the bounding box and label
+    st.image(image, caption='Processed Image', use_column_width=True)
+    st.markdown(f"<h3 style='text-align: center;'>{name} is recognized</h3>", unsafe_allow_html=True)