Update app.py

app.py

@@ -126,12 +126,24 @@ def plot_confusion_matrix(y_true, y_pred, title):
     plt.xlabel('Predicted')
     plt.ylabel('Actual')
     plt.colorbar()
+
+    # Add labels for x and y ticks (Actual/Predicted class labels)
     plt.xticks([0, 1], labels=[0, 1])
     plt.yticks([0, 1], labels=[0, 1])
+
+    # Annotate the confusion matrix
+    thresh = cm.max() / 2  # Define threshold to choose text color (black or white)
+    for i in range(cm.shape[0]):
+        for j in range(cm.shape[1]):
+            plt.text(j, i, format(cm[i, j], 'd'),
+                     ha="center", va="center",
+                     color="white" if cm[i, j] > thresh else "black")
+
     plt.tight_layout()
     plt.savefig(f"{title}.png")
     return Image.open(f"{title}.png")
 
+
 def identical_train_test_split(output_emb, output_raw, labels, percentage_idx):
     N = output_emb.shape[0]  # Get the total number of samples
     
@@ -233,8 +245,8 @@ def process_hdf5_file(uploaded_file, percentage_idx):
         print(f'test_data_emb: {test_data_emb.shape}')
         pred_raw = classify_based_on_distance(train_data_raw, train_labels, test_data_raw)
         pred_emb = classify_based_on_distance(train_data_emb, train_labels, test_data_emb)
-        print(f'pred_emb: {pred_emb}')
-        print(f'actual labels: {test_labels}')
+        #print(f'pred_emb: {pred_emb}')
+        #print(f'actual labels: {test_labels}')
         # Step 9: Generate confusion matrices for both raw and embeddings
         raw_cm_image = plot_confusion_matrix(test_labels, pred_raw, title="Confusion Matrix (Raw Channels)")
         emb_cm_image = plot_confusion_matrix(test_labels, pred_emb, title="Confusion Matrix (Embeddings)")