| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152 |
- import unittest
- from dataloaders import neo
- import mne
- import numpy as np
- class TestDataloader(unittest.TestCase):
- def test_load_sample_data(self):
- root_path = './tests/data'
- sessions = {'cylinder': ['eeg-data/2', 'eeg-data/6'], 'ball': ['eeg-data/4', 'eeg-data/5']}
- event_id = {'rest': 0, 'cylinder': 1, 'ball': 2}
- raw = neo.raw_preprocessing(root_path, sessions, unify_label=True)
- events, event_id = mne.events_from_annotations(raw, event_id=event_id)
- events, events_cnt = np.unique(events[:, -1], return_counts=True)
- self.assertTrue(np.allclose(events_cnt, (300, 150, 150)))
- def test_load_session(self):
- root_path = './tests/data'
- sessions = {'cylinder': ['eeg-data/2', 'eeg-data/6'], 'ball': ['eeg-data/4']}
- raws = neo.load_sessions(root_path, sessions)
- # test if interleaved
- sess_f = tuple(f for f, r in raws)
- self.assertEqual(len(raws), 3)
- self.assertTupleEqual(sess_f, ('cylinder', 'ball', 'cylinder'))
- def test_event_parser(self):
- # fixed length
- fs = 100
- test_event = np.array([[0, 0, 4], [100, 0, 4], [600, 0, 3], [700, 0, 3], [1000, 0, 4], [1100, 0, 4]])
- gt = np.array([[0, 400, 0], [600, 400, 2], [1000, 400, 0]])
- ret = neo.reconstruct_events(test_event, fs, 'ball', trial_duration=4)
- self.assertTrue(np.allclose(ret, gt))
- # varing length
- gt = np.array([[0, 600, 0], [600, 400, 2], [1000, 100, 0]])
- ret = neo.reconstruct_events(test_event, fs, 'ball', trial_duration=None)
- self.assertTrue(np.allclose(ret, gt))
- # change indices
- gt = np.array([[0, 400, 2], [600, 400, 0], [1000, 400, 2]])
- ret = neo.reconstruct_events(test_event, fs, 'ball', mov_trial_ind=[4], rest_trial_ind=[2, 3], trial_duration=4)
- self.assertTrue(np.allclose(ret, gt))
- # use original indices
- gt = np.array([[0, 400, 4], [600, 400, 3], [1000, 400, 4]])
- ret = neo.reconstruct_events(test_event, fs, None, trial_duration=4, use_original_label=True, mov_trial_ind=[3], rest_trial_ind=[4])
- self.assertTrue(np.allclose(ret, gt))
- # use original indices,
- gt = np.array([[0, 400, 0], [600, 400, 1], [1000, 400, 0]])
- test_event = np.array([[0, 0, 0], [100, 0, 0], [600, 0, 1], [700, 0, 1], [1000, 0, 0], [1100, 0, 0]])
- ret = neo.reconstruct_events(test_event, fs, None, trial_duration=4, use_original_label=True, mov_trial_ind=[1], rest_trial_ind=[0])
- self.assertTrue(np.allclose(ret, gt))
|
