orient_express 2.3.1__tar.gz → 2.4.1__tar.gz

{orient_express-2.3.1 → orient_express-2.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.3.1
+Version: 2.4.1
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -46,11 +46,13 @@ pip install orient_express
 ```
 
 For local development:
+
 ```bash
 pip install -e .
 ```
 
 Or with Poetry:
+
 ```bash
 poetry install
 ```
@@ -102,8 +104,8 @@ vertex_model.deploy_to_endpoint(
 
 # remote prediction API depends on the endpoint container deployed with the model
 predictions = vertex_model.remote_predict(
-    [{"image": "https://storage.googleapis.com/ssm-media-uploads/example.jpg"}], 
-    endpoint_name="my-classifier-endpoint"
+    "my-classifier-endpoint",
+    [{"image": "https://storage.googleapis.com/ssm-media-uploads/example.jpg"}],
 )
 ```
 
@@ -179,19 +181,20 @@ predictions = local_predictor.predict(X_test)
 ### Platform Support Matrix
 
 | Platform | Architecture | ONNX Runtime Package | CUDA Available |
-|----------|--------------|---------------------|----------------|
-| Linux    | x86_64       | onnxruntime-gpu     | Yes            |
-| Linux    | aarch64      | onnxruntime         | No             |
-| Windows  | x64 (AMD64)  | onnxruntime-gpu     | Yes            |
-| Windows  | ARM64        | onnxruntime         | No             |
-| macOS    | x86_64       | onnxruntime         | No             |
-| macOS    | arm64        | onnxruntime         | No             |
+| -------- | ------------ | -------------------- | -------------- |
+| Linux    | x86_64       | onnxruntime-gpu      | Yes            |
+| Linux    | aarch64      | onnxruntime          | No             |
+| Windows  | x64 (AMD64)  | onnxruntime-gpu      | Yes            |
+| Windows  | ARM64        | onnxruntime          | No             |
+| macOS    | x86_64       | onnxruntime          | No             |
+| macOS    | arm64        | onnxruntime          | No             |
 
 The appropriate package is installed automatically based on your platform.
 
 ### Selecting CPU vs CUDA Execution
 
 When loading a predictor, use the `device` parameter to specify the execution provider:
+
 ```python
 from orient_express.predictors import ObjectDetectionPredictor
 
@@ -203,6 +206,7 @@ predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cuda")
 ```
 
 When using a Vertex AI model:
+
 ```python
 # CPU inference
 predictor = model.get_local_predictor(device="cpu")
@@ -540,11 +544,8 @@ index = build_vector_index(
     num_workers=8,               # parallel image loading
 )
 
-# Aggregate to one centroid per label
-centroid_index = index.aggregate()
-
 # Save and load
-centroid_index.dump("/path/to/artifact_dir")
+index.dump("/path/to/artifact_dir")
 
 from orient_express.predictors import get_predictor
 loaded_index = get_predictor("/path/to/artifact_dir")
@@ -560,16 +561,23 @@ batch_results = loaded_index.search_batch(query_matrix, k=5)
 
 #### Multi-label support
 
-Vectors can have multiple labels. This is useful when a single visual cluster maps to multiple SKUs:
+Vectors can have composite labels (use tuples). This is useful when a single visual cluster maps to multiple things:
 
 ```python
 index = VectorIndex(
     vectors=feature_matrix,
-    labels=[["sku_101", "sku_102"], ["sku_103"], ["sku_104", "sku_105"]],
+    labels=[("sku_101", "sku_102"), ("sku_103",)],
 )
+```
 
-# aggregate() computes one centroid per unique label
-aggregated = index.aggregate()  # 5 vectors, one per SKU
+#### Per-label aggregation
+
+Vector indices in which labels are not unique can be aggregated so that each label has a single centroid.
+If `per_label=True` and the labels are composite (tuples), then the labels will be unpacked and aggregated separately.
+
+```python
+aggregated = index.aggregate(per_label=True)  # 3 vectors, one per label element ["sku_101", "sku_102", "sku_103"]
+aggregated = index.aggregate(per_label=False)  # 2 vectors, one per composite label  [("sku_101", "sku_102"), ("sku_103")]
 ```
 
 #### Output Structure
@@ -585,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.3.1 → orient_express-2.4.1}/README.md

@@ -19,11 +19,13 @@ pip install orient_express
 ```
 
 For local development:
+
 ```bash
 pip install -e .
 ```
 
 Or with Poetry:
+
 ```bash
 poetry install
 ```
@@ -75,8 +77,8 @@ vertex_model.deploy_to_endpoint(
 
 # remote prediction API depends on the endpoint container deployed with the model
 predictions = vertex_model.remote_predict(
-    [{"image": "https://storage.googleapis.com/ssm-media-uploads/example.jpg"}], 
-    endpoint_name="my-classifier-endpoint"
+    "my-classifier-endpoint",
+    [{"image": "https://storage.googleapis.com/ssm-media-uploads/example.jpg"}],
 )
 ```
 
@@ -152,19 +154,20 @@ predictions = local_predictor.predict(X_test)
 ### Platform Support Matrix
 
 | Platform | Architecture | ONNX Runtime Package | CUDA Available |
-|----------|--------------|---------------------|----------------|
-| Linux    | x86_64       | onnxruntime-gpu     | Yes            |
-| Linux    | aarch64      | onnxruntime         | No             |
-| Windows  | x64 (AMD64)  | onnxruntime-gpu     | Yes            |
-| Windows  | ARM64        | onnxruntime         | No             |
-| macOS    | x86_64       | onnxruntime         | No             |
-| macOS    | arm64        | onnxruntime         | No             |
+| -------- | ------------ | -------------------- | -------------- |
+| Linux    | x86_64       | onnxruntime-gpu      | Yes            |
+| Linux    | aarch64      | onnxruntime          | No             |
+| Windows  | x64 (AMD64)  | onnxruntime-gpu      | Yes            |
+| Windows  | ARM64        | onnxruntime          | No             |
+| macOS    | x86_64       | onnxruntime          | No             |
+| macOS    | arm64        | onnxruntime          | No             |
 
 The appropriate package is installed automatically based on your platform.
 
 ### Selecting CPU vs CUDA Execution
 
 When loading a predictor, use the `device` parameter to specify the execution provider:
+
 ```python
 from orient_express.predictors import ObjectDetectionPredictor
 
@@ -176,6 +179,7 @@ predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cuda")
 ```
 
 When using a Vertex AI model:
+
 ```python
 # CPU inference
 predictor = model.get_local_predictor(device="cpu")
@@ -513,11 +517,8 @@ index = build_vector_index(
     num_workers=8,               # parallel image loading
 )
 
-# Aggregate to one centroid per label
-centroid_index = index.aggregate()
-
 # Save and load
-centroid_index.dump("/path/to/artifact_dir")
+index.dump("/path/to/artifact_dir")
 
 from orient_express.predictors import get_predictor
 loaded_index = get_predictor("/path/to/artifact_dir")
@@ -533,16 +534,23 @@ batch_results = loaded_index.search_batch(query_matrix, k=5)
 
 #### Multi-label support
 
-Vectors can have multiple labels. This is useful when a single visual cluster maps to multiple SKUs:
+Vectors can have composite labels (use tuples). This is useful when a single visual cluster maps to multiple things:
 
 ```python
 index = VectorIndex(
     vectors=feature_matrix,
-    labels=[["sku_101", "sku_102"], ["sku_103"], ["sku_104", "sku_105"]],
+    labels=[("sku_101", "sku_102"), ("sku_103",)],
 )
+```
 
-# aggregate() computes one centroid per unique label
-aggregated = index.aggregate()  # 5 vectors, one per SKU
+#### Per-label aggregation
+
+Vector indices in which labels are not unique can be aggregated so that each label has a single centroid.
+If `per_label=True` and the labels are composite (tuples), then the labels will be unpacked and aggregated separately.
+
+```python
+aggregated = index.aggregate(per_label=True)  # 3 vectors, one per label element ["sku_101", "sku_102", "sku_103"]
+aggregated = index.aggregate(per_label=False)  # 2 vectors, one per composite label  [("sku_101", "sku_102"), ("sku_103")]
 ```
 
 #### Output Structure
@@ -558,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.3.1 → orient_express-2.4.1}/orient_express/predictors/__init__.py

@@ -99,6 +99,7 @@ def load_vector_index(dir: str, metadata: dict | None = None):
     data = np.load(artifact_path, allow_pickle=False)
     vectors = data["vectors"]
     labels = json.loads(str(data["labels_json"]))
+    labels = [tuple(label) if isinstance(label, list) else label for label in labels]
     index = VectorIndex(vectors=vectors, labels=labels)
     index.model_path = artifact_path
     return index

{orient_express-2.3.1 → orient_express-2.4.1}/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.1.2"
+        return "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:v2.4.1"
 
     def get_serving_container_health_route(self, model_name):
         return f"/v1/models/{model_name}"

{orient_express-2.3.1 → orient_express-2.4.1}/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.3.1 → orient_express-2.4.1}/orient_express/predictors/vector_index.py

@@ -2,6 +2,7 @@ import os
 import json
 from dataclasses import dataclass
 import warnings
+from typing import Any
 
 import yaml
 import numpy as np
@@ -37,7 +38,7 @@ class VectorIndex(Predictor):
     def __init__(
         self,
         vectors: np.ndarray,
-        labels: list[list],
+        labels: list[int | str | tuple],
         normalize: bool = False,
     ):
         if vectors.ndim != 2:
@@ -53,8 +54,16 @@ class VectorIndex(Predictor):
             # Avoid division by zero for zero vectors
             norms = np.where(norms == 0, 1, norms)
             vectors = vectors / norms
+
         self.vectors = vectors
         self.labels = labels
+        self.label_to_idx = {}
+
+        for i, label in enumerate(labels):
+            if label not in self.label_to_idx:
+                self.label_to_idx[label] = []
+            self.label_to_idx[label].append(i)
+
         self.model_path = None
 
     def search(self, query: np.ndarray, k: int = 1) -> list[SearchResult]:
@@ -86,21 +95,36 @@ class VectorIndex(Predictor):
             )
         return all_results
 
+    def get_by_idx(self, idx: int) -> np.ndarray:
+        return self.vectors[idx]
+
+    def get_by_idxs(self, idxs: list[int]) -> np.ndarray:
+        return self.vectors[idxs]
+
+    def get_by_label(self, label: Any) -> np.ndarray:
+        return self.vectors[self.label_to_idx[label]]
+
+    def get_by_labels(self, labels: list[Any]) -> np.ndarray:
+        idxs = []
+        for label in labels:
+            idxs.extend(self.label_to_idx[label])
+        return self.vectors[idxs]
+
     def aggregate(self, per_label: bool = False) -> "VectorIndex":
         """
         Aggregate the vectors into a single centroid per unique label group.
 
         Args:
             per_label: if True, create one centroid per individual label
-                (splitting label groups apart). If False (default), create
-                one centroid per unique label group.
+                (splitting label groups (tuples) apart). If False (default),
+                create one centroid per unique label group.
         """
         key_to_indices: dict = {}
-        for i, label_group in enumerate(self.labels):
-            if per_label:
-                keys = [label for label in label_group]
+        for i, label in enumerate(self.labels):
+            if per_label and isinstance(label, tuple):
+                keys = label
             else:
-                keys = [tuple(sorted(label_group))]
+                keys = [label]
             for key in keys:
                 if key not in key_to_indices:
                     key_to_indices[key] = []
@@ -119,7 +143,7 @@ class VectorIndex(Predictor):
             if norm > 0:
                 centroid = centroid / norm
             centroids[i] = centroid
-            new_labels.append([key] if per_label else list(key))
+            new_labels.append(key)
 
         return VectorIndex(vectors=centroids, labels=new_labels)
 
@@ -161,12 +185,9 @@ class VectorIndex(Predictor):
         return self.vectors.shape[0]
 
     def __repr__(self) -> str:
-        unique_labels = set()
-        for group in self.labels:
-            unique_labels.update(group)
         return (
             f"VectorIndex({len(self)} vectors, dim={self.dim}, "
-            f"{len(unique_labels)} unique labels)"
+            f"{len(set(self.labels))} unique labels)"
         )
 
 
@@ -197,11 +218,10 @@ def _collate_pil(batch):
 
 def build_vector_index(
     crops: list[Image.Image | str | CropSpec],
-    labels: list,
+    labels: list[int | str | tuple],
     feature_extractor,
     batch_size: int = 128,
     normalize: bool = True,
-    multi_label: bool = False,
     num_workers: int = 0,
 ) -> VectorIndex:
     """Build a VectorIndex by extracting features from crops.
@@ -210,14 +230,12 @@ def build_vector_index(
         crops: images to extract features from. Each element can be a PIL
             Image (used directly), a file path string (loaded as a whole
             image), or a CropSpec (loaded and cropped to the specified bbox).
-        labels: one label per crop. If multi_label is False, each element is a
-            single label. If multi_label is True, each element should be a list
-            of labels.
+        labels: one label per crop. labels can be any hashable type, includeing
+            iterables like tuples.
         feature_extractor: anything with a .predict(list[Image]) method that
             returns objects with a .feature attribute (e.g. FeatureExtractionPredictor).
         batch_size: number of crops to process at once.
         normalize: whether to L2-normalize the resulting feature vectors.
-        multi_label: if True, labels are already lists of labels per crop.
         num_workers: number of DataLoader workers for parallel image loading.
     """
     if len(crops) != len(labels):
@@ -225,11 +243,6 @@ def build_vector_index(
             f"crops length ({len(crops)}) must match labels length ({len(labels)})"
         )
 
-    if multi_label:
-        index_labels = [list(group) for group in labels]
-    else:
-        index_labels = [[label] for label in labels]
-
     loader = DataLoader(
         _CropDataset(crops),
         batch_size=batch_size,
@@ -243,4 +256,4 @@ def build_vector_index(
         all_features.extend([r.feature for r in results])
 
     vectors = np.vstack(all_features)
-    return VectorIndex(vectors=vectors, labels=index_labels, normalize=normalize)
+    return VectorIndex(vectors=vectors, labels=labels, normalize=normalize)

{orient_express-2.3.1 → orient_express-2.4.1}/orient_express/utils/image_processor.py

@@ -74,6 +74,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.3.1 → orient_express-2.4.1}/orient_express/vertex.py

@@ -1,5 +1,6 @@
 import os
 import tempfile
+import warnings
 
 import yaml
 import joblib
@@ -287,7 +288,6 @@ def get_vertex_model(
 ):
     vertex_init(project_name, region)
     models = aiplatform.Model.list(filter=f"display_name={model_name}")
-
     if not models:
         if raise_exception:
             raise Exception(
@@ -295,19 +295,26 @@ def get_vertex_model(
             )
         else:
             return None
+    if len(models) > 1:
+        warnings.warn(
+            f"Multiple models found with name '{model_name}'. Using the latest one."
+        )
+
+    latest_model = sorted(models, key=lambda x: x.update_time, reverse=True)[0]
+    resource_name = latest_model.resource_name
 
     if version is None:
-        latest_model = sorted(models, key=lambda x: x.update_time, reverse=True)[0]
         return VertexModel(
             latest_model, model_name, project_name, region, int(latest_model.version_id)
         )
 
-    for model in models:
-        if int(model.version_id) == version:
-            return VertexModel(model, model_name, project_name, region, version)
-
-    if raise_exception:
-        raise Exception(
-            f"Model '{model_name}' with version '{version}' not found in registry for project '{project_name}' region '{region}'"
-        )
-    return None
+    if version is not None:
+        try:
+            model = aiplatform.Model(model_name=resource_name, version=str(version))
+        except Exception:
+            if raise_exception:
+                raise Exception(
+                    f"Failed to fetch model '{model_name}' with version '{version}' in registry for project '{project_name}' region '{region}'"
+                )
+            return None
+        return VertexModel(model, model_name, project_name, region, version)

{orient_express-2.3.1 → orient_express-2.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.3.1"
+version = "2.4.1"
 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"