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

# Set Streamlit configuration to disable deprecation warnings
st.set_option('deprecation.showPyplotGlobalUse', False)

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

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

# Create a file uploader to upload images in JPG or PNG format
uploaded_images = st.file_uploader("Upload images", type=["jpg", "png"], accept_multiple_files=True)

# List to store selected file names
selected_file_names = []

# List to store selected images
selected_images = []

# Process uploaded images if any
if uploaded_images:
    # Add a "Select All" checkbox in the sidebar for convenience
    select_all = st.sidebar.checkbox("Select All", False)
    
    # Iterate over each uploaded image
    for idx, img in enumerate(uploaded_images):
        image = Image.open(img)
        checkbox_key = f"{img.name}_checkbox_{idx}"  # Unique key for each checkbox
        
        # Display thumbnail image and checkbox in sidebar
        st.sidebar.image(image, caption=f"{img.name} ({img.size / 1024.0:.1f} KB)", width=40)
        selected = st.sidebar.checkbox(f"Select {img.name}", value=select_all, key=checkbox_key)
        
        # Add selected images to the list
        if selected:
            selected_images.append(image)
            selected_file_names.append(img.name)

# Button to start emotion prediction
if st.button("Predict Emotions") and selected_images:
    # Predict emotion for each selected image using the pipeline
    results = [pipe(image) for image in selected_images]
    emotions = [result[0]["label"].split("_")[-1].capitalize() for result in results]

    # Display images and predicted emotions
    for i, (image, result) in enumerate(zip(selected_images, results)):
        st.image(image, caption=f"Predicted emotion: {emotions[i]}", use_column_width=True)
        st.write(f"Emotion Scores for Image #{i+1}")
        st.write(f"{emotions[i]}: {result[0]['score']:.4f}")
        st.write(f"Original File Name: {selected_file_names[i]}")

    # Calculate emotion statistics
    emotion_counts = pd.Series(emotions).value_counts()
    total_faces = len(selected_images)

    # Define a color map for emotions
    color_map = {
        'Neutral': '#B38B6D',
        'Happy': '#FFFF00',
        'Sad': '#0000FF',
        'Angry': '#FF0000',
        'Disgust': '#008000',
        'Surprise': '#FFA500',
        'Fear': '#000000'
    }

    # Plot pie chart for emotion distribution
    st.write("Emotion Distribution (Pie Chart):")
    fig_pie, ax_pie = plt.subplots()
    pie_colors = [color_map.get(emotion, '#999999') for emotion in emotion_counts.index]
    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(f"Total Faces Analyzed: {total_faces}")
    st.pyplot(fig_pie)

    # Plot bar chart for emotion distribution
    st.write("Emotion Distribution (Bar Chart):")
    fig_bar, ax_bar = plt.subplots()
    bar_colors = [color_map.get(emotion, '#999999') for emotion in emotion_counts.index]
    emotion_counts.plot(kind='bar', color=bar_colors, ax=ax_bar)
    ax_bar.set_xlabel('Emotion')
    ax_bar.set_ylabel('Count')
    ax_bar.set_title(f"Emotion Distribution - Total Faces Analyzed: {total_faces}")
    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, int(i.get_height()), ha='center', va='bottom')
    st.pyplot(fig_bar)