Added sleep trains options on the tool

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)

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)