Update app.py

app.py

@@ -1,147 +1,64 @@
 import streamlit as st
-from keras_preprocessing.image import ImageDataGenerator
-from keras.models import Sequential, load_model
-from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D
-import matplotlib.pyplot as plt
+from keras.models import load_model
+import cv2
 import numpy as np
-import os
-import random
-from sklearn.metrics import accuracy_score, confusion_matrix
-import seaborn as sns
-
-# Function to plot the confusion matrix
-def plot_confusion_matrix(cm, class_labels):
-    plt.figure(figsize=(8, 6))
-    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_labels, yticklabels=class_labels)
-    plt.ylabel('True label')
-    plt.xlabel('Predicted label')
-    st.pyplot()
-
-# Function to plot the training loss and accuracy
-def plot_training_history(history):
-    epochs = range(1, len(history.history['loss']) + 1)
-    plt.figure(figsize=(10, 6))
-    
-    plt.subplot(1, 2, 1)
-    plt.plot(epochs, history.history['loss'], 'y', label='Training loss')
-    plt.plot(epochs, history.history['val_loss'], 'r', label='Validation loss')
-    plt.title('Training and validation loss')
-    plt.xlabel('Epochs')
-    plt.ylabel('Loss')
-    plt.legend()
-    
-    plt.subplot(1, 2, 2)
-    plt.plot(epochs, history.history['accuracy'], 'y', label='Training acc')
-    plt.plot(epochs, history.history['val_accuracy'], 'r', label='Validation acc')
-    plt.title('Training and validation accuracy')
-    plt.xlabel('Epochs')
-    plt.ylabel('Accuracy')
-    plt.legend()
-    
-    st.pyplot()
-
-# Streamlit UI layout
-st.title('Emotion Detection from Facial Expressions')
-
-# Paths
-train_data_dir = 'data/train/'
-validation_data_dir = 'data/test/'
-
-# Hyperparameters
-IMG_HEIGHT = 48
-IMG_WIDTH = 48
-batch_size = 32
-epochs = 50
-
-# Image generators
-train_datagen = ImageDataGenerator(rescale=1./255,
-                                   rotation_range=30,
-                                   shear_range=0.3,
-                                   zoom_range=0.3,
-                                   horizontal_flip=True,
-                                   fill_mode='nearest')
-
-validation_datagen = ImageDataGenerator(rescale=1./255)
-
-train_generator = train_datagen.flow_from_directory(train_data_dir,
-                                                    color_mode='grayscale',
-                                                    target_size=(IMG_HEIGHT, IMG_WIDTH),
-                                                    batch_size=batch_size,
-                                                    class_mode='categorical',
-                                                    shuffle=True)
-
-validation_generator = validation_datagen.flow_from_directory(validation_data_dir,
-                                                              color_mode='grayscale',
-                                                              target_size=(IMG_HEIGHT, IMG_WIDTH),
-                                                              batch_size=batch_size,
-                                                              class_mode='categorical',
-                                                              shuffle=True)
+import time
+from PIL import Image
 
+# Load the pre-trained model
+model = load_model('models/emotion_detection_model_50epochs.h5', compile=False)
 class_labels = ['Angry', 'Disgust', 'Fear', 'Happy', 'Neutral', 'Sad', 'Surprise']
 
-# Model definition
-model = Sequential([
-    Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(48, 48, 1)),
-    Conv2D(64, kernel_size=(3, 3), activation='relu'),
-    MaxPooling2D(pool_size=(2, 2)),
-    Dropout(0.1),
-    Conv2D(128, kernel_size=(3, 3), activation='relu'),
-    MaxPooling2D(pool_size=(2, 2)),
-    Dropout(0.1),
-    Conv2D(256, kernel_size=(3, 3), activation='relu'),
-    MaxPooling2D(pool_size=(2, 2)),
-    Dropout(0.1),
-    Flatten(),
-    Dense(512, activation='relu'),
-    Dropout(0.2),
-    Dense(7, activation='softmax')
-])
-
-model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
-st.text(model.summary())
-
-train_path = "data/train/"
-test_path = "data/test/"
-
-num_train_imgs = sum([len(files) for r, d, files in os.walk(train_path)])
-num_test_imgs = sum([len(files) for r, d, files in os.walk(test_path)])
-
-# Train the model if button is pressed
-if st.button('Train Model'):
-    history = model.fit(train_generator,
-                        steps_per_epoch=num_train_imgs // batch_size,
-                        epochs=epochs,
-                        validation_data=validation_generator,
-                        validation_steps=num_test_imgs // batch_size)
-    model.save('models/emotion_detection_model_50epochs.h5')
-    st.success("Model trained and saved successfully!")
+# Function to preprocess each frame for prediction
+def preprocess_frame(frame):
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)  # Convert to grayscale
+    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(48, 48))
+
+    if len(faces) > 0:
+        (x, y, w, h) = faces[0]  # Use the first detected face
+        face = gray[y:y+h, x:x+w]
+        face = cv2.resize(face, (48, 48))
+        face = face.astype('float32') / 255.0  # Normalize pixel values
+        face = np.expand_dims(face, axis=0)
+        face = np.expand_dims(face, axis=-1)  # Reshape for the model (48,48,1)
+        return face, (x, y, w, h)
+    return None, None
 
-    # Plot training history
-    plot_training_history(history)
-
-# Load model and predict
-if st.button('Load Model and Test'):
-    my_model = load_model('models/emotion_detection_model_50epochs.h5', compile=False)
-    test_img, test_lbl = validation_generator.__next__()
-    predictions = my_model.predict(test_img)
-    
-    predictions = np.argmax(predictions, axis=1)
-    test_labels = np.argmax(test_lbl, axis=1)
-    
-    accuracy = accuracy_score(y_true=test_labels, y_pred=predictions)
-    st.write(f"Test Accuracy: {accuracy}")
-
-    # Confusion Matrix
-    cm = confusion_matrix(y_true=test_labels, y_pred=predictions)
-    plot_confusion_matrix(cm, class_labels)
-
-    # Display a random prediction
-    n = random.randint(0, test_img.shape[0] - 1)
-    image = test_img[n]
-    original_label = class_labels[test_labels[n]]
-    predicted_label = class_labels[predictions[n]]
-    
-    plt.imshow(image[:, :, 0], cmap='gray')
-    plt.title(f"Original Label: {original_label} | Predicted Label: {predicted_label}")
-    plt.axis("off")
-    st.pyplot()
+# Streamlit UI layout
+st.title("Real-time Emotion Detection")
+
+run = st.checkbox('Run Camera')
+
+FRAME_WINDOW = st.image([])
+
+# Start the camera and predict emotion
+cap = cv2.VideoCapture(0)  # Open default camera (change index if necessary)
+
+if run:
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Preprocess the frame
+        face, face_coords = preprocess_frame(frame)
+        if face is not None:
+            # Make emotion prediction
+            predictions = model.predict(face)
+            emotion_label = class_labels[np.argmax(predictions)]
+
+            # Draw a rectangle around the face and display the emotion label
+            (x, y, w, h) = face_coords
+            cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
+            cv2.putText(frame, emotion_label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+
+        # Convert BGR image to RGB
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        FRAME_WINDOW.image(frame_rgb)
+
+        # Small delay for smooth output
+        time.sleep(0.03)
+else:
+    cap.release()
+    cv2.destroyAllWindows()