118 lines
4.8 KiB
Python
118 lines
4.8 KiB
Python
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) |