Spaces:
Runtime error
Runtime error
| 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)}") |