Spaces:
Running
Running
import streamlit as st | |
import cv2 | |
import torch | |
# Load your YOLO model | |
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov8_Medium.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 | |