Update app.py

app.py

@@ -3,19 +3,24 @@ import numpy as np
 import cv2
 from PIL import Image
 import streamlit as st
+import kagglehub  # Import kagglehub to load the AffectNet dataset
 
-# Load ONNX model
-onnx_model = ort.InferenceSession("onnx_model.onnx")
+# Download the AffectNet dataset
+path = kagglehub.dataset_download("fatihkgg/affectnet-yolo-format")
+print("Path to AffectNet dataset:", path)
 
-# Emotion labels (same as the model's output classes)
+# Emotion labels for AffectNet
 emotion_labels = ["Anger", "Disgust", "Fear", "Happy", "Sadness", "Surprise", "Neutral"]
 
+# Load ONNX model
+onnx_model = ort.InferenceSession("onnx_model.onnx")
+
 # 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
+# Preprocess image function for ONNX model
 def preprocess_image(image):
     """Preprocess image to match model input requirements"""
     # Convert the image to grayscale
@@ -48,8 +53,9 @@ def predict_emotion_onnx(onnx_model, image_input):
     return emotion_labels[predicted_class], probabilities[predicted_class]
 
 # Streamlit interface
-st.title("Emotion Recognition with ONNX")
+st.title("Emotion Recognition with ONNX and AffectNet")
 
+# File uploader to upload images
 uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
@@ -65,3 +71,4 @@ if uploaded_file is not None:
     # Display the predicted emotion and probability
     st.write(f"Predicted Emotion: {emotion_label}")
     st.write(f"Confidence: {probability:.2f}")
+