Update app.py

app.py

@@ -4,9 +4,13 @@ from transformers import pipeline
 import pandas as pd
 import matplotlib.pyplot as plt
 
-# Initialize session state for results if not already present
+# Initialize session state for results, image names, and image sizes if not already present
 if 'results' not in st.session_state:
     st.session_state['results'] = []
+if 'image_names' not in st.session_state:
+    st.session_state['image_names'] = []
+if 'image_sizes' not in st.session_state:
+    st.session_state['image_sizes'] = []
 
 # Disable PyplotGlobalUseWarning
 st.set_option('deprecation.showPyplotGlobalUse', False)
@@ -20,12 +24,13 @@ st.title("Emotion Recognition with vit-face-expression")
 # Upload images
 uploaded_images = st.file_uploader("Upload images", type=["jpg", "png"], accept_multiple_files=True)
 
-# Store selected file names
-selected_file_names = []
-
 # Display thumbnail images alongside file names and sizes in the sidebar
 selected_images = []
 if uploaded_images:
+    # Reset the image names and sizes lists each time new images are uploaded
+    st.session_state['image_names'] = [img.name for img in uploaded_images]
+    st.session_state['image_sizes'] = [round(img.size / 1024.0, 1) for img in uploaded_images]
+
     # Add a "Select All" checkbox in the sidebar
     select_all = st.sidebar.checkbox("Select All", False)
     
@@ -38,57 +43,46 @@ if uploaded_images:
         
         if selected:
             selected_images.append(image)
-            selected_file_names.append(img.name)
 
 if st.button("Predict Emotions") and selected_images:
     # Predict emotion for each selected image using the pipeline
     st.session_state['results'] = [pipe(image) for image in selected_images]
-    # Display images and predicted emotions
-    for i, (image, result) in enumerate(zip(selected_images, st.session_state['results'])):
-        predicted_class = result[0]["label"]
-        predicted_emotion = predicted_class.split("_")[-1].capitalize()
-        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: {selected_file_names[i] if i < len(selected_file_names) else 'Unknown'}")
 
 # Generate DataFrame from results
 if st.button("Generate DataFrame"):
-    # Access the results from the session state
+    # Access the results, image names, and sizes from the session state
     results = st.session_state['results']
+    image_names = st.session_state['image_names']
+    image_sizes = st.session_state['image_sizes']
     if results:
-        st.write("results info inner loop:", results)
-        # Initialize an empty dictionary to store the scores
-        emotion_scores = {
-            'Neutral': [],
-            'Happy': [],
-            'Surprise': [],
-            'Disgust': [],
-            'Angry': [],
-            # Add other emotions if necessary
-            'Sad': [],  # Add this if you have 'sad' scores in your results
-            'Fear': []   # Add this if you have 'fear' scores in your results
-            
-        }
+        # Initialize an empty list to store all the data
+        data = []
 
-        # Iterate over the results and populate the dictionary
-        for result_set in results:
+        # Iterate over the results and populate the list with dictionaries
+        for i, result_set in enumerate(results):
             # Initialize a dictionary for the current set with zeros
-            current_scores = {emotion: 0 for emotion in emotion_scores.keys()}
+            current_data = {
+                'Neutral': 0,
+                'Happy': 0,
+                'Surprise': 0,
+                'Disgust': 0,
+                'Angry': 0,
+                # Add other emotions if necessary
+                'Image Name': image_names[i],
+                'Image Size (KB)': image_sizes[i]
+            }
             
             for result in result_set:
                 # Capitalize the label and update the score in the current set
                 emotion = result['label'].capitalize()
                 score = round(result['score'], 4)  # Round the score to 4 decimal places
-                current_scores[emotion] = score
+                current_data[emotion] = score
             
-            # Add the current scores to the emotion_scores dictionary
-            for emotion, score in current_scores.items():
-                emotion_scores[emotion].append(score)
+            # Append the current data to the data list
+            data.append(current_data)
 
-        # Convert the dictionary into a pandas DataFrame
-        df_emotions = pd.DataFrame(emotion_scores)
+        # Convert the list of dictionaries into a pandas DataFrame
+        df_emotions = pd.DataFrame(data)
 
         # Display the DataFrame
         st.write(df_emotions)