import torch
import gradio as gr
from PIL import Image
import io

# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='fire.pt')  # Load custom model

def detect_objects(image):
    # Run the YOLOv5 model
    results = model(image)

    # Save the results image
    results_image = results.render()[0]  # Render returns a list, we take the first element

    # Convert the numpy array result to an image
    results_image = Image.fromarray(results_image)

    # Save to a buffer
    buf = io.BytesIO()
    results_image.save(buf, format='JPEG')
    byte_im = buf.getvalue()

    return results_image

# Gradio interface
interface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="YOLOv5 Image Detection",
    description="Upload an image to detect objects using YOLOv5."
)

# Launch the Gradio app
interface.launch()