Working for CHI

Author: eeddy

libemg/discrete.py

@@ -0,0 +1,87 @@
+import numpy as np
+import torch.nn.functional as F
+import torch 
+from playsound import playsound
+import matplotlib.pyplot as plt
+from libemg.feature_extractor import FeatureExtractor
+from libemg.utils import get_windows
+import pyautogui
+import time 
+import statistics
+
+class DiscreteControl:
+    """
+    A class for detecting gestures using the Teager-Kaiser energy operator on EMG signals.
+    """
+    
+    def __init__(self, odh, window_size, increment, threshold=100, buffer=20, subject=None, model=None):
+        self.odh = odh
+        self.window_size = window_size
+        self.increment = increment
+        self.threshold = threshold
+        self.buffer_size = buffer
+        self.subject = subject
+        self.model = model 
+    
+    def run(self):
+        """
+        Main loop for gesture detection.
+        Continuously monitors EMG data and detects gestures based on energy thresholds.
+        """
+        gesture_mapping = ['Nothing', 'Close', 'Flexion', 'Extension', 'Open', 'Pinch']
+        expected_count = 250
+
+        while True:
+            buffer = []
+
+            # Get and process EMG data
+            data, counts = self.odh.get_data(self.window_size)
+            if counts['emg'][0][0] >= expected_count:
+                data, counts = self.odh.get_data(250)
+                emg = data['emg'][::-1]
+                feats = self.get_features([emg], 10, 5, None, None)
+                pred, _ = self.predict(feats[0])
+                buffer.append(pred)
+                mode_pred = statistics.mode(buffer[-20:])
+                if mode_pred != 0: 
+                    print(str(time.time()) + ' ' + gesture_mapping[mode_pred])
+                    self.key_press(mode_pred, gesture_mapping)
+                    self.odh.reset()
+                    expected_count = 250
+                    buffer = []
+                else:
+                    expected_count += 10
+
+    def key_press(self, pred, mapping):
+        if mapping[pred] == 'Close':
+            pyautogui.press('c')
+        elif mapping[pred] == 'Flexion':
+            pyautogui.press('f')
+        elif mapping[pred] == 'Extension':
+            pyautogui.press('e')
+        elif mapping[pred] == 'Open':
+            pyautogui.press('o')
+        elif mapping[pred] == 'Pinch':
+            pyautogui.press('p')
+        playsound('Other/click.wav')
+    
+    def predict(self, gest, device='cpu'):
+        g_tensor = torch.tensor([gest], dtype=torch.float32).to(device)
+        with torch.no_grad(): 
+            output = self.model.forward_once(g_tensor)
+            pred = output.argmax(dim=1).item()
+            prob = F.softmax(output, dim=1).max().item()
+        return pred, prob
+    
+    def get_features(self, data, window_size, window_inc, feats, feat_dic):
+        fe = FeatureExtractor()
+        data = np.array([get_windows(d, window_size, window_inc) for d in data], dtype='object')
+        if feats is None:
+            return data 
+        if feat_dic is not None:
+            feats = np.array([fe.extract_features(feats, d, array=True, feature_dic=feat_dic) for d in data], dtype='object')
+        else:
+            feats = np.array([fe.extract_features(feats, np.array(d, dtype='float'), array=True) for d in data], dtype='object')
+        feats = np.nan_to_num(feats, copy=True, nan=0, posinf=0, neginf=0)
+        # expected shape: (NFiles,) -> (Time, channel)
+        return feats

libemg/emg_predictor.py

@@ -27,7 +27,6 @@
 import re
 from matplotlib.animation import FuncAnimation
 from functools import partial
-from typing import Callable
 
 from libemg.utils import get_windows
 from libemg.environments.controllers import RegressorController, ClassifierController
@@ -308,34 +307,21 @@ def add_majority_vote(self, num_samples=5):
 
     def add_velocity(self, train_windows, train_labels,
                      velocity_metric_handle = None,
-                     velocity_mapping_handle: str | None | Callable[[int], int] = None):
+                     velocity_mapping_handle = None):
         """Adds velocity (i.e., proportional) control where a multiplier is generated for the level of contraction intensity.
 
         Note, that when using this optional, ramp contractions should be captured for training. 
 
         Parameters
         -----------
-        train_windows: np.ndarray
-            The training windows extracted from the offline data handler.
-        train_labels: np.ndarray
-            The labels associated with the train windows. Allows for per class proportional control mapping. 
-        velocity_mapping_handle: function or a string (valid options: "SIGMOID", "SQUARED", "LOG", "RELU") 
-            A function that maps the proportionality (bounded between 0-1) to some other value. 
         """
-        if isinstance(velocity_mapping_handle, str):
-            if not velocity_mapping_handle in ['SIGMOID', 'SQUARED', 'LOG', 'RELU']:
-                print("Invalid velocity mapping... Defaulting to linear.")
-            else:
-                if velocity_mapping_handle == 'SQUARED':
-                    self.velocity_mapping_handle = lambda x: x**2
-                # TODO: Fill out rest 
-        else:
-            self.velocity_metric_handle = velocity_metric_handle
-            self.velocity_mapping_handle = velocity_mapping_handle
-        
+        self.velocity_metric_handle = velocity_metric_handle
+        self.velocity_mapping_handle = velocity_mapping_handle
         self.velocity = True
+
         self.th_min_dic, self.th_max_dic = self._set_up_velocity_control(train_windows, train_labels)
 
+
 
     '''
     ---------------------- Private Helper Functions ----------------------
@@ -634,7 +620,10 @@ def __init__(self,
             ["adapt_flag", (1,1), np.int32],
             ["active_flag", (1,1), np.int8]
         ]
-        smm_items.extend(required_smm_items)
+        current_smm_tags = [item[0] for item in smm_items]
+        for smm_item in required_smm_items:
+            if smm_item[0] not in current_smm_tags:
+                smm_items.append(smm_item)
         self.smm = smm
         self.smm_items = smm_items
 
@@ -763,7 +752,7 @@ def _run_helper(self):
                     model_input = None
                     for mod in self.odh.modalities:
                         # todo: features for each modality can be different
-                        mod_features = fe.extract_features(self.features, window[mod], array=True)
+                        mod_features = fe.extract_features(self.features, window[mod], feature_dic=self.predictor.feature_params, array=True)
                         if model_input is None:
                             model_input = mod_features
                         else:
@@ -963,13 +952,8 @@ def insert_classifier_output(data):
                                                 insert_classifier_output)
             self.options['model_smm_writes'] += 1
 
-        if self.output_format == "predictions":
-            message = str(prediction) + calculated_velocity + '\n'
-        elif self.output_format == "probabilities":
-            message = ' '.join([f'{i:.2f}' for i in probabilities[0]]) + calculated_velocity + " " + str(time_stamp)
-        else:
-            raise ValueError(f"Unexpected value for output_format. Accepted values are 'predictions' and 'probabilities'. Got: {self.output_format}.")
-
+        message = str(prediction) + " " + str(np.abs(np.array(window['emg'])).mean(axis=2).mean()) + str(calculated_velocity)
+    
         if not self.tcp:
             self.sock.sendto(bytes(message, 'utf-8'), (self.ip, self.port))
         else:
@@ -997,7 +981,6 @@ def visualize(self, max_len=50, legend=None):
         cmap = cm.get_cmap('turbo', num_classes)
 
         controller = ClassifierController(output_format=self.output_format, num_classes=num_classes, ip=self.ip, port=self.port)
-        controller.start()
 
         if legend is not None:
             for i in range(num_classes):
@@ -1184,7 +1167,6 @@ def visualize(self, max_len = 50, legend = False):
         ax.set_ylabel('Prediction')
 
         controller = RegressorController(ip=self.ip, port=self.port)
-        controller.start()
 
         # Wait for controller to start receiving data
         predictions = None