Update app.py

app.py

@@ -2,157 +2,149 @@ import os
 import cv2
 import mediapipe as mp
 import streamlit as st
+from PIL import Image
 import numpy as np
+import face_recognition
 import pandas as pd
 from datetime import datetime
 
-# Initialize face detector
-face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-recognizer = cv2.face.LBPHFaceRecognizer_create()
-
 # Initialize MediaPipe
 mp_face_detection = mp.solutions.face_detection
 mp_drawing = mp.solutions.drawing_utils
 
+# Function to create user database directory if it doesn't exist
 def initialize_database():
     if not os.path.exists("database"):
         os.makedirs("database")
     if not os.path.exists("database/records.csv"):
-        df = pd.DataFrame(columns=['name', 'image_path', 'label', 'date_added'])
+        df = pd.DataFrame(columns=['name', 'image_path', 'encoding', 'date_added'])
         df.to_csv("database/records.csv", index=False)
-    if not os.path.exists("database/recognizer"):
-        os.makedirs("database/recognizer")
 
-def prepare_training_data():
-    df = pd.read_csv("database/records.csv")
-    if df.empty:
-        return None
-    
-    faces = []
-    labels = []
-    label_names = {}
-    
-    for index, row in df.iterrows():
-        image_path = row['image_path']
-        label = row['label']
-        name = row['name']
-        label_names[label] = name
-        
-        if os.path.exists(image_path):
-            image = cv2.imread(image_path)
-            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-            face_rects = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
-            
-            for (x, y, w, h) in face_rects:
-                face = gray[y:y+h, x:x+w]
-                faces.append(face)
-                labels.append(label)
-    
-    if faces:
-        return faces, np.array(labels), label_names
+# Function to compute face encoding
+def get_face_encoding(image):
+    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    encodings = face_recognition.face_encodings(rgb_image)
+    if encodings:
+        return encodings[0]
     return None
 
-def train_recognizer():
-    data = prepare_training_data()
-    if data is not None:
-        faces, labels, label_names = data
-        recognizer.train(faces, labels)
-        recognizer.save("database/recognizer/trained_model.yml")
-        return True
-    return False
-
+# Function to add face to database
 def add_to_database(image, name):
     initialize_database()
     
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
-    
-    if len(faces) == 0:
+    # Get face encoding
+    encoding = get_face_encoding(image)
+    if encoding is None:
         return None, "No face detected in the image"
     
+    # Save image to database folder
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     image_path = f"database/{name}_{timestamp}.jpg"
     cv2.imwrite(image_path, image)
     
+    # Update records
     df = pd.read_csv("database/records.csv")
-    new_label = len(df) if not df.empty else 0
     new_record = pd.DataFrame({
         'name': [name],
         'image_path': [image_path],
-        'label': [new_label],
+        'encoding': [encoding.tolist()],
         'date_added': [datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
     })
     df = pd.concat([df, new_record], ignore_index=True)
     df.to_csv("database/records.csv", index=False)
-    
-    train_recognizer()
-    
     return image_path, "Successfully added to database"
 
-def get_status_color(confidence):
-    if confidence >= 70:
-        return (0, 255, 0)  # Green for high confidence
-    elif confidence >= 50:
-        return (0, 255, 255)  # Yellow for medium confidence
-    else:
-        return (0, 0, 255)  # Red for low confidence
+def find_best_match(face_encoding, known_encodings, known_names, tolerance=0.6):
+    """
+    Find the best match for a face encoding among known encodings
+    Returns name and confidence score of the best match, or None if no match found
+    """
+    if not known_encodings or not known_names:
+        return None, 0
+    
+    # Calculate face distances
+    face_distances = face_recognition.face_distance(known_encodings, face_encoding)
+    
+    if len(face_distances) > 0:
+        best_match_index = np.argmin(face_distances)
+        best_match_distance = face_distances[best_match_index]
+        
+        # Convert distance to confidence score (0-100%)
+        confidence = (1 - best_match_distance) * 100
+        
+        # Only return match if distance is below tolerance
+        if best_match_distance <= tolerance:
+            return known_names[best_match_index], confidence
+    
+    return None, 0
 
+# Function to detect faces and perform recognition
 def detect_and_recognize_faces(image, output_folder="output"):
     if not os.path.exists(output_folder):
         os.makedirs(output_folder)
     
-    try:
-        recognizer.read("database/recognizer/trained_model.yml")
-    except:
-        return image, [], ["No trained model available"]
-    
-    df = pd.read_csv("database/records.csv")
-    label_to_name = dict(zip(df['label'], df['name']))
+    # Convert image to RGB
+    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
+    # Get face locations and encodings
+    face_locations = face_recognition.face_locations(rgb_image, model="cnn")
+    face_encodings = face_recognition.face_encodings(rgb_image, face_locations)
     
     face_images = []
     face_results = []
+    recognition_details = []
     
-    for (x, y, w, h) in faces:
-        face_image = image[y:y+h, x:x+w]
+    # Load database
+    try:
+        df = pd.read_csv("database/records.csv")
+        if not df.empty:
+            known_encodings = [np.array(eval(enc)) for enc in df['encoding']]
+            known_names = df['name'].tolist()
+        else:
+            known_encodings = []
+            known_names = []
+    except Exception as e:
+        st.error(f"Error loading database: {str(e)}")
+        known_encodings = []
+        known_names = []
+    
+    # Process each detected face
+    for idx, (face_location, face_encoding) in enumerate(zip(face_locations, face_encodings)):
+        top, right, bottom, left = face_location
+        
+        # Extract and save face image
+        face_image = image[top:bottom, left:right]
         face_images.append(face_image)
         
-        face_path = os.path.join(output_folder, f"face_{len(face_images)}.jpg")
-        cv2.imwrite(face_path, face_image)
+        face_output_path = os.path.join(output_folder, f"face_{idx+1}.jpg")
+        cv2.imwrite(face_output_path, face_image)
         
-        face_gray = gray[y:y+h, x:x+w]
-        try:
-            label, confidence = recognizer.predict(face_gray)
-            confidence = min(100, 100 * (1 - (confidence / 300)))
-            
-            if confidence >= 50:
-                name = label_to_name.get(label, "Unknown")
-                result = f"Match: {name}\nConfidence: {confidence:.1f}%"
-                if confidence < 70:
-                    result += "\nPlease Verify (Not Sure)"
-            else:
-                result = "No match in database\nPlease Verify"
-            
-            color = get_status_color(confidence)
-                
-        except:
-            result = "Recognition error"
-            color = (0, 0, 255)
+        # Find best match
+        matched_name, confidence = find_best_match(face_encoding, known_encodings, known_names)
+        
+        if matched_name and confidence > 0:
+            result = f"Match: {matched_name}\nConfidence: {confidence:.1f}%"
+            color = (36, 255, 12)  # Green for match
+        else:
+            result = "No match in database"
+            color = (0, 0, 255)    # Red for no match
             
         face_results.append(result)
+        recognition_details.append({
+            'location': (left, top, right, bottom),
+            'color': color,
+            'text': result.split('\n')[0]
+        })
         
-        # Draw rectangle
-        cv2.rectangle(image, (x, y), (x+w, y+h), color, 2)
-        
-        # Draw text with multiple lines
-        text_lines = result.split('\n')
-        y_offset = y - 10
-        for line in text_lines:
-            cv2.putText(image, line, (x, y_offset),
-                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
-            y_offset -= 20  # Move up for next line
+    # Draw results on image
+    for detail in recognition_details:
+        left, top, right, bottom = detail['location']
+        # Draw rectangle around face
+        cv2.rectangle(image, (left, top), (right, bottom), detail['color'], 2)
+        # Draw text above face
+        cv2.putText(image, detail['text'], 
+                    (left, max(0, top - 10)),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, detail['color'], 2)
     
     return image, face_images, face_results
 
@@ -178,6 +170,7 @@ st.header("Face Detection and Recognition")
 uploaded_file = st.file_uploader("Choose an image for recognition", type=["jpg", "jpeg", "png"], key="recognition_uploader")
 
 if uploaded_file is not None:
+    # Load and process image
     file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
     image = cv2.imdecode(file_bytes, 1)
     
@@ -191,9 +184,10 @@ if uploaded_file is not None:
              caption="Processed Image", 
              use_column_width=True)
     
+    # Display extracted faces and results
     if face_images:
         st.subheader("Detected Faces")
-        cols = st.columns(min(len(face_images), 4))
+        cols = st.columns(min(len(face_images), 4))  # Limit to 4 columns max
         for idx, (face, result, col) in enumerate(zip(face_images, face_results, cols)):
             with col:
                 st.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), 
@@ -207,17 +201,20 @@ if uploaded_file is not None:
 if st.sidebar.checkbox("Show Database Contents"):
     try:
         df = pd.read_csv("database/records.csv")
-        st.sidebar.dataframe(df)
+        display_df = df.drop('encoding', axis=1)  # Don't show encodings
+        st.sidebar.dataframe(display_df)
     except:
         st.sidebar.write("Database is empty or not initialized.")
 
 # Add clear database button
 if st.sidebar.button("Clear Database"):
     try:
+        # Remove all files in database directory
         for file in os.listdir("database"):
             file_path = os.path.join("database", file)
             if os.path.isfile(file_path):
                 os.remove(file_path)
+        # Reinitialize database
         initialize_database()
         st.sidebar.success("Database cleared successfully!")
     except Exception as e: