-
Notifications
You must be signed in to change notification settings - Fork 0
/
process_image.py
47 lines (34 loc) · 1.5 KB
/
process_image.py
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
37
38
39
40
41
42
43
44
45
46
47
# import necessary libraries
import cv2
from ultralytics import YOLO
# initialize the detector with weights
detector = YOLO("./model/traffic_sign_detector.pt", task="detect")
# path to the image
img_path = "./data/input/stop_sign.jpg"
# to visualize detections
# reading the image
image = cv2.imread(img_path)
# detector.predict returns a list of detection objects
detections = detector.predict(image)
for detection in detections:
# Get class indices and class names
class_ids = detection.boxes.cls # cls stores the class IDs
# Iterate over bounding boxes
for i, bbox in enumerate(detection.boxes):
# Get the bounding box coordinates
x1, y1, x2, y2 = bbox.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
# Draw the bounding box
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 3)
# Get class name using the class ID
class_id = int(class_ids[i])
class_name = detection.names[class_id] # detection.names holds the class names
# Display the class name on the image
cv2.putText(image, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Show the image with detection
cv2.imshow("Detected", image)
# Break on 'q' key
if cv2.waitKey(10000) & 0xFF == ord("q"):
break
# uncomment the following line to detect and save the annotated image without visualizing it
# results = detector(img_path, save=True)