Modified cvep_pomdp.py and other files to work with the public releas…

…e of the cvep dataset

pomdp_bci/config/cvep.json

@@ -1,7 +1,7 @@
 {
-  "n_subs": 10,
-  "sub_list": ["juan", "ludo", "simon", "noemie", "emilie", "marcel", "max", "quentin", "rebai", "felix"],
+  "n_subs": 11,
+  "excluded_subs": [2],
   "ch_slice": ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18", "19", "20", "21", "22", "23", "24", "25", "26", "27", "28", "29", "30", "31", "32"],
   "sfreq": 500,
   "max_augmentation": 3
-}
+}

pomdp_bci/cvep_pomdp.py

@@ -28,7 +28,7 @@
                             get_code_prediction, make_preds_accumul_aggresive
 
 
-def fit_clf(win_data, win_labels):
+def fit_clf(win_data, win_labels, sfreq, algo='EEGnet_patchembeddingdilation'):
     """
     Return a fit classifier using the selected data and architecture
 
@@ -58,7 +58,7 @@ def fit_clf(win_data, win_labels):
     win_labels = np.squeeze(win_labels[index])
 
     # Initialize NN
-    win_samples = int(code_win_len * params['sfreq'])
+    win_samples = int(code_win_len * sfreq)
     n_channels = win_data.shape[1]  # Number of channels in the data (for channels last)
 
     if algo == 'EEGnet_patchembeddingdilation':
@@ -125,12 +125,13 @@ def fit_clf(win_data, win_labels):
 
 for dataset, algo in itertools.product(datasets, algos):
     # Load dataset-specific parameters
-    with open(f'{config/dataset}.json', 'r') as dataset_params:
+    with open(f'config/{dataset}.json', 'r') as dataset_params:
         params = json.loads(dataset_params.read())
 
     downsample = int(params['sfreq'] / 250)  # After downsampling, sfreq should be 250Hz
 
-    sub_list = params['sub_list']
+    excluded_subs = params['excluded_subs']
+    sub_list = [sub_n for sub_n in range(params['n_subs']) if sub_n not in excluded_subs]
     score_dict = {sub: {} for sub in sub_list}
     metadata_dict = {}
 
@@ -191,7 +192,7 @@ def fit_clf(win_data, win_labels):
         data_cv /= cal_std + 1e-8
 
         # Fit
-        pomdp_clf = fit_clf(data_cal, labels_cal, algorithm=algo)
+        pomdp_clf = fit_clf(data_cal, labels_cal, sfreq=params['sfreq'], algo=algo)
         print('Data was fit')
 
         # Predict the codes of validation data (split in 10 to avoid OOM from the GPU)

pomdp_bci/utils/utils_problem.py

@@ -84,7 +84,7 @@ def load_data(subject, dataset, eeg_path=None, ch_keep=[]):
         event_id = nakanishi.event_id
 
     elif dataset == 'cvep':
-        filename = f"{subject}_mseqwhite.set"
+        filename = f"P{subject}_whitemseq.set"
 
         file_path = os.path.join(eeg_path, filename)
         raw = mne.io.read_raw_eeglab(file_path, preload=True, verbose=False)
@@ -114,8 +114,6 @@ def load_data(subject, dataset, eeg_path=None, ch_keep=[]):
 
     # CVEP needs the montage manually set
     if dataset == 'cvep':
-        montage = get_liveamp_montage(eeg_path)
-        raw.set_montage(montage)
         raw = raw.drop_channels(['21', '10'])
 
     mne.set_eeg_reference(raw, 'average', copy=False, verbose=False)