Merge pull request #28 from 5sControl/dev

Dev

connection/ModelPredictionsReceiver.py

@@ -11,7 +11,6 @@ def __init__(self, server_url: str, logger: Logger) -> None:
         self._server_url = server_url
         self._logger = logger
         self._port = configs["port"]
-        print("server_url - ", self._server_url)
 
     @staticmethod
     def _convert_image2bytes(image: np.array, format='PNG') -> io.BytesIO:

model_image/ObjectDetectionModel.py

@@ -4,15 +4,16 @@
 
 
 class YOLOv8ObjDetectionModel:
-    def __init__(self, path: str, conf_thresh: float, iou_thresh: float, classes: list) -> None:
-        self.model = YOLO(path)
+    def __init__(self, model_path: str, conf_thresh: float, iou_thresh: float, classes: list, img_size: int) -> None:
+        self._model = YOLO(model_path)
         self.conf_thresh = conf_thresh
         self.iou_thresh = iou_thresh
         self.classes = classes
+        self.img_size = img_size
 
     @torch.no_grad()
     def __call__(self, img: np.array, classes: list = None) -> list:
-        results = self.model(
+        results = self._model(
             source=img,
             conf=self.conf_thresh,
             iou=self.iou_thresh,

model_image/app.py

@@ -1,8 +1,7 @@
 from PIL import Image
 from flask import Flask, jsonify, request
-from flask_configs.load_configs import *
+from flask_configs.load_configs import bottle_model_configs, box_model_configs, human_model_configs
 from ObjectDetectionModel import YOLOv8ObjDetectionModel
-from YOLORObjectDetectionModel import YOLORObjectDetectionModel
 import numpy as np
 import colorlog
 import logging
@@ -24,9 +23,9 @@
 logger.propagate = False
 
 app = Flask(__name__)
-human_model = YOLORObjectDetectionModel(HUMAN_MODEL_PATH, CONF_PATH, CONF_THRES, IOU_THRES, CLASSES, logger=logger)
-box_model = YOLOv8ObjDetectionModel(BOX_MODEL_PATH, CONF_THRES, IOU_THRES, CLASSES)
-bottle_model = YOLORObjectDetectionModel(HUMAN_MODEL_PATH, CONF_PATH, 0.2, IOU_THRES, [39], logger=logger)
+human_model = YOLOv8ObjDetectionModel(**human_model_configs)
+box_model = YOLOv8ObjDetectionModel(**box_model_configs)
+bottle_model = YOLOv8ObjDetectionModel(**bottle_model_configs)
 
 
 convert_bytes2image = lambda bytes: np.array(Image.open(io.BytesIO(bytes)), dtype=np.uint8)

model_image/flask_configs/flask_confs.json

@@ -1,12 +1,30 @@
 {
-    "classes": [
-        0
-    ],
-    "iou_thres": 0.5,
-    "conf_thres": 0.5,
-    "config_path": "weigths/yolor_csp_x.cfg",
-    "box_detect_model": "weigths/min_max_v0.3.10.pt",
-    "human_detect_model": "weigths/min_max_v0.4.4b.pt",
-    "img_size": 640,
+    "box_model": {
+        "classes": [
+            0
+        ],
+        "iou_thresh": 0.5,
+        "conf_thresh": 0.5,
+        "model_path": "weigths/min_max_v0.3.10_box.pt",
+        "img_size": 640
+    },
+    "human_model": {
+        "classes": [
+            0
+        ],
+        "iou_thresh": 0.5,
+        "conf_thresh": 0.5,
+        "model_path": "weigths/min_max_v0.5.2_bottle.pt",
+        "img_size": 640
+    },
+    "bottle_model": {
+        "classes": [
+            0
+        ],
+        "iou_thresh": 0.5,
+        "conf_thresh": 0.6,
+        "model_path": "weigths/min_max_v0.5.2_human.pt",
+        "img_size": 640
+    },
     "port": 5000
 }

model_image/flask_configs/load_configs.py

@@ -3,11 +3,6 @@
 
 with open("flask_configs/flask_confs.json", "r") as conf:
     configs = json.load(conf)
-    CONF_THRES = configs.get("conf_thres")
-    IOU_THRES = configs.get("iou_thres")
-    BOX_MODEL_PATH = configs.get("box_detect_model")
-    HUMAN_MODEL_PATH = configs.get("human_detect_model")
-    CONF_PATH = configs.get("config_path")
-    CLASSES = configs.get("classes")
-    IMG_SIZE = configs.get("img_size")
-    PORT = configs.get("port")
+    box_model_configs = configs["box_model"]
+    human_model_configs = configs["human_model"]
+    bottle_model_configs = configs["bottle_model"]

model_image/requirements.txt

@@ -3,4 +3,5 @@ httplib2==0.22.0
 numpy==1.22.3
 opencv_python==4.7.0.72
 ultralytics==8.0.136
-Flask==2.2.2
+Flask==2.2.2
+py-cpuinfo==9.0.0