facefeature / app.py
xtlyxt's picture
Update app.py
ead6b33 verified
raw
history blame
5.73 kB
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',
)