Added inference code

portiloop/capture.py

@@ -399,7 +399,7 @@ def _capture_process(p_data_o, p_msg_io, duration, frequency, python_clock, time
 
 
 class Capture:
-    def __init__(self):
+    def __init__(self, quantInferenceClass):
         # {now.strftime('%m_%d_%Y_%H_%M_%S')}
         self.filename = EDF_PATH / 'recording.edf'
         self._p_capture = None
@@ -415,6 +415,9 @@ class Capture:
         self.custom_fir_cutoff = 30
         self.filter = True
         self.record = False
+        self.detect = False
+        self.stimulate = False
+        self.threshold = 0.5
         self.lsl = False
         self.display = False
         self.python_clock = True
@@ -430,6 +433,8 @@ class Capture:
         self._t_capture = None
         self.channel_states = ['disabled', 'disabled', 'disabled', 'disabled', 'disabled', 'disabled', 'disabled', 'disabled']
         
+        self.quantInferenceClass = quantInferenceClass
+        
         # widgets ===============================
         
         # CHANNELS ------------------------------
@@ -554,6 +559,12 @@ class Capture:
             disabled=False
         )
         
+        self.b_threshold = widgets.FloatText(
+            value=self.threshold,
+            description='Threshold:',
+            disabled=True
+        )
+        
         self.b_polyak_mean = widgets.FloatText(
             value=self.polyak_mean,
             description='Polyak mean:',
@@ -610,6 +621,20 @@ class Capture:
             indent=False
         )
         
+        self.b_detect = widgets.Checkbox(
+            value=self.detect,
+            description='Detect',
+            disabled=False,
+            indent=False
+        )
+        
+        self.b_stimulate = widgets.Checkbox(
+            value=self.stimulate,
+            description='Stimulate',
+            disabled=True,
+            indent=False
+        )
+
         self.b_record = widgets.Checkbox(
             value=self.record,
             description='Record EDF',
@@ -636,8 +661,10 @@ class Capture:
         self.b_capture.observe(self.on_b_capture, 'value')
         self.b_clock.observe(self.on_b_clock, 'value')
         self.b_frequency.observe(self.on_b_frequency, 'value')
+        self.b_threshold.observe(self.on_b_threshold, 'value')
         self.b_duration.observe(self.on_b_duration, 'value')
         self.b_filter.observe(self.on_b_filter, 'value')
+        self.b_detect.observe(self.on_b_detect, 'value')
         self.b_record.observe(self.on_b_record, 'value')
         self.b_lsl.observe(self.on_b_lsl, 'value')
         self.b_display.observe(self.on_b_display, 'value')
@@ -668,7 +695,8 @@ class Capture:
                               self.b_filename,
                               self.b_power_line,
                               self.b_clock,
-                              widgets.HBox([self.b_filter, self.b_record, self.b_lsl, self.b_display]),
+                              widgets.HBox([self.b_filter, self.b_detect, self.b_stimulate, self.b_record, self.b_lsl, self.b_display]),
+                              self.b_threshold,
                               self.b_accordion_filter,
                               self.b_capture]))
 
@@ -677,6 +705,7 @@ class Capture:
         self.b_duration.disabled = False
         self.b_filename.disabled = False
         self.b_filter.disabled = False
+        self.b_detect.disabled = False
         self.b_record.disabled = False
         self.b_record.lsl = False
         self.b_display.disabled = False
@@ -694,12 +723,16 @@ class Capture:
         self.b_custom_fir.disabled = False
         self.b_custom_fir_order.disabled = not self.custom_fir
         self.b_custom_fir_cutoff.disabled = not self.custom_fir
+        self.b_stimulate.disabled = not self.detect
+        self.b_threshold.disabled = not self.detect
     
     def disable_buttons(self):
         self.b_frequency.disabled = True
         self.b_duration.disabled = True
         self.b_filename.disabled = True
         self.b_filter.disabled = True
+        self.b_stimulate.disabled = True
+        self.b_filter.disabled = True
         self.b_record.disabled = True
         self.b_record.lsl = True
         self.b_display.disabled = True
@@ -717,6 +750,7 @@ class Capture:
         self.b_custom_fir.disabled = True
         self.b_custom_fir_order.disabled = True
         self.b_custom_fir_cutoff.disabled = True
+        self.b_threshold.disabled = True
     
     def on_b_radio_ch2(self, value):
         self.channel_states[1] = value['new']
@@ -751,7 +785,7 @@ class Capture:
                 warnings.warn("Capture already running, operation aborted.")
                 return
             self._t_capture = Thread(target=self.start_capture,
-                                args=(self.filter, self.record, self.lsl, self.display, 500, self.python_clock))
+                                args=(self.filter, self.detect, self.quantInferenceClass, self.record, self.lsl, self.display, 500, self.python_clock))
             self._t_capture.start()
         elif val == 'Stop':
             with self._lock_msg_out:
@@ -790,6 +824,13 @@ class Capture:
         else:
             self.b_frequency.value = self.frequency
             
+    def on_b_threshold(self, value):
+        val = value['new']
+        if val >= 0 and val <= 1:
+            self.threshold = val
+        else:
+            self.b_threshold.value = self.threshold
+            
     def on_b_filename(self, value):
         val = value['new']
         if val != '':
@@ -844,6 +885,15 @@ class Capture:
         val = value['new']
         self.filter = val
     
+    def on_b_stimulate(self, value):
+        val = value['new']
+        self.stimulate = val
+    
+    def on_b_detect(self, value):
+        val = value['new']
+        self.detect = val
+        self.enable_buttons()
+    
     def on_b_record(self, value):
         val = value['new']
         self.record = val
@@ -894,6 +944,8 @@ class Capture:
 
     def start_capture(self,
                       filter,
+                      detect,
+                      quantInferenceClass,
                       record,
                       lsl,
                       viz,
@@ -918,6 +970,9 @@ class Capture:
                                 alpha_avg=self.polyak_mean,
                                 alpha_std=self.polyak_std,
                                 epsilon=self.epsilon)
+            
+        if detect:
+            infer = quantInferenceClass()
 
         self._p_capture = mp.Process(target=_capture_process,
                                      args=(p_data_o,
@@ -975,6 +1030,15 @@ class Capture:
             
             filtered_point = n_array.tolist()
             
+            if detect:
+                results = infer.add_datapoints(filtered_points)
+                
+                for r in results:
+                    print(r >= threshold)
+                
+                if stimulate and True:
+                    print('stimulation')
+            
             if lsl:
                 lsl_outlet.push_sample(filtered_point[-1])
             

portiloop/inference.py

@@ -0,0 +1,95 @@
+from pycoral.utils import edgetpu
+import time
+from abc import ABC, abstractmethod
+from pathlib import Path
+import numpy as np
+
+DEFAULT_MODEL_PATH = str(Path(__file__).parent / "models/portiloop_model_quant.tflite")
+print(DEFAULT_MODEL_PATH)
+
+class AbstractQuantizedModelForInference(ABC):
+    @abstractmethod
+    def add_datapoints(self, input_float):
+        return NotImplemented
+
+class QuantizedModelForInference(AbstractQuantizedModelForInference):
+    def __init__(self, num_models_parallel=8, window_size=54, seq_stride=42, model_path=None, verbose=False, channel=2):
+        model_path = DEFAULT_MODEL_PATH if model_path is None else model_path
+        self.verbose = verbose
+        self.channel = channel
+        self.num_models_parallel = num_models_parallel
+        
+        self.interpreters = []
+        for i in range(self.num_models_parallel):
+            self.interpreters.append(edgetpu.make_interpreter(model_path))
+            self.interpreters[i].allocate_tensors()
+        self.interpreter_counter = 0
+        
+        self.input_details = self.interpreters[0].get_input_details()
+        self.output_details = self.interpreters[0].get_output_details()
+        
+        self.buffer = []
+        self.seq_stride = seq_stride
+        self.window_size = window_size 
+        
+        self.stride_counters = [np.floor((self.seq_stride / self.num_models_parallel) * i) for i in range(self.num_models_parallel)]
+        for idx, i in enumerate(self.stride_counters[1:]):
+            self.stride_counters[idx+1] = i - self.stride_counters[idx]
+        self.current_stride_counter = self.stride_counters[0] - 1
+        
+        
+    def add_datapoints(self, inputs_float):
+        res = []
+        for inp in inputs_float:
+            result = self.add_datapoint(inp)
+            if result is not None:
+                res.append(result)
+        return res
+    
+        
+    def add_datapoint(self, input_float):
+        input_float = input_float[self.channel-1]
+        result = None
+        self.buffer.append(input_float)
+        if len(self.buffer) > self.window_size:
+            self.buffer = self.buffer[1:]
+            self.current_stride_counter += 1
+            if self.current_stride_counter == self.stride_counter[self.interpreter_counter]:
+                result = self.call_model(self.interpreter_counter, self.buffer)
+                self.interpreter_counter += 1
+                self.interpreter_counter %= self.num_model_parallel
+                self.current_stride_counter = 0
+        return result
+            
+                
+        
+    def call_model(self, idx, input_float=None):
+        if input_float is None:
+            # For debuggin purposes
+            input_shape = input_details[0]['shape']
+            input = np.array(np.random.random_sample(input_shape), dtype=np.int8)
+        else:
+            # Convert float input to Int
+            input_scale, input_zero_point = input_details[0]["quantization"]
+            input = np.asarray(input_float) / input_scale + input_zero_point
+            input = input.astype(input_details[0]["dtype"])
+
+        interpreter.set_tensor(input_details[0]['index'], input)
+        if self.verbose:
+            start_time = time.time()
+
+        interpreter.invoke()
+
+        if self.verbose:
+            end_time = time.time()
+
+        output = interpreter.get_tensor(output_details[0]['index'])
+        output_scale, output_zero_point = input_details[0]["quantization"]
+        output = float(output - output_zero_point) * output_scale
+
+        if self.verbose:
+            print(f"Computed output {output} in {end_time - start_time} seconds")
+
+        return output
+    
+    

portiloop/notebooks/tests.ipynb

@@ -2,16 +2,47 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "7b2fc5da",
    "metadata": {
     "scrolled": false
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "910f8e489b6341119f4d6e17a5b2aedc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(Accordion(children=(GridBox(children=(Label(value='CH1'), Label(value='CH2'), Label(value='CH3'…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Process Process-1:\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/usr/lib/python3.7/multiprocessing/process.py\", line 297, in _bootstrap\n",
+      "    self.run()\n",
+      "  File \"/usr/lib/python3.7/multiprocessing/process.py\", line 99, in run\n",
+      "    self._target(*self._args, **self._kwargs)\n",
+      "  File \"/home/mendel/software/portiloop-software/portiloop/capture.py\", line 325, in _capture_process\n",
+      "    assert data == [0x3E], \"The communication with the ADS cannot be established.\"\n",
+      "AssertionError: The communication with the ADS cannot be established.\n"
+     ]
+    }
+   ],
    "source": [
     "from portiloop.capture import Capture\n",
+    "from portiloop.inference import QuantizedModelForInference\n",
     "\n",
-    "cap = Capture()"
+    "cap = Capture(QuantizedModelForInference)"
    ]
   }
  ],

setup.py

@@ -14,6 +14,7 @@ setup(
                       'python-periphery',
                       'spidev',
                       'pylsl-coral',
-                      'scipy'
+                      'scipy',
+                      'pycoral'
                      ]
 )