Merge pull request #1 from KernelA/shared-memory

Shared memory

.gitignore

@@ -136,4 +136,7 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
-.vscode/
+.vscode/
+*.jpg
+test.ipynb
+**/nms/*.c

MANIFEST.in

@@ -0,0 +1 @@
+global-exclude *.c

README.md

@@ -1,25 +1,59 @@
-# video-detection
+# YOLO video detection in the TouchDesigner
 
 ## Description
 
 
 ## Requirements
 
-1. Python 3.10 or higher.
+1. Python 3.9 or higher. It is more preferable to have same version as in the TouchDesigner.
 
 ## How to run
 
-Install dependencies:
+### Install dependencies
+
+For GPU:
+```
+pip install -r. /requirements.gpu.txt
+```
+s
+For CPU:
 ```
-pip install -r. /requirements.txt
+pip install -r. /requirements.cpu.txt
 ```
 
 For development:
 ```
-pip install -r ./requirements.txt -r ./requirements.dev.txt
+pip install -r ./requirements.dev.txt
+```
+
+### Compile Cython extension
+
+You need C compiler to compile extension. [See](https://docs.cython.org/en/latest/src/quickstart/install.html#installing-cython)
+
+```
+python -m build --no-isolation -w
+```
+
+Install binary distribution:
+```
+pip install ./dist/*.whl
 ```
 
+## Usecases
+
+### Video detection
 Run:
 ```
 python ./main.py -c <checkpoint_path> -i <video_path> -o <video_path>
 ```
+
+**If you run script in the directory with source code see modification of `sys.path` in the beginner of the file.**
+
+### With TouchDesigner
+
+Run server processing:
+```
+python ./processing.py -p <path_to_model>
+```
+
+Open `touch_designer.py` in the TouchDesigner as SCript TOP.

log_settings.yaml

@@ -1,10 +1,10 @@
 version: 1
 formatters:
   default_console_thread:
-    (): video_detection.log_set.utc_formatter.UTCFormatter
+    (): yolo_models.log_set.utc_formatter.UTCFormatter
     format: '%(asctime)s %(levelname)s %(threadName)s %(module)s %(funcName)s %(message)s'
   default_console_process:
-    (): video_detection.log_set.utc_formatter.UTCFormatter
+    (): yolo_models.log_set.utc_formatter.UTCFormatter
     format: '%(asctime)s %(levelname)s %(processName)s %(module)s %(funcName)s %(message)s'
 handlers:
   console:
@@ -13,7 +13,7 @@ handlers:
     formatter: default_console_thread
     stream: ext://sys.stdout
 loggers:
-  custom:
+  server_processing:
     level: DEBUG
     handlers: [console]
     propagate: no

main.py

@@ -8,7 +8,7 @@
 from ultralytics import YOLO
 from ultralytics.yolo.engine.results import Results
 
-from video_detection.log_set import init_logging
+from yolo_models.log_set import init_logging
 
 
 def main(args):

processing.py

@@ -0,0 +1,210 @@
+import argparse
+import enum
+from multiprocessing import shared_memory
+import time
+import sys
+import pathlib
+import os
+import logging
+
+for i, path in enumerate(sys.path):
+    if pathlib.Path(path).name == "site-packages":
+        index = i
+        break
+else:
+    raise RuntimeError("Cannot find required dir")
+
+sys.path.insert(0, sys.path[index])
+del index
+
+from matplotlib import cm
+import numpy as np
+import torch
+import cv2
+
+from yolo_models.detection import PyTorchYoloDetector
+from yolo_models.log_set import init_logging
+
+# Sync with touch designer code
+@enum.unique
+class BufferStates(enum.IntEnum):
+    SERVER = 0
+    CLIENT = 1
+    SERVER_ALIVE = 2
+
+@enum.unique
+class States(bytes, enum.Enum):
+    NULL_STATE = b'0'
+    READY_SERVER_MESSAGE = b'1'
+    READY_CLIENT_MESSAGE = b'2'
+    IS_SERVER_ALIVE = b'3'
+
+@enum.unique
+class ParamsIndex(enum.IntEnum):
+    IOU_THRESH = 0
+    SCORE_THRESH = 1
+    TOP_K = 2
+    ETA = 3
+    IMAGE_WIDTH = 4
+    IMAGE_HEIGHT = 5
+    IMAGE_CHANNELS = 6
+    SHARED_ARRAY_MEM_NAME = 7
+    SHARD_STATE_MEM_NAME = 8
+    IMAGE_DTYPE = 9
+
+
+class ProcessingServer:
+    def __init__(self,
+                 update_shared_mem_name: str,
+                 params_shared_mem_name: str,
+                 array_shared_mem_name: str,
+                 image_with: int,
+                 image_height: int,
+                 num_channels: int,
+                 image_dtype):
+        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
+        self._image_height = image_height
+        self._num_channels = num_channels
+        self._image_dtype = image_dtype
+        self._sh_mem_update = None
+        self._sh_mem_params = None
+        self._sh_mem_array = None
+        self._shared_array = None
+        self._update_shared_mem_name = update_shared_mem_name
+        self._params_shared_mem_name = params_shared_mem_name
+        self._array_shared_mem_name = array_shared_mem_name
+
+    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))
+
+        params = [None] * len(ParamsIndex)
+        params[ParamsIndex.ETA] = 1.0
+        params[ParamsIndex.IOU_THRESH] = 0.5
+        params[ParamsIndex.SCORE_THRESH] = 0.5
+        params[ParamsIndex.TOP_K] = 0
+        params[ParamsIndex.IMAGE_WIDTH] = self._image_width
+        params[ParamsIndex.IMAGE_HEIGHT] = self._image_height
+        params[ParamsIndex.IMAGE_CHANNELS] = self._num_channels
+        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
+
+        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)
+        self._shared_array = np.ndarray(
+            (self._image_height, self._image_width, self._num_channels),
+            dtype=self._image_dtype, buffer=self._sh_mem_array.buf)
+
+    def __enter__(self):
+        self.init_mem()
+
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self._logging.info("Stop processing")
+        self.dispose()
+
+    def dispose(self):
+        if self._sh_mem_update is not None:
+            self._logging.info("Free update shared memory")
+            self._sh_mem_update.buf[BufferStates.SERVER_ALIVE] = States.NULL_STATE.value[0]
+            self._sh_mem_update.close()
+            self._sh_mem_update.unlink()
+
+        del self._shared_array
+
+        if self._sh_mem_array is not None:
+            self._logging.info("Free array shared memory")
+            self._sh_mem_array.close()
+            self._sh_mem_array.unlink()
+
+        if self._sh_mem_params is not None:
+            self._logging.info("Free params shared memory")
+            self._sh_mem_params.shm.close()
+            self._sh_mem_params.shm.unlink()
+
+    def start_processing(self):
+        self._sh_mem_update.buf[BufferStates.SERVER] = States.NULL_STATE.value[0]
+        self._sh_mem_update.buf[BufferStates.CLIENT] = States.NULL_STATE.value[0]
+        self._sh_mem_update.buf[BufferStates.SERVER_ALIVE] = States.IS_SERVER_ALIVE.value[0]
+
+        self._logging.info("Awaiting message")
+
+        while True:
+            while self._sh_mem_update.buf[BufferStates.SERVER] != States.READY_SERVER_MESSAGE.value[0]:
+                time.sleep(1e-3)
+
+            self._sh_mem_update.buf[BufferStates.SERVER] = States.NULL_STATE.value[0]
+
+            self.process(self._shared_array, self._sh_mem_params)
+            self._sh_mem_update.buf[BufferStates.CLIENT] = States.READY_CLIENT_MESSAGE.value[0]
+
+    def process(self, image: np.ndarray, params: list):
+        pass
+
+
+class DetectorSever(ProcessingServer):
+    def __init__(self,
+                 path_to_model: str,
+                 update_shared_mem_name: str,
+                 params_shared_mem_name: str,
+                 array_shared_mem_name: str,
+                 image_with: int,
+                 image_height: int,
+                 num_channels: int,
+                image_dtype):
+        super().__init__(update_shared_mem_name, params_shared_mem_name, array_shared_mem_name, image_with, image_height, num_channels, image_dtype)
+        self._detector = PyTorchYoloDetector(path_to_model,
+                                             device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+                                             trace=True,
+                                             numpy_post_process=True)
+
+    def process(self, image: np.ndarray, params: list):
+        det_info = self._detector.predict(
+            image,
+            score_threshold=params[ParamsIndex.SCORE_THRESH],
+            nms_threshold=params[ParamsIndex.IOU_THRESH],
+            max_k=params[ParamsIndex.TOP_K],
+            eta=params[ParamsIndex.ETA])
+
+        for xyxy in det_info.xyxy_boxes:
+            cv2.rectangle(image, xyxy[:2], xyxy[2:], (255, 0, 0), thickness=1)
+
+def main(args):
+    with DetectorSever(args.checkpoint_path,
+                  args.shared_update_mem_name,
+                  args.shared_params_mem_name,
+                  args.shared_array_mem_name,
+                  args.image_width,
+                  args.image_height,
+                  args.num_channels,
+                  args.image_type) as det:
+        det.start_processing()
+
+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("-c", "--num_channels", type=int, default=3, help="A number of channels in image. 3 for RGB")
+    parser.add_argument("--image_type", type=str, choices=["uint8", "float32"], default="float32", help="Image dtype")
+    parser.add_argument("--shared_array_mem_name", type=str, default="array", help="Name of shared memory for array")
+    parser.add_argument("--shared_update_mem_name", type=str, default="update_info", help="Name of shared memory for transfering information about updates")
+    parser.add_argument("--shared_params_mem_name", type=str, default="params", 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", help="A path to settings for logging")
+
+    args = parser.parse_args()
+
+    if not os.path.isfile(args.checkpoint_path):
+        raise FileNotFoundError(f"Cannot find: '{args.checkpoint_path}'")
+
+    init_logging(log_config=args.log_config)
+    main(args)
+

pyproject.toml

@@ -1,14 +1,18 @@
 [build-system]
 requires = [
-    "setuptools>=51.1",
-    "wheel>=0.36"
+    "setuptools>=63.0",
+    "wheel>=0.40"
 ]
 build-backend = "setuptools.build_meta"
 
+[tool.isort]
+skip = ["setup.py"]
+
 [project]
-name = "video_detection"
+name = "yolo_models"
 authors =[ { name="None"  }]
-requires-python = ">=3.10"
+requires-python = ">=3.9"
 version = "0.0.1"
+dynamic= ["dependencies"]
 readme = "README.md"
 classifiers = [ "Intended Audience :: Developers", "Operating System :: OS Independent" ,"License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3" ]

requirements.base.txt

@@ -0,0 +1,9 @@
+pyyaml
+ultralytics~=8.0.114
+av~=10.0.0
+onnxsim~=0.4.33
+onnx~=1.14.0
+ipykernel
+cython~=3.0.0
+build~=0.10.0
+-r ./requirements.inference.txt

requirements.cpu.txt

@@ -0,0 +1,2 @@
+-r ./requirements.base.txt
+-r ./requirements.torch.cpu.txt

requirements.dev.txt

@@ -1,3 +1,5 @@
 autopep8
 pylint
-pytest
+pytest
+setuptools>=63.0
+wheel>=0.40

requirements.gpu.txt

@@ -0,0 +1,2 @@
+-r ./requirements.base.txt
+-r ./requirements.torch.gpu.txt

requirements.inference.txt

@@ -0,0 +1,4 @@
+opencv-python-headless~=4.8.0.76
+# onnxruntime-openvino~=1.11.0
+onnxruntime-gpu~=1.15.1
+matplotlib~=3.7.2

requirements.torch.cpu.txt

@@ -0,0 +1,2 @@
+torch~=2.0.1 
+torchvision~=0.15.2 

requirements.torch.gpu.txt

@@ -0,0 +1,3 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch~=2.0.1 
+torchvision~=0.15.2