Update app.py

app.py

@@ -1,18 +1,19 @@
-import sqlite3
-import streamlit as st
+import os
 import cv2
 import numpy as np
-import os
-from keras.models import load_model
+import sqlite3
+import streamlit as st
 from datetime import datetime
 from PIL import Image
+from keras.models import load_model
 
-# Database setup
+# Constants
 DATABASE_NAME = "emotion_recognition.db"
 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()
@@ -44,20 +45,16 @@ def fetch_recent_activity():
     conn.close()
     return rows
 
-# Load pre-trained emotion detection model
+# Load Emotion Detection Model
 @st.cache_resource
 def load_emotion_model():
-    try:
-        model = load_model('CNN_Model_acc_75.h5')
-        return model
-    except Exception as e:
-        st.error("Error loading the emotion detection model. Please ensure the file exists and is accessible.")
-        return None
+    model = load_model('CNN_Model_acc_75.h5')
+    return model
 
 emotion_model = load_emotion_model()
-emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise'] if emotion_model else []
+emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
 
-# Initialize LBPH face recognizer
+# Initialize LBPH Face Recognizer
 face_recognizer = cv2.face.LBPHFaceRecognizer_create()
 
 def train_recognizer():
@@ -65,11 +62,11 @@ def train_recognizer():
     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):  # Skip files, process only directories
+        if not os.path.isdir(person_dir):
             continue
         for filename in os.listdir(person_dir):
             filepath = os.path.join(person_dir, filename)
-            if not filepath.lower().endswith(('.jpg', '.jpeg', '.png')):  # Ensure valid image formats
+            if not filepath.lower().endswith(('.jpg', '.jpeg', '.png')):
                 continue
             image = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
             if image is not None:
@@ -83,17 +80,19 @@ def train_recognizer():
     face_recognizer.train(faces, np.array(labels))
     return label_ids_rev
 
+label_ids_rev = train_recognizer()
 
-# Sidebar options
+# 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")
-    capture_button = st.button("Capture Face")
+    use_camera = st.checkbox("Use Camera to Capture Face")
+    capture_button = st.button("Capture Face" if use_camera else "Upload Image")
 
-    if capture_button and name:
+    if use_camera and capture_button and name:
         cap = cv2.VideoCapture(0)
         st.write("Capturing face... Look into the camera.")
         captured_faces = []
@@ -132,6 +131,29 @@ if sidebar_choice == "Register New Face":
         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)
+                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()
@@ -141,7 +163,7 @@ elif sidebar_choice == "View Recent Activity":
     else:
         st.write("No recent activity found.")
 
-else:  # Emotion Detection
+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')
@@ -156,25 +178,24 @@ else:  # Emotion Detection
             label_id, confidence = face_recognizer.predict(face_resized)
             label = label_ids_rev.get(label_id, "Unknown")
 
-            if emotion_model:
-                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])]
+            # 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)
+            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, _ = process_frame(image)
-            st.image(frame, caption="Processed Image with Face Annotations")
-
+            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.")