Merge pull request #5 from Portiloop/milo/filetypes_and_staging

portiloop/src/demo/offline.py

@@ -1,13 +1,12 @@
-import matplotlib.pyplot as plt
 import numpy as np
 from portiloop.src.detection import SleepSpindleRealTimeDetector
-plt.switch_backend('agg')
+from portiloop.src.stimulation import UpStateDelayer
 from portiloop.src.processing import FilterPipeline
-from portiloop.src.demo.utils import compute_output_table, xdf2array, offline_detect, offline_filter, OfflineSleepSpindleRealTimeStimulator
+from portiloop.src.demo.utils import compute_output_table, sleep_stage, xdf2array, offline_detect, offline_filter, OfflineSleepSpindleRealTimeStimulator
 import gradio as gr
 
 
-def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
+def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, stimulation_phase="Fast", buffer_time=0.25):
     # Get the options from the checkbox group
     offline_filtering = 0 in detect_filter_opts
     lacourse = 1 in detect_filter_opts
@@ -30,6 +29,7 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
     # Read the xdf file to a numpy array
     print("Loading xdf file...")
     data_whole, columns = xdf2array(xdf_file.name, int(channel_num))
+
     # Do the offline filtering of the data
     if offline_filtering:
         print("Filtering offline...")
@@ -39,13 +39,18 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
         data_whole = np.concatenate((data_whole, offline_filtered_data), axis=1)
         columns.append("offline_filtered_signal")
 
+    # Do the sleep staging approximation
+    if wamsley or lacourse:
+        print("Sleep staging...")
+        mask = sleep_stage(data_whole[:, columns.index("offline_filtered_signal")], threshold=150, group_size=100)
+
     #  Do Wamsley's method
     if wamsley:
         print("Running Wamsley detection...")
         wamsley_data = offline_detect("Wamsley", \
             data_whole[:, columns.index("offline_filtered_signal")],\
                 data_whole[:, columns.index("time_stamps")],\
-                    freq)
+                    freq, mask)
         wamsley_data = np.expand_dims(wamsley_data, axis=1)
         data_whole = np.concatenate((data_whole, wamsley_data), axis=1)
         columns.append("wamsley_spindles")
@@ -56,7 +61,7 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
         lacourse_data = offline_detect("Lacourse", \
             data_whole[:, columns.index("offline_filtered_signal")],\
                 data_whole[:, columns.index("time_stamps")],\
-                    freq)
+                    freq, mask)
         lacourse_data = np.expand_dims(lacourse_data, axis=1)
         data_whole = np.concatenate((data_whole, lacourse_data), axis=1)
         columns.append("lacourse_spindles")
@@ -72,12 +77,17 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
     if online_detection:
         detector = SleepSpindleRealTimeDetector(threshold=threshold, channel=1) # always 1 because we have only one channel
         stimulator = OfflineSleepSpindleRealTimeStimulator()
+        if stimulation_phase != "Fast":
+            stimulation_delayer = UpStateDelayer(freq, stimulation_phase == 'Peak', time_to_buffer=buffer_time, stimulate=lambda: None)
+            stimulator.add_delayer(stimulation_delayer)
+            
 
     if online_filtering or online_detection:
         print("Running online filtering and detection...")
 
         points = []
         online_activations = []
+        delayed_stims = []
 
         # Go through the data
         for index, point in enumerate(data):
@@ -93,6 +103,13 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
                 # Detect the spindles
                 result = detector.detect([filtered_point])
 
+                if stimulation_phase != "Fast":
+                    delayed_stim = stimulation_delayer.step_timesteps(filtered_point[0])
+                    if delayed_stim:
+                        delayed_stims.append(1)
+                    else:
+                        delayed_stims.append(0)
+
                 # Stimulate if necessary
                 stim = stimulator.stimulate(result)
                 if stim:
@@ -112,6 +129,12 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq):
         data_whole = np.concatenate((data_whole, online_activations), axis=1)
         columns.append("online_stimulations")
 
+        if stimulation_phase != "Fast":
+            delayed_stims = np.array(delayed_stims)
+            delayed_stims = np.expand_dims(delayed_stims, axis=1)
+            data_whole = np.concatenate((data_whole, delayed_stims), axis=1)
+            columns.append("delayed_stimulations")
+
     print("Saving output...")
     # Output the data to a csv file
     np.savetxt("output.csv", data_whole, delimiter=",", header=",".join(columns), comments="")

portiloop/src/demo/phase_demo.py

@@ -0,0 +1,63 @@
+import gradio as gr
+
+from portiloop.src.demo.offline import run_offline
+
+        
+def on_upload_file(file):
+    # Check if file extension is .xdf
+    if file.name.split(".")[-1] != "xdf":
+        raise gr.Error("Please upload a .xdf file.")
+    else:
+        return file.name
+
+
+def main():
+    with gr.Blocks(title="Portiloop") as demo:
+        gr.Markdown("# Portiloop Demo")
+        gr.Markdown("This Demo takes as input an XDF file coming from the Portiloop EEG device and allows you to convert it to CSV and perform the following actions:: \n * Filter the data offline \n * Perform offline spindle detection using Wamsley or Lacourse. \n * Simulate the Portiloop online filtering and spindle detection with different parameters.")
+        gr.Markdown("Upload your XDF file and click **Run Inference** to start the processing...")
+
+        with gr.Row():
+            xdf_file_button = gr.UploadButton(label="Click to Upload", type="file", file_count="single")
+            xdf_file_static = gr.File(label="XDF File", type='file', interactive=False)
+
+            xdf_file_button.upload(on_upload_file, xdf_file_button, xdf_file_static)
+
+            # Make a checkbox group for the options
+            detect_filter = gr.CheckboxGroup(['Offline Filtering', 'Lacourse Detection', 'Wamsley Detection', 'Online Filtering', 'Online Detection'], type='index', label="Filtering/Detection options")
+            
+            # Options for phase stimulation
+            with gr.Row():  
+                # Dropwdown for phase
+                phase = gr.Dropdown(choices=["Peak", "Fast", "Valley"], value="Peak", label="Phase", interactive=True)
+                buffer_time = gr.Slider(0, 1, value=0.3, step=0.01, label="Buffer Time", interactive=True)
+                
+            # Threshold value
+            threshold = gr.Slider(0, 1, value=0.82, step=0.01, label="Threshold", interactive=True)
+            # Detection Channel
+            detect_channel = gr.Dropdown(choices=["1", "2", "3", "4", "5", "6", "7", "8"], value="2", label="Detection Channel in XDF recording", interactive=True) 
+            # Frequency
+            freq = gr.Dropdown(choices=["100", "200", "250", "256", "500", "512", "1000", "1024"], value="250", label="Sampling Frequency (Hz)", interactive=True)
+
+        with gr.Row():
+            output_array = gr.File(label="Output CSV File")
+            output_table = gr.Markdown(label="Output Table")
+
+        run_inference = gr.Button(value="Run Inference")
+        run_inference.click(
+            fn=run_offline, 
+            inputs=[
+                xdf_file_static, 
+                detect_filter,
+                threshold, 
+                detect_channel,
+                freq,
+                phase,
+                buffer_time], 
+            outputs=[output_array, output_table])
+
+    demo.queue()
+    demo.launch(share=False)
+
+if __name__ == "__main__":
+    main() 

portiloop/src/demo/test_offline.py

@@ -2,7 +2,9 @@ import itertools
 import unittest
 from portiloop.src.demo.offline import run_offline
 from pathlib import Path
+import matplotlib.pyplot as plt
 
+from portiloop.src.demo.utils import sleep_stage, xdf2array
 
 class TestOffline(unittest.TestCase):
 
@@ -21,7 +23,7 @@ class TestOffline(unittest.TestCase):
         all_options_iterator = itertools.product(*map(combinatorial_config.get, keys))
         all_options_dicts = [dict(zip(keys, values)) for values in all_options_iterator]
         self.filtered_options = [value for value in all_options_dicts if (value['online_detection'] and value['online_filtering']) or not value['online_detection']]
-        self.xdf_file = Path(__file__).parents[3] / "test_xdf" / "test_file.xdf"
+        self.xdf_file = Path(__file__).parents[3] / "test_file.xdf"
 
 
     def test_all_options(self):
@@ -30,17 +32,20 @@ class TestOffline(unittest.TestCase):
                 self.assertTrue(config['online_filtering'])
 
     def test_single_option(self):
+
+        # Test options correspond to an index in the possible checkbox group options
+        test_options = [0, 1, 2]
+
         res = list(run_offline(
                 self.xdf_file,
-                offline_filtering=True,
-                online_filtering=True,
-                online_detection=True,
-                wamsley=True,
-                lacourse=True,
+                test_options,
                 threshold=0.5,
                 channel_num=2,
-                freq=250))
+                freq=250,
+                stimulation_phase="Peak",
+                buffer_time=0.3))
         print(res)
+        pass
 
     def tearDown(self):
         pass

portiloop/src/demo/utils.py

@@ -13,6 +13,32 @@ STREAM_NAMES = {
 }
 
 
+def sleep_stage(data, threshold=150, group_size=2):
+    """Sleep stage approximation using a threshold and a group size.
+        Returns a numpy array containing all indices in the input data which CAN be used for offline detection. 
+        These indices can then be used to reconstruct the signal from the original data.
+    """
+    # Find all indexes where the signal is above or below the threshold
+    above = np.where(data > threshold)
+    below = np.where(data < -threshold)
+    indices = np.concatenate((above, below), axis=1)[0]
+
+    indices = np.sort(indices)
+    # Get all the indices where the difference between two consecutive indices is larger than 100
+    groups = np.where(np.diff(indices) <= group_size)[0] + 1
+    # Get the important indices
+    important_indices = indices[groups]
+    # Get all the indices between the important indices
+    group_filler = [np.arange(indices[groups[n] - 1] + 1, index) for n, index in enumerate(important_indices)]
+    # Create flat array from fillers
+    group_filler = np.concatenate(group_filler)
+    # Append all group fillers to the indices
+    masked_indices = np.sort(np.concatenate((indices, group_filler)))
+    unmasked_indices = np.setdiff1d(np.arange(len(data)), masked_indices)
+
+    return unmasked_indices
+
+
 class OfflineSleepSpindleRealTimeStimulator(Stimulator):
     def __init__(self):
         self.last_detected_ts = time.time()
@@ -87,15 +113,19 @@ def xdf2array(xdf_path, channel):
     return np.array(csv_list), columns
     
 
-def offline_detect(method, data, timesteps, freq):
+def offline_detect(method, data, timesteps, freq, mask):
+    # Extract only the interesting elements from the mask
+    data_masked = data[mask]
+
     # Get the spindle data from the offline methods
     time = np.arange(0, len(data)) / freq
+    time_masked = time[mask] 
     if method == "Lacourse":
         detector = DetectSpindle(method='Lacourse2018')
-        spindles, _, _ = detect_Lacourse2018(data, freq, time, detector)
+        spindles, _, _ = detect_Lacourse2018(data_masked, freq, time_masked, detector)
     elif method == "Wamsley":
         detector = DetectSpindle(method='Wamsley2012')
-        spindles, _, _ = detect_Wamsley2012(data, freq, time, detector)
+        spindles, _, _ = detect_Wamsley2012(data_masked, freq, time_masked, detector)
     else:
         raise ValueError("Invalid method")
 
@@ -134,18 +164,25 @@ def offline_filter(signal, freq):
 def compute_output_table(online_stimulation, lacourse_spindles, wamsley_spindles):
     # Count the number of spindles detected by each method
     online_stimulation_count = np.sum(online_stimulation)
-    lacourse_spindles_count = sum([1 for index, spindle in enumerate(lacourse_spindles) if spindle == 1 and lacourse_spindles[index - 1] == 0])
-    wamsley_spindles_count = sum([1 for index, spindle in enumerate(wamsley_spindles) if spindle == 1 and wamsley_spindles[index - 1] == 0])
-
-    # Count how many spindles were detected by both online and lacourse
-    both_online_lacourse = sum([1 for index, spindle in enumerate(online_stimulation) if spindle == 1 and lacourse_spindles[index] == 1])
-    # Count how many spindles were detected by both online and wamsley
-    both_online_wamsley = sum([1 for index, spindle in enumerate(online_stimulation) if spindle == 1 and wamsley_spindles[index] == 1])
+    if lacourse_spindles is not None:
+        lacourse_spindles_count = sum([1 for index, spindle in enumerate(lacourse_spindles) if spindle == 1 and lacourse_spindles[index - 1] == 0])
+        # Count how many spindles were detected by both online and lacourse
+        both_online_lacourse = sum([1 for index, spindle in enumerate(online_stimulation) if spindle == 1 and lacourse_spindles[index] == 1])
+    
+    if wamsley_spindles is not None:
+        wamsley_spindles_count = sum([1 for index, spindle in enumerate(wamsley_spindles) if spindle == 1 and wamsley_spindles[index - 1] == 0])
+        # Count how many spindles were detected by both online and wamsley
+        both_online_wamsley = sum([1 for index, spindle in enumerate(online_stimulation) if spindle == 1 and wamsley_spindles[index] == 1])
+    
+    
 
     # Create markdown table with the results
     table = "| Method | Detected spindles | Overlap with Portiloop |\n"
     table += "| --- | --- | --- |\n"
     table += f"| Online | {online_stimulation_count} | {online_stimulation_count} |\n"
-    table += f"| Lacourse | {lacourse_spindles_count} | {both_online_lacourse} |\n"
-    table += f"| Wamsley | {wamsley_spindles_count} | {both_online_wamsley} |\n"
+    if lacourse_spindles is not None:
+        table += f"| Lacourse | {lacourse_spindles_count} | {both_online_lacourse} |\n"
+    if wamsley_spindles is not None:
+        table += f"| Wamsley | {wamsley_spindles_count} | {both_online_wamsley} |\n"
     return table
+    

portiloop/src/stimulation.py

@@ -3,6 +3,8 @@ from enum import Enum
 import time
 from threading import Thread, Lock
 from pathlib import Path
+import numpy as np
+import matplotlib.pyplot as plt
 
 from portiloop.src import ADS
 
@@ -146,20 +148,18 @@ class SleepSpindleRealTimeStimulator(Stimulator):
 
 # Class that delays stimulation to always stimulate peak or through
 class UpStateDelayer:
-    def __init__(self, sample_freq, spindle_freq, peak, time_to_buffer): 
+    def __init__(self, sample_freq, peak, time_to_buffer, stimulate=None): 
         '''
         args:
             sample_freq: int -> Sampling frequency of signal in Hz
             time_to_wait: float -> Time to wait to build buffer in seconds
         '''
         # Get number of timesteps for a whole spindle
-        self.spindle_timesteps = (1/spindle_freq) * sample_freq # s * 
         self.sample_freq = sample_freq
-        self.buffer_size = 1.5 * self.spindle_timesteps
         self.peak = peak
         self.buffer = []
         self.time_to_buffer = time_to_buffer
-        self.stimulate = None
+        self.stimulate = stimulate
         
         self.state = States.NO_SPINDLE
 
@@ -192,10 +192,39 @@ class UpStateDelayer:
                 return True
             return False
 
+    def step_timesteps(self, point):
+        '''
+        Step the delayer, ads a point to buffer if necessary.
+        Returns True if stimulation is actually done
+        '''
+        if self.state == States.NO_SPINDLE:
+            return False
+        elif self.state == States.BUFFERING:
+            self.buffer.append(point)
+            # If we are done buffering, move on to the waiting stage
+            if len(self.buffer) >= self.time_to_buffer * self.sample_freq:
+                # Compute the necessary time to wait
+                self.time_to_wait = self.compute_time_to_wait()
+                self.state = States.DELAYING
+                self.buffer = []
+                self.delaying_counter = 0
+            return False
+        elif self.state == States.DELAYING:
+            # Check if we are done delaying
+            self.delaying_counter += 1
+            if self.delaying_counter >= self.time_to_wait * self.sample_freq:
+                # Actually stimulate the patient after the delay
+                if self.stimulate is not None:
+                    self.stimulate()
+                # Reset state
+                self.time_to_wait = -1
+                self.state = States.NO_SPINDLE
+                return True
+            return False
+
     def detected(self):
         if self.state == States.NO_SPINDLE:
             self.state = States.BUFFERING
-            self.time_started = time.time()
 
     def compute_time_to_wait(self):
         """
@@ -208,8 +237,27 @@ class UpStateDelayer:
         # Returns the index of the last peak in the buffer
         peaks, _ = find_peaks(self.buffer, prominence=1)
 
+        # Make a figure to show the peaks
+        if False:
+            plt.figure()
+            plt.plot(self.buffer)
+            for peak in peaks:
+                plt.axvline(x=peak)
+            plt.plot(np.zeros_like(self.buffer), "--", color="gray")
+            plt.show()
+
+        if len(peaks) == 0:
+            print("No peaks found, increase buffer size")
+            return (self.sample_freq / 10) * (1.0 / self.sample_freq)
+
+        # Compute average distance between each peak
+        avg_dist = np.mean(np.diff(peaks))
+
         # Compute the time until next peak and return it
-        return (len(self.buffer) - peaks[-1]) * (1 / self.sample_freq)
+        if (avg_dist < len(self.buffer) - peaks[-1]):
+            print("Average distance between peaks is smaller than the time to last peak, decrease buffer size")
+            return (len(self.buffer) - peaks[-1]) * (1.0 / self.sample_freq)
+        return (avg_dist - (len(self.buffer) - peaks[-1])) * (1.0 / self.sample_freq)
 
 class States(Enum):
     NO_SPINDLE = 0