Update app.py

app.py

@@ -1,19 +1,15 @@
 import streamlit as st
+from streamlit_webrtc import webrtc_streamer, VideoTransformerBase, VideoFrame
 import cv2
 import numpy as np
 from datetime import datetime
-import torch
-from facenet_pytorch import MTCNN, InceptionResnetV1
 from keras.models import load_model
-from PIL import Image
 import sqlite3
 import os
-import tempfile
 
-# SQLite Database
+# Database Initialization
 DB_NAME = "emotion_detection.db"
 
-# Initialize SQLite Database
 def initialize_database():
     conn = sqlite3.connect(DB_NAME)
     cursor = conn.cursor()
@@ -30,173 +26,58 @@ def initialize_database():
 
 initialize_database()
 
-# Load the emotion detection model
+# Load emotion detection model
 @st.cache_resource
 def load_emotion_model():
-    model = load_model('CNN_Model_acc_75.h5')  # Ensure the file exists
-    return model
+    return load_model('CNN_Model_acc_75.h5')
 
 emotion_model = load_emotion_model()
 emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
 
-# Initialize FaceNet model and MTCNN
-facenet = InceptionResnetV1(pretrained='vggface2').eval()
-mtcnn = MTCNN()
-
-# Directory for known faces
-KNOWN_FACES_DIR = "known_faces"
-if not os.path.exists(KNOWN_FACES_DIR):
-    os.makedirs(KNOWN_FACES_DIR)
-
-# Load known faces and embeddings
-known_faces = []
-known_names = []
-
-def load_known_faces():
-    for image_name in os.listdir(KNOWN_FACES_DIR):
-        if image_name.endswith(('.jpg', '.jpeg', '.png')):
-            image_path = os.path.join(KNOWN_FACES_DIR, image_name)
-            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()
-
-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
-            min_distance = distance
-            name = known_names[idx]
-    return name
-
-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
-            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)
-
-            with torch.no_grad():
-                face_embedding = facenet(face_tensor).numpy()
-
-            name = recognize_face(face_embedding)
-
-            # Save record in SQLite
-            if name != "Unknown":
-                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                conn = sqlite3.connect(DB_NAME)
-                cursor = conn.cursor()
-                cursor.execute("""
-                INSERT INTO face_data (name, emotion, timestamp)
-                VALUES (?, ?, ?)
-                """, (name, emotion, timestamp))
-                conn.commit()
-                conn.close()
-
-            # 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 Transformer for Streamlit WebRTC
+class EmotionDetector(VideoTransformerBase):
+    def __init__(self):
+        self.model = emotion_model
+
+    def transform(self, frame: VideoFrame) -> VideoFrame:
+        img = frame.to_ndarray(format="bgr24")
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        faces = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml").detectMultiScale(
+            gray, scaleFactor=1.1, minNeighbors=5, minSize=(48, 48)
+        )
+
+        for (x, y, w, h) in faces:
+            face = gray[y:y + h, x:x + w]
+            face_resized = cv2.resize(face, (48, 48))
+            face_normalized = face_resized / 255.0
+            face_reshaped = np.reshape(face_normalized, (1, 48, 48, 1))
+            
+            prediction = self.model.predict(face_reshaped)
+            emotion = emotion_labels[np.argmax(prediction[0])]
+            
+            # Draw bounding box and label
+            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
+            cv2.putText(img, emotion, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+
+        return VideoFrame.from_ndarray(img, format="bgr24")
 
 # Sidebar menu
-menu = st.sidebar.selectbox("Menu", ["Home", "Register New Face", "View Records"])
+menu = st.sidebar.selectbox("Menu", ["Home", "View Records"])
 
 if menu == "Home":
-    st.title("Emotion Detection")
-    st.write("Choose input source to start detection.")
-    upload_choice = st.radio("Choose Input Source", ["Camera", "Upload Image", "Upload Video"])
-
-    if upload_choice == "Camera":
-        cap = cv2.VideoCapture(0)
-        if not cap.isOpened():
-            st.error("Unable to access the camera.")
-        else:
-            while True:
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                frame, result_text = process_frame(frame)
-                st.image(frame, channels="BGR")
-                st.write(result_text)
-            cap.release()
-
-    elif upload_choice == "Upload Image":
-        uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
-        if uploaded_image:
-            image = Image.open(uploaded_image)
-            frame = np.array(image)
-            frame, result_text = process_frame(frame)
-            st.image(frame)
-            st.write(result_text)
-
-    elif upload_choice == "Upload Video":
-        uploaded_video = st.file_uploader("Upload Video", type=["mp4", "mov", "avi"])
-        if uploaded_video:
-            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-                temp_file.write(uploaded_video.read())
-                video_source = cv2.VideoCapture(temp_file.name)
-                while video_source.isOpened():
-                    ret, frame = video_source.read()
-                    if not ret:
-                        break
-                    frame, result_text = process_frame(frame)
-                    st.image(frame, channels="BGR")
-                    st.write(result_text)
-                video_source.release()
-
-elif menu == "Register New Face":
-    st.title("Register New Face")
-    name = st.text_input("Enter Name")
-    if st.button("Capture Image"):
-        cap = cv2.VideoCapture(0)
-        if not cap.isOpened():
-            st.error("Unable to access the camera.")
-        else:
-            ret, frame = cap.read()
-            if ret:
-                image_path = os.path.join(KNOWN_FACES_DIR, f"{name}.jpg")
-                cv2.imwrite(image_path, frame)
-                st.success(f"Face registered successfully for {name}")
-                load_known_faces()  # Refresh known faces
-            cap.release()
+    st.title("Real-Time Emotion Detection")
+    st.write("Using your camera for real-time emotion detection.")
+    
+    webrtc_streamer(
+        key="emotion-detection",
+        video_transformer_factory=EmotionDetector,
+        media_stream_constraints={"video": True, "audio": False},
+    )
 
 elif menu == "View Records":
     st.title("View Records")
     st.subheader("Recent Records")
+    
     conn = sqlite3.connect(DB_NAME)
     cursor = conn.cursor()
     cursor.execute("SELECT name, emotion, timestamp FROM face_data ORDER BY timestamp DESC LIMIT 5")