Merge pull request #7 from Portiloop/milo/algotool-trains

portiloop/src/demo/demo.py

@@ -29,10 +29,15 @@ def main():
             # Threshold value
             threshold = gr.Slider(0, 1, value=0.82, step=0.01, label="Threshold", interactive=True)
             # Detection Channel
+        
+        with gr.Row():
             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)
 
+            # Detect trains dropdown
+            detect_trains = gr.Dropdown(choices=["All Spindles", "Isolated & First", "Trains"], value="All Spindles", label="Detection mode:", interactive=True)
+
         with gr.Row():
             output_array = gr.File(label="Output CSV File")
             output_table = gr.Markdown(label="Output Table")
@@ -45,7 +50,8 @@ def main():
                 detect_filter,
                 threshold, 
                 detect_channel,
-                freq], 
+                freq,
+                detect_trains], 
             outputs=[output_array, output_table])
 
     demo.queue()

portiloop/src/demo/offline.py

@@ -2,11 +2,11 @@ import numpy as np
 from portiloop.src.detection import SleepSpindleRealTimeDetector
 from portiloop.src.stimulation import UpStateDelayer
 from portiloop.src.processing import FilterPipeline
-from portiloop.src.demo.utils import compute_output_table, sleep_stage, xdf2array, offline_detect, offline_filter, OfflineSleepSpindleRealTimeStimulator
+from portiloop.src.demo.utils import OfflineIsolatedSpindleRealTimeStimulator, OfflineSpindleTrainRealTimeStimulator, 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, stimulation_phase="Fast", buffer_time=0.25):
+def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, detect_trains, 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
@@ -76,7 +76,14 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, stim
     # Create the detector
     if online_detection:
         detector = SleepSpindleRealTimeDetector(threshold=threshold, channel=1) # always 1 because we have only one channel
-        stimulator = OfflineSleepSpindleRealTimeStimulator()
+
+        if detect_trains == "All Spindles":
+            stimulator = OfflineSleepSpindleRealTimeStimulator()
+        elif detect_trains == "Trains":
+            stimulator = OfflineSpindleTrainRealTimeStimulator()
+        elif detect_trains == "Isolated & First":
+            stimulator = OfflineIsolatedSpindleRealTimeStimulator()
+
         if stimulation_phase != "Fast":
             stimulation_delayer = UpStateDelayer(freq, stimulation_phase == 'Peak', time_to_buffer=buffer_time, stimulate=lambda: None)
             stimulator.add_delayer(stimulation_delayer)
@@ -139,9 +146,11 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, stim
     # Output the data to a csv file
     np.savetxt("output.csv", data_whole, delimiter=",", header=",".join(columns), comments="")
 
+    # Compute the overlap of online stimulations with the 
 
     output_table = compute_output_table(
-        data_whole[:, columns.index("online_stimulations")] if online_detection else data_whole[:, columns.index("online_stimulations_portiloop")],
+        data_whole[:, columns.index("online_stimulations")],
+        data_whole[:, columns.index("online_stimulations_portiloop")],
         data_whole[:, columns.index("lacourse_spindles")] if lacourse else None, 
         data_whole[:, columns.index("wamsley_spindles")] if wamsley else None,)
 

portiloop/src/demo/utils.py

@@ -39,6 +39,7 @@ def sleep_stage(data, threshold=150, group_size=2):
     return unmasked_indices
 
 
+
 class OfflineSleepSpindleRealTimeStimulator(Stimulator):
     def __init__(self):
         self.last_detected_ts = time.time()
@@ -70,6 +71,54 @@ class OfflineSleepSpindleRealTimeStimulator(Stimulator):
         self.delayer = delayer
         self.delayer.stimulate = lambda: True
 
+
+class OfflineSpindleTrainRealTimeStimulator(OfflineSleepSpindleRealTimeStimulator):
+    def __init__(self):
+        super().__init__()
+        self.max_spindle_train_t = 6.0
+    
+    def stimulate(self, detection_signal):
+        self.index += 1
+        stim = False
+        for sig in detection_signal:
+            # We detect a stimulation
+            if sig:
+                # Record time of stimulation
+                ts = self.index
+                
+                elapsed = ts - self.last_detected_ts
+                # Check if time since last stimulation is long enough
+                if self.wait_timesteps < elapsed < int(self.max_spindle_train_t * 250):
+                    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
+    
+class OfflineIsolatedSpindleRealTimeStimulator(OfflineSpindleTrainRealTimeStimulator):
+    def stimulate(self, detection_signal):
+        self.index += 1
+        stim = False
+        for sig in detection_signal:
+            # We detect a stimulation
+            if sig:
+                # Record time of stimulation
+                ts = self.index
+                
+                elapsed = ts - self.last_detected_ts
+                # Check if time since last stimulation is long enough
+                if int(self.max_spindle_train_t * 250) < elapsed:
+                    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 xdf2array(xdf_path, channel):
     xdf_data, _ = pyxdf.load_xdf(xdf_path)
 
@@ -162,7 +211,14 @@ def offline_filter(signal, freq):
 
     return signal
 
-def compute_output_table(online_stimulation, lacourse_spindles, wamsley_spindles):
+def compute_output_table(irl_online_stimulations, online_stimulation, lacourse_spindles, wamsley_spindles, time_overlap_s=2.0):
+
+
+    # Count the number of spindles in this run which overlap with spindles found IRL
+    irl_spindles_count = sum(irl_online_stimulations)
+    both_online_irl = sum([1 for index, spindle in enumerate(online_stimulation)\
+         if spindle == 1 and 1 in irl_online_stimulations[index - int((time_overlap_s / 2) * 250):index + int((time_overlap_s / 2) * 250)]])
+
     # Count the number of spindles detected by each method
     online_stimulation_count = np.sum(online_stimulation)
     if lacourse_spindles is not None:
@@ -175,12 +231,11 @@ def compute_output_table(online_stimulation, lacourse_spindles, wamsley_spindles
         # 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 = "| Method | # of Detected spindles | Overlap with Online (in tool) |\n"
     table += "| --- | --- | --- |\n"
-    table += f"| Online | {online_stimulation_count} | {online_stimulation_count} |\n"
+    table += f"| Online in Tool | {online_stimulation_count} | {online_stimulation_count} |\n"
+    table += f"| Online detection IRL | {irl_spindles_count} | {both_online_irl} |\n"
     if lacourse_spindles is not None:
         table += f"| Lacourse | {lacourse_spindles_count} | {both_online_lacourse} |\n"
     if wamsley_spindles is not None: