Add heuristic compression of action series on key/scroll/mouse down&up

utils/ducktrack.py

@@ -1,9 +1,13 @@
-import sys, pathlib; sys.path.append(str(pathlib.Path(__file__).parents[1]))
+import sys, pathlib;
+
+sys.path.append(str(pathlib.Path(__file__).parents[1]))
 
 import os
 import json
+from typing import List
 from desktop_env.envs.desktop_env import Action, MouseClick
 
+
 class DuckTrackEventActionConverter:
     def __init__(self, human_readable: str, compress_move: bool = True):
         self.human_readable = human_readable
@@ -13,7 +17,7 @@ class DuckTrackEventActionConverter:
         """Converts an enum to its string representation if HUMAN_READABLE is True, otherwise returns its value."""
         return enum.name if self.human_readable else enum.value
 
-    def compress_mouse_move(self, data: list[dict], index: int):
+    def compress_mouse_move(self, data: List[dict], index: int):
         """Compresses consecutive mouse move events into first and last move events."""
         first_move, last_move = data[index], data[index]
         while index < len(data) and data[index]["action"] == "move":
@@ -31,9 +35,9 @@ class DuckTrackEventActionConverter:
         mouse_button = event["button"]
         mouse_pressed = event["pressed"]
 
-        if mouse_pressed == True:
+        if mouse_pressed:
             action["action_type"] = self.enum_to_str(Action.MOUSE_DOWN)
-        elif mouse_pressed == False:
+        elif not mouse_pressed:
             action["action_type"] = self.enum_to_str(Action.MOUSE_UP)
         else:
             raise NotImplementedError(mouse_pressed)
@@ -65,7 +69,8 @@ class DuckTrackEventActionConverter:
             down = True
 
         return {"action_type": self.enum_to_str(Action.CLICK),
-                "click_type": self.enum_to_str(MouseClick.WHEEL_DOWN) if down else self.enum_to_str(MouseClick.WHEEL_UP)}
+                "click_type": self.enum_to_str(MouseClick.WHEEL_DOWN) if down else self.enum_to_str(
+                    MouseClick.WHEEL_UP)}
 
     def event_to_action(self, event: dict):
         """Converts an event to its corresponding action based on the event type."""
@@ -88,7 +93,7 @@ class DuckTrackEventActionConverter:
             raise FileNotFoundError(ducktrack_event_file)
 
         # set to default
-        if compress_move == None:
+        if compress_move is None:
             compress_move = self.compress_move
 
         with open(ducktrack_event_file, 'r') as file:
@@ -97,22 +102,99 @@ class DuckTrackEventActionConverter:
         result = {"action": [], "event": []}
         index = 0
 
+        # Compress the mouse move events
         while index < len(data):
             event = data[index]
             if event["action"] == "move" and compress_move:
                 first_move, last_move, index = self.compress_mouse_move(data, index)
-                result["action"].extend([self.event_to_action(first_move), self.event_to_action(last_move)])
+                result["action"].extend([self.event_to_action(last_move)])
-                result["event"].extend([first_move, last_move])
+                result["event"].extend([last_move])
             else:
                 result["action"].append(self.event_to_action(event))
                 result["event"].append(event)
                 index += 1
 
+        # Compress the key down and key up actions
+        # todo: handling the key down and key up events
+        _new_actions = []
+        _action = list(result["action"])
+        idx = 0
+
+        while True:
+            if idx >= len(_action):
+                break
+
+            if _action[idx]["action_type"] == self.enum_to_str(Action.KEY_DOWN):
+                typed_text = []
+                while idx < len(_action) and _action[idx]["action_type"] in [self.enum_to_str(Action.KEY_DOWN), self.enum_to_str(Action.KEY_UP)] and len(_action[idx]["key"]) == 1:
+                    if _action[idx]["action_type"] == self.enum_to_str(Action.KEY_DOWN):
+                        typed_text.append(chr(_action[idx]["key"][0]))
+                    idx += 1
+                if typed_text:
+                    _new_actions.append({"action_type": self.enum_to_str(Action.TYPE), "text": typed_text})
+                else:
+                    _new_actions.append(_action[idx])
+                    idx += 1
+            else:
+                _new_actions.append(_action[idx])
+                idx += 1
+
+        result["action"] = _new_actions
+
+        # Compress the scroll up and scroll down events
+        # todo: handling the key down and key up events
+        _new_actions = []
+        _action = list(result["action"])
+        idx = 0
+
+        while True:
+            if idx >= len(_action):
+                break
+
+            if _action[idx]["action_type"] == self.enum_to_str(Action.CLICK) and _action[idx]["click_type"] in [self.enum_to_str(MouseClick.WHEEL_UP), self.enum_to_str(MouseClick.WHEEL_DOWN)]:
+                typed_text = []
+                while idx < len(_action) and _action[idx]["action_type"] == self.enum_to_str(Action.CLICK) and _action[idx]["click_type"] in [self.enum_to_str(MouseClick.WHEEL_UP), self.enum_to_str(MouseClick.WHEEL_DOWN)]:
+                    if _action[idx]["click_type"] == self.enum_to_str(MouseClick.WHEEL_UP):
+                        typed_text.append("UP")
+                        idx += 1
+                    elif _action[idx]["click_type"] == self.enum_to_str(MouseClick.WHEEL_DOWN):
+                        typed_text.append("DOWN")
+                        idx += 1
+                _new_actions.append({"action_type": self.enum_to_str(Action.CLICK), "click_type": "SCROLL", "text": typed_text})
+            else:
+                _new_actions.append(_action[idx])
+                idx += 1
+
+        result["action"] = _new_actions
+
+        # Compress the mouse down and mouse up actions
+        # todo: handling the key down and key up events
+        _new_actions = []
+        _action = list(result["action"])
+        idx = 0
+
+        while True:
+            if idx >= len(_action):
+                break
+            if _action[idx]["action_type"] == self.enum_to_str(Action.MOUSE_DOWN):
+                if idx + 1 < len(_action) and _action[idx+1]["action_type"] == self.enum_to_str(Action.MOUSE_UP):
+                    _new_actions.append({"action_type": self.enum_to_str(Action.CLICK), "click_type": _action[idx]["click_type"]})
+                    idx += 2
+                else:
+                    _new_actions.append(_action[idx])
+                    idx += 1
+            else:
+                _new_actions.append(_action[idx])
+                idx += 1
+
+        result["action"] = _new_actions
+
         with open(out_file, "w") as f:
             json.dump(result, f)
 
+
 if __name__ == "__main__":
-    converter = DuckTrackEventActionConverter(human_readable=False)
+    converter = DuckTrackEventActionConverter(human_readable=True)
     converter.ducktrack_event_file_to_action(ducktrack_event_file="sample.jsonl",
                                              out_file="output.json",
                                              compress_move=True)