Update app.py

app.py

@@ -4,7 +4,7 @@ import numpy as np
 from ultralytics import YOLO
 
 # Load the YOLO model
-model = YOLO('yolov8_motorbikes_detection_s.pt')  # Ensure the model file is in the root directory of your Space
+model = YOLO('yolov8_Medium.pt')  # Ensure the model file is in the root directory of your Space
 
 def run_yolo(image):
     # Run the model on the image and get results
@@ -69,7 +69,8 @@ def process_video(uploaded_file):
     
     # Create a video writer to save the processed frames
     height, width, _ = frames[0].shape
-    out = cv2.VideoWriter('processed_video.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
+    output_path = 'processed_video.mp4'
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
 
     for frame in frames:
         out.write(frame)  # Write each processed frame to the video
@@ -80,7 +81,13 @@ def process_video(uploaded_file):
     progress_bar.progress(100)
     progress_text.text('Processing: 100%')
     st.success('Video processing complete!')
-
+    
+    st.video(output_path)  # Display the processed video
+    
+    # Create a download button for the processed video
+    with open(output_path, 'rb') as f:
+        video_bytes = f.read()
+    st.download_button(label='Download Processed Video', data=video_bytes, file_name='processed_video.mp4', mime='video/mp4')
 
 def main():
     st.title("Motorbike Violation Detection")
@@ -103,7 +110,7 @@ def main():
         elif uploaded_file.type == "video/mp4":
             # Process the video
             process_video(uploaded_file)  # Process the video and save the output
-            st.video('processed_video.mp4')  # Display the processed video
+            
 
 if __name__ == "__main__":
     main()