Merge pull request #16 from 5sControl/dev

Dev
5sControl · Sep 20, 2023 · b6b5ff0 · b6b5ff0
2 parents cdc33d8 + c00c0d3
commit b6b5ff0
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 17 deletions.
diff --git a/idle_model/IdleObjectDetectionModel.py b/idle_model/IdleObjectDetectionModel.py
@@ -4,10 +4,11 @@
 from yolor.utils.datasets import letterbox
 from yolor.utils.general import non_max_suppression, scale_coords
 import time
+from logging import Logger
 
 
 class IdleObjectDetectionModel:
-    def __init__(self, model_path: str, config_path: str, conf_thresh: float, iou_thresh: float, classes: list, logger) -> None:
+    def __init__(self, model_path: str, config_path: str, conf_thresh: float, iou_thresh: float, classes: list, logger: Logger, **kwargs) -> None:
         self.model, self.device = get_model(model_path, config_path)
         self.conf_thresh = conf_thresh
         self.iou_thresh = iou_thresh

diff --git a/idle_model/app.py b/idle_model/app.py
@@ -1,6 +1,6 @@
 from PIL import Image
 from flask import Flask, jsonify, request
-from configs.load_configs import *
+from configs import configs
 from IdleObjectDetectionModel import IdleObjectDetectionModel
 import numpy as np
 import io
@@ -25,12 +25,8 @@
 
 app = Flask(__name__)
 model = IdleObjectDetectionModel(
-    MODEL_PATH,
-    CONF_PATH,
-    CONF_THRES,
-    IOU_THRES,
-    CLASSES,
-    logger
+    **configs,
+    logger=logger
 )
 
 convert_bytes2image = lambda bytes: np.array(Image.open(io.BytesIO(bytes)), dtype=np.uint8)

diff --git a/idle_model/configs/__init__.py b/idle_model/configs/__init__.py
@@ -0,0 +1 @@
+from .load_configs import configs
diff --git a/idle_model/configs/confs.json b/idle_model/configs/confs.json
@@ -2,9 +2,9 @@
     "classes": [
         67
     ],
-    "iou_thres": 0.5,
-    "conf_thres": 0.6,
+    "iou_thresh": 0.5,
+    "conf_thresh": 0.7,
     "model_path": "weights/idle_v0.4.4.pt",
-    "conf_path": "weights/yolor_csp_x.cfg",
+    "config_path": "weights/yolor_csp_x.cfg",
     "port": 5001
 }
diff --git a/idle_model/configs/load_configs.py b/idle_model/configs/load_configs.py
@@ -3,9 +3,3 @@
 
 with open("configs/confs.json", "r") as conf:
     configs = json.load(conf)
-CONF_THRES = configs.get("conf_thres")
-CONF_PATH = configs.get("conf_path")
-IOU_THRES = configs.get("iou_thres")
-MODEL_PATH = configs.get("model_path")
-CLASSES = configs.get("classes")
-PORT = configs.get("port")