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()