Update app.py

app.py

@@ -61,23 +61,41 @@ 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, light_mode=True):
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+from PIL import Image
+
+def plot_confusion_matrix_beamPred(cm, classes, title, save_path, dark_mode=None):
+    """
+    Plot confusion matrix and adjust colors based on light/dark mode settings.
+    :param cm: Confusion matrix data.
+    :param classes: List of class labels.
+    :param title: Plot title.
+    :param save_path: Path to save the plot.
+    :param dark_mode: Boolean to toggle between light and dark modes. If None, use the current theme.
+    """
+    
+    # If dark_mode is None, try detecting it from rcParams (matplotlib theme)
+    if dark_mode is None:
+        dark_mode = plt.rcParams['axes.facecolor'] == '#333333'  # Check if dark background is set
+    
     # Compute the average F1-score
     avg_f1 = compute_f1_score(cm)
 
     # 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:
+    if dark_mode:
         plt.style.use('dark_background')  # Use dark mode styling
         text_color = 'white'
         cmap = 'cividis'  # Dark-mode-friendly colormap
+    else:
+        plt.style.use('default')  # Use default (light) mode styling
+        text_color = 'black'
+        cmap = 'Blues'  # Light-mode-friendly colormap
 
     plt.figure(figsize=(10, 10))
 
-    # Plot the confusion matrix with a dark-mode compatible colormap
+    # Plot the confusion matrix with the selected colormap
     sns.heatmap(cm, cmap=cmap, cbar=True, linecolor='white', vmin=0, vmax=cm.max(), alpha=0.85)
 
     # Add F1-score to the title
@@ -98,7 +116,6 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path, light_mode=Tru
     # Return the saved image
     return Image.open(save_path)
 
-
 def compute_average_confusion_matrix(folder):
     confusion_matrices = []
     max_num_labels = 0