Update app.py

app.py

@@ -84,18 +84,21 @@ def split_dataset(channels, labels, percentage_idx):
 def euclidean_distance(x, centroid):
     return np.linalg.norm(x - centroid)
 
-# Function to classify test data based on distance to class centroids
+import torch
+
 def classify_based_on_distance(train_data, train_labels, test_data):
-    centroid_0 = np.mean(train_data[train_labels == 0], axis=0)
-    centroid_1 = np.mean(train_data[train_labels == 1], axis=0)
+    # Compute the centroids for the two classes
+    centroid_0 = train_data[train_labels == 0].mean(dim=0)  # Use torch.mean
+    centroid_1 = train_data[train_labels == 1].mean(dim=0)  # Use torch.mean
     
     predictions = []
     for test_point in test_data:
+        # Compute Euclidean distance between the test point and each centroid
         dist_0 = euclidean_distance(test_point, centroid_0)
         dist_1 = euclidean_distance(test_point, centroid_1)
         predictions.append(0 if dist_0 < dist_1 else 1)
     
-    return np.array(predictions)
+    return torch.tensor(predictions)  # Return predictions as a PyTorch tensor
 
 # Function to generate confusion matrix plot
 def plot_confusion_matrix(y_true, y_pred, title):