Upload 3 files

127262b verified about 2 years ago

1.38 kB

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"
    )

	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"
	)