1import torch
2from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
3from transformers.generation import GenerationConfig
4import json
5import re
6import os
7import numpy as np
8from PIL import Image, ImageDraw
9from qwen_vl_utils import process_vision_info
10
11
12
13# load model
14model_name_or_path = "microsoft/GUI-Actor-Verifier-2B"
15model = Qwen2VLForConditionalGeneration.from_pretrained(
16 model_name_or_path,
17 device_map="cuda:0",
18 trust_remote_code=True,
19 torch_dtype=torch.bfloat16,
20 attn_implementation="flash_attention_2"
21 ).eval()
22output_len = 1
23
24tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
25processor = AutoProcessor.from_pretrained(model_name_or_path)
26
27def draw_annotations(img, point_in_pixel, bbox, output_path='test.png', color='red', size=1):
28 draw = ImageDraw.Draw(img)
29
30 # Draw the ground truth bounding box in green
31 if bbox:
32 # Assuming bbox format is [x1, y1, x2, y2]
33 draw.rectangle(bbox, outline="yellow", width=4)
34
35 # Draw a small circle around the predicted point in red
36 if point_in_pixel:
37 # Create a small rectangle around the point (5 pixels in each direction)
38 radius = np.ceil(8 * size).astype(int)
39 circle_bbox = [
40 point_in_pixel[0] - radius, # x1
41 point_in_pixel[1] - radius, # y1
42 point_in_pixel[0] + radius, # x2
43 point_in_pixel[1] + radius # y2
44 ]
45 draw.ellipse(circle_bbox, outline=color, width=np.ceil(4 * size).astype(int))
46
47 return img
48
49def ground_only_positive(model, tokenizer, processor, instruction, image, point):
50 if isinstance(image, str):
51 image_path = image
52 image = Image.open(image_path)
53 else:
54 image_path = image_to_temp_filename(image)
55 assert os.path.exists(image_path) and os.path.isfile(image_path), "Invalid input image path."
56
57 width, height = image.size
58 image = draw_annotations(image, point, None, output_path=None, size=height/1000 * 1.2)
59
60 prompt_origin = "Please observe the screenshot and exame whether the hollow red circle accurately placed on the intended position in the image: '{}'. Answer True or False."
61 full_prompt = prompt_origin.format(instruction)
62
63 messages = [
64 {
65 "role": "user",
66 "content": [
67 {
68 "type": "image",
69 "image": image,
70 },
71 {"type": "text", "text": full_prompt},
72 ],
73 }
74 ]
75 # Preparation for inference
76 text_input = processor.apply_chat_template(
77 messages, tokenize=False, add_generation_prompt=True
78 )
79 image_inputs, video_inputs = process_vision_info(messages)
80 inputs = processor(
81 text=[text_input],
82 images=image_inputs,
83 videos=video_inputs,
84 padding=True,
85 return_tensors="pt",
86 )
87 inputs = inputs.to("cuda:0")
88
89 generated_ids = model.generate(
90 **inputs,
91 max_new_tokens=output_len,
92 do_sample=False,
93 temperature=0.0
94 )
95
96 generated_ids_trimmed = [
97 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
98 ]
99 response = processor.batch_decode(
100 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
101 )[0]
102
103 print(response)
104 matches = re.findall(r'\b(?:True|False)\b', response)
105 if not len(matches):
106 answer = 'Error Format'
107 else:
108 answer = matches[-1]
109 return answer
110
111# given the image path and instruction and coorindate
112instruction = 'close this window'
113image = Image.open('test.png')
114width, height = image.size
115point = [int(0.9709 * width), int(0.1548, * height)] # The point should be in pixels
116answer = ground_only_positive(model, tokenizer, processor, instruction, image, point) # output True or False