Spaces:
Running
Running
File size: 2,202 Bytes
be14aa6 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# Import machine learning
from ultralytics import YOLO
# Import computer vision
import cv2
from PIL import Image
def predict(detection_model: YOLO, image_path: str) -> Image:
"""
The function predicts with the given model and returns the image with the bounding boxes
Args:
detection_model (YOLO): The YOLO model
image_path (str): The path to the image
Returns:
Image: The image with the bounding boxes
"""
# Read the image
img = cv2.imread(image_path)
# Predict with the model
output = detection_model.predict(image_path, verbose=False)
# If output is not None then draw the bounding boxes
if output:
# Loop through the output
for bbox in output:
for i, p in enumerate(bbox):
# Converting the p to cpu
p = p.to("cpu")
# Conveting to numpy
p = p.numpy()
# Extracting the coords
coords = p.boxes.xyxy[0]
# Extracting the coords
xmin = coords[0]
ymin = coords[1]
xmax = coords[2]
ymax = coords[3]
# Converting to int
xmin = int(xmin)
ymin = int(ymin)
xmax = int(xmax)
ymax = int(ymax)
# Extracting the prob
prob = p.boxes.conf[0]
# Drawing the bounding box
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
# Drawing the text with the probability
cv2.putText(
img,
f"{prob:.2f}",
(xmin, ymin - 10),
cv2.FONT_HERSHEY_SIMPLEX,
0.9,
(36, 255, 12),
2,
)
# Converting the image to RGB
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Converting the image to PIL Image object
pil_image = Image.fromarray(img_rgb)
# Return the image
return pil_image
# If output is None then return the original image
return Image.open(image_path)
|