1import cv2 as cv
2import gdown
3from ultralytics import YOLO
4
5# Download YOLO Checkpoint from google drive
6url = (
7 "https://drive.google.com/file/d/1q-CNPubqyz4OQaPsH5nc5eS2Buy-Fkug/view?usp=sharing"
8)
9output = "yolo11n.pt"
10md5 = "md5:261474e91b15f5ef14a63c21ce6c0cbb"
11gdown.cached_download(url, output, hash=md5, fuzzy=True)
12
13# Load the image from disk
14image = cv.imread("image.png")
15
16# Load the YOLO-Model
17model = YOLO("yolo11n.pt") # pretrained YOLO11n model
18
19# Run inference on the image
20results = model([image])
21
22# Iterate over all results (only one in this case)
23for result in results:
24 # Iterate over all boxes for current result
25 for box in result.boxes:
26 # Get the name of the detection
27 cls_id = box.cls.item()
28 cls_name = result.names[cls_id]
29
30 # Get the coordinates
31 x1, y1, x2, y2 = box.xyxy.cpu()[0]
32
33 # Draw a nice frame for visual reference
34 cv.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
35 cv.rectangle(
36 image, (int(x1), int(y1 - 16)), (int(x2), int(y1)), (0, 0, 200), -1
37 ) # Filled bar on top
38 cv.rectangle(
39 image, (int(x1), int(y1 - 16)), (int(x2), int(y1)), (0, 0, 255), 2
40 ) # Filled bar on top
41
42 # Put the class label on top
43 cv.putText(
44 image, cls_name, (int(x1 + 4), int(y1 - 4)), 1, 1, (255, 255, 255), 1
45 )
46
47# Display the image and wait for user input
48cv.imshow("Image", image)
49cv.imwrite("result.png", image)
50cv.waitKey(0)