app.py

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 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.to_csv("database/records.csv", index=False)

# 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

# Function to add face to database
def add_to_database(image, name):
    initialize_database()
    
    # 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_record = pd.DataFrame({
        'name': [name],
        'image_path': [image_path],
        '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)
    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

# 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)
    
    # 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 = []
    
    # 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_output_path = os.path.join(output_folder, f"face_{idx+1}.jpg")
        cv2.imwrite(face_output_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_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)
    
    return image, face_images, face_results

# Streamlit UI
st.title("Face Recognition System")

# Sidebar for database management
st.sidebar.header("Database Management")
upload_for_db = st.sidebar.file_uploader("Add face to database", type=["jpg", "jpeg", "png"], key="db_uploader")
if upload_for_db:
    person_name = st.sidebar.text_input("Enter person's name")
    if st.sidebar.button("Add to Database") and person_name:
        file_bytes = np.asarray(bytearray(upload_for_db.read()), dtype=np.uint8)
        img = cv2.imdecode(file_bytes, 1)
        image_path, message = add_to_database(img, person_name)
        if image_path:
            st.sidebar.success(f"Added {person_name} to database!")
        else:
            st.sidebar.error(message)

# Main interface
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)
    
    with st.spinner("Processing image..."):
        output_folder = "output"
        detected_image, face_images, face_results = detect_and_recognize_faces(image, output_folder)
    
    # Display results
    st.subheader("Recognition Results")
    st.image(cv2.cvtColor(detected_image, cv2.COLOR_BGR2RGB), 
             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
        for idx, (face, result, col) in enumerate(zip(face_images, face_results, cols)):
            with col:
                st.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), 
                        caption=f"Face {idx+1}", 
                        use_column_width=True)
                st.write(result)
    else:
        st.warning("No faces detected in the image.")

# Display database contents
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)
    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:
        st.sidebar.error(f"Error clearing database: {str(e)}")