Update app.py

app.py

@@ -61,36 +61,44 @@ 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
 
-def plot_confusion_matrix_beamPred(cm, classes, title, save_path):
+def plot_confusion_matrix_beamPred(cm, classes, title, save_path, light_mode=True):
     # Compute the average F1-score
     avg_f1 = compute_f1_score(cm)
 
-    # Set dark mode styling
-    plt.style.use('dark_background')
+    # Choose the color scheme based on the mode
+    if light_mode:
+        plt.style.use('default')  # Use default (light) mode styling
+        text_color = 'black'
+        cmap = 'Blues'  # Light-mode-friendly colormap
+    else:
+        plt.style.use('dark_background')  # Use dark mode styling
+        text_color = 'white'
+        cmap = 'cividis'  # Dark-mode-friendly colormap
+
     plt.figure(figsize=(10, 10))
-    
+
     # Plot the confusion matrix with a dark-mode compatible colormap
-    #sns.heatmap(cm, cmap="magma", cbar=True, linecolor='white', vmin=0, vmax=cm.max(), alpha=0.85)
-    sns.heatmap(cm, cmap="cividis", cbar=True, linecolor='white', vmin=0, vmax=cm.max(), alpha=0.85)
-    
+    sns.heatmap(cm, cmap=cmap, cbar=True, linecolor='white', vmin=0, vmax=cm.max(), alpha=0.85)
+
     # Add F1-score to the title
-    plt.title(f"{title}\n(F1 Score: {avg_f1:.3f})", color="white", fontsize=14)
-    
+    plt.title(f"{title}\n(F1 Score: {avg_f1:.3f})", color=text_color, fontsize=14)
+
     tick_marks = np.arange(len(classes))
-    plt.xticks(tick_marks, classes, color="white", fontsize=14)  # White text for dark mode
-    plt.yticks(tick_marks, classes, color="white", fontsize=14)  # White text for dark mode
+    plt.xticks(tick_marks, classes, color=text_color, fontsize=14)  # Adjust text color based on the mode
+    plt.yticks(tick_marks, classes, color=text_color, fontsize=14)  # Adjust text color based on the mode
 
-    plt.ylabel('True label', color="white", fontsize=14)
-    plt.xlabel('Predicted label', color="white", fontsize=14)
+    plt.ylabel('True label', color=text_color, fontsize=14)
+    plt.xlabel('Predicted label', color=text_color, fontsize=14)
     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, transparent=True)  # Transparent to blend with the site background
     plt.close()
-    
+
     # Return the saved image
     return Image.open(save_path)
 
+
 def compute_average_confusion_matrix(folder):
     confusion_matrices = []
     max_num_labels = 0
@@ -140,7 +148,7 @@ 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):
+def plot_confusion_matrix_from_csv(csv_file_path, title, save_path, light_mode=True):
     # Load CSV file
     data = pd.read_csv(csv_file_path)
     
@@ -153,29 +161,37 @@ def plot_confusion_matrix_from_csv(csv_file_path, title, save_path):
     
     # Compute F1-score
     f1 = f1_score(y_true, y_pred, average='macro')  # Macro-average F1-score
+
+    # Set styling based on light or dark mode
+    if light_mode:
+        plt.style.use('default')  # Light mode styling
+        text_color = 'black'
+        cmap = 'Blues'  # Light-mode-friendly colormap
+    else:
+        plt.style.use('dark_background')  # Dark mode styling
+        text_color = 'white'
+        cmap = 'magma'  # Dark-mode-friendly colormap
     
-    # Set dark mode styling
-    plt.style.use('dark_background')
     plt.figure(figsize=(5, 5))
-    
-    # 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')
-    
+
+    # Plot the confusion matrix with the chosen colormap
+    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap, 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.title(f"{title}\n(F1 Score: {f1:.3f})", color=text_color, fontsize=14)
+
+    # Customize tick labels for light/dark mode
+    plt.xticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color=text_color, fontsize=10)
+    plt.yticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color=text_color, fontsize=10)
+
+    plt.ylabel('True label', color=text_color, fontsize=12)
+    plt.xlabel('Predicted label', color=text_color, 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, transparent=True)  # Use transparent to blend with the website
     plt.close()
-    
+
     # Return the saved image
     return Image.open(save_path)
 
@@ -308,35 +324,45 @@ def classify_based_on_distance(train_data, train_labels, test_data):
     
     return torch.tensor(predictions)  # Return predictions as a PyTorch tensor
 
-def plot_confusion_matrix(y_true, y_pred, title):
+def plot_confusion_matrix(y_true, y_pred, title, save_path="confusion_matrix.png", light_mode=True):
     cm = confusion_matrix(y_true, y_pred)
-    
+
     # Calculate F1 Score
     f1 = f1_score(y_true, y_pred, average='weighted')
 
-    plt.style.use('dark_background')
+    # Choose the color scheme based on the mode
+    if light_mode:
+        plt.style.use('default')  # Light mode styling
+        text_color = 'black'
+        cmap = 'Blues'  # Light-mode-friendly colormap
+    else:
+        plt.style.use('dark_background')  # Dark mode styling
+        text_color = 'white'
+        cmap = 'magma'  # Dark-mode-friendly colormap
+
     plt.figure(figsize=(5, 5))
-    
-    # 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')
-    
+
+    # Plot the confusion matrix with a colormap compatible for the mode
+    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap, 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.title(f"{title}\n(F1 Score: {f1:.3f})", color=text_color, fontsize=14)
+
+    # Customize tick labels for light/dark mode
+    plt.xticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color=text_color, fontsize=10)
+    plt.yticks([0.5, 1.5], labels=['Class 0', 'Class 1'], color=text_color, fontsize=10)
+
+    plt.ylabel('True label', color=text_color, fontsize=12)
+    plt.xlabel('Predicted label', color=text_color, 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.savefig(save_path, transparent=True)  # Transparent to blend with website background
     plt.close()
-    
+
     # Return the saved image
-    return Image.open(f"{title}.png")
+    return Image.open(save_path)
     
 def identical_train_test_split(output_emb, output_raw, labels, train_percentage):
     N = output_emb.shape[0]