Update app.py

app.py

@@ -1,49 +1,42 @@
 import streamlit as st
 import cv2
-import numpy as np
+import tempfile
 from ultralytics import YOLO
-from tempfile import NamedTemporaryFile
-import os
-
-# Initialize YOLOv8 model
-model = YOLO("yolov8n.pt")
+import numpy as np
 
-# Streamlit app title and creator information
-st.markdown("[Created by Engr. Hamesh Raj](https://www.linkedin.com/in/datascientisthameshraj/)")
-st.title("🎥 YOLOv8 Object Detection on Videos")
+# Title and Description
+st.title("🔍 YOLOv8 Object Detection on Video")
+st.write("Upload a video file to detect objects using the YOLOv8 model. You can download the processed video with bounding boxes around detected objects.")
 
 # Sidebar for video upload
-st.sidebar.header("Upload Video")
-uploaded_video = st.sidebar.file_uploader("Choose a video...", type=["mp4", "mov", "avi", "mkv"])
+uploaded_file = st.sidebar.file_uploader("Upload a video", type=["mp4", "avi", "mov"])
 
-if uploaded_video is not None:
+if uploaded_file is not None:
     # Save the uploaded video to a temporary file
-    temp_video = NamedTemporaryFile(delete=False)
-    temp_video.write(uploaded_video.read())
-    video_path = temp_video.name
+    temp_input_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+    temp_input_file.write(uploaded_file.read())
+    temp_input_file.close()
 
     # Display the uploaded video
-    st.sidebar.video(uploaded_video)
+    st.video(temp_input_file.name)
 
-    # Submit button to process the video
+    # Process video button
     if st.sidebar.button("Submit"):
-        st.subheader("Processing Video...")
+        # Load YOLOv8 model
+        model = YOLO("yolov8n.pt")
+
+        # Open the input video file
+        cap = cv2.VideoCapture(temp_input_file.name)
 
-        # Open the video file
-        cap = cv2.VideoCapture(video_path)
+        # Get video properties
         frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         fps = cap.get(cv2.CAP_PROP_FPS)
 
-        # Create a temporary file to save the output video
-        temp_output_video = NamedTemporaryFile(delete=False, suffix='.mp4')
-        output_video_path = temp_output_video.name
-
         # Define codec and create VideoWriter object
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
+        temp_output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+        out = cv2.VideoWriter(temp_output_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
 
-        # Process each frame of the video
         while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
@@ -52,35 +45,34 @@ if uploaded_video is not None:
             # Perform object detection
             results = model(frame)
 
-            # Draw bounding boxes on the frame
-            for result in results:
-                for box in result.boxes:
-                    x1, y1, x2, y2 = map(int, box.xyxy[0])
-                    conf = box.conf[0]
-                    cls = box.cls[0]
-                    label = f'{model.names[int(cls)]} {conf:.2f}'
-
-                    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
-                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
-
+            # Access detection results and draw bounding boxes
+            if results:
+                for result in results:
+                    boxes = result.boxes  # Access boxes attribute
+                    for box in boxes:
+                        x1, y1, x2, y2 = map(int, box.xyxy[0])
+                        conf = box.conf[0]
+                        cls = box.cls[0]
+                        label = f'{model.names[int(cls)]} {conf:.2f}'
+
+                        # Draw bounding boxes on the frame
+                        cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
+                        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+
+            # Write the frame with bounding boxes to the output video
             out.write(frame)
 
+        # Release resources
         cap.release()
         out.release()
 
         # Display the processed video
-        st.subheader("Processed Video")
-        st.video(output_video_path)
-
-        # Download button for the processed video
-        with open(output_video_path, "rb") as file:
-            st.download_button(
-                label="Download Processed Video",
-                data=file,
-                file_name="processed_video.mp4",
-                mime="video/mp4"
-            )
-
-        # Clean up temporary files
-        os.remove(video_path)
-        os.remove(output_video_path)
+        st.video(temp_output_file.name)
+
+        # Provide download link for processed video
+        st.sidebar.download_button(
+            label="Download Processed Video",
+            data=open(temp_output_file.name, "rb"),
+            file_name="processed_video.mp4",
+            mime="video/mp4"
+        )