TheKnight115's picture
Update app.py
5d6a1b5 verified
raw
history blame
790 Bytes
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