Update app.py

app.py

@@ -4,73 +4,160 @@ import mediapipe as mp
 import streamlit as st
 from PIL import Image
 import numpy as np
+from deepface import DeepFace
+import pandas as pd
+from datetime import datetime
 
-# Initialize the MediaPipe Face Detection and Drawing utilities
+# Initialize MediaPipe
 mp_face_detection = mp.solutions.face_detection
 mp_drawing = mp.solutions.drawing_utils
 
-# Function to detect faces and save each detected face as an image
-def detect_and_save_faces(image, output_folder="output"):
+# 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', 'date_added'])
+        df.to_csv("database/records.csv", index=False)
+
+# Function to add face to database
+def add_to_database(image, name):
+    initialize_database()
+    
+    # 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],
+        '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
+
+# 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 for MediaPipe
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
+    
     # Detect faces
     with mp_face_detection.FaceDetection(min_detection_confidence=0.5) as face_detection:
         results = face_detection.process(image_rgb)
-
         face_images = []
-        # Draw face detections and save each face
+        face_results = []
+        
         if results.detections:
+            df = pd.read_csv("database/records.csv")
+            
             for idx, detection in enumerate(results.detections):
                 bboxC = detection.location_data.relative_bounding_box
                 h, w, _ = image.shape
-
-                # Get the bounding box coordinates
+                
+                # Get coordinates
                 x_min = int(bboxC.xmin * w)
                 y_min = int(bboxC.ymin * h)
                 box_width = int(bboxC.width * w)
                 box_height = int(bboxC.height * h)
-
-                # Extract the face from the image
+                
+                # Extract face
                 face_image = image[y_min:y_min + box_height, x_min:x_min + box_width]
                 face_images.append(face_image)
-
-                # Save the detected face
+                
+                # Save detected face
                 face_output_path = os.path.join(output_folder, f"face_{idx+1}.jpg")
                 cv2.imwrite(face_output_path, face_image)
-
-        # Draw bounding boxes on the original image
-        for detection in results.detections:
-            mp_drawing.draw_detection(image, detection)
-
-    return image, face_images
-
-# Streamlit UI for face detection
-st.title("Face Detection and Extraction")
-
-# File uploader
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+                
+                # Perform face recognition if database is not empty
+                if not df.empty:
+                    try:
+                        matches = []
+                        for _, row in df.iterrows():
+                            try:
+                                result = DeepFace.verify(
+                                    img1_path=face_output_path,
+                                    img2_path=row['image_path'],
+                                    model_name='VGG-Face',
+                                    distance_metric='cosine'
+                                )
+                                if result['verified']:
+                                    matches.append((row['name'], result['distance']))
+                            except Exception as e:
+                                continue
+                        
+                        if matches:
+                            # Get the best match (lowest distance)
+                            best_match = min(matches, key=lambda x: x[1])
+                            face_results.append(f"Match found: {best_match[0]}")
+                            # Draw name on image
+                            cv2.putText(image, best_match[0], (x_min, y_min - 10),
+                                      cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
+                        else:
+                            face_results.append("No match found")
+                    except Exception as e:
+                        face_results.append(f"Recognition error: {str(e)}")
+                
+                # Draw detection box
+                mp_drawing.draw_detection(image, detection)
+        
+        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"])
+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 = add_to_database(img, person_name)
+        st.sidebar.success(f"Added {person_name} to database!")
+
+# Main interface for face detection and recognition
+st.header("Face Detection and Recognition")
+uploaded_file = st.file_uploader("Choose an image for recognition", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Convert the uploaded file to an OpenCV image
+    # Convert uploaded file to image
     file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
     image = cv2.imdecode(file_bytes, 1)
-
-    # Run face detection and save faces
+    
+    # Process image
     output_folder = "output"
-    detected_image, face_images = detect_and_save_faces(image, output_folder)
-
-    # Display the original image with detected faces
-    st.subheader("Original Image with Detected Faces")
-    st.image(cv2.cvtColor(detected_image, cv2.COLOR_BGR2RGB), caption="Detected Faces", use_column_width=True)
-
-    # Display the extracted faces
-    st.subheader("Extracted Faces")
+    detected_image, face_images, face_results = detect_and_recognize_faces(image, output_folder)
+    
+    # Display results
+    st.subheader("Results")
+    st.image(cv2.cvtColor(detected_image, cv2.COLOR_BGR2RGB), 
+             caption="Detected and Recognized Faces", 
+             use_column_width=True)
+    
+    # Display extracted faces and recognition results
     if face_images:
-        for idx, face in enumerate(face_images):
-            st.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), caption=f"Face {idx+1}", use_column_width=False)
+        cols = st.columns(len(face_images))
+        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.write("No faces detected.")
+        st.write("No faces detected.")
+
+# Display database contents
+if st.sidebar.checkbox("Show Database Contents"):
+    try:
+        df = pd.read_csv("database/records.csv")
+        st.sidebar.dataframe(df)
+    except:
+        st.sidebar.write("Database is empty or not initialized.")