Add support for freeform SHM'ed output (#33)

Author: janjongboom

edge_impulse_linux/runner.py

@@ -13,15 +13,17 @@ def now():
     return round(time.time() * 1000)
 
 class ImpulseRunner:
-    def __init__(self, model_path: str, timeout: int = 5):
+    def __init__(self, model_path: str, timeout: int = 30, allow_shm = True):
         self._model_path = model_path
         self._tempdir = None
         self._runner = None
         self._client = None
         self._ix = 0
         self._debug = False
         self._hello_resp = None
-        self._shm = None
+        self._allow_shm = allow_shm
+        self._input_shm = None
+        self._freeform_output_shm = []
         self._timeout = timeout
 
     def init(self, debug=False):
@@ -57,17 +59,32 @@ def init(self, debug=False):
 
         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)
-            }
+        if self._allow_shm:
+            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._input_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)
+                }
+
+            if ('freeform_output_shm' in hello_resp.keys()):
+                for output_shm in hello_resp['freeform_output_shm']:
+                    shm_name = output_shm['name']
+                    # python does not want the leading slash
+                    shm_name = shm_name.lstrip('/')
+                    shm = shared_memory.SharedMemory(name=shm_name)
+                    self._freeform_output_shm.append({
+                        'index': output_shm['index'],
+                        'shm': shm,
+                        'type': output_shm['type'],
+                        'elements': output_shm['elements'],
+                        'array': np.ndarray((output_shm['elements'],), dtype=np.float32, buffer=shm.buf)
+                    })
 
         return self._hello_resp
 
@@ -88,18 +105,23 @@ def stop(self):
             # 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
+        if self._input_shm is not None:
+            self._input_shm['shm'].close()
+            resource_tracker.unregister(self._input_shm['shm']._name, "shared_memory")
+            self._input_shm = None
+
+        for shm in self._freeform_output_shm:
+            shm['shm'].close()
+            resource_tracker.unregister(shm['shm']._name, "shared_memory")
+        self._freeform_output_shm = []
 
     def hello(self):
         msg = {"hello": 1}
         return self.send_msg(msg)
 
     def classify(self, data):
-        if self._shm:
-            self._shm['array'][:] = data
+        if self._input_shm:
+            self._input_shm['array'][:] = data
 
             msg = {
                 "classify_shm": {
@@ -113,6 +135,13 @@ def classify(self, data):
             msg["debug"] = True
 
         send_resp = self.send_msg(msg)
+
+        if 'result' in send_resp and 'freeform' in send_resp['result'] and send_resp['result']['freeform'] == 'shm':
+            freeform = []
+            for shm in self._freeform_output_shm:
+                freeform.append(shm['array'].tolist())
+            send_resp['result']['freeform'] = freeform
+
         return send_resp
 
     def set_threshold(self, obj):

examples/audio/classify.py

@@ -41,7 +41,7 @@ def main(argv):
     with AudioImpulseRunner(modelfile) as runner:
         try:
             model_info = runner.init()
-            # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+            # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
             labels = model_info['model_parameters']['labels']
             print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
 

examples/custom/classify.py

@@ -54,7 +54,7 @@ def main(argv):
     runner = ImpulseRunner(modelfile)
     try:
         model_info = runner.init()
-        # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+        # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
 
         print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
 

examples/image/classify-full-frame.py

@@ -71,7 +71,7 @@ def main(argv):
     with ImageImpulseRunner(modelfile) as runner:
         try:
             model_info = runner.init()
-            # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+            # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
 
             print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
             labels = model_info['model_parameters']['labels']

examples/image/classify-image.py

@@ -43,7 +43,7 @@ def main(argv):
     with ImageImpulseRunner(modelfile) as runner:
         try:
             model_info = runner.init()
-            # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+            # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
 
             print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
             labels = model_info['model_parameters']['labels']

examples/image/classify-video.py

@@ -48,7 +48,7 @@ def main(argv):
     with ImageImpulseRunner(modelfile) as runner:
         try:
             model_info = runner.init()
-            # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+            # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
             print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
             labels = model_info['model_parameters']['labels']
 

examples/image/classify.py

@@ -75,7 +75,7 @@ def main(argv):
     with ImageImpulseRunner(modelfile) as runner:
         try:
             model_info = runner.init()
-            # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+            # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
             print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
             labels = model_info['model_parameters']['labels']
             if len(args)>= 2:

examples/image/set-thresholds.py

@@ -44,7 +44,7 @@ def main(argv):
     with ImageImpulseRunner(modelfile) as runner:
         try:
             model_info = runner.init()
-            # model_info = runner.init(debug=True, timeout=10) # to get debug print out and set longer timeout
+            # model_info = runner.init(debug=True, timeout=60) # to get debug print out and set longer timeout
 
             print('Loaded runner for "' + model_info['project']['owner'] + ' / ' + model_info['project']['name'] + '"')
             if not 'thresholds' in model_info['model_parameters']: