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

# Load the trained YOLOv8 model for human detection
model = YOLO("best.pt")  # Replace with the correct path to your human detection model

# Define prediction function
def detect_humans(image):
    results = model(image)  # Perform inference
    result_img = results[0].plot()  # Draw bounding boxes
    return Image.fromarray(result_img.astype(np.uint8))  # Return image with boxes

# Build Gradio interface
interface = gr.Interface(
    fn=detect_humans,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(type="pil", label="Detected Humans"),
    title="Human Detection (YOLOv8)",
    description="Upload an image. The model will detect humans using your trained YOLOv8 model."
)

# Launch the interface
interface.launch(debug=True)