feat(model): revert inference code on YOLO8

Author: eugenos-programos

connection/ModelPredictionsReceiver.py

@@ -41,13 +41,10 @@ def predict(self, img: np.array) -> np.array:
             img[img.shape[0] // 2:, :, :]
         ]
         preds = [self._predict(imgs[0]), self._predict(imgs[1])]
-        print(preds)
         if len(preds[1]):
             preds[1][:, 1] += img.shape[0] // 2
             preds[1][:, 3] += img.shape[0] // 2
         result = np.append(*preds)
         if len(result.shape) == 1:
             result = np.expand_dims(result, 0)
-        print(result.shape)
-        print(result)
         return result

idle_model/IdleObjectDetectionModel.py

@@ -2,25 +2,24 @@
 import numpy as np
 from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
+from ultralytics import YOLO
 
 
 class IdleObjectDetectionModel:
     def __init__(self, path: str, conf_thresh, iou_thresh, classes) -> None:
-        self.model = DetrForObjectDetection.from_pretrained(path)
-        self.model.eval()
-        self.processor = DetrImageProcessor.from_pretrained(path)
+        self.model = YOLO(path)
         self.conf_thresh = conf_thresh
         self.iou_thresh = iou_thresh
         self.classes = classes
 
     @torch.no_grad()
     def __call__(self, img: np.array) -> np.array:
-        inputs = self.processor(images=img, return_tensors="pt")
-        outputs = self.model(**inputs)
-        target_sizes = torch.tensor([img.shape[:-1]])
-        results = self.processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=self.conf_thresh)[0]
-        pred = []
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            if self.model.config.id2label[label.item()] == 'cell phone':
-                pred.append([round(i, 2) for i in box.tolist()] + [score.item()])
-        return np.array(pred)
+        output = self.model(
+            source=img,
+            conf=self.conf_thresh,
+            iou=self.iou_thresh,
+            max_det=600,
+            classes=self.classes,
+            verbose=False
+        )[0].boxes
+        return output.xyxy

idle_model/configs/confs.json

@@ -4,6 +4,6 @@
     ],
     "iou_thres": 0.5,
     "conf_thres": 0.8,
-    "model_path": "./weights",
+    "model_path": "weights/yolov8l.pt",
     "port": 5001
 }

main.py

@@ -43,3 +43,4 @@
             img = utils.put_rectangle(img, preds[:, :4], preds[:, 4])
             reporter.send_report(reporter.create_report(img, str(start_tracking)))
         prev_preds = preds
+    time.sleep(2)