Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import tensorflow as tf
+import torch
 import numpy as np
 from scipy.spatial.distance import cosine
 import cv2
@@ -7,8 +7,10 @@ import os
 
 RECOGNITION_THRESHOLD = 0.3
 
+# Assuming the PyTorch model is a ResNet (or similar) and has been trained accordingly
 # Load the embedding model
-embedding_model = tf.keras.models.load_model('full_mode2.pth')
+embedding_model = torch.load('full_mode2.pth')
+embedding_model.eval()  # Set the model to evaluation mode
 
 # Database to store embeddings and user IDs
 user_embeddings = {}
@@ -16,13 +18,15 @@ user_embeddings = {}
 # Preprocess the image
 def preprocess_image(image):
     image = cv2.resize(image, (375, 375))  # Resize image
-    image = tf.keras.applications.resnet50.preprocess_input(image)
-    return np.expand_dims(image, axis=0)
+    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
 
 # Generate embedding
 def generate_embedding(image):
     preprocessed_image = preprocess_image(image)
-    return embedding_model.predict(preprocessed_image)[0]
+    with torch.no_grad():  # No need to track gradients
+        return embedding_model(preprocessed_image).numpy()[0]
 
 # Register new user
 def register_user(image, user_id):
@@ -55,6 +59,10 @@ def recognize_user(image):
     except Exception as e:
         return f"Error during recognition: {str(e)}"
         
+def main():
+    # ... (rest of your Gradio setup code)
+
+        
 def main():
     with gr.Blocks() as demo:
         gr.Markdown("Facial Recognition System")