def __init__(self, user_id, module_id, device_type, rmq_address, rmq_user, rmq_pwd, step_size):
self.user_id = user_id
self.module_id = module_id
self.device_type = device_type
self.rmq_address = rmq_address
self.rmq_user = rmq_user
self.rmq_pwd = rmq_pwd
self.eeg_subscriber = None
self.mu_publisher = None
self.routing_keys = {
_EEG: _ROUTING_KEY % (user_id, device_type, _EEG),
_MU: _ROUTING_KEY % (user_id, device_type, _MU),
}
self.preprocessor = None
self.eeg_data = np.zeros((0, 8))
self.count = 0
self.eyeblinks_remover = EyeBlinksRemover()
self.step_size = step_size
self.started_fit = False
class PreprocessingModule(object):
def __init__(self, user_id, module_id, device_type, rmq_address, rmq_user, rmq_pwd, step_size):
self.user_id = user_id
self.module_id = module_id
self.device_type = device_type
self.rmq_address = rmq_address
self.rmq_user = rmq_user
self.rmq_pwd = rmq_pwd
self.eeg_subscriber = None
self.mu_publisher = None
self.routing_keys = {
_EEG: _ROUTING_KEY % (user_id, device_type, _EEG),
_MU: _ROUTING_KEY % (user_id, device_type, _MU),
}
self.preprocessor = None
self.eeg_data = np.zeros((0, 8))
self.count = 0
self.eyeblinks_remover = EyeBlinksRemover()
self.step_size = step_size
self.started_fit = False
def initialize(self):
"""
Initialize EEG preprocessor, publisher, and subscriber
"""
self.mu_publisher = PikaPublisher(self.rmq_address, self.rmq_user, self.rmq_pwd)
self.eeg_subscriber = PikaSubscriber(self.rmq_address, self.rmq_user, self.rmq_pwd)
self.eeg_subscriber.connect()
self.mu_publisher.connect()
self.mu_publisher.register(self.routing_keys[_MU])
self.eeg_subscriber.subscribe(self.routing_keys[_EEG])
def start(self):
_LOGGER.info("[Module %s] Starting Preprocessing. Routing " "keys: %s" % (self.module_id, self.routing_keys))
self.eeg_subscriber.consume_messages(self.routing_keys[_EEG], self._preprocess)
def refit_ica(self):
t = Thread(target=self.eyeblinks_remover.fit, args=(self.eeg_data[1000:],))
t.start()
# self.eyeblinks_remover.fit(self.eeg_data[1000:])
def _preprocess(self, ch, method, properties, body):
eeg = json.loads(body)
self.eeg_data = np.vstack([self.eeg_data, get_raw(eeg)])
# self.count += len(eeg)
self.count += self.step_size
print(self.count)
if (self.count >= 5000 and not self.started_fit) or self.count % 10000 == 0:
_LOGGER.info("refitting...")
self.started_fit = True
self.refit_ica()
timestamp = eeg[-1]["timestamp"]
eeg = from_raw(self.eyeblinks_remover.transform(get_raw(eeg)))
process = preprocess_stft(eeg, _METADATA)
mu_left = process["left"][-1]
mu_right = process["right"][-1]
data = {"timestamp": timestamp, "left": mu_left, "right": mu_right}
_LOGGER.debug("--> mu: %s" % data)
if self.eyeblinks_remover.fitted:
self.mu_publisher.publish(self.routing_keys[_MU], data)
