Update app.py

app.py

@@ -25,7 +25,7 @@ def beam_prediction_task(data_percentage, task_complexity):
     raw_cm = compute_average_confusion_matrix(raw_folder)
     if raw_cm is not None:
         raw_cm_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
-        plot_confusion_matrix_beamPred(raw_cm, classes=np.arange(raw_cm.shape[0]), title=f"Raw Confusion Matrix ({data_percentage}% data, {task_complexity} beams)", save_path=raw_cm_path)
+        plot_confusion_matrix_beamPred(raw_cm, classes=np.arange(raw_cm.shape[0]), title=f"Raw Confusion Matrix\n({data_percentage}% data, {task_complexity} beams)", save_path=raw_cm_path)
         raw_img = Image.open(raw_cm_path)
     else:
         raw_img = None
@@ -68,7 +68,7 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
 
     # Set dark mode styling
     plt.style.use('dark_background')
-    plt.figure(figsize=(5, 5))
+    plt.figure(figsize=(7, 7), facecolor='#2f2f2f')
     
     # Plot the confusion matrix with a dark-mode compatible colormap
     sns.heatmap(cm, cmap="magma", cbar=True, linecolor='white')
@@ -80,20 +80,17 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     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)
     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.savefig(save_path, facecolor='#2f2f2f', 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 = []
     max_num_labels = 0