Spaces:
Running
Running
| # 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) | |