|
import streamlit as st |
|
from PIL import Image |
|
from transformers import pipeline |
|
import pandas as pd |
|
import matplotlib.pyplot as plt |
|
|
|
|
|
st.set_option('deprecation.showPyplotGlobalUse', False) |
|
|
|
|
|
pipe = pipeline("image-classification", model="trpakov/vit-face-expression", top_k=None) |
|
|
|
|
|
st.title("Emotion Recognition with vit-face-expression") |
|
|
|
|
|
uploaded_images = st.file_uploader("Upload images", type=["jpg", "png"], accept_multiple_files=True) |
|
|
|
|
|
selected_file_names = [] |
|
|
|
|
|
selected_images = [] |
|
if uploaded_images: |
|
|
|
|
|
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}" |
|
|
|
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) |
|
selected_file_names.append(img.name) |
|
|
|
if st.button("Predict Emotions") and selected_images: |
|
emotions = [] |
|
if len(selected_images) == 2: |
|
|
|
results = [pipe(image) for image in selected_images] |
|
|
|
|
|
col1, col2 = st.columns(2) |
|
for i in range(2): |
|
predicted_class = results[i][0]["label"] |
|
predicted_emotion = predicted_class.split("_")[-1].capitalize() |
|
emotions.append(predicted_emotion) |
|
col = col1 if i == 0 else col2 |
|
col.image(selected_images[i], caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True) |
|
col.write(f"Emotion Scores: {predicted_emotion}: {results[i][0]['score']:.4f}") |
|
|
|
col.write(f"Original File Name: {selected_file_names[i]}") |
|
|
|
|
|
st.write("Keys and Values of all results:") |
|
col1, col2 = st.columns(2) |
|
for i, result in enumerate(results): |
|
col = col1 if i == 0 else col2 |
|
col.write(f"Keys and Values of results[{i}]:") |
|
for res in result: |
|
label = res["label"] |
|
score = res["score"] |
|
col.write(f"{label}: {score:.4f}") |
|
else: |
|
|
|
results = [pipe(image) for image in selected_images] |
|
|
|
|
|
for i, (image, result) in enumerate(zip(selected_images, results)): |
|
predicted_class = result[0]["label"] |
|
predicted_emotion = predicted_class.split("_")[-1].capitalize() |
|
emotions.append(predicted_emotion) |
|
st.image(image, caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True) |
|
st.write(f"Emotion Scores for #{i+1} Image") |
|
st.write(f"{predicted_emotion}: {result[0]['score']:.4f}") |
|
|
|
st.write(f"Original File Name: {selected_file_names[i] if i < len(selected_file_names) else 'Unknown'}") |
|
|
|
|
|
emotion_counts = pd.Series(emotions).value_counts() |
|
|
|
|
|
color_map = { |
|
'Neutral': '#B38B6D', |
|
'Happy': '#FFFF00', |
|
'Sad': '#0000FF', |
|
'Angry': '#FF0000', |
|
'Disgust': '#008000', |
|
'Surprise': '#FFA500', |
|
'Fear': '#000000' |
|
|
|
} |
|
|
|
|
|
total_faces = len(selected_images) |
|
|
|
|
|
pie_colors = [color_map.get(emotion, '#999999') for emotion in emotion_counts.index] |
|
|
|
|
|
st.write("Emotion Distribution (Pie Chart):") |
|
fig_pie, ax_pie = plt.subplots() |
|
|
|
ax_pie.pie(emotion_counts, labels=emotion_counts.index, autopct='%1.1f%%', startangle=140, colors=pie_colors, textprops={'color': 'white', 'weight': 'bold'}) |
|
|
|
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) |
|
|
|
|
|
bar_colors = [color_map.get(emotion, '#999999') for emotion in emotion_counts.index] |
|
|
|
|
|
st.write("Emotion Distribution (Bar Chart):") |
|
fig_bar, ax_bar = plt.subplots() |
|
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) |
|
|