Add support for EIM files with shared memory, add support for setting thresholds (#31)

* Add support for shared memory inference

* Cleanup

* Remove set-thresholds.py

* Add set_threshold example

* Remove debug statement

* image: missing importing sys package

* runner: small simplification

* set-thresholds: make the example interactive

* image: more optimized image conversion

Previous versions were ~10 slower than the new one.

---------

Co-authored-by: Mateusz Majchrzycki <[email protected]>

Author: janjongboom

.gitignore

@@ -8,3 +8,4 @@ build
 .vscode
 *.eim
 .act-secrets
+*.png

README.md

@@ -32,7 +32,7 @@ This library lets you run machine learning models and collect sensor data on Lin
     $ pip3 install -r requirements.txt
     ```
 
-    For the computer vision examples you'll want `opencv-python>=4.5.1.48`
+    For the computer vision examples you'll want `opencv-python>=4.5.1.48,<5`
     Note on macOS on apple silicon, you will need to use a later version,
     4.10.0.84 tested and installs cleanly
 

edge_impulse_linux/image.py

@@ -1,10 +1,15 @@
 #!/usr/bin/env python
 
 import numpy as np
-import cv2
+import sys
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
+
 from edge_impulse_linux.runner import ImpulseRunner
 import math
-import psutil
 
 class ImageImpulseRunner(ImpulseRunner):
     def __init__(self, model_path: str):
@@ -44,7 +49,7 @@ def classify(self, data):
 
     # This returns images in RGB format (not BGR)
     def get_frames(self, videoDeviceId = 0):
-        if psutil.OSX or psutil.MACOS:
+        if sys.platform == "darwin":
             print('Make sure to grant the this script access to your webcam.')
             print('If your webcam is not responding, try running "tccutil reset Camera" to reset the camera access privileges.')
 
@@ -57,7 +62,7 @@ def get_frames(self, videoDeviceId = 0):
 
     # This returns images in RGB format (not BGR)
     def classifier(self, videoDeviceId = 0):
-        if psutil.OSX or psutil.MACOS:
+        if sys.platform == "darwin":
             print('Make sure to grant the this script access to your webcam.')
             print('If your webcam is not responding, try running "tccutil reset Camera" to reset the camera access privileges.')
 
@@ -137,9 +142,7 @@ def get_features_from_image_auto_studio_settings(self, img):
             raise Exception(
                 'Runner has not initialized, please call init() first')
         if self.resizeMode == 'not-reported':
-            raise Exception(
-                'Model file "' + self._model_path + '" does not report the image resize mode\n'
-                'Please update the model file via edge-impulse-linux-runner --download')
+            self.resizeMode = 'squash'
         return get_features_from_image_with_studio_mode(img, self.resizeMode, self.dim[0], self.dim[1], self.isGrayscale)
 
 
@@ -233,17 +236,10 @@ def get_features_from_image_with_studio_mode(img, mode, output_width, output_hei
 
     if is_grayscale:
         resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2GRAY)
-        pixels = np.array(resized_img).flatten().tolist()
-
-        for p in pixels:
-            features.append((p << 16) + (p << 8) + p)
+        features = (resized_img.astype(np.uint32) * 0x010101).flatten().tolist()
     else:
-        pixels = np.array(resized_img).flatten().tolist()
-
-        for ix in range(0, len(pixels), 3):
-            r = pixels[ix + 0]
-            g = pixels[ix + 1]
-            b = pixels[ix + 2]
-            features.append((r << 16) + (g << 8) + b)
+        # Use numpy's vectorized operations for RGB feature encoding
+        pixels = resized_img.astype(np.uint32)
+        features = ((pixels[..., 0] << 16) | (pixels[..., 1] << 8) | pixels[..., 2]).flatten().tolist()
 
     return features, resized_img

edge_impulse_linux/runner.py

@@ -6,12 +6,12 @@
 import signal
 import socket
 import json
-
+from multiprocessing import shared_memory, resource_tracker
+import numpy as np
 
 def now():
     return round(time.time() * 1000)
 
-
 class ImpulseRunner:
     def __init__(self, model_path: str):
         self._model_path = model_path
@@ -20,6 +20,8 @@ def __init__(self, model_path: str):
         self._client = None
         self._ix = 0
         self._debug = False
+        self._hello_resp = None
+        self._shm = None
 
     def init(self, debug=False):
         if not os.path.exists(self._model_path):
@@ -50,27 +52,71 @@ def init(self, debug=False):
         self._client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
         self._client.connect(socket_path)
 
-        return self.hello()
+        hello_resp = self._hello_resp = self.hello()
+
+        if ('features_shm' in hello_resp.keys()):
+            shm_name = hello_resp['features_shm']['name']
+            # python does not want the leading slash
+            shm_name = shm_name.lstrip('/')
+            shm = shared_memory.SharedMemory(name=shm_name)
+            self._shm = {
+                'shm': shm,
+                'type': hello_resp['features_shm']['type'],
+                'elements': hello_resp['features_shm']['elements'],
+                'array': np.ndarray((hello_resp['features_shm']['elements'],), dtype=np.float32, buffer=shm.buf)
+            }
+
+        return self._hello_resp
+
+    def __del__(self):
+        self.stop()
 
     def stop(self):
-        if self._tempdir:
+        if self._tempdir is not None:
             shutil.rmtree(self._tempdir)
+            self._tempdir = None
 
-        if self._client:
+        if self._client is not None:
             self._client.close()
+            self._client = None
 
-        if self._runner:
+        if self._runner is not None:
             os.kill(self._runner.pid, signal.SIGINT)
             # todo: in Node we send a SIGHUP after 0.5sec if process has not died, can we do this somehow here too?
+            self._runner = None
+
+        if self._shm is not None:
+            self._shm['shm'].close()
+            resource_tracker.unregister(self._shm['shm']._name, "shared_memory")
+            self._shm = None
 
     def hello(self):
         msg = {"hello": 1}
         return self.send_msg(msg)
 
     def classify(self, data):
-        msg = {"classify": data}
+        if self._shm:
+            self._shm['array'][:] = data
+
+            msg = {
+                "classify_shm": {
+                    "elements": len(data),
+                }
+            }
+        else:
+            msg = {"classify": data}
+
         if self._debug:
             msg["debug"] = True
+
+        send_resp = self.send_msg(msg)
+        return send_resp
+
+    def set_threshold(self, obj):
+        if not 'id' in obj:
+            raise Exception('set_threshold requires an object with an "id" field')
+
+        msg = { 'set_threshold': obj }
         return self.send_msg(msg)
 
     def send_msg(self, msg):

examples/image/classify-full-frame.py

@@ -2,7 +2,11 @@
 
 import device_patches       # Device specific patches for Jetson Nano (needs to be before importing cv2)
 
-import cv2
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
 import os
 import sys, getopt
 import signal

examples/image/classify-image.py

@@ -2,7 +2,11 @@
 
 import device_patches       # Device specific patches for Jetson Nano (needs to be before importing cv2)  # noqa: F401
 
-import cv2
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
 import os
 import sys
 import getopt

examples/image/classify-video.py

@@ -2,11 +2,14 @@
 
 import device_patches       # Device specific patches for Jetson Nano (needs to be before importing cv2)
 
-import cv2
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
 import os
 import time
 import sys, getopt
-import numpy as np
 from edge_impulse_linux.image import ImageImpulseRunner
 
 runner = None

examples/image/classify.py

@@ -2,7 +2,11 @@
 
 import device_patches       # Device specific patches for Jetson Nano (needs to be before importing cv2)
 
-import cv2
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
 import os
 import sys, getopt
 import signal

examples/image/resize_demo.py

@@ -1,5 +1,9 @@
 import numpy as np
-import cv2
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
 from edge_impulse_linux.image import get_features_from_image_with_studio_mode
 
 

examples/image/set-thresholds.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+
+import device_patches       # Device specific patches for Jetson Nano (needs to be before importing cv2)  # noqa: F401
+
+try:
+    import cv2
+except ImportError:
+    print('Missing OpenCV, install via `pip3 install "opencv-python>=4.5.1.48,<5"`')
+    exit(1)
+import os
+import sys
+import getopt
+import json
+from edge_impulse_linux.image import ImageImpulseRunner
+
+runner = None
+
+def help():
+    print('python set-thresholds.py <path_to_model.eim> <path_to_image.jpg>')
+
+def main(argv):
+    try:
+        opts, args = getopt.getopt(argv, "h", ["--help"])
+    except getopt.GetoptError:
+        help()
+        sys.exit(2)
+
+    for opt, arg in opts:
+        if opt in ('-h', '--help'):
+            help()
+            sys.exit()
+
+    if len(args) != 2:
+        help()
+        sys.exit(2)
+
+    model = args[0]
+
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    modelfile = os.path.join(dir_path, model)
+
+    print('MODEL: ' + modelfile)
+
+    with ImageImpulseRunner(modelfile) as runner:
+        try:
+            model_info = runner.init()
+            # model_info = runner.init(debug=True) # to get debug print out
+
+            print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
+            if not 'thresholds' in model_info['model_parameters']:
+                print('This model does not expose any thresholds, build a new Linux deployment (.eim file) to get configurable thresholds')
+                exit(1)
+
+            print('Thresholds:')
+            for threshold in model_info['model_parameters']['thresholds']:
+                print('    -', json.dumps(threshold))
+
+            # Example output for an object detection model:
+            # Thresholds:
+            #    - {"id": 3, "min_score": 0.20000000298023224, "type": "object_detection"}
+
+            img = cv2.imread(args[1])
+            if img is None:
+                print('Failed to load image', args[1])
+                exit(1)
+
+            # imread returns images in BGR format, so we need to convert to RGB
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            # this mode uses the same settings used in studio to crop and resize the input
+            features, cropped = runner.get_features_from_image_auto_studio_settings(img)
+
+            print("Which threshold would you like to change? (id)")
+            while True:
+                try:
+                    threshold_id = int(input('Enter threshold ID: '))
+                    if threshold_id not in [t['id'] for t in model_info['model_parameters']['thresholds']]:
+                        print('Invalid threshold ID, try again')
+                        continue
+                    break
+                except ValueError:
+                    print('Invalid input, please enter a number')
+
+            print("Enter a new threshold value (between 0.0 and 1.0):")
+            while True:
+                try:
+                    new_threshold = float(input('New threshold value: '))
+                    if new_threshold < 0.0 or new_threshold > 1.0:
+                        print('Invalid threshold value, must be between 0.0 and 1.0')
+                        continue
+                    break
+                except ValueError:
+                    print('Invalid input, please enter a number')
+
+            # dynamically override the thresold from 0.2 -> 0.8
+            runner.set_threshold({
+                'id': threshold_id,
+                'min_score': new_threshold,
+            })
+
+            res = runner.classify(features)
+            print('classify response', json.dumps(res, indent=4))
+
+        finally:
+            if (runner):
+                runner.stop()
+
+if __name__ == "__main__":
+   main(sys.argv[1:])