mprateek's picture
Update app.py
494a291 verified
raw
history blame contribute delete
No virus
2.76 kB
import streamlit as st
from PIL import Image
from transformers import pipeline
import pandas as pd
import matplotlib.pyplot as plt
# Configuration to suppress warnings in Streamlit
st.set_option('deprecation.showPyplotGlobalUse', False)
# Initialize the image classification pipeline with a specific model
pipe = pipeline("image-classification", model="trpakov/vit-face-expression")
# Setting the title of the Streamlit application
st.title("Emotion Recognition App by Prateek Mohan")
# Interface for uploading an image
uploaded_image = st.file_uploader("Please upload an image", type=["jpg", "png"], accept_multiple_files=False)
color_map = {
'Neutral': '#B38B6D',
'Happy': '#FFFF00',
'Sad': '#0000FF',
'Angry': '#FF0000',
'Disgust': '#008000',
'Surprise': '#FFA500',
'Fear': '#000000'
}
# Process the image immediately after upload
if uploaded_image:
# Load and display the image
image = Image.open(uploaded_image)
st.image(image, caption="Uploaded Image", use_column_width=True)
# Emotion prediction using the pre-trained model
result = pipe(image)
predicted_class = result[0]["label"]
predicted_emotion = predicted_class.split("_")[-1].capitalize()
emotion_score = result[0]["score"]
# Displaying the predicted emotion and its confidence score
st.write(f"Predicted Emotion: {predicted_emotion}")
st.write(f"Confidence Score: {emotion_score:.4f}")
# Data preparation for visualization
emotion_counts = pd.Series([predicted_emotion])
# Color mapping for different emotions
color_map = {
'Neutral': '#B38B6D',
'Happy': '#FFFF00',
'Sad': '#0000FF',
'Angry': '#FF0000',
'Disgust': '#008000',
'Surprise': '#FFA500',
'Fear': '#000000'
}
# pie_colors = [color_map.get(emotion, '#999999') for emotion in emotion_counts.index]
# # Pie chart visualization
# fig_pie, ax_pie = plt.subplots(figsize=(4, 4))
# 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("Emotion Distribution")
# st.pyplot(fig_pie)
# # Bar chart visualization
# fig_bar, ax_bar = plt.subplots(figsize=(4, 4))
# emotion_counts.plot(kind='bar', color=pie_colors, ax=ax_bar)
# ax_bar.set_xlabel('Emotion')
# ax_bar.set_ylabel('Count')
# ax_bar.set_title("Emotion Count")
# 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, str(int(i.get_height())), ha='center', va='bottom')
# st.pyplot(fig_bar)