|
4 | 4 | import numpy as np |
5 | 5 | import cv2 |
6 | 6 |
|
7 | | -def image_annotate(image:str) -> Image.Image: |
8 | | - # load the input image |
9 | | - image = cv2.imread(image) |
10 | 7 |
|
11 | | - # load pre-trained vision model |
12 | | - model = YOLO("yolo12s.pt") |
| 8 | +def image_annotate(image: str) -> Image.Image: |
| 9 | + # load the input image |
| 10 | + image = cv2.imread(image) |
13 | 11 |
|
14 | | - # run object detection on the image |
15 | | - result = model(image)[0] |
| 12 | + # load pre-trained vision model |
| 13 | + model = YOLO("yolo12s.pt") |
16 | 14 |
|
17 | | - # convert YOLO output to a Supervision-compatible detections format |
18 | | - detections = sv.Detections.from_ultralytics(result) |
| 15 | + # run object detection on the image |
| 16 | + result = model(image)[0] |
19 | 17 |
|
20 | | - # initialize a box annotator for drawing detection bounding boxes |
21 | | - box_annotator = sv.BoxAnnotator() |
| 18 | + # convert YOLO output to a Supervision-compatible detections format |
| 19 | + detections = sv.Detections.from_ultralytics(result) |
22 | 20 |
|
23 | | - # annotate the image with detected objects |
24 | | - annotated_image = box_annotator.annotate( |
25 | | - scene=image.copy(), |
26 | | - detections=detections) |
27 | | - |
28 | | - # convert BGR to RGB for correct display |
29 | | - annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB) |
| 21 | + # initialize a box annotator for drawing detection bounding boxes |
| 22 | + box_annotator = sv.BoxAnnotator() |
30 | 23 |
|
31 | | - # convert the annotated NumPy array (RGB) to a PIL Image object |
32 | | - annotated_pil_image = Image.fromarray(annotated_image_rgb) |
| 24 | + # annotate the image with detected objects |
| 25 | + annotated_image = box_annotator.annotate(scene=image.copy(), detections=detections) |
| 26 | + |
| 27 | + # convert BGR to RGB for correct display |
| 28 | + annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB) |
| 29 | + |
| 30 | + # convert the annotated NumPy array (RGB) to a PIL Image object |
| 31 | + annotated_pil_image = Image.fromarray(annotated_image_rgb) |
| 32 | + |
| 33 | + # return annotated image |
| 34 | + return annotated_pil_image |
33 | 35 |
|
34 | | - # return annotated image |
35 | | - return annotated_pil_image |
36 | 36 |
|
37 | 37 | if __name__ == "__main__": |
38 | | - annotate_img = image_annotate("/content/sample_data/Furry.png") |
39 | | - annotate_img.show() |
| 38 | + annotate_img = image_annotate("/content/sample_data/Furry.png") |
| 39 | + annotate_img.show() |
0 commit comments