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

aissemble_inference_core/client/registry.py

@@ -0,0 +1,201 @@
+###
+# #%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%
+###
+"""Module registry for dynamic OIP module discovery.
+
+This module provides a registry that discovers and loads OIP modules
+via Python entry points. Modules can register:
+- Runtimes (MLServer-compatible model runtimes)
+- Translators (OIP protocol translators)
+- Builders (task-specific inference builders)
+
+Entry point groups:
+- inference.runtimes: Model runtime implementations
+- inference.translators: Data format translators
+- inference.builders: Inference builder implementations
+"""
+
+import importlib.metadata
+import logging
+from typing import Any, Dict, List, Type
+
+logger = logging.getLogger(__name__)
+
+
+class ModuleRegistry:
+    """Registry for dynamically discovered OIP modules.
+
+    This registry uses Python entry points to discover installed modules.
+    Modules register themselves by declaring entry points in their
+    pyproject.toml files.
+
+    Example pyproject.toml entry points:
+        [project.entry-points."inference.runtimes"]
+        yolo = "aissemble_inference_yolo:YOLORuntime"
+
+        [project.entry-points."inference.translators"]
+        yolo = "aissemble_inference_yolo:YOLOTranslator"
+
+    Usage:
+        registry = ModuleRegistry.instance()
+        registry.discover()
+
+        # Get a specific runtime
+        runtime_class = registry.get_runtime("yolo")
+
+        # List all available modules
+        available = registry.list_available()
+    """
+
+    _instance: "ModuleRegistry | None" = None
+
+    def __init__(self):
+        """Initialize an empty registry."""
+        self._runtimes: Dict[str, Type[Any]] = {}
+        self._translators: Dict[str, Type[Any]] = {}
+        self._builders: Dict[str, Type[Any]] = {}
+        self._loaded = False
+
+    @classmethod
+    def instance(cls) -> "ModuleRegistry":
+        """Get the singleton registry instance.
+
+        Returns:
+            The global ModuleRegistry instance
+        """
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset the singleton instance (primarily for testing)."""
+        cls._instance = None
+
+    def discover(self) -> None:
+        """Load all installed OIP modules via entry points.
+
+        This method scans for entry points in the following groups:
+        - inference.runtimes: Model runtime implementations
+        - inference.translators: Data format translators
+        - inference.builders: Inference builder implementations
+
+        Entry points are loaded lazily - the actual class is only
+        imported when accessed via get_* methods.
+        """
+        if self._loaded:
+            return
+
+        entry_point_groups = [
+            ("inference.runtimes", self._runtimes),
+            ("inference.translators", self._translators),
+            ("inference.builders", self._builders),
+        ]
+
+        for group, registry in entry_point_groups:
+            try:
+                eps = importlib.metadata.entry_points(group=group)
+                for ep in eps:
+                    try:
+                        registry[ep.name] = ep.load()
+                        logger.debug(f"Loaded {group} entry point: {ep.name}")
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to load {group} entry point '{ep.name}': {e}"
+                        )
+            except Exception as e:
+                logger.warning(f"Failed to scan entry points for group '{group}': {e}")
+
+        self._loaded = True
+
+    def get_runtime(self, name: str) -> Type[Any] | None:
+        """Get a runtime class by name.
+
+        Args:
+            name: The registered name of the runtime
+
+        Returns:
+            The runtime class, or None if not found
+        """
+        self.discover()
+        return self._runtimes.get(name)
+
+    def get_translator(self, name: str) -> Type[Any] | None:
+        """Get a translator class by name.
+
+        Args:
+            name: The registered name of the translator
+
+        Returns:
+            The translator class, or None if not found
+        """
+        self.discover()
+        return self._translators.get(name)
+
+    def get_builder(self, name: str) -> Type[Any] | None:
+        """Get a builder class by name.
+
+        Args:
+            name: The registered name of the builder
+
+        Returns:
+            The builder class, or None if not found
+        """
+        self.discover()
+        return self._builders.get(name)
+
+    def list_available(self) -> Dict[str, List[str]]:
+        """List all discovered modules.
+
+        Returns:
+            Dictionary mapping category to list of available module names
+        """
+        self.discover()
+        return {
+            "runtimes": list(self._runtimes.keys()),
+            "translators": list(self._translators.keys()),
+            "builders": list(self._builders.keys()),
+        }
+
+    def register_runtime(self, name: str, runtime_class: Type[Any]) -> None:
+        """Manually register a runtime class.
+
+        Args:
+            name: Name to register the runtime under
+            runtime_class: The runtime class to register
+        """
+        self._runtimes[name] = runtime_class
+
+    def register_translator(self, name: str, translator_class: Type[Any]) -> None:
+        """Manually register a translator class.
+
+        Args:
+            name: Name to register the translator under
+            translator_class: The translator class to register
+        """
+        self._translators[name] = translator_class
+
+    def register_builder(self, name: str, builder_class: Type[Any]) -> None:
+        """Manually register a builder class.
+
+        Args:
+            name: Name to register the builder under
+            builder_class: The builder class to register
+        """
+        self._builders[name] = builder_class

aissemble_inference_core/client/results/__init__.py

@@ -0,0 +1,29 @@
+###
+# #%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 aissemble_inference_core.client.results.object_detection_result import (
+    BoundingBox,
+    Detection,
+    ObjectDetectionResult,
+)
+from aissemble_inference_core.client.results.summarization_result import (
+    SummarizationResult,
+)
+
+__all__ = ["BoundingBox", "Detection", "ObjectDetectionResult", "SummarizationResult"]

aissemble_inference_core/client/results/object_detection_result.py

@@ -0,0 +1,155 @@
+###
+# #%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 dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class BoundingBox:
+    """Represents a bounding box for object detection.
+
+    Attributes:
+        x1: Left x-coordinate
+        y1: Top y-coordinate
+        x2: Right x-coordinate
+        y2: Bottom y-coordinate
+    """
+
+    x1: float
+    y1: float
+    x2: float
+    y2: float
+
+    @property
+    def width(self) -> float:
+        """Calculate the width of the bounding box."""
+        return self.x2 - self.x1
+
+    @property
+    def height(self) -> float:
+        """Calculate the height of the bounding box."""
+        return self.y2 - self.y1
+
+    @property
+    def area(self) -> float:
+        """Calculate the area of the bounding box."""
+        return self.width * self.height
+
+    @property
+    def center(self) -> tuple[float, float]:
+        """Calculate the center point of the bounding box."""
+        return ((self.x1 + self.x2) / 2, (self.y1 + self.y2) / 2)
+
+
+@dataclass
+class Detection:
+    """Represents a single detected object.
+
+    Attributes:
+        bbox: Bounding box coordinates
+        label: Class label of the detected object
+        confidence: Confidence score (0-1)
+        label_id: Optional numeric ID for the label
+    """
+
+    bbox: BoundingBox
+    label: str
+    confidence: float
+    label_id: int | None = None
+
+
+@dataclass
+class ObjectDetectionResult:
+    """Value object representing the result of an object detection inference.
+
+    Attributes:
+        detections: List of detected objects
+        image_width: Width of the input image
+        image_height: Height of the input image
+    """
+
+    detections: list[Detection]
+    image_width: int
+    image_height: int
+
+    def filter_by_confidence(self, threshold: float) -> "ObjectDetectionResult":
+        """Filter detections by confidence threshold.
+
+        Args:
+            threshold: Minimum confidence score (0-1)
+
+        Returns:
+            New ObjectDetectionResult with filtered detections
+        """
+        filtered = [d for d in self.detections if d.confidence >= threshold]
+        return ObjectDetectionResult(
+            detections=filtered,
+            image_width=self.image_width,
+            image_height=self.image_height,
+        )
+
+    def filter_by_label(self, labels: list[str]) -> "ObjectDetectionResult":
+        """Filter detections by label.
+
+        Args:
+            labels: List of labels to keep
+
+        Returns:
+            New ObjectDetectionResult with filtered detections
+        """
+        filtered = [d for d in self.detections if d.label in labels]
+        return ObjectDetectionResult(
+            detections=filtered,
+            image_width=self.image_width,
+            image_height=self.image_height,
+        )
+
+    def draw_on_image(self, image: Any) -> Any:
+        """Draw bounding boxes on the image.
+
+        Args:
+            image: PIL Image or numpy array
+
+        Returns:
+            Image with drawn bounding boxes
+
+        Raises:
+            ImportError: If PIL is not installed
+        """
+        try:
+            from PIL import Image, ImageDraw
+        except ImportError as e:
+            raise ImportError(
+                "PIL (Pillow) is required for drawing. Install with: pip install Pillow"
+            ) from e
+
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        draw = ImageDraw.Draw(image)
+        for detection in self.detections:
+            bbox = detection.bbox
+            draw.rectangle(
+                [(bbox.x1, bbox.y1), (bbox.x2, bbox.y2)], outline="red", width=2
+            )
+            text = f"{detection.label}: {detection.confidence:.2f}"
+            draw.text((bbox.x1, bbox.y1 - 10), text, fill="red")
+
+        return image

aissemble_inference_core/client/results/summarization_result.py

@@ -0,0 +1,78 @@
+###
+# #%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%
+###
+"""Value objects for text summarization results."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class SummarizationResult:
+    """Result of a text summarization inference.
+
+    This class encapsulates the output from a summarization model, including
+    the generated summary text and optional metadata about the operation.
+
+    Attributes:
+        summary: The generated summary text
+        original_length: Character count of the original input text
+        summary_length: Character count of the generated summary
+        model_name: Optional name of the model that generated the summary
+        input_tokens: Optional count of input tokens (for token-based models)
+        output_tokens: Optional count of output tokens (for token-based models)
+    """
+
+    summary: str
+    original_length: int
+    summary_length: int
+    model_name: str | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+
+    @property
+    def compression_ratio(self) -> float:
+        """Calculate the compression ratio (original length / summary length).
+
+        Returns:
+            Ratio of original to summary length. Higher values indicate more compression.
+            Returns 0.0 if summary is empty.
+
+        Example:
+            >>> result = SummarizationResult(
+            ...     summary="Short summary",
+            ...     original_length=1000,
+            ...     summary_length=100
+            ... )
+            >>> result.compression_ratio
+            10.0
+        """
+        if self.summary_length == 0:
+            return 0.0
+        return self.original_length / self.summary_length
+
+    def __repr__(self) -> str:
+        """Return a detailed string representation."""
+        return (
+            f"SummarizationResult("
+            f"summary={self.summary[:50]!r}..., "
+            f"original_length={self.original_length}, "
+            f"summary_length={self.summary_length}, "
+            f"compression_ratio={self.compression_ratio:.2f}x"
+            f")"
+        )

aissemble_inference_core/client/translator.py

@@ -0,0 +1,57 @@
+###
+# #%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 abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from aissemble_inference_core.client.oip_adapter import OipRequest, OipResponse
+
+I = TypeVar("I")
+O = TypeVar("O")
+
+
+class Translator(ABC, Generic[I, O]):
+    """
+    Directly inspired by DJL, this is the ONLY class that knows about the OIP json schema,
+    tensor shapes, inputs, parameters, framing, etc. It is intended to be pluggable and
+    testable in isolation. It also provides tolerable reader behavior, where unknown
+    fields can be ignored.
+    """
+
+    @abstractmethod
+    def preprocess(self, input: I) -> OipRequest:
+        """
+        Preprocesses the input to create an OipRequest.
+
+        :param input: The input data of type I.
+        :return: An OipRequest object.
+        """
+        # TODO: Implement OipRequest creation
+        pass
+
+    @abstractmethod
+    def postprocess(self, response: OipResponse) -> O:
+        """
+        Postprocesses the OipResponse to produce the output.
+
+        :param response: The OipResponse object.
+        :return: The output data of type O.
+        """
+        # TODO: Implement OipResponse creation
+        pass

aissemble_inference_core/client/translators/__init__.py

@@ -0,0 +1,34 @@
+###
+# #%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 aissemble_inference_core.client.translators.object_detection_translator import (
+    DefaultObjectDetectionTranslator,
+)
+from aissemble_inference_core.client.translators.summarization_translator import (
+    DefaultSummarizationTranslator,
+)
+from aissemble_inference_core.client.translators.tensorflow_object_detection_translator import (
+    TensorFlowObjectDetectionTranslator,
+)
+
+__all__ = [
+    "DefaultObjectDetectionTranslator",
+    "DefaultSummarizationTranslator",
+    "TensorFlowObjectDetectionTranslator",
+]

aissemble_inference_core/client/translators/_image_utils.py

@@ -0,0 +1,75 @@
+###
+# #%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 image preprocessing in translators."""
+
+import base64
+from io import BytesIO
+from typing import Any
+
+
+def encode_image_for_oip(input_data: Any) -> tuple[str, int, int]:
+    """Encode image to base64 string for OIP transport.
+
+    Accepts various input formats:
+    - PIL Image
+    - numpy array
+    - file path (string)
+    - bytes
+
+    Args:
+        input_data: Image in various formats
+
+    Returns:
+        Tuple of (base64_string, width, height)
+
+    Raises:
+        ImportError: If PIL is not installed
+        ValueError: If input type is not supported
+    """
+    try:
+        from PIL import Image
+    except ImportError as e:
+        raise ImportError(
+            "PIL (Pillow) is required for image handling. "
+            "Install with: pip install Pillow"
+        ) from e
+
+    if isinstance(input_data, str):
+        image = Image.open(input_data)
+    elif isinstance(input_data, bytes):
+        image = Image.open(BytesIO(input_data))
+    elif hasattr(input_data, "mode"):
+        image = input_data
+    else:
+        import numpy as np
+
+        if isinstance(input_data, np.ndarray):
+            image = Image.fromarray(input_data)
+        else:
+            raise ValueError(f"Unsupported input type: {type(input_data)}")
+
+    width, height = image.size
+
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    image_bytes = buffer.getvalue()
+    encoded = base64.b64encode(image_bytes).decode("utf-8")
+
+    return encoded, width, height