Update app.py

app.py

@@ -112,7 +112,7 @@ def plot_confusion_matrix(y_true, y_pred, title):
     plt.savefig(f"{title}.png")
     return Image.open(f"{title}.png")
 
-def identical_train_test_split(output_emb, output_raw, percentage):
+def identical_train_test_split(output_emb, output_raw, labels, percentage):
     N = output_emb.shape[0]  # Get the total number of samples
     
     # Generate the indices for shuffling and splitting
@@ -200,6 +200,7 @@ def process_hdf5_file(uploaded_file, percentage_idx):
 
         train_data_emb, test_data_emb, train_data_raw, test_data_raw, train_labels, test_labels = identical_train_test_split(output_emb.view(len(output_emb),-1),
                                                                                                                              output_raw.view(len(output_raw),-1),
+                                                                                                                             labels,
                                                                                                                              percentage_idx)
         
         # Step 8: Perform classification using the Euclidean distance for both raw and embeddings