import torch
import streamlit as st
from PIL import Image
import io

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

# Streamlit interface
st.title("YOLOv5 Image Detection")
st.write("Upload an image to detect objects using YOLOv5")

# File uploader for image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Convert the uploaded file to a PIL Image
    image = Image.open(uploaded_file)

    # 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()

    # Display the input and output images side by side
    st.image(image, caption='Uploaded Image', use_column_width=True)
    st.image(results_image, caption='Detected Image', use_column_width=True)

    # Provide a download button for the output image
    st.download_button(
        label="Download Output Image",
        data=byte_im,
        file_name="output.jpg",
        mime="image/jpeg"
    )