File size: 855 Bytes
d19b3e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import cv2
import torch
import plasma.functional as f
from ultralytics import YOLO
class YOLORunner(f.Pipe):
def __init__(self, model: YOLO, image_size, conf_thrs, height_ratio, device, verbose):
super().__init__(image_size=image_size, conf_thrs=conf_thrs, device=device, height_ratio=height_ratio, verbose=verbose)
self.model = model
def run(self, image):
results = self.model.predict(source=image, imgsz=self.image_size, conf=self.conf_thrs, verbose=self.verbose)
results = results[0].boxes.xyxy.cpu().numpy().astype(int)
# expand the height of the boxes
height = results[:, 3] - results[:, 1]
results[:, 1] = (results[:, 1] - height * self.height_ratio).clip(0)
results[:, 3] = (results[:, 3] + height * self.height_ratio).clip(0, image.shape[0])
return results |