Update utils.py

utils.py

@@ -1,58 +1,49 @@
-"""
-utils.py
-
-This module contains utility functions for:
- - Loading and processing images
- - Object detection with YOLO
- - OCR with EasyOCR / PaddleOCR
- - Image annotation and bounding box manipulation
- - Captioning / semantic parsing of detected icons
-"""
-
+# from ultralytics import YOLO
 import os
 import io
 import base64
 import time
+from PIL import Image, ImageDraw, ImageFont
 import json
-import sys
-import re
-from typing import Tuple, List
+import requests
+# utility function
+import os
 
-import torch
-import numpy as np
+
+import json
+import sys
+import os
 import cv2
-from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+# %matplotlib inline
 from matplotlib import pyplot as plt
-
 import easyocr
 from paddleocr import PaddleOCR
-import supervision as sv
-import torchvision.transforms as T
-from torchvision.transforms import ToPILImage
-from torchvision.ops import box_convert
-
-# Optional: import AzureOpenAI if used
-from openai import AzureOpenAI
-
-# Initialize OCR readers
 reader = easyocr.Reader(['en'])
 paddle_ocr = PaddleOCR(
-    lang='en',  # other languages available
+    lang='en',  # other lang also available
     use_angle_cls=False,
-    use_gpu=False,  # using cuda might conflict with PyTorch in the same process
+    use_gpu=False,  # using cuda will conflict with pytorch in the same process
     show_log=False,
     max_batch_size=1024,
     use_dilation=True,  # improves accuracy
     det_db_score_mode='slow',  # improves accuracy
-    rec_batch_num=1024
-)
+    rec_batch_num=1024)
+import time
+import base64
+
+import os
+import ast
+import torch
+from typing import Tuple, List
+from torchvision.ops import box_convert
+import re
+from torchvision.transforms import ToPILImage
+import supervision as sv
+import torchvision.transforms as T
 
 
 def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
-    """
-    Loads the captioning model and processor.
-    Supports either BLIP2 or Florence-2 models.
-    """
     if not device:
         device = "cuda" if torch.cuda.is_available() else "cpu"
     if model_name == "blip2":
@@ -60,53 +51,45 @@ def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2
         processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
         if device == 'cpu':
             model = Blip2ForConditionalGeneration.from_pretrained(
-                model_name_or_path, device_map=None, torch_dtype=torch.float32
-            )
+            model_name_or_path, device_map=None, torch_dtype=torch.float32
+        ) 
         else:
             model = Blip2ForConditionalGeneration.from_pretrained(
-                model_name_or_path, device_map=None, torch_dtype=torch.float16
-            ).to(device)
+            model_name_or_path, device_map=None, torch_dtype=torch.float16
+        ).to(device)
     elif model_name == "florence2":
-        from transformers import AutoProcessor, AutoModelForCausalLM
+        from transformers import AutoProcessor, AutoModelForCausalLM 
         processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
         if device == 'cpu':
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True
-            )
+            model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True)
         else:
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True
-            ).to(device)
+            model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True).to(device)
     return {'model': model.to(device), 'processor': processor}
 
 
 def get_yolo_model(model_path):
-    """
-    Loads a YOLO model from a given model_path using ultralytics.
-    """
     from ultralytics import YOLO
+    # Load the model.
     model = YOLO(model_path)
     return model
 
 
 @torch.inference_mode()
 def get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=None, batch_size=32):
-    # Ensure batch_size is an integer
-    if batch_size is None:
-        batch_size = 32
+    # Number of samples per batch, --> 256 roughly takes 23 GB of GPU memory for florence model
 
     to_pil = ToPILImage()
     if starting_idx:
         non_ocr_boxes = filtered_boxes[starting_idx:]
     else:
         non_ocr_boxes = filtered_boxes
-    cropped_pil_images = []
-    for coord in non_ocr_boxes:
-        xmin, xmax = int(coord[0] * image_source.shape[1]), int(coord[2] * image_source.shape[1])
-        ymin, ymax = int(coord[1] * image_source.shape[0]), int(coord[3] * image_source.shape[0])
+    croped_pil_image = []
+    for i, coord in enumerate(non_ocr_boxes):
+        xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
+        ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
         cropped_image = image_source[ymin:ymax, xmin:xmax, :]
-        cropped_pil_images.append(to_pil(cropped_image))
-    
+        croped_pil_image.append(to_pil(cropped_image))
+
     model, processor = caption_model_processor['model'], caption_model_processor['processor']
     if not prompt:
         if 'florence' in model.config.name_or_path:
@@ -116,29 +99,17 @@ def get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_
     
     generated_texts = []
     device = model.device
-    for i in range(0, len(cropped_pil_images), batch_size):
-        batch = cropped_pil_images[i:i + batch_size]
+    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)
+            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)
+            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=100,
-                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
-            )
+            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)
@@ -147,57 +118,51 @@ def get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_
 
 
 
-
 def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor):
-    """
-    Generates parsed textual content for detected icons using the phi3_v model variant.
-    """
     to_pil = ToPILImage()
     if ocr_bbox:
         non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
     else:
         non_ocr_boxes = filtered_boxes
-    cropped_pil_images = []
-    for coord in non_ocr_boxes:
-        xmin, xmax = int(coord[0] * image_source.shape[1]), int(coord[2] * image_source.shape[1])
-        ymin, ymax = int(coord[1] * image_source.shape[0]), int(coord[3] * image_source.shape[0])
+    croped_pil_image = []
+    for i, coord in enumerate(non_ocr_boxes):
+        xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
+        ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
         cropped_image = image_source[ymin:ymax, xmin:xmax, :]
-        cropped_pil_images.append(to_pil(cropped_image))
+        croped_pil_image.append(to_pil(cropped_image))
 
     model, processor = caption_model_processor['model'], caption_model_processor['processor']
     device = model.device
-    messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}]
+    messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}] 
     prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
     batch_size = 5  # Number of samples per batch
     generated_texts = []
 
-    for i in range(0, len(cropped_pil_images), batch_size):
-        images = cropped_pil_images[i:i+batch_size]
+    for i in range(0, len(croped_pil_image), batch_size):
+        images = croped_pil_image[i:i+batch_size]
         image_inputs = [processor.image_processor(x, return_tensors="pt") for x in images]
-        inputs = {'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
+        inputs ={'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
         texts = [prompt] * len(images)
-        for idx, txt in enumerate(texts):
-            inp = processor._convert_images_texts_to_inputs(image_inputs[idx], txt, return_tensors="pt")
-            inputs['input_ids'].append(inp['input_ids'])
-            inputs['attention_mask'].append(inp['attention_mask'])
-            inputs['pixel_values'].append(inp['pixel_values'])
-            inputs['image_sizes'].append(inp['image_sizes'])
-        max_len = max(x.shape[1] for x in inputs['input_ids'])
-        for idx, v in enumerate(inputs['input_ids']):
-            pad_tensor = processor.tokenizer.pad_token_id * torch.ones(1, max_len - v.shape[1], dtype=torch.long)
-            inputs['input_ids'][idx] = torch.cat([pad_tensor, v], dim=1)
-            pad_att = torch.zeros(1, max_len - v.shape[1], dtype=torch.long)
-            inputs['attention_mask'][idx] = torch.cat([pad_att, inputs['attention_mask'][idx]], dim=1)
+        for i, txt in enumerate(texts):
+            input = processor._convert_images_texts_to_inputs(image_inputs[i], txt, return_tensors="pt")
+            inputs['input_ids'].append(input['input_ids'])
+            inputs['attention_mask'].append(input['attention_mask'])
+            inputs['pixel_values'].append(input['pixel_values'])
+            inputs['image_sizes'].append(input['image_sizes'])
+        max_len = max([x.shape[1] for x in inputs['input_ids']])
+        for i, v in enumerate(inputs['input_ids']):
+            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)
+            inputs['attention_mask'][i] = torch.cat([torch.zeros(1, max_len - v.shape[1], dtype=torch.long), inputs['attention_mask'][i]], dim=1)
         inputs_cat = {k: torch.concatenate(v).to(device) for k, v in inputs.items()}
 
-        generation_args = {
-            "max_new_tokens": 25,
-            "temperature": 0.01,
-            "do_sample": False,
-        }
-        generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
-        # Remove input tokens from the generated sequence
+        generation_args = { 
+            "max_new_tokens": 25, 
+            "temperature": 0.01, 
+            "do_sample": False, 
+        } 
+        generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args) 
+        # # remove input tokens 
         generate_ids = generate_ids[:, inputs_cat['input_ids'].shape[1]:]
         response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
         response = [res.strip('\n').strip() for res in response]
@@ -205,19 +170,7 @@ def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, captio
 
     return generated_texts
 
-
 def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
-    """
-    Removes overlapping bounding boxes based on IoU and optionally considers OCR boxes.
-    
-    Args:
-        boxes: Tensor of bounding boxes (in xyxy format).
-        iou_threshold: IoU threshold to determine overlaps.
-        ocr_bbox: Optional list of OCR bounding boxes.
-    
-    Returns:
-        Filtered boxes as a torch.Tensor.
-    """
     assert ocr_bbox is None or isinstance(ocr_bbox, List)
 
     def box_area(box):
@@ -231,30 +184,39 @@ def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
         return max(0, x2 - x1) * max(0, y2 - y1)
 
     def IoU(box1, box2):
-        inter = intersection_area(box1, box2)
-        union = box_area(box1) + box_area(box2) - inter + 1e-6
-        ratio1 = inter / box_area(box1) if box_area(box1) > 0 else 0
-        ratio2 = inter / box_area(box2) if box_area(box2) > 0 else 0
-        return max(inter / union, ratio1, ratio2)
+        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):
-        inter = intersection_area(box1, box2)
-        return (inter / box_area(box1)) > 0.95
+        # 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 in enumerate(boxes):
+        # 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:
-                # Only add the box if it does not overlap with any OCR box
-                if not any(IoU(box1, box3) > iou_threshold and not is_inside(box1, box3) for box3 in 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)
@@ -262,17 +224,11 @@ def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
 
 
 def remove_overlap_new(boxes, iou_threshold, ocr_bbox=None):
-    """
-    Removes overlapping boxes with OCR priority.
-    
-    Args:
-        boxes: List of dictionaries, each with keys: 'type', 'bbox', 'interactivity', 'content'.
-        iou_threshold: IoU threshold for removal.
-        ocr_bbox: List of OCR box dictionaries.
-    
-    Returns:
-        A list of filtered box dictionaries.
-    """
+    '''
+    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):
@@ -286,130 +242,132 @@ def remove_overlap_new(boxes, iou_threshold, ocr_bbox=None):
         return max(0, x2 - x1) * max(0, y2 - y1)
 
     def IoU(box1, box2):
-        inter = intersection_area(box1, box2)
-        union = box_area(box1) + box_area(box2) - inter + 1e-6
-        ratio1 = inter / box_area(box1) if box_area(box1) > 0 else 0
-        ratio2 = inter / box_area(box2) if box_area(box2) > 0 else 0
-        return max(inter / union, ratio1, ratio2)
+        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):
-        inter = intersection_area(box1, box2)
-        return (inter / box_area(box1)) > 0.80
+        # 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.80
 
+    # 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:
+                # keep yolo boxes + prioritize ocr label
                 box_added = False
                 for box3_elem in ocr_bbox:
-                    box3 = box3_elem['bbox']
-                    if is_inside(box3, box1):
-                        try:
-                            filtered_boxes.append({
-                                'type': 'text',
-                                'bbox': box1_elem['bbox'],
-                                'interactivity': True,
-                                'content': box3_elem['content']
-                            })
-                            filtered_boxes.remove(box3_elem)
-                        except Exception:
+                    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)
+                                # print('remove ocr bbox:', 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
-                    elif is_inside(box1, box3):
-                        box_added = True
-                        break
                 if not box_added:
-                    filtered_boxes.append({
-                        'type': 'icon',
-                        'bbox': box1_elem['bbox'],
-                        'interactivity': True,
-                        'content': None
-                    })
+                    filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': None})
+                            
             else:
                 filtered_boxes.append(box1)
-    return filtered_boxes  # Optionally, you could return torch.tensor(filtered_boxes) if needed
+    return filtered_boxes # torch.tensor(filtered_boxes)
 
 
 def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
-    """
-    Loads an image and applies transformations.
-    
-    Returns:
-        image: Original image as a NumPy array.
-        image_transformed: Transformed tensor.
-    """
-    transform = T.Compose([
-        T.RandomResize([800], max_size=1333),
-        T.ToTensor(),
-        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ])
+    transform = T.Compose(
+        [
+            T.RandomResize([800], max_size=1333),
+            T.ToTensor(),
+            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
     image_source = Image.open(image_path).convert("RGB")
     image = np.asarray(image_source)
     image_transformed, _ = transform(image_source, None)
     return image, image_transformed
 
 
-def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str],
-             text_scale: float, text_padding=5, text_thickness=2, thickness=3) -> Tuple[np.ndarray, dict]:
-    """
-    Annotates an image with bounding boxes and labels.
-    """
-    # Validate phrases input
-    phrases = [str(phrase) if not isinstance(phrase, str) else phrase for phrase in phrases]
+def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], text_scale: float, 
+             text_padding=5, text_thickness=2, thickness=3) -> np.ndarray:
+    """    
+    This function annotates an image with bounding boxes and labels.
 
+    Parameters:
+    image_source (np.ndarray): The source image to be annotated.
+    boxes (torch.Tensor): A tensor containing bounding box coordinates. in cxcywh format, pixel scale
+    logits (torch.Tensor): A tensor containing confidence scores for each bounding box.
+    phrases (List[str]): A list of labels for each bounding box.
+    text_scale (float): The scale of the text to be displayed. 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
+
+    Returns:
+    np.ndarray: The annotated image.
+    """
     h, w, _ = image_source.shape
     boxes = boxes * torch.Tensor([w, h, w, h])
     xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
     xywh = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xywh").numpy()
     detections = sv.Detections(xyxy=xyxy)
 
-    labels = [f"{phrase}" for phrase in phrases]
+    labels = [f"{phrase}" for phrase in range(boxes.shape[0])]
 
-    from util.box_annotator import BoxAnnotator
-    box_annotator = BoxAnnotator(text_scale=text_scale, text_padding=text_padding,
-                                 text_thickness=text_thickness, thickness=thickness)
+    from util.box_annotator import BoxAnnotator 
+    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
     annotated_frame = image_source.copy()
-    annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w, h))
+    annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w,h))
 
     label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
     return annotated_frame, label_coordinates
 
 
-
 def predict(model, image, caption, box_threshold, text_threshold):
+    """ Use huggingface model to replace the original model
     """
-    Uses a Hugging Face model to perform grounded object detection.
-    
-    Args:
-        model: Dictionary with 'model' and 'processor'.
-        image: Input PIL image.
-        caption: Caption text.
-        box_threshold: Confidence threshold for boxes.
-        text_threshold: Threshold for text detection.
-    
-    Returns:
-        boxes, logits, phrases from the detection.
-    """
-    model_obj, processor = model['model'], model['processor']
-    device = model_obj.device
+    model, processor = model['model'], model['processor']
+    device = model.device
 
     inputs = processor(images=image, text=caption, return_tensors="pt").to(device)
     with torch.no_grad():
-        outputs = model_obj(**inputs)
+        outputs = model(**inputs)
 
     results = processor.post_process_grounded_object_detection(
         outputs,
         inputs.input_ids,
-        box_threshold=box_threshold,
-        text_threshold=text_threshold,
+        box_threshold=box_threshold, # 0.4,
+        text_threshold=text_threshold, # 0.3,
         target_sizes=[image.size[::-1]]
     )[0]
     boxes, logits, phrases = results["boxes"], results["scores"], results["labels"]
@@ -417,109 +375,78 @@ def predict(model, image, caption, box_threshold, text_threshold):
 
 
 def predict_yolo(model, image_path, box_threshold, imgsz, scale_img, iou_threshold=0.7):
+    """ Use huggingface model to replace the original model
     """
-    Uses a YOLO model for object detection.
-    
-    Args:
-        model: YOLO model instance.
-        image_path: Path to the image.
-        box_threshold: Confidence threshold.
-        imgsz: Image size for scaling (if scale_img is True).
-        scale_img: Boolean flag to scale the image.
-        iou_threshold: IoU threshold for non-max suppression.
-    
-    Returns:
-        Bounding boxes, confidence scores, and placeholder phrases.
-    """
-    kwargs = {
-        'conf': box_threshold,  # Confidence threshold
-        'iou': iou_threshold,   # IoU threshold
-        'verbose': False
-    }
+    # model = model['model']
     if scale_img:
-        kwargs['imgsz'] = imgsz
+        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))]
 
-    results = model.predict(image_path, **kwargs)
-    boxes = results[0].boxes.xyxy
-    conf = results[0].boxes.conf
-    return boxes, conf, [str(i) for i in range(len(boxes))]
+    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, scale_img=False,
-                        imgsz=None, batch_size=None):
-    """
-    Processes an image to generate semantic (SOM) labels.
-    
-    Args:
-        img_path: Path to the image.
-        model: YOLO model for detection.
-        BOX_TRESHOLD: Confidence threshold for box prediction.
-        output_coord_in_ratio: If True, output coordinates in ratio.
-        ocr_bbox: OCR bounding boxes.
-        text_scale, text_padding: Parameters for drawing annotations.
-        draw_bbox_config: Custom configuration for bounding box drawing.
-        caption_model_processor: Dictionary with caption model and processor.
-        ocr_text: List of OCR-detected texts.
-        use_local_semantics: Whether to use local semantic processing.
-        iou_threshold: IoU threshold for filtering overlaps.
-        prompt: Optional caption prompt.
-        scale_img: Whether to scale the image.
-        imgsz: Image size for YOLO.
-        batch_size: Batch size for captioning.
-    
-    Returns:
-        Encoded annotated image, label coordinates, and filtered boxes.
+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
     """
     image_source = Image.open(img_path).convert("RGB")
     w, h = image_source.size
     if not imgsz:
         imgsz = (h, w)
-    # Run YOLO detection
-    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
-    )
+    # 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 = np.asarray(image_source)
+    image_source = np.asarray(image_source)
     phrases = [str(i) for i in range(len(phrases))]
 
-    # Process OCR bounding boxes (if any)
+    # annotate the image with labels
+    h, w, _ = image_source.shape
     if ocr_bbox:
         ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
-        ocr_bbox = ocr_bbox.tolist()
+        ocr_bbox=ocr_bbox.tolist()
     else:
         print('no ocr bbox!!!')
         ocr_bbox = None
+    # 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()]
+    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 filtered boxes so that boxes with 'content' == None are at the end
+    # 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_tensor = torch.tensor([box['bbox'] for box in filtered_boxes_elem])
+    filtered_boxes = torch.tensor([box['bbox'] for box in filtered_boxes_elem])
 
-    if batch_size is None:
-        batch_size = 32
-
-    # Generate parsed icon semantics if required
+    
+    # get parsed icon local semantics
     if use_local_semantics:
         caption_model = caption_model_processor['model']
-        if 'phi3_v' in caption_model.config.model_type:
-            parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes_tensor, ocr_bbox, image_source_np, caption_model_processor)
+        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_tensor, starting_idx, image_source_np, caption_model_processor, prompt=prompt, batch_size=batch_size)
+            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 boxes with no OCR content with parsed icon content
-        for box in filtered_boxes_elem:
-            if box['content'] is None and parsed_content_icon:
+        # 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}")
@@ -528,72 +455,51 @@ def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD=0.01, output_coord_in
         ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
         parsed_content_merged = ocr_text
 
-    filtered_boxes_cxcywh = box_convert(boxes=filtered_boxes_tensor, in_fmt="xyxy", out_fmt="cxcywh")
-    phrases = [i for i in range(len(filtered_boxes_cxcywh))]
+    filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
+
+    phrases = [i for i in range(len(filtered_boxes))]
     
-    # Annotate image with bounding boxes and labels
+    # draw boxes
     if draw_bbox_config:
-        annotated_frame, label_coordinates = annotate(
-            image_source=image_source_np, boxes=filtered_boxes_cxcywh, logits=logits, phrases=phrases, **draw_bbox_config
-        )
+        annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
     else:
-        annotated_frame, label_coordinates = annotate(
-            image_source=image_source_np, boxes=filtered_boxes_cxcywh, logits=logits, phrases=phrases,
-            text_scale=text_scale, text_padding=text_padding
-        )
+        annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding)
     
     pil_img = Image.fromarray(annotated_frame)
     buffered = io.BytesIO()
     pil_img.save(buffered, format="PNG")
     encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
-    
     if output_coord_in_ratio:
-        label_coordinates = {k: [v[0] / w, v[1] / h, v[2] / w, v[3] / h] for k, v in label_coordinates.items()}
+        # h, w, _ = image_source.shape
+        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, filtered_boxes_elem
 
 
 def get_xywh(input):
-    """
-    Converts a bounding box from a list of two points into (x, y, width, height).
-    """
-    x, y = input[0][0], input[0][1]
-    w = input[2][0] - input[0][0]
-    h = input[2][1] - input[0][1]
-    return int(x), int(y), int(w), int(h)
-
+    x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1]
+    x, y, w, h = int(x), int(y), int(w), int(h)
+    return x, y, w, h
 
 def get_xyxy(input):
-    """
-    Converts a bounding box from a list of two points into (x, y, x2, y2).
-    """
-    x, y = input[0][0], input[0][1]
-    x2, y2 = input[2][0], input[2][1]
-    return int(x), int(y), int(x2), int(y2)
-
+    x, y, xp, yp = input[0][0], input[0][1], input[2][0], input[2][1]
+    x, y, xp, yp = int(x), int(y), int(xp), int(yp)
+    return x, y, xp, yp
 
 def get_xywh_yolo(input):
-    """
-    Converts a YOLO-style bounding box (x1, y1, x2, y2) into (x, y, width, height).
-    """
-    x, y = input[0], input[1]
-    w = input[2] - input[0]
-    h = input[3] - input[1]
-    return int(x), int(y), int(w), int(h)
+    x, y, w, h = input[0], input[1], input[2] - input[0], input[3] - input[1]
+    x, y, w, h = int(x), int(y), int(w), int(h)
+    return x, y, w, h
+    
 
 
-def check_ocr_box(image_path, display_img=True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None, use_paddleocr=False):
-    """
-    Runs OCR on the given image using PaddleOCR or EasyOCR and optionally displays annotated results.
-    
-    Returns:
-        A tuple containing:
-            - A tuple (text, bounding boxes)
-            - The goal_filtering parameter (unchanged)
-    """
+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:
-        text_threshold = 0.5 if easyocr_args is None else easyocr_args.get('text_threshold', 0.5)
+        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]
         conf = [item[1] for item in result]
         coord = [item[0] for item in result if item[1][1] > text_threshold]
@@ -602,21 +508,26 @@ def check_ocr_box(image_path, display_img=True, output_bb_format='xywh', goal_fi
         if easyocr_args is None:
             easyocr_args = {}
         result = reader.readtext(image_path, **easyocr_args)
+        # print('goal filtering pred:', result[-5:])
         coord = [item[0] for item in result]
         text = [item[1] for item in result]
-    
+    # read the image using cv2
     if display_img:
         opencv_img = cv2.imread(image_path)
         opencv_img = cv2.cvtColor(opencv_img, cv2.COLOR_RGB2BGR)
         bb = []
         for item in coord:
             x, y, a, b = get_xywh(item)
+            # print(x, y, a, b)
             bb.append((x, y, a, b))
-            cv2.rectangle(opencv_img, (x, y), (x + a, y + b), (0, 255, 0), 2)
+            cv2.rectangle(opencv_img, (x, y), (x+a, y+b), (0, 255, 0), 2)
+        
+        # Display the image
         plt.imshow(opencv_img)
     else:
         if output_bb_format == 'xywh':
             bb = [get_xywh(item) for item in coord]
         elif output_bb_format == 'xyxy':
             bb = [get_xyxy(item) for item in coord]
-    return (text, bb), goal_filtering
+        # print('bounding box!!!', bb)
+    return (text, bb), goal_filtering