parse_gui_bls/1/model.py

import base64
import json
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any

import cv2
import numpy as np
import triton_python_backend_utils as pb_utils  # type: ignore

_BOX_PATTERN = re.compile(r"\[(\-?\d+),(\-?\d+)\]\[(\-?\d+),(\-?\d+)\]")


def _decode_string(value: Any) -> str:
    if isinstance(value, np.bytes_):
        return value.tobytes().decode("utf-8")
    if isinstance(value, bytes):
        return value.decode("utf-8")
    return str(value)


def _parse_box(raw_box: Any):
    if isinstance(raw_box, (list, tuple)) and len(raw_box) == 4:
        return [int(raw_box[0]), int(raw_box[1]), int(raw_box[2]), int(raw_box[3])]

    if isinstance(raw_box, str):
        m = _BOX_PATTERN.match(raw_box.strip())
        if m is not None:
            return [int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4))]

    return None


class TritonPythonModel:
    def initialize(self, args):
        self.vision_dim = 768
        self.bge_dim = 1024
        self.siglip_batch_size = int(os.environ.get("SIGLIP_BATCH_SIZE", "8"))
        self.bge_batch_size = int(os.environ.get("BGE_M3_BATCH_SIZE", "16"))
        self.siglip_request_workers = max(1, int(os.environ.get("SIGLIP_REQUEST_WORKERS", "4")))
        self.bge_request_workers = max(1, int(os.environ.get("BGE_REQUEST_WORKERS", "4")))
        self.top_level_workers = max(1, int(os.environ.get("PARSE_GUI_TOPLEVEL_WORKERS", "2")))
        self.top_executor = ThreadPoolExecutor(max_workers=self.top_level_workers)
        self.siglip_executor = ThreadPoolExecutor(max_workers=self.siglip_request_workers)
        self.bge_executor = ThreadPoolExecutor(max_workers=self.bge_request_workers)

    def _request_merge_detect(self, image_bytes_obj):
        req = pb_utils.InferenceRequest(
            model_name="merge_detect",
            requested_output_names=["OD_DATA_LIST"],
            inputs=[pb_utils.Tensor("IMAGE_BYTES", np.array([image_bytes_obj], dtype=np.object_))],
        )
        resp = req.exec()
        if resp.has_error():
            raise pb_utils.TritonModelException(resp.error().message())
        return resp

    def _request_siglip_engine(self, pixel_values: np.ndarray):
        preferred_memory = pb_utils.PreferredMemory(
            pb_utils.TRITONSERVER_MEMORY_CPU, 0
        )
        req = pb_utils.InferenceRequest(
            model_name="siglip_engine",
            requested_output_names=["image_features"],
            inputs=[pb_utils.Tensor("pixel_values", pixel_values.astype(np.float32, copy=False))],
            preferred_memory=preferred_memory,
        )
        resp = req.exec()
        if resp.has_error():
            raise pb_utils.TritonModelException(resp.error().message())
        out = pb_utils.get_output_tensor_by_name(resp, "image_features").as_numpy()
        return out.astype(np.float32, copy=False)

    def _request_bge(self, texts: list[str]):
        if len(texts) == 0:
            return np.empty((0, self.bge_dim), dtype=np.float32)
        tensor = np.array([[t.encode("utf-8")] for t in texts], dtype=np.object_)
        preferred_memory = pb_utils.PreferredMemory(
            pb_utils.TRITONSERVER_MEMORY_CPU, 0
        )
        req = pb_utils.InferenceRequest(
            model_name="bge_m3",
            requested_output_names=["EMBEDDINGS"],
            inputs=[pb_utils.Tensor("TEXTS", tensor)],
            preferred_memory=preferred_memory,
        )
        resp = req.exec()
        if resp.has_error():
            raise pb_utils.TritonModelException(resp.error().message())
        out = pb_utils.get_output_tensor_by_name(resp, "EMBEDDINGS").as_numpy()
        return out.astype(np.float32, copy=False)

    @staticmethod
    def _prepare_siglip_pixel_values(crop_images: list[np.ndarray]) -> np.ndarray:
        pixel_values = []
        for crop in crop_images:
            im = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
            im = cv2.resize(im, (224, 224), interpolation=cv2.INTER_AREA)
            x = im.astype(np.float32) / 255.0
            x = (x - 0.5) / 0.5
            x = np.transpose(x, (2, 0, 1))
            pixel_values.append(x)
        if len(pixel_values) == 0:
            return np.empty((0, 3, 224, 224), dtype=np.float32)
        return np.stack(pixel_values, axis=0).astype(np.float32)

    def _siglip_embeddings_from_crops(self, crop_images: list[np.ndarray]):
        if len(crop_images) == 0:
            return np.empty((0, self.vision_dim), dtype=np.float32)
        chunks: list[np.ndarray] = []
        for i in range(0, len(crop_images), self.siglip_batch_size):
            chunk = crop_images[i : i + self.siglip_batch_size]
            chunks.append(self._prepare_siglip_pixel_values(chunk))

        if len(chunks) == 1 or self.siglip_request_workers <= 1:
            embeddings = [self._request_siglip_engine(chunks[0])]
            return np.concatenate(embeddings, axis=0).astype(np.float32, copy=False)

        chunk_results: list[np.ndarray | None] = [None] * len(chunks)
        future_to_idx = {
            self.siglip_executor.submit(self._request_siglip_engine, pixel_values): idx
            for idx, pixel_values in enumerate(chunks)
        }
        for future, idx in future_to_idx.items():
            chunk_results[idx] = future.result()

        embeddings = [r for r in chunk_results if r is not None]
        return np.concatenate(embeddings, axis=0).astype(np.float32, copy=False)

    def _embed_bge_for_caption_and_text(
        self,
        captions: list[str],
        texts: list[str],
    ) -> tuple[np.ndarray, np.ndarray]:
        n = len(captions)
        function_embeddings = np.zeros((n, self.bge_dim), dtype=np.float32)
        text_embeddings = np.zeros((n, self.bge_dim), dtype=np.float32)

        unique_text_to_targets: dict[str, list[tuple[str, int]]] = {}

        for idx, caption in enumerate(captions):
            normalized_caption = caption.strip()
            if normalized_caption.lower() == "unknown":
                continue
            unique_text_to_targets.setdefault(normalized_caption, []).append(("function", idx))

        for idx, text in enumerate(texts):
            normalized_text = text.strip()
            if len(normalized_text) == 0:
                continue
            unique_text_to_targets.setdefault(normalized_text, []).append(("text", idx))

        request_texts = list(unique_text_to_targets.keys())

        chunks: list[list[str]] = []
        for i in range(0, len(request_texts), self.bge_batch_size):
            chunk_texts = request_texts[i : i + self.bge_batch_size]
            chunks.append(chunk_texts)

        if len(chunks) == 0:
            return function_embeddings, text_embeddings

        chunk_results: list[np.ndarray | None] = [None] * len(chunks)
        if len(chunks) == 1 or self.bge_request_workers <= 1:
            chunk_results[0] = self._request_bge(chunks[0])
        else:
            future_to_idx = {
                self.bge_executor.submit(self._request_bge, chunk_texts): idx for idx, chunk_texts in enumerate(chunks)
            }
            for future, idx in future_to_idx.items():
                chunk_results[idx] = future.result()

        for idx, chunk_texts in enumerate(chunks):
            chunk_embeddings = chunk_results[idx]
            if chunk_embeddings is None or chunk_embeddings.size == 0:
                continue
            for j, text_key in enumerate(chunk_texts):
                if j >= chunk_embeddings.shape[0]:
                    break
                for target, row_idx in unique_text_to_targets.get(text_key, []):
                    if target == "function":
                        function_embeddings[row_idx] = chunk_embeddings[j]
                    else:
                        text_embeddings[row_idx] = chunk_embeddings[j]

        return function_embeddings, text_embeddings

    def execute(self, requests):
        logger = pb_utils.Logger
        responses = []
        st = time.time()

        for request in requests:
            image_bytes_tensor = pb_utils.get_input_tensor_by_name(request, "IMAGE_BYTES")
            image_obj = image_bytes_tensor.as_numpy().reshape(-1)[0]
            image_base64 = _decode_string(image_obj)
            image_bytes = base64.b64decode(image_base64)

            image_np = np.frombuffer(image_bytes, np.uint8)
            image_bgr = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
            if image_bgr is None:
                raise pb_utils.TritonModelException("Failed to decode IMAGE_BYTES.")

            merge_resp = self._request_merge_detect(image_obj)
            od_raw = pb_utils.get_output_tensor_by_name(merge_resp, "OD_DATA_LIST").as_numpy()
            gui_items_raw = json.loads(_decode_string(od_raw.reshape(-1)[0]))

            image_h, image_w = image_bgr.shape[:2]
            gui_items = []
            crop_images = []
            for od in gui_items_raw:
                box = _parse_box(od.get("box", [0, 0, 0, 0]))
                if box is None:
                    continue

                left = max(0, min(int(box[0]), image_w - 1))
                top = max(0, min(int(box[1]), image_h - 1))
                right = max(0, min(int(box[2]), image_w))
                bottom = max(0, min(int(box[3]), image_h))
                if right <= left or bottom <= top:
                    continue

                crop = image_bgr[top:bottom, left:right]
                if crop.size == 0:
                    continue

                caption = str(od.get("caption", "Unknown")).strip()
                if len(caption) == 0:
                    caption = "Unknown"
                gui_items.append(
                    {
                        "box": [left, top, right, bottom],
                        "text": str(od.get("text", "")),
                        "caption": caption,
                    }
                )
                crop_images.append(crop.copy())

            captions = [str(item["caption"]) for item in gui_items]
            texts = [str(item["text"]) for item in gui_items]
            if self.top_level_workers > 1:
                future_siglip = self.top_executor.submit(self._siglip_embeddings_from_crops, crop_images)
                future_bge = self.top_executor.submit(self._embed_bge_for_caption_and_text, captions, texts)
                siglip_embeddings = future_siglip.result()
                function_embedding, text_embedding = future_bge.result()
            else:
                siglip_embeddings = self._siglip_embeddings_from_crops(crop_images)
                function_embedding, text_embedding = self._embed_bge_for_caption_and_text(captions, texts)

            n = len(gui_items)

            if n == 0:
                bbox = np.empty((0, 4), dtype=np.float32)
                vision_embedding = np.empty((0, self.vision_dim), dtype=np.float32)
                function_embedding = np.empty((0, self.bge_dim), dtype=np.float32)
                text_embedding = np.empty((0, self.bge_dim), dtype=np.float32)
            else:
                width = float(max(image_w, 1))
                height = float(max(image_h, 1))
                bbox = np.array(
                    [
                        [
                            float(item["box"][0]) / width,
                            float(item["box"][1]) / height,
                            float(item["box"][2]) / width,
                            float(item["box"][3]) / height,
                        ]
                        for item in gui_items
                    ],
                    dtype=np.float32,
                )
                np.clip(bbox, 0.0, 1.0, out=bbox)

                vision_embedding = np.zeros((n, self.vision_dim), dtype=np.float32)
                if siglip_embeddings.ndim == 2 and siglip_embeddings.shape[0] > 0:
                    n_copy = min(n, siglip_embeddings.shape[0])
                    d_copy = min(self.vision_dim, siglip_embeddings.shape[1])
                    vision_embedding[:n_copy, :d_copy] = siglip_embeddings[:n_copy, :d_copy]

            responses.append(
                pb_utils.InferenceResponse(
                    output_tensors=[
                        pb_utils.Tensor("bbox", bbox),
                        pb_utils.Tensor("text_embedding", text_embedding),
                        pb_utils.Tensor("function_embedding", function_embedding),
                        pb_utils.Tensor("vision_embedding", vision_embedding),
                    ]
                )
            )

        logger.log_info(f"parse_gui_bls execute duration : {int((time.time() - st) * 1000)} ms")
        return responses

    def finalize(self):
        try:
            self.top_executor.shutdown(wait=False, cancel_futures=True)
        except Exception:
            pass
        try:
            self.siglip_executor.shutdown(wait=False, cancel_futures=True)
        except Exception:
            pass
        try:
            self.bge_executor.shutdown(wait=False, cancel_futures=True)
        except Exception:
            pass