Update app.py

app.py

@@ -155,31 +155,37 @@ st.write("selected_images info:", selected_images)
 st.write("selected_file_names info:", selected_file_names)
 st.write("results info:", results)
         
+# Generate DataFrame from results
+if st.button("Generate DataFrame") and results:
+    # Initialize an empty dictionary to store the scores
+    emotion_scores = {
+        'Neutral': [],
+        'Happy': [],
+        'Surprise': [],
+        'Disgust': [],
+        'Angry': [],
+        'Sad': [],  # Add this if you have 'sad' scores in your results
+        'Fear': []   # Add this if you have 'fear' scores in your results
+    }
 
-# Generate DataFrame button
-if st.button("Generate DataFrame") and selected_images:
-    # Create a list to store data for DataFrame
-    df_data = []
-
-    # Iterate through selected images to gather data
-    for result in results:
-        # Extract predicted emotions and scores
-        emotion_scores = {res["label"].split("_")[-1].capitalize(): res["score"] for res in result}
-
-        # Append emotion scores to the list
-        df_data.append({
-            "Neutral": f"{emotion_scores.get('Neutral', 0.0):.4f}",
-            "Happy": f"{emotion_scores.get('Happy', 0.0):.4f}",
-            "Sad": f"{emotion_scores.get('Sad', 0.0):.4f}",
-            "Angry": f"{emotion_scores.get('Angry', 0.0):.4f}",
-            "Disgust": f"{emotion_scores.get('Disgust', 0.0):.4f}",
-            "Surprise": f"{emotion_scores.get('Surprise', 0.0):.4f}",
-            "Fear": f"{emotion_scores.get('Fear', 0.0):.4f}"  # Add this line if 'Fear' is a possible label
-        })
-
-    # Create DataFrame
-    df = pd.DataFrame(df_data)
-
-    # Display DataFrame
-    st.write(df)
-
+    # Iterate over the results and populate the dictionary
+    for result_set in results:
+        for result in result_set:
+            # Capitalize the label and add the score to the corresponding list
+            emotion = result['label'].capitalize()
+            score = result['score']
+            if emotion in emotion_scores:
+                emotion_scores[emotion].append(score)
+            else:
+                # If the emotion is not in the dictionary, initialize a new list
+                emotion_scores[emotion] = [score]
+
+    # Convert the dictionary into a pandas DataFrame
+    df_emotions = pd.DataFrame(emotion_scores)
+
+    # Display the DataFrame
+    st.write(df_emotions)
+
+    # Optional: Save the DataFrame to a CSV file
+    df_emotions.to_csv('emotion_scores.csv', index=False)
+    st.success('DataFrame generated and saved as emotion_scores.csv')