TheKnight115 commited on
Commit
7d35b07
·
verified ·
1 Parent(s): a393422

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -18
app.py CHANGED
@@ -1,26 +1,35 @@
1
  import streamlit as st
2
  import cv2
3
- import torch
 
4
 
5
- # Load your YOLO model
6
- model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov8_Medium.pt')
7
 
8
- # Function for YOLO detection
9
- def detect_objects(image):
10
  results = model(image)
11
  return results
12
 
13
- # Streamlit UI
14
- st.title("Motorbike Violation Detection")
15
- uploaded_file = st.file_uploader("Upload an image or video", type=["jpg", "jpeg", "png", "mp4"])
16
 
17
- if uploaded_file is not None:
18
- if uploaded_file.type == "video/mp4":
19
- # Process video here
20
- st.video(uploaded_file)
21
- # Add video processing code
22
- else:
23
- # Process image
24
- image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
25
- results = detect_objects(image)
26
- st.image(results.render()[0]) # Display results
 
 
 
 
 
 
 
 
 
 
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()