import streamlit as st
from PIL import Image
from transformers import pipeline

# Create an image classification pipeline with scores
pipe = pipeline("image-classification", model="trpakov/vit-face-expression", top_k=None)

# Define emotion labels
emotion_labels = ["Neutral", "Sad", "Angry", "Surprised", "Happy"]

# Streamlit app
st.title("Emotion Recognition with vit-face-expression")

# Slider example
x = st.slider('Select a value')
st.write(f"{x} squared is {x * x}")

# Upload images
uploaded_images = st.file_uploader("Upload images", type=["jpg", "png"], accept_multiple_files=True)

if st.button("Predict Emotions") and uploaded_images:
    if len(uploaded_images) == 2:
        # Open the uploaded images
        images = [Image.open(img) for img in uploaded_images]

        # Predict emotion for each image using the pipeline
        results = [pipe(image) for image in images]

        # Display images and predicted emotions side by side
        col1, col2 = st.columns(2)
        for i in range(2):
            predicted_class = results[i][0]["label"]
            predicted_emotion = predicted_class.split("_")[-1].capitalize()
            col = col1 if i == 0 else col2
            col.image(images[i], caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True)
            col.write(f"Emotion Scores for {predicted_emotion}: {results[i][0]['score']:.4f}")
            
            # Display scores for other categories
            st.write(f"Emotion Scores for other categories (Image {i+1}):")
            for label, score in zip(emotion_labels, results[i][0]["score"]):
                if label.lower() != predicted_emotion.lower():  # Exclude the predicted emotion
                    st.write(f"{label}: {score:.4f}")
    else:
        # Open the uploaded images
        images = [Image.open(img) for img in uploaded_images]

        # Predict emotion for each image using the pipeline
        results = [pipe(image) for image in images]

        # Display images and predicted emotions
        for i, result in enumerate(results):
            predicted_class = result[0]["label"]
            predicted_emotion = predicted_class.split("_")[-1].capitalize()
            st.image(images[i], caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True)
            st.write(f"Emotion Scores for Image {i+1}:")
            st.write(f"{predicted_emotion}: {result[0]['score']:.4f}")