Update app.py

app.py

@@ -357,13 +357,20 @@ def plot_confusion_matrix(y_true, y_pred, title, save_path="confusion_matrix.png
     plt.xlabel('Predicted label', color=text_color, fontsize=12)
     plt.tight_layout()
 
-    # Save the plot as an image
-    plt.savefig(save_path, transparent=True)  # Transparent to blend with website background
+    # Save the plot as an image, ensure the image is fully written
+    plt.savefig(save_path, bbox_inches='tight', transparent=True)
     plt.close()
 
+    # Ensure the file is properly written and synced to disk before opening
+    time.sleep(0.5)  # Add a small delay to ensure file writing is complete
+    
+    # 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)
-    
+
 def identical_train_test_split(output_emb, output_raw, labels, train_percentage):
     N = output_emb.shape[0]
     indices = torch.randperm(N)