Update app.py

app.py

@@ -2,31 +2,20 @@ import sqlite3
 import streamlit as st
 import cv2
 import numpy as np
-import time
 import os
 from keras.models import load_model
+from datetime import datetime
 from PIL import Image
-import tempfile
-
-# 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)
 
 # Database setup
 DATABASE_NAME = "emotion_recognition.db"
+KNOWN_FACES_DIR = "known_faces"
+if not os.path.exists(KNOWN_FACES_DIR):
+    os.makedirs(KNOWN_FACES_DIR)
 
 def init_db():
     conn = sqlite3.connect(DATABASE_NAME)
     cursor = conn.cursor()
-    cursor.execute('''
-        CREATE TABLE IF NOT EXISTS registered_faces (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            name TEXT NOT NULL,
-            image BLOB NOT NULL
-        )
-    ''')
     cursor.execute('''
         CREATE TABLE IF NOT EXISTS attendance_log (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -40,21 +29,6 @@ def init_db():
 
 init_db()
 
-def register_face(name, image):
-    conn = sqlite3.connect(DATABASE_NAME)
-    cursor = conn.cursor()
-    cursor.execute("INSERT INTO registered_faces (name, image) VALUES (?, ?)", (name, image))
-    conn.commit()
-    conn.close()
-
-def fetch_registered_faces():
-    conn = sqlite3.connect(DATABASE_NAME)
-    cursor = conn.cursor()
-    cursor.execute("SELECT id, name FROM registered_faces")
-    rows = cursor.fetchall()
-    conn.close()
-    return rows
-
 def log_attendance(name, emotion):
     conn = sqlite3.connect(DATABASE_NAME)
     cursor = conn.cursor()
@@ -70,41 +44,77 @@ def fetch_recent_activity():
     conn.close()
     return rows
 
-# Load the emotion model
+# Load pre-trained emotion detection model
 @st.cache_resource
 def load_emotion_model():
     model = load_model('CNN_Model_acc_75.h5')
     return model
 
-model = load_emotion_model()
-
-# Emotion labels
+emotion_model = load_emotion_model()
 emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
 
+# Initialize LBPH face recognizer
+face_recognizer = cv2.face.LBPHFaceRecognizer_create()
+
+def train_recognizer():
+    faces = []
+    labels = []
+    for name in os.listdir(KNOWN_FACES_DIR):
+        for filename in os.listdir(os.path.join(KNOWN_FACES_DIR, name)):
+            filepath = os.path.join(KNOWN_FACES_DIR, name, filename)
+            image = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
+            faces.append(image)
+            labels.append(name)
+    label_ids = {name: idx for idx, name in enumerate(set(labels))}
+    label_ids_rev = {idx: name for name, idx in label_ids.items()}
+    labels = [label_ids[label] for label in labels]
+    face_recognizer.train(faces, np.array(labels))
+    return label_ids_rev
+
+label_ids_rev = train_recognizer()
+
 # Sidebar options
-sidebar_choice = st.sidebar.selectbox("Choose an option", ["Emotion Detection", "Register New Face", "View Registered Faces", "Recent Activity"])
+sidebar_choice = st.sidebar.selectbox("Choose an option", ["Emotion Detection", "Register New Face", "View Recent Activity"])
 
+# Main App Logic
 if sidebar_choice == "Register New Face":
     st.header("Register New Face")
     name = st.text_input("Enter Name")
-    uploaded_image = st.file_uploader("Upload Face Image", type=["png", "jpg", "jpeg"])
-    if name and uploaded_image:
-        image = np.array(Image.open(uploaded_image))
-        _, buffer = cv2.imencode('.jpg', image)
-        register_face(name, buffer.tobytes())
-        st.success(f"Successfully registered {name}!")
-
-elif sidebar_choice == "View Registered Faces":
-    st.header("Registered Faces")
-    faces = fetch_registered_faces()
-    if faces:
-        for face_id, name in faces:
-            st.write(f"ID: {face_id}, Name: {name}")
-    else:
-        st.write("No faces registered yet.")
-
-elif sidebar_choice == "Recent Activity":
-    st.header("Recent Activity (Attendance Log)")
+    capture_button = st.button("Capture Face via Camera")
+    if capture_button and name:
+        cap = cv2.VideoCapture(0)
+        st.write("Capturing face... Look into the camera.")
+        captured_faces = []
+        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+
+        while len(captured_faces) < 5:
+            ret, frame = cap.read()
+            if not ret:
+                st.error("Error capturing video")
+                break
+            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+            faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(100, 100))
+            for (x, y, w, h) in faces:
+                face_roi = gray_frame[y:y + h, x:x + w]
+                captured_faces.append(face_roi)
+                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
+            cv2.imshow("Face Registration", frame)
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+        cap.release()
+        cv2.destroyAllWindows()
+
+        # Save faces
+        person_dir = os.path.join(KNOWN_FACES_DIR, name)
+        if not os.path.exists(person_dir):
+            os.makedirs(person_dir)
+        for i, face in enumerate(captured_faces):
+            cv2.imwrite(os.path.join(person_dir, f"{name}_{i}.jpg"), face)
+        label_ids_rev = train_recognizer()
+        st.success(f"{name} has been registered successfully!")
+
+elif sidebar_choice == "View Recent Activity":
+    st.header("Recent Activity")
     logs = fetch_recent_activity()
     if logs:
         for name, emotion, timestamp in logs:
@@ -113,58 +123,52 @@ elif sidebar_choice == "Recent Activity":
         st.write("No recent activity found.")
 
 else:  # Emotion Detection
-    st.sidebar.write("Emotion Labels: Angry, Fear, Happy, Neutral, Sad, Surprise")
-
-    upload_choice = st.radio("Choose input source", ["Upload Image", "Upload Video", "Camera"])
+    st.header("Emotion Detection with Face Recognition")
+    mode = st.radio("Choose mode", ["Image", "Camera"])
+    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
 
     def process_frame(frame):
         gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        faces = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml').detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-
+        faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(100, 100))
         result_text = ""
         for (x, y, w, h) in faces:
-            roi_gray = gray_frame[y:y+h, x:x+w]
-            roi_color = frame[y:y+h, x:x+w]
-            face_roi = cv2.resize(roi_color, (48, 48))
-            face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB) / 255.0
-            face_roi = np.expand_dims(face_roi, axis=0)
-
-            predictions = model.predict(face_roi)
-            emotion = emotion_labels[np.argmax(predictions[0])]
-
-            label = "Unknown"  # Placeholder for face recognition (add later)
+            face_roi = gray_frame[y:y + h, x:x + w]
+            face_resized = cv2.resize(face_roi, (150, 150))
+            label_id, confidence = face_recognizer.predict(face_resized)
+            label = label_ids_rev.get(label_id, "Unknown")
+
+            # Emotion Detection
+            face_color = cv2.resize(frame[y:y + h, x:x + w], (48, 48)) / 255.0
+            face_color = np.expand_dims(cv2.cvtColor(face_color, cv2.COLOR_BGR2RGB), axis=0)
+            emotion_prediction = emotion_model.predict(face_color)
+            emotion = emotion_labels[np.argmax(emotion_prediction[0])]
+
+            # Log Attendance
             log_attendance(label, emotion)
 
+            # Annotate frame
             result_text = f"{label} is feeling {emotion}"
-            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, 1, (0, 255, 0), 2)
+            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, 1, (0, 255, 0), 2)
         return frame, result_text
 
-    if upload_choice == "Upload Image":
-        uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
+    if mode == "Image":
+        uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
         if uploaded_image:
             image = np.array(Image.open(uploaded_image))
             frame, result_text = process_frame(image)
-            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", "avi", "mkv"])
-        if uploaded_video:
-            with tempfile.NamedTemporaryFile(delete=False) as tfile:
-                tfile.write(uploaded_video.read())
-                video_source = cv2.VideoCapture(tfile.name)
-                while True:
-                    ret, frame = video_source.read()
-                    if not ret:
-                        break
-                    frame, result_text = process_frame(frame)
-                    st.image(frame, channels="BGR", use_column_width=True)
-
-    elif upload_choice == "Camera":
-        image = st.camera_input("Take a picture")
-        if image:
-            frame = np.array(Image.open(image))
+            st.image(frame, caption=result_text)
+
+    elif mode == "Camera":
+        cap = cv2.VideoCapture(0)
+        st.write("Press 'q' to exit.")
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
             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)
+            cv2.imshow("Emotion Detection", frame)
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+        cap.release()
+        cv2.destroyAllWindows()