modify verification app

app.py

@@ -7,7 +7,6 @@ import gradio as gr
 from PIL import Image
 from facenet_pytorch import MTCNN
 from transformers import ViTImageProcessor, ViTModel
-import pickle
 import time
 
 # Define the ViT class
@@ -43,91 +42,59 @@ transform = transforms.Compose([
     transforms.ToTensor()
 ])
 
-# Load the database of embeddings
-with open('face_database_me.pkl', 'rb') as f:
-    database = pickle.load(f)
-
 def cosine_similarity(embedding1, embedding2):
     similarity = torch.nn.functional.cosine_similarity(embedding1.flatten().unsqueeze(0), embedding2.flatten().unsqueeze(0))
     return similarity.item()
 
-def compare_embeddings(embedding, database, threshold=0.9):
-    best_match = None
-    best_similarity = threshold
-    for name, db_embeddings in database.items():
-        for db_embedding in db_embeddings:
-            db_embedding = torch.tensor(db_embedding).to(device)
-            similarity = cosine_similarity(embedding, db_embedding)
-            if similarity > best_similarity:
-                best_match = name
-                best_similarity = similarity
-    if best_match is not None:
-        return best_match, best_similarity
-    return None, None
-
-def align_faces(frame):
+def align_face(frame):
     # Convert the frame to a PIL image if it's a numpy array
     if isinstance(frame, np.ndarray):
         frame = Image.fromarray(frame)
     boxes, _ = mtcnn.detect(frame)
-    aligned_faces = []
-    if boxes is not None:
+    if boxes is not None and len(boxes) > 0:
         faces = mtcnn(frame)
-        if faces is not None:
-            for face in faces:
-                # Convert the face tensor to PIL Image
-                face = transforms.ToPILImage()(face)
-                aligned_faces.append(face)
-    return aligned_faces, boxes
+        if faces is not None and len(faces) > 0:
+            face = faces[0]
+            # Convert the face tensor to PIL Image
+            face = transforms.ToPILImage()(face)
+            return face
+    return None
 
-def draw_annotations(frame, detections, names=None):
-    if detections is None:
-        return frame
-    if names is None:
-        names = ["Unknown"] * len(detections)
-    for i, detection in enumerate(detections):
-        x1, y1, x2, y2 = map(int, detection)
-        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        if names[i]:
-            cv2.putText(frame, names[i], (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
-    return frame
-
-def process_image(image):
+def process_images(image1, image2):
     start_time = time.time()
     
-    frame = np.array(image)
-    aligned_faces, boxes = align_faces(frame)
+    frame1 = np.array(image1)
+    frame2 = np.array(image2)
+    
+    face1 = align_face(frame1)
+    face2 = align_face(frame2)
     
-    names = []
-    if aligned_faces is not None:
-        for face in aligned_faces:
-            face = transform(face).unsqueeze(0).to(device)
-            with torch.no_grad():
-                embedding = model(face)
-            name, similarity = compare_embeddings(embedding, database)
-            if name is not None:
-                names.append(f"{name} ({similarity:.2f})")
-            else:
-                names.append("Unknown")
-        annotated_image = draw_annotations(frame, boxes, names)
-        result = "Face recognition complete."
-    else:
-        annotated_image = frame
-        result = "No faces detected."
+    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)
+    
+    with torch.no_grad():
+        embedding1 = model(face1)
+        embedding2 = model(face2)
+    
+    similarity = cosine_similarity(embedding1, embedding2)
     
     end_time = time.time()
     inference_time = end_time - start_time
-    result += f" Inference time: {inference_time:.2f} seconds"
     
-    return annotated_image, result
+    result = f"Similarity: {similarity:.2f}\nInference time: {inference_time:.2f} seconds"
+    
+    return (frame1, frame2), result
 
 # Create the Gradio interface
 iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),  # Ensure the input type matches what the function expects
-    outputs=[gr.Image(type="numpy"), gr.Textbox()],
-    title="Face Detection and Recognition with MTCNN and ViT",
-    description="Upload an image and the model will detect and recognize faces in it."
+    fn=process_images,
+    inputs=[gr.Image(type="pil"), gr.Image(type="pil")],
+    outputs=[gr.Gallery(), gr.Textbox()],
+    title="Face Verification with MTCNN and ViT",
+    description="Upload two images and the model will verify if the faces in both images are of the same person."
 )
 
 # Launch the interface