Update app.py

app.py

@@ -15,7 +15,7 @@ from sklearn.metrics import f1_score
 import seaborn as sns
 
 #################### BEAM PREDICTION #########################}
-def beam_prediction_task(data_percentage, task_complexity):
+def beam_prediction_task(data_percentage, task_complexity, user_mode="light"):
     # Folder naming convention based on input_type, data_percentage, and task_complexity
     raw_folder = f"images/raw_{data_percentage/100:.1f}_{task_complexity}"
     embeddings_folder = f"images/embedding_{data_percentage/100:.1f}_{task_complexity}"
@@ -24,7 +24,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_beamPred(raw_cm, classes=np.arange(raw_cm.shape[0]), title=f"Raw Confusion Matrix\n({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\n({data_percentage}% data, {task_complexity} beams)", save_path=raw_cm_path, user_mode=user_mode)
         raw_img = Image.open(raw_cm_path)
     else:
         raw_img = None
@@ -33,15 +33,13 @@ 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_beamPred(embeddings_cm, classes=np.arange(embeddings_cm.shape[0]), title=f"Embeddings Confusion Matrix\n({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\n({data_percentage}% data, {task_complexity} beams)", save_path=embeddings_cm_path, user_mode=user_mode)
         embeddings_img = Image.open(embeddings_cm_path)
     else:
         embeddings_img = None
 
     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
@@ -61,30 +59,12 @@ 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
 
-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
-    
+def plot_confusion_matrix_beamPred(cm, classes, title, save_path, user_mode="light"):
     # Compute the average F1-score
     avg_f1 = compute_f1_score(cm)
 
-    # Choose the color scheme based on the mode
-    if dark_mode:
+    # Choose the color scheme based on the user's mode
+    if user_mode == 'dark':
         plt.style.use('dark_background')  # Use dark mode styling
         text_color = 'white'
         cmap = 'cividis'  # Dark-mode-friendly colormap
@@ -95,7 +75,7 @@ def plot_confusion_matrix_beamPred(cm, classes, title, save_path, dark_mode=None
 
     plt.figure(figsize=(10, 10))
 
-    # Plot the confusion matrix with the selected colormap
+    # Plot the confusion matrix with a colormap compatible for the mode
     sns.heatmap(cm, cmap=cmap, cbar=True, linecolor='white', vmin=0, vmax=cm.max(), alpha=0.85)
 
     # Add F1-score to the title
@@ -579,12 +559,21 @@ with gr.Blocks(css="""
                 task_complexity_dropdown = gr.Dropdown(label="Task Complexity (Number of Beams)", choices=[16, 32, 64, 128, 256], value=16)
 
         with gr.Row():
+            
+            mode_input = gr.Textbox(visible=False)  # Hidden input to capture user mode
+            gr.Markdown("""
+            <script>
+            const userPrefersDark = window.matchMedia && window.matchMedia('(prefers-color-scheme: dark)').matches;
+            document.querySelector('input[name="mode_input"]').value = userPrefersDark ? 'dark' : 'light';
+            </script>
+            """)
+            
             raw_img_bp = gr.Image(label="Raw Channels", type="pil", width=300, height=500)
             embeddings_img_bp = gr.Image(label="Embeddings", type="pil", width=300, height=500)
 
         # Update the confusion matrices whenever sliders change
-        data_percentage_slider.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
-        task_complexity_dropdown.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
+        data_percentage_slider.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown, mode_input], outputs=[raw_img_bp, embeddings_img_bp])
+        task_complexity_dropdown.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown, mode_input], outputs=[raw_img_bp, embeddings_img_bp])
 
         # Add a conclusion section at the bottom
         gr.Markdown("""