import streamlit as st import cv2 import numpy as np from ultralytics import YOLO # Load the YOLO model model = YOLO('yolov8_Medium.pt') # Use the file name as it will be in the root directory def run_yolo(image): # Run the model results = model(image) return results def main(): st.title("Motorbike Violation Detection") # 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()