import gradio as gr
import face_recognition
from PIL import Image, ImageDraw
import os
import random

# Path to the data folder containing sample images
DATA_FOLDER = "data/"

# 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):
    """
    Perform face recognition on the selected image.
    """
    if image_option == "Upload My Image" and uploaded_image is not None:
        img = uploaded_image
        status = "Processed the uploaded image."
    else:
        img_path = random.choice(sample_images)
        img = Image.open(img_path)
        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")

    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(recognize_faces, inputs=[image_option, uploaded_image], outputs=[output_image, status])

if __name__ == "__main__":
    demo.launch()