Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from PIL import Image
 from transformers import pipeline
+import pandas as pd
+import matplotlib.pyplot as plt
 
 # Create an image classification pipeline with scores
 pipe = pipeline("image-classification", model="trpakov/vit-face-expression", top_k=None)
@@ -28,6 +30,7 @@ if uploaded_images:
             selected_images.append(image)
 
 if st.button("Predict Emotions") and selected_images:
+    emotions = []
     if len(selected_images) == 2:
         # Predict emotion for each selected image using the pipeline
         results = [pipe(image) for image in selected_images]
@@ -37,6 +40,7 @@ if st.button("Predict Emotions") and selected_images:
         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}")
@@ -61,8 +65,28 @@ if st.button("Predict Emotions") and 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}")
             # Use the index to get the corresponding filename
             st.write(f"Original File Name: {uploaded_images[i].name if i < len(uploaded_images) else 'Unknown'}")
+
+    # Calculate emotion statistics
+    emotion_counts = pd.Series(emotions).value_counts()
+
+    # Plot pie chart
+    st.write("Emotion Distribution (Pie Chart):")
+    plt.figure(figsize=(8, 6))
+    plt.pie(emotion_counts, labels=emotion_counts.index, autopct='%1.1f%%', startangle=140)
+    plt.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
+    st.pyplot()
+
+    # Plot bar chart
+    st.write("Emotion Distribution (Bar Chart):")
+    plt.figure(figsize=(10, 6))
+    emotion_counts.plot(kind='bar', color='skyblue')
+    plt.xlabel('Emotion')
+    plt.ylabel('Count')
+    plt.title('Emotion Distribution')
+    st.pyplot()