Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import numpy as np
 from sklearn.neighbors import KNeighborsClassifier
+import matplotlib.pyplot as plt
 
 def get_user_data_train():
     data_points = []
@@ -33,6 +34,23 @@ def knn_classification(X_train, y_train, X_test, k_value):
     predictions = knn_classifier.predict(X_test)
     return predictions
 
+def plot_training_and_test_data(X_train, y_train, X_test, predictions):
+    unique_labels = np.unique(y_train)
+
+    # Plot training data
+    for label in unique_labels:
+        indices = np.where(y_train == label)
+        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)
+
+    plt.xlabel('X-coordinate')
+    plt.ylabel('Y-coordinate')
+    plt.title('Training and Test Data with Predicted Labels')
+    plt.legend()
+    st.pyplot()
+
 def main():
     st.title("K-Nearest Neighbor Classification App")
 
@@ -48,6 +66,9 @@ def main():
     # Perform k-nearest neighbor classification
     predictions = knn_classification(X_train, y_train, X_test, k_value)
 
+    # Plot training and test data
+    plot_training_and_test_data(X_train, y_train, X_test, predictions)
+
     # Display results
     st.subheader("Results:")
     st.write("User-defined Data Points for Testing:")