Update app.py

app.py

@@ -43,8 +43,7 @@ def plot_training_and_test_data(X_train, y_train, X_test, predictions):
         plt.scatter(X_train[indices, 0], X_train[indices, 1], label=f'Training ({label})')
 
     # Plot test data with predicted labels
-    numerical_predictions = np.arange(len(np.unique(predictions)))
-    plt.scatter(X_test[:, 0], X_test[:, 1], label=f'Test (Predicted Labels)', marker='x', c=numerical_predictions[predictions])
+    plt.scatter(X_test[:, 0], X_test[:, 1], label=f'Test (Predicted Labels)', marker='x', c=predictions)
 
     plt.xlabel('X-coordinate')
     plt.ylabel('Y-coordinate')