Update app.py

app.py

@@ -1,26 +1,35 @@
 import streamlit as st
 import cv2
-import torch
+import numpy as np
+from ultralytics import YOLO
 
-# Load your YOLO model
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov8_Medium.pt')
+# Load the YOLO model
+model = YOLO('yolov8_Medium.pt')  # Use the file name as it will be in the root directory
 
-# Function for YOLO detection
-def detect_objects(image):
+def run_yolo(image):
+    # Run the model
     results = model(image)
     return results
 
-# Streamlit UI
-st.title("Motorbike Violation Detection")
-uploaded_file = st.file_uploader("Upload an image or video", type=["jpg", "jpeg", "png", "mp4"])
+def main():
+    st.title("Motorbike Violation Detection")
 
-if uploaded_file is not None:
-    if uploaded_file.type == "video/mp4":
-        # Process video here
-        st.video(uploaded_file)
-        # Add video processing code
-    else:
-        # Process image
-        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
-        results = detect_objects(image)
-        st.image(results.render()[0])  # Display results
+    # Upload file
+    uploaded_file = st.file_uploader("Choose an image or video...", type=["jpg", "jpeg", "png", "mp4"])
+
+    if uploaded_file is not None:
+        if uploaded_file.type in ["image/jpeg", "image/png", "image/jpg"]:
+            # Process the image
+            image = np.array(cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1))
+            results = run_yolo(image)
+            
+            # Display the results
+            st.image(results[0].plot(), caption='Detected Image', use_column_width=True)
+
+        elif uploaded_file.type == "video/mp4":
+            # Process the video
+            video_bytes = uploaded_file.read()
+            st.video(video_bytes)
+
+if __name__ == "__main__":
+    main()