| from ultralytics import YOLO |
| from PIL import Image |
| import numpy as np |
| import cv2 as cv |
|
|
|
|
| class detectPipeline(): |
| def __init__(self) -> None: |
| self.model = YOLO('yolo_v8_nano_model.pt') |
| self.class_names = {i: chr(65 + i) for i in range(26)} |
|
|
|
|
| def detect_signs(self, img_path: str): |
| |
| img = Image.open(img_path).convert('RGB') |
| img_array = np.array(img) |
|
|
| |
| detections = self.model(img_array)[0] |
| sign_detections = [] |
| for sign in detections.boxes.data.tolist(): |
| x1, y1, x2, y2, score, class_id = sign |
| sign_detections.append([int(x1), int(y1), int(x2), int(y2), score, int(class_id)]) |
| return sign_detections |
|
|
| def drawDetections2Image(self, img_path, detections): |
| img = Image.open(img_path).convert('RGB') |
| img = np.array(img) |
| for bbox in detections: |
| x1, y1, x2, y2, score, class_id = bbox |
| cv.rectangle(img, pt1=(x1, y1), pt2=(x2, y2), color=(0, 255, 0), thickness=25) |
| cv.putText(img, text=f'{self.class_names[class_id]} ({round(score*100, 2)}%)', org=(x1, y1-20), fontFace=cv.FONT_HERSHEY_SIMPLEX, fontScale=3.5, |
| color=(0, 0, 255), lineType=cv.LINE_AA, thickness=10) |
| img_detections = np.array(img) |
| return img_detections |
|
|
