Transformers

box_annotator.py

@@ -0,0 +1,262 @@
+from typing import List, Optional, Union, Tuple
+
+import cv2
+import numpy as np
+
+from supervision.detection.core import Detections
+from supervision.draw.color import Color, ColorPalette
+
+
+class BoxAnnotator:
+    """
+    A class for drawing bounding boxes on an image using detections provided.
+
+    Attributes:
+        color (Union[Color, ColorPalette]): The color to draw the bounding box,
+            can be a single color or a color palette
+        thickness (int): The thickness of the bounding box lines, default is 2
+        text_color (Color): The color of the text on the bounding box, default is white
+        text_scale (float): The scale of the text on the bounding box, default is 0.5
+        text_thickness (int): The thickness of the text on the bounding box,
+            default is 1
+        text_padding (int): The padding around the text on the bounding box,
+            default is 5
+
+    """
+
+    def __init__(
+        self,
+        color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,
+        thickness: int = 3, # 1 for seeclick 2 for mind2web and 3 for demo
+        text_color: Color = Color.BLACK,
+        text_scale: float = 0.5, # 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
+        text_thickness: int = 2, #1, # 2 for demo
+        text_padding: int = 10,
+        avoid_overlap: bool = True,
+    ):
+        self.color: Union[Color, ColorPalette] = color
+        self.thickness: int = thickness
+        self.text_color: Color = text_color
+        self.text_scale: float = text_scale
+        self.text_thickness: int = text_thickness
+        self.text_padding: int = text_padding
+        self.avoid_overlap: bool = avoid_overlap
+
+    def annotate(
+        self,
+        scene: np.ndarray,
+        detections: Detections,
+        labels: Optional[List[str]] = None,
+        skip_label: bool = False,
+        image_size: Optional[Tuple[int, int]] = None,
+    ) -> np.ndarray:
+        """
+        Draws bounding boxes on the frame using the detections provided.
+
+        Args:
+            scene (np.ndarray): The image on which the bounding boxes will be drawn
+            detections (Detections): The detections for which the
+                bounding boxes will be drawn
+            labels (Optional[List[str]]): An optional list of labels
+                corresponding to each detection. If `labels` are not provided,
+                corresponding `class_id` will be used as label.
+            skip_label (bool): Is set to `True`, skips bounding box label annotation.
+        Returns:
+            np.ndarray: The image with the bounding boxes drawn on it
+
+        Example:
+            ```python
+            import supervision as sv
+
+            classes = ['person', ...]
+            image = ...
+            detections = sv.Detections(...)
+
+            box_annotator = sv.BoxAnnotator()
+            labels = [
+                f"{classes[class_id]} {confidence:0.2f}"
+                for _, _, confidence, class_id, _ in detections
+            ]
+            annotated_frame = box_annotator.annotate(
+                scene=image.copy(),
+                detections=detections,
+                labels=labels
+            )
+            ```
+        """
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        for i in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[i].astype(int)
+            class_id = (
+                detections.class_id[i] if detections.class_id is not None else None
+            )
+            idx = class_id if class_id is not None else i
+            color = (
+                self.color.by_idx(idx)
+                if isinstance(self.color, ColorPalette)
+                else self.color
+            )
+            cv2.rectangle(
+                img=scene,
+                pt1=(x1, y1),
+                pt2=(x2, y2),
+                color=color.as_bgr(),
+                thickness=self.thickness,
+            )
+            if skip_label:
+                continue
+
+            text = (
+                f"{class_id}"
+                if (labels is None or len(detections) != len(labels))
+                else labels[i]
+            )
+
+            text_width, text_height = cv2.getTextSize(
+                text=text,
+                fontFace=font,
+                fontScale=self.text_scale,
+                thickness=self.text_thickness,
+            )[0]
+
+            if not self.avoid_overlap:
+                text_x = x1 + self.text_padding
+                text_y = y1 - self.text_padding
+
+                text_background_x1 = x1
+                text_background_y1 = y1 - 2 * self.text_padding - text_height
+
+                text_background_x2 = x1 + 2 * self.text_padding + text_width
+                text_background_y2 = y1
+                # text_x = x1 - self.text_padding - text_width
+                # text_y = y1 + self.text_padding + text_height
+                # text_background_x1 = x1 - 2 * self.text_padding - text_width
+                # text_background_y1 = y1
+                # text_background_x2 = x1
+                # text_background_y2 = y1 + 2 * self.text_padding + text_height
+            else:
+                text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2 = get_optimal_label_pos(self.text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size)
+
+            cv2.rectangle(
+                img=scene,
+                pt1=(text_background_x1, text_background_y1),
+                pt2=(text_background_x2, text_background_y2),
+                color=color.as_bgr(),
+                thickness=cv2.FILLED,
+            )
+            # import pdb; pdb.set_trace()
+            box_color = color.as_rgb()
+            luminance = 0.299 * box_color[0] + 0.587 * box_color[1] + 0.114 * box_color[2]
+            text_color = (0,0,0) if luminance > 160 else (255,255,255)
+            cv2.putText(
+                img=scene,
+                text=text,
+                org=(text_x, text_y),
+                fontFace=font,
+                fontScale=self.text_scale,
+                # color=self.text_color.as_rgb(),
+                color=text_color,
+                thickness=self.text_thickness,
+                lineType=cv2.LINE_AA,
+            )
+        return scene
+    
+
+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, return_max=True):
+    intersection = intersection_area(box1, box2)
+    union = box_area(box1) + box_area(box2) - intersection
+    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
+    if return_max:
+        return max(intersection / union, ratio1, ratio2)
+    else:
+        return intersection / union
+
+
+def get_optimal_label_pos(text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size):
+    """ check overlap of text and background detection box, and get_optimal_label_pos, 
+        pos: str, position of the text, must be one of 'top left', 'top right', 'outer left', 'outer right' TODO: if all are overlapping, return the last one, i.e. outer right
+        Threshold: default to 0.3
+    """
+
+    def get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size):
+        is_overlap = False
+        for i in range(len(detections)):
+            detection = detections.xyxy[i].astype(int)
+            if IoU([text_background_x1, text_background_y1, text_background_x2, text_background_y2], detection) > 0.3:
+                is_overlap = True
+                break
+        # check if the text is out of the image
+        if text_background_x1 < 0 or text_background_x2 > image_size[0] or text_background_y1 < 0 or text_background_y2 > image_size[1]:
+            is_overlap = True
+        return is_overlap
+    
+    # if pos == 'top left':
+    text_x = x1 + text_padding
+    text_y = y1 - text_padding
+
+    text_background_x1 = x1
+    text_background_y1 = y1 - 2 * text_padding - text_height
+
+    text_background_x2 = x1 + 2 * text_padding + text_width
+    text_background_y2 = y1
+    is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
+    if not is_overlap:
+        return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
+    
+    # elif pos == 'outer left':
+    text_x = x1 - text_padding - text_width
+    text_y = y1 + text_padding + text_height
+
+    text_background_x1 = x1 - 2 * text_padding - text_width
+    text_background_y1 = y1
+
+    text_background_x2 = x1
+    text_background_y2 = y1 + 2 * text_padding + text_height
+    is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
+    if not is_overlap:
+        return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
+    
+
+    # elif pos == 'outer right':
+    text_x = x2 + text_padding
+    text_y = y1 + text_padding + text_height
+
+    text_background_x1 = x2
+    text_background_y1 = y1
+
+    text_background_x2 = x2 + 2 * text_padding + text_width
+    text_background_y2 = y1 + 2 * text_padding + text_height
+
+    is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
+    if not is_overlap:
+        return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
+
+    # elif pos == 'top right':
+    text_x = x2 - text_padding - text_width
+    text_y = y1 - text_padding
+
+    text_background_x1 = x2 - 2 * text_padding - text_width
+    text_background_y1 = y1 - 2 * text_padding - text_height
+
+    text_background_x2 = x2
+    text_background_y2 = y1
+
+    is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
+    if not is_overlap:
+        return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
+
+    return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2

config.json

@@ -0,0 +1,5 @@
+{
+  "box_threshold": 0.05,
+  "model_type": "omniparserv2",
+  "transformers_version": "4.48.1"
+}

configuration_omniparser.py

@@ -0,0 +1,13 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class OmniparserConfig(PretrainedConfig):
+    model_type = "omniparserv2"
+
+    def __init__(
+        self,
+        box_threshold: float = 0.05,
+        **kwargs,
+    ):
+        self.box_threshold = box_threshold
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bbcbf73561f6bc5d0a17ea6a2081feed2d1304e87602d8c502d9a5c4bd85576
+size 16

modelling_omniparser.py

@@ -0,0 +1,351 @@
+from transformers import PreTrainedModel
+from transformers import AutoProcessor, AutoModelForCausalLM 
+from .configuration_omniparser import OmniparserConfig
+from ultralytics import YOLO
+from torchvision.transforms import ToPILImage
+from torchvision.ops import box_convert
+import base64
+import supervision as sv
+from .box_annotator import BoxAnnotator 
+import os
+import torch
+import numpy as np
+from PIL import Image
+import io
+import cv2
+import easyocr
+from typing import List, Dict
+from dataclasses import dataclass
+from transformers.modeling_outputs import ModelOutput
+import json
+
+@dataclass
+class OmniparserOutput(ModelOutput):
+    """
+    Output type of [`OmniparserModel`].
+
+    Args:
+        annotated_image (`str`):
+            The image with bounding boxes and labels drawn on it, returned as base64 encoded string.
+        parsed_content_list (`List[Dict]`):
+            Lists of detected elements with their properties (type, bbox, interactivity, content, source).
+    """
+    annotated_image: str = None
+    parsed_content_list: List[Dict] = None
+
+
+class OmniparserModel(PreTrainedModel):
+    config_class = OmniparserConfig
+    _keys_to_ignore_on_load_missing = [r"caption_model\..*"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._device_type = 'cuda' if torch.cuda.is_available() else 'cpu'
+        current_file_directory = os.path.dirname(os.path.abspath(__file__))
+        self.icon_detect_path = os.path.join(current_file_directory, 'icon_detect/model.pt')
+        icon_caption_path = os.path.join(current_file_directory, 'icon_caption')
+        caption_model_processor = self.get_caption_model_processor(model_name_or_path=icon_caption_path)
+        self.caption_model, self.caption_model_processor = caption_model_processor['model'], caption_model_processor['processor']
+        self.reader = easyocr.Reader(['en'])
+        self.to_pil = ToPILImage()
+        self.som_model_inited = False
+
+    def forward(self, image_source):
+        # Lazy load YOLO model as autoload calls .eval which breaks YOLO unless init in forward
+        if not self.som_model_inited:
+            self.som_model = YOLO(self.icon_detect_path)
+            self.som_model_inited = True
+
+        if image_source is None:
+            raise ValueError("No image provided")
+        if isinstance(image_source, str):
+            image_source = Image.open(image_source)
+        image_source = image_source.convert("RGB")
+
+        # Scan for text
+        ocr_text, ocr_bbox = self.check_ocr_box(image_source, easyocr_args={'text_threshold': 0.8})
+        # Yolo + merge text and YOLO + florence
+        annotated_image, parsed_content_list = self.get_som_labeled_img(image_source, 
+                                                                        self.som_model,
+                                                                        box_threshold=self.config.box_threshold,
+                                                                        ocr_bbox=ocr_bbox,
+                                                                        ocr_text=ocr_text)
+
+        return OmniparserOutput(
+            annotated_image=annotated_image,
+            parsed_content_list=parsed_content_list
+        )
+
+    def get_caption_model_processor(self, model_name_or_path):
+        processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
+        if self._device_type == 'cpu':
+            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(self._device_type)
+        return {'model': model.to(self._device_type), 'processor': processor}
+
+    def check_ocr_box(self, image_source, easyocr_args):
+        image_np = np.array(image_source)
+        w, h = image_source.size
+        result = self.reader.readtext(image_np, **easyocr_args)
+        coord = [item[0] for item in result]
+        text = [item[1] for item in result]
+        bb = [self.get_xyxy(item) for item in coord]
+        return text, bb
+
+    def get_som_labeled_img(self, image_source, 
+                           model, 
+                           box_threshold,
+                           ocr_bbox, 
+                           ocr_text,
+                           iou_threshold=0.7,
+                           prompt=None,
+                           batch_size=128):
+        w, h = image_source.size
+        box_overlay_ratio = max(w, h) / 3200
+        draw_bbox_config = {
+            'text_scale': 0.8 * box_overlay_ratio,
+            'text_thickness': max(int(2 * box_overlay_ratio), 1),
+            'text_padding': max(int(3 * box_overlay_ratio), 1),
+            'thickness': max(int(3 * box_overlay_ratio), 1),
+        }
+        xyxy, logits, phrases = self.predict_yolo(yolo_model=model, image=image_source, box_threshold=box_threshold, 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))]
+
+        # merge text and YOLO boxes
+        ocr_bbox_elem = None
+        if ocr_bbox:
+            ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
+            ocr_bbox=ocr_bbox.tolist()
+            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 self.int_box_area(box, w, h) > 0] 
+        xyxy_elem = [{'type': 'icon', 'bbox': box, 'interactivity': True, 'content': None, 'source': 'box_yolo_content_yolo'} for box in xyxy.tolist() if self.int_box_area(box, w, h) > 0]
+        filtered_boxes = self.remove_overlap_new(boxes=xyxy_elem, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox_elem)
+        
+        # sort the filtered_boxes so that ocr first, yolo last
+        filtered_boxes_elem = sorted(filtered_boxes, key=lambda x: x['content'] is None)
+        # get the index of the first yolo box
+        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 local semantics using florence on yolo boxes
+        # TODO: fix that should not be using starting_idx ... instead use a where content is none ... O(n) rather than using sort
+        parsed_content_icon = self.get_parsed_content_icon(filtered_boxes, starting_idx, image_source, 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
+
+        # draw boxes
+        filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
+        phrases = [i for i in range(len(filtered_boxes))]
+        annotated_frame, _ = self.annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
+
+        pil_img = Image.fromarray(annotated_frame)
+        buffered = io.BytesIO()
+        pil_img.save(buffered, format="PNG")
+        encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
+        assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
+
+        return encoded_image, filtered_boxes_elem
+
+    @staticmethod
+    def predict_yolo(yolo_model, image, box_threshold, iou_threshold):
+        """ Use huggingface model to replace the original model
+        """
+        result = yolo_model.predict(
+            source=image,
+            conf=box_threshold,
+            iou=iou_threshold,
+            verbose=False
+        )
+        boxes = result[0].boxes.xyxy
+        conf = result[0].boxes.conf
+        phrases = [str(i) for i in range(len(boxes))]
+
+        return boxes, conf, phrases
+
+    @staticmethod
+    def remove_overlap_new(boxes, iou_threshold, ocr_bbox):
+        '''
+        boxes format: [{'type': 'icon', 'bbox':[x,y], 'interactivity':True, 'content':None }, ...]
+        ocr_bbox format: [{'type': 'text', 'bbox':[x,y], 'interactivity':False, 'content':str }, ...]
+
+        '''
+        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.80
+
+        filtered_boxes = []
+        if ocr_bbox:
+            filtered_boxes.extend(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:
+                if ocr_bbox:
+                    box_added = False
+                    ocr_labels = ''
+                    for box3_elem in ocr_bbox:
+                        if not box_added:
+                            box3 = box3_elem['bbox']
+                            if is_inside(box3, box1): # ocr inside icon
+                                # delete the box3_elem from ocr_bbox
+                                try:
+                                    # gather all ocr labels
+                                    ocr_labels += box3_elem['content'] + ' '
+                                    filtered_boxes.remove(box3_elem)
+                                except:
+                                    continue
+                            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
+                                box_added = True
+                                break
+                            else:
+                                continue
+                    if not box_added:
+                        if ocr_labels:
+                            filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': ocr_labels, 'source':'box_yolo_content_ocr'})
+                        else:
+                            filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': None, 'source':'box_yolo_content_yolo'})
+                else:
+                    filtered_boxes.append(box1)
+        return filtered_boxes
+
+    @torch.inference_mode()
+    def get_parsed_content_icon(self, filtered_boxes, starting_idx, image_source, prompt=None, batch_size=128):
+        # Number of samples per batch, --> 128 roughly takes 4 GB of GPU memory for florence v2 model
+        non_ocr_boxes = filtered_boxes[starting_idx:]
+        cropped_pil_image = []
+        for i, coord in enumerate(non_ocr_boxes):
+            try:
+                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_image = cv2.resize(cropped_image, (64, 64))
+                cropped_pil_image.append(self.to_pil(cropped_image))
+            except:
+                continue
+
+        caption_model, caption_model_processor = self.caption_model, self.caption_model_processor
+        if not prompt:
+            prompt = "<CAPTION>"
+        
+        generated_texts = []
+        device = caption_model.device
+        for i in range(0, len(cropped_pil_image), batch_size):
+            batch = cropped_pil_image[i:i+batch_size]
+            if caption_model.device.type == 'cuda':
+                inputs = caption_model_processor(images=batch, text=[prompt]*len(batch), return_tensors="pt", do_resize=False).to(device=device, dtype=torch.float16)
+            else:
+                inputs = caption_model_processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
+            generated_ids = caption_model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=20,num_beams=1, do_sample=False, early_stopping=False)
+            generated_text = caption_model_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
+
+    @staticmethod
+    def annotate(image_source, boxes, logits, phrases, text_scale, 
+                text_padding=5, text_thickness=2, thickness=3):
+        """    
+        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 range(boxes.shape[0])]
+
+        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))
+
+        label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
+        return annotated_frame, label_coordinates
+    
+    @staticmethod
+    def int_box_area(box, w, h):
+        x1, y1, x2, y2 = box
+        int_box = [int(x1*w), int(y1*h), int(x2*w), int(y2*h)]
+        area = (int_box[2] - int_box[0]) * (int_box[3] - int_box[1])
+        return area
+
+    @staticmethod
+    def get_xyxy(input):
+        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
+
+
+class OmniparserForCausalLM(OmniparserModel):
+    """OmniParser model with a causal language modeling interface."""
+    _keys_to_ignore_on_load_missing = [r"caption_model\..*"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        """Required by HF generate() - passes image_source from kwargs to forward."""
+        return {"image_source": kwargs.get("image_source", None)}
+    
+    def generate(self, image_source=None, **kwargs):
+        """Process image and return formatted text output of detected elements."""
+        outputs = super().forward(image_source=image_source)
+        res = json.dumps(outputs.parsed_content_list)
+        return res
+
+# Register the model with Auto classes
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+AutoConfig.register("omniparserv2", OmniparserConfig)
+AutoModel.register(OmniparserConfig, OmniparserModel)
+AutoModelForCausalLM.register(OmniparserConfig, OmniparserForCausalLM)