Update app.py

app.py

@@ -324,49 +324,36 @@ 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, save_path="confusion_matrix.png", light_mode=True):
+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')
 
-    # 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.style.use('dark_background')
     plt.figure(figsize=(5, 5))
-
-    # 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')
-
+    
+    # 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}\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.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.tight_layout()
-
-    # Save the plot as an image, ensure the image is fully written
-    plt.savefig(save_path, bbox_inches='tight', transparent=True)
+    
+    # Save the plot as an image
+    plt.savefig(f"{title}.png", transparent=True)  # Use transparent to blend with the dark mode website
     plt.close()
-
-    # Check if the file exists and can be opened
-    if not os.path.exists(save_path):
-        raise FileNotFoundError(f"File {save_path} not found.")
-
+    
     # Return the saved image
-    return Image.open(save_path)
+    return Image.open(f"{title}.png")
+
 
 def identical_train_test_split(output_emb, output_raw, labels, train_percentage):
     N = output_emb.shape[0]
@@ -574,7 +561,7 @@ with gr.Blocks(css="""
         # Add a conclusion section at the bottom
         gr.Markdown("""
             <div class="explanation-box">
-                **Conclusions**: By adjusting the data percentage and task complexity, you can observe how the LWM generalizes well to unseen scenarios and handles various complexities in the beam prediction task.
+                <b>Conclusions<b>: LWM embeddings offer such high generalization that with just a limited number of training samples, we can get high performances.
             </div>
         """)
 
@@ -636,7 +623,7 @@ with gr.Blocks(css="""
         # Add a conclusion section at the bottom
         gr.Markdown("""
             <div class="explanation-box">
-                **Conclusions**: With this task, you can evaluate how well LWM embeddings perform on LoS/NLoS classification tasks, and compare it to the performance of raw channels in identifying these features.
+                <b>Conclusions<b>: LWM CLS embeddings, although very small (raw channels size / 32), offer a rich and holistic knowledge about channels, making them suitable for a task like LoS/NLoS classfication, specifically with very limited data.  
             </div>
         """)