Merge pull request #13 from KernelA/new-params

New parameters added :minor:

Author: KernelA

.gitignore

@@ -140,4 +140,5 @@ cython_debug/
 *.jpg
 test.ipynb
 **/nms/*.c
-wheelhouse/
+wheelhouse/
+class_colors.json

processing.py

@@ -1,11 +1,14 @@
 import argparse
 import os
+import logging
 
 from yolo_models.processing.detection_processing import DetectorProcess
 from yolo_models.log_set import init_logging
 
 
 def main(args):
+    logger = logging.getLogger()
+
     with DetectorProcess(args.checkpoint_path,
                          args.shared_update_mem_name,
                          args.shared_params_mem_name,
@@ -14,16 +17,29 @@ def main(args):
                          args.image_height,
                          args.num_channels,
                          args.image_type) as det:
+
+        if os.path.exists(args.class_colormap):
+            logger.info("Load class colors from '%s'", args.class_colormap)
+            det.load_class_colormap(args.class_colormap)
+        else:
+            logger.info("Cannot find a file with class colors. Create it with defaults colors.")
+            det.save_class_colormap(args.class_colormap)
+
         det.start_processing()
 
 
+def check_file(path: str):
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"Cannot find: '{path}'")
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
-    parser.add_argument("-iw", "--image_width", type=int, default=640,
-                        help="Width of image to transfer between processes")
-    parser.add_argument("-ih", "--image_height", type=int, default=640,
-                        help="Height of image to transfer between processes")
+    parser.add_argument("-iw", "--image_width", type=int, default=1280,
+                        help="A maximum width of image to transfer between processes")
+    parser.add_argument("-ih", "--image_height", type=int, default=1280,
+                        help="A maximum height of image to transfer between processes")
     parser.add_argument("-c", "--num_channels", type=int, default=3,
                         help="A number of channels in image. 3 for RGB")
     parser.add_argument("--image_type", type=str,
@@ -36,13 +52,12 @@ def main(args):
                         help="Name of shared memory for transfering of parameters")
     parser.add_argument("-p", "--checkpoint_path", type=str,
                         required=True, help="A path to model .pt")
-    parser.add_argument("--log_config", type=str, default="log_settings.yaml",
+    parser.add_argument("--log_config", type=str, default=None,
                         help="A path to settings for logging")
+    parser.add_argument("--class_colormap", type=str, default="class_colors.json",
+                        help="A path to json with colors for each class")
 
     args = parser.parse_args()
-
-    if not os.path.isfile(args.checkpoint_path):
-        raise FileNotFoundError(f"Cannot find: '{args.checkpoint_path}'")
-
+    check_file(args.checkpoint_path)
     init_logging(log_config=args.log_config)
     main(args)

requirements.inference.txt

@@ -1,3 +1,2 @@
 opencv-python-headless>=4.7.0.0
 onnxruntime>=1.11.0
-matplotlib~=3.7

setup.py

@@ -46,10 +46,10 @@ def get_req(filepath: str):
 setup(install_requires=get_req("requirements.inference.txt"),
       version=get_version(),
       packages=find_packages(include=[f"{PACKAGE_NAME}*"]),
+      package_data={"": [f"{PACKAGE_NAME}/log_set/*.yaml"]},
       extras_require={
-          "torch": get_req("requirements.torch.gpu.txt")
-},
-    ext_modules=cythonize(exts,
-                          compiler_directives={"language_level": "3"},
-                          )
-)
+          "torch": get_req("requirements.torch.gpu.txt")},
+      ext_modules=cythonize(exts,
+                            compiler_directives={"language_level": "3"},
+                            )
+      )

tests/test_color_conv.py

@@ -0,0 +1,11 @@
+import pytest
+
+from yolo_models.processing.color_utils import hex_to_rgb, rgb_to_hex
+
+
+@pytest.mark.parametrize("r", [0, 128, 255])
+@pytest.mark.parametrize("g", [0, 128, 255])
+@pytest.mark.parametrize("b", [0, 128, 255])
+def test_rgb_seq(r, g, b):
+    color = (r, g, b)
+    assert hex_to_rgb(rgb_to_hex(color)) == color

touch_designer.py

@@ -1,15 +1,16 @@
+import time
 from multiprocessing import shared_memory
 
-import numpy as np
 import cv2
-
-from yolo_models.processing.info import BufferStates, States, ParamsIndex
+import numpy as np
+from yolo_models.processing.info import (BufferStates, DrawInfo, ParamsIndex,
+                                         States)
 
 # Sync with external process
 SHARED_MEM_PARAMS_LIST = shared_memory.ShareableList(name="params")
 
-WIDTH = SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_WIDTH]
-HEIGHT = SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_HEIGHT]
+MAX_IMAGE_WIDTH = SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_WIDTH]
+MAX_IMAGE_HEIGHT = SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_HEIGHT]
 DTYPE = SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_DTYPE]
 NUM_CHANNELS = SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_CHANNELS]
 EXIT = False
@@ -21,28 +22,33 @@
 SHARED_MEM_ARRAY = shared_memory.SharedMemory(
     name=SHARED_MEM_PARAMS_LIST[ParamsIndex.SHARED_ARRAY_MEM_NAME], create=False)
 
-ARRAY = np.ndarray((WIDTH, HEIGHT, NUM_CHANNELS), dtype=DTYPE, buffer=SHARED_MEM_ARRAY.buf)
+ARRAY = np.ndarray((MAX_IMAGE_WIDTH, MAX_IMAGE_HEIGHT, NUM_CHANNELS),
+                   dtype=DTYPE, buffer=SHARED_MEM_ARRAY.buf)
 
 
 def onSetupParameters(scriptOp):
     page = scriptOp.appendCustomPage("Detection")
     p = page.appendFloat("Nms", label="Intersection ratio between bboxes")
-    p.default = 0.5
     p.min = 0
     p.max = 1
-    p = page.appendFloat("Score", label="Minimum score for object")
     p.default = 0.5
+    p = page.appendFloat("Score", label="Minimum score for object")
     p.min = 0.1
     p.max = 1
+    p.default = 0.5
     p = page.appendInt("Maxk", label="Maximum number of objects to detect based on score")
-    p.default = 5
     p.min = 0
     p.max = 1000
+    p.default = 5
     p = page.appendFloat("Eta", label="Filtering")
-    p.default = 1.0
     p.min = 0
     p.clampMin = True
     p.max = 1000
+    p.default = 1.0
+    p = page.appendToggle("Drawtext", label="Draw text labels")
+    p.default = False
+    p = page.appendToggle("Drawscore", label="Draw score labels")
+    p.default = False
     return
 
 
@@ -71,31 +77,65 @@ def onCook(scriptOp):
 
     video_in = scriptOp.inputs[0]
     # By default, the image is flipped up. We flip it early
-    frame = video_in.numpyArray(delayed=True, writable=False)
+    image = video_in.numpyArray(delayed=True, writable=False)
 
-    if frame is None:
+    if image is None:
+        scriptOp.clear()
         return
 
-    if COLOR_CONVERSION is not None:
-        image = cv2.cvtColor(frame, COLOR_CONVERSION)
+    if image.shape[0] > MAX_IMAGE_HEIGHT:
+        debug("Too large image height")
+        scriptOp.clear()
+        return
 
-    image = cv2.resize(image, (WIDTH, HEIGHT))
+    if image.shape[1] > MAX_IMAGE_WIDTH:
+        debug("Too large image width")
+        scriptOp.clear()
+        return
+
+    if COLOR_CONVERSION is not None:
+        image = cv2.cvtColor(image, COLOR_CONVERSION)
 
+    SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_HEIGHT] = image.shape[0]
+    SHARED_MEM_PARAMS_LIST[ParamsIndex.IMAGE_WIDTH] = image.shape[1]
     SHARED_MEM_PARAMS_LIST[ParamsIndex.SCORE_THRESH] = scriptOp.par.Score.eval()
     SHARED_MEM_PARAMS_LIST[ParamsIndex.IOU_THRESH] = scriptOp.par.Nms.eval()
     SHARED_MEM_PARAMS_LIST[ParamsIndex.TOP_K] = scriptOp.par.Maxk.eval()
     SHARED_MEM_PARAMS_LIST[ParamsIndex.ETA] = scriptOp.par.Eta.eval()
-    np.copyto(ARRAY, image)
+    draw_info = DrawInfo.DRAW_BBOX
+
+    if scriptOp.par.Drawtext.eval():
+        draw_info |= DrawInfo.DRAW_TEXT
+
+    if scriptOp.par.Drawscore.eval():
+        draw_info |= DrawInfo.DRAW_CONF
+
+    SHARED_MEM_PARAMS_LIST[ParamsIndex.DRAW_INFO] = int(draw_info)
 
+    ARRAY[:image.shape[0], :image.shape[1]] = image
     SHARED_MEM_UPDATE_STATES.buf[BufferStates.SERVER] = States.READY_SERVER_MESSAGE.value[0]
 
+    start_time = time.monotonic()
+
+    is_skip = False
+
     while not EXIT and SHARED_MEM_UPDATE_STATES.buf[BufferStates.SERVER_ALIVE] == States.IS_SERVER_ALIVE.value[0] and SHARED_MEM_UPDATE_STATES.buf[BufferStates.CLIENT] != States.READY_CLIENT_MESSAGE.value[0]:
-        pass
+        time.sleep(1e-3)
+        elapsed = time.monotonic() - start_time
+
+        if elapsed > 1:
+            debug("Too long processing copy frame as is")
+            is_skip = True
+            break
 
     if SHARED_MEM_UPDATE_STATES.buf[BufferStates.SERVER_ALIVE] != States.IS_SERVER_ALIVE.value[0]:
+        scriptOp.clear()
         raise ValueError("Server process died")
 
-    scriptOp.copyNumpyArray(ARRAY)
-    SHARED_MEM_UPDATE_STATES.buf[BufferStates.CLIENT] = States.NULL_STATE.value[0]
+    if is_skip:
+        scriptOp.clear()
+        return
 
+    scriptOp.copyNumpyArray(ARRAY[:image.shape[0], :image.shape[1]])
+    SHARED_MEM_UPDATE_STATES.buf[BufferStates.CLIENT] = States.NULL_STATE.value[0]
     return

yolo_models/log_set/init_log.py

@@ -1,12 +1,16 @@
 import os
 from typing import Optional
 import logging.config
+import pkg_resources
 
 import yaml
 
 
 def init_logging(log_config: Optional[str] = None, log_env_var: str = "LOG_CONFIG"):
-    path_to_config = log_config if log_config is not None else os.environ[log_env_var]
+    path_to_config = log_config if log_config is not None else os.environ.get(log_env_var)
+
+    if path_to_config is None:
+        path_to_config = pkg_resources.resource_filename(__name__, "log_settings.yaml")
 
     with open(path_to_config, "rb") as file:
         log_config = yaml.safe_load(file)

log_settings.yaml renamed to yolo_models/log_set/log_settings.yaml

@@ -4,7 +4,8 @@
 
 import numpy as np
 
-from .info import ParamsIndex, BufferStates, States
+from .info import ParamsIndex, BufferStates, States, DrawInfo
+
 
 class BaseProcessServer:
     def __init__(self,
@@ -15,7 +16,7 @@ def __init__(self,
                  image_height: int,
                  num_channels: int,
                  image_dtype):
-        self._logging  = logging.getLogger("server_processing")
+        self._logging = logging.getLogger("server_processing")
         dtype_size = np.dtype(image_dtype).itemsize
         self._image_size_bytes = image_height * image_with * num_channels * dtype_size
         self._image_width = image_with
@@ -30,9 +31,21 @@ def __init__(self,
         self._params_shared_mem_name = params_shared_mem_name
         self._array_shared_mem_name = array_shared_mem_name
 
+    def _get_shared_mem_update(self, create: bool):
+        return shared_memory.SharedMemory(
+            name=self._update_shared_mem_name, create=create, size=len(BufferStates))
+
+    def _get_shared_mem_array(self, create: bool):
+        return shared_memory.SharedMemory(
+            name=self._array_shared_mem_name, create=create, size=self._image_size_bytes)
+
     def init_mem(self):
         assert self._sh_mem_update is None,  "Memory already initialized"
-        self._sh_mem_update = shared_memory.SharedMemory(name=self._update_shared_mem_name, create=True, size=len(BufferStates))
+        try:
+            self._sh_mem_update = self._get_shared_mem_update(True)
+        except FileExistsError:
+            self._logging.warning("Cannot create new shared memory. Use existed")
+            self._sh_mem_update = self._get_shared_mem_update(False)
 
         params = [None] * len(ParamsIndex)
         params[ParamsIndex.ETA] = 1.0
@@ -45,11 +58,17 @@ def init_mem(self):
         params[ParamsIndex.SHARED_ARRAY_MEM_NAME] = self._array_shared_mem_name
         params[ParamsIndex.SHARD_STATE_MEM_NAME] = self._update_shared_mem_name
         params[ParamsIndex.IMAGE_DTYPE] = self._image_dtype
+        params[ParamsIndex.DRAW_INFO] = int(DrawInfo.DRAW_BBOX)
 
         self._sh_mem_params = shared_memory.ShareableList(
             name=self._params_shared_mem_name, sequence=params
         )
-        self._sh_mem_array = shared_memory.SharedMemory(name=self._array_shared_mem_name, create=True, size=self._image_size_bytes)
+
+        try:
+            self._sh_mem_array = self._get_shared_mem_array(True)
+        except FileExistsError:
+            self._sh_mem_array = self._get_shared_mem_array(False)
+
         self._shared_array = np.ndarray(
             (self._image_height, self._image_width, self._num_channels),
             dtype=self._image_dtype, buffer=self._sh_mem_array.buf)
@@ -95,7 +114,14 @@ def start_processing(self):
 
             self._sh_mem_update.buf[BufferStates.SERVER] = States.NULL_STATE.value[0]
 
-            self.process(self._shared_array, self._sh_mem_params)
+            image_height = self._sh_mem_params[ParamsIndex.IMAGE_HEIGHT]
+            image_width = self._sh_mem_params[ParamsIndex.IMAGE_WIDTH]
+
+            actual_image = self._shared_array[:image_height, :image_width]
+
+            self.process(actual_image, self._sh_mem_params)
+            self._shared_array[:image_height, :image_width] = actual_image
+
             self._sh_mem_update.buf[BufferStates.CLIENT] = States.READY_CLIENT_MESSAGE.value[0]
 
     def process(self, image: np.ndarray, params: list):

yolo_models/processing/base_processing.py

@@ -0,0 +1,18 @@
+from typing import Tuple
+import re
+
+HEX_REGEX = re.compile("^#[0-9a-f]{6}$", re.IGNORECASE)
+
+
+def check_regex_color(hex_color: str):
+    if HEX_REGEX.match(hex_color) is None:
+        raise ValueError(f"Incorrect hex color: '{hex_color}'")
+
+
+def rgb_to_hex(rgb: Tuple[int]) -> str:
+    return f"#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}"
+
+
+def hex_to_rgb(hex_color: str) -> Tuple[int]:
+    check_regex_color(hex_color)
+    return tuple(int(hex_color[i: i + 2], 16) for i in range(1, len(hex_color), 2))