Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -4,13 +4,29 @@ import numpy as np
|
|
4 |
from ultralytics import YOLO
|
5 |
|
6 |
# Load the YOLO model
|
7 |
-
model = YOLO('yolov8_Medium.pt') #
|
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 |
|
@@ -23,8 +39,11 @@ def main():
|
|
23 |
image = np.array(cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1))
|
24 |
results = run_yolo(image)
|
25 |
|
26 |
-
#
|
27 |
-
|
|
|
|
|
|
|
28 |
|
29 |
elif uploaded_file.type == "video/mp4":
|
30 |
# Process the video
|
|
|
4 |
from ultralytics import YOLO
|
5 |
|
6 |
# Load the YOLO model
|
7 |
+
model = YOLO('yolov8_Medium.pt') # Ensure the model file is in the root directory of your Space
|
8 |
|
9 |
def run_yolo(image):
|
10 |
+
# Run the model on the image and get results
|
11 |
results = model(image)
|
12 |
return results
|
13 |
|
14 |
+
def process_results(results, image):
|
15 |
+
# Draw bounding boxes and labels on the image
|
16 |
+
boxes = results[0].boxes # Get boxes from results
|
17 |
+
for box in boxes:
|
18 |
+
# Get the box coordinates and label
|
19 |
+
x1, y1, x2, y2 = map(int, box.xyxy[0]) # Convert to integer coordinates
|
20 |
+
conf = box.conf[0] # Confidence score
|
21 |
+
cls = int(box.cls[0]) # Class index
|
22 |
+
label = model.names[cls] # Get class name from index
|
23 |
+
|
24 |
+
# Draw rectangle and label on the image
|
25 |
+
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2) # Blue box
|
26 |
+
cv2.putText(image, f"{label} {conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
|
27 |
+
|
28 |
+
return image
|
29 |
+
|
30 |
def main():
|
31 |
st.title("Motorbike Violation Detection")
|
32 |
|
|
|
39 |
image = np.array(cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1))
|
40 |
results = run_yolo(image)
|
41 |
|
42 |
+
# Process the results and draw boxes on the image
|
43 |
+
processed_image = process_results(results, image)
|
44 |
+
|
45 |
+
# Display the processed image
|
46 |
+
st.image(processed_image, caption='Detected Image', use_column_width=True)
|
47 |
|
48 |
elif uploaded_file.type == "video/mp4":
|
49 |
# Process the video
|