# SPDX-FileCopyrightText: The Docling Contributors # SPDX-License-Identifier: MIT """KServe v2 remote implementation for object-detection models.""" from __future__ import annotations import logging import time from pathlib import Path from typing import TYPE_CHECKING, List, Optional, Union import numpy as np from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.kserve_transport_utils import resolve_kserve_transport_base_url from docling.datamodel.object_detection_engine_options import ( ApiKserveV2ObjectDetectionEngineOptions, ) from docling.exceptions import OperationNotAllowed from docling.models.inference_engines.common import KserveV2Client, KserveV2HttpClient from docling.models.inference_engines.object_detection.base import ( ObjectDetectionEngineInput, ObjectDetectionEngineOutput, ) from docling.models.inference_engines.object_detection.hf_base import ( HfObjectDetectionEngineBase, ) if TYPE_CHECKING: from docling.datamodel.stage_model_specs import EngineModelConfig _log = logging.getLogger(__name__) class ApiKserveV2ObjectDetectionEngine(HfObjectDetectionEngineBase): """Remote object-detection engine backed by KServe v2-compatible serving.""" def __init__( self, *, enable_remote_services: bool, options: ApiKserveV2ObjectDetectionEngineOptions, model_config: Optional[EngineModelConfig] = None, accelerator_options: AcceleratorOptions, artifacts_path: Optional[Union[Path, str]] = None, ): super().__init__( options=options, model_config=model_config, accelerator_options=accelerator_options, artifacts_path=artifacts_path, ) self.options: ApiKserveV2ObjectDetectionEngineOptions = options self._kserve_client: Optional[KserveV2Client] = None self._input_images_name: Optional[str] = None self._input_orig_target_sizes_name: Optional[str] = None self._output_labels_name: Optional[str] = None self._output_boxes_name: Optional[str] = None self._output_scores_name: Optional[str] = None if not enable_remote_services: raise OperationNotAllowed( "Connections to remote services are only allowed when set explicitly. " "pipeline_options.enable_remote_services=True." ) def _resolve_model_name(self) -> str: if self.options.model_name: return self.options.model_name return self._repo_id.replace("/", "--") def _resolve_model_version(self) -> Optional[str]: return self.options.model_version def _resolve_tensor_names(self) -> tuple[str, str, str, str, str]: if self._kserve_client is None: raise RuntimeError("KServe v2 client is not initialized.") metadata = self._kserve_client.get_model_metadata() if len(metadata.inputs) < 2: raise RuntimeError( "Expected object-detection model metadata to expose at least 2 inputs " f"(images, orig_target_sizes), got {len(metadata.inputs)}." ) if len(metadata.outputs) < 3: raise RuntimeError( "Expected object-detection model metadata to expose at least 3 outputs " f"(labels, boxes, scores), got {len(metadata.outputs)}." ) input_images_name = metadata.inputs[0].name input_orig_target_sizes_name = metadata.inputs[1].name output_labels_name = metadata.outputs[0].name output_boxes_name = metadata.outputs[1].name output_scores_name = metadata.outputs[2].name return ( input_images_name, input_orig_target_sizes_name, output_labels_name, output_boxes_name, output_scores_name, ) def initialize(self) -> None: """Initialize preprocessor/labels and prepare remote client.""" _log.info("Initializing KServe v2 object-detection engine") revision = self._model_config.revision or "main" model_folder = self._resolve_model_folder( repo_id=self._repo_id, revision=revision ) self._processor = self._load_preprocessor(model_folder) self._id_to_label = self._load_label_mapping(model_folder) if self._kserve_client is not None: self._kserve_client.close() base_url = resolve_kserve_transport_base_url( url=self.options.url, transport=self.options.transport, ) if self.options.transport == "http": self._kserve_client = KserveV2HttpClient( base_url=base_url, model_name=self._resolve_model_name(), model_version=self._resolve_model_version(), timeout=self.options.timeout, headers=self.options.headers, use_binary_data=self.options.use_binary_data, ) else: from docling.models.inference_engines.common.kserve_v2_grpc import ( KserveV2GrpcClient, ) self._kserve_client = KserveV2GrpcClient( base_url=base_url, model_name=self._resolve_model_name(), model_version=self._resolve_model_version(), timeout=self.options.timeout, metadata=self.options.grpc_metadata, use_tls=self.options.grpc_use_tls, max_message_bytes=self.options.grpc_max_message_bytes, use_binary_data=self.options.use_binary_data, grpc_channel_args=list(self.options.grpc_channel_args), ) ( self._input_images_name, self._input_orig_target_sizes_name, self._output_labels_name, self._output_boxes_name, self._output_scores_name, ) = self._resolve_tensor_names() self._initialized = True _log.info( "KServe v2 object-detection engine ready (inputs=[%s, %s], outputs=[%s, %s, %s])", self._input_images_name, self._input_orig_target_sizes_name, self._output_labels_name, self._output_boxes_name, self._output_scores_name, ) def predict_batch( self, input_batch: List[ObjectDetectionEngineInput] ) -> List[ObjectDetectionEngineOutput]: """Run inference on a batch of images against a KServe v2 endpoint.""" if not input_batch: return [] if not self._initialized: raise RuntimeError("Engine not initialized. Call initialize() first.") # Type narrowing: _initialized guarantees these are non-None assert self._processor is not None assert self._kserve_client is not None assert self._input_images_name is not None assert self._input_orig_target_sizes_name is not None assert self._output_labels_name is not None assert self._output_boxes_name is not None assert self._output_scores_name is not None if _log.isEnabledFor(logging.DEBUG): _t_preproc_start = time.time() _t_preproc_mono = time.monotonic() images = [item.image.convert("RGB") for item in input_batch] processed_inputs = self._processor(images=images, return_tensors="np") pixel_values = np.asarray(processed_inputs["pixel_values"]) orig_sizes = np.asarray( [[image.width, image.height] for image in images], dtype=np.int64, ) if _log.isEnabledFor(logging.DEBUG): _log.debug( "PIPELINE_PROFILING KServe object-detection HF preprocessor: batch_size=%d start=%.3f end=%.3f duration=%.3fs", len(input_batch), _t_preproc_start, time.time(), time.monotonic() - _t_preproc_mono, ) outputs = self._kserve_client.infer( inputs={ self._input_images_name: pixel_values, self._input_orig_target_sizes_name: orig_sizes, }, output_names=[ self._output_labels_name, self._output_boxes_name, self._output_scores_name, ], request_parameters=self.options.request_parameters, ) try: labels_batch = outputs[self._output_labels_name] boxes_batch = outputs[self._output_boxes_name] scores_batch = outputs[self._output_scores_name] except KeyError as exc: raise RuntimeError( "Missing one or more expected KServe v2 outputs: " f"{self._output_labels_name}, " f"{self._output_boxes_name}, " f"{self._output_scores_name}" ) from exc if len(labels_batch) != len(input_batch): raise RuntimeError( "KServe v2 output batch size mismatch for labels: " f"expected {len(input_batch)}, got {len(labels_batch)}" ) batch_outputs: List[ObjectDetectionEngineOutput] = [] for idx, input_item in enumerate(input_batch): batch_outputs.append( self._build_output( input_item=input_item, labels=labels_batch[idx], scores=scores_batch[idx], boxes=boxes_batch[idx], apply_score_threshold=True, ) ) return batch_outputs def close(self) -> None: if self._kserve_client is None: return self._kserve_client.close() def __del__(self) -> None: try: self.close() except Exception: pass