app.py

import gradio as gr
from ultralytics import YOLO
from PIL import Image
import os



# Load the trained YOLOv8 model
model = YOLO("best.pt")

# Define the prediction function
def predict(image):
    results = model(image)  # Run YOLOv8 model on the uploaded image
    results_img = results[0].plot()  # Get image with bounding boxes
    return Image.fromarray(results_img)

# Get example images from the images folder
def get_example_images():
    examples = []
    image_folder = "images"
    for filename in os.listdir(image_folder):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            examples.append(os.path.join(image_folder, filename))
    return examples

# Create Gradio interface
interface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Image(type="pil"),
    title="Helmet Detection with YOLO",
    description="Upload an image to detect helmets.",
    examples=get_example_images()
)

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