import gradio as gr
import cv2
import numpy as np
from PIL import Image
from yolov8 import YOLOv8Face

model = YOLOv8Face('weights/yolov8n-face.onnx')


def detect_and_blur_faces(image, blur_style):
    boxes, scores, classids, landmarks = model.detect(image)

    output_image = image.copy()

    for i, box in enumerate(boxes):
        x1, y1, w, h = [int(val) for val in box]
        x2, y2 = x1 + w, y1 + h

        face = output_image[y1:y2, x1:x2]
        blurred_face = cv2.GaussianBlur(face, (99, 99), 30)

        if blur_style == 'Oval':
            mask = np.zeros((y2-y1, x2-x1, 3), dtype=np.uint8)
            ellipse_mask = cv2.ellipse(mask, (w//2, h//2), (w//2, h//2), 0, 0, 360, (255, 255, 255), -1)
            blurred_face = np.where(ellipse_mask==np.array([255, 255, 255]), blurred_face, face)
        
        output_image[y1:y2, x1:x2] = blurred_face

    return output_image


# Set up the Gradio interface.
image_input = gr.inputs.Image(shape=(None, None))
blur_style = gr.inputs.Radio(['Rectangle', 'Oval'], label="Blur Style")
image_output = gr.outputs.Image(type='numpy')

gr.Interface(fn=detect_and_blur_faces, inputs=[image_input, blur_style], outputs=image_output, title="Face Detection and Blurring").launch()