Update app.py

app.py

@@ -1,7 +1,40 @@
 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')
@@ -9,29 +42,60 @@ model = YOLO('models/best.pt')
 # Streamlit UI
 st.title("Object Detection with Ultralytics")
 
-st.markdown(
-    """
-    <style>
-    #MainMenu {visibility: hidden;}
-    footer {visibility: hidden;}
-    .ezrtsby0 {visibility: hidden;}
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
-# Upload image
-uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
+# 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:
-    # Read the uploaded image
-    image = Image.open(uploaded_file)
+    # 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
-    result = model.predict(source=image, conf=0.2,save=True)
-    img_save_path = "output/"
+        # 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
     
-    for r in result : r.save(filename= img_save_path + uploaded_file.name) 
 
-    # Display the output image
-    st.image(img_save_path+uploaded_file.name, caption="Detected Objects", use_column_width=True)
+    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)
+