tensorplate-python 0.1.3__py3-none-any.whl

tensorplate/preprocess.py

@@ -0,0 +1,165 @@
+"""Client-side image preprocessing for detector models.
+
+Decodes an image (path, bytes, or HWC uint8 ndarray), letterboxes it to
+the model's input size, normalizes, orders channels, and lays it out as an
+NCHW float32 input tensor. A :class:`LetterboxTransform` records the scale
+and padding so detections map back to source-image pixels.
+
+numpy and Pillow (the ``tensorplate-python[vision]`` extra) are imported
+lazily so the core package stays importable without them.
+"""
+
+from __future__ import annotations
+
+import io
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from tensorplate.tensors import DType, Layout, TensorInput
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+def _clamp(value: float, low: float, high: float) -> float:
+    return max(low, min(high, value))
+
+
+@dataclass(frozen=True)
+class LetterboxTransform:
+    """Maps model-input pixel coordinates back to source-image pixels."""
+
+    src_height: int
+    src_width: int
+    scale_x: float
+    scale_y: float
+    pad_x: float
+    pad_y: float
+    input_height: int
+    input_width: int
+
+    def map_box_to_source(
+        self, x1: float, y1: float, x2: float, y2: float
+    ) -> tuple[float, float, float, float]:
+        """Map a box from model-input pixels to source pixels (clamped to bounds)."""
+        return (
+            _clamp((x1 - self.pad_x) / self.scale_x, 0.0, float(self.src_width)),
+            _clamp((y1 - self.pad_y) / self.scale_y, 0.0, float(self.src_height)),
+            _clamp((x2 - self.pad_x) / self.scale_x, 0.0, float(self.src_width)),
+            _clamp((y2 - self.pad_y) / self.scale_y, 0.0, float(self.src_height)),
+        )
+
+
+@dataclass(frozen=True)
+class PreprocessConfig:
+    """Preprocessing options. Defaults target YOLO-style NCHW float32 input."""
+
+    input_size: tuple[int, int] = (640, 640)  # (height, width)
+    letterbox: bool = True
+    channel_order: str = "rgb"  # "rgb" or "bgr"
+    channels_first: bool = True  # NCHW when True, else NHWC
+    scale: float = 1.0 / 255.0
+    mean: tuple[float, ...] | None = None
+    std: tuple[float, ...] | None = None
+    dtype: DType = DType.FLOAT32
+    pad_value: int = 114
+    input_name: str = "images"
+
+
+def preprocess(
+    image: str | bytes | Path | np.ndarray,
+    config: PreprocessConfig | None = None,
+) -> tuple[TensorInput, LetterboxTransform]:
+    """Preprocess ``image`` into a model input tensor and a back-mapping transform.
+
+    ``image`` may be a filesystem path, encoded image bytes, or an HWC
+    uint8 ndarray (assumed RGB). Decoding path/bytes requires Pillow.
+    """
+    import numpy as np
+
+    cfg = config if config is not None else PreprocessConfig()
+    if cfg.channel_order not in ("rgb", "bgr"):
+        raise ValueError(f"channel_order must be 'rgb' or 'bgr', got {cfg.channel_order!r}")
+
+    source = _decode_to_hwc_rgb_uint8(image)
+    input_height, input_width = cfg.input_size
+    resized, scale_x, scale_y, pad_x, pad_y = _resize(source, cfg)
+    transform = LetterboxTransform(
+        src_height=int(source.shape[0]),
+        src_width=int(source.shape[1]),
+        scale_x=scale_x,
+        scale_y=scale_y,
+        pad_x=pad_x,
+        pad_y=pad_y,
+        input_height=input_height,
+        input_width=input_width,
+    )
+
+    array = resized.astype("float32")
+    if cfg.channel_order == "bgr":
+        array = array[:, :, ::-1]
+    array = array * cfg.scale
+    if cfg.mean is not None:
+        array = array - np.asarray(cfg.mean, dtype="float32")
+    if cfg.std is not None:
+        array = array / np.asarray(cfg.std, dtype="float32")
+    if cfg.channels_first:
+        array = np.transpose(array, (2, 0, 1))
+    array = array[np.newaxis, ...]
+
+    tensor = TensorInput.from_numpy(cfg.input_name, array, dtype=cfg.dtype, layout=Layout.ROW_MAJOR)
+    return tensor, transform
+
+
+def _decode_to_hwc_rgb_uint8(image: str | bytes | Path | np.ndarray) -> np.ndarray:
+    if isinstance(image, (str, Path)):
+        return _decode_pil(image)
+    if isinstance(image, (bytes, bytearray)):
+        return _decode_pil(io.BytesIO(bytes(image)))
+    import numpy as np
+
+    array = np.asarray(image)
+    if array.ndim != 3 or array.shape[2] != 3:
+        raise ValueError(
+            f"ndarray image must be HWC with 3 channels, got shape {tuple(array.shape)}"
+        )
+    if array.dtype != np.uint8:
+        raise ValueError(f"ndarray image must have dtype uint8, got {array.dtype}")
+    return np.ascontiguousarray(array)
+
+
+def _decode_pil(source: str | Path | io.BytesIO) -> np.ndarray:
+    try:
+        from PIL import Image
+    except ModuleNotFoundError as exc:  # pragma: no cover - exercised only without Pillow
+        raise RuntimeError(
+            "image decoding requires Pillow; install `tensorplate-python[vision]`"
+        ) from exc
+    import numpy as np
+
+    with Image.open(source) as handle:
+        return np.asarray(handle.convert("RGB"))
+
+
+def _resize(
+    source: np.ndarray, cfg: PreprocessConfig
+) -> tuple[np.ndarray, float, float, float, float]:
+    import numpy as np
+    from PIL import Image
+
+    src_height, src_width = int(source.shape[0]), int(source.shape[1])
+    input_height, input_width = cfg.input_size
+    pil = Image.fromarray(source)
+    if cfg.letterbox:
+        scale = min(input_width / src_width, input_height / src_height)
+        new_width = max(1, round(src_width * scale))
+        new_height = max(1, round(src_height * scale))
+        resized = np.asarray(pil.resize((new_width, new_height), Image.Resampling.BILINEAR))
+        canvas = np.full((input_height, input_width, 3), cfg.pad_value, dtype="uint8")
+        pad_x = (input_width - new_width) // 2
+        pad_y = (input_height - new_height) // 2
+        canvas[pad_y : pad_y + new_height, pad_x : pad_x + new_width] = resized
+        return canvas, float(scale), float(scale), float(pad_x), float(pad_y)
+    resized = np.asarray(pil.resize((input_width, input_height), Image.Resampling.BILINEAR))
+    return resized, input_width / src_width, input_height / src_height, 0.0, 0.0

tensorplate/serving.py

@@ -0,0 +1,291 @@
+"""Synchronous client for the TensorPlate v0.1 serving HTTP envelope."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from tensorplate.client import (
+    DEFAULT_TIMEOUT_S,
+    ResolvedEndpoint,
+    http_request,
+    resolve_serving_url,
+)
+from tensorplate.errors import (
+    ErrorCode,
+    ProtocolError,
+    ServingError,
+    UnsupportedSchemaVersionError,
+)
+from tensorplate.tensors import DType, Layout, TensorInput, TensorOutput
+
+#: Serving envelope schema_version this SDK speaks.
+SCHEMA_VERSION = "0.1"
+
+_HEALTH_PATH = "/health"
+
+
+@dataclass(frozen=True)
+class Timing:
+    """Optional per-request timing reported by the worker."""
+
+    queue_latency_ns: int | None = None
+    execution_latency_ns: int | None = None
+    total_latency_ns: int | None = None
+
+
+@dataclass(frozen=True)
+class InferResult:
+    """Parsed v0.1 ``success`` response."""
+
+    request_id: str
+    outputs: tuple[TensorOutput, ...]
+    correlation_id: str | None = None
+    timing: Timing | None = None
+
+    def output(self, name: str) -> TensorOutput:
+        """Return the output tensor named ``name`` or raise ``KeyError``."""
+        for tensor in self.outputs:
+            if tensor.name == name:
+                return tensor
+        raise KeyError(f"no output named {name!r}")
+
+
+@dataclass(frozen=True)
+class HealthSnapshot:
+    """Parsed ``GET /health`` response."""
+
+    state: str
+    endpoint: str
+    backend: str
+    active_model_id: str | None = None
+    last_error_code: ErrorCode | None = None
+    last_error_message: str | None = None
+    queue_depth: int | None = None
+    in_flight: int | None = None
+
+    @property
+    def is_ready(self) -> bool:
+        """True only for the ``ready`` state (``degraded`` is not ready)."""
+        return self.state == "ready"
+
+
+class ServingClient:
+    """Synchronous client for the v0.1 ``/infer`` serving HTTP envelope.
+
+    The endpoint is resolved once at construction with the same precedence
+    as ``tensorplate infer`` (explicit URL, then CLI profile, then
+    read-only agent discovery, then the loopback default). The client is
+    model-class-neutral so higher-level helpers can build on it.
+    """
+
+    def __init__(
+        self,
+        serving_url: str | None = None,
+        *,
+        profile: str | None = None,
+        config_path: str | None = None,
+        timeout: float = DEFAULT_TIMEOUT_S,
+        discover: bool = True,
+    ) -> None:
+        self._endpoint = resolve_serving_url(
+            serving_url,
+            profile=profile,
+            config_path=config_path,
+            timeout=timeout,
+            discover=discover,
+        )
+        self._timeout = timeout
+
+    @property
+    def endpoint(self) -> ResolvedEndpoint:
+        """The resolved serving endpoint and how it was resolved."""
+        return self._endpoint
+
+    @staticmethod
+    def tensor_input(
+        name: str,
+        data: bytes,
+        dtype: DType | str,
+        shape: Sequence[int],
+        *,
+        layout: Layout | str = Layout.ROW_MAJOR,
+    ) -> TensorInput:
+        """Build a :class:`TensorInput` from raw bytes and metadata."""
+        return TensorInput(
+            name=name,
+            dtype=DType(dtype),
+            shape=tuple(shape),
+            data=data,
+            layout=Layout(layout),
+        )
+
+    def infer(
+        self,
+        endpoint: str,
+        inputs: Sequence[TensorInput],
+        *,
+        deadline_ms: int | None = None,
+        correlation_id: str | None = None,
+    ) -> InferResult:
+        """Run a single synchronous inference against ``endpoint``.
+
+        Raises :class:`~tensorplate.errors.ServingError` for a typed
+        ``failure`` envelope, :class:`~tensorplate.errors.TransportError`
+        / :class:`~tensorplate.errors.RequestTimeoutError` for transport
+        failures, and :class:`~tensorplate.errors.ProtocolError` /
+        :class:`~tensorplate.errors.UnsupportedSchemaVersionError` for a
+        malformed or unsupported response.
+        """
+        if not endpoint:
+            raise ValueError("infer endpoint must be non-empty")
+        if not inputs:
+            raise ValueError("infer requires at least one input tensor")
+        if deadline_ms is not None and (
+            not isinstance(deadline_ms, int) or isinstance(deadline_ms, bool) or deadline_ms < 1
+        ):
+            raise ValueError("deadline_ms must be an integer greater than or equal to 1")
+        if correlation_id is not None and not correlation_id:
+            raise ValueError("correlation_id must be non-empty when provided")
+        request_id = str(uuid.uuid4())
+        request_body: dict[str, object] = {
+            "schema_version": SCHEMA_VERSION,
+            "request_id": request_id,
+            "endpoint": endpoint,
+            "inputs": [tensor.to_named_input() for tensor in inputs],
+        }
+        if deadline_ms is not None:
+            request_body["deadline_ms"] = deadline_ms
+        headers: dict[str, str] = {}
+        if correlation_id is not None:
+            request_body["metadata"] = {"correlation_id": correlation_id}
+            headers["X-Correlation-Id"] = correlation_id
+        payload = json.dumps(request_body).encode("utf-8")
+        status_code, raw = http_request(
+            "POST", self._endpoint.url, body=payload, headers=headers, timeout=self._timeout
+        )
+        return _parse_infer_response(raw, status_code, self._endpoint.url)
+
+    def health(self) -> HealthSnapshot:
+        """Read ``GET /health`` and parse the worker's readiness snapshot.
+
+        The body is parsed regardless of the HTTP status (``ready`` and
+        ``degraded`` return 200; the rest return 503), so inspect
+        :attr:`HealthSnapshot.state`, not just reachability.
+        """
+        url = f"{self._endpoint.origin}{_HEALTH_PATH}"
+        _status_code, raw = http_request("GET", url, timeout=self._timeout)
+        payload = _decode_json_object(raw, url)
+        _require_supported_schema(payload)
+        return _parse_health(payload, url)
+
+
+def _decode_json_object(raw: bytes, url: str) -> dict[str, object]:
+    try:
+        parsed = json.loads(raw)
+    except ValueError as exc:
+        raise ProtocolError(f"serving response from {url!r} is not valid JSON: {exc}") from exc
+    if not isinstance(parsed, dict):
+        raise ProtocolError(f"serving response from {url!r} is not a JSON object")
+    return parsed
+
+
+def _require_supported_schema(payload: dict[str, object]) -> None:
+    version = payload.get("schema_version")
+    if version != SCHEMA_VERSION:
+        raise UnsupportedSchemaVersionError(
+            version if isinstance(version, str) else None, SCHEMA_VERSION
+        )
+
+
+def _parse_infer_response(raw: bytes, status_code: int, url: str) -> InferResult:
+    payload = _decode_json_object(raw, url)
+    _require_supported_schema(payload)
+    status = payload.get("status")
+    if status == "failure":
+        raise _serving_error(payload)
+    if status != "success":
+        raise ProtocolError(
+            f"serving response from {url!r} has unexpected status {status!r} (HTTP {status_code})"
+        )
+    outputs_obj = payload.get("outputs")
+    if not isinstance(outputs_obj, list) or not outputs_obj:
+        raise ProtocolError(f"serving success response from {url!r} is missing non-empty 'outputs'")
+    outputs = tuple(TensorOutput.from_named_output(item) for item in outputs_obj)
+    request_id = payload.get("request_id")
+    if not isinstance(request_id, str):
+        raise ProtocolError(f"serving success response from {url!r} is missing 'request_id'")
+    correlation_id = payload.get("correlation_id")
+    return InferResult(
+        request_id=request_id,
+        outputs=outputs,
+        correlation_id=correlation_id if isinstance(correlation_id, str) else None,
+        timing=_parse_timing(payload.get("timing")),
+    )
+
+
+def _serving_error(payload: dict[str, object]) -> ServingError:
+    request_id = payload.get("request_id")
+    request_id = request_id if isinstance(request_id, str) else None
+    error_obj = payload.get("error")
+    if not isinstance(error_obj, dict):
+        return ServingError(
+            ErrorCode.INTERNAL,
+            "serving worker returned a failure without a typed error",
+            request_id=request_id,
+        )
+    try:
+        code = ErrorCode(error_obj.get("code"))
+    except ValueError:
+        code = ErrorCode.INTERNAL
+    message_obj = error_obj.get("message")
+    context_obj = error_obj.get("context")
+    return ServingError(
+        code,
+        message_obj if isinstance(message_obj, str) else "serving worker returned an error",
+        context=context_obj if isinstance(context_obj, str) else None,
+        request_id=request_id,
+    )
+
+
+def _parse_health(payload: dict[str, object], url: str) -> HealthSnapshot:
+    state = payload.get("state")
+    if not isinstance(state, str):
+        raise ProtocolError(f"/health response from {url!r} is missing 'state'")
+    endpoint = payload.get("endpoint")
+    backend = payload.get("backend")
+    code_obj = payload.get("last_error_code")
+    try:
+        last_error_code = ErrorCode(code_obj) if code_obj is not None else None
+    except ValueError:
+        last_error_code = None
+    return HealthSnapshot(
+        state=state,
+        endpoint=endpoint if isinstance(endpoint, str) else "",
+        backend=backend if isinstance(backend, str) else "",
+        active_model_id=_optional_str(payload.get("active_model_id")),
+        last_error_code=last_error_code,
+        last_error_message=_optional_str(payload.get("last_error_message")),
+        queue_depth=_optional_int(payload.get("queue_depth")),
+        in_flight=_optional_int(payload.get("in_flight")),
+    )
+
+
+def _parse_timing(obj: object) -> Timing | None:
+    if not isinstance(obj, dict):
+        return None
+    return Timing(
+        queue_latency_ns=_optional_int(obj.get("queue_latency_ns")),
+        execution_latency_ns=_optional_int(obj.get("execution_latency_ns")),
+        total_latency_ns=_optional_int(obj.get("total_latency_ns")),
+    )
+
+
+def _optional_int(value: object) -> int | None:
+    return value if isinstance(value, int) else None
+
+
+def _optional_str(value: object) -> str | None:
+    return value if isinstance(value, str) else None