version 1.5

.gitignore

@@ -1,5 +1,6 @@
 weights/icon_caption_blip2
 weights/icon_caption_florence
 weights/icon_detect/
+weights/icon_detect_v1_5/
 .gradio
 __pycache__/

README.md

@@ -7,12 +7,13 @@
 [![arXiv](https://img.shields.io/badge/Paper-green)](https://arxiv.org/abs/2408.00203)
 [![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 
-📢 [[Project Page](https://microsoft.github.io/OmniParser/)] [[Blog Post](https://www.microsoft.com/en-us/research/articles/omniparser-for-pure-vision-based-gui-agent/)] [[Models](https://huggingface.co/microsoft/OmniParser)] [[Huggingface demo](https://huggingface.co/spaces/microsoft/OmniParser)]
+📢 [[Project Page](https://microsoft.github.io/OmniParser/)] [[Blog Post](https://www.microsoft.com/en-us/research/articles/omniparser-for-pure-vision-based-gui-agent/)] [[Models](https://huggingface.co/microsoft/OmniParser)] [huggingface space](https://huggingface.co/spaces/microsoft/OmniParser)
 
 **OmniParser** is a comprehensive method for parsing user interface screenshots into structured and easy-to-understand elements, which significantly enhances the ability of GPT-4V to generate actions that can be accurately grounded in the corresponding regions of the interface. 
 
 ## News
-- [2024/10] OmniParser is the #1 trending model on huggingface model hub (starting 10/29/2024). 
+- [2024/11] OmniParser V1.5 is out! It features: 1) new icon detection model trained on cleaned data, and 2) improved bbox merging logic. 
+- [2024/10] OmniParser was the #1 trending model on huggingface model hub (starting 10/29/2024). 
 - [2024/10] Feel free to checkout our demo on [huggingface space](https://huggingface.co/spaces/microsoft/OmniParser)! (stay tuned for OmniParser + Claude Computer Use)
 - [2024/10] Both Interactive Region Detection Model and Icon functional description model are released! [Hugginface models](https://huggingface.co/microsoft/OmniParser)
 - [2024/09] OmniParser achieves the best performance on [Windows Agent Arena](https://microsoft.github.io/WindowsAgentArena/)! 

utils.py

@@ -75,10 +75,12 @@ def get_yolo_model(model_path):
 
 
 @torch.inference_mode()
-def get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=None):
+def get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=None, batch_size=32):
+    # Number of samples per batch, --> 256 roughly takes 23 GB of GPU memory for florence model
+
     to_pil = ToPILImage()
-    if ocr_bbox:
-        non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
+    if starting_idx:
+        non_ocr_boxes = filtered_boxes[starting_idx:]
     else:
         non_ocr_boxes = filtered_boxes
     croped_pil_image = []
@@ -94,25 +96,24 @@ def get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_mode
             prompt = "<CAPTION>"
         else:
             prompt = "The image shows"
-
-    batch_size = 10  # Number of samples per batch
+    
     generated_texts = []
     device = model.device
-
     for i in range(0, len(croped_pil_image), batch_size):
+        start = time.time()
         batch = croped_pil_image[i:i+batch_size]
         if model.device.type == 'cuda':
             inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device, dtype=torch.float16)
         else:
             inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
         if 'florence' in model.config.name_or_path:
-            generated_ids = model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=1024,num_beams=3, do_sample=False)
+            generated_ids = model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=100,num_beams=3, do_sample=False)
         else:
             generated_ids = model.generate(**inputs, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True, num_return_sequences=1) # temperature=0.01, do_sample=True,
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
         generated_text = [gen.strip() for gen in generated_text]
         generated_texts.extend(generated_text)
-
+    
     return generated_texts
 
 
@@ -192,6 +193,12 @@ def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
             ratio1, ratio2 = 0, 0
         return max(intersection / union, ratio1, ratio2)
 
+    def is_inside(box1, box2):
+        # return box1[0] >= box2[0] and box1[1] >= box2[1] and box1[2] <= box2[2] and box1[3] <= box2[3]
+        intersection = intersection_area(box1, box2)
+        ratio1 = intersection / box_area(box1)
+        return ratio1 > 0.95
+
     boxes = boxes.tolist()
     filtered_boxes = []
     if ocr_bbox:
@@ -201,18 +208,104 @@ def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
         # if not any(IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2) for j, box2 in enumerate(boxes) if i != j):
         is_valid_box = True
         for j, box2 in enumerate(boxes):
+            # keep the smaller box
             if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
                 is_valid_box = False
                 break
         if is_valid_box:
             # add the following 2 lines to include ocr bbox
             if ocr_bbox:
-                if not any(IoU(box1, box3) > iou_threshold for k, box3 in enumerate(ocr_bbox)):
+                # only add the box if it does not overlap with any ocr bbox
+                if not any(IoU(box1, box3) > iou_threshold and not is_inside(box1, box3) for k, box3 in enumerate(ocr_bbox)):
                     filtered_boxes.append(box1)
             else:
                 filtered_boxes.append(box1)
     return torch.tensor(filtered_boxes)
 
+
+def remove_overlap_new(boxes, iou_threshold, ocr_bbox=None):
+    '''
+    ocr_bbox format: [{'type': 'text', 'bbox':[x,y], 'interactivity':False, 'content':str }, ...]
+    boxes format: [{'type': 'icon', 'bbox':[x,y], 'interactivity':True, 'content':None }, ...]
+
+    '''
+    assert ocr_bbox is None or isinstance(ocr_bbox, List)
+
+    def box_area(box):
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    def intersection_area(box1, box2):
+        x1 = max(box1[0], box2[0])
+        y1 = max(box1[1], box2[1])
+        x2 = min(box1[2], box2[2])
+        y2 = min(box1[3], box2[3])
+        return max(0, x2 - x1) * max(0, y2 - y1)
+
+    def IoU(box1, box2):
+        intersection = intersection_area(box1, box2)
+        union = box_area(box1) + box_area(box2) - intersection + 1e-6
+        if box_area(box1) > 0 and box_area(box2) > 0:
+            ratio1 = intersection / box_area(box1)
+            ratio2 = intersection / box_area(box2)
+        else:
+            ratio1, ratio2 = 0, 0
+        return max(intersection / union, ratio1, ratio2)
+
+    def is_inside(box1, box2):
+        # return box1[0] >= box2[0] and box1[1] >= box2[1] and box1[2] <= box2[2] and box1[3] <= box2[3]
+        intersection = intersection_area(box1, box2)
+        ratio1 = intersection / box_area(box1)
+        return ratio1 > 0.95
+
+    # boxes = boxes.tolist()
+    filtered_boxes = []
+    if ocr_bbox:
+        filtered_boxes.extend(ocr_bbox)
+    # print('ocr_bbox!!!', ocr_bbox)
+    for i, box1_elem in enumerate(boxes):
+        box1 = box1_elem['bbox']
+        is_valid_box = True
+        for j, box2_elem in enumerate(boxes):
+            # keep the smaller box
+            box2 = box2_elem['bbox']
+            if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
+                is_valid_box = False
+                break
+        if is_valid_box:
+            # add the following 2 lines to include ocr bbox
+            if ocr_bbox:
+                # only add the box if it does not overlap with any ocr bbox
+                box_added = False
+                for box3_elem in ocr_bbox:
+                    if not box_added:
+                        box3 = box3_elem['bbox']
+                        if is_inside(box3, box1): # ocr inside icon
+                            box_added = True
+                            # delete the box3_elem from ocr_bbox
+                            try:
+                                filtered_boxes.append({'type': 'text', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': box3_elem['content']})
+                                filtered_boxes.remove(box3_elem)
+                            except:
+                                continue
+                            break
+                        elif is_inside(box1, box3): # icon inside ocr
+                            box_added = True
+                            try:
+                                filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': None})
+                                filtered_boxes.remove(box3_elem)
+                            except:
+                                continue
+                            break
+                        else:
+                            continue
+                if not box_added:
+                    filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': None})
+                            
+            else:
+                filtered_boxes.append(box1)
+    return filtered_boxes # torch.tensor(filtered_boxes)
+
+
 def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
     transform = T.Compose(
         [
@@ -280,17 +373,23 @@ def predict(model, image, caption, box_threshold, text_threshold):
     return boxes, logits, phrases
 
 
-def predict_yolo(model, image_path, box_threshold, imgsz):
+def predict_yolo(model, image_path, box_threshold, imgsz, scale_img, iou_threshold=0.7):
     """ Use huggingface model to replace the original model
     """
     # model = model['model']
-    
-    result = model.predict(
-    source=image_path,
-    conf=box_threshold,
-    imgsz=imgsz
-    # iou=0.5, # default 0.7
-    )
+    if scale_img:
+        result = model.predict(
+        source=image_path,
+        conf=box_threshold,
+        imgsz=imgsz,
+        iou=iou_threshold, # default 0.7
+        )
+    else:
+        result = model.predict(
+        source=image_path,
+        conf=box_threshold,
+        iou=iou_threshold, # default 0.7
+        )
     boxes = result[0].boxes.xyxy#.tolist() # in pixel space
     conf = result[0].boxes.conf
     phrases = [str(i) for i in range(len(boxes))]
@@ -298,19 +397,15 @@ def predict_yolo(model, image_path, box_threshold, imgsz):
     return boxes, conf, phrases
 
 
-def get_som_labeled_img(img_path, 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,imgsz=640):
+def get_som_labeled_img(img_path, 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=None):
     """ ocr_bbox: list of xyxy format bbox
     """
-    TEXT_PROMPT = "clickable buttons on the screen"
-    # BOX_TRESHOLD = 0.02 # 0.05/0.02 for web and 0.1 for mobile
-    TEXT_TRESHOLD = 0.01 # 0.9 # 0.01
     image_source = Image.open(img_path).convert("RGB")
     w, h = image_source.size
-    # import pdb; pdb.set_trace()
-    if False: # TODO
-        xyxy, logits, phrases = predict(model=model, image=image_source, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD)
-    else:
-        xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD, imgsz=imgsz)
+    if not imgsz:
+        imgsz = (h, w)
+    # print('image size:', w, h)
+    xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD, imgsz=imgsz, scale_img=scale_img, iou_threshold=0.1)
     xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device)
     image_source = np.asarray(image_source)
     phrases = [str(i) for i in range(len(phrases))]
@@ -323,7 +418,20 @@ def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_
     else:
         print('no ocr bbox!!!')
         ocr_bbox = None
-    filtered_boxes = remove_overlap(boxes=xyxy, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox)
+    # filtered_boxes = remove_overlap(boxes=xyxy, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox)
+    # starting_idx = len(ocr_bbox)
+    # print('len(filtered_boxes):', len(filtered_boxes), starting_idx)
+
+    ocr_bbox_elem = [{'type': 'text', 'bbox':box, 'interactivity':False, 'content':txt} for box, txt in zip(ocr_bbox, ocr_text)]
+    xyxy_elem = [{'type': 'icon', 'bbox':box, 'interactivity':True, 'content':None} for box in xyxy.tolist()]
+    filtered_boxes = remove_overlap_new(boxes=xyxy_elem, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox_elem)
+    
+    # sort the filtered_boxes so that the one with 'content': None is at the end, and get the index of the first 'content': None
+    filtered_boxes_elem = sorted(filtered_boxes, key=lambda x: x['content'] is None)
+    # get the index of the first 'content': None
+    starting_idx = next((i for i, box in enumerate(filtered_boxes_elem) if box['content'] is None), -1)
+    filtered_boxes = torch.tensor([box['bbox'] for box in filtered_boxes_elem])
+
     
     # get parsed icon local semantics
     if use_local_semantics:
@@ -331,10 +439,14 @@ def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_
         if 'phi3_v' in caption_model.config.model_type: 
             parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor)
         else:
-            parsed_content_icon = get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=prompt)
+            parsed_content_icon = get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=prompt,batch_size=batch_size)
         ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
         icon_start = len(ocr_text)
         parsed_content_icon_ls = []
+        # fill the filtered_boxes_elem None content with parsed_content_icon in order
+        for i, box in enumerate(filtered_boxes_elem):
+            if box['content'] is None:
+                box['content'] = parsed_content_icon.pop(0)
         for i, txt in enumerate(parsed_content_icon):
             parsed_content_icon_ls.append(f"Icon Box ID {str(i+icon_start)}: {txt}")
         parsed_content_merged = ocr_text + parsed_content_icon_ls
@@ -361,7 +473,7 @@ def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_
         label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()}
         assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
 
-    return encoded_image, label_coordinates, parsed_content_merged
+    return encoded_image, label_coordinates, filtered_boxes_elem
 
 
 def get_xywh(input):
@@ -383,9 +495,14 @@ def get_xywh_yolo(input):
 
 def check_ocr_box(image_path, display_img = True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None, use_paddleocr=False):
     if use_paddleocr:
+        if easyocr_args is None:
+            text_threshold = 0.5
+        else:
+            text_threshold = easyocr_args['text_threshold']
         result = paddle_ocr.ocr(image_path, cls=False)[0]
-        coord = [item[0] for item in result]
-        text = [item[1][0] for item in result]
+        conf = [item[1] for item in result]
+        coord = [item[0] for item in result if item[1][1] > text_threshold]
+        text = [item[1][0] for item in result if item[1][1] > text_threshold]
     else:  # EasyOCR
         if easyocr_args is None:
             easyocr_args = {}

weights/convert_safetensor_to_pt.py

@@ -1,9 +1,27 @@
 import torch
 from ultralytics.nn.tasks import DetectionModel
 from safetensors.torch import load_file
+import argparse
+import yaml
 
-tensor_dict = load_file("weights/icon_detect/model.safetensors")
+# accept args to specify v1 or v1_5
+parser = argparse.ArgumentParser(description='Specify version v1 or v1_5')
+parser.add_argument('--version', choices=['v1', 'v1_5'], required=True, help='Specify the version: v1 or v1_5')
+args = parser.parse_args()
 
-model = DetectionModel('weights/icon_detect/model.yaml')
-model.load_state_dict(tensor_dict)
-torch.save({'model':model}, 'weights/icon_detect/best.pt')
+if args.version == 'v1':
+    tensor_dict = load_file("weights/icon_detect/model.safetensors")
+    model = DetectionModel('weights/icon_detect/model.yaml')
+    model.load_state_dict(tensor_dict)
+    torch.save({'model':model}, 'weights/icon_detect/best.pt')
+elif args.version == 'v1_5':
+    print("Converting v1_5")
+    tensor_dict = load_file("weights/icon_detect_v1_5/model.safetensors")
+    model = DetectionModel('weights/icon_detect_v1_5/model.yaml')
+    model.load_state_dict(tensor_dict)
+    save_dict = {'model':model}
+
+    with open("weights/icon_detect_v1_5/train_args.yaml", 'r') as file:
+        train_args = yaml.safe_load(file)
+    save_dict.update(train_args)
+    torch.save(save_dict, 'weights/icon_detect_v1_5/best.pt')