Spaces:
Sleeping
Sleeping
| 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) | |
| # Streamlit app | |
| st.title("Emotion Recognition with vit-face-expression") | |
| # Slider example | |
| #x = st.slider('Select a value') | |
| #st.write(f"{x} squared is {x * x}") | |
| # 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: | |
| 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=False, key=checkbox_key) | |
| if selected: | |
| 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] | |
| # Display images and predicted emotions side by side | |
| 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}") | |
| # Use the index to get the corresponding filename | |
| col.write(f"Original File Name: {uploaded_images[i].name}") | |
| # Display the keys and values of all results | |
| 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: | |
| # Predict emotion for each selected image using the pipeline | |
| results = [pipe(image) for image in selected_images] | |
| # Display images and predicted emotions | |
| 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() | |