try:
    from ultralytics import YOLO
except ImportError:
    import os
    os.system("pip install ultralytics==8.3.15")
    from ultralytics import YOLO


import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np


# Load YOLOv8 face model (make sure it's in the same directory)
model = YOLO("yolov8n-face.pt")

# Image detection
def detect_faces(image, conf):
    results = model.predict(source=image, conf=conf)
    result_img = results[0].plot()
    return Image.fromarray(result_img)

# Webcam detection
def detect_faces_live(frame, conf):
    results = model.predict(source=frame, conf=conf)
    result_frame = results[0].plot()
    return Image.fromarray(result_frame)

# Gradio interface
title = "👤 Real-Time Face Detection with YOLOv8"
description = """
Upload an image or use your webcam to detect faces in real-time.  
Powered by **Ultralytics YOLOv8** and **Gradio**.
"""

with gr.Blocks(title=title) as app:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)

    conf_slider = gr.Slider(0.1, 1.0, value=0.5, label="Confidence Threshold")

    with gr.Tab("📸 Image Upload"):
        gr.Interface(
            fn=detect_faces,
            inputs=[gr.Image(type="pil", label="Upload Image"), conf_slider],
            outputs=gr.Image(label="Detected Faces"),
            live=False
        )

    with gr.Tab("🎥 Live Webcam"):
        gr.Interface(
            fn=detect_faces_live,
            inputs=[gr.Image(type="numpy", streaming=True), conf_slider],
            outputs=gr.Image(label="Detected Faces"),
            live=True
        )

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