from transformers import AutoModelForCausalLM, AutoProcessor
from PIL import Image


def postprocess(text: str, image_size: tuple[int]):
    """Function that decodes model's generation into action json.

    Args:
        text: single generated sample
        image_size: corresponding image size
    """
    point_pattern = r"<loc_(\d+)>,<loc_(\d+)>"

    try:
        location = re.findall(point_pattern, text)[0]
        if len(location) > 0:
            point = [int(loc) for loc in location]

    except Exception:
        point = (0, 0)

    return point

# Load model and processor
model = AutoModelForCausalLM.from_pretrained("HongxinLi/GoClick-Large", trust_remote_code=True)
processor = AutoProcessor.from_pretrained("HongxinLi/GoClick-Large", trust_remote_code=True)

# Load UI screenshot
image = Image.open("ui_screenshot.png")

# Stage 1: Planning

# Functionality Grounding (For AutoGUI FuncPred Benchmark)
planning_prompt = f"Locate the element according to its detailed functionality description. {goal_info} (Output the center coordinates of the target)"

# Intent Grounding (For RefExp, MOTIF, and VisualWebBench Action Grounding)
planning_prompt = f"I want to {goal_info}. Please locate the target element I should interact with. (Output the center coordinates of the target)"

# Description Grounding (For ScreenSpot/v2 and VisualWebBench Element Grounding))
planning_prompt = f"Where is the {goal_info} element? (Output the center coordinates of the target)"


inputs = processor(
    images=image,
    text=prompt,
    return_tensors="pt",
    do_resize=True,
).to(model.device, dtype=model.dtype)

outputs = model.generate(
            **inputs,
            do_sample= False,
            max_new_tokens=max_new_tokens,
            use_cache=True
        )

text_output = processor.tokenizer.batch_decode(outputs, skip_special_tokens=False)[0]
text_output = postprocess(text_output, img_size)

@misc{li2026goclicklightweightelementgrounding,
      title={GoClick: Lightweight Element Grounding Model for Autonomous GUI Interaction}, 
      author={Hongxin Li and Yuntao Chen and Zhaoxiang Zhang},
      year={2026},
      eprint={2604.23941},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2604.23941}, 
}