Update app.py

app.py

@@ -1,89 +1,43 @@
-import cv2
-import numpy as np
 import onnxruntime as ort
+import numpy as np
+import cv2
 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
+# Load ONNX model
+onnx_model = ort.InferenceSession("onnx_model.onnx")
 
-# 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
+# Preprocess image function
+def preprocess_image(image):
+    """Preprocess image to match model input requirements"""
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)  # Convert to grayscale
+    image_resized = cv2.resize(image, (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 image
+    image_input = np.expand_dims(image_input, axis=3)  # Add channel dimension
+    image_input = image_input.astype(np.float32) / 255.0  # Normalize
     return image_input
 
-# 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
+# Predict emotion using the 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
+# Streamlit interface
+st.title("Emotion Recognition with ONNX")
 
-# 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"
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
-# 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 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
+    # Preprocess the image
+    image_input = preprocess_image(image)
 
-    # Combine predictions from both models
-    final_emotion = combine_predictions(affectnet_emotion, onnx_emotion)
+    # Predict the emotion
+    emotion_prediction = predict_emotion_onnx(onnx_model, image_input)
 
-    # Display the final emotion prediction
-    st.write(f"Detected Emotion: {final_emotion}")
+    # Display the prediction
+    st.write(f"Predicted Emotion: {emotion_prediction[0]}")