Finished setting up the gradio demo

.gitignore

@@ -2,6 +2,9 @@
 .vscode/
 .idea/
 
+# Output from the demo
+output.csv
+
 # Vagrant
 .vagrant/
 

portiloop/src/demo/demo.py

@@ -0,0 +1,137 @@
+import gradio as gr
+import matplotlib.pyplot as plt
+import time
+import numpy as np
+import pandas as pd
+from portiloop.src.demo.demo_stimulator import DemoSleepSpindleRealTimeStimulator
+from portiloop.src.detection import SleepSpindleRealTimeDetector
+
+from portiloop.src.stimulation import UpStateDelayer
+plt.switch_backend('agg')
+from portiloop.src.processing import FilterPipeline
+
+
+def do_treatment(csv_file, filtering, threshold, detect_channel, freq, spindle_freq, spindle_detection_mode, time_to_buffer):
+
+    # Read the csv file to a numpy array
+    data_whole = np.loadtxt(csv_file.name, delimiter=',')
+
+    # Get the data from the selected channel
+    detect_channel = int(detect_channel)
+    freq = int(freq)
+    data = data_whole[:, detect_channel - 1]
+
+    # Create the detector and the stimulator
+    detector = SleepSpindleRealTimeDetector(threshold=threshold, channel=1) # always 1 because we have only one channel
+    stimulator = DemoSleepSpindleRealTimeStimulator()
+    if spindle_detection_mode != 'Fast':
+        delayer = UpStateDelayer(freq, spindle_freq, spindle_detection_mode == 'Peak', time_to_buffer=time_to_buffer)
+        stimulator.add_delayer(delayer)
+    
+    # Create the filtering pipeline
+    if filtering:
+        filter = FilterPipeline(nb_channels=1, sampling_rate=freq)
+
+    # Plotting variables
+    points = []
+    activations = []
+    delayed_activations = []
+
+    # Go through the data
+    for index, point in enumerate(data):
+        # Step the delayer if exists
+        if spindle_detection_mode != 'Fast':
+            delayed = delayer.step(point)
+            if delayed:
+                delayed_activations.append(1)
+            else:
+                delayed_activations.append(0)
+
+        # Filter the data
+        if filtering:
+            filtered_point = filter.filter(np.array([point]))
+        else:
+            filtered_point = point
+        
+        filtered_point = filtered_point.tolist()
+
+        # Detect the spindles
+        result = detector.detect([filtered_point])
+
+        # Stimulate if necessary
+        stim = stimulator.stimulate(result)
+        if stim:
+            activations.append(1)
+        else:
+            activations.append(0)
+        
+        # Add data to plotting buffer
+        points.append(filtered_point[0])
+
+        # Plot the data
+        if index % (10 * freq) == 0:
+            plt.close()
+            fig = plt.figure(figsize=(20, 10))
+            plt.clf()
+            plt.plot(points[-10 * freq:], label="Data")
+            # Draw vertical lines for activations
+            for index in get_activations(activations[-10 * freq:]):
+                plt.axvline(x=index, color='r', label="Fast Stimulation")
+            if spindle_detection_mode != 'Fast':
+                for index in get_activations(delayed_activations[-10 * freq:]):
+                    plt.axvline(x=index, color='g', label="Delayed Stimulation")
+            yield fig, None
+
+    # Put all points and activations back in numpy arrays
+    points = np.array(points)
+    activations = np.array(activations)
+    delayed_activations = np.array(delayed_activations)
+    # Concatenate with the original data
+    data_whole = np.concatenate((data_whole, points.reshape(-1, 1), activations.reshape(-1, 1), delayed_activations.reshape(-1, 1)), axis=1)
+    # Output the data to a csv file
+    np.savetxt('output.csv', data_whole, delimiter=',')
+
+    yield None, "output.csv"
+        
+# Function to return a list of all indexes where activations have happened
+def get_activations(activations):
+    return [i for i, x in enumerate(activations) if x == 1]
+    
+
+with gr.Blocks() as demo:
+    gr.Markdown("Enter your csv file and click **Run Inference** to get the output.")
+
+    # Row containing all inputs:
+    with gr.Row():
+        # CSV file
+        csv_file = gr.UploadButton(label="CSV File", file_count="single")
+        # Filtering (Boolean)
+        filtering = gr.Checkbox(label="Filtering (On/Off)", value=True)
+        # Threshold value
+        threshold = gr.Slider(0, 1, value=0.82, step=0.01, label="Threshold", interactive=True)
+        # Detection Channel
+        detect_column = gr.Dropdown(choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], value="1", label="Detection Column", interactive=True) 
+        # Frequency
+        freq = gr.Dropdown(choices=["100", "200", "250", "256", "500", "512", "1000", "1024"], value="250", label="Frequency", interactive=True)
+        # Spindle Frequency
+        spindle_freq = gr.Slider(10, 16, value=12, step=1, label="Spindle Frequency", interactive=True)
+        # Spindle Detection Mode
+        spindle_detection_mode = gr.Dropdown(choices=["Fast", "Peak", "Valley"], value="Peak", label="Spindle Detection Mode", interactive=True)
+        # Time to buffer
+        time_to_buffer = gr.Slider(0, 1, value=0, step=0.01, label="Time to Buffer", interactive=True)
+
+    # Output plot
+    output_plot = gr.Plot()
+    # Output file
+    output_array = gr.File(label="Output CSV File")
+
+    # Row containing all buttons:
+    with gr.Row():
+        # Run inference button
+        run_inference = gr.Button(value="Run Inference")
+        # Reset button
+        reset = gr.Button(value="Reset", variant="secondary")
+    run_inference.click(fn=do_treatment, inputs=[csv_file, filtering, threshold, detect_column, freq, spindle_freq, spindle_detection_mode, time_to_buffer], outputs=[output_plot, output_array])
+
+demo.queue()
+demo.launch()

portiloop/src/demo/demo_stimulator.py

@@ -0,0 +1,30 @@
+import time
+from portiloop.src.stimulation import Stimulator
+
+
+class DemoSleepSpindleRealTimeStimulator(Stimulator):
+    def __init__(self):
+        self.last_detected_ts = time.time()
+        self.wait_t = 0.4  # 400 ms
+    
+    def stimulate(self, detection_signal):
+        stim = False
+        for sig in detection_signal:
+            # We detect a stimulation
+            if sig:
+                # Record time of stimulation
+                ts = time.time()
+                
+                # Check if time since last stimulation is long enough
+                if ts - self.last_detected_ts > self.wait_t:
+                    if self.delayer is not None:
+                        # If we have a delayer, notify it
+                        self.delayer.detected()
+                    stim = True
+
+                self.last_detected_ts = ts
+        return stim
+
+    def add_delayer(self, delayer):
+        self.delayer = delayer
+        self.delayer.stimulate = lambda: True

portiloop/src/detection.py

@@ -5,6 +5,8 @@ from portiloop.src import ADS
 
 if ADS:
     from pycoral.utils import edgetpu
+else:
+    import tensorflow as tf
 import numpy as np
 
 
@@ -53,7 +55,10 @@ class SleepSpindleRealTimeDetector(Detector):
         
         self.interpreters = []
         for i in range(self.num_models_parallel):
-            self.interpreters.append(edgetpu.make_interpreter(model_path))
+            if ADS:
+                self.interpreters.append(edgetpu.make_interpreter(model_path))
+            else:
+                self.interpreters.append(tf.lite.Interpreter(model_path=model_path))
             self.interpreters[i].allocate_tensors()
         self.interpreter_counter = 0
         
@@ -76,6 +81,10 @@ class SleepSpindleRealTimeDetector(Detector):
         super().__init__(threshold)
 
     def detect(self, datapoints):
+        """
+        Takes datapoints as input and outputs a detection signal.
+        datapoints is a list of lists of n channels: may contain several datapoints.
+        """
         res = []
         for inp in datapoints:
             result = self.add_datapoint(inp)

portiloop/src/stimulation.py

@@ -142,7 +142,7 @@ class SleepSpindleRealTimeStimulator(Stimulator):
 
     def add_delayer(self, delayer):
         self.delayer = delayer
-        self.delayer.stimulate = lambda x: self.send_stimulation("DELAY_STIM", True)
+        self.delayer.stimulate = lambda: self.send_stimulation("DELAY_STIM", True)
 
 # Class that delays stimulation to always stimulate peak or through
 class UpStateDelayer:
@@ -182,7 +182,7 @@ class UpStateDelayer:
             return False
         elif self.state == States.DELAYING:
             # Check if we are done delaying
-            if time.time() - self.time_started >= self.time_to_wait():
+            if time.time() - self.time_started >= self.time_to_wait:
                 # Actually stimulate the patient after the delay
                 if self.stimulate is not None:
                     self.stimulate()