Update app.py

app.py

@@ -1,74 +1,65 @@
-import onnxruntime as ort
-import numpy as np
 import cv2
-from PIL import Image
+import numpy as np
+import onnxruntime as ort
 import streamlit as st
-import kagglehub  # Import kagglehub to load the AffectNet dataset
-
-# Download the AffectNet dataset
-path = kagglehub.dataset_download("fatihkgg/affectnet-yolo-format")
-print("Path to AffectNet dataset:", path)
-
-# Emotion labels for AffectNet
-emotion_labels = ["Anger", "Disgust", "Fear", "Happy", "Sadness", "Surprise", "Neutral"]
-
-# Load ONNX model
-onnx_model = ort.InferenceSession("onnx_model.onnx")
+from PIL import Image
 
-# 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)
+# Load the ONNX model
+def load_model(model_path='onnx_model.onnx'):
+    # Load the ONNX model
+    model = ort.InferenceSession(model_path)
+    return model
 
-# Preprocess image function for ONNX model
+# Preprocess the image
 def preprocess_image(image):
-    """Preprocess image to match model input requirements"""
-    # Convert the image to grayscale
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-    
-    # Resize image to 48x48 (model's expected input size)
-    image_resized = cv2.resize(image, (48, 48))
-    
-    # Add batch dimension and channels (for grayscale: 1 channel)
-    image_input = np.expand_dims(image_resized, axis=0)  # Add batch dimension (1, 48, 48)
-    image_input = np.expand_dims(image_input, axis=1)  # Add channel dimension (1, 1, 48, 48)
-    
-    # Normalize the image
-    image_input = image_input.astype(np.float32) / 255.0
-    
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # Convert to BGR (OpenCV format)
+    image_resized = cv2.resize(image, (224, 224))  # Resize to 224x224
+    image_input = np.expand_dims(image_resized, axis=0)  # Add batch dimension
+    image_input = image_input.transpose(0, 3, 1, 2)  # Change dimensions to (1, 3, 224, 224)
+    image_input = image_input.astype(np.float32) / 255.0  # Normalize the image
     return image_input
 
+# Map the raw output to emotions
+def get_emotion_from_output(output):
+    emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happiness', 'Sadness', 'Surprise', 'Neutral']
+    # Get the index of the highest value in the output (i.e., predicted emotion)
+    emotion_index = np.argmax(output)
+    confidence = output[0][emotion_index]  # Confidence of the prediction
+    emotion = emotion_labels[emotion_index]  # Corresponding emotion label
+    return emotion, confidence
+
 # Predict emotion using the ONNX model
-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})
-    
-    # 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]
+def predict_emotion_onnx(model, image_input):
+    # Get the input name and output name for the ONNX model
+    input_name = model.get_inputs()[0].name
+    output_name = model.get_outputs()[0].name
+    # Run the model
+    prediction = model.run([output_name], {input_name: image_input})
+    return prediction[0]
 
-# Streamlit interface
-st.title("Emotion Recognition with ONNX and AffectNet")
+# Streamlit UI
+st.title("Emotion Detection")
 
-# File uploader to upload images
+# Upload an image
 uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 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)
 
+    # Load model
+    onnx_model = load_model()
+
     # Preprocess the image
     image_input = preprocess_image(image)
 
-    # Predict the emotion
-    emotion_label, probability = predict_emotion_onnx(onnx_model, image_input)
+    # Get emotion prediction
+    emotion_prediction = predict_emotion_onnx(onnx_model, image_input)
 
-    # Display the predicted emotion and probability
-    st.write(f"Predicted Emotion: {emotion_label}")
-    st.write(f"Confidence: {probability:.2f}")
+    # Get the emotion label and confidence
+    emotion_label, confidence = get_emotion_from_output(emotion_prediction)
 
+    # Display the predicted emotion and confidence
+    st.write(f"Predicted Emotion: {emotion_label}")
+    st.write(f"Confidence: {confidence:.2f}")