Update app.py

app.py

@@ -1,89 +1,89 @@
 import cv2
 import numpy as np
-import os
 import onnxruntime as ort
-import streamlit as st
 from PIL import Image
+import streamlit as st
+
+# Step 1: Preprocess Image for ONNX Model (224x224 input required)
+def preprocess_for_onnx(image):
+    """Preprocess image for the ONNX emotion model."""
+    image_resized = cv2.resize(image, (224, 224))  # Resize to match ONNX input
+    image_normalized = image_resized.astype(np.float32) / 255.0  # Normalize image
+    image_input = np.transpose(image_normalized, (2, 0, 1))  # Change dimension to (C, H, W)
+    image_input = np.expand_dims(image_input, axis=0)  # Add batch dimension
+    return image_input
 
-# 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
+# Step 2: Preprocess Image for AffectNet Model (48x48 grayscale input)
+def preprocess_for_affectnet(image):
+    """Preprocess image for AffectNet (grayscale, 48x48)."""
+    image_resized = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)  # Convert to grayscale
+    image_resized = cv2.resize(image_resized, (48, 48))  # Resize to 48x48
     image_input = np.expand_dims(image_resized, axis=0)  # Add batch dimension
-    image_input = np.expand_dims(image_input, axis=0)  # Add channel dimension (for grayscale)
-    image_input = image_input.astype(np.float32) / 255.0  # Normalize the image
+    image_input = np.expand_dims(image_input, axis=0)  # Add channel dimension for grayscale
+    image_input = image_input.astype(np.float32) / 255.0  # Normalize 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(emotion):
-    """Map predicted emotion index to a label"""
-    emotion_map = {
-        0: "Anger",
-        1: "Disgust",
-        2: "Fear",
-        3: "Happiness",
-        4: "Sadness",
-        5: "Surprise",
-        6: "Neutral"
-    }
-    return emotion_map.get(emotion, "Unknown")
-
-# Display emotion with smile detection
-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 this is the correct path
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Model file {model_path} not found!")
-    emotion_model = ort.InferenceSession(model_path)
-    return emotion_model
-
-# Predict emotion using the ONNX model
-def predict_emotion(image_input):
-    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})
-    return np.argmax(prediction[0])
-
-# Streamlit app code
-st.title("Emotion Recognition App")
-
-# Upload an image
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
-# If an image is uploaded
+# Step 3: Load ONNX Model
+def load_onnx_model(model_path):
+    """Load the ONNX emotion detection model."""
+    return ort.InferenceSession(model_path)
+
+# Step 4: Predict Emotion with ONNX Model
+def predict_emotion_onnx(onnx_model, image_input):
+    """Predict emotion using the ONNX model."""
+    input_name = onnx_model.get_inputs()[0].name
+    output_name = onnx_model.get_outputs()[0].name
+    prediction = onnx_model.run([output_name], {input_name: image_input})
+    return prediction
+
+# Step 5: Predict Emotion with AffectNet (Placeholder)
+def predict_emotion_affectnet(image_input):
+    """Predict emotion using AffectNet (Placeholder function)."""
+    # Placeholder for the actual AffectNet prediction
+    # This would involve loading a model trained on AffectNet and making predictions.
+    emotion = "happy"  # Example: This would be the predicted emotion
+    return emotion
+
+# Step 6: Combine Predictions from Both Models
+def combine_predictions(affectnet_emotion, onnx_emotion):
+    """Combine predictions from AffectNet and ONNX models."""
+    if affectnet_emotion == onnx_emotion:
+        return affectnet_emotion  # If both models agree, return the same emotion
+    else:
+        return "Uncertain"  # If they disagree, return "Uncertain"
+
+# Step 7: Main Application Logic
+# Load the ONNX model from the application directory
+onnx_model_path = 'onnx_model.onnx'  # Path to your ONNX model
+onnx_model = load_onnx_model(onnx_model_path)
+
+# Step 8: Streamlit App Logic (User Interaction)
+st.title('Emotion Detection App')
+
+# Upload image via Streamlit interface
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+# If an image is uploaded, proceed with emotion prediction
 if uploaded_file is not None:
     # Open and display the uploaded image
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Preprocess the image
-    image_input = preprocess_image(image)
+    # Preprocess image for AffectNet (grayscale)
+    image_input_affectnet = preprocess_for_affectnet(image)
+
+    # Preprocess image for ONNX model (RGB)
+    image_input_onnx = preprocess_for_onnx(np.array(image))
+
+    # Predict emotion using AffectNet model
+    affectnet_emotion = predict_emotion_affectnet(image_input_affectnet)
+
+    # Predict emotion using ONNX model
+    onnx_prediction = predict_emotion_onnx(onnx_model, image_input_onnx)
+    onnx_emotion = onnx_prediction[0][0]  # Assuming the model outputs a single emotion label
 
-    # Predict the emotion
-    emotion_label = predict_emotion_with_smile(image_input)
+    # Combine predictions from both models
+    final_emotion = combine_predictions(affectnet_emotion, onnx_emotion)
 
-    # Display the predicted emotion
-    st.write(f"Detected Emotion: {emotion_label}")
+    # Display the final emotion prediction
+    st.write(f"Detected Emotion: {final_emotion}")