203 lines
8.2 KiB
Python
203 lines
8.2 KiB
Python
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
|
|
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:
|
|
action["action_type"] = self.enum_to_str(Action.MOUSE_DOWN)
|
|
elif not mouse_pressed:
|
|
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
|
|
|
|
# TODO: NEED to be test to match the scroll up and down with our action, e.g. scroll here once is equal to scroll 10 or scroll 20?
|
|
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 is 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
|
|
|
|
# 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(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=True)
|
|
converter.ducktrack_event_file_to_action(ducktrack_event_file="sample.jsonl",
|
|
out_file="output.json",
|
|
compress_move=True)
|