import onnxruntime as ort
import numpy as np
import cv2
from PIL import Image
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")

# 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 for ONNX model
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
    
    return image_input

# 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]

# Streamlit interface
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:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Preprocess the image
    image_input = preprocess_image(image)

    # Predict the emotion
    emotion_label, probability = 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}")