Spaces:
Sleeping
Sleeping
| 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/modelv2/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) | |