Spaces:
Sleeping
Sleeping
File size: 1,342 Bytes
8be6929 |
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 |
from super_gradients.training import models
import cv2 as cv
import numpy as np
from PIL import Image
import matplotlib.pyplot as pl
class predictPipeline():
def __init__(self) -> None:
# Load model
self.model = models.get('yolo_nas_m',
num_classes=1,
checkpoint_path='yolo_nas_m_model.pth')
def detect(self, img_path):
image = Image.open(img_path).convert('RGB')
img_array = np.array(image)
preds = self.model.predict(img_array, conf=0.5)[0].prediction
bboxes_coordinates = []
for idx, bbox in enumerate(preds.bboxes_xyxy):
bboxes_coordinates.append([int(num) for num in bbox] + [round(preds.confidence[idx]*100, 2)])
return bboxes_coordinates
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 = bbox
cv.rectangle(img, pt1=(x1, y1), pt2=(x2, y2), color=(0, 255, 0), thickness=2)
cv.putText(img, text=f'{score}%', org=(x1, y1-2), fontFace=cv.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0, 0, 255), lineType=cv.LINE_AA)
img_detections = np.array(img)
return img_detections
|