Update app.py

app.py

@@ -1,51 +1,68 @@
 import streamlit as st
-import numpy as np
-from PIL import Image
 import onnxruntime as ort
 import cv2
+import numpy as np
+from PIL import Image
 
-# Set the page config
+# Set page configuration
 st.set_page_config(page_title="Emotion Recognition App", layout="centered")
-
 st.title("Emotion Recognition App")
 
-# Upload an image
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
 # Load the ONNX model using onnxruntime
 @st.cache_resource
 def load_model():
-    # Path to the uploaded ONNX model (should be the name of the model file you uploaded)
-    model_path = "emotion_model.onnx"
+    model_path = "onnx_model.onnx"  # Ensure this is the correct path to your uploaded ONNX model
     return ort.InferenceSession(model_path)
 
-# Load the model
+# Load the emotion detection model
 emotion_model = load_model()
 
-# Class labels for facial emotions (based on the training dataset)
-emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happiness', 'Sadness', 'Surprise', 'Neutral']
+# Process the uploaded image
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
-# Preprocess image to match model input requirements
 def preprocess_image(image):
-    # Convert image to grayscale and resize to match the input size expected by the model
-    image_gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-    image_resized = cv2.resize(image_gray, (64, 64))  # The model expects 64x64 input
-    image_normalized = image_resized / 255.0  # Normalize to [0, 1] range
-    image_reshaped = np.expand_dims(image_normalized, axis=0)  # Add batch dimension
-    image_reshaped = np.expand_dims(image_reshaped, axis=0)  # Add channel dimension (1 channel for grayscale)
-    return image_reshaped.astype(np.float32)
+    """Preprocess image to match model input requirements"""
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    image_resized = cv2.resize(image, (224, 224))  # Resize image to model input size
+    image_input = np.transpose(image_resized, (2, 0, 1))  # Change data format for the model
+    image_input = image_input.astype(np.float32) / 255.0  # Normalize the image
+    image_input = np.expand_dims(image_input, axis=0)  # Add batch dimension
+    return image_input
 
-# Process the uploaded image
+def predict_emotion(image_input):
+    """Run inference and predict the emotion"""
+    input_name = emotion_model.get_inputs()[0].name
+    output_name = emotion_model.get_outputs()[0].name
+    prediction = emotion_model.run([output_name], {input_name: image_input})
+    emotion = np.argmax(prediction[0])  # Get the class with the highest probability
+    return emotion
+
+# Define a function to display emotion text
+def display_emotion(emotion):
+    """Map emotion index to a human-readable emotion"""
+    emotion_map = {
+        0: "Anger",
+        1: "Disgust",
+        2: "Fear",
+        3: "Happiness",
+        4: "Sadness",
+        5: "Surprise",
+        6: "Neutral"
+    }
+    return emotion_map.get(emotion, "Unknown")
+
+# If an image is uploaded
 if uploaded_file is not None:
-    # Open and preprocess the image
-    image = Image.open(uploaded_file).convert("RGB")
-    processed_image = preprocess_image(image)
+    # Open and display the uploaded image
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Perform inference
-    input_name = emotion_model.get_inputs()[0].name
-    outputs = emotion_model.run(None, {input_name: processed_image})
-    predicted_class = np.argmax(outputs[0], axis=1)[0]  # Get the index of the highest probability
-    predicted_emotion = emotion_labels[predicted_class]
+    # Preprocess the image
+    image_input = preprocess_image(image)
+
+    # Predict the emotion
+    emotion = predict_emotion(image_input)
+    emotion_label = display_emotion(emotion)
 
-    # Display the results
-    st.image(image, caption=f"Detected Emotion: {predicted_emotion}", use_column_width=True)
+    # Display the predicted emotion
+    st.write(f"Detected Emotion: {emotion_label}")