Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -1,26 +1,35 @@
|
|
1 |
import streamlit as st
|
2 |
import cv2
|
3 |
-
import
|
|
|
4 |
|
5 |
-
# Load
|
6 |
-
model =
|
7 |
|
8 |
-
|
9 |
-
|
10 |
results = model(image)
|
11 |
return results
|
12 |
|
13 |
-
|
14 |
-
st.title("Motorbike Violation Detection")
|
15 |
-
uploaded_file = st.file_uploader("Upload an image or video", type=["jpg", "jpeg", "png", "mp4"])
|
16 |
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import streamlit as st
|
2 |
import cv2
|
3 |
+
import numpy as np
|
4 |
+
from ultralytics import YOLO
|
5 |
|
6 |
+
# Load the YOLO model
|
7 |
+
model = YOLO('yolov8_Medium.pt') # Use the file name as it will be in the root directory
|
8 |
|
9 |
+
def run_yolo(image):
|
10 |
+
# Run the model
|
11 |
results = model(image)
|
12 |
return results
|
13 |
|
14 |
+
def main():
|
15 |
+
st.title("Motorbike Violation Detection")
|
|
|
16 |
|
17 |
+
# Upload file
|
18 |
+
uploaded_file = st.file_uploader("Choose an image or video...", type=["jpg", "jpeg", "png", "mp4"])
|
19 |
+
|
20 |
+
if uploaded_file is not None:
|
21 |
+
if uploaded_file.type in ["image/jpeg", "image/png", "image/jpg"]:
|
22 |
+
# Process the image
|
23 |
+
image = np.array(cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1))
|
24 |
+
results = run_yolo(image)
|
25 |
+
|
26 |
+
# Display the results
|
27 |
+
st.image(results[0].plot(), caption='Detected Image', use_column_width=True)
|
28 |
+
|
29 |
+
elif uploaded_file.type == "video/mp4":
|
30 |
+
# Process the video
|
31 |
+
video_bytes = uploaded_file.read()
|
32 |
+
st.video(video_bytes)
|
33 |
+
|
34 |
+
if __name__ == "__main__":
|
35 |
+
main()
|