Merge pull request #9 from 5sControl/dev

Update algorithm

machine_control_utils/HTTPLIB2Capture.py

@@ -24,9 +24,6 @@ def get_snapshots(self, n_images=1) -> list:
                 continue
             imgs[counter] = img
             counter += 1
-            time.sleep(0.2)
-        if n_images == 1:
-            return imgs[0] 
         return imgs
 
     def try_get_snapshot(self):

machine_control_utils/model.py

@@ -1,28 +1,29 @@
 import numpy as np
 from yolo_utils.openvino_functional import detect
 from openvino.runtime import Core
-from deep_sort_realtime.deepsort_tracker import DeepSort
-
-
-class Model:
-    def __init__(self, object_detector_path: str, person_conf: float, n_images: int):
-        self.object_detector = YoloDetector(object_detector_path, person_conf)
-        self.object_tracker = DeepSort(max_age=3)
-        self.n_images = n_images
-
-    def predict(self, img):
-        boxes, confs = self.object_detector.predict_boxes(img)
-        detections: list = []
-        for (x1, y1, x2, y2), conf in zip(boxes, confs):
-            detections.append(((x1, y1, x2 - x1, y2 - y1), conf, "person"))
-        tracks = self.object_tracker.update_tracks(detections, frame=img)
-        return tracks
-
-
-class YoloDetector(Model):
-    def __init__(self, model_path: str, person_conf: float):
+# from deep_sort_realtime.deepsort_tracker import DeepSort
+
+
+# class Model:
+#     def __init__(self, object_detector_path: str, person_conf: float, n_images: int):
+#         self.object_detector = YoloDetector(object_detector_path, person_conf)
+#         self.object_tracker = DeepSort(max_age=3)
+#         self.n_images = n_images
+#
+#     def predict(self, img):
+#         boxes, confs = self.object_detector.predict_boxes(img)
+#         detections: list = []
+#         for (x1, y1, x2, y2), conf in zip(boxes, confs):
+#             detections.append(((x1, y1, x2 - x1, y2 - y1), conf, "person"))
+#         tracks = self.object_tracker.update_tracks(detections, frame=img)
+#         return tracks
+
+
+class YoloDetector():
+    def __init__(self, model_path: str, person_conf: float, n_images: int):
         self.model = self.init_model(model_path)
         self.person_conf = person_conf
+        self.n_images = n_images
 
     @staticmethod
     def init_model(model_path):
@@ -31,7 +32,7 @@ def init_model(model_path):
         model = core.compile_model(read_model)
         return model
 
-    def predict_boxes(self, img):
+    def predict(self, img):
         res = detect(np.array(img), self.model)[0]["det"]
         boxes, confs = [], []
         if len(res):

main.py

@@ -1,4 +1,4 @@
-from machine_control_utils.model import Model
+from machine_control_utils.model import YoloDetector
 from run import run_camera, run_local, run_example
 import warnings
 from dotenv import load_dotenv
@@ -9,10 +9,10 @@
 
 
 model_path = "model/yolov8x.xml"
-person_conf = 0.45
-n_images = 10
+person_conf = 0.5
+n_images = 1
 
-model = Model(model_path, person_conf, n_images)
+model = YoloDetector(model_path, person_conf, n_images)
 
 
 run_camera(model)

run.py

@@ -7,11 +7,11 @@
     get_intersection,
     get_areas,
 )
-from machine_control_utils.model import Model
+from machine_control_utils.model import YoloDetector
 
 
-def run_machine_control(model: Model, img, area_values):
-    tracks = model.predict(img)
+def run_machine_control(model: YoloDetector, img, area_values):
+    boxes, confs = model.predict(img)
 
     for i, a_val in enumerate(area_values):
         x1, y1, x2, y2 = a_val.coords
@@ -22,10 +22,7 @@ def run_machine_control(model: Model, img, area_values):
                     cv2.FONT_HERSHEY_SIMPLEX, 1, (100, 0, 200), 1, cv2.LINE_AA)
 
         in_area = False
-        for track in tracks:
-            if not track.is_confirmed():
-                continue
-            x1, y1, x2, y2 = map(lambda x: int(x), track.to_ltrb())
+        for (x1, y1, x2, y2), conf in zip(boxes, confs):
             small_box = x1, y1, x2, y2
             small_box_plot = x1, y1, x2 - x1, y2 - y1
 
@@ -62,7 +59,7 @@ def run_machine_control(model: Model, img, area_values):
     # cv2.imshow("img", img)
 
 
-def run_local(model: Model):
+def run_local(model: YoloDetector):
     cap = cv2.VideoCapture(0)
     succes, img = cap.read()
     area_values = get_areas(img.shape)
@@ -76,14 +73,14 @@ def run_local(model: Model):
     cv2.destroyAllWindows()
 
 
-def run_camera(model: Model):
+def run_camera(model: YoloDetector):
     username = os.environ.get("username")
     password = os.environ.get("password")
     source = os.environ.get("camera_url")
 
     dataset = HTTPLIB2Capture(source, username=username, password=password)
 
-    img = dataset.get_snapshots()
+    img = dataset.get_snapshots()[0]
     area_values = get_areas(img.shape)
 
     while True:
@@ -92,7 +89,7 @@ def run_camera(model: Model):
             run_machine_control(model, img, area_values)
 
 
-def run_example(model: Model):
+def run_example(model: YoloDetector):
     img = cv2.imread("test_image.jpg")
     area_values = get_areas(img.shape)
     run_machine_control(model, img, area_values)

yolo_utils/openvino_functional.py

@@ -1,8 +1,8 @@
 from typing import Tuple, Dict
 import cv2
 
-from legacy.yolo_utils.plotting import colors
-from legacy.yolo_utils import ops
+from yolo_utils.plotting import colors
+from yolo_utils import ops
 import torch
 import numpy as np
 import random