-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
292 lines (246 loc) · 10.3 KB
/
train.py
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
import cortex
from cortex import Cortex
import os
from dotenv import load_dotenv
class Train():
"""
A class to use BCI API to control the training of the mental command detections.
Attributes
----------
c : Cortex
Cortex communicate with Emotiv Cortex Service
Methods
-------
start():
to start a training process from starting a websocket
subscribe_data(streams):
To subscribe one or more data streams
unload_profile(profile_name):
To unload an profile
load_profile(profile_name):
To load an profile for training
train_mc_action(status):
To control training of mentalCommand action
Callbacks functions
-------------------
on_create_session_done(*args, **kwargs):
to handle create_session_done which inform when session created successfully
on_query_profile_done(*args, **kwargs):
to handle query_profile_done which inform when query_profile done
on_load_unload_profile_done(*args, **kwargs):
to handle load_unload_profile_done which inform when profile is loaded or unloaded successfully
on_save_profile_done(*args, **kwargs):
to handle save_profile_done which inform when profile is saved successfully
on_new_data_labels(*args, **kwargs):
to handle new_data_labels which inform when sys event is subscribed successfully
on_new_sys_data(*args, **kwargs):
to handle new_sys_data which inform when sys event is streamed
"""
def __init__(self, app_client_id, app_client_secret, **kwargs):
self.c = Cortex(app_client_id, app_client_secret, debug_mode=True, **kwargs)
self.c.bind(create_session_done=self.on_create_session_done)
self.c.bind(query_profile_done=self.on_query_profile_done)
self.c.bind(load_unload_profile_done=self.on_load_unload_profile_done)
self.c.bind(save_profile_done=self.on_save_profile_done)
self.c.bind(new_data_labels=self.on_new_data_labels)
self.c.bind(new_sys_data=self.on_new_sys_data)
self.c.bind(inform_error=self.on_inform_error)
def start(self, profile_name, actions, headsetId=''):
"""
To start training process as below workflow
(1) check access right -> authorize -> connect headset->create session
(2) query profile -> get current profile -> load/create profile -> subscribe sys
(3) start and accept MC action training in the action list one by one
Parameters
----------
profile_name : string, required
name of profile
actions : list, required
list of actions which will be trained
headsetId: string , optional
id of wanted headet which you want to work with it.
If the headsetId is empty, the first headset in list will be set as wanted headset
Returns
-------
None
"""
if profile_name == '':
raise ValueError('Empty profile_name. The profile_name cannot be empty.')
self.profile_name = profile_name
self.actions = actions
self.action_idx = 0
self.c.set_wanted_profile(profile_name)
if headsetId != '':
self.c.set_wanted_headset(headsetId)
self.c.open()
def subscribe_data(self, streams):
"""
To subscribe to one or more data streams
'com': Mental command
'fac' : Facial expression
'sys': training event
Parameters
----------
streams : list, required
list of streams. For example, ['sys']
Returns
-------
None
"""
self.c.sub_request(streams)
def load_profile(self, profile_name):
"""
To load an existed profile or create new profile for training
Parameters
----------
profile_name : str, required
profile name
Returns
-------
None
"""
status = 'load'
self.c.setup_profile(profile_name, status)
def unload_profile(self, profile_name):
"""
To unload an existed profile or create new profile for training
Parameters
----------
profile_name : str, required
profile name
Returns
-------
None
"""
self.c.setup_profile(profile_name, 'unload')
def save_profile(self, profile_name):
"""
To save a profile
Parameters
----------
profile_name : str, required
profile name
Returns
-------
None
"""
self.c.setup_profile(profile_name, 'save')
def get_active_action(self, profile_name):
self.c.get_mental_command_active_action(profile_name)
def get_command_brain_map(self, profile_name):
self.c.get_mental_command_brain_map(profile_name)
def get_training_threshold(self):
self.c.get_mental_command_training_threshold(profile_name)
def train_mc_action(self, status):
"""
To control the training of the mental command action.
Make sure the headset is at good contact quality. You need to focus during 8 seconds for training an action.
For simplicity, the example will train action by action in the actions list
Parameters
----------
status : string, required
to control training: there are 5 types: start, accept, reject, erase, reset
Returns
-------
None
"""
if self.action_idx < len(self.actions):
action = self.actions[self.action_idx]
print('train_mc_action: -----------------------------------: '+ action + ":" + status)
self.c.train_request(detection='mentalCommand',
action=action,
status=status)
else:
# save profile after training
print('train_mc_action: -----------------------------------: Done')
self.c.setup_profile(self.profile_name, 'save')
self.action_idx = 0 # reset action_idx
# callbacks functions
def on_create_session_done(self, *args, **kwargs):
print('on_create_session_done')
self.c.query_profile()
def on_query_profile_done(self, *args, **kwargs):
print('on_query_profile_done')
self.profile_lists = kwargs.get('data')
if self.profile_name in self.profile_lists:
# the profile is existed
self.c.get_current_profile()
else:
# create profile
self.c.setup_profile(self.profile_name, 'create')
def on_load_unload_profile_done(self, *args, **kwargs):
is_loaded = kwargs.get('isLoaded')
if is_loaded == True:
# subscribe sys stream to receive Training Event
self.subscribe_data(['sys'])
else:
print('The profile ' + self.profile_name + ' is unloaded')
self.profile_name = ''
# close socket
self.c.close()
def on_save_profile_done (self, *args, **kwargs):
print('Save profile ' + self.profile_name + " successfully.")
# You can test some advanced bci such as active actions, brain map, and training threshold. before unload profile
self.unload_profile(self.profile_name)
def on_new_sys_data (self, *args, **kwargs):
data = kwargs.get('data')
train_event = data[1]
action = self.actions[self.action_idx]
print('on_new_sys_data: ' + action +" : " + train_event)
if train_event == 'MC_Succeeded':
# train action successful. you can accept the training to complete or reject the training
self.train_mc_action('accept')
elif train_event == 'MC_Failed':
self.train_mc_action("reject")
elif train_event == 'MC_Completed' or train_event == 'MC_Rejected':
# training complete. Move to next action
self.action_idx = self.action_idx + 1
self.train_mc_action('start')
def on_new_data_labels(self, *args, **kwargs):
data = kwargs.get('data')
print('on_new_data_labels')
print(data)
if data['streamName'] == 'sys':
# subscribe sys event successfully
# start training
print('on_new_data_labels: start training ')
self.train_mc_action('start')
def on_inform_error(self, *args, **kwargs):
error_data = kwargs.get('error_data')
error_code = error_data['code']
error_message = error_data['message']
print(error_data)
if error_code == cortex.ERR_PROFILE_ACCESS_DENIED:
# disconnect headset for next use
print('Get error ' + error_message + ". Disconnect headset to fix this issue for next use.")
self.c.disconnect_headset()
# -----------------------------------------------------------
#
# GETTING STARTED
# - Please reference to https://emotiv.gitbook.io/cortex-api/ first.
# - Connect your headset with dongle or bluetooth. You can see the headset via Emotiv Launcher
# - Please make sure the your_app_client_id and your_app_client_secret are set before starting running.
# - The function on_create_session_done, on_query_profile_done, on_load_unload_profile_done will help
# handle create and load an profile automatically . So you should not modify them
# - The functions on_new_data_labels(), on_new_sys_data() will help to control action by action training.
# You can modify these functions to control the training such as: reject an training, use advanced bci api.
# RESULT
# - train mental command action
#
# -----------------------------------------------------------
def main():
# Load environment variables from .env file
load_dotenv()
# Please fill your application clientId and clientSecret before running script
your_app_client_id = os.environ['CLIENT_ID']
your_app_client_secret = os.environ['CLIENT_SECRET']
print(your_app_client_id)
# Init Train
t=Train(your_app_client_id, your_app_client_secret)
profile_name = 'pgame' # set your profile name. If the profile is not exited it will be created.
# list actions which you want to train
actions = ['neutral', 'open', 'shut','lift', 'drop', '']
t.start(profile_name, actions)
if __name__ =='__main__':
main()
# -----------------------------------------------------------