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

# Load your YOLO model
model_path = 'best_hat_detection_model_100.pt'  # Your model file on Hugging Face Space
model = YOLO(model_path)

# Function to perform object detection and return the image with bounding boxes
def predict_safety_hat(image):
    # Convert image from numpy array (Gradio format) to a format YOLO can process
    image = np.array(image)
    
    # Run inference using the YOLO model
    results = model(image)
    
    # Get the image with predictions drawn (bounding boxes, etc.)
    result_image = results[0].plot()  # YOLOv8 returns an image with the bounding boxes drawn
    
    # Convert the result image back to PIL format for Gradio
    return Image.fromarray(result_image)

# Gradio Interface
inputs = gr.Image(label="Input kép", type="numpy")
outputs = gr.Image(label="Output kép")

title = "Safety Hat Detection | Biztonsági sisak detektálás"
description = "Tölts fel egy képet és a rendszer bekeretezi a sisakot"

gr.Interface(fn=predict_safety_hat, inputs=inputs, outputs=outputs, title=title, description=description).launch()