Create model.py

model.py

@@ -0,0 +1,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