aissemble-inference-core 1.5.0__py3-none-any.whl

aissemble_inference_core/client/translators/_tensor_utils.py

@@ -0,0 +1,89 @@
+###
+# #%L
+# aiSSEMBLE::Open Inference Protocol::Core
+# %%
+# Copyright (C) 2024 Booz Allen Hamilton Inc.
+# %%
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# #L%
+###
+"""Shared utilities for tensor validation and extraction."""
+
+from typing import Any
+
+from aissemble_inference_core.client.oip_adapter import TensorData
+
+
+def validate_required_tensors(
+    outputs: dict[str, TensorData], required: list[str], format_name: str
+) -> None:
+    """Validate that all required tensors are present in outputs.
+
+    Args:
+        outputs: Dictionary mapping tensor names to TensorData
+        required: List of required tensor names
+        format_name: Name of the format for error messages (e.g., "TensorFlow")
+
+    Raises:
+        ValueError: If any required tensors are missing
+    """
+    missing = [name for name in required if name not in outputs]
+    if missing:
+        raise ValueError(
+            f"{format_name} format requires tensors {missing} but server returned: "
+            f"{list(outputs.keys())}"
+        )
+
+
+def validate_tensor_lengths_match(
+    *tensors: list[Any], names: list[str] | None = None
+) -> None:
+    """Validate that all tensors have the same length.
+
+    Args:
+        *tensors: Variable number of tensor data lists
+        names: Optional names for error messages
+
+    Raises:
+        ValueError: If tensor lengths don't match
+    """
+    lengths = [len(t) for t in tensors]
+    if len(set(lengths)) > 1:
+        if names:
+            details = ", ".join(f"{len(t)} {name}" for t, name in zip(tensors, names))
+            raise ValueError(f"Tensor length mismatch: {details}. Expected all equal.")
+        else:
+            raise ValueError(f"Tensor length mismatch: {lengths}. Expected all equal.")
+
+
+def unbatch_tensor(tensor: TensorData, num_items: int | None = None) -> list[Any]:
+    """Extract data from batched tensor [1, N, ...] -> [N, ...] and apply limit.
+
+    Args:
+        tensor: TensorData with potentially batched data
+        num_items: Optional limit on number of items to extract
+
+    Returns:
+        Unbatched list of items
+    """
+    data = tensor.data
+
+    # Unbatch if [1, N] or [1, N, M] shape
+    if isinstance(data, list) and data and isinstance(data[0], list):
+        data = data[0]
+
+    # Apply limit if specified
+    if num_items is not None:
+        data = data[:num_items]
+
+    return data

aissemble_inference_core/client/translators/object_detection_translator.py

@@ -0,0 +1,212 @@
+###
+# #%L
+# aiSSEMBLE::Open Inference Protocol::Core
+# %%
+# Copyright (C) 2024 Booz Allen Hamilton Inc.
+# %%
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# #L%
+###
+from typing import Any
+
+from aissemble_inference_core.client.oip_adapter import (
+    OipRequest,
+    OipResponse,
+    TensorData,
+)
+from aissemble_inference_core.client.results import (
+    BoundingBox,
+    Detection,
+    ObjectDetectionResult,
+)
+from aissemble_inference_core.client.translator import Translator
+from aissemble_inference_core.client.translators._image_utils import (
+    encode_image_for_oip,
+)
+from aissemble_inference_core.client.translators._tensor_utils import (
+    validate_required_tensors,
+    validate_tensor_lengths_match,
+)
+
+
+class DefaultObjectDetectionTranslator(Translator[Any, ObjectDetectionResult]):
+    """A reasonable default translator for object detection payloads.
+
+    This translator handles:
+    - Auto-encoding images to base64 bytes for OIP compatibility
+    - Converting various image formats (PIL, numpy, file paths)
+    - Parsing bounding box outputs in common formats
+    - Creating ObjectDetectionResult with proper metadata
+
+    The translator expects the model to output:
+    - bboxes: [N, 4] tensor with coordinates (x1, y1, x2, y2)
+    - labels: [N] tensor with class labels
+    - scores: [N] tensor with confidence scores
+
+    Thread Safety:
+    This translator is thread-safe and stateless. Image dimensions are stored in
+    request parameters and retrieved from response parameters.
+    """
+
+    def __init__(
+        self,
+        input_name: str = "image",
+        bbox_output_name: str = "bboxes",
+        label_output_name: str = "labels",
+        score_output_name: str = "scores",
+    ):
+        """Initialize the translator with configurable tensor names.
+
+        Args:
+            input_name: Name of the input tensor (default: "image")
+            bbox_output_name: Name of bounding box output tensor (default: "bboxes")
+            label_output_name: Name of label output tensor (default: "labels")
+            score_output_name: Name of score output tensor (default: "scores")
+        """
+        self.input_name = input_name
+        self.bbox_output_name = bbox_output_name
+        self.label_output_name = label_output_name
+        self.score_output_name = score_output_name
+
+    def preprocess(self, input_data: Any) -> OipRequest:  # noqa: A003
+        """Preprocess image input into an OipRequest.
+
+        Accepts various input formats:
+        - PIL Image
+        - numpy array
+        - file path (string)
+        - bytes
+
+        Args:
+            input_data: Image data in supported format
+
+        Returns:
+            OipRequest with encoded image tensor and dimensions in parameters
+        """
+        image_bytes, width, height = encode_image_for_oip(input_data)
+
+        tensor = TensorData(
+            name=self.input_name,
+            shape=[1, len(image_bytes)],
+            datatype="BYTES",
+            data=[[image_bytes]],
+        )
+
+        # Store image dimensions in request parameters for stateless operation
+        return OipRequest(
+            inputs=[tensor],
+            parameters={"_image_width": width, "_image_height": height},
+        )
+
+    def postprocess(self, response: OipResponse) -> ObjectDetectionResult:
+        """Postprocess OipResponse into ObjectDetectionResult.
+
+        Args:
+            response: OIP response containing detection outputs and dimensions
+
+        Returns:
+            ObjectDetectionResult with parsed detections
+
+        Raises:
+            ValueError: If required tensors are missing or dimensions not available
+        """
+        outputs = {out.name: out for out in response.outputs}
+
+        # Validate required tensors are present
+        required = [
+            self.bbox_output_name,
+            self.label_output_name,
+            self.score_output_name,
+        ]
+        validate_required_tensors(outputs, required, "Default object detection")
+
+        # Get image dimensions from response parameters
+        params = response.parameters or {}
+        image_width = params.get("_image_width")
+        image_height = params.get("_image_height")
+
+        if image_width is None or image_height is None:
+            raise ValueError(
+                "Image dimensions not found in response parameters. "
+                "Ensure the OIP adapter preserves request parameters in the response."
+            )
+
+        bbox_tensor = outputs[self.bbox_output_name]
+        bboxes = self._extract_bboxes(bbox_tensor)
+        labels = self._extract_tensor_data(outputs[self.label_output_name])
+        scores = self._extract_tensor_data(outputs[self.score_output_name])
+
+        # Validate all tensors have same length
+        validate_tensor_lengths_match(
+            bboxes, labels, scores, names=["bboxes", "labels", "scores"]
+        )
+
+        detections = []
+        for bbox, label, score in zip(bboxes, labels, scores):
+            detection = Detection(
+                bbox=BoundingBox(
+                    x1=float(bbox[0]),
+                    y1=float(bbox[1]),
+                    x2=float(bbox[2]),
+                    y2=float(bbox[3]),
+                ),
+                label=str(label),
+                confidence=float(score),
+            )
+            detections.append(detection)
+
+        return ObjectDetectionResult(
+            detections=detections,
+            image_width=image_width,
+            image_height=image_height,
+        )
+
+    def _extract_bboxes(self, tensor: TensorData) -> list[list[float]]:
+        """Extract bounding boxes from tensor, reshaping if necessary.
+
+        OIP may flatten [N, 4] tensor data into a flat list. This method
+        reconstructs the bounding box structure based on the tensor shape.
+
+        Args:
+            tensor: TensorData with bounding box data
+
+        Returns:
+            List of [x1, y1, x2, y2] coordinate lists
+        """
+        data = self._extract_tensor_data(tensor)
+        if not data:
+            return []
+
+        if isinstance(data[0], list):
+            return data
+
+        shape = tensor.shape
+        if len(shape) == 2 and shape[1] == 4:
+            num_boxes = shape[0]
+            return [data[i * 4 : (i + 1) * 4] for i in range(num_boxes)]
+
+        return data
+
+    def _extract_tensor_data(self, tensor: TensorData) -> list[Any]:
+        """Extract the actual data from a tensor, flattening if necessary.
+
+        Args:
+            tensor: TensorData object
+
+        Returns:
+            Flattened list of tensor values
+        """
+        data = tensor.data
+        if isinstance(data, list) and len(data) > 0 and isinstance(data[0], list):
+            return data[0]
+        return data

aissemble_inference_core/client/translators/summarization_translator.py

@@ -0,0 +1,147 @@
+###
+# #%L
+# aiSSEMBLE::Open Inference Protocol::Core
+# %%
+# Copyright (C) 2024 Booz Allen Hamilton Inc.
+# %%
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# #L%
+###
+"""Translator for text summarization tasks."""
+
+from aissemble_inference_core.client.oip_adapter import (
+    OipRequest,
+    OipResponse,
+    TensorData,
+)
+from aissemble_inference_core.client.results.summarization_result import (
+    SummarizationResult,
+)
+from aissemble_inference_core.client.translator import Translator
+
+
+class DefaultSummarizationTranslator(Translator[str, SummarizationResult]):
+    """Default translator for text summarization inference.
+
+    This translator converts input text to OIP format and parses OIP responses
+    into SummarizationResult objects. It is the only component aware of the
+    OIP JSON schema and tensor details for summarization.
+
+    The translator is configurable to work with different model implementations
+    that may use different tensor names or parameter conventions.
+    """
+
+    def __init__(
+        self,
+        input_name: str = "text",
+        output_name: str = "summary",
+    ):
+        """Initialize the summarization translator.
+
+        Args:
+            input_name: Name of the input tensor (default: "text")
+            output_name: Name of the output tensor (default: "summary")
+        """
+        self.input_name = input_name
+        self.output_name = output_name
+        self._original_text: str = ""
+
+    def preprocess(self, text: str) -> OipRequest:
+        """Convert input text to OIP request format.
+
+        Args:
+            text: The text to summarize
+
+        Returns:
+            OipRequest with text as BYTES tensor
+
+        Raises:
+            ValueError: If text is empty or invalid
+        """
+        if not text or not isinstance(text, str):
+            raise ValueError("Input text must be a non-empty string")
+
+        self._original_text = text
+
+        # Encode text as UTF-8 bytes
+        text_bytes = text.encode("utf-8")
+        text_str = text_bytes.decode("utf-8")
+
+        tensor = TensorData(
+            name=self.input_name,
+            shape=[1],
+            datatype="BYTES",
+            data=[text_str],
+        )
+
+        return OipRequest(inputs=[tensor])
+
+    def postprocess(self, response: OipResponse) -> SummarizationResult:
+        """Convert OIP response to SummarizationResult.
+
+        Args:
+            response: OIP response containing summary tensor
+
+        Returns:
+            SummarizationResult with extracted summary and metadata
+
+        Raises:
+            ValueError: If response format is invalid
+        """
+        outputs = {out.name: out for out in response.outputs}
+
+        if self.output_name not in outputs:
+            raise ValueError(
+                f"Expected output tensor '{self.output_name}' not found in response. "
+                f"Available: {list(outputs.keys())}"
+            )
+
+        summary_tensor = outputs[self.output_name]
+        summary_data = self._extract_text(summary_tensor)
+
+        if not summary_data:
+            raise ValueError("Summary output is empty")
+
+        return SummarizationResult(
+            summary=summary_data,
+            original_length=len(self._original_text),
+            summary_length=len(summary_data),
+            model_name=response.model_name,
+        )
+
+    def _extract_text(self, tensor: TensorData) -> str:
+        """Extract text string from tensor data.
+
+        Args:
+            tensor: TensorData containing text
+
+        Returns:
+            Extracted text string
+        """
+        if not tensor.data:
+            return ""
+
+        # Handle nested list structure [[text]] or flat [text]
+        data = tensor.data
+        if isinstance(data, list) and len(data) > 0:
+            if isinstance(data[0], list) and len(data[0]) > 0:
+                text = data[0][0]
+            else:
+                text = data[0]
+
+            if isinstance(text, str):
+                return text
+            elif isinstance(text, bytes):
+                return text.decode("utf-8")
+
+        return str(data)

aissemble_inference_core/client/translators/tensorflow_object_detection_translator.py

@@ -0,0 +1,231 @@
+###
+# #%L
+# aiSSEMBLE::Open Inference Protocol::Core
+# %%
+# Copyright (C) 2024 Booz Allen Hamilton Inc.
+# %%
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# #L%
+###
+"""TensorFlow Object Detection API translator for OIP object detection.
+
+This translator handles the TensorFlow Object Detection API / TensorFlow Serving
+tensor format, commonly used with models exported from TensorFlow Model Zoo.
+
+Key format characteristics:
+- Normalized coordinates [0, 1] instead of pixel coordinates
+- TensorFlow Serving-style tensor names (detection_boxes, detection_classes, etc.)
+- Batched shape conventions [1, N, 4] instead of [N, 4]
+- Integer class IDs requiring mapping to label names
+- Coordinate order: (ymin, xmin, ymax, xmax) instead of (x1, y1, x2, y2)
+
+Despite these differences from the default YOLO-style format, end users use the
+SAME client API - the translator handles all tensor-level complexity invisibly.
+
+Thread Safety:
+This translator is thread-safe and stateless. Image dimensions are stored in
+request parameters and retrieved from response parameters.
+"""
+
+import logging
+from typing import Any
+
+from aissemble_inference_core.client.oip_adapter import (
+    OipRequest,
+    OipResponse,
+    TensorData,
+)
+from aissemble_inference_core.client.results import (
+    BoundingBox,
+    Detection,
+    ObjectDetectionResult,
+)
+from aissemble_inference_core.client.translator import Translator
+from aissemble_inference_core.client.translators._image_utils import (
+    encode_image_for_oip,
+)
+from aissemble_inference_core.client.translators._tensor_utils import (
+    unbatch_tensor,
+    validate_required_tensors,
+    validate_tensor_lengths_match,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class TensorFlowObjectDetectionTranslator(Translator[Any, ObjectDetectionResult]):
+    """Translator for TensorFlow Object Detection API / TensorFlow Serving format.
+
+    This translator handles the standard TensorFlow Object Detection API output format,
+    commonly used by models exported from TensorFlow Model Zoo and served via TensorFlow
+    Serving, KServe, or other TF-compatible inference backends.
+
+    Expected output tensors:
+    - detection_boxes: [1, N, 4] with NORMALIZED coordinates (ymin, xmin, ymax, xmax)
+    - detection_classes: [1, N] with integer class IDs
+    - detection_scores: [1, N] with confidence scores
+    - num_detections: [1] with count of valid detections
+
+    Key differences from default YOLO-style format:
+    1. Normalized coords [0,1] vs pixel coords
+    2. Different coord order: (ymin, xmin, ymax, xmax) vs (x1, y1, x2, y2)
+    3. Batched shape [1, N, 4] vs unbatched [N, 4]
+    4. Integer class IDs vs string labels
+    5. Extra num_detections tensor
+
+    Despite these differences, users interact with the same ObjectDetectionResult
+    type, demonstrating complete tensor abstraction.
+    """
+
+    def __init__(
+        self,
+        class_names: dict[int, str] | None = None,
+    ):
+        """Initialize translator with class ID to name mapping.
+
+        Args:
+            class_names: Optional mapping of class IDs to human-readable names.
+                        If None, uses class IDs as labels (e.g., "class_17").
+        """
+        self._class_names = class_names or {}
+
+    def preprocess(self, input_data: Any) -> OipRequest:
+        """Preprocess image input into OipRequest.
+
+        Accepts PIL Image, numpy array, file path, or bytes.
+
+        Args:
+            input_data: Image data in supported format
+
+        Returns:
+            OipRequest with encoded image tensor and dimensions in parameters
+        """
+        image_bytes, width, height = encode_image_for_oip(input_data)
+
+        tensor = TensorData(
+            name="image",
+            shape=[1, len(image_bytes)],
+            datatype="BYTES",
+            data=[[image_bytes]],
+        )
+
+        # Store image dimensions in request parameters for stateless operation
+        return OipRequest(
+            inputs=[tensor],
+            parameters={"_image_width": width, "_image_height": height},
+        )
+
+    def postprocess(self, response: OipResponse) -> ObjectDetectionResult:
+        """Postprocess OipResponse into ObjectDetectionResult.
+
+        Handles the TensorFlow Object Detection API tensor format:
+        - Validates required tensors are present
+        - Extracts batch dimension [1, N, ...] -> [N, ...]
+        - Converts normalized coords to pixel coords
+        - Reorders coords from (ymin, xmin, ymax, xmax) to (x1, y1, x2, y2)
+        - Maps class IDs to names
+        - Respects num_detections count
+
+        Args:
+            response: OIP response with TensorFlow format tensors and dimensions
+
+        Returns:
+            ObjectDetectionResult - same as default translator!
+
+        Raises:
+            ValueError: If required tensors are missing or malformed
+        """
+        outputs = {out.name: out for out in response.outputs}
+
+        # Validate required tensors are present
+        required = ["detection_boxes", "detection_classes", "detection_scores"]
+        validate_required_tensors(outputs, required, "TensorFlow Object Detection API")
+
+        # Get image dimensions from response parameters
+        params = response.parameters or {}
+        image_width = params.get("_image_width")
+        image_height = params.get("_image_height")
+
+        if image_width is None or image_height is None:
+            raise ValueError(
+                "Image dimensions not found in response parameters. "
+                "Ensure the OIP adapter preserves request parameters in the response."
+            )
+
+        # Extract num_detections if present, otherwise use all detections
+        num_detections = self._get_num_detections(outputs)
+
+        # Extract and unbatch tensors [1, N] -> [N]
+        boxes = unbatch_tensor(outputs["detection_boxes"], num_detections)
+        class_ids = [
+            int(cid)
+            for cid in unbatch_tensor(outputs["detection_classes"], num_detections)
+        ]
+        scores = [
+            float(s)
+            for s in unbatch_tensor(outputs["detection_scores"], num_detections)
+        ]
+
+        # Validate all tensors have same length
+        validate_tensor_lengths_match(
+            boxes, class_ids, scores, names=["boxes", "classes", "scores"]
+        )
+
+        detections = []
+        for box, class_id, score in zip(boxes, class_ids, scores):
+            # Validate normalized coordinate range
+            ymin_norm, xmin_norm, ymax_norm, xmax_norm = box
+            if not all(0 <= coord <= 1 for coord in box):
+                logger.warning(
+                    f"Normalized coordinates outside [0,1] range: {box}. "
+                    f"Server may be misconfigured or model may be wrong type."
+                )
+
+            # Convert normalized coords to pixel coords and reorder
+            x1 = xmin_norm * image_width
+            y1 = ymin_norm * image_height
+            x2 = xmax_norm * image_width
+            y2 = ymax_norm * image_height
+
+            # Map class ID to name
+            label = self._class_names.get(class_id, f"class_{class_id}")
+
+            detection = Detection(
+                bbox=BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2),
+                label=label,
+                confidence=score,
+            )
+            detections.append(detection)
+
+        return ObjectDetectionResult(
+            detections=detections,
+            image_width=image_width,
+            image_height=image_height,
+        )
+
+    def _get_num_detections(self, outputs: dict[str, TensorData]) -> int | None:
+        """Extract number of valid detections from num_detections tensor if present.
+
+        Args:
+            outputs: Dictionary of output tensors
+
+        Returns:
+            Number of detections or None if tensor not present
+        """
+        if "num_detections" not in outputs:
+            return None
+
+        num_det_data = outputs["num_detections"].data
+        if isinstance(num_det_data, list):
+            return int(num_det_data[0]) if num_det_data else None
+        return int(num_det_data)