Update app.py

app.py

@@ -71,14 +71,15 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     plt.figure(figsize=(5, 5))
     
     # 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')
+    sns.heatmap(cm, annot=True, fmt=".2f", cmap="magma", cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
     
     # Add F1-score to the title
     plt.title(f"{title}\n(F1 Score: {avg_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)
-    plt.yticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color="white", fontsize=10)
+    tick_marks = np.arange(len(classes))
+    plt.xticks(tick_marks, classes, color="white", fontsize=10)  # White text for dark mode
+    plt.yticks(tick_marks, classes, color="white", fontsize=10)  # White text for dark mode
+
     
     plt.ylabel('True label', color="white", fontsize=12)
     plt.xlabel('Predicted label', color="white", fontsize=12)
@@ -92,6 +93,7 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     return Image.open(save_path)
 
 
+
 def compute_average_confusion_matrix(folder):
     confusion_matrices = []
     max_num_labels = 0