Update app.py

app.py

@@ -21,11 +21,17 @@ except ImportError as e:
 class FaceRecognitionSystem:
     def __init__(self):
         self.setup_database()
-        global FACE_NET_AVAILABLE  # Declare global at the start of method
+        global FACE_NET_AVAILABLE
         
         if FACE_NET_AVAILABLE:
             try:
-                self.mtcnn = MTCNN(keep_all=True, device='cpu')
+                # Initialize MTCNN with simpler settings
+                self.mtcnn = MTCNN(
+                    keep_all=False,  # Only detect one face for simplicity
+                    min_face_size=40,
+                    thresholds=[0.6, 0.7, 0.7],
+                    device='cpu'
+                )
                 self.resnet = InceptionResnetV1(pretrained='vggface2', classify=False, device='cpu').eval()
                 print("✅ FaceNet models loaded successfully")
             except Exception as e:
@@ -79,19 +85,31 @@ class FaceRecognitionSystem:
                 pil_image = Image.open(image).convert('RGB')
             elif isinstance(image, np.ndarray):
                 # numpy array
-                pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+                if len(image.shape) == 3:
+                    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+                else:
+                    return None, "Invalid image format"
             else:
                 # Handle Gradio file object
                 pil_image = Image.open(image).convert('RGB')
             
-            # Detect faces and extract embeddings
-            faces = self.mtcnn(pil_image)
-            if faces is None:
+            # Detect face and get cropped face
+            face_tensor, prob = self.mtcnn(pil_image, return_prob=True)
+            
+            if face_tensor is None:
                 return None, "No face detected in the image"
             
-            # Get embedding for the first face
-            embedding = self.resnet(faces.unsqueeze(0))
-            embedding_np = embedding.detach().numpy().flatten()
+            if prob < 0.9:  # Confidence threshold
+                return None, "Low confidence face detection"
+            
+            # Ensure proper tensor dimensions [3, 160, 160]
+            if face_tensor.dim() == 4:
+                face_tensor = face_tensor.squeeze(0)  # Remove batch dimension if present
+            
+            # Get embedding
+            with torch.no_grad():
+                embedding = self.resnet(face_tensor.unsqueeze(0))  # Add batch dimension
+                embedding_np = embedding.detach().numpy().flatten()
             
             return embedding_np, "Face embedding extracted successfully"
             
@@ -155,7 +173,7 @@ class FaceRecognitionSystem:
         # Find the best match
         best_match = None
         min_distance = float('inf')
-        threshold = 0.8  # Similarity threshold
+        threshold = 1.0  # Similarity threshold
         
         for student_id, name, class_name, embedding_blob in students:
             known_embedding = np.frombuffer(embedding_blob, dtype=np.float32)
@@ -192,15 +210,7 @@ class FaceRecognitionSystem:
                 result_message = f"⚠️ {name} already marked today (ID: {student_id})"
             
             # Convert image for display
-            if isinstance(image, str):
-                display_image = cv2.imread(image)
-            elif isinstance(image, np.ndarray):
-                display_image = image.copy()
-            else:
-                # Handle Gradio file object
-                display_image = np.array(Image.open(image))
-                if len(display_image.shape) == 3 and display_image.shape[2] == 3:
-                    display_image = cv2.cvtColor(display_image, cv2.COLOR_RGB2BGR)
+            display_image = self.prepare_display_image(image)
             
             # Add text to image
             if display_image is not None:
@@ -210,12 +220,32 @@ class FaceRecognitionSystem:
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                 cv2.putText(display_image, f"Class: {class_name}", (10, 90), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+                cv2.putText(display_image, f"Similarity: {1-distance:.3f}", (10, 120), 
+                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                 
-                display_image = cv2.cvtColor(display_image, cv2.COLOR_BGR2RGB)
                 return display_image, result_message
             
         return image, "❌ No matching student found"
     
+    def prepare_display_image(self, image):
+        """Convert image to display format"""
+        try:
+            if isinstance(image, str):
+                display_image = cv2.imread(image)
+            elif isinstance(image, np.ndarray):
+                display_image = image.copy()
+                if len(display_image.shape) == 3 and display_image.shape[2] == 3:
+                    display_image = cv2.cvtColor(display_image, cv2.COLOR_RGB2BGR)
+            else:
+                # Handle Gradio file object
+                display_image = np.array(Image.open(image))
+                if len(display_image.shape) == 3 and display_image.shape[2] == 3:
+                    display_image = cv2.cvtColor(display_image, cv2.COLOR_RGB2BGR)
+            
+            return display_image
+        except:
+            return None
+    
     def get_all_students(self):
         """Get all registered students"""
         self.cursor.execute('SELECT name, student_id, class FROM students ORDER BY name')
@@ -240,6 +270,12 @@ class FaceRecognitionSystem:
             ''')
         return self.cursor.fetchall()
 
+# Import torch after FaceNet imports to avoid conflicts
+try:
+    import torch
+except ImportError:
+    torch = None
+
 # Initialize the system
 face_system = FaceRecognitionSystem()