Update faceapp.py

faceapp.py

@@ -1,11 +1,12 @@
 import streamlit as st
 from PIL import Image
-from transformers import ViTForImageClassification, ViTImageProcessor
+from transformers import pipeline
 
-# Load the model
-model_name = "trpakov/vit-face-expression"
-model = ViTForImageClassification.from_pretrained(model_name)
-image_processor = ViTImageProcessor.from_pretrained(model_name)
+# 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")
@@ -14,20 +15,42 @@ st.title("Emotion Recognition with vit-face-expression")
 x = st.slider('Select a value')
 st.write(f"{x} squared is {x * x}")
 
-# Upload image
-uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png"])
-
-if uploaded_image:
-    image = Image.open(uploaded_image)
-    inputs = image_processor(images=image, return_tensors="pt")
-    pixel_values = inputs.pixel_values
-
-    # Predict emotion
-    with torch.no_grad():
-        outputs = model(pixel_values)
-        predicted_class = torch.argmax(outputs.logits, dim=1).item()
-
-    emotion_labels = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
-    predicted_emotion = emotion_labels[predicted_class]
-
-    st.image(image, caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True)
+# 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}")