Update app.py

app.py

@@ -1,27 +1,11 @@
-import streamlit as st
-import onnxruntime as ort
 import cv2
 import numpy as np
+import onnxruntime as ort  # Import onnxruntime for model inference
+import streamlit as st
 from PIL import Image
 
-# Set page configuration
-st.set_page_config(page_title="Emotion Recognition App", layout="centered")
-st.title("Emotion Recognition App")
-
-# Load the ONNX model using onnxruntime
-@st.cache_resource
-def load_model():
-    model_path = "onnx_model.onnx"  # Ensure this is the correct path to your uploaded ONNX model
-    return ort.InferenceSession(model_path)
-
-# Load the emotion detection model
-emotion_model = load_model()
-
-# Process the uploaded image
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
-def preprocess_image(image):
-    """Preprocess image to match model input requirements"""
+# Preprocess image to match model input requirements
+def preprocess_image(image, face_landmarks=None):
     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)  # Convert image to grayscale
     image_resized = cv2.resize(image, (48, 48))  # Resize image to 48x48
     image_input = np.expand_dims(image_resized, axis=0)  # Add batch dimension
@@ -29,27 +13,48 @@ def preprocess_image(image):
     image_input = image_input.astype(np.float32) / 255.0  # Normalize the image
     return image_input
 
+# Check if smile is present in the facial landmarks
+def check_for_smile(face_landmarks):
+    """Simple rule to check for smile based on landmarks"""
+    mouth = face_landmarks['bottom_lip'] + face_landmarks['top_lip']
+    mouth_distance = np.linalg.norm(np.array(mouth[0]) - np.array(mouth[-1]))
+    if mouth_distance > 30:  # This threshold might need adjustment
+        return True
+    return False
+
+# Display emotion with post-processing to check for smiles
+def display_emotion_with_smile(emotion, face_landmarks=None):
+    if emotion == 6 and face_landmarks:  # 'Neutral' is typically 6 in the emotion_map
+        if check_for_smile(face_landmarks):
+            return "Happiness"  # Override neutral with happiness if a smile is detected
+    return display_emotion(emotion)  # Otherwise return the normal emotion
+
+# Predict emotion with smile detection
+def predict_emotion_with_smile(image_input, face_landmarks=None):
+    """Run inference and predict the emotion, considering smile detection"""
+    emotion = predict_emotion(image_input)  # Normal emotion prediction
+    emotion_label = display_emotion_with_smile(emotion, face_landmarks)
+    return emotion_label
+
+# Load ONNX model
+def load_model():
+    model_path = '"onnx_model.onnx"'  # Make sure to set the correct path
+    emotion_model = ort.InferenceSession(model_path)
+    return emotion_model
+
+# Predict emotion using the ONNX model
 def predict_emotion(image_input):
-    """Run inference and predict the emotion"""
+    emotion_model = load_model()
     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
+    return np.argmax(prediction[0])
+
+# Streamlit app code
+st.title("Emotion Recognition App")
 
-# 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")
+# Upload an image
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 # If an image is uploaded
 if uploaded_file is not None:
@@ -61,8 +66,7 @@ if uploaded_file is not None:
     image_input = preprocess_image(image)
 
     # Predict the emotion
-    emotion = predict_emotion(image_input)
-    emotion_label = display_emotion(emotion)
+    emotion_label = predict_emotion_with_smile(image_input)
 
     # Display the predicted emotion
     st.write(f"Detected Emotion: {emotion_label}")