Update app.py

app.py

@@ -62,7 +62,6 @@ def compute_f1_score(cm):
     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)
@@ -70,20 +69,22 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
     # Update title to include average F1-score
     full_title = f"{title} (Avg F1-Score: {avg_f1:.2f})"
     
-    # Plot the confusion matrix
+    # Plot the confusion matrix with dark mode adjustments
     plt.figure(figsize=(8, 6))
-    plt.imshow(cm, interpolation='nearest', cmap='coolwarm')
-    plt.title(full_title)
+    plt.imshow(cm, interpolation='nearest', cmap='coolwarm')  # Dark mode color scheme
+    plt.title(full_title, color='white', pad=20)  # Add padding to prevent title clipping, white text for dark mode
     plt.colorbar()
     
     tick_marks = np.arange(len(classes))
-    plt.xticks(tick_marks, classes, rotation=45)
-    plt.yticks(tick_marks, classes)
+    plt.xticks(tick_marks, classes, rotation=45, color='white')  # White text for dark mode
+    plt.yticks(tick_marks, classes, color='white')  # White text for dark mode
 
-    plt.tight_layout()
-    plt.ylabel('True label')
-    plt.xlabel('Predicted label')
-    plt.savefig(save_path)
+    plt.tight_layout(pad=2.0)  # Add padding to prevent axis label clipping
+    plt.ylabel('True label', color='white')  # White text for dark mode
+    plt.xlabel('Predicted label', color='white')  # White text for dark mode
+
+    # Save the plot with a black background for dark mode
+    plt.savefig(save_path, facecolor='black')
     plt.close()
 
 def compute_average_confusion_matrix(folder):
@@ -135,7 +136,6 @@ percentage_values_los = np.linspace(0.001, 1, 20) * 100  # 20 percentage values
 from sklearn.metrics import f1_score
 import seaborn as sns
 
-# Function to compute confusion matrix, F1-score and plot it with dark mode style
 def plot_confusion_matrix_from_csv(csv_file_path, title, save_path):
     # Load CSV file
     data = pd.read_csv(csv_file_path)
@@ -147,29 +147,30 @@ def plot_confusion_matrix_from_csv(csv_file_path, title, save_path):
     # Compute confusion matrix
     cm = confusion_matrix(y_true, y_pred)
     
-    # Compute F1-score
-    f1 = f1_score(y_true, y_pred, average='macro')  # Macro-average F1-score
-    
-    # Set dark mode styling
-    plt.style.use('dark_background')
+    # Calculate F1 Score
+    f1 = f1_score(y_true, y_pred, average='weighted')
+
+    # Plot the confusion matrix with dark mode colors
     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)  # Display F1-Score in title, add padding for better visibility
+    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
     
-    # 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} (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)
+    # 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", fontsize=12)
-    plt.xlabel('Predicted label', color="white", fontsize=12)
-    plt.tight_layout()
+    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 of labels
     
     # 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='black')  # Set background to black for dark mode
     plt.close()
     
     # Return the saved image
@@ -304,32 +305,39 @@ def classify_based_on_distance(train_data, train_labels, test_data):
     
     return torch.tensor(predictions)  # Return predictions as a PyTorch tensor
 
-# Function to generate confusion matrix plot
 def plot_confusion_matrix(y_true, y_pred, title):
     cm = confusion_matrix(y_true, y_pred)
+    
+    # Calculate F1 Score
+    f1 = f1_score(y_true, y_pred, average='weighted')
+    
+    # Plot the confusion matrix with dark mode colors
     plt.figure(figsize=(5, 5))
-    plt.imshow(cm, cmap='Blues')
-    plt.title(title)
-    plt.xlabel('Predicted')
-    plt.ylabel('Actual')
+    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()
-
-    # Add labels for x and y ticks (Actual/Predicted class labels)
-    plt.xticks([0, 1], labels=[0, 1])
-    plt.yticks([0, 1], labels=[0, 1])
-
+    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  # Define threshold to choose text color (black or white)
+    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.tight_layout()
-    plt.savefig(f"{title}.png")
+    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
+    plt.close()
+    
     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