Update app.py

app.py

@@ -5,9 +5,6 @@ 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")
 
@@ -19,42 +16,27 @@ st.write(f"{x} squared is {x * x}")
 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 keys and values of all results for the current image
-            st.write(f"Keys and Values of results[{i}] (Image {i+1}):")
-            for key, value in results[i][0].items():
-                st.write(f"Key: {key}, Value: {value}")
-    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}")
-
-            # Display keys and values of all results for the current image
-            st.write(f"Keys and Values of results[{i}] (Image {i+1}):")
-            for key, value in result[0].items():
-                st.write(f"Key: {key}, Value: {value}")
+    # 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]
+
+    # Determine display style based on the number of images
+    num_images = len(uploaded_images)
+    display_style = "column" if num_images > 2 else "row"
+
+    # 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)
+
+        # Display scores for all categories
+        st.write(f"Emotion Scores for Image {i+1}:")
+        for emotion, score in result[0]["score"].items():
+            st.write(f"{emotion.capitalize()}: {score:.4f}")
+
+        # If only two images and not the last image, display side by side
+        if num_images == 2 and i == 0:
+            st.write("----")  # Divider between images if side by side