277 lines
11 KiB
Python
277 lines
11 KiB
Python
import argparse
|
|
import json
|
|
import math
|
|
import re
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
from datasets import load_dataset
|
|
from PIL.Image import Image
|
|
from PIL.Image import open as open_img
|
|
from tqdm import tqdm
|
|
from transformers import AutoModelForImageTextToText, AutoProcessor
|
|
from transformers.modeling_utils import PreTrainedModel
|
|
from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
|
|
from transformers.processing_utils import ProcessorMixin
|
|
|
|
INSTRUCTION_LOCALIZATION: str = "Localize an element on the GUI image according to my instructions and output a click position as Click(x, y) with x num pixels from the left edge and y num pixels from the top edge."
|
|
INSTRUCTION_LOCALIZATION_TOOLCALL: str = "Localize an element on the GUI image according to my instructions and output a click position. You must output a valid JSON format."
|
|
|
|
|
|
def load_screenspot(dataset_id: str, subset: str = "test"):
|
|
dataset = load_dataset(dataset_id)
|
|
return dataset[subset]
|
|
|
|
|
|
def l1(dx: float, dy: float) -> float:
|
|
"""Return L1 length of a vector"""
|
|
return abs(dx) + abs(dy)
|
|
|
|
|
|
def l2(dx: float, dy: float) -> float:
|
|
"""Return L2 length of a vector"""
|
|
return (dx**2 + dy**2) ** 0.5
|
|
|
|
|
|
def point_to_rectangle_dist(x: float, y: float, rectangle: tuple, distance_type="L2"):
|
|
"""Compute the distance of a predicted point to the closest edge of the bbox. If the point is in the bbox, then return 0."""
|
|
x1, y1, x2, y2 = rectangle # x1,y1 is top-left, x2,y2 is bottom-right
|
|
|
|
# Check if the point is inside the rectangle
|
|
if x1 <= x <= x2 and y1 <= y <= y2:
|
|
return 0
|
|
|
|
# Calculate the closest point on the rectangle
|
|
closest_x = max(x1, min(x, x2))
|
|
closest_y = max(y1, min(y, y2))
|
|
|
|
# Calculate the distance
|
|
dx = x - closest_x
|
|
dy = y - closest_y
|
|
if distance_type == "L1":
|
|
return l1(dx, dy)
|
|
elif distance_type == "L2":
|
|
return l2(dx, dy)
|
|
else:
|
|
raise ValueError("Invalid distance type. Use 'L1' or 'L2'.")
|
|
|
|
|
|
def is_in_bbox(bbox: tuple, x: float, y: float) -> bool:
|
|
"""Check if a point is inside a bounding box."""
|
|
x_top_left, y_top_left, x_bottom_right, y_bottom_right = bbox
|
|
return x_top_left <= x <= x_bottom_right and y_top_left <= y <= y_bottom_right
|
|
|
|
|
|
def assemble_message(image, instruction, use_tool_call: bool = True):
|
|
system_message = {
|
|
"role": "system",
|
|
"content": '[{"name": "click_action", "description": "Click at specific coordinates on the screen.", "parameters": {"additionalProperties": false, "description": "Click at specific coordinates on the screen.", "properties": {"action": {"const": "click", "default": "click", "title": "Action", "type": "string"}, "x": {"description": "The x coordinate, number of pixels from the left edge.", "title": "X", "type": "integer"}, "y": {"description": "The y coordinate, number of pixels from the top edge.", "title": "Y", "type": "integer"}}, "required": ["action", "x", "y"], "title": "ClickAction", "type": "object"}, "strict": true}]',
|
|
}
|
|
|
|
user_message = {
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image",
|
|
"image": image,
|
|
},
|
|
{
|
|
"type": "text",
|
|
"text": f"{INSTRUCTION_LOCALIZATION_TOOLCALL if use_tool_call else INSTRUCTION_LOCALIZATION}\n{instruction}",
|
|
},
|
|
],
|
|
}
|
|
|
|
messages = [system_message, user_message] if use_tool_call else [user_message]
|
|
return messages
|
|
|
|
|
|
def do_smart_resize(image: Image, image_processor: ProcessorMixin) -> tuple[Image, int, int]:
|
|
"""Do a QWEN2.5-VL smart resize using parameters of an image-processor"""
|
|
resized_height, resized_width = smart_resize(
|
|
image.height,
|
|
image.width,
|
|
factor=image_processor.patch_size * image_processor.merge_size,
|
|
min_pixels=image_processor.min_pixels,
|
|
max_pixels=image_processor.max_pixels,
|
|
)
|
|
return image.resize(size=(resized_width, resized_height), resample=None), resized_height, resized_width
|
|
|
|
|
|
def inference(
|
|
model: PreTrainedModel, processor: ProcessorMixin, dataset, smart_resize: bool = True, use_toolcall: bool = True
|
|
):
|
|
"""Gather raw inference results from the model"""
|
|
results = []
|
|
for i, sample in enumerate(tqdm(dataset, "running inference requests")):
|
|
bbox = sample["bbox"]
|
|
instruction = sample["instruction"]
|
|
image = sample["image"] # this seems to be a pnd , maybe jpg artifacts cause the difference?
|
|
image_shape_raw = (image.height, image.width)
|
|
message = assemble_message(image=image, instruction=instruction)
|
|
|
|
# Preparation for inference
|
|
if smart_resize:
|
|
image, resized_height, resized_width = do_smart_resize(
|
|
image=image, image_processor=processor.image_processor
|
|
)
|
|
else:
|
|
resized_height, resized_width = image_shape_raw
|
|
text = processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
|
|
|
|
# compress to JPEG, which is needed for highest possible performance
|
|
buffer = BytesIO()
|
|
image.convert("RGB").save(buffer, format="JPEG", quality=90)
|
|
image = open_img(buffer)
|
|
|
|
inputs = processor(
|
|
text=[text],
|
|
images=image,
|
|
padding=True,
|
|
return_tensors="pt",
|
|
)
|
|
inputs = inputs.to("cuda")
|
|
|
|
# Inference: Generation of the output
|
|
generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False)
|
|
generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
|
|
output_text = processor.batch_decode(
|
|
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
|
|
)
|
|
# print(output_text)
|
|
if use_toolcall:
|
|
try:
|
|
content = json.loads(output_text[0])
|
|
prediction_raw = f"Click({content['arguments']['x']}, {content['arguments']['y']})"
|
|
except Exception as e:
|
|
print(f"Error parsing tool call, using message content instead if available: {repr(e)}")
|
|
prediction_raw = output_text[0]
|
|
else:
|
|
prediction_raw = output_text[0]
|
|
|
|
results.append(
|
|
{
|
|
"sample_id": i,
|
|
"ground_truth": tuple(bbox),
|
|
"prediction_raw": prediction_raw,
|
|
"image_shape_raw": image_shape_raw,
|
|
"img_shape_processed": (resized_height, resized_width),
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
def get_sample_result(result: dict):
|
|
"""Postprocess a inference result and compute metrics for this sample."""
|
|
raw_height, raw_width = result["image_shape_raw"]
|
|
height, width = result["img_shape_processed"]
|
|
has_resized_image = height != raw_height or width != raw_width
|
|
try:
|
|
bbox = result["ground_truth"]
|
|
prediction_raw = result["prediction_raw"]
|
|
match = re.match(r"Click\((\d+),\s*(\d+)\)", prediction_raw)
|
|
assert match is not None
|
|
predicted_x = float(match.group(1)) / width
|
|
predicted_y = float(match.group(2)) / height
|
|
|
|
except Exception as e:
|
|
sample_metric = {
|
|
"sample_id": result["sample_id"],
|
|
"has_correct_format": False,
|
|
"has_resized_image": has_resized_image,
|
|
"click_in_box": False,
|
|
"click_l1_dist_to_bbox": 2, # Longest possible L1 distance in the unit square
|
|
"click_l2_dist_to_bbox": math.sqrt(2), # Longest possible L2 distance in the unit square
|
|
}
|
|
|
|
sample_metric = {
|
|
"sample_id": result["sample_id"],
|
|
"has_correct_format": True,
|
|
"has_resized_image": has_resized_image,
|
|
"click_in_box": True if is_in_bbox(bbox, x=predicted_x, y=predicted_y) else False,
|
|
"click_l1_dist_to_bbox": point_to_rectangle_dist(
|
|
predicted_x, predicted_y, bbox, "L1"
|
|
), # Longest possible L1 distance in the unit square
|
|
"click_l2_dist_to_bbox": point_to_rectangle_dist(
|
|
predicted_x, predicted_y, bbox, "L2"
|
|
), # Longest possible L2 distance in the unit square
|
|
}
|
|
return sample_metric
|
|
|
|
|
|
def aggregate_metrics(sample_metrics):
|
|
"""Aggregate per-sample metrics into metrics for the entire dataset."""
|
|
aggregated_metrics = {}
|
|
aggregated_metrics["click_accuracy"] = np.mean([r["click_in_box"] for r in sample_metrics])
|
|
|
|
for threshold in [0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]:
|
|
aggregated_metrics[f"click_accuracy_p{threshold}"] = np.mean(
|
|
[r["click_l2_dist_to_bbox"] < threshold for r in sample_metrics]
|
|
)
|
|
|
|
aggregated_metrics["avg_click_l1_dist_to_bbox"] = np.mean([r["click_l1_dist_to_bbox"] for r in sample_metrics])
|
|
aggregated_metrics["avg_click_l2_dist_to_bbox"] = np.mean([r["click_l2_dist_to_bbox"] for r in sample_metrics])
|
|
aggregated_metrics["format_accuracy"] = np.mean([r["has_correct_format"] for r in sample_metrics])
|
|
aggregated_metrics["has_resized_image"] = np.mean([r["has_resized_image"] for r in sample_metrics])
|
|
return aggregated_metrics
|
|
|
|
|
|
def evaluate_results(results: list[dict]):
|
|
"""Do evaluate based on the raw model outputs."""
|
|
per_sample_metrics = []
|
|
for result in results:
|
|
metric_dict = get_sample_result(result)
|
|
per_sample_metrics.append(metric_dict)
|
|
aggregated = aggregate_metrics(per_sample_metrics)
|
|
return aggregated
|
|
|
|
|
|
def main(
|
|
model_id: str = "Hcompany/Holo1-3B",
|
|
dataset_id: str = "rootsautomation/ScreenSpot",
|
|
outfile: str = "results.json",
|
|
use_toolcall: bool = True,
|
|
):
|
|
model = AutoModelForImageTextToText.from_pretrained(model_id)
|
|
processor = AutoProcessor.from_pretrained(model_id)
|
|
dataset = load_screenspot(dataset_id)
|
|
results = inference(model.cuda(), processor, dataset, use_toolcall=use_toolcall)
|
|
metrics = evaluate_results(results)
|
|
with open(outfile, "w") as fp:
|
|
json.dump(metrics, fp)
|
|
for metric, value in metrics.items():
|
|
print(f"{metric}:\t{value}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run the main function with model and dataset IDs.")
|
|
|
|
parser.add_argument(
|
|
"--model_id",
|
|
type=str,
|
|
default="Hcompany/Holo1-3B",
|
|
help="The identifier for the model to use (default: Hcompany/Holo1-3B)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dataset_id",
|
|
type=str,
|
|
default="rootsautomation/ScreenSpot",
|
|
help="The identifier for the dataset to use (default: rootsautomation/ScreenSpot)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--outfile",
|
|
type=str,
|
|
default="result.json",
|
|
help="Output json-file containing the aggregated metrics.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--use_toolcall",
|
|
type=bool,
|
|
default=True,
|
|
help="Enable or disable tool call prompting",
|
|
)
|
|
args = parser.parse_args()
|
|
main(model_id=args.model_id, dataset_id=args.dataset_id, outfile=args.outfile, use_toolcall=args.use_toolcall) |