change the similarity algorithm

app.py

@@ -1,7 +1,6 @@
 import torch
 from torch import nn
 import torchvision.transforms as transforms
-import cv2
 import numpy as np
 import gradio as gr
 from PIL import Image
@@ -36,13 +35,14 @@ model.to(device)
 # Initialize MTCNN for face detection
 mtcnn = MTCNN(keep_all=True, device=device)
 
-# Define the transformation
-transform = transforms.Compose([
+# Define the transformation with normalization
+val_test_transform_vit = transforms.Compose([
     transforms.Resize((224, 224)),
-    transforms.ToTensor()
+    transforms.ToTensor(),
+    transforms.Normalize(mean=processor.image_mean, std=processor.image_std)
 ])
 
-def compare_faces(embedding1, embedding2, threshold=1.0):
+def compare_faces(embedding1, embedding2, threshold=0.6):  # Adjusted threshold
     dist = np.linalg.norm(embedding1.cpu().numpy() - embedding2.cpu().numpy())
     return dist, dist < threshold
 
@@ -76,8 +76,8 @@ def process_images(image1, image2):
     if face1 is None or face2 is None:
         return None, "Face not detected in one or both images."
     
-    face1 = transform(face1).unsqueeze(0).to(device)
-    face2 = transform(face2).unsqueeze(0).to(device)
+    face1 = val_test_transform_vit(face1).unsqueeze(0).to(device)
+    face2 = val_test_transform_vit(face2).unsqueeze(0).to(device)
     
     with torch.no_grad():
         embedding1 = model(face1)
@@ -86,7 +86,7 @@ def process_images(image1, image2):
     embedding1 = l2_normalize(embedding1)
     embedding2 = l2_normalize(embedding2)
     
-    distance, is_match = compare_faces(embedding1, embedding2)
+    distance, is_match = compare_faces(embedding1, embedding2, threshold=0.88)  
     
     end_time = time.time()
     inference_time = end_time - start_time