Views
No views yet
pip install -r .requirements.txt
1import torch
2from PIL import Image, ImageDraw
3from qwen_vl_utils import process_vision_info
4from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
5import os
6
7model = Qwen2VLForConditionalGeneration.from_pretrained(
8 'FRank62Wu/ShowUI-Narrator', torch_dtype="auto", device_map="cuda"
9)
10
11
12# processor = AutoProcessor.from_pretrained('FRank62Wu/ShowUI-Narrator') # load from local dir
13
14image_processor_kwargs = {
15 "size": {
16 "shortest_edge": 56*56,
17 "longest_edge": 720*28*28
18 }
19}
20
21processor = AutoProcessor.from_pretrained(
22 'FRank62Wu/ShowUI-Narrator',
23 **image_processor_kwargs
24)
25
26processor.tokenizer.pad_token = processor.tokenizer.eos_token1import os
2import base64
3from PIL import Image
4from io import BytesIO
5import copy
6
7import cv2
8from ultralytics import YOLO
9
10def image_to_base64(img_path):
11 with open(img_path, "rb") as img_file:
12 encoded_img = base64.b64encode(img_file.read()).decode("utf-8")
13 return encoded_img
14
15check_point_path = './ShowUI_Action_Narrator_cursor_detect/best.pt'
16
17class Screenshots_processor:
18 def __init__(self, img_path, max_size, delta, check_point_path):
19 self.img_path = img_path
20 self.cursor_model = YOLO(check_point_path)
21 self.scs = []
22 self.crop_scs =[]
23 self.max_size = max_size
24 self.delta = delta
25
26 def create_crop(self):
27 for each in sorted(os.listdir(self.img_path)):
28 if each.endswith('.jsonl') or '_crop' in each:
29 continue
30 else:
31 each = os.path.join(self.img_path, each)
32 self.scs.append(each)
33
34 frame_x, frame_y = [], []
35 for idx, image_path in enumerate(self.scs):
36 results = self.cursor_model(image_path)
37 img = Image.open(image_path)
38 width, height = img.size
39 img.close()
40
41 for result in results:
42 if result.boxes.xywh.size(0) > 0:
43 boxes = result.boxes
44 xywh_tensor = boxes.xywh
45 x, y = xywh_tensor[0][0].item(), xywh_tensor[0][1].item()
46 frame_x.append(x)
47 frame_y.append(y)
48 else:
49 print('Cursor not detected')
50
51 if len(frame_x) == 0 or len(frame_y) ==0:
52 self.crop_scs = copy.deepcopy(self.scs)
53 return self.crop_scs
54
55 elif (len(frame_x) <= 1) or (max(frame_x)- min(frame_x))>=self.max_size or (max(frame_y)- min(frame_y))>=self.max_size:
56 print('add margin')
57 mid_x, mid_y = sum(frame_x) // len(frame_x), sum(frame_y) // len(frame_y)
58 margin_= self.max_size + self.delta
59 for idx, each in enumerate(sorted(self.scs)):
60 image_path = each
61 image1 = Image.open(image_path).convert('RGB')
62 file_name_tail = image_path.split('/')[-1]
63 save_path = image_path.replace(file_name_tail, f'{idx}_crop.jpg')
64
65 x1 = max(0, min(width - margin_, mid_x - margin_ // 2))
66 y1 = max(0, min(height - margin_, mid_y - margin_ // 2))
67 x2 = min(x1 + margin_, width)
68 y2 = min(y1 + margin_, height)
69 start_crop = image1.crop((x1, y1, x2, y2))
70
71 start_crop.save(save_path)
72 self.crop_scs.append(save_path)
73 image1.close()
74 return self.crop_scs, self.scs
75
76 else:
77 mid_x, mid_y = sum(frame_x) // len(frame_x), sum(frame_y) // len(frame_y)
78 margin = self.max_size
79 margin_ = self.max_size
80 x1 = max(0, min(width - margin, mid_x - margin // 2))
81 y1 = max(0, min(height - margin, mid_y - margin // 2))
82 x2 = min(x1 + margin, width)
83 y2 = min(y1 + margin, height)
84 for idx, each in enumerate(sorted(self.scs)):
85 image_path = each
86 image1 = Image.open(image_path).convert('RGB')
87 file_name_tail = image_path.split('/')[-1]
88 save_path = image_path.replace(file_name_tail, f'{idx}_crop.jpg')
89
90 x1 = max(0, min(width - margin_, mid_x - margin_ // 2))
91 y1 = max(0, min(height - margin_, mid_y - margin_ // 2))
92 x2 = min(x1 + margin_, width)
93 y2 = min(y1 + margin_, height)
94 start_crop = image1.crop((x1, y1, x2, y2))
95
96 start_crop.save(save_path)
97 self.crop_scs.append(save_path)
98 image1.close()
99 return self.crop_scs, self.scs
100
101
102
103class Videoscreen_processor:
104 def __init__(self, vid_path, fps, max_size, delta, check_point_path):
105 self.vid_path = vid_path
106 self.fps = fps
107 self.cursor_model = YOLO(check_point_path)
108 self.scs = []
109 self.crop_scs =[]
110 self.max_size = max_size
111 self.delta = delta
112
113
114
115 def sample_from_video(self):
116
117 video_path_tail = self.vid_path.split('/')[-1]
118 cap = cv2.VideoCapture(self.vid_path)
119 if not cap.isOpened():
120 print("Error: Could not open video.")
121 return []
122 video_fps = cap.get(cv2.CAP_PROP_FPS) # fps
123 print(video_fps)
124 total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
125 frame_interval = int(video_fps // self.fps)
126 frame_count = 0
127 frame_num = 0
128 while True:
129 ret, frame = cap.read()
130 if not ret:
131 break
132 if frame_count>1:
133 break
134 if frame_num % frame_interval == 0:
135 frame_count = frame_num // frame_interval
136 image_path = os.path.join(self.vid_path.replace(video_path_tail, f"frame_{frame_count}.jpg"))
137 self.scs.append(image_path)
138 frame_count += 1
139 cv2.imwrite(image_path, frame)
140 frame_num += 1
141 cap.release()
142
143 frame_x, frame_y = [], []
144 for idx, image_path in enumerate(self.scs):
145 results = self.cursor_model(image_path)
146 img = Image.open(image_path)
147 width, height = img.size
148 img.close()
149
150 for result in results:
151 if result.boxes.xywh.size(0) > 0:
152 boxes = result.boxes
153 xywh_tensor = boxes.xywh
154 x, y = xywh_tensor[0][0].item(), xywh_tensor[0][1].item()
155 frame_x.append(x)
156 frame_y.append(y)
157 else:
158 print('Cursor not detected')
159
160 if len(frame_x) == 0 or len(frame_y) ==0:
161 self.crop_scs = copy.deepcopy(self.scs)
162 return self.crop_scs, self.crop_scs
163
164 elif (len(frame_x) <= 1) or (max(frame_x)- min(frame_x))>=self.max_size or (max(frame_y)- min(frame_y))>=self.max_size:
165 print('add margin')
166 mid_x, mid_y = sum(frame_x) // len(frame_x), sum(frame_y) // len(frame_y)
167 margin_= self.max_size + self.delta
168 for idx, each in enumerate(sorted(self.scs)):
169 image_path = each
170 image1 = Image.open(image_path).convert('RGB')
171 file_name_tail = image_path.split('/')[-1]
172 save_path = image_path.replace(file_name_tail, f'{idx}_crop.jpg')
173
174 x1 = max(0, min(width - margin_, mid_x - margin_ // 2))
175 y1 = max(0, min(height - margin_, mid_y - margin_ // 2))
176 x2 = min(x1 + margin_, width)
177 y2 = min(y1 + margin_, height)
178 start_crop = image1.crop((x1, y1, x2, y2))
179
180 start_crop.save(save_path)
181 self.crop_scs.append(save_path)
182 image1.close()
183 return self.crop_scs, self.scs
184
185 else:
186 mid_x, mid_y = sum(frame_x) // len(frame_x), sum(frame_y) // len(frame_y)
187 margin = self.max_size
188 x1 = max(0, min(width - margin, mid_x - margin // 2))
189 y1 = max(0, min(height - margin, mid_y - margin // 2))
190 x2 = min(x1 + margin, width)
191 y2 = min(y1 + margin, height)
192 for idx, each in enumerate(sorted(self.scs)):
193 image_path = each
194 image1 = Image.open(image_path).convert('RGB')
195 file_name_tail = image_path.split('/')[-1].replace('frame_','').replace('.png','')
196 save_path = image_path.replace(file_name_tail, f'{idx}_crop.jpg')
197
198 x1 = max(0, min(width - margin_, mid_x - margin_ // 2))
199 y1 = max(0, min(height - margin_, mid_y - margin_ // 2))
200 x2 = min(x1 + margin_, width)
201 y2 = min(y1 + margin_, height)
202 start_crop = image1.crop((x1, y1, x2, y2))
203
204 start_crop.save(save_path)
205 self.crop_scs.append(save_path)
206 image1.close()
207 return self.crop_scs, self.scs
208
2091Cursor_detector = Screenshots_processor('./storage/folder_to_screenshots',512, 128, check_point_path)
2
3cropped_imgs_list, original_imgs_list = Cursor_detector.create_crop()
41"""load model"""
2import torch
3from PIL import Image, ImageDraw
4from qwen_vl_utils import process_vision_info
5from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
6import os
7import json
8import codecs
9import argparse
10import random
11import re
12
13
14max_pixels_temp = 160*28*28
15max_pixels_narr = 760*28*28
16min_pixels_narr = 240*28*28
17
18
19
20
21model = Qwen2VLForConditionalGeneration.from_pretrained(
22 'FRank62Wu/ShowUI-Narrator', torch_dtype="auto", device_map="cuda"
23)
24
25
26processor = AutoProcessor.from_pretrained('FRank62Wu/ShowUI-Narrator')
27processor.tokenizer.pad_token = processor.tokenizer.eos_token
28
29
30_SYSTEM_PROMPT='For the given video frames of a GUI action, The frames are decribed in the format of <0> to <{N}>.'
31
32
33
34_SYSTEM_PROMPT_NARR='''You are an ai assistant to narrate the action of the user for the video frames in the following detail.
35'Action': The type of action
36'Element': The target of the action
37'Source': The starting position (Applicable for action type: Drag)
38'Destination': The ending position (Applicable for action type: Drag)
39'Purpose': The intended result of the action
40The Action include left click, right click, double click, drag, or Keyboard type.
41'''
42
43
44Action_no_reference_grounding = [
45 'Describe the start frame and the end frame of the action in this video?',
46 'When Did the action happened in this video? Tell me the start frame and the end frame.',
47 'Locate the start and the end frame of the action in this video',
48 "Observe the cursor in this GUI video, marking start and end frame of the action in video frames."
49]
50
51
52Dense_narration_query = ['Narrate the action in the given video.',
53 'Describe the action of the user in the given frames',
54 'Describe the action in this video.',
55 'Narrate the action detail of the user in the video.']
561path_to_data =''
2
3query = _SYSTEM_PROMPT.format(N=9) + ' ' + random.choice(Action_no_reference_grounding)
4messages = [
5 {
6 'role': 'user',
7 'content': [
8 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/0_crop.png","max_pixels": max_pixels_temp},
9 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/1_crop.png","max_pixels": max_pixels_temp},
10 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/2_crop.png","max_pixels": max_pixels_temp},
11 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/3_crop.png","max_pixels": max_pixels_temp},
12 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/4_crop.png","max_pixels": max_pixels_temp},
13 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/5_crop.png","max_pixels": max_pixels_temp},
14 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/6_crop.png","max_pixels": max_pixels_temp},
15 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/7_crop.png","max_pixels": max_pixels_temp},
16 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/8_crop.png","max_pixels": max_pixels_temp},
17 {'type':"image", "image": f"{path_to_data}/storage/test_benchmark_Act2Cap/303/9_crop.png","max_pixels": max_pixels_temp},
18 {'type':"text",'text': query},
19 ]
20 }
21 ]
22
23
24
25## round_1 for temporal grounding
26text = processor.apply_chat_template(
27 messages, tokenize=False, add_generation_prompt=True,
28 )
29
30image_inputs, video_inputs = process_vision_info(messages)
31inputs = processor(
32 text=[text],
33 images=image_inputs,
34 videos=video_inputs,
35 padding=True,
36 return_tensors="pt",
37 )
38inputs = inputs.to("cuda")
39generated_ids = model.generate(**inputs, max_new_tokens=128)
40generated_ids_trimmed = [
41 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
42]
43output_text = processor.batch_decode(
44 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
45)[0]
46
47print(output_text)
48>>> Output: <6> and <8>
1# round_2 for dense narration caption
2try:
3 matches = re.search(r"<(\w+)>.*?<(\w+)>", output_text)
4 s1, e1 = int(matches.group(1)), int(matches.group(2))
5except:
6 s1, e1 =0, 9
7
8
9query = _SYSTEM_PROMPT_NARR + ' ' + random.choice(Dense_narration_query)
10
11selected_images = []
12
13if e1-s1<3:
14 pixels_narr = max_pixels_narr
15else:
16 max_pixel_per_image = int(760*3/(e1- s1 +1))*28*28
17 pixels_narr = max_pixel_per_image
18
19
20for idx, each in enumerate(messages[0]['content']):
21 if idx >= s1 and idx <= e1:
22 new_image = each.copy()
23 new_image['max_pixels'] =pixels_narr
24 selected_images.append(new_image)
25
26
27messages = [
28 {
29 'role': 'user',
30 'content':selected_images+ [{'type':"text",'text': query},
31 ]
32 }
33 ]
34
35text = processor.apply_chat_template(
36 messages, tokenize=False, add_generation_prompt=True,
37 )
38
39image_inputs, video_inputs = process_vision_info(messages)
40inputs = processor(
41 text=[text],
42 images=image_inputs,
43 videos=video_inputs,
44 padding=True,
45 return_tensors="pt",
46 )
47inputs = inputs.to("cuda")
48generated_ids = model.generate(**inputs, max_new_tokens=128)
49generated_ids_trimmed = [
50 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
51]
52output_text_narration = processor.batch_decode(
53 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
54)[0]
55
56print(output_text_narration)
57 >>> Output: {"Action": "double click", "Element": "sc2 trans shape button", "Source": null, "Destination": null, "Purpose": " Select the SC2 Trans Shape."}