fixing verification algorithm

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 faceViT4.pth filter=lfs diff=lfs merge=lfs -text
+faceNet6.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,49 +1,30 @@
 import torch
-from torch import nn
 import torchvision.transforms as transforms
 import numpy as np
 import gradio as gr
-from PIL import Image
-from facenet_pytorch import MTCNN
-from transformers import ViTImageProcessor, ViTModel
+from PIL import Image, ImageDraw
+from facenet_pytorch import MTCNN, InceptionResnetV1
 import time
 
-# Define the ViT class
-class ViT(nn.Module):
-    def __init__(self, base_model):
-        super(ViT, self).__init__()
-        self.base_model = base_model
-    
-    def forward(self, x):
-        x = self.base_model(x).pooler_output 
-        return x
-
-# Load the model and processor
-model_name = "google/vit-base-patch16-224"
-processor = ViTImageProcessor.from_pretrained(model_name)   
-base_model = ViTModel.from_pretrained("WinKawaks/vit-small-patch16-224")
-model = ViT(base_model)
-model.load_state_dict(torch.load('faceViT4.pth'))
-
-# Set the model to evaluation mode
-model.eval()
-
-# Check if CUDA is available and move the model to GPU if it is
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
+# Initialize MTCNN for face detection with smaller face size detection
+mtcnn = MTCNN(keep_all=True, device='cuda' if torch.cuda.is_available() else 'cpu', min_face_size=12)
 
-# Initialize MTCNN for face detection
-mtcnn = MTCNN(keep_all=True, device=device)
+# Load the pre-trained FaceNet model
+facenet = InceptionResnetV1(pretrained='vggface2').eval().to('cuda' if torch.cuda.is_available() else 'cpu')
+model_path = r'D:\BRI BRAIN\faceNet6.pth'
+model_state_dict = torch.load(model_path)
+facenet.load_state_dict(model_state_dict)
+facenet.eval()  # Set the model to evaluation mode
 
 # Define the transformation with normalization
-val_test_transform_vit = transforms.Compose([
-    transforms.Resize((224, 224)),
+val_test_transform = transforms.Compose([
+    transforms.Resize((160, 160)),  # FaceNet expects 160x160 input
     transforms.ToTensor(),
-    transforms.Normalize(mean=processor.image_mean, std=processor.image_std)
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 ])
 
-def compare_faces(embedding1, embedding2, threshold=0.6):  # Adjusted threshold
-    dist = np.linalg.norm(embedding1.cpu().numpy() - embedding2.cpu().numpy())
+def compare_faces(embedding1, embedding2, threshold=0.5):  # Adjusted threshold
+    dist = np.linalg.norm(embedding1 - embedding2)
     return dist, dist < threshold
 
 def align_face(frame):
@@ -57,11 +38,16 @@ def align_face(frame):
             face = faces[0]
             # Convert the face tensor to PIL Image
             face = transforms.ToPILImage()(face)
-            return face
-    return None
+            return face, boxes[0]
+    return None, None
+
+def draw_bounding_box(image, box):
+    draw = ImageDraw.Draw(image)
+    draw.rectangle(box.tolist(), outline="red", width=3)
+    return image
 
 def l2_normalize(tensor):
-    norm = torch.norm(tensor, p=2, dim=1, keepdim=True)
+    norm = np.linalg.norm(tensor, ord=2, axis=1, keepdims=True)
     return tensor / norm
 
 def process_images(image1, image2):
@@ -70,37 +56,44 @@ def process_images(image1, image2):
     frame1 = np.array(image1)
     frame2 = np.array(image2)
     
-    face1 = align_face(frame1)
-    face2 = align_face(frame2)
+    face1, box1 = align_face(frame1)
+    face2, box2 = align_face(frame2)
     
     if face1 is None or face2 is None:
         return None, "Face not detected in one or both images."
     
-    face1 = val_test_transform_vit(face1).unsqueeze(0).to(device)
-    face2 = val_test_transform_vit(face2).unsqueeze(0).to(device)
+    face1 = val_test_transform(face1).unsqueeze(0).to('cuda' if torch.cuda.is_available() else 'cpu')
+    face2 = val_test_transform(face2).unsqueeze(0).to('cuda' if torch.cuda.is_available() else 'cpu')
     
     with torch.no_grad():
-        embedding1 = model(face1)
-        embedding2 = model(face2)
+        embedding1 = facenet(face1).cpu().numpy()
+        embedding2 = facenet(face2).cpu().numpy()
     
     embedding1 = l2_normalize(embedding1)
     embedding2 = l2_normalize(embedding2)
     
-    distance, is_match = compare_faces(embedding1, embedding2, threshold=0.88)  
+    distance, is_match = compare_faces(embedding1, embedding2, threshold=0.2)
+    
+    # Calculate confidence
+    confidence = max(0.0, 1.0 - distance / 1.0)  # Ensure confidence is between 0 and 1
     
     end_time = time.time()
     inference_time = end_time - start_time
     
-    result = f"Distance: {distance:.2f}\nMatch: {is_match}\nInference time: {inference_time:.2f} seconds"
+    # Draw bounding boxes on the original images
+    image1_with_box = draw_bounding_box(image1, box1)
+    image2_with_box = draw_bounding_box(image2, box2)
+    
+    result = f"Distance: {distance:.2f}\nMatch: {is_match}\nConfidence: {confidence:.2f}\nInference time: {inference_time:.2f} seconds"
     
-    return (image1, image2), result
+    return [image1_with_box, image2_with_box], result
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=process_images,
     inputs=[gr.Image(type="pil"), gr.Image(type="pil")],
     outputs=[gr.Gallery(), gr.Textbox()],
-    title="Face Verification with MTCNN and ViT",
+    title="Face Verification with FaceNet",
     description="Upload two images and the model will verify if the faces in both images are of the same person."
 )
 

faceNet6.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc80c77dda898bfdb928a761bdd91522482160f7fc5d341573ce79de80b38d56
+size 112013482