import streamlit as st
from PIL import Image
from transformers import pipeline
import pandas as pd
import matplotlib.pyplot as plt

# Configuration to suppress warnings in Streamlit
st.set_option('deprecation.showPyplotGlobalUse', False)

# Initialize the image classification pipeline with a specific model
pipe = pipeline("image-classification", model="trpakov/vit-face-expression")

# Setting the title of the Streamlit application
st.title("Emotion Recognition App by Prateek Mohan")

# Interface for uploading an image
uploaded_image = st.file_uploader("Please upload an image", type=["jpg", "png"], accept_multiple_files=False)
color_map = {
        'Neutral': '#B38B6D',  
        'Happy': '#FFFF00',    
        'Sad': '#0000FF',     
        'Angry': '#FF0000',    
        'Disgust': '#008000',  
        'Surprise': '#FFA500', 
        'Fear': '#000000'      
       
    }
# Process the image immediately after upload
if uploaded_image:
    # Load and display the image
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Emotion prediction using the pre-trained model
    result = pipe(image)
    predicted_class = result[0]["label"]
    predicted_emotion = predicted_class.split("_")[-1].capitalize()
    emotion_score = result[0]["score"]

    # Displaying the predicted emotion and its confidence score
    st.write(f"Predicted Emotion: {predicted_emotion}")
    st.write(f"Confidence Score: {emotion_score:.4f}")

    # Data preparation for visualization
    emotion_counts = pd.Series([predicted_emotion])

    # Color mapping for different emotions
    color_map = {
        'Neutral': '#B38B6D',
        'Happy': '#FFFF00',
        'Sad': '#0000FF',
        'Angry': '#FF0000',
        'Disgust': '#008000',
        'Surprise': '#FFA500',
        'Fear': '#000000'
    }
    # pie_colors = [color_map.get(emotion, '#999999') for emotion in emotion_counts.index]  
    # # Pie chart visualization
    # fig_pie, ax_pie = plt.subplots(figsize=(4, 4))
    # ax_pie.pie(emotion_counts, labels=emotion_counts.index, autopct='%1.1f%%', startangle=140, colors=pie_colors)
    # ax_pie.axis('equal')
    # ax_pie.set_title("Emotion Distribution")
    # st.pyplot(fig_pie)

    # # Bar chart visualization
    # fig_bar, ax_bar = plt.subplots(figsize=(4, 4))
    # emotion_counts.plot(kind='bar', color=pie_colors, ax=ax_bar)
    # ax_bar.set_xlabel('Emotion')
    # ax_bar.set_ylabel('Count')
    # ax_bar.set_title("Emotion Count")
    # ax_bar.yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    # for i in ax_bar.patches:
    #     ax_bar.text(i.get_x() + i.get_width() / 2, i.get_height() + 0.1, str(int(i.get_height())), ha='center', va='bottom')
    # st.pyplot(fig_bar)