2
0

test_validation.py 2.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. import unittest
  2. import os
  3. import numpy as np
  4. from glob import glob
  5. import shutil
  6. from bci_core import utils as ana_utils
  7. from bci_core.online import model_loader
  8. from training import train_model, model_saver
  9. from dataloaders import library_ieeg
  10. from online_sim import simulation
  11. from validation import val_by_epochs
  12. class TestOnlineSim(unittest.TestCase):
  13. @classmethod
  14. def setUpClass(cls):
  15. root_path = './tests/data'
  16. cls.event_id = {'ball': 2, 'rest': 0}
  17. raw = library_ieeg.raw_preprocessing(os.path.join(root_path, 'ecog-data/1', 'bp_mot_t_h.mat'), finger_model='ball')
  18. raw = raw.pick_channels([raw.info['ch_names'][i] for i in [5,6,7,12,13,14,20,21]])
  19. cls.raw = raw
  20. # split into 2 pieces
  21. t_min, t_max = raw.times[0], raw.times[-1]
  22. t_mid = raw.times[len(raw.times) // 2]
  23. raw_train = raw.copy().crop(tmin=t_min, tmax=t_mid, include_tmax=True)
  24. cls.raw_val = raw.copy().crop(tmin=t_mid, tmax=t_max)
  25. # reconstruct single event for validation
  26. if cls.raw_val.annotations.onset[0] > t_mid:
  27. # correct time by first timestamp
  28. cls.raw_val.annotations.onset -= t_mid
  29. # train with the first half
  30. model = train_model(raw_train, event_id=cls.event_id, model_type='baseline')
  31. model_saver(model, './tests/data/', 'baseline', 'test', cls.event_id)
  32. cls.model_path = glob(os.path.join('./tests/data/', 'test', '*.pkl'))[0]
  33. @classmethod
  34. def tearDownClass(cls) -> None:
  35. shutil.rmtree(os.path.join('./tests/data/', 'test'))
  36. return super().tearDownClass()
  37. def test_event_metric(self):
  38. event_gt = np.array([[0, 0, 0], [5, 0, 1], [7, 0, 0], [9, 0, 2]])
  39. event_pred = np.array([[1, 0, 0], [4, 0, 1], [6, 0, 1], [7, 0, 0], [10, 0, 1], [11, 0, 2]])
  40. fs = 1
  41. precision, recall, f1_score = ana_utils.event_metric(event_gt, event_pred, fs, ignore_event=(0,))
  42. self.assertEqual(f1_score, 2 / 3)
  43. self.assertEqual(precision, 1 / 2)
  44. self.assertEqual(recall, 1)
  45. def test_sim(self):
  46. model = model_loader(self.model_path,
  47. state_change_threshold=0.7,
  48. state_trans_prob=0.7)
  49. metric_hmm, metric_nohmm, fig_pred = simulation(self.raw, self.event_id, model=model, epoch_length=1., step_length=0.1)
  50. fig_pred.savefig('./tests/data/pred.pdf')
  51. print(metric_hmm)
  52. self.assertTrue(metric_hmm[-2] > 0.9) # f1-score (with hmm)
  53. self.assertTrue(metric_nohmm[-2] < 0.5) # f1-score (without hmm)
  54. def test_val_model(self):
  55. metrices, fig_conf = val_by_epochs(self.raw_val, self.model_path, self.event_id, 1.)
  56. fig_conf.savefig('./tests/data/conf.pdf')
  57. self.assertGreater(metrices[0], 0.85)
  58. self.assertGreater(metrices[1], 0.7)
  59. self.assertGreater(metrices[2], 0.7)
  60. if __name__ == '__main__':
  61. unittest.main()