NightPrince/Helmet-Detect-model

1import gradio as gr2from ultralytics import YOLO3from PIL import Image4 5# Load the trained YOLOv8 model6model = YOLO("best.pt")7 8# Define the prediction function9def predict(image):10    results = model(image)  # Run YOLOv8 model on the uploaded image11    results_img = results[0].plot()  # Get image with bounding boxes12    return Image.fromarray(results_img)13 14# Create Gradio interface15interface = gr.Interface(16    fn=predict, 17    inputs=gr.Image(type="pil"), 18    outputs=gr.Image(type="pil"),19    title="Helmet Detection with YOLOv8",20    description="Upload an image to detect helmets."21)22 23# Launch Gradio app24interface.launch(share=True)25