orient_express 2.4.0__tar.gz → 2.4.2__tar.gz

{orient_express-2.4.0 → orient_express-2.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.4.0
+Version: 2.4.2
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -593,6 +593,191 @@ class SearchResult:
 
 ---
 
+## Deployed Endpoint APIs
+
+When you upload a model with orient-express and deploy it to a Vertex AI endpoint, the actual HTTP API exposed by the endpoint is determined by the serving container image — not by your Python predictor code. Orient-express ships two such images:
+
+- `image-onnx` — serves the built-in ONNX predictor types (classification, detection, segmentation).
+- `xgboost-scikit-learn` — serves any joblib-loadable model (sklearn pipelines, xgboost, etc.).
+
+This section documents the request/response shape each image's endpoint exposes once deployed.
+
+### How the docker images connect to GCP endpoints
+
+The deployment flow is:
+
+1. **Train + export.** You build a predictor locally (e.g. `ClassificationPredictor("model.onnx", classes)`) or train a sklearn/xgboost model.
+2. **Upload.** `upload_model` / `upload_model_joblib` pushes the artifacts to GCS under `gs://<bucket>/models/<model_name>/<version>/` and registers a Vertex AI Model with `serving_container_image_uri` pointing at one of orient-express's images.
+3. **Deploy.** `vertex_model.deploy_to_endpoint(...)` (or the Vertex console) attaches the registered model to a Vertex AI Endpoint. Vertex starts the container with `AIP_STORAGE_URI` set to the GCS path from step 2, plus `MODEL_NAME` set to the model's display name.
+4. **Serve.** The container downloads the artifacts on startup, instantiates the right predictor via metadata, and listens on `/v1/models/<MODEL_NAME>:predict`.
+5. **Call.** Clients POST to `https://<region>-aiplatform.googleapis.com/v1/projects/<project>/locations/<region>/endpoints/<endpoint_id>:predict` with a Bearer token. Vertex routes the request into the container and returns the JSON response.
+
+Every endpoint accepts the same envelope:
+
+```json
+{
+  "instances": [...],
+  "parameters": {...}
+}
+```
+
+What goes in `instances` / `parameters`, and what comes back in `predictions`, is per-image and per-predictor — covered below.
+
+### ONNX Image Endpoint
+
+Image: `us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:<tag>`
+
+Common request shape across all ONNX predictor types:
+
+```json
+{
+  "instances": [
+    {"image": "<http(s) URL | gs:// URI | base64 | data URI>"}
+  ],
+  "parameters": {
+    "confidence": 0.5
+  }
+}
+```
+
+Common response envelope:
+
+```json
+{
+  "predictions": [
+    {"status": "success", ...predictor-specific fields...}
+  ]
+}
+```
+
+`status` values: `"success"`, `"failed to download image"`, or `"failed to get debug image"`. Malformed payloads return top-level `{"error": "Failed to decode input"}` instead of `predictions`.
+
+The predictor-specific fields differ by model type — covered in the subsections below.
+
+#### Classification
+
+For models uploaded as `ClassificationPredictor`. `parameters.confidence` is **not** honored (the predictor always returns the top class).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "class": "cat",
+  "score": 0.95,
+  "class_scores": {"cat": 0.95, "dog": 0.03, "bird": 0.02}
+}
+```
+
+No `debug_image` for classification (nothing meaningful to draw).
+
+#### Multi-label classification
+
+For models uploaded as `MultiLabelClassificationPredictor`. `parameters.confidence` is the per-class threshold for inclusion in `classes` (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": {
+    "classes": ["contains_cat", "contains_bird"],
+    "class_scores": {"contains_cat": 0.95, "contains_dog": 0.03, "contains_bird": 0.82}
+  },
+  "debug_image": null
+}
+```
+
+`debug_image` is always `null` for multi-label.
+
+#### Object detection
+
+For models uploaded as `BoundingBoxPredictor`. `parameters.confidence` filters detections below the threshold (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": [
+    {"class": "person", "score": 0.92, "bbox": {"x1": 100.5, "y1": 50.2, "x2": 300.8, "y2": 400.1}}
+  ],
+  "debug_image": "<base64 JPEG with boxes overlaid>"
+}
+```
+
+`bbox` coordinates are in pixels of the original (EXIF-corrected) image. `predictions` is an empty list when nothing clears the confidence threshold.
+
+#### Instance segmentation
+
+For models uploaded as `InstanceSegmentationPredictor`. `parameters.confidence` filters detections (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": [
+    {"class": "person", "score": 0.89, "bbox": {"x1": 100.5, "y1": 50.2, "x2": 300.8, "y2": 400.1}}
+  ],
+  "debug_image": "<base64 JPEG with masks overlaid>"
+}
+```
+
+Per-instance mask arrays are **not** included in the response by default (too large). The annotated mask overlay is baked into `debug_image`.
+
+#### Semantic segmentation
+
+For models uploaded as `SemanticSegmentationPredictor`. `parameters.confidence` is the per-pixel threshold above which a class is considered "valid" (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": {
+    "class_mask": "<base64 PNG, uint8, per-pixel class id>",
+    "valid_mask": "<base64 PNG, uint8, 0=below threshold, 1=above>"
+  },
+  "debug_image": "<base64 JPEG with color-coded overlay>"
+}
+```
+
+`class_mask` always paints every pixel with the argmax winner. `valid_mask` tells you which pixels actually cleared the confidence threshold — AND them together client-side to get the "real" segmentation.
+
+### XGBoost / scikit-learn Endpoint
+
+Image: `us-west1-docker.pkg.dev/shiftsmart-api/orient-express/xgboost-scikit-learn:<tag>`
+
+For models uploaded via `upload_model_joblib` — sklearn pipelines, xgboost models, or anything `joblib.load`-able with a `.predict(DataFrame)` method.
+
+Request shape — `instances` is a list of dicts, one row per input:
+
+```json
+{
+  "instances": [
+    {"pclass": 1, "sex": "female", "age": 29, "fare": 100.0, "embarked": "S"},
+    {"pclass": 3, "sex": "male", "age": 35, "fare": 8.05, "embarked": "S"}
+  ]
+}
+```
+
+The server constructs `pd.DataFrame(instances)` and calls `model.predict(df)` on it. The columns your pipeline expects must be present in each instance dict.
+
+Response shape — one prediction per input row:
+
+```json
+{
+  "predictions": [0, 1]
+}
+```
+
+Each element is whatever your `.predict()` returns — a class label for classifiers, a numeric value for regressors, an array for multi-output models.
+
+`parameters` is ignored — there's no per-request configuration for this image.
+
+---
+
 ## Color Schemes
 
 For predictors that support annotation (`BoundingBoxPredictor`, `InstanceSegmentationPredictor`, `SemanticSegmentationPredictor`), you can set a custom color scheme:

{orient_express-2.4.0 → orient_express-2.4.2}/README.md

@@ -566,6 +566,191 @@ class SearchResult:
 
 ---
 
+## Deployed Endpoint APIs
+
+When you upload a model with orient-express and deploy it to a Vertex AI endpoint, the actual HTTP API exposed by the endpoint is determined by the serving container image — not by your Python predictor code. Orient-express ships two such images:
+
+- `image-onnx` — serves the built-in ONNX predictor types (classification, detection, segmentation).
+- `xgboost-scikit-learn` — serves any joblib-loadable model (sklearn pipelines, xgboost, etc.).
+
+This section documents the request/response shape each image's endpoint exposes once deployed.
+
+### How the docker images connect to GCP endpoints
+
+The deployment flow is:
+
+1. **Train + export.** You build a predictor locally (e.g. `ClassificationPredictor("model.onnx", classes)`) or train a sklearn/xgboost model.
+2. **Upload.** `upload_model` / `upload_model_joblib` pushes the artifacts to GCS under `gs://<bucket>/models/<model_name>/<version>/` and registers a Vertex AI Model with `serving_container_image_uri` pointing at one of orient-express's images.
+3. **Deploy.** `vertex_model.deploy_to_endpoint(...)` (or the Vertex console) attaches the registered model to a Vertex AI Endpoint. Vertex starts the container with `AIP_STORAGE_URI` set to the GCS path from step 2, plus `MODEL_NAME` set to the model's display name.
+4. **Serve.** The container downloads the artifacts on startup, instantiates the right predictor via metadata, and listens on `/v1/models/<MODEL_NAME>:predict`.
+5. **Call.** Clients POST to `https://<region>-aiplatform.googleapis.com/v1/projects/<project>/locations/<region>/endpoints/<endpoint_id>:predict` with a Bearer token. Vertex routes the request into the container and returns the JSON response.
+
+Every endpoint accepts the same envelope:
+
+```json
+{
+  "instances": [...],
+  "parameters": {...}
+}
+```
+
+What goes in `instances` / `parameters`, and what comes back in `predictions`, is per-image and per-predictor — covered below.
+
+### ONNX Image Endpoint
+
+Image: `us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:<tag>`
+
+Common request shape across all ONNX predictor types:
+
+```json
+{
+  "instances": [
+    {"image": "<http(s) URL | gs:// URI | base64 | data URI>"}
+  ],
+  "parameters": {
+    "confidence": 0.5
+  }
+}
+```
+
+Common response envelope:
+
+```json
+{
+  "predictions": [
+    {"status": "success", ...predictor-specific fields...}
+  ]
+}
+```
+
+`status` values: `"success"`, `"failed to download image"`, or `"failed to get debug image"`. Malformed payloads return top-level `{"error": "Failed to decode input"}` instead of `predictions`.
+
+The predictor-specific fields differ by model type — covered in the subsections below.
+
+#### Classification
+
+For models uploaded as `ClassificationPredictor`. `parameters.confidence` is **not** honored (the predictor always returns the top class).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "class": "cat",
+  "score": 0.95,
+  "class_scores": {"cat": 0.95, "dog": 0.03, "bird": 0.02}
+}
+```
+
+No `debug_image` for classification (nothing meaningful to draw).
+
+#### Multi-label classification
+
+For models uploaded as `MultiLabelClassificationPredictor`. `parameters.confidence` is the per-class threshold for inclusion in `classes` (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": {
+    "classes": ["contains_cat", "contains_bird"],
+    "class_scores": {"contains_cat": 0.95, "contains_dog": 0.03, "contains_bird": 0.82}
+  },
+  "debug_image": null
+}
+```
+
+`debug_image` is always `null` for multi-label.
+
+#### Object detection
+
+For models uploaded as `BoundingBoxPredictor`. `parameters.confidence` filters detections below the threshold (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": [
+    {"class": "person", "score": 0.92, "bbox": {"x1": 100.5, "y1": 50.2, "x2": 300.8, "y2": 400.1}}
+  ],
+  "debug_image": "<base64 JPEG with boxes overlaid>"
+}
+```
+
+`bbox` coordinates are in pixels of the original (EXIF-corrected) image. `predictions` is an empty list when nothing clears the confidence threshold.
+
+#### Instance segmentation
+
+For models uploaded as `InstanceSegmentationPredictor`. `parameters.confidence` filters detections (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": [
+    {"class": "person", "score": 0.89, "bbox": {"x1": 100.5, "y1": 50.2, "x2": 300.8, "y2": 400.1}}
+  ],
+  "debug_image": "<base64 JPEG with masks overlaid>"
+}
+```
+
+Per-instance mask arrays are **not** included in the response by default (too large). The annotated mask overlay is baked into `debug_image`.
+
+#### Semantic segmentation
+
+For models uploaded as `SemanticSegmentationPredictor`. `parameters.confidence` is the per-pixel threshold above which a class is considered "valid" (default `0.5`).
+
+Per-image response:
+
+```json
+{
+  "status": "success",
+  "predictions": {
+    "class_mask": "<base64 PNG, uint8, per-pixel class id>",
+    "valid_mask": "<base64 PNG, uint8, 0=below threshold, 1=above>"
+  },
+  "debug_image": "<base64 JPEG with color-coded overlay>"
+}
+```
+
+`class_mask` always paints every pixel with the argmax winner. `valid_mask` tells you which pixels actually cleared the confidence threshold — AND them together client-side to get the "real" segmentation.
+
+### XGBoost / scikit-learn Endpoint
+
+Image: `us-west1-docker.pkg.dev/shiftsmart-api/orient-express/xgboost-scikit-learn:<tag>`
+
+For models uploaded via `upload_model_joblib` — sklearn pipelines, xgboost models, or anything `joblib.load`-able with a `.predict(DataFrame)` method.
+
+Request shape — `instances` is a list of dicts, one row per input:
+
+```json
+{
+  "instances": [
+    {"pclass": 1, "sex": "female", "age": 29, "fare": 100.0, "embarked": "S"},
+    {"pclass": 3, "sex": "male", "age": 35, "fare": 8.05, "embarked": "S"}
+  ]
+}
+```
+
+The server constructs `pd.DataFrame(instances)` and calls `model.predict(df)` on it. The columns your pipeline expects must be present in each instance dict.
+
+Response shape — one prediction per input row:
+
+```json
+{
+  "predictions": [0, 1]
+}
+```
+
+Each element is whatever your `.predict()` returns — a class label for classifiers, a numeric value for regressors, an array for multi-output models.
+
+`parameters` is ignored — there's no per-request configuration for this image.
+
+---
+
 ## Color Schemes
 
 For predictors that support annotation (`BoundingBoxPredictor`, `InstanceSegmentationPredictor`, `SemanticSegmentationPredictor`), you can set a custom color scheme:

{orient_express-2.4.0 → orient_express-2.4.2}/orient_express/predictors/predictor.py

@@ -46,7 +46,7 @@ class ImagePredictor(Predictor):
         self.model_path = model_path
 
     def get_serving_container_image_uri(self):
-        return "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:v2.4.0"
+        return "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:v2.4.2"
 
     def get_serving_container_health_route(self, model_name):
         return f"/v1/models/{model_name}"

{orient_express-2.4.0 → orient_express-2.4.2}/orient_express/predictors/semantic_segmentation.py

@@ -7,17 +7,19 @@ import torch
 import torch.nn.functional as F
 
 from .predictor import OnnxSessionWrapper, ImagePredictor
-from ..utils.image_processor import pil_to_opencv, opencv_to_pil
+from ..utils.image_processor import mask_to_base64, pil_to_opencv, opencv_to_pil
 
 
 @dataclass
 class SemanticSegmentationPrediction:
     class_mask: np.ndarray
+    valid_mask: np.ndarray
     conf_masks: np.ndarray
 
     def to_dict(self, include_conf_masks: bool = False):
         dict_repr = {
-            "class_mask": self.class_mask,
+            "class_mask": mask_to_base64(self.class_mask.astype(np.uint8)),
+            "valid_mask": mask_to_base64(self.valid_mask.astype(np.uint8)),
         }
         if include_conf_masks:
             dict_repr["conf_masks"] = self.conf_masks.tolist()
@@ -60,31 +62,38 @@ class SemanticSegmentationPredictor(ImagePredictor):
     backend_model = OnnxSemanticSegmentation
 
     def predict(
-        self, images: list[Image.Image]
+        self, images: list[Image.Image], confidence: float = 0.5
     ) -> list[SemanticSegmentationPrediction]:
         if not images:
             return []
         raw_outputs = self.model(images)
         outputs = []
         for masks in raw_outputs:
-            class_mask = np.argmax(masks, axis=0)
+            class_mask = np.argmax(masks, axis=0).astype(np.uint8)
+            valid_mask = np.max(masks, axis=0) >= confidence
             outputs.append(
                 SemanticSegmentationPrediction(
                     class_mask=class_mask,
+                    valid_mask=valid_mask,
                     conf_masks=masks,
                 )
             )
         return outputs
 
     def get_annotated_image(
-        self, image: Image.Image, mask: np.array, mask_opacity: float = 0.3
+        self,
+        image: Image.Image,
+        prediction: SemanticSegmentationPrediction,
+        mask_opacity: float = 0.3,
     ) -> Image.Image:
         opencv_image = pil_to_opencv(image)
         mask_overlay = opencv_image.copy()
 
         for class_id, class_name in self.classes.items():
             fill_color = self.color_scheme.get(class_name, (120, 120, 120))
-            mask_overlay[mask == class_id] = fill_color[:3]
+            mask_overlay[
+                (prediction.class_mask == class_id) & prediction.valid_mask
+            ] = fill_color[:3]
 
         annotated_image = cv2.addWeighted(
             mask_overlay, mask_opacity, opencv_image, 1 - mask_opacity, 0

{orient_express-2.4.0 → orient_express-2.4.2}/orient_express/utils/image_processor.py

@@ -1,6 +1,9 @@
 import base64
+import ipaddress
 import logging
+import socket
 from io import BytesIO
+from urllib.parse import urlparse
 
 import cv2
 import numpy as np
@@ -9,6 +12,34 @@ import requests
 
 from .gs import get_gcs_from_http_url, read_file_bytes
 
+DOWNLOAD_TIMEOUT_SECONDS = 30
+
+
+class UnsafeUrlError(ValueError):
+    pass
+
+
+def validate_url(http_url):
+    # Reject URLs that could be used for SSRF: the request runs server-side
+    # with the service account's network access, so only allow http(s) URLs
+    # that resolve to public IPs (blocks the GCE metadata server, localhost,
+    # and private/internal ranges).
+    parsed = urlparse(http_url)
+    if parsed.scheme not in ("http", "https"):
+        raise UnsafeUrlError(f"unsupported URL scheme: {parsed.scheme}")
+    if not parsed.hostname:
+        raise UnsafeUrlError("URL has no hostname")
+    try:
+        addr_infos = socket.getaddrinfo(parsed.hostname, None)
+    except socket.gaierror as e:
+        raise UnsafeUrlError(f"could not resolve host: {parsed.hostname}") from e
+    for addr_info in addr_infos:
+        ip = ipaddress.ip_address(addr_info[4][0])
+        if not ip.is_global:
+            raise UnsafeUrlError(
+                f"URL host {parsed.hostname} resolves to non-public address {ip}"
+            )
+
 
 def read_image_from_url(http_url, http_as_gcs=False) -> Image.Image:
     # Extract GSC URI from http link and download the file directly.
@@ -19,7 +50,8 @@ def read_image_from_url(http_url, http_as_gcs=False) -> Image.Image:
         if gs_uri:
             return read_image_from_gs(gs_uri)
 
-    response = requests.get(http_url)
+    validate_url(http_url)
+    response = requests.get(http_url, timeout=DOWNLOAD_TIMEOUT_SECONDS)
     response.raise_for_status()
     image = Image.open(BytesIO(response.content))
     return image
@@ -74,6 +106,12 @@ def image_to_base64(image):
     return base64.b64encode(bytes_content).decode("utf-8")
 
 
+def mask_to_base64(mask: np.ndarray) -> str:
+    buffered = BytesIO()
+    Image.fromarray(mask).save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
 def image_to_bytes(image):
     buffered = BytesIO()
     if image.mode == "RGBA":

{orient_express-2.4.0 → orient_express-2.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.4.0"
+version = "2.4.2"
 description = "A library to simplify model deployment to Vertex AI"
 authors = ["Alexey Zankevich <alex.zankevich@shiftsmart.com>", "Ian Myers <ian.myers@shiftsmart.com>"]
 readme = "README.md"
@@ -25,6 +25,9 @@ onnxruntime = {version = "^1.20", markers = "sys_platform == 'darwin' or platfor
 [tool.poetry.group.dev.dependencies]
 black = "24.10.0"
 pytest = "8.3.3"
+ipykernel = "^7.2.0"
+
+[tool.poetry.group.server.dependencies]
 kserve = "^0.16.0"
 python-json-logger = "^4.0.0"