orient_express 2.3.0__tar.gz → 2.4.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.3.0
+Version: 2.4.0
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -21,6 +21,7 @@ Requires-Dist: pyyaml (>=6.0,<7.0)
 Requires-Dist: scikit-learn (>=1.5.0,<2.0.0)
 Requires-Dist: torch (>=2.5.0)
 Requires-Dist: torchvision (>=0.20.0)
+Requires-Dist: tqdm (>=4.67.1)
 Requires-Dist: xgboost (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
 
@@ -45,11 +46,13 @@ pip install orient_express
 ```
 
 For local development:
+
 ```bash
 pip install -e .
 ```
 
 Or with Poetry:
+
 ```bash
 poetry install
 ```
@@ -101,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"}],
 )
 ```
 
@@ -178,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
 
@@ -202,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")
@@ -539,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")
@@ -559,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",)],
 )
+```
+
+#### Per-label aggregation
 
-# aggregate() computes one centroid per unique label
-aggregated = index.aggregate()  # 5 vectors, one per SKU
+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

{orient_express-2.3.0 → orient_express-2.4.0}/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",)],
 )
+```
+
+#### Per-label aggregation
 
-# aggregate() computes one centroid per unique label
-aggregated = index.aggregate()  # 5 vectors, one per SKU
+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

{orient_express-2.3.0 → orient_express-2.4.0}/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.0 → orient_express-2.4.0}/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.0"
 
     def get_serving_container_health_route(self, model_name):
         return f"/v1/models/{model_name}"

{orient_express-2.3.0 → orient_express-2.4.0}/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,38 +95,63 @@ class VectorIndex(Predictor):
             )
         return all_results
 
-    def aggregate(self) -> "VectorIndex":
+    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 label.
+        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 (tuples) apart). If False (default),
+                create one centroid per unique label group.
         """
-        label_to_indices: dict = {}
-        for i, label_group in enumerate(self.labels):
-            for label in label_group:
-                if label not in label_to_indices:
-                    label_to_indices[label] = []
-                label_to_indices[label].append(i)
-
-        unique_labels = sorted(label_to_indices.keys(), key=str)
+        key_to_indices: dict = {}
+        for i, label in enumerate(self.labels):
+            if per_label and isinstance(label, tuple):
+                keys = label
+            else:
+                keys = [label]
+            for key in keys:
+                if key not in key_to_indices:
+                    key_to_indices[key] = []
+                key_to_indices[key].append(i)
+
+        sorted_keys = sorted(key_to_indices.keys(), key=str)
         centroids = np.empty(
-            (len(unique_labels), self.vectors.shape[1]), dtype=self.vectors.dtype
+            (len(sorted_keys), self.vectors.shape[1]), dtype=self.vectors.dtype
         )
 
-        for i, label in enumerate(unique_labels):
-            indices = label_to_indices[label]
+        new_labels = []
+        for i, key in enumerate(sorted_keys):
+            indices = key_to_indices[key]
             centroid = self.vectors[indices].mean(axis=0)
             norm = np.linalg.norm(centroid)
             if norm > 0:
                 centroid = centroid / norm
             centroids[i] = centroid
+            new_labels.append(key)
 
-        new_labels = [[label] for label in unique_labels]
         return VectorIndex(vectors=centroids, labels=new_labels)
 
     def get_serving_container_image_uri(self) -> str:
         warnings.warn(
             "VectorIndex does not support serving via a container. Returning incompatible image URI."
         )
-        return "us-west1-docker.pkg.dev/shiftsmartinc/orient-express/image-onnx:v2.1.2"
+        return "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:v2.1.2"
 
     def get_serving_container_health_route(self, model_name) -> str:
         return f"/v1/models/{model_name}"
@@ -151,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)"
         )
 
 
@@ -187,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.
@@ -200,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):
@@ -215,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,
@@ -233,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.0 → orient_express-2.4.0}/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.0 → orient_express-2.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.3.0"
+version = "2.4.0"
 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"
@@ -18,6 +18,7 @@ Pillow = ">=10.0.1"
 gcsfs = "^2024.10.0"
 scikit-learn = "^1.5.0"
 xgboost = "^2.0.0"
+tqdm = ">=4.67.1"
 onnxruntime-gpu = {version = "^1.20", markers = "sys_platform == 'linux' and platform_machine == 'x86_64' or sys_platform == 'win32' and platform_machine == 'AMD64'"}
 onnxruntime = {version = "^1.20", markers = "sys_platform == 'darwin' or platform_machine == 'aarch64' or platform_machine == 'arm64' or sys_platform == 'win32' and platform_machine == 'ARM64'"}