Update app.py

app.py

@@ -1,36 +1,61 @@
 import gradio as gr
+import face_recognition
+from PIL import Image, ImageDraw
 import os
-from PIL import Image
 import random
 
-data_folder = "data/"  # Pre-existing images
-existing_images = [os.path.join(data_folder, img) for img in os.listdir(data_folder) if img.endswith(('png', 'jpg', 'jpeg'))]
+# Path to the data folder containing sample images
+DATA_FOLDER = "data/"
 
-def process_image(image_option, uploaded_image):
+# Load sample images from the data folder
+sample_images = [os.path.join(DATA_FOLDER, img) for img in os.listdir(DATA_FOLDER) if img.lower().endswith(('png', 'jpg', 'jpeg'))]
+
+def recognize_faces(image_option, uploaded_image):
     """
-    Process the selected image (either uploaded or from data folder)
+    Perform face recognition on the selected image.
     """
     if image_option == "Upload My Image" and uploaded_image is not None:
         img = uploaded_image
-        msg = "Uploaded image processed!"
+        status = "Processed the uploaded image."
     else:
-        img_path = random.choice(existing_images)  # Pick a random existing image
+        img_path = random.choice(sample_images)
         img = Image.open(img_path)
-        msg = f"Using pre-existing image: {img_path}"
-    
-    return img, msg
+        status = f"Processed a sample image: {os.path.basename(img_path)}"
+
+    # Convert PIL image to numpy array
+    img_array = face_recognition.load_image_file(img)
+
+    # Find all face locations and face encodings in the image
+    face_locations = face_recognition.face_locations(img_array)
+    face_encodings = face_recognition.face_encodings(img_array, face_locations)
+
+    # Convert back to PIL image for drawing
+    pil_image = Image.fromarray(img_array)
+    draw = ImageDraw.Draw(pil_image)
+
+    # Iterate over each face found in the image
+    for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
+        # Draw a box around the face
+        draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255), width=2)
+
+    # Clean up the drawing library
+    del draw
 
+    return pil_image, status
+
+# Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("## Face Recognition Test Run")
-    
-    image_option = gr.Radio(["Upload My Image", "Use Sample Image"], label="Select an Option")
-    uploaded_image = gr.Image(label="Upload Image (if selected)", type="pil", interactive=True)
-    
+
+    with gr.Row():
+        image_option = gr.Radio(["Upload My Image", "Use Sample Image"], label="Select an Option")
+        uploaded_image = gr.Image(label="Upload Image (if selected)", type="pil", interactive=True)
+
     submit = gr.Button("Process Image")
     output_image = gr.Image(label="Processed Image")
     status = gr.Textbox(label="Status")
-    
-    submit.click(process_image, inputs=[image_option, uploaded_image], outputs=[output_image, status])
-    
+
+    submit.click(recognize_faces, inputs=[image_option, uploaded_image], outputs=[output_image, status])
+
 if __name__ == "__main__":
     demo.launch()