Update app.py

app.py

@@ -2,8 +2,6 @@ import streamlit as st
 import numpy as np
 from sklearn.neighbors import KNeighborsClassifier
 import matplotlib.pyplot as plt
-st.set_option('deprecation.showPyplotGlobalUse', False)
-
 
 def get_user_data_train():
     data_points = []
@@ -45,7 +43,8 @@ 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
-    plt.scatter(X_test[:, 0], X_test[:, 1], label=f'Test (Predicted Labels)', marker='x', c=predictions)
+    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.xlabel('X-coordinate')
     plt.ylabel('Y-coordinate')