File size: 1,383 Bytes
79f8119 |
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 |
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):
# Data Preprocessing
img = Image.open(img_path).convert('RGB')
img_array = np.array(img)
# Making detections using YOLOv8 Nano
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
|