Add files via upload

Author: snailpt

prepare_dataset.py

@@ -0,0 +1,81 @@
+"""
+# Function: Read files and preprocess them
+# Steps:
+# 1. Import data from the gdf file provided before the competition, 
+#    remove unwanted channels, and select required events.
+# 2. Select desired time segments for slicing; treat each segment (4s) as one sample.
+# 3. Import labels from the mat file provided after the competition, 
+#    ensuring they correspond with epochs and their numbers match.
+# 4. Save the resulting data in a new mat file, 
+#    preparing it for use in the subsequent main.py.
+"""
+
+import mne
+import numpy as np
+import scipy.signal as signal
+from scipy.io import savemat
+import scipy.io as sio
+import numpy as np
+
+def changeGdf2Mat(dir_path, mode="train"):
+    '''
+    read data from GDF files and store as mat files
+
+    Parameters
+    ----------
+    dir_path : str
+        GDF file dir path.
+    mode : str, optional
+        change train dataset or eval dataset. The default is "train".
+
+    Returns
+    -------
+    None.
+
+    '''
+    mode_str = ''
+    if mode=="train":
+        mode_str = 'T'
+    else:
+        mode_str = 'E'
+    for nSub in range(1, 10):
+        # Load the gdf file
+        data_filename = dir_path+'BCICIV_2a_gdf/A0{}{}.gdf'.format(nSub, mode_str)
+        raw = mne.io.read_raw_gdf(data_filename)  
+    
+        # Select the events of interest
+        events, event_dict = mne.events_from_annotations(raw) 
+        if mode=="train":
+            # train dataset are labeled
+            event_id = {'Left': event_dict['769'],
+                        'Right': event_dict['770'], 
+                        'Foot': event_dict['771'],
+                        'Tongue': event_dict['772']}  
+        else:
+            # evaluate dataset are labeled as 'Unknnow'
+            event_id = {'Unknown': event_dict['783']}
+            
+        # 选取我们关心的四个类别对应的事件，这里events[:, 2]是指events中的第三列，即事件的编号。
+        selected_events = events[np.isin(events[:, 2], list(event_id.values()))]  
+        
+        # remove EOG channels
+        raw.info['bads'] += ['EOG-left', 'EOG-central', 'EOG-right']
+        picks = mne.pick_types(raw.info, meg=False, eeg=True, eog=False, stim=False, exclude='bads')
+        # Epoch the data
+        # using 4s (1000 sample point ) segmentation
+        epochs = mne.Epochs(raw, selected_events, event_id, picks=picks,tmin=0, tmax=3.996, preload=True, baseline=None)
+        
+        filtered_data = epochs.get_data()
+        label_filename = dir_path + 'true_labels/'+'A0{}{}.mat'.format(nSub, mode_str)
+        mat = sio.loadmat(label_filename)  # load target mat file
+        labels = mat['classlabel']
+       
+        # Save the data and labels to a .mat file
+        result_filename = 'mymat_raw/A0{}{}.mat'.format(nSub, mode_str)
+        savemat(result_filename, {'data': filtered_data, 'label': labels})
+
+dir_path = './'
+# prepare train dataset
+changeGdf2Mat(dir_path, 'train')
+# prepare test dataset
+changeGdf2Mat(dir_path, 'eval')