Update app.py

app.py

@@ -47,11 +47,41 @@ def beam_prediction_task(data_percentage, task_complexity):
 
     return raw_img, embeddings_img
 
+from sklearn.metrics import f1_score
+
+# Function to compute the F1-score based on the confusion matrix
+def compute_f1_score(cm):
+    # Compute precision and recall
+    TP = np.diag(cm)
+    FP = np.sum(cm, axis=0) - TP
+    FN = np.sum(cm, axis=1) - TP
+    
+    precision = TP / (TP + FP)
+    recall = TP / (TP + FN)
+    
+    # Handle division by zero in precision or recall
+    precision = np.nan_to_num(precision)
+    recall = np.nan_to_num(recall)
+    
+    # Compute F1 score
+    f1 = 2 * (precision * recall) / (precision + recall)
+    f1 = np.nan_to_num(f1)  # Replace NaN with 0
+    return np.mean(f1)  # Return the mean F1-score across all classes
+
+# Function to plot and save confusion matrix with F1-score in the title
 def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
+    # Compute the average F1-score
+    avg_f1 = compute_f1_score(cm)
+    
+    # Update title to include average F1-score
+    full_title = f"{title} (Avg F1-Score: {avg_f1:.2f})"
+    
+    # Plot the confusion matrix
     plt.figure(figsize=(8, 6))
     plt.imshow(cm, interpolation='nearest', cmap='coolwarm')
-    plt.title(title)
+    plt.title(full_title)
     plt.colorbar()
+    
     tick_marks = np.arange(len(classes))
     plt.xticks(tick_marks, classes, rotation=45)
     plt.yticks(tick_marks, classes)
@@ -61,6 +91,22 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     plt.xlabel('Predicted label')
     plt.savefig(save_path)
     plt.close()
+
+
+#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):