Update app.py

app.py

@@ -5,13 +5,21 @@ from scipy.spatial.distance import cosine
 import cv2
 import os
 
-RECOGNITION_THRESHOLD = 0.1
+# K-mean Clustering
+from sklearn.cluster import KMeans
+import matplotlib.pyplot as plt
+
+# Threshold 
+RECOGNITION_THRESHOLD = 0.1  # Adjust as needed
+n_clusters = 5  # You can adjust this based on your data
+kmeans = KMeans(n_clusters=n_clusters)
 
 # Load the embedding model
 embedding_model = tf.keras.models.load_model('embedding_model.h5')
 
 # Database to store embeddings and user IDs
-user_embeddings = {}
+user_embeddings = []
+user_ids = []
 
 # Preprocess the image
 def preprocess_image(image):
@@ -28,7 +36,8 @@ def generate_embedding(image):
 def register_user(image, user_id):
     try:
         embedding = generate_embedding(image)
-        user_embeddings[user_id] = embedding
+        user_embeddings.append(embedding)
+        user_ids.append(user_id)
         return f"User {user_id} registered successfully."
     except Exception as e:
         return f"Error during registration: {str(e)}"
@@ -37,23 +46,29 @@ def register_user(image, user_id):
 def recognize_user(image):
     try:
         new_embedding = generate_embedding(image)
-        min_distance = float('inf')
-        recognized_user_id = "Unknown"
-        
-        for user_id, embedding in user_embeddings.items():
-            distance = cosine(new_embedding, embedding)
-            print(f"Distance for {user_id}: {distance}")  # Debug: Print distances
-            if distance < min_distance:
-                min_distance = distance
-                recognized_user_id = user_id
-        
-        print(f"Min distance: {min_distance}")  # Debug: Print minimum distance
+        # Update the KMeans model
+        kmeans.fit(user_embeddings)
+        cluster_label = kmeans.predict([new_embedding])[0]
+        distances = kmeans.transform([new_embedding])[0]
+        min_distance = np.min(distances)
+
         if min_distance > RECOGNITION_THRESHOLD:
             return "User not recognized."
-        else:
-            return f"Recognized User: {recognized_user_id}"
+
+        # Find the user ID(s) in the closest cluster
+        recognized_user_ids = [user_ids[i] for i, label in enumerate(kmeans.labels_) if label == cluster_label]
+        return f"Recognized User(s): {', '.join(recognized_user_ids)}"
     except Exception as e:
         return f"Error during recognition: {str(e)}"
+
+def plot_clusters():
+    # Assuming embeddings are 2-dimensional
+    plt.figure(figsize=(8, 6))
+    plt.scatter(*zip(*user_embeddings), c=kmeans.labels_)
+    plt.title('User Embeddings Clustered by K-Means')
+    plt.xlabel('Embedding Dimension 1')
+    plt.ylabel('Embedding Dimension 2')
+    plt.show()
         
 def main():
     with gr.Blocks() as demo: