Update app.py

app.py

@@ -65,27 +65,32 @@ def compute_f1_score(cm):
 def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     # Compute the average F1-score
     avg_f1 = compute_f1_score(cm)
+
+    # Set dark mode styling
+    plt.style.use('dark_background')
+    plt.figure(figsize=(5, 5))
     
-    # Update title to include average F1-score
-    full_title = f"{title} (Avg F1-Score: {avg_f1:.2f})"
+    # Plot the confusion matrix with a dark-mode compatible colormap
+    sns.heatmap(cm, annot=True, fmt="d", cmap="magma", cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
     
-    # Plot the confusion matrix with dark mode adjustments
-    plt.figure(figsize=(8, 6))
-    plt.imshow(cm, interpolation='nearest', cmap='coolwarm')  # Dark mode color scheme
-    plt.title(full_title, color='white', pad=20)  # Add padding to prevent title clipping, white text for dark mode
-    plt.colorbar()
+    # Add F1-score to the title
+    plt.title(f"{title}\n(F1 Score: {avg_f1:.3f})", color="white", fontsize=14)
     
-    tick_marks = np.arange(len(classes))
-    plt.xticks(tick_marks, classes, rotation=45, color='white')  # White text for dark mode
-    plt.yticks(tick_marks, classes, color='white')  # White text for dark mode
-
-    plt.tight_layout(pad=2.0)  # Add padding to prevent axis label clipping
-    plt.ylabel('True label', color='white')  # White text for dark mode
-    plt.xlabel('Predicted label', color='white')  # White text for dark mode
-
-    # Save the plot with a black background for dark mode
-    plt.savefig(save_path, facecolor='black')
+    # Customize tick labels for dark mode
+    plt.xticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color="white", fontsize=10)
+    plt.yticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color="white", fontsize=10)
+    
+    plt.ylabel('True label', color="white", fontsize=12)
+    plt.xlabel('Predicted label', color="white", fontsize=12)
+    plt.tight_layout()
+    
+    # Save the plot as an image
+    plt.savefig(save_path, transparent=True)  # Use transparent to blend with the dark mode website
     plt.close()
+    
+    # Return the saved image
+    return Image.open(save_path)
+
 
 def compute_average_confusion_matrix(folder):
     confusion_matrices = []
@@ -158,7 +163,7 @@ def plot_confusion_matrix_from_csv(csv_file_path, title, save_path):
     sns.heatmap(cm, annot=True, fmt="d", cmap="magma", cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
     
     # Add F1-score to the title
-    plt.title(f"{title} (F1 Score: {f1:.3f})", color="white", fontsize=14)
+    plt.title(f"{title}\n(F1 Score: {f1:.3f})", color="white", fontsize=14)
     
     # Customize tick labels for dark mode
     plt.xticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color="white", fontsize=10)