Merge remote-tracking branch 'origin/main'

utils/ducktrack.py

@@ -0,0 +1,118 @@
+import sys, pathlib; sys.path.append(str(pathlib.Path(__file__).parents[1]))
+
+import os
+import json
+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
+        self.compress_move = compress_move
+
+    def enum_to_str(self, enum):
+        """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):
+        """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":
+            last_move = data[index]
+            index += 1
+        return first_move, last_move, index
+
+    def move_event_to_action(self, event: dict):
+        return {"action_type": self.enum_to_str(Action.MOUSE_MOVE),
+                "x": event["x"],
+                "y": event["y"]}
+
+    def click_event_to_action(self, event: dict):
+        action = {}
+        mouse_button = event["button"]
+        mouse_pressed = event["pressed"]
+
+        if mouse_pressed == True:
+            action["action_type"] = self.enum_to_str(Action.MOUSE_DOWN)
+        elif mouse_pressed == False:
+            action["action_type"] = self.enum_to_str(Action.MOUSE_UP)
+        else:
+            raise NotImplementedError(mouse_pressed)
+
+        if mouse_button == "left":
+            action["click_type"] = self.enum_to_str(MouseClick.LEFT)
+        elif mouse_button == "right":
+            action["click_type"] = self.enum_to_str(MouseClick.RIGHT)
+        elif mouse_button == "middle":
+            action["click_type"] = self.enum_to_str(MouseClick.MIDDLE)
+        else:
+            raise NotImplementedError(mouse_button)
+
+        return action
+
+    def press_event_to_action(self, event: dict):
+        return {"action_type": self.enum_to_str(Action.KEY_DOWN),
+                "key": [ord(c) for c in event["name"]]}
+
+    def release_event_to_action(self, event: dict):
+        return {"action_type": self.enum_to_str(Action.KEY_UP),
+                "key": [ord(c) for c in event["name"]]}
+
+    def scroll_event_to_action(self, event: dict):
+        # TODO: need to confirm if df < 0 means scroll up or down
+        if event["dy"] < 0:
+            down = False
+        else:
+            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)}
+
+    def event_to_action(self, event: dict):
+        """Converts an event to its corresponding action based on the event type."""
+        if event["action"] == "move":
+            return self.move_event_to_action(event)
+        elif event["action"] == "click":
+            return self.click_event_to_action(event)
+        elif event["action"] == "press":
+            return self.press_event_to_action(event)
+        elif event["action"] == "release":
+            return self.release_event_to_action(event)
+        elif event["action"] == "scroll":
+            return self.scroll_event_to_action(event)
+        else:
+            raise NotImplementedError(event["action"])
+
+    def ducktrack_event_file_to_action(self, ducktrack_event_file: str, out_file: str, compress_move: bool = None):
+        """Converts DuckTrack event data to a list of actions and saves them to a file."""
+        if not os.path.exists(ducktrack_event_file):
+            raise FileNotFoundError(ducktrack_event_file)
+        
+        # set to default
+        if compress_move == None:
+            compress_move = self.compress_move
+        
+        with open(ducktrack_event_file, 'r') as file:
+            data = [json.loads(line) for line in file]
+        
+        result = {"action": [], "event": []}
+        index = 0
+        
+        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["event"].extend([first_move, last_move])
+            else:
+                result["action"].append(self.event_to_action(event))
+                result["event"].append(event)
+                index += 1
+        
+        with open(out_file, "w") as f:
+            json.dump(result, f)
+
+if __name__ == "__main__":
+    converter = DuckTrackEventActionConverter(human_readable=False)
+    converter.ducktrack_event_file_to_action(ducktrack_event_file="sample.jsonl",
+                                             out_file="output.json",
+                                             compress_move=True)