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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -4
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') # 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
 
@@ -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
- # 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
 
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