xdf import testing and import markers

Author: patricefortin

data/dyad-example-synthetic.xdf data/dyad-example-with-markers.xdfdata/dyad-example-synthetic.xdf renamed to data/dyad-example-with-markers.xdf

@@ -10,11 +10,11 @@ class XDFImport():
     Read an XDF file and enable to export stream in a convenient format (e.g., an EEG stream into an mne.Raw instance).
 
     Arguments:
-      file_path: Path to LSL data (i.e., XDF file). Can be absolute or relative.
+      file_path: Path to XDF file (LSL data recorded with LabRecorder). Can be absolute or relative.
       stream_type: Define which type of stream the user is looking to convert.
       stream_matches: List of the stream index(es) in the XDF the user wishes to convert (can be `str` which the class will try to match to the name of an existing stream or an `int` which will be interpreted as such). Do not set to convert all of the request type
+      mne_type_map: Dict to map stream types to mne channel types
       scale: Scaling factor or 'auto' for automatic scaling, None for no scaling.
-      save_FIF_path: Boolean indicating whether to save the converted data to FIF format.
     """
 
     file_path: str
@@ -25,18 +25,16 @@ class XDFImport():
 
     def __init__(self,
         file_path: str,
-        stream_type: str = 'EEG',
-        stream_matches: list = None,
+        select_type: str = None,
+        select_matches: list = None,
         mne_type_map: dict = None,
         scale: float | str = None,
-        save_FIF_path: bool = None,
         verbose: bool = False,
         convert_to_mne: bool = True,
         ):
 
         self.file_path = file_path 
         self.scale = scale
-        self.save_FIF_path = save_FIF_path
 
         self.selected_stream_indices = []
         self.map_id_to_idx = {}
@@ -48,44 +46,74 @@ def __init__(self,
 
         self.available_streams = [XDFStream(stream, mne_type_map=mne_type_map) for stream in streams]
 
-
-        if verbose: 
-            self.print_available_streams()
-
         self.map_streams()
 
         # Prepare the "selected_streams" list
-        if stream_matches is None:
-            self.select_streams_by_type(stream_type)
+        if select_matches is not None:
+            self.select_streams_by_matches(select_matches)
         else:
-            self.select_streams_by_matches(stream_matches)
+            self.select_streams_by_type(select_type)
+
+        if verbose: 
+            print(self)
 
         # Create MNE objects
         if convert_to_mne:
             self.convert_streams_to_mne()
 
+    @property
+    def raws_dict(self):
+        ret = dict()
+        for stream in self.selected_data_streams:
+            ret[stream.name] = stream.raw
+        return ret
+    
+    @property
+    def raws(self):
+        return [stream.raw for stream in self.selected_data_streams]
+    
+    @property
+    def annotations_dict(self):
+        ret = dict()
+        for stream in self.selected_markers_streams:
+            ret[stream.name] = stream.annotations
+        return ret
+    
+    @property
+    def annotations(self):
+        return [stream.annotations for stream in self.selected_markers_streams]
+    
+    @property
+    def annotations_flat(self):
+        ret = []
+        for x in self.annotations:
+            ret += x    
+        return ret
+    
     @property
     def selected_streams(self):
         return [self.available_streams[idx] for idx in self.selected_stream_indices]
 
+    @property
+    def selected_data_streams(self):
+        return [self.available_streams[idx] for idx in self.selected_stream_indices if self.available_streams[idx].is_mne_raw_compatible]
+    
+    @property
+    def selected_markers_streams(self):
+        return [self.available_streams[idx] for idx in self.selected_stream_indices if self.available_streams[idx].is_mne_annotations_compatible]
+    
     @property
     def selected_stream_names(self):
         return [stream.name for stream in self.selected_streams]
 
     @property
-    def raw_all(self):
-        ret = dict()
-        for stream in self.selected_streams:
-            ret[stream.name] = stream.raw
-        return ret
+    def selected_data_stream_names(self):
+        return [stream.name for stream in self.selected_data_streams]
+    
+    @property
+    def selected_markers_stream_names(self):
+        return [stream.name for stream in self.selected_markers_streams]
 
-    def print_available_streams(self):
-        print(f"List of available streams in XDF file {self.file_path}:")
-        for stream in self.available_streams:
-            print(f"  Stream id {stream.id} of type '{stream.type}' with name '{stream.name}'")
-            print(f"    Channel names: {','.join(stream.ch_names)}")
-            print(f"    Channel types: {','.join(stream.ch_types)}")
-
     def get_streams_for_type(self, stream_type: str):
         return [stream for stream in self.available_streams if stream.type == stream_type]
 
@@ -113,8 +141,13 @@ def select_streams_by_type(self, stream_type: str) -> list:
 
         if self.verbose:
             print(f"Looking for streams of type '{stream_type}'")
-
-        self.selected_stream_indices = self.get_stream_indices_for_type(stream_type)
+        
+        if stream_type is not None:
+            self.selected_stream_indices = self.get_stream_indices_for_type(stream_type)
+        else:
+            # Use all the mne raw compatible
+            # TODO this if is complicated
+            self.selected_stream_indices = [idx for idx, stream in enumerate(self.available_streams) if stream.is_mne_compatible]
 
         # Assert that we have found at least one real stream of selected type
         if len(self.selected_stream_indices) == 0:
@@ -156,7 +189,7 @@ def convert_streams_to_mne(self):
         """
 
         # Find if all the stream have unique names (true if any stream name is duplicated)
-        names = [stream.name for stream in self.selected_streams]
+        names = self.selected_stream_names
 
         for stream in self.selected_streams:
             if self.verbose:
@@ -168,18 +201,14 @@ def convert_streams_to_mne(self):
 
             stream.convert_to_mne(self.scale, append_stream_id=has_duplicate_names)
 
-            # Save file is asked too
-            if self.save_FIF_path is not None: 
-                stream.save_fif_file(self.save_FIF_path)
-
         if self.verbose:
             print("Convertion done.")
 
-    def save_fif_files(self, dir_path):
-        return [stream.save_fif_file(dir_path) for stream in self.selected_streams]
+    def save_to_fif_files(self, dir_path):
+        return [stream.save_to_fif_file(dir_path) for stream in self.selected_data_streams]
 
     def rename_channels(self, new_names):
-        return [stream.rename_channels(new_names) for stream in self.selected_streams]
+        return [stream.rename_channels(new_names) for stream in self.selected_data_streams]
 
     def set_montage(self, montage):
         """
@@ -190,12 +219,22 @@ def set_montage(self, montage):
             montage: A path to a local Dig montage or a mne standard montage.
         """
         if self.verbose or True:
-            print(f"Setting '{montage}' as the montage for streams: {','.join(self.selected_stream_names)}")
+            print(f"Setting '{montage}' as the montage for streams: {','.join(self.selected_data_stream_names)}")
 
 
-        for stream in self.selected_streams:
+        for stream in self.selected_data_streams:
             try: 
                 stream.set_montage(montage)
             except ValueError as e:
                 warnings.warn(f"Invalid montage given to mne.set_montage(): {montage}")
                 raise e
+
+    def __str__(self):
+        available_streams_str = "\n  ".join([str(stream) for stream in self.available_streams])
+        selected_streams_str = ",".join([stream.name for stream in self.selected_streams])
+        return f"""XDFImport with {len(self.available_streams)} available streams and {len(self.selected_streams)} selected streams (loaded from {self.file_path})
+Available streams:
+  {available_streams_str}
+Selected streams: [{selected_streams_str}]
+"""
+    

hypyp/xdf/xdf_import.py

@@ -1,3 +1,4 @@
+from typing import List
 import os
 
 import mne
@@ -6,12 +7,13 @@
 
 class XDFStream():
     stream: StreamData
-    metadata: dict
     metadata_desc: dict
-    raw: mne.io.RawArray
+    raw: mne.io.RawArray | None
+    annotations: mne.Annotations | None
+    time_offset: float
 
     @staticmethod
-    def stream_type_to_mne_type(stream_type: str, type_map: dict = None):
+    def stream_type_to_mne_ch_type(stream_type: str, type_map: dict = None):
         if type_map is not None and stream_type in type_map.keys():
             return type_map[stream_type]
 
@@ -39,12 +41,12 @@ def stream_type_to_mne_type(stream_type: str, type_map: dict = None):
     @staticmethod
     def get_mne_ch_types(stream_type: str, ch_names: list, type_map: dict = None):
         ch_types = []
-        base_type = XDFStream.stream_type_to_mne_type(stream_type, type_map=type_map)
+        type_from_stream = XDFStream.stream_type_to_mne_ch_type(stream_type, type_map=type_map)
         for ch_name in ch_names:
-            ch_type = base_type
+            ch_type = type_from_stream
             if stream_type.startswith('Acc'):
                 ch_type = 'misc'
-            if ch_name.startswith('Marker'):
+            if ch_name.startswith('Markers'):
                 ch_type = 'stim'
             if ch_name.startswith('Acc'):
                 ch_type = 'misc'
@@ -61,9 +63,14 @@ def __init__(self, stream, mne_type_map: dict = None):
         self.stream = stream
         self.mne_type_map = mne_type_map
         self.metadata_desc = {}
+        self.raw = None
+        self.annotations = None
         if self.stream["info"]["desc"] and self.stream["info"]["desc"][0]:
             self.metadata_desc = self.stream["info"]["desc"][0]
-
+        
+        self.time_offset = float(self.stream['info']['created_at'][0])
+        
+        
     @property
     def name(self):
         return self.stream["info"]["name"][0]
@@ -80,10 +87,28 @@ def id(self):
     def srate(self):
         return float(self.stream["info"]["nominal_srate"][0])
 
+    @property
+    def is_mne_compatible(self):
+        return self.is_mne_raw_compatible or self.is_mne_annotations_compatible
+
+    @property
+    def is_mne_raw_compatible(self):
+        if self.type.lower() == 'quality':
+            return False
+        return self.srate > 0
+
+    @property
+    def is_mne_annotations_compatible(self):
+        return self.srate == 0.0
+
     @property
     def time_series(self):
         # Here we check wether the data is in the correct shape () and transpose it if necessary
         # ! We assume that no EEG recoding would have more channels than sample point !
+        # TODO debug time stamps
+        # print("BBBBBBBBBBBBBBBBBBB")
+        # print(self.time_offset)  # Time stamps of the markers
+        # print(self.stream['time_stamps'])  # Time stamps of the markers
         if self.stream["time_series"].shape[0] > self.stream["time_series"].shape[1]:
             return self.stream["time_series"].T
 
@@ -118,12 +143,45 @@ def ch_names(self):
     def ch_types(self):
         return XDFStream.get_mne_ch_types(self.type, self.ch_names, self.mne_type_map)
 
+    def create_mne_annotations(self):
+        if self.type == 'Markers':
+            timestamps = self.stream['time_stamps']  # Time stamps of the markers
+            descriptions = self.stream['time_series']  # Marker descriptions (the event labels)
+
+            # Ensure timestamps are in seconds (if they aren't already)
+            timestamps = np.array(timestamps) - self.time_offset
+            onset_times = timestamps
+            duration = np.zeros_like(onset_times)
+
+            # Create the description for each marker (can be customized, here we use the marker text)
+            description = [str(desc[0]) for desc in descriptions]  # Flatten description if necessary
+
+            # Create the annotations object
+            self.annotations = mne.Annotations(onset=onset_times, duration=duration, description=description)
+            # TODO debug time stamps
+            # print("AAAAAAAAAAAAAAAAAAAAAAAAAaaa")
+            # print(onset_times)
+
+        else:
+            self.markers = None
+
+    def get_unique_name(self, append_stream_id: bool = False):
+        unique_stream_name = self.name
+        if append_stream_id:
+            unique_stream_name += f'-{self.id}'
+        return unique_stream_name
+        
+        
     def convert_to_mne(self, scale: float | str | None, append_stream_id: bool = False, verbose: bool = False):
-        mne_info = self.create_mne_info(append_stream_id=append_stream_id, verbose=verbose)
-        self.raw = self.create_mne_raw(mne_info, scale, verbose=verbose)
-        self.unique_stream_name = mne_info['subject_info']['his_id']
+        self.unique_name = self.get_unique_name(append_stream_id=append_stream_id)
+        if self.is_mne_raw_compatible:
+            mne_info = self.create_mne_info(verbose=verbose)
+            self.raw = self.create_mne_raw(mne_info, scale, verbose=verbose)
+
+        if self.is_mne_annotations_compatible:
+            self.create_mne_annotations()
 
-    def create_mne_info(self, append_stream_id: bool = False, verbose: bool = False):
+    def create_mne_info(self, verbose: bool = False):
         """
         Create a mne.info object from the XDF's EEG stream metadata.
 
@@ -139,14 +197,9 @@ def create_mne_info(self, append_stream_id: bool = False, verbose: bool = False)
 
         mne_info = mne.create_info(ch_names, ch_types=ch_types, sfreq=self.srate, verbose=verbose)
 
-        unique_stream_name = self.name
-        if append_stream_id:
-            unique_stream_name += f'-{self.id}'
-        
         mne_info["subject_info"] = {
-            "id": unique_stream_name,
             "id": self.id, # integer identifier of the subject
-            "his_id": self.name, # string identifier of the subject
+            "his_id": self.unique_name, # string identifier of the subject
         }
 
         return mne_info
@@ -222,10 +275,17 @@ def rename_channels(self, new_names):
 
         self.raw.rename_channels(mapping)
 
-    def save_fif_file(self, dir_path: str):
+    def save_to_fif_file(self, dir_path: str):
         os.makedirs(dir_path, exist_ok=True)
-        save_file_path = os.path.join(dir_path, f"{self.unique_stream_name}_raw.fif")
+        save_file_path = os.path.join(dir_path, f"{self.unique_name}_raw.fif")
         self.raw.save(save_file_path, overwrite=True)
         return save_file_path
 
-    
+    
+    def __str__(self):
+        ch_names = ','.join(self.ch_names)
+        ch_types = ','.join(self.ch_types)
+        return f"""Stream id {self.id} of type '{self.type}' with name '{self.name}'
+    Sampling Rate: {self.srate}
+    Channel names: [{ch_names}]
+    Channel types: [{ch_types}]"""