import gradio as gr
import cv2
import mediapipe as mp
import numpy as np
from PIL import Image, ImageDraw

# Initialize MediaPipe Face Detection
mp_face_detection = mp.solutions.face_detection
mp_drawing = mp.solutions.drawing_utils

def detect_faces(image):
    # Convert to OpenCV format
    img = np.array(image.convert("RGB"))
    img_rgb = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    # Create a face detection object
    with mp_face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5) as face_detection:
        results = face_detection.process(img)

        # Draw results
        if results.detections:
            for detection in results.detections:
                bboxC = detection.location_data.relative_bounding_box
                h, w, _ = img.shape
                x, y, width, height = (
                    int(bboxC.xmin * w),
                    int(bboxC.ymin * h),
                    int(bboxC.width * w),
                    int(bboxC.height * h),
                )

                # Draw bounding box
                cv2.rectangle(img_rgb, (x, y), (x + width, y + height), (0, 255, 0), 3)

                # Add label "Unknown"
                cv2.rectangle(img_rgb, (x, y - 25), (x + 100, y), (0, 255, 0), -1)
                cv2.putText(img_rgb, "Unknown", (x + 5, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)

    # Convert back to PIL for Gradio
    result_image = Image.fromarray(cv2.cvtColor(img_rgb, cv2.COLOR_BGR2RGB))
    return result_image


# Gradio Interface
demo = gr.Interface(
    fn=detect_faces,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Image(label="Detected Faces"),
    title="Simple Face Detection (Lightweight)",
    description="Detects faces using MediaPipe and labels them as 'Unknown'. Future-ready for adding identity recognition.",
)

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