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