Update app.py

app.py

@@ -7,6 +7,14 @@ import streamlit as st
 # Load ONNX model
 onnx_model = ort.InferenceSession("onnx_model.onnx")
 
+# Emotion labels (same as the model's output classes)
+emotion_labels = ["Anger", "Disgust", "Fear", "Happy", "Sadness", "Surprise", "Neutral"]
+
+# Softmax function to convert logits to probabilities
+def softmax(logits):
+    exp_logits = np.exp(logits - np.max(logits))  # Stability trick
+    return exp_logits / np.sum(exp_logits)
+
 # Preprocess image function
 def preprocess_image(image):
     """Preprocess image to match model input requirements"""
@@ -30,7 +38,14 @@ def predict_emotion_onnx(onnx_model, image_input):
     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
+    
+    # Apply softmax to the output logits
+    probabilities = softmax(prediction[0][0])  # We assume batch size of 1
+    
+    # Get the predicted emotion label (index of the highest probability)
+    predicted_class = np.argmax(probabilities)
+    
+    return emotion_labels[predicted_class], probabilities[predicted_class]
 
 # Streamlit interface
 st.title("Emotion Recognition with ONNX")
@@ -45,7 +60,8 @@ if uploaded_file is not None:
     image_input = preprocess_image(image)
 
     # Predict the emotion
-    emotion_prediction = predict_emotion_onnx(onnx_model, image_input)
+    emotion_label, probability = predict_emotion_onnx(onnx_model, image_input)
 
-    # Display the prediction
-    st.write(f"Predicted Emotion: {emotion_prediction[0]}")
+    # Display the predicted emotion and probability
+    st.write(f"Predicted Emotion: {emotion_label}")
+    st.write(f"Confidence: {probability:.2f}")