.

main.py

@@ -118,6 +118,17 @@ KEYS = ['\t', '\n', '\r', ' ', '!', '"', '#', '$', '%', '&', "'", '(',
         'shift', 'shiftleft', 'shiftright', 'sleep', 'space', 'stop', 'subtract', 'tab',
         'up', 'volumedown', 'volumemute', 'volumeup', 'win', 'winleft', 'winright', 'yen',
         'command', 'option', 'optionleft', 'optionright']
+
+KEYMAPPING = {
+    'arrowup': 'up',
+    'arrowdown': 'down',
+    'arrowleft': 'left',
+    'arrowright': 'right',
+    'meta': 'command',
+    'contextmenu': 'apps',
+    'control': 'ctrl',
+}
+
 INVALID_KEYS = ['f13', 'f14', 'f15', 'f16', 'f17', 'f18', 'f19', 'f20',
                 'f21', 'f22', 'f23', 'f24', 'select', 'separator', 'execute']
 VALID_KEYS = [key for key in KEYS if key not in INVALID_KEYS]
@@ -174,6 +185,8 @@ def prepare_model_inputs(
     }
     for key in keys_down:
         key = key.lower()
+        if key in KEYMAPPING:
+            key = KEYMAPPING[key]
         if key in stoi:
             inputs['key_events'][stoi[key]] = 1
         else:

online_data_generation.py

@@ -56,6 +56,45 @@ autoencoder = autoencoder.to(device)
 # Global flag for graceful shutdown
 running = True
 
+
+KEYMAPPING = {
+    'arrowup': 'up',
+    'arrowdown': 'down',
+    'arrowleft': 'left',
+    'arrowright': 'right',
+    'meta': 'command',
+    'contextmenu': 'apps',
+    'control': 'ctrl',
+}
+
+KEYS = ['\t', '\n', '\r', ' ', '!', '"', '#', '$', '%', '&', "'", '(',
+        ')', '*', '+', ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7',
+        '8', '9', ':', ';', '<', '=', '>', '?', '@', '[', '\\', ']', '^', '_', '`',
+        'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
+        'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '{', '|', '}', '~',
+        'accept', 'add', 'alt', 'altleft', 'altright', 'apps', 'backspace',
+        'browserback', 'browserfavorites', 'browserforward', 'browserhome',
+        'browserrefresh', 'browsersearch', 'browserstop', 'capslock', 'clear',
+        'convert', 'ctrl', 'ctrlleft', 'ctrlright', 'decimal', 'del', 'delete',
+        'divide', 'down', 'end', 'enter', 'esc', 'escape', 'execute', 'f1', 'f10',
+        'f11', 'f12', 'f13', 'f14', 'f15', 'f16', 'f17', 'f18', 'f19', 'f2', 'f20',
+        'f21', 'f22', 'f23', 'f24', 'f3', 'f4', 'f5', 'f6', 'f7', 'f8', 'f9',
+        'final', 'fn', 'hanguel', 'hangul', 'hanja', 'help', 'home', 'insert', 'junja',
+        'kana', 'kanji', 'launchapp1', 'launchapp2', 'launchmail',
+        'launchmediaselect', 'left', 'modechange', 'multiply', 'nexttrack',
+        'nonconvert', 'num0', 'num1', 'num2', 'num3', 'num4', 'num5', 'num6',
+        'num7', 'num8', 'num9', 'numlock', 'pagedown', 'pageup', 'pause', 'pgdn',
+        'pgup', 'playpause', 'prevtrack', 'print', 'printscreen', 'prntscrn',
+        'prtsc', 'prtscr', 'return', 'right', 'scrolllock', 'select', 'separator',
+        'shift', 'shiftleft', 'shiftright', 'sleep', 'space', 'stop', 'subtract', 'tab',
+        'up', 'volumedown', 'volumemute', 'volumeup', 'win', 'winleft', 'winright', 'yen',
+        'command', 'option', 'optionleft', 'optionright']
+INVALID_KEYS = ['f13', 'f14', 'f15', 'f16', 'f17', 'f18', 'f19', 'f20',
+                'f21', 'f22', 'f23', 'f24', 'select', 'separator', 'execute']
+VALID_KEYS = [key for key in KEYS if key not in INVALID_KEYS]
+itos = VALID_KEYS
+stoi = {key: i for i, key in enumerate(itos)}
+
 def signal_handler(sig, frame):
     """Handle Ctrl+C and other termination signals"""
     global running
@@ -459,10 +498,27 @@ def format_trajectory_for_processing(trajectory):
         # Extract input data
         inputs = entry.get("inputs", {})
         key_events = []
+        down_keys = set([])
         for key in inputs.get("keys_down", []):
-            key_events.append(("keydown", key))
+            key = key.lower()
+            if key in KEYMAPPING:
+                print (f"Key {key} mapped to {KEYMAPPING[key]}")
+                key = KEYMAPPING[key]
+            if key not in stoi:
+                print (f"Key {key} not found in stoi")
+            if key not in down_keys and key in stoi:
+                down_keys.add(key)
+                key_events.append(("keydown", key))
         for key in inputs.get("keys_up", []):
-            key_events.append(("keyup", key))
+            key = key.lower()
+            if key in KEYMAPPING:
+                print (f"Key {key} mapped to {KEYMAPPING[key]}")
+                key = KEYMAPPING[key]
+            if key not in stoi:
+                print (f"Key {key} not found in stoi")
+            if key in down_keys and key in stoi:
+                down_keys.remove(key)
+                key_events.append(("keyup", key))
         event = {
             "pos": (inputs.get("x"), inputs.get("y")),
             "left_click": inputs.get("is_left_click", False),