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