Update app.py

app.py

@@ -1,211 +1,108 @@
 import os
 import cv2
 import numpy as np
-import sqlite3
 import streamlit as st
 from datetime import datetime
-from PIL import Image
-from keras.models import load_model
+from tensorflow.keras.models import load_model
 
-# Constants
-DATABASE_NAME = "emotion_recognition.db"
+# Directories
 KNOWN_FACES_DIR = "known_faces"
-if not os.path.exists(KNOWN_FACES_DIR):
-    os.makedirs(KNOWN_FACES_DIR)
-
-# Initialize Database
-def init_db():
-    conn = sqlite3.connect(DATABASE_NAME)
-    cursor = conn.cursor()
-    cursor.execute('''
-        CREATE TABLE IF NOT EXISTS attendance_log (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            name TEXT NOT NULL,
-            emotion TEXT NOT NULL,
-            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
-        )
-    ''')
-    conn.commit()
-    conn.close()
-
-init_db()
-
-def log_attendance(name, emotion):
-    conn = sqlite3.connect(DATABASE_NAME)
-    cursor = conn.cursor()
-    cursor.execute("INSERT INTO attendance_log (name, emotion) VALUES (?, ?)", (name, emotion))
-    conn.commit()
-    conn.close()
+EMOTION_MODEL_PATH = "emotion_model.h5"
+CASCADE_PATH = "haarcascade_frontalface_default.xml"
 
-def fetch_recent_activity():
-    conn = sqlite3.connect(DATABASE_NAME)
-    cursor = conn.cursor()
-    cursor.execute("SELECT name, emotion, timestamp FROM attendance_log ORDER BY timestamp DESC LIMIT 10")
-    rows = cursor.fetchall()
-    conn.close()
-    return rows
-
-# Load Emotion Detection Model
-@st.cache_resource
-def load_emotion_model():
-    model = load_model('CNN_Model_acc_75.h5')
-    return model
-
-emotion_model = load_emotion_model()
-emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
+# Constants
+IMG_SIZE = (200, 200)
 
-# Initialize LBPH Face Recognizer
+# Load models
+emotion_model = load_model(EMOTION_MODEL_PATH)
+face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
 face_recognizer = cv2.face.LBPHFaceRecognizer_create()
 
+# Helper Functions
+def load_emotion_labels():
+    return ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
+
 def train_recognizer():
     faces = []
     labels = []
-    for name in os.listdir(KNOWN_FACES_DIR):
-        person_dir = os.path.join(KNOWN_FACES_DIR, name)
-        if not os.path.isdir(person_dir):
+    label_map = {}
+    for idx, person_name in enumerate(os.listdir(KNOWN_FACES_DIR)):
+        person_path = os.path.join(KNOWN_FACES_DIR, person_name)
+        if not os.path.isdir(person_path):
             continue
-        for filename in os.listdir(person_dir):
-            filepath = os.path.join(person_dir, filename)
-            if not filepath.lower().endswith(('.jpg', '.jpeg', '.png')):
-                continue
-            image = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
-            if image is not None:
-                faces.append(image)
-                labels.append(name)
-    if not faces or not labels:
-        raise ValueError("No valid training data found in the known faces directory.")
-    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]
+        label_map[idx] = person_name
+        for filename in os.listdir(person_path):
+            filepath = os.path.join(person_path, filename)
+            if filepath.lower().endswith(('.jpg', '.jpeg', '.png')):
+                img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
+                if img is not None:
+                    faces.append(img)
+                    labels.append(idx)
+    if len(faces) == 0:
+        st.warning("No valid training data found. Add faces first.")
+        return {}
     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 Recent Activity"])
-
-# Main App Logic
-if sidebar_choice == "Register New Face":
-    st.header("Register New Face")
-    name = st.text_input("Enter Name")
-    use_camera = st.checkbox("Use Camera to Capture Face")
-    capture_button = st.button("Capture Face" if use_camera else "Upload Image")
-
-    if use_camera and 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:  # Capture 5 images
-            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()
-
-        if captured_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 with {len(captured_faces)} captured images!")
-        else:
-            st.warning("No faces captured. Please try again.")
-
-    elif not use_camera and capture_button and name:
-        uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
-        if uploaded_image:
-            image = np.array(Image.open(uploaded_image))
-            gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-            faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(100, 100))
-
-            if len(faces) > 0:
-                person_dir = os.path.join(KNOWN_FACES_DIR, name)
-                if not os.path.exists(person_dir):
-                    os.makedirs(person_dir)
+    return {v: k for k, v in label_map.items()}
+
+def detect_faces(image):
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
+    return gray, faces
+
+def detect_emotions(face_img):
+    resized_face = cv2.resize(face_img, (48, 48))
+    normalized_face = resized_face / 255.0
+    reshaped_face = np.expand_dims(normalized_face, axis=(0, -1))
+    emotion_probabilities = emotion_model.predict(reshaped_face)
+    emotion_idx = np.argmax(emotion_probabilities)
+    return load_emotion_labels()[emotion_idx]
+
+# Streamlit App
+st.title("Face Recognition and Emotion Detection")
+st.sidebar.title("Options")
+option = st.sidebar.selectbox("Choose an action", ["Home", "Register New Face", "Recognize Faces"])
+
+# Train the recognizer initially
+if option != "Register New Face":
+    label_map = train_recognizer()
+
+if option == "Home":
+    st.write("Use the sidebar to register new faces or recognize them.")
+
+elif option == "Register New Face":
+    person_name = st.text_input("Enter the person's name")
+    capture_mode = st.radio("Select input method", ["Use Camera", "Upload Image"])
+    
+    if person_name and st.button("Register Face"):
+        person_dir = os.path.join(KNOWN_FACES_DIR, person_name)
+        os.makedirs(person_dir, exist_ok=True)
+        if capture_mode == "Use Camera":
+            st.warning("Switch to a device with a camera for this option.")
+        elif capture_mode == "Upload Image":
+            uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+            if uploaded_file:
+                img = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)
+                gray, faces = detect_faces(img)
                 for (x, y, w, h) in faces:
-                    face_roi = gray_image[y:y + h, x:x + w]
-                    face_filename = os.path.join(person_dir, f"{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg")
-                    cv2.imwrite(face_filename, face_roi)
-                label_ids_rev = train_recognizer()
-                st.success(f"Face for {name} has been registered successfully!")
-            else:
-                st.warning("No face detected. Try another image.")
-        else:
-            st.warning("Please upload an image.")
-
-elif sidebar_choice == "View Recent Activity":
-    st.header("Recent Activity")
-    logs = fetch_recent_activity()
-    if logs:
-        for name, emotion, timestamp in logs:
-            st.write(f"Name: {name}, Emotion: {emotion}, Timestamp: {timestamp}")
-    else:
-        st.write("No recent activity found.")
-
-else:
-    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 = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(100, 100))
-        result_text = ""
+                    face_img = gray[y:y+h, x:x+w]
+                    resized_img = cv2.resize(face_img, IMG_SIZE)
+                    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+                    filepath = os.path.join(person_dir, f"{timestamp}.jpg")
+                    cv2.imwrite(filepath, resized_img)
+                st.success("Face registered successfully!")
+        label_map = train_recognizer()
+
+elif option == "Recognize Faces":
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        img = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)
+        gray, faces = detect_faces(img)
         for (x, y, w, h) in faces:
-            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(label, emotion)
-            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)
-        return frame, result_text
-
-    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=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)
-            cv2.imshow("Emotion Detection", frame)
-            if cv2.waitKey(1) & 0xFF == ord('q'):
-                break
-        cap.release()
-        cv2.destroyAllWindows()
+            face_img = gray[y:y+h, x:x+w]
+            resized_img = cv2.resize(face_img, IMG_SIZE)
+            label, confidence = face_recognizer.predict(resized_img)
+            name = label_map.get(label, "Unknown")
+            emotion = detect_emotions(face_img)
+            cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 2)
+            cv2.putText(img, f"{name}, {emotion}", (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+        st.image(img, channels="BGR")