Create screenspot_eval.py

screenspot_eval.py

@@ -0,0 +1,277 @@
+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)