Add `handler.py` and `requirements.txt` (#23)

- Add `handler.py` and `requirements.txt` (acd5c25b4e94f3dff29ee1a1e3b3764c75d3e5f7)
- Update handler.py (1ebdf336dc12fef6aba071bb3b5a905272ede53d)
- Update README.md (3655c7c72dc109866aa38022f07735f7cc475b7b)
- Apply suggestions from

@ThomasDh-C
code review (54dd9132c195511be5b6f4dbc46f14b32934ee4f)


Co-authored-by: Alvaro Bartolome <[email protected]>

README.md

@@ -1,7 +1,9 @@
 ---
 library_name: transformers
 license: mit
-pipeline_tag: image-text-to-text
+tags:
+- endpoint-template
+- custom_code
 ---
 📢 [[GitHub Repo](https://github.com/microsoft/OmniParser/tree/master)] [[OmniParser V2 Blog Post](https://www.microsoft.com/en-us/research/articles/omniparser-v2-turning-any-llm-into-a-computer-use-agent/)] [Huggingface demo](https://huggingface.co/spaces/microsoft/OmniParser-v2)
 
@@ -28,5 +30,4 @@ This model hub includes a finetuned version of YOLOv8 and a finetuned Florence-2
 - While OmniParser only converts screenshot image into texts, it can be used to construct an GUI agent based on LLMs that is actionable. When developing and operating the agent using OmniParser, the developers need to be responsible and follow common safety standard. 
 
 # License
-Please note that icon_detect model is under AGPL license, and icon_caption is under MIT license. Please refer to the LICENSE file in the folder of each model. 
-
+Please note that icon_detect model is under AGPL license, and icon_caption is under MIT license. Please refer to the LICENSE file in the folder of each model.

handler.py

@@ -0,0 +1,645 @@
+import base64
+import io
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+import cv2
+import easyocr
+import numpy as np
+import torch
+from PIL import Image
+from PIL.Image import Image as ImageType
+from supervision.detection.core import Detections
+from supervision.draw.color import Color, ColorPalette
+from torchvision.ops import box_convert
+from torchvision.transforms import ToPILImage
+from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers.image_utils import load_image
+from ultralytics import YOLO
+
+# NOTE: here so that it's downloaded before hand so that the endpoint it not stuck listening, whilst the required
+# files are still being downloaded
+easyocr.Reader(["en"])
+
+
+class EndpointHandler:
+    def __init__(self, model_dir: str = "/repository") -> None:
+        self.device = (
+            torch.device("cuda") if torch.cuda.is_available()
+            else (torch.device("mps") if torch.backends.mps.is_available()
+                  else torch.device("cpu"))
+        )
+
+        # bounding box detection model
+        self.yolo = YOLO(f"{model_dir}/icon_detect/model.pt")
+
+        # captioning model
+        self.processor = AutoProcessor.from_pretrained(
+            "microsoft/Florence-2-base", trust_remote_code=True
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            f"{model_dir}/icon_caption",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+        ).to(self.device)
+
+        # ocr
+        self.ocr = easyocr.Reader(["en"])
+
+        # box annotator
+        self.annotator = BoxAnnotator()
+
+    def __call__(self, data: Dict[str, Any]) -> Any:
+        # data should contain the following:
+        #  "inputs": {
+        #    "image": url/base64,
+        #    (optional) "image_size": {"w": int, "h": int},
+        #    (optional) "bbox_threshold": float,
+        #    (optional) "iou_threshold": float,
+        #  }
+        data = data.pop("inputs")
+
+        # read image from either url or base64 encoding
+        image = load_image(data["image"])
+
+        ocr_texts, ocr_bboxes = self.check_ocr_bboxes(
+            image,
+            out_format="xyxy",
+            ocr_kwargs={"text_threshold": 0.8},
+        )
+        annotated_image, filtered_bboxes_out = self.get_som_labeled_img(
+            image,
+            image_size=data.get("image_size", None),
+            ocr_texts=ocr_texts,
+            ocr_bboxes=ocr_bboxes,
+            bbox_threshold=data.get("bbox_threshold", 0.05),
+            iou_threshold=data.get("iou_threshold", None),
+        )
+        return {
+            "image": annotated_image,
+            "bboxes": filtered_bboxes_out,
+        }
+
+    def check_ocr_bboxes(
+        self,
+        image: ImageType,
+        out_format: Literal["xywh", "xyxy"] = "xywh",
+        ocr_kwargs: Optional[Dict[str, Any]] = {},
+    ) -> Tuple[List[str], List[List[int]]]:
+        if image.mode == "RBGA":
+            image = image.convert("RGB")
+
+        result = self.ocr.readtext(np.array(image), **ocr_kwargs)  # type: ignore
+        texts = [str(item[1]) for item in result]
+        bboxes = [
+            self.coordinates_to_bbox(item[0], format=out_format) for item in result
+        ]
+        return (texts, bboxes)
+
+    @staticmethod
+    def coordinates_to_bbox(
+        coordinates: np.ndarray, format: Literal["xywh", "xyxy"] = "xywh"
+    ) -> List[int]:
+        match format:
+            case "xywh":
+                return [
+                    int(coordinates[0][0]),
+                    int(coordinates[0][1]),
+                    int(coordinates[2][0] - coordinates[0][0]),
+                    int(coordinates[2][1] - coordinates[0][1]),
+                ]
+            case "xyxy":
+                return [
+                    int(coordinates[0][0]),
+                    int(coordinates[0][1]),
+                    int(coordinates[2][0]),
+                    int(coordinates[2][1]),
+                ]
+
+    @staticmethod
+    def bbox_area(bbox: List[int], w: int, h: int) -> int:
+        bbox = [bbox[0] * w, bbox[1] * h, bbox[2] * w, bbox[3] * h]
+        return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+
+    @staticmethod
+    def remove_bbox_overlap(
+        xyxy_bboxes: List[Dict[str, Any]],
+        ocr_bboxes: Optional[List[Dict[str, Any]]] = None,
+        iou_threshold: Optional[float] = 0.7,
+    ) -> List[Dict[str, Any]]:
+        filtered_bboxes = []
+        if ocr_bboxes is not None:
+            filtered_bboxes.extend(ocr_bboxes)
+
+        for i, bbox_outter in enumerate(xyxy_bboxes):
+            bbox_left = bbox_outter["bbox"]
+            valid_bbox = True
+
+            for j, bbox_inner in enumerate(xyxy_bboxes):
+                if i == j:
+                    continue
+
+                bbox_right = bbox_inner["bbox"]
+                if (
+                    intersection_over_union(
+                        bbox_left,
+                        bbox_right,
+                    )
+                    > iou_threshold  # type: ignore
+                ) and (area(bbox_left) > area(bbox_right)):
+                    valid_bbox = False
+                    break
+
+            if valid_bbox is False:
+                continue
+
+            if ocr_bboxes is None:
+                filtered_bboxes.append(bbox_outter)
+                continue
+
+            box_added = False
+            ocr_labels = []
+            for ocr_bbox in ocr_bboxes:
+                if not box_added:
+                    bbox_right = ocr_bbox["bbox"]
+                    if overlap(bbox_right, bbox_left):
+                        try:
+                            ocr_labels.append(ocr_bbox["content"])
+                            filtered_bboxes.remove(ocr_bbox)
+                        except Exception:
+                            continue
+                    elif overlap(bbox_left, bbox_right):
+                        box_added = True
+                        break
+
+            if not box_added:
+                filtered_bboxes.append(
+                    {
+                        "type": "icon",
+                        "bbox": bbox_outter["bbox"],
+                        "interactivity": True,
+                        "content": " ".join(ocr_labels) if ocr_labels else None,
+                    }
+                )
+
+        return filtered_bboxes
+
+    def get_som_labeled_img(
+        self,
+        image: ImageType,
+        image_size: Optional[Dict[Literal["w", "h"], int]] = None,
+        ocr_texts: Optional[List[str]] = None,
+        ocr_bboxes: Optional[List[List[int]]] = None,
+        bbox_threshold: float = 0.01,
+        iou_threshold: Optional[float] = None,
+        caption_prompt: Optional[str] = None,
+        caption_batch_size: int = 64,  # ~2GiB of GPU VRAM (can be increased to 128 which is ~4GiB of GPU VRAM)
+    ) -> Tuple[str, List[Dict[str, Any]]]:
+        if image.mode == "RBGA":
+            image = image.convert("RGB")
+
+        w, h = image.size
+        if image_size is None:
+            imgsz = {"h": h, "w": w}
+        else:
+            imgsz = [image_size.get("h", h), image_size.get("w", w)]
+
+        out = self.yolo.predict(
+            image,
+            imgsz=imgsz,
+            conf=bbox_threshold,
+            iou=iou_threshold or 0.7,
+            verbose=False,
+        )[0]
+        if out.boxes is None:
+            raise RuntimeError(
+                "YOLO prediction failed to produce the bounding boxes..."
+            )
+
+        xyxy_bboxes = out.boxes.xyxy
+        xyxy_bboxes = xyxy_bboxes / torch.Tensor([w, h, w, h]).to(xyxy_bboxes.device)
+        image_np = np.asarray(image)  # type: ignore
+
+        if ocr_bboxes:
+            ocr_bboxes = torch.tensor(ocr_bboxes) / torch.Tensor([w, h, w, h])  # type: ignore
+            ocr_bboxes = ocr_bboxes.tolist()  # type: ignore
+
+        ocr_bboxes = [
+            {
+                "type": "text",
+                "bbox": bbox,
+                "interactivity": False,
+                "content": text,
+                "source": "box_ocr_content_ocr",
+            }
+            for bbox, text in zip(ocr_bboxes, ocr_texts)  # type: ignore
+            if self.bbox_area(bbox, w, h) > 0
+        ]
+        xyxy_bboxes = [
+            {
+                "type": "icon",
+                "bbox": bbox,
+                "interactivity": True,
+                "content": None,
+                "source": "box_yolo_content_yolo",
+            }
+            for bbox in xyxy_bboxes.tolist()
+            if self.bbox_area(bbox, w, h) > 0
+        ]
+
+        filtered_bboxes = self.remove_bbox_overlap(
+            xyxy_bboxes=xyxy_bboxes,
+            ocr_bboxes=ocr_bboxes,  # type: ignore
+            iou_threshold=iou_threshold or 0.7,
+        )
+
+        filtered_bboxes_out = sorted(
+            filtered_bboxes, key=lambda x: x["content"] is None
+        )
+        starting_idx = next(
+            (
+                idx
+                for idx, bbox in enumerate(filtered_bboxes_out)
+                if bbox["content"] is None
+            ),
+            -1,
+        )
+
+        filtered_bboxes = torch.tensor([box["bbox"] for box in filtered_bboxes_out])
+        non_ocr_bboxes = filtered_bboxes[starting_idx:]
+
+        bbox_images = []
+        for _, coordinates in enumerate(non_ocr_bboxes):
+            try:
+                xmin, xmax = (
+                    int(coordinates[0] * image_np.shape[1]),
+                    int(coordinates[2] * image_np.shape[1]),
+                )
+                ymin, ymax = (
+                    int(coordinates[1] * image_np.shape[0]),
+                    int(coordinates[3] * image_np.shape[0]),
+                )
+                cropped_image = image_np[ymin:ymax, xmin:xmax, :]
+                cropped_image = cv2.resize(cropped_image, (64, 64))
+                bbox_images.append(ToPILImage()(cropped_image))
+            except Exception:
+                continue
+
+        if caption_prompt is None:
+            caption_prompt = "<CAPTION>"
+
+        captions = []
+        for idx in range(0, len(bbox_images), caption_batch_size):  # type: ignore
+            batch = bbox_images[idx : idx + caption_batch_size]  # type: ignore
+            inputs = self.processor(
+                images=batch,
+                text=[caption_prompt] * len(batch),
+                return_tensors="pt",
+                do_resize=False,
+            )
+            if self.device.type in {"cuda", "mps"}:
+                inputs = inputs.to(device=self.device, dtype=torch.float16)
+
+            with torch.inference_mode():
+                generated_ids = self.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_texts = self.processor.batch_decode(
+                generated_ids, skip_special_tokens=True
+            )
+            captions.extend([text.strip() for text in generated_texts])
+
+        ocr_texts = [f"Text Box ID {idx}: {text}" for idx, text in enumerate(ocr_texts)]  # type: ignore
+        for _, bbox in enumerate(filtered_bboxes_out):
+            if bbox["content"] is None:
+                bbox["content"] = captions.pop(0)
+
+        filtered_bboxes = box_convert(
+            boxes=filtered_bboxes, in_fmt="xyxy", out_fmt="cxcywh"
+        )
+
+        annotated_image = image_np.copy()
+        bboxes_annotate = filtered_bboxes * torch.Tensor([w, h, w, h])
+        xyxy_annotate = box_convert(
+            bboxes_annotate, in_fmt="cxcywh", out_fmt="xyxy"
+        ).numpy()
+        detections = Detections(xyxy=xyxy_annotate)
+        labels = [str(idx) for idx in range(bboxes_annotate.shape[0])]
+
+        annotated_image = self.annotator.annotate(
+            scene=annotated_image,
+            detections=detections,
+            labels=labels,
+            image_size=(w, h),
+        )
+        assert w == annotated_image.shape[1] and h == annotated_image.shape[0]
+
+        out_image = Image.fromarray(annotated_image)
+        out_buffer = io.BytesIO()
+        out_image.save(out_buffer, format="PNG")
+        encoded_image = base64.b64encode(out_buffer.getvalue()).decode("ascii")
+
+        return encoded_image, filtered_bboxes_out
+
+
+def area(bbox: List[int]) -> int:
+    return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+
+
+def intersection_area(bbox_left: List[int], bbox_right: List[int]) -> int:
+    return max(
+        0, min(bbox_left[2], bbox_right[2]) - min(bbox_left[0], bbox_right[0])
+    ) * max(0, min(bbox_left[3], bbox_right[3]) - min(bbox_left[1], bbox_right[1]))
+
+
+def intersection_over_union(bbox_left: List[int], bbox_right: List[int]) -> float:
+    intersection = intersection_area(bbox_left, bbox_right)
+    bbox_left_area = area(bbox_left)
+    bbox_right_area = area(bbox_right)
+    union = bbox_left_area + bbox_right_area - intersection + 1e-6
+
+    ratio_left, ratio_right = 0, 0
+    if bbox_left_area > 0 and bbox_right_area > 0:
+        ratio_left = intersection / bbox_left_area
+        ratio_right = intersection / bbox_right_area
+    return max(intersection / union, ratio_left, ratio_right)
+
+
+def overlap(bbox_left: List[int], bbox_right: List[int]) -> bool:
+    intersection = intersection_area(bbox_left, bbox_right)
+    ratio_left = intersection / area(bbox_left)
+    return ratio_left > 0.80
+
+
+class BoxAnnotator:
+    def __init__(
+        self,
+        color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,  # type: ignore
+        thickness: int = 3,
+        text_color: Color = Color.BLACK,  # type: ignore
+        text_scale: float = 0.5,
+        text_thickness: int = 2,
+        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:
+        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
+            else:
+                (
+                    text_x,
+                    text_y,
+                    text_background_x1,
+                    text_background_y1,
+                    text_background_x2,
+                    text_background_y2,
+                ) = self.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,
+            )
+            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=text_color,
+                thickness=self.text_thickness,
+                lineType=cv2.LINE_AA,
+            )
+        return scene
+
+    @staticmethod
+    def get_optimal_label_pos(
+        text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size
+    ):
+        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 (
+                    intersection_over_union(
+                        [
+                            text_background_x1,
+                            text_background_y1,
+                            text_background_x2,
+                            text_background_y2,
+                        ],
+                        detection,
+                    )
+                    > 0.3
+                ):
+                    is_overlap = True
+                    break
+            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
+
+        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,
+            )
+
+        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,
+            )
+
+        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,
+            )
+
+        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,
+        )

requirements.txt

@@ -0,0 +1,7 @@
+easyocr
+einops==0.8.0
+opencv-python
+opencv-python-headless
+supervision==0.18.0
+timm
+ultralytics==8.3.70