Update app.py

app.py

@@ -2,149 +2,157 @@ 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', 'encoding', 'date_added'])
+        df = pd.DataFrame(columns=['name', 'image_path', 'label', 'date_added'])
         df.to_csv("database/records.csv", index=False)
+    if not os.path.exists("database/recognizer"):
+        os.makedirs("database/recognizer")
 
-# 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]
+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
     return None
 
-# Function to add face to database
+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
+
 def add_to_database(image, name):
     initialize_database()
     
-    # Get face encoding
-    encoding = get_face_encoding(image)
-    if encoding is None:
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
+    
+    if len(faces) == 0:
         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],
-        'encoding': [encoding.tolist()],
+        'label': [new_label],
         '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 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
+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
 
-# 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)
     
-    # Convert image to RGB
-    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    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']))
     
-    # 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)
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
     
     face_images = []
     face_results = []
-    recognition_details = []
     
-    # 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]
+    for (x, y, w, h) in faces:
+        face_image = image[y:y+h, x:x+w]
         face_images.append(face_image)
         
-        face_output_path = os.path.join(output_folder, f"face_{idx+1}.jpg")
-        cv2.imwrite(face_output_path, face_image)
+        face_path = os.path.join(output_folder, f"face_{len(face_images)}.jpg")
+        cv2.imwrite(face_path, face_image)
         
-        # 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_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)
             
         face_results.append(result)
-        recognition_details.append({
-            'location': (left, top, right, bottom),
-            'color': color,
-            'text': result.split('\n')[0]
-        })
         
-    # 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)
+        # 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
     
     return image, face_images, face_results
 
@@ -170,7 +178,6 @@ 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)
     
@@ -184,10 +191,9 @@ 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))  # Limit to 4 columns max
+        cols = st.columns(min(len(face_images), 4))
         for idx, (face, result, col) in enumerate(zip(face_images, face_results, cols)):
             with col:
                 st.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), 
@@ -201,20 +207,17 @@ if uploaded_file is not None:
 if st.sidebar.checkbox("Show Database Contents"):
     try:
         df = pd.read_csv("database/records.csv")
-        display_df = df.drop('encoding', axis=1)  # Don't show encodings
-        st.sidebar.dataframe(display_df)
+        st.sidebar.dataframe(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: