wehub-resource-sync
272 lines
11 KiB
Python
272 lines
11 KiB
Python
from gui_agents.s3.core.mllm import LMMAgent
|
|
from gui_agents.s3.memory.procedural_memory import PROCEDURAL_MEMORY
|
|
from gui_agents.s3.utils.common_utils import (
|
|
call_llm_formatted,
|
|
split_thinking_response,
|
|
compress_image,
|
|
)
|
|
from gui_agents.s3.utils.formatters import (
|
|
THOUGHTS_ANSWER_TAG_FORMATTER,
|
|
)
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from io import BytesIO
|
|
from typing import Dict
|
|
import base64
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class BehaviorNarrator:
|
|
def __init__(self, engine_params):
|
|
self.judge_agent = LMMAgent(engine_params=engine_params)
|
|
|
|
@staticmethod
|
|
def extract_mouse_action(action: str) -> list[str]:
|
|
mouse_actions = []
|
|
for sub_action in action.split(";"):
|
|
sub_action = sub_action.strip()
|
|
if (
|
|
sub_action.startswith("pyautogui.click")
|
|
or sub_action.startswith("pyautogui.moveTo")
|
|
or sub_action.startswith("pyautogui.dragTo")
|
|
):
|
|
mouse_actions.append(sub_action)
|
|
return mouse_actions
|
|
|
|
@staticmethod
|
|
def mark_action(mouse_actions: list[str], img: Image):
|
|
draw = ImageDraw.Draw(img)
|
|
font = ImageFont.load_default(25)
|
|
|
|
drag_start_width, drag_start_height = None, None
|
|
|
|
for mouse_action in mouse_actions:
|
|
width, height = mouse_action.split("(")[1].strip(")").split(", ")[:2]
|
|
width, height = int(width), int(height)
|
|
|
|
# Clamp coordinates within bounds
|
|
width = max(0, min(img.width - 1, width))
|
|
height = max(0, min(img.height - 1, height))
|
|
|
|
def place_text(label, color, x, y):
|
|
bbox = draw.textbbox((0, 0), label, font=font)
|
|
text_w, text_h = (
|
|
bbox[2] - bbox[0],
|
|
bbox[3] - bbox[1],
|
|
) # Measure text size
|
|
offset_x, offset_y = -5, 5 # Default offset
|
|
if x + offset_x + text_w > img.width: # Out of bounds on right
|
|
offset_x = -text_w - 5
|
|
if y + offset_y + text_h > img.height: # Out of bounds on bottom
|
|
offset_y = -text_h - 5
|
|
if x + offset_x < 0: # Out of bounds on left
|
|
offset_x = 5
|
|
if y + offset_y < 0: # Out of bounds on top
|
|
offset_y = 5
|
|
draw.text((x + offset_x, y + offset_y), label, fill=color, font=font)
|
|
|
|
if mouse_action.startswith("pyautogui.click"):
|
|
draw.circle((width, height), radius=3, fill=(255, 0, 0))
|
|
place_text("Click", (255, 0, 0), width, height)
|
|
if mouse_action.startswith("pyautogui.moveTo"):
|
|
draw.circle((width, height), radius=3, fill=(0, 0, 255))
|
|
place_text("MoveTo", (0, 0, 255), width, height)
|
|
drag_start_height, drag_start_width = height, width
|
|
if mouse_action.startswith("pyautogui.dragTo"):
|
|
draw.line(
|
|
[(drag_start_width, drag_start_height), (width, height)],
|
|
fill=(0, 255, 0),
|
|
width=2,
|
|
)
|
|
draw.circle((width, height), radius=3, fill=(0, 255, 0))
|
|
place_text("DragTo", (0, 255, 0), width, height)
|
|
|
|
@staticmethod
|
|
def get_mouse_action_representation(mouse_actions: list[str]) -> str:
|
|
"""
|
|
Returns a string representation of the mouse action for the given action.
|
|
"""
|
|
assert (
|
|
len(mouse_actions) <= 2
|
|
), f"Multiple mouse action types found: {mouse_actions}"
|
|
if len(mouse_actions) == 1:
|
|
action = mouse_actions[0]
|
|
if action.startswith("pyautogui.click"):
|
|
return "The red circle labeled 'Click' marks the position where the mouse was clicked."
|
|
elif action.startswith("pyautogui.moveTo"):
|
|
return "The blue circle labeled 'MoveTo' marks the position where the mouse was moved to."
|
|
else:
|
|
raise ValueError(f"Unknown single action type: {action}")
|
|
else:
|
|
assert mouse_actions[0].startswith("pyautogui.moveTo") and mouse_actions[
|
|
1
|
|
].startswith("pyautogui.dragTo")
|
|
return "The blue circle labeled 'MoveTo' marks the starting position of the mouse.\nThe green circle labeled 'DragTo' marks the ending position.\nThe green line illustrates the mouse's drag path."
|
|
|
|
@staticmethod
|
|
def get_zoomed_image(
|
|
image_bytes: bytes,
|
|
x: int,
|
|
y: int,
|
|
width: int = 300,
|
|
height: int = 300,
|
|
upscaling: bool = False,
|
|
scale: int = 4,
|
|
add_bounding_box: bool = False,
|
|
) -> bytes:
|
|
"""Returns a zoomed image centered around (x, y) coordinates.
|
|
|
|
Args:
|
|
image_bytes (bytes): The original image in bytes.
|
|
x (int): The x-coordinate of the center point.
|
|
y (int): The y-coordinate of the center point.
|
|
width (int): The width of the zoomed area.
|
|
height (int): The height of the zoomed area.
|
|
padding (int): Extra padding around the zoomed area.
|
|
upscaling (bool): Whether to upscale and enhance the zoomed image.
|
|
scale (int): The upscaling factor if upscaling is True.
|
|
add_bounding_box (bool): Whether to add a bounding box around the zoomed area in the original image.
|
|
|
|
Returns:
|
|
bytes: The zoomed image in bytes.
|
|
bytes: The original image with bounding box in bytes (if add_bounding_box is True). Otherwise, returns original bytes.
|
|
"""
|
|
# Find zoom dimensions
|
|
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
|
cx, cy = x - width // 2, y - height // 2 # Center coordinates
|
|
W, H = img.size
|
|
left = min(max(cx, 0), W - width)
|
|
top = min(max(cy, 0), H - height)
|
|
right = left + width
|
|
bottom = top + height
|
|
zoomed_img = img.crop((left, top, right, bottom))
|
|
# Add noticeable bounding box to original image
|
|
if add_bounding_box:
|
|
draw_img = img.copy()
|
|
draw = ImageDraw.Draw(draw_img)
|
|
draw.rectangle([left, top, right, bottom], outline="red", width=3)
|
|
original_with_box_bytes = compress_image(
|
|
image=draw_img
|
|
) # Compress to reduce size
|
|
else:
|
|
original_with_box_bytes = image_bytes
|
|
if upscaling:
|
|
# Upscale and enhance zoomed image
|
|
zoomed_img = cv2.cvtColor(
|
|
np.array(zoomed_img), cv2.COLOR_RGB2BGR
|
|
) # PIL -> OpenCV
|
|
zoomed_img = cv2.resize(
|
|
zoomed_img, None, fx=scale, fy=scale, interpolation=cv2.INTER_LANCZOS4
|
|
)
|
|
zoomed_img = cv2.fastNlMeansDenoisingColored(
|
|
zoomed_img, None, 5, 5, 7, 21
|
|
) # light denoise (helps with JPEG speckle)
|
|
zoomed_img = Image.fromarray(
|
|
cv2.cvtColor(zoomed_img, cv2.COLOR_BGR2RGB)
|
|
) # OpenCV -> PIL
|
|
zoomed_img_bytes = compress_image(image=zoomed_img) # Compress to reduce size
|
|
return zoomed_img_bytes, original_with_box_bytes
|
|
|
|
def judge(
|
|
self,
|
|
screenshot_num: int,
|
|
before_img_bytes: bytes,
|
|
after_img_bytes: bytes,
|
|
pyautogui_action: str,
|
|
) -> Dict[str, str]:
|
|
if pyautogui_action == "DONE":
|
|
return {
|
|
"fact_thoughts": "The agent has indicated that it is done with the task.",
|
|
"fact_answer": "The agent has indicated that it is done with the task.",
|
|
}
|
|
elif pyautogui_action == "FAIL":
|
|
return {
|
|
"fact_thoughts": "The agent has indicated that it is impossible to proceed further with the task.",
|
|
"fact_answer": "The agent has indicated that it is impossible to proceed further with the task.",
|
|
}
|
|
# Prepare ANNOTATED BEFORE image
|
|
mouse_actions = BehaviorNarrator.extract_mouse_action(pyautogui_action)
|
|
before_img = Image.open(BytesIO(before_img_bytes))
|
|
BehaviorNarrator.mark_action(mouse_actions, before_img)
|
|
out_buffer = BytesIO()
|
|
before_img.save(out_buffer, format="PNG")
|
|
marked_before_img_bytes = out_buffer.getvalue()
|
|
marked_before_img_message = {
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{base64.b64encode(marked_before_img_bytes).decode('utf-8')}",
|
|
"detail": "high",
|
|
},
|
|
}
|
|
if mouse_actions:
|
|
coords = mouse_actions[-1].split("(")[1].strip(")").split(", ")
|
|
x, y = int(coords[0]), int(coords[1])
|
|
zoomed_after_img_bytes, marked_after_img_bytes = (
|
|
BehaviorNarrator.get_zoomed_image(
|
|
image_bytes=after_img_bytes,
|
|
x=x,
|
|
y=y,
|
|
width=300,
|
|
height=300,
|
|
scale=4,
|
|
upscaling=True,
|
|
add_bounding_box=True,
|
|
)
|
|
)
|
|
after_img_message = {
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{base64.b64encode(marked_after_img_bytes).decode('utf-8')}",
|
|
"detail": "high",
|
|
},
|
|
}
|
|
zoomed_after_img_message = {
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{base64.b64encode(zoomed_after_img_bytes).decode('utf-8')}",
|
|
"detail": "high",
|
|
},
|
|
}
|
|
else:
|
|
after_img_message = {
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{base64.b64encode(after_img_bytes).decode('utf-8')}",
|
|
"detail": "high",
|
|
},
|
|
}
|
|
zoomed_after_img_message = None
|
|
|
|
fact_message = [
|
|
{
|
|
"role": "system",
|
|
"content": PROCEDURAL_MEMORY.BEHAVIOR_NARRATOR_SYSTEM_PROMPT,
|
|
}
|
|
]
|
|
fact_message_content = [
|
|
{"type": "text", "text": "BEFORE:"},
|
|
marked_before_img_message,
|
|
{"type": "text", "text": f"Agent Action: {pyautogui_action}"},
|
|
{"type": "text", "text": "AFTER:"},
|
|
after_img_message,
|
|
]
|
|
if zoomed_after_img_message:
|
|
fact_message_content += [
|
|
{"type": "text", "text": "ZOOMED AFTER:"},
|
|
zoomed_after_img_message,
|
|
]
|
|
fact_message += [{"role": "user", "content": fact_message_content}]
|
|
fact_response = call_llm_formatted(
|
|
self.judge_agent,
|
|
[THOUGHTS_ANSWER_TAG_FORMATTER],
|
|
messages=fact_message,
|
|
temperature=0.0,
|
|
)
|
|
fact_answer, fact_thoughts = split_thinking_response(fact_response)
|
|
|
|
result = {
|
|
"fact_thoughts": fact_thoughts,
|
|
"fact_answer": f"Fact Caption from Screenshot {screenshot_num}: {fact_answer}",
|
|
}
|
|
return result
|