makeitfr commited on
Commit
2111212
·
verified ·
1 Parent(s): 448c739

Upload OmniParser/util/utils.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. OmniParser/util/utils.py +628 -0
OmniParser/util/utils.py ADDED
@@ -0,0 +1,628 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # from ultralytics import YOLO
2
+ import os
3
+ import io
4
+ import base64
5
+ import time
6
+ from PIL import Image, ImageDraw, ImageFont
7
+ import json
8
+ import requests
9
+ # utility function
10
+ import os
11
+ from openai import AzureOpenAI
12
+
13
+ import json
14
+ import sys
15
+ import os
16
+ import cv2
17
+ import numpy as np
18
+ # %matplotlib inline
19
+ from matplotlib import pyplot as plt
20
+ import easyocr
21
+ from paddleocr import PaddleOCR
22
+ reader = easyocr.Reader(['en'])
23
+ paddle_ocr = PaddleOCR(lang='en', use_angle_cls=False)
24
+ import time
25
+ import base64
26
+
27
+ import os
28
+ import ast
29
+ import torch
30
+ from typing import Tuple, List, Union
31
+ from torchvision.ops import box_convert
32
+ import re
33
+ from torchvision.transforms import ToPILImage
34
+ import supervision as sv
35
+ import torchvision.transforms as T
36
+ from util.box_annotator import BoxAnnotator
37
+
38
+
39
+ def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
40
+ if not device:
41
+ device = "cuda" if torch.cuda.is_available() else "cpu"
42
+ if model_name == "blip2":
43
+ from transformers import Blip2Processor, Blip2ForConditionalGeneration
44
+ processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
45
+ if device == 'cpu':
46
+ model = Blip2ForConditionalGeneration.from_pretrained(
47
+ model_name_or_path, device_map=None, torch_dtype=torch.float32
48
+ )
49
+ else:
50
+ model = Blip2ForConditionalGeneration.from_pretrained(
51
+ model_name_or_path, device_map=None, torch_dtype=torch.float16
52
+ ).to(device)
53
+ else:
54
+ raise ValueError(f"Model {model_name} not supported. Only 'blip2' is available.")
55
+ return {'model': model.to(device), 'processor': processor}
56
+
57
+
58
+ def get_yolo_model(model_path):
59
+ from ultralytics import YOLO
60
+ # Load the model.
61
+ model = YOLO(model_path)
62
+ return model
63
+
64
+
65
+ @torch.inference_mode()
66
+ def get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=None, batch_size=128):
67
+ # Number of samples per batch, --> 128 roughly takes 4 GB of GPU memory for blip2 model
68
+ to_pil = ToPILImage()
69
+ if starting_idx:
70
+ non_ocr_boxes = filtered_boxes[starting_idx:]
71
+ else:
72
+ non_ocr_boxes = filtered_boxes
73
+ croped_pil_image = []
74
+ for i, coord in enumerate(non_ocr_boxes):
75
+ try:
76
+ xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
77
+ ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
78
+ cropped_image = image_source[ymin:ymax, xmin:xmax, :]
79
+ cropped_image = cv2.resize(cropped_image, (64, 64))
80
+ croped_pil_image.append(to_pil(cropped_image))
81
+ except Exception as e:
82
+ print(f"[WARNING] Failed to crop image at {i}: {str(e)}")
83
+ continue
84
+
85
+ model, processor = caption_model_processor['model'], caption_model_processor['processor']
86
+ if not prompt:
87
+ prompt = "The image shows"
88
+
89
+ # Reduce batch size on CPU to prevent crashes
90
+ device = model.device
91
+ if device.type == 'cpu':
92
+ batch_size = max(1, min(4, batch_size)) # Use max 4 on CPU
93
+ print(f"[Caption] Running on CPU, reducing batch_size to {batch_size}")
94
+
95
+ print(f"[Caption] Processing {len(croped_pil_image)} cropped images with batch_size={batch_size}")
96
+ generated_texts = []
97
+ try:
98
+ for i in range(0, len(croped_pil_image), batch_size):
99
+ start = time.time()
100
+ batch = croped_pil_image[i:i+batch_size]
101
+ print(f"[Caption] Batch {i//batch_size + 1}: processing {len(batch)} images...")
102
+
103
+ try:
104
+ t1 = time.time()
105
+ if model.device.type == 'cuda':
106
+ inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt", do_resize=False).to(device=device, dtype=torch.float16)
107
+ else:
108
+ inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
109
+ print(f"[Caption] Inputs prepared in {time.time()-t1:.2f}s")
110
+
111
+ t2 = time.time()
112
+ generated_ids = model.generate(**inputs, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True, num_return_sequences=1)
113
+ print(f"[Caption] Generation done in {time.time()-t2:.2f}s")
114
+
115
+ generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
116
+ generated_text = [gen.strip() for gen in generated_text]
117
+ generated_texts.extend(generated_text)
118
+ print(f"[Caption] Batch complete in {time.time()-start:.2f}s")
119
+ except Exception as e:
120
+ print(f"[ERROR] Batch generation failed: {str(e)}")
121
+ import traceback
122
+ print(traceback.format_exc())
123
+ # Add placeholder captions for this batch instead of crashing
124
+ for j in range(len(batch)):
125
+ generated_texts.append(f"Item {len(generated_texts)+1}")
126
+ print(f"[WARNING] Using placeholder captions for this batch")
127
+ except Exception as e:
128
+ print(f"[ERROR] Caption processing failed: {str(e)}")
129
+ import traceback
130
+ print(traceback.format_exc())
131
+ # Return placeholder captions if generation fails
132
+ generated_texts = [f"Item {i+1}" for i in range(len(croped_pil_image))]
133
+
134
+ return generated_texts
135
+
136
+
137
+
138
+ def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor):
139
+ to_pil = ToPILImage()
140
+ if ocr_bbox:
141
+ non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
142
+ else:
143
+ non_ocr_boxes = filtered_boxes
144
+ croped_pil_image = []
145
+ for i, coord in enumerate(non_ocr_boxes):
146
+ xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
147
+ ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
148
+ cropped_image = image_source[ymin:ymax, xmin:xmax, :]
149
+ croped_pil_image.append(to_pil(cropped_image))
150
+
151
+ model, processor = caption_model_processor['model'], caption_model_processor['processor']
152
+ device = model.device
153
+ messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}]
154
+ prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
155
+
156
+ batch_size = 5 # Number of samples per batch
157
+ generated_texts = []
158
+
159
+ for i in range(0, len(croped_pil_image), batch_size):
160
+ images = croped_pil_image[i:i+batch_size]
161
+ image_inputs = [processor.image_processor(x, return_tensors="pt") for x in images]
162
+ inputs ={'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
163
+ texts = [prompt] * len(images)
164
+ for i, txt in enumerate(texts):
165
+ input = processor._convert_images_texts_to_inputs(image_inputs[i], txt, return_tensors="pt")
166
+ inputs['input_ids'].append(input['input_ids'])
167
+ inputs['attention_mask'].append(input['attention_mask'])
168
+ inputs['pixel_values'].append(input['pixel_values'])
169
+ inputs['image_sizes'].append(input['image_sizes'])
170
+ max_len = max([x.shape[1] for x in inputs['input_ids']])
171
+ for i, v in enumerate(inputs['input_ids']):
172
+ inputs['input_ids'][i] = torch.cat([processor.tokenizer.pad_token_id * torch.ones(1, max_len - v.shape[1], dtype=torch.long), v], dim=1)
173
+ inputs['attention_mask'][i] = torch.cat([torch.zeros(1, max_len - v.shape[1], dtype=torch.long), inputs['attention_mask'][i]], dim=1)
174
+ inputs_cat = {k: torch.concatenate(v).to(device) for k, v in inputs.items()}
175
+
176
+ generation_args = {
177
+ "max_new_tokens": 25,
178
+ "temperature": 0.01,
179
+ "do_sample": False,
180
+ }
181
+ generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
182
+ # # remove input tokens
183
+ generate_ids = generate_ids[:, inputs_cat['input_ids'].shape[1]:]
184
+ response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
185
+ response = [res.strip('\n').strip() for res in response]
186
+ generated_texts.extend(response)
187
+
188
+ return generated_texts
189
+
190
+ def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
191
+ assert ocr_bbox is None or isinstance(ocr_bbox, List)
192
+
193
+ def box_area(box):
194
+ return (box[2] - box[0]) * (box[3] - box[1])
195
+
196
+ def intersection_area(box1, box2):
197
+ x1 = max(box1[0], box2[0])
198
+ y1 = max(box1[1], box2[1])
199
+ x2 = min(box1[2], box2[2])
200
+ y2 = min(box1[3], box2[3])
201
+ return max(0, x2 - x1) * max(0, y2 - y1)
202
+
203
+ def IoU(box1, box2):
204
+ intersection = intersection_area(box1, box2)
205
+ union = box_area(box1) + box_area(box2) - intersection + 1e-6
206
+ if box_area(box1) > 0 and box_area(box2) > 0:
207
+ ratio1 = intersection / box_area(box1)
208
+ ratio2 = intersection / box_area(box2)
209
+ else:
210
+ ratio1, ratio2 = 0, 0
211
+ return max(intersection / union, ratio1, ratio2)
212
+
213
+ def is_inside(box1, box2):
214
+ # return box1[0] >= box2[0] and box1[1] >= box2[1] and box1[2] <= box2[2] and box1[3] <= box2[3]
215
+ intersection = intersection_area(box1, box2)
216
+ ratio1 = intersection / box_area(box1)
217
+ return ratio1 > 0.95
218
+
219
+ boxes = boxes.tolist()
220
+ filtered_boxes = []
221
+ if ocr_bbox:
222
+ filtered_boxes.extend(ocr_bbox)
223
+ # print('ocr_bbox!!!', ocr_bbox)
224
+ for i, box1 in enumerate(boxes):
225
+ # if not any(IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2) for j, box2 in enumerate(boxes) if i != j):
226
+ is_valid_box = True
227
+ for j, box2 in enumerate(boxes):
228
+ # keep the smaller box
229
+ if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
230
+ is_valid_box = False
231
+ break
232
+ if is_valid_box:
233
+ # add the following 2 lines to include ocr bbox
234
+ if ocr_bbox:
235
+ # only add the box if it does not overlap with any ocr bbox
236
+ if not any(IoU(box1, box3) > iou_threshold and not is_inside(box1, box3) for k, box3 in enumerate(ocr_bbox)):
237
+ filtered_boxes.append(box1)
238
+ else:
239
+ filtered_boxes.append(box1)
240
+ return torch.tensor(filtered_boxes)
241
+
242
+
243
+ def remove_overlap_new(boxes, iou_threshold, ocr_bbox=None):
244
+ '''
245
+ ocr_bbox format: [{'type': 'text', 'bbox':[x,y], 'interactivity':False, 'content':str }, ...]
246
+ boxes format: [{'type': 'icon', 'bbox':[x,y], 'interactivity':True, 'content':None }, ...]
247
+
248
+ '''
249
+ assert ocr_bbox is None or isinstance(ocr_bbox, List)
250
+
251
+ def box_area(box):
252
+ return (box[2] - box[0]) * (box[3] - box[1])
253
+
254
+ def intersection_area(box1, box2):
255
+ x1 = max(box1[0], box2[0])
256
+ y1 = max(box1[1], box2[1])
257
+ x2 = min(box1[2], box2[2])
258
+ y2 = min(box1[3], box2[3])
259
+ return max(0, x2 - x1) * max(0, y2 - y1)
260
+
261
+ def IoU(box1, box2):
262
+ intersection = intersection_area(box1, box2)
263
+ union = box_area(box1) + box_area(box2) - intersection + 1e-6
264
+ if box_area(box1) > 0 and box_area(box2) > 0:
265
+ ratio1 = intersection / box_area(box1)
266
+ ratio2 = intersection / box_area(box2)
267
+ else:
268
+ ratio1, ratio2 = 0, 0
269
+ return max(intersection / union, ratio1, ratio2)
270
+
271
+ def is_inside(box1, box2):
272
+ # return box1[0] >= box2[0] and box1[1] >= box2[1] and box1[2] <= box2[2] and box1[3] <= box2[3]
273
+ intersection = intersection_area(box1, box2)
274
+ ratio1 = intersection / box_area(box1)
275
+ return ratio1 > 0.80
276
+
277
+ # boxes = boxes.tolist()
278
+ filtered_boxes = []
279
+ if ocr_bbox:
280
+ filtered_boxes.extend(ocr_bbox)
281
+ # print('ocr_bbox!!!', ocr_bbox)
282
+ for i, box1_elem in enumerate(boxes):
283
+ box1 = box1_elem['bbox']
284
+ is_valid_box = True
285
+ for j, box2_elem in enumerate(boxes):
286
+ # keep the smaller box
287
+ box2 = box2_elem['bbox']
288
+ if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
289
+ is_valid_box = False
290
+ break
291
+ if is_valid_box:
292
+ if ocr_bbox:
293
+ # keep yolo boxes + prioritize ocr label
294
+ box_added = False
295
+ ocr_labels = ''
296
+ for box3_elem in ocr_bbox:
297
+ if not box_added:
298
+ box3 = box3_elem['bbox']
299
+ if is_inside(box3, box1): # ocr inside icon
300
+ # box_added = True
301
+ # delete the box3_elem from ocr_bbox
302
+ try:
303
+ # gather all ocr labels
304
+ ocr_labels += box3_elem['content'] + ' '
305
+ filtered_boxes.remove(box3_elem)
306
+ except:
307
+ continue
308
+ # break
309
+ elif is_inside(box1, box3): # icon inside ocr, don't added this icon box, no need to check other ocr bbox bc no overlap between ocr bbox, icon can only be in one ocr box
310
+ box_added = True
311
+ break
312
+ else:
313
+ continue
314
+ if not box_added:
315
+ if ocr_labels:
316
+ filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': ocr_labels, 'source':'box_yolo_content_ocr'})
317
+ else:
318
+ filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': None, 'source':'box_yolo_content_yolo'})
319
+ else:
320
+ filtered_boxes.append(box1)
321
+ return filtered_boxes # torch.tensor(filtered_boxes)
322
+
323
+
324
+ def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
325
+ transform = T.Compose(
326
+ [
327
+ T.RandomResize([800], max_size=1333),
328
+ T.ToTensor(),
329
+ T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
330
+ ]
331
+ )
332
+ image_source = Image.open(image_path).convert("RGB")
333
+ image = np.asarray(image_source)
334
+ image_transformed, _ = transform(image_source, None)
335
+ return image, image_transformed
336
+
337
+
338
+ def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], text_scale: float,
339
+ text_padding=5, text_thickness=2, thickness=3) -> np.ndarray:
340
+ """
341
+ This function annotates an image with bounding boxes and labels.
342
+
343
+ Parameters:
344
+ image_source (np.ndarray): The source image to be annotated.
345
+ boxes (torch.Tensor): A tensor containing bounding box coordinates. in cxcywh format, pixel scale
346
+ logits (torch.Tensor): A tensor containing confidence scores for each bounding box.
347
+ phrases (List[str]): A list of labels for each bounding box.
348
+ text_scale (float): The scale of the text to be displayed. 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
349
+
350
+ Returns:
351
+ np.ndarray: The annotated image.
352
+ """
353
+ h, w, _ = image_source.shape
354
+ boxes = boxes * torch.Tensor([w, h, w, h])
355
+ xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
356
+ xywh = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xywh").numpy()
357
+ detections = sv.Detections(xyxy=xyxy)
358
+
359
+ labels = [f"{phrase}" for phrase in range(boxes.shape[0])]
360
+
361
+ box_annotator = BoxAnnotator(text_scale=text_scale, text_padding=text_padding,text_thickness=text_thickness,thickness=thickness) # 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
362
+ annotated_frame = image_source.copy()
363
+ annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w,h))
364
+
365
+ label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
366
+ return annotated_frame, label_coordinates
367
+
368
+
369
+ def predict(model, image, caption, box_threshold, text_threshold):
370
+ """ Use huggingface model to replace the original model
371
+ """
372
+ model, processor = model['model'], model['processor']
373
+ device = model.device
374
+
375
+ inputs = processor(images=image, text=caption, return_tensors="pt").to(device)
376
+ with torch.no_grad():
377
+ outputs = model(**inputs)
378
+
379
+ results = processor.post_process_grounded_object_detection(
380
+ outputs,
381
+ inputs.input_ids,
382
+ box_threshold=box_threshold, # 0.4,
383
+ text_threshold=text_threshold, # 0.3,
384
+ target_sizes=[image.size[::-1]]
385
+ )[0]
386
+ boxes, logits, phrases = results["boxes"], results["scores"], results["labels"]
387
+ return boxes, logits, phrases
388
+
389
+
390
+ def predict_yolo(model, image, box_threshold, imgsz, scale_img, iou_threshold=0.7):
391
+ """ Use huggingface model to replace the original model
392
+ """
393
+ # model = model['model']
394
+ if scale_img:
395
+ result = model.predict(
396
+ source=image,
397
+ conf=box_threshold,
398
+ imgsz=imgsz,
399
+ iou=iou_threshold, # default 0.7
400
+ )
401
+ else:
402
+ result = model.predict(
403
+ source=image,
404
+ conf=box_threshold,
405
+ iou=iou_threshold, # default 0.7
406
+ )
407
+ boxes = result[0].boxes.xyxy#.tolist() # in pixel space
408
+ conf = result[0].boxes.conf
409
+ phrases = [str(i) for i in range(len(boxes))]
410
+
411
+ return boxes, conf, phrases
412
+
413
+ def int_box_area(box, w, h):
414
+ x1, y1, x2, y2 = box
415
+ int_box = [int(x1*w), int(y1*h), int(x2*w), int(y2*h)]
416
+ area = (int_box[2] - int_box[0]) * (int_box[3] - int_box[1])
417
+ return area
418
+
419
+ def get_som_labeled_img(image_source: Union[str, Image.Image], model=None, BOX_TRESHOLD=0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None, scale_img=False, imgsz=None, batch_size=128, save_cropped_images=False, cropped_images_dir='cropped_images'):
420
+ """Process either an image path or Image object
421
+
422
+ Args:
423
+ image_source: Either a file path (str) or PIL Image object
424
+ ...
425
+ """
426
+ if isinstance(image_source, str):
427
+ image_source = Image.open(image_source)
428
+ image_source = image_source.convert("RGB") # for CLIP
429
+ w, h = image_source.size
430
+ if not imgsz:
431
+ imgsz = (h, w)
432
+ # print('image size:', w, h)
433
+ xyxy, logits, phrases = predict_yolo(model=model, image=image_source, box_threshold=BOX_TRESHOLD, imgsz=imgsz, scale_img=scale_img, iou_threshold=0.1)
434
+ xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device)
435
+ image_source = np.asarray(image_source)
436
+ phrases = [str(i) for i in range(len(phrases))]
437
+
438
+ # annotate the image with labels
439
+ if ocr_bbox:
440
+ ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
441
+ ocr_bbox=ocr_bbox.tolist()
442
+ else:
443
+ print('no ocr bbox!!!')
444
+ ocr_bbox = None
445
+
446
+ if ocr_bbox is not None:
447
+ ocr_bbox_elem = [{'type': 'text', 'bbox':box, 'interactivity':False, 'content':txt, 'source': 'box_ocr_content_ocr'} for box, txt in zip(ocr_bbox, ocr_text) if int_box_area(box, w, h) > 0]
448
+ else:
449
+ ocr_bbox_elem = []
450
+ xyxy_elem = [{'type': 'icon', 'bbox':box, 'interactivity':True, 'content':None} for box in xyxy.tolist() if int_box_area(box, w, h) > 0]
451
+ filtered_boxes = remove_overlap_new(boxes=xyxy_elem, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox_elem)
452
+
453
+ # sort the filtered_boxes so that the one with 'content': None is at the end, and get the index of the first 'content': None
454
+ filtered_boxes_elem = sorted(filtered_boxes, key=lambda x: x['content'] is None)
455
+ # get the index of the first 'content': None
456
+ starting_idx = next((i for i, box in enumerate(filtered_boxes_elem) if box['content'] is None), -1)
457
+ filtered_boxes = torch.tensor([box['bbox'] for box in filtered_boxes_elem])
458
+ print('len(filtered_boxes):', len(filtered_boxes), starting_idx)
459
+
460
+ # get parsed icon local semantics
461
+ time1 = time.time()
462
+ if use_local_semantics and caption_model_processor is not None:
463
+ caption_model = caption_model_processor['model']
464
+ try:
465
+ print("[SOM] Starting caption generation...")
466
+ if 'phi3_v' in caption_model.config.model_type:
467
+ print("[SOM] Using Phi3V model")
468
+ parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor)
469
+ else:
470
+ print("[SOM] Using Caption model")
471
+ parsed_content_icon = get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=prompt,batch_size=batch_size)
472
+ print("[SOM] Caption generation complete")
473
+ except Exception as e:
474
+ print(f"[ERROR] Caption generation failed: {str(e)}")
475
+ import traceback
476
+ print(traceback.format_exc())
477
+ # Use placeholder captions if generation fails
478
+ num_icons = len(filtered_boxes) - (starting_idx if starting_idx > 0 else 0)
479
+ parsed_content_icon = [f"Icon {i}" for i in range(num_icons)]
480
+ elif use_local_semantics:
481
+ # Use OCR text if no caption model is available
482
+ print("[SOM] Caption model disabled, using OCR text only")
483
+
484
+ # Save cropped images if requested
485
+ if save_cropped_images:
486
+ print(f"[SOM] Saving cropped UI images to {cropped_images_dir}...")
487
+ from pathlib import Path
488
+ Path(cropped_images_dir).mkdir(parents=True, exist_ok=True)
489
+ saved_count = 0
490
+ for i, bbox in enumerate(filtered_boxes):
491
+ try:
492
+ # Convert from ratio coordinates to pixel coordinates
493
+ xmin = int(bbox[0] * image_source.shape[1])
494
+ xmax = int(bbox[2] * image_source.shape[1])
495
+ ymin = int(bbox[1] * image_source.shape[0])
496
+ ymax = int(bbox[3] * image_source.shape[0])
497
+
498
+ # Ensure coordinates are within bounds
499
+ xmin = max(0, xmin)
500
+ ymin = max(0, ymin)
501
+ xmax = min(image_source.shape[1], xmax)
502
+ ymax = min(image_source.shape[0], ymax)
503
+
504
+ # Crop the image
505
+ cropped_image = image_source[ymin:ymax, xmin:xmax, :]
506
+
507
+ # Save the cropped image
508
+ pil_image = Image.fromarray(cropped_image)
509
+ save_path = os.path.join(cropped_images_dir, f"crop_{i:04d}.png")
510
+ pil_image.save(save_path)
511
+ saved_count += 1
512
+ except Exception as e:
513
+ print(f"[WARNING] Failed to save crop {i}: {str(e)}")
514
+ print(f"[SOM] Saved {saved_count} cropped UI images")
515
+
516
+ parsed_content_icon = []
517
+ for i, bbox in enumerate(filtered_boxes):
518
+ # Check if any OCR text intersects with this bbox
519
+ label = f"Icon {i}"
520
+ if ocr_text and ocr_bbox is not None and len(ocr_text) > 0:
521
+ for j, ocr_b in enumerate(ocr_bbox):
522
+ # Simple intersection check
523
+ if (bbox[0] < ocr_b[2] and bbox[2] > ocr_b[0] and
524
+ bbox[1] < ocr_b[3] and bbox[3] > ocr_b[1]):
525
+ label = ocr_text[j]
526
+ break
527
+ parsed_content_icon.append(label)
528
+
529
+ ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
530
+ icon_start = len(ocr_text)
531
+ parsed_content_icon_ls = []
532
+ # fill the filtered_boxes_elem None content with parsed_content_icon in order
533
+ for i, box in enumerate(filtered_boxes_elem):
534
+ if box['content'] is None:
535
+ box['content'] = parsed_content_icon.pop(0)
536
+ for i, txt in enumerate(parsed_content_icon):
537
+ parsed_content_icon_ls.append(f"Icon Box ID {str(i+icon_start)}: {txt}")
538
+ parsed_content_merged = ocr_text + parsed_content_icon_ls
539
+ else:
540
+ ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
541
+ parsed_content_merged = ocr_text
542
+ print('time to get parsed content:', time.time()-time1)
543
+
544
+ if len(filtered_boxes) > 0:
545
+ filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
546
+ else:
547
+ # No boxes detected, return empty results
548
+ print('No boxes detected in the image')
549
+ annotated_frame = image_source.copy()
550
+ label_coordinates = []
551
+ pil_img = Image.fromarray(annotated_frame)
552
+ buffered = io.BytesIO()
553
+ pil_img.save(buffered, format="PNG")
554
+ buffered.seek(0)
555
+ pil_str = base64.b64encode(buffered.getvalue()).decode()
556
+ return pil_str, label_coordinates, []
557
+
558
+ phrases = [i for i in range(len(filtered_boxes))]
559
+
560
+ # draw boxes
561
+ if draw_bbox_config:
562
+ annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
563
+ else:
564
+ annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding)
565
+
566
+ pil_img = Image.fromarray(annotated_frame)
567
+ buffered = io.BytesIO()
568
+ pil_img.save(buffered, format="PNG")
569
+ encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
570
+ if output_coord_in_ratio:
571
+ label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()}
572
+ assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
573
+
574
+ return encoded_image, label_coordinates, filtered_boxes_elem
575
+
576
+
577
+ def get_xywh(input):
578
+ x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1]
579
+ x, y, w, h = int(x), int(y), int(w), int(h)
580
+ return x, y, w, h
581
+
582
+ def get_xyxy(input):
583
+ x, y, xp, yp = input[0][0], input[0][1], input[2][0], input[2][1]
584
+ x, y, xp, yp = int(x), int(y), int(xp), int(yp)
585
+ return x, y, xp, yp
586
+
587
+ def get_xywh_yolo(input):
588
+ x, y, w, h = input[0], input[1], input[2] - input[0], input[3] - input[1]
589
+ x, y, w, h = int(x), int(y), int(w), int(h)
590
+ return x, y, w, h
591
+
592
+ def check_ocr_box(image_source: Union[str, Image.Image], display_img = True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None, use_paddleocr=False):
593
+ if isinstance(image_source, str):
594
+ image_source = Image.open(image_source)
595
+ if image_source.mode == 'RGBA':
596
+ # Convert RGBA to RGB to avoid alpha channel issues
597
+ image_source = image_source.convert('RGB')
598
+ image_np = np.array(image_source)
599
+ w, h = image_source.size
600
+ if use_paddleocr:
601
+ if easyocr_args is None:
602
+ text_threshold = 0.5
603
+ else:
604
+ text_threshold = easyocr_args['text_threshold']
605
+ result = paddle_ocr.ocr(image_np, cls=False)[0]
606
+ coord = [item[0] for item in result if item[1][1] > text_threshold]
607
+ text = [item[1][0] for item in result if item[1][1] > text_threshold]
608
+ else: # EasyOCR
609
+ if easyocr_args is None:
610
+ easyocr_args = {}
611
+ result = reader.readtext(image_np, **easyocr_args)
612
+ coord = [item[0] for item in result]
613
+ text = [item[1] for item in result]
614
+ if display_img:
615
+ opencv_img = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
616
+ bb = []
617
+ for item in coord:
618
+ x, y, a, b = get_xywh(item)
619
+ bb.append((x, y, a, b))
620
+ cv2.rectangle(opencv_img, (x, y), (x+a, y+b), (0, 255, 0), 2)
621
+ # matplotlib expects RGB
622
+ plt.imshow(cv2.cvtColor(opencv_img, cv2.COLOR_BGR2RGB))
623
+ else:
624
+ if output_bb_format == 'xywh':
625
+ bb = [get_xywh(item) for item in coord]
626
+ elif output_bb_format == 'xyxy':
627
+ bb = [get_xyxy(item) for item in coord]
628
+ return (text, bb), goal_filtering