Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PIL import Image | |
| from transformers import pipeline | |
| # Create an image classification pipeline with scores | |
| pipe = pipeline("image-classification", model="trpakov/vit-face-expression", top_k=None) | |
| # Define emotion labels | |
| emotion_labels = ["Neutral", "Sad", "Angry", "Surprised", "Happy"] | |
| # 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) | |
| if st.button("Predict Emotions") and uploaded_images: | |
| if len(uploaded_images) == 2: | |
| # Open the uploaded images | |
| images = [Image.open(img) for img in uploaded_images] | |
| # Predict emotion for each image using the pipeline | |
| results = [pipe(image) for image in 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() | |
| col = col1 if i == 0 else col2 | |
| col.image(images[i], caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True) | |
| col.write(f"Emotion Scores for {predicted_emotion}: {results[i][0]['score']:.4f}") | |
| # Display scores for other categories | |
| st.write(f"Emotion Scores for other categories (Image {i+1}):") | |
| for label, score in zip(emotion_labels, results[i][0]["score"]): | |
| if label.lower() != predicted_emotion.lower(): # Exclude the predicted emotion | |
| st.write(f"{label}: {score:.4f}") | |
| else: | |
| # Open the uploaded images | |
| images = [Image.open(img) for img in uploaded_images] | |
| # Predict emotion for each image using the pipeline | |
| results = [pipe(image) for image in images] | |
| # Display images and predicted emotions | |
| for i, result in enumerate(results): | |
| predicted_class = result[0]["label"] | |
| predicted_emotion = predicted_class.split("_")[-1].capitalize() | |
| st.image(images[i], caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True) | |
| st.write(f"Emotion Scores for Image {i+1}:") | |
| st.write(f"{predicted_emotion}: {result[0]['score']:.4f}") | |