Update app.py

app.py

@@ -1,35 +1,37 @@
 import gradio as gr
 import torch
-import torchvision
-import transformers
+from torchvision import transforms
+from PIL import Image
 import numpy as np
 from scipy.spatial.distance import cosine
-import cv2
-import os
-
 
+# Constants
 RECOGNITION_THRESHOLD = 0.3
 
-
-# Load the embedding model
-embedding_model = torch.load('final_modelnew.pth', map_location=torch.device('cpu'))
-embedding_model.eval()  # Set the model to evaluation mode
+# Load the model
+model_path = 'final_modelnew.pth'
+model = torch.load(model_path, map_location=torch.device('cpu'))
+model.eval()  # Set the model to evaluation mode
 
 # Database to store embeddings and user IDs
 user_embeddings = {}
 
 # Preprocess the image
 def preprocess_image(image):
-    image = cv2.resize(image, (375, 375))  # Resize image
-    image = image / 255.0  # Normalize pixel values
-    image = np.transpose(image, (2, 0, 1))  # Change from HWC to CHW format
-    return torch.tensor(image, dtype=torch.float32).unsqueeze(0)  # Add batch dimension
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+    ])
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
+    image = transform(image).unsqueeze(0)
+    return image
 
 # Generate embedding
 def generate_embedding(image):
     preprocessed_image = preprocess_image(image)
     with torch.no_grad():  # No need to track gradients
-        return embedding_model(preprocessed_image).numpy()[0]
+        embedding = model(preprocessed_image)
+        return embedding.numpy()[0]
 
 # Register new user
 def register_user(image, user_id):
@@ -49,12 +51,10 @@ def recognize_user(image):
         
         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
+
         if min_distance > RECOGNITION_THRESHOLD:
             return "User not recognized."
         else:
@@ -62,11 +62,10 @@ def recognize_user(image):
     except Exception as e:
         return f"Error during recognition: {str(e)}"
 
-        
 def main():
     with gr.Blocks() as demo:
         gr.Markdown("Facial Recognition System")
-        
+
         with gr.Tab("Register"):
             with gr.Row():
                 img_register = gr.Image()
@@ -85,4 +84,4 @@ def main():
     demo.launch(share=True)
 
 if __name__ == "__main__":
-    main()
+    main()