Update app.py

app.py

@@ -34,7 +34,7 @@ def beam_prediction_task(data_percentage, task_complexity):
     embeddings_cm = compute_average_confusion_matrix(embeddings_folder)
     if embeddings_cm is not None:
         embeddings_cm_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
-        plot_confusion_matrix_beamPred(embeddings_cm, classes=np.arange(embeddings_cm.shape[0]), title=f"Embeddings Confusion Matrix ({data_percentage}% data, {task_complexity} beams)", save_path=embeddings_cm_path)
+        plot_confusion_matrix_beamPred(embeddings_cm, classes=np.arange(embeddings_cm.shape[0]), title=f"Embeddings Confusion Matrix\n({data_percentage}% data, {task_complexity} beams)", save_path=embeddings_cm_path)
         embeddings_img = Image.open(embeddings_cm_path)
     else:
         embeddings_img = None
@@ -71,7 +71,7 @@ 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, cmap="magma", cbar=False, linecolor='white')
+    sns.heatmap(cm, cmap="magma", cbar=True, linecolor='white')
     
     # Add F1-score to the title
     plt.title(f"{title}\n(F1 Score: {avg_f1:.3f})", color="white", fontsize=14)
@@ -316,34 +316,31 @@ def plot_confusion_matrix(y_true, y_pred, title):
     
     # Calculate F1 Score
     f1 = f1_score(y_true, y_pred, average='weighted')
-    
-    # Plot the confusion matrix with dark mode colors
+
+    plt.style.use('dark_background')
     plt.figure(figsize=(5, 5))
-    plt.imshow(cm, interpolation='nearest', cmap='coolwarm')  # Dark mode color scheme
-    plt.title(f"{title}\nF1-Score: {f1:.2f}", color='white', pad=20)  # Add padding for the title to prevent clipping
-    plt.colorbar()
-    plt.xticks([0, 1], labels=['Class 0', 'Class 1'], color='white')  # White text for dark mode
-    plt.yticks([0, 1], labels=['Class 0', 'Class 1'], color='white')  # White text for dark mode
-    
-    # Annotate the confusion matrix
-    thresh = cm.max() / 2
-    for i in range(cm.shape[0]):
-        for j in range(cm.shape[1]):
-            plt.text(j, i, format(cm[i, j], 'd'),
-                     ha="center", va="center",
-                     color="white" if cm[i, j] > thresh else "black")
-
-    plt.ylabel('True label', color='white')  # White text for dark mode
-    plt.xlabel('Predicted label', color='white')  # White text for dark mode
-    plt.tight_layout(pad=2.0)  # Add padding to prevent clipping
-    
-    # Save the plot
-    plt.savefig(f"{title}.png", facecolor='black')  # Set background to black for dark mode
+    
+    # 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')
+    
+    # Add F1-score to the title
+    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)
+    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(f"{title}.png", transparent=True)  # Use transparent to blend with the dark mode website
     plt.close()
     
+    # Return the saved image
     return Image.open(f"{title}.png")
-
-
+    
 def identical_train_test_split(output_emb, output_raw, labels, percentage):
     N = output_emb.shape[0]  # Get the total number of samples