Spaces:
Running
Running
File size: 1,075 Bytes
eaa4e30 7d35b07 eaa4e30 7d35b07 eaa4e30 7d35b07 eaa4e30 7d35b07 eaa4e30 7d35b07 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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()
|