Spaces:
Sleeping
Sleeping
File size: 3,096 Bytes
7c9bdfb 5c12faa 7c9bdfb 5c12faa 7c9bdfb 5c12faa 7c9bdfb 5c12faa 7c9bdfb 5c12faa 7c9bdfb 5c12faa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import streamlit as st
from ultralytics import YOLO
import os
from os import listdir, remove
from PIL import Image
import cv2
import requests
def download_file(url, destination_folder):
"""
Download a file from a URL and save it to a local destination.
Args:
url (str): The URL of the file to download.
destination_folder (str): The local folder where the file should be saved.
"""
# Create the destination folder if it doesn't exist
if not os.path.exists(destination_folder):
os.makedirs(destination_folder)
# Extract filename from URL
filename = url.split('/')[-1]
destination_path = os.path.join(destination_folder, filename)
# Download the file
with open(destination_path, 'wb') as f:
response = requests.get(url)
f.write(response.content)
return destination_path
url = 'https://yolo-v8-models.s3.ap-south-1.amazonaws.com/best.pt'
destination_folder = 'models'
downloaded_file = download_file(url, destination_folder)
# Load the model
model = YOLO('models/best.pt')
# Streamlit UI
st.title("Object Detection with Ultralytics")
# Upload image or video
uploaded_file = st.file_uploader("Choose an image or video", type=["jpg", "jpeg", "png","mp4"])
# Demo section
st.header("Demo")
col1, col2 = st.columns(2)
# Display image in first column
with col1:
st.image("demoimg.jpeg", caption="Annotated Image", use_column_width=True)
# Display video in second column with adjusted width
with col2:
st.write(f'<div style="width: {300};height :500">', unsafe_allow_html=True)
st.video("demovideo.mp4")
st.write('</div>', unsafe_allow_html=True)
if uploaded_file is not None:
# Check if the uploaded file is a video
if uploaded_file.type.startswith("video/"):
# Progress bar to show the progress of object detection
progress_bar = st.progress(0)
st.header(uploaded_file.name)
# Perform object detection
with st.spinner('Performing object detection...'):
for percent_complete in range(100):
result = model.predict(source=uploaded_file, conf=0.2, save=True ,stream=True)
for i in result : i.save("video.mp4")
progress_bar.progress(percent_complete + 1)
st.success(f"Video saved successfully ")
# Perform object detection
else:
# Read the uploaded image
image = Image.open(uploaded_file)
img_name = "converted_image.jpg"
image.save(img_name)
# Perform object detection
result = model.predict(source=img_name, conf=0.2, save=True)
# Save the output image
img_save_path = "output/"
os.makedirs(img_save_path, exist_ok=True)
for r in result:
r.save(filename=os.path.join(img_save_path, uploaded_file.name))
st.success("Detected Object")
# Display the output image
st.image(os.path.join(img_save_path, uploaded_file.name), caption="Detected Objects", use_column_width=True)
|