kraken/backend/tests/test_training.py

import os
import training
import unittest
import joblib
from glob import glob
from dataloaders import neo
from bci_core.feature_extractors import FeatExtractor
from bci_core.model import baseline_model, stacking_riemann, ChannelScaler
import shutil
from sklearn.utils.validation import check_is_fitted, NotFittedError
from sklearn.pipeline import Pipeline
from sklearn.ensemble import StackingClassifier


class TestTraining(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        root_path = './tests/data'
        sessions = {'cylinder': ['eeg-data/2', 'eeg-data/6'], 'ball': ['eeg-data/4', 'eeg-data/5']}
        cls.event_id = {'rest': 0, 'cylinder': 1, 'ball': 2}

        raw = neo.raw_preprocessing(root_path, sessions, unify_label=True)
        raw.drop_channels(['T3', 'T4', 'A1', 'A2', 'T5', 'T6', 'M1', 'M2', 'Fp1', 'Fp2', 'F7', 'F8', 'O1', 'Oz', 'O2', 'F3', 'F4', 'Fz'])
        cls.raw = raw
    
    def test_training_baseline(self):
        model = training.train_model(self.raw, self.event_id, model_type='baseline')
        check_is_fitted(model)

    def test_saver(self):
        feat_ext = FeatExtractor(1000, lfb_bands=[(15, 30), [30, 45]], hg_bands=[(55, 95), (105, 145)])
        model_riemann = stacking_riemann(12, 12, 1, 1)
        model_baseline = baseline_model(1)
        scaler = ChannelScaler()
        event_id = {'1': 5, '0': 3}
        training.model_saver([feat_ext, scaler, model_riemann, model_baseline], './tests/data', 'baseline', 'f77cbe10a8de473992542e9f4e913a66', event_id)
        self.assertTrue(os.path.isdir(os.path.join('./tests/data', 'f77cbe10a8de473992542e9f4e913a66')))

        model_file = glob(os.path.join('./tests/data', 
                                       'f77cbe10a8de473992542e9f4e913a66', 
                                       '*.pkl'))
        
        self.assertEqual(len(model_file), 1)

        name = os.path.normpath(model_file[0]).split(os.sep)
        class_name, events, date = name[-1].split('_')
        print(class_name, events, date)
        self.assertTrue(class_name == 'baseline')
        self.assertTrue(events == '0+1')
        # load model
        feat, scaler, model_riem, model_base = joblib.load(model_file[0])
        self.assertTrue(isinstance(feat, FeatExtractor))
        self.assertTrue(isinstance(scaler, ChannelScaler))
        self.assertTrue(isinstance(model_riem, StackingClassifier))
        self.assertTrue(isinstance(model_base, Pipeline))

        shutil.rmtree(os.path.join('./tests/data', 'f77cbe10a8de473992542e9f4e913a66'))