Update app.py

app.py

@@ -31,7 +31,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(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 ({data_percentage}% data, {task_complexity} beams)", save_path=raw_cm_path)
         raw_img = Image.open(raw_cm_path)
     else:
         raw_img = None
@@ -40,14 +40,28 @@ 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(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 ({data_percentage}% data, {task_complexity} beams)", save_path=embeddings_cm_path)
         embeddings_img = Image.open(embeddings_cm_path)
     else:
         embeddings_img = None
 
     return raw_img, embeddings_img
 
+def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
+    plt.figure(figsize=(8, 6))
+    plt.imshow(cm, interpolation='nearest', cmap='coolwarm')
+    plt.title(title)
+    plt.colorbar()
+    tick_marks = np.arange(len(classes))
+    plt.xticks(tick_marks, classes, rotation=45)
+    plt.yticks(tick_marks, classes)
 
+    plt.tight_layout()
+    plt.ylabel('True label')
+    plt.xlabel('Predicted label')
+    plt.savefig(save_path)
+    plt.close()
+    
 # Function to compute the average confusion matrix across CSV files in a folder
 def compute_average_confusion_matrix(folder):
     confusion_matrices = []