Update app.py

app.py

@@ -1,11 +1,26 @@
-from ultralytics import YOLO
-model = YOLO("yolov8m.pt")
-
-cap = cv2.VideoCapture(tfile.name)
-while cap.isOpened():
-    ret, frame = cap.read()
-    if not ret:
-        break
-    results = model(frame)
-    annotated_frame = results[0].plot()
-    stframe.image(annotated_frame, channels="BGR")
+import streamlit as st
+import cv2
+import torch
+
+# Load your YOLO model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/yolov8_motorbike_detect_v2_Check.pt')
+
+# Function for YOLO detection
+def detect_objects(image):
+    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"])
+
+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