app.py

import streamlit as st
from PIL import Image
from transformers import pipeline
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import xlsxwriter  # Required for handling images in Excel

# Initialize session state for results, image names, and image sizes if not already present
if 'results' not in st.session_state:
    st.session_state['results'] = []
if 'image_names' not in st.session_state:
    st.session_state['image_names'] = []
if 'image_sizes' not in st.session_state:
    st.session_state['image_sizes'] = []

# Disable PyplotGlobalUseWarning
st.set_option('deprecation.showPyplotGlobalUse', False)

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

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

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

# Display thumbnail images alongside file names and sizes in the sidebar
selected_images = []
if uploaded_images:
    # Reset the image names and sizes lists each time new images are uploaded
    st.session_state['image_names'] = [img.name for img in uploaded_images]
    st.session_state['image_sizes'] = [round(img.size / 1024.0, 1) for img in uploaded_images]

    # Add a "Select All" checkbox in the sidebar
    select_all = st.sidebar.checkbox("Select All", False)
    
    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)         
        
        if selected:
            selected_images.append(image)

if st.button("Predict Emotions") and selected_images:
    # Predict emotion for each selected image using the pipeline
    st.session_state['results'] = [pipe(image) for image in selected_images]

# Generate DataFrame from results
if st.button("Generate HeatMap & DataFrame"):
    # Access the results, image names, and sizes from the session state
    results = st.session_state['results']
    image_names = st.session_state['image_names']
    image_sizes = st.session_state['image_sizes']
    if results:
        # Initialize an empty list to store all the data
        data = []

        # Iterate over the results and populate the list with dictionaries
        for i, result_set in enumerate(results):
            # Initialize a dictionary for the current set with zeros
            current_data = {
                'Happy': 0,
                'Surprise': 0,
                'Neutral': 0,
                'Sad': 0,
                'Disgust': 0,
                'Angry': 0,
                'Fear': 0,
                'Image Name': image_names[i],
                'Image Size (KB)': f"{image_sizes[i]:.1f}"  # Format the size to one decimal place
            }
            
            for result in result_set:
                # Capitalize the label and update the score in the current set
                emotion = result['label'].capitalize()
                score = round(result['score'], 4)  # Round the score to 4 decimal places
                current_data[emotion] = score
            
            # Append the current data to the data list
            data.append(current_data)

        # Convert the list of dictionaries into a pandas DataFrame
        df_emotions = pd.DataFrame(data)

        # Display the DataFrame
        st.write(df_emotions)

        # Plotting the heatmap for the first seven columns
        plt.figure(figsize=(10, 10))
        sns.heatmap(df_emotions.iloc[:, :7], annot=True, fmt=".1f", cmap='viridis')
        plt.title('Heatmap of Emotion Scores')
        plt.xlabel('Emotion Categories')
        plt.ylabel('Data Points')
        st.pyplot(plt)

        # Save the DataFrame to a CSV file
        df_emotions.to_csv('emotion_scores.csv', index=False)
        st.success('DataFrame generated and saved as emotion_scores.csv')

        with open('emotion_scores.csv', 'r') as f:
            csv_file = f.read()

        st.download_button(
            label='Download Emotion Scores as CSV',
            data=csv_file,
            file_name='emotion_scores.csv',
            mime='text/csv',
        )

# Save the DataFrame to an Excel file with images
if st.button("Save to Excel"):
    # Create a new Excel writer object
    writer = pd.ExcelWriter('emotion_scores.xlsx', engine='xlsxwriter')
    df_emotions.to_excel(writer, index=False)

    # Access the xlsxwriter workbook and worksheet objects
    workbook  = writer.book
    worksheet = writer.sheets['Sheet1']

    # Iterate over the images and insert them into the 'Image View' column
    for idx, image_path in enumerate(st.session_state['image_names']):
        # Your logic to define the image_path
        worksheet.insert_image(f'J{idx + 2}', image_path)  # 'J' is the 10th column

    # Close the writer and save the Excel file
    writer.save()
    st.session_state['excel_saved'] = True
    st.success('DataFrame generated and saved as emotion_scores.xlsx')

# Check if the Excel file was saved and provide a download button
if st.session_state.get('excel_saved', False):
    with open('emotion_scores.xlsx', 'rb') as f:
        excel_file = f.read()

    st.download_button(
        label='Download Emotion Scores as Excel',
        data=excel_file,
        file_name='emotion_scores.xlsx',
        mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
    )