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 import io # 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] # Initialize an empty DataFrame outside of the button press condition df_emotions = pd.DataFrame() # 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) # Add a placeholder for the 'Image View' column df_emotions['Image View'] = [''] * len(df_emotions) # Add a sequence of numbers for the 'Image Num' column df_emotions['Image Num'] = list(range(len(df_emotions))) # 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 without the 'Image View' and 'Image Num' columns df_emotions.drop(columns=['Image View', 'Image Num']).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', ) # Create a BytesIO buffer for the Excel file output = io.BytesIO() # Create a new Excel writer object using the buffer as the file writer = pd.ExcelWriter(output, engine='xlsxwriter') df_emotions.to_excel(writer, index=False, header=True) # Access the xlsxwriter workbook and worksheet objects workbook = writer.book worksheet = writer.sheets['Sheet1'] # Set the column width and row height worksheet.set_column('A:G', 8) # Set width for columns A-G worksheet.set_column('H:H', 22) # Set width for column H (Image Name) worksheet.set_column('I:I', 14) # Set width for column I (Image Size) worksheet.set_column('J:J', 12) # Set width for column J (Image View) worksheet.set_column('K:K', 12) # Set width for column K (Image Num) for row_num in range(len(df_emotions) + 1): # +1 to include the header row worksheet.set_row(row_num, 52) # Set the row height to 38 # Iterate over the images and insert them into the 'Image View' column for idx, image in enumerate(selected_images): # Convert the image to a format that can be inserted into Excel image_stream = io.BytesIO() image.save(image_stream, format='JPEG') # image.save(image_stream, format='JPEG') # Save the image as JPEG; or PNG image_stream.seek(0) # Calculate the scaling factor to fit the image inside the cell cell_width = 64 scale_factor = cell_width / image.width # Insert the image into the cell worksheet.insert_image(f'J{idx + 2}', 'image.jpg', { #or image.png 'image_data': image_stream, 'x_scale': scale_factor, 'y_scale': scale_factor, 'x_offset': 2, 'y_offset': 2, 'positioning': 1 }) # Close the writer object writer.close() # Rewind the buffer output.seek(0) # Use Streamlit's download button to offer the Excel file for download st.download_button( label='Download Emotion Scores as Excel', data=output, file_name='emotion_scores.xlsx', mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', )