orient_express 2.0__tar.gz → 2.1__tar.gz

orient_express-2.0/README.md → orient_express-2.1/PKG-INFO

@@ -1,3 +1,31 @@
+Metadata-Version: 2.4
+Name: orient_express
+Version: 2.1
+Summary: A library to simplify model deployment to Vertex AI
+Author: Alexey Zankevich
+Author-email: alex.zankevich@shiftsmart.com
+Requires-Python: >=3.10,<3.13
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: Pillow (==10.0.1)
+Requires-Dist: gcsfs (>=2024.10.0,<2025.0.0)
+Requires-Dist: google-cloud-aiplatform
+Requires-Dist: google-cloud-storage
+Requires-Dist: kserve (==0.16.0)
+Requires-Dist: onnxruntime (>=1.20,<2.0) ; sys_platform == "darwin" or platform_machine == "aarch64" or platform_machine == "arm64" or sys_platform == "win32" and platform_machine == "ARM64"
+Requires-Dist: onnxruntime-gpu (>=1.20,<2.0) ; sys_platform == "linux" and platform_machine == "x86_64" or sys_platform == "win32" and platform_machine == "AMD64"
+Requires-Dist: opencv-python-headless (==4.11.0.86)
+Requires-Dist: pandas
+Requires-Dist: python-json-logger (==2.0.7)
+Requires-Dist: pyyaml (>=6.0,<7.0)
+Requires-Dist: scikit-learn (==1.5.2)
+Requires-Dist: torch (>=2.5.0,<3.0.0)
+Requires-Dist: torchvision (>=0.20.0,<0.21.0)
+Requires-Dist: xgboost (==2.1.2)
+Description-Content-Type: text/markdown
+
 # Orient Express
 
 A library to accelerate model deployments to Vertex AI directly from colab notebooks
@@ -18,6 +46,16 @@ Both workflows handle versioning, artifact storage in GCS, and integration with
 pip install orient_express
 ```
 
+For local development:
+```bash
+pip install -e .
+```
+
+Or with Poetry:
+```bash
+poetry install
+```
+
 ## Workflows
 
 ### ONNX Image Model Workflow
@@ -137,6 +175,43 @@ local_predictor = vertex_model.get_local_predictor()
 predictions = local_predictor.predict(X_test)
 ```
 
+## ONNX Runtime and Device Support
+
+### 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             |
+
+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
+
+# CPU inference (works on all platforms)
+predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cpu")
+
+# CUDA inference (requires Linux x64 or Windows x64 with CUDA drivers)
+predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cuda")
+```
+
+When using a Vertex AI model:
+```python
+# CPU inference
+predictor = model.get_local_predictor(device="cpu")
+
+# CUDA inference
+predictor = model.get_local_predictor(device="cuda")
+```
+
 ### Pinning Model Versions
 
 By default, `get_vertex_model` returns the most recently updated version. To pin to a specific version:
@@ -213,6 +288,53 @@ class ClassificationPrediction:
 
 </details>
 
+### MultiLabelClassificationPredictor
+
+<details>
+<summary>Click to expand</summary>
+
+For image multi-label classification models that output a set of binary class probabilities.
+
+#### ONNX Graph Requirements
+
+|             |                                                              |
+| ----------- | ------------------------------------------------------------ |
+| **Inputs**  | `images`: `[batch, height, width, 3]` uint8 RGB              |
+| **Outputs** | `scores`: `[batch, num_classes]` float32 class probabilities |
+
+The graph must handle normalization internally. No target_sizes input is needed.
+
+#### Usage
+
+```python
+from orient_express.predictors import MultiLabelClassificationPredictor
+
+predictor = MultiLabelClassificationPredictor(
+    onnx_path="classifier.onnx",
+    classes={1: "contains_cat", 2: "contains_dog", 3: "contains_bird"}
+)
+
+predictions = predictor.predict(images, confidence=0.5)
+# Returns: list[MultiLabelClassificationPrediction]
+```
+
+#### Output Structure
+
+```python
+@dataclass
+class MultiLabelClassificationPrediction:
+    classes: list[str]             # Predicted class names based on confidence threshold
+    class_scores: dict[str, float] # Scores for all classes
+
+# to_dict() output:
+{
+    "classes": ["contains_cat", "contains_bird"],
+    "class_scores": {"contains_cat": 0.95, "contains_dog": 0.03, "contains_bird": 0.82}
+}
+```
+
+</details>
+
 ### BoundingBoxPredictor
 
 <details>
@@ -650,3 +772,4 @@ python deploy.py --run-type scheduled
 ```
 
 </details>
+

orient_express-2.0/PKG-INFO → orient_express-2.1/README.md

@@ -1,20 +1,3 @@
-Metadata-Version: 2.4
-Name: orient_express
-Version: 2.0
-Summary: A library to simplify model deployment to Vertex AI
-Author: Alexey Zankevich
-Author-email: alex.zankevich@shiftsmart.com
-Requires-Python: >=3.10,<3.13
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: google-cloud-aiplatform
-Requires-Dist: google-cloud-storage
-Requires-Dist: pandas
-Requires-Dist: pyyaml (>=6.0,<7.0)
-Description-Content-Type: text/markdown
-
 # Orient Express
 
 A library to accelerate model deployments to Vertex AI directly from colab notebooks
@@ -35,6 +18,16 @@ Both workflows handle versioning, artifact storage in GCS, and integration with
 pip install orient_express
 ```
 
+For local development:
+```bash
+pip install -e .
+```
+
+Or with Poetry:
+```bash
+poetry install
+```
+
 ## Workflows
 
 ### ONNX Image Model Workflow
@@ -154,6 +147,43 @@ local_predictor = vertex_model.get_local_predictor()
 predictions = local_predictor.predict(X_test)
 ```
 
+## ONNX Runtime and Device Support
+
+### 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             |
+
+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
+
+# CPU inference (works on all platforms)
+predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cpu")
+
+# CUDA inference (requires Linux x64 or Windows x64 with CUDA drivers)
+predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cuda")
+```
+
+When using a Vertex AI model:
+```python
+# CPU inference
+predictor = model.get_local_predictor(device="cpu")
+
+# CUDA inference
+predictor = model.get_local_predictor(device="cuda")
+```
+
 ### Pinning Model Versions
 
 By default, `get_vertex_model` returns the most recently updated version. To pin to a specific version:
@@ -230,6 +260,53 @@ class ClassificationPrediction:
 
 </details>
 
+### MultiLabelClassificationPredictor
+
+<details>
+<summary>Click to expand</summary>
+
+For image multi-label classification models that output a set of binary class probabilities.
+
+#### ONNX Graph Requirements
+
+|             |                                                              |
+| ----------- | ------------------------------------------------------------ |
+| **Inputs**  | `images`: `[batch, height, width, 3]` uint8 RGB              |
+| **Outputs** | `scores`: `[batch, num_classes]` float32 class probabilities |
+
+The graph must handle normalization internally. No target_sizes input is needed.
+
+#### Usage
+
+```python
+from orient_express.predictors import MultiLabelClassificationPredictor
+
+predictor = MultiLabelClassificationPredictor(
+    onnx_path="classifier.onnx",
+    classes={1: "contains_cat", 2: "contains_dog", 3: "contains_bird"}
+)
+
+predictions = predictor.predict(images, confidence=0.5)
+# Returns: list[MultiLabelClassificationPrediction]
+```
+
+#### Output Structure
+
+```python
+@dataclass
+class MultiLabelClassificationPrediction:
+    classes: list[str]             # Predicted class names based on confidence threshold
+    class_scores: dict[str, float] # Scores for all classes
+
+# to_dict() output:
+{
+    "classes": ["contains_cat", "contains_bird"],
+    "class_scores": {"contains_cat": 0.95, "contains_dog": 0.03, "contains_bird": 0.82}
+}
+```
+
+</details>
+
 ### BoundingBoxPredictor
 
 <details>
@@ -667,4 +744,3 @@ python deploy.py --run-type scheduled
 ```
 
 </details>
-

{orient_express-2.0 → orient_express-2.1}/orient_express/predictors/__init__.py

@@ -17,9 +17,13 @@ from .semantic_segmentation import (
     SemanticSegmentationPrediction,
 )
 from .classification import ClassificationPredictor, ClassificationPrediction
+from .multi_label_classification import (
+    MultiLabelClassificationPredictor,
+    MultiLabelClassificationPrediction,
+)
 
 
-def get_predictor(dir: str):
+def get_predictor(dir: str, device: str = "cpu"):
     downloaded_files = os.listdir(dir)
     metadata_path = get_metadata_path(dir)
     if not os.path.exists(metadata_path):
@@ -43,20 +47,30 @@ def get_predictor(dir: str):
             joblib_path = os.path.join(dir, metadata["model_file"])
             return joblib.load(joblib_path)
         elif model_type == InstanceSegmentationPredictor.model_type:
-            return load_image_predictor(InstanceSegmentationPredictor, dir, metadata)
+            return load_image_predictor(
+                InstanceSegmentationPredictor, dir, metadata, device
+            )
         elif model_type == BoundingBoxPredictor.model_type:
-            return load_image_predictor(BoundingBoxPredictor, dir, metadata)
+            return load_image_predictor(BoundingBoxPredictor, dir, metadata, device)
         elif model_type == SemanticSegmentationPredictor.model_type:
-            return load_image_predictor(SemanticSegmentationPredictor, dir, metadata)
+            return load_image_predictor(
+                SemanticSegmentationPredictor, dir, metadata, device
+            )
         elif model_type == ClassificationPredictor.model_type:
-            return load_image_predictor(ClassificationPredictor, dir, metadata)
+            return load_image_predictor(ClassificationPredictor, dir, metadata, device)
+        elif model_type == MultiLabelClassificationPredictor.model_type:
+            return load_image_predictor(
+                MultiLabelClassificationPredictor, dir, metadata, device
+            )
         else:
             raise Exception(f"Unknown model_type '{model_type}'")
 
 
-def load_image_predictor(model_type: type[ImagePredictor], dir: str, metadata: dict):
+def load_image_predictor(
+    model_type: type[ImagePredictor], dir: str, metadata: dict, device: str = "cpu"
+):
     onnx_path = os.path.join(dir, metadata["model_file"])
     if "classes" not in metadata:
         raise Exception("No classes defined in metadata.yaml")
     classes = metadata["classes"]
-    return model_type(onnx_path, classes)
+    return model_type(onnx_path, classes, device)

{orient_express-2.0 → orient_express-2.1}/orient_express/predictors/classification.py

@@ -31,7 +31,7 @@ class OnnxClassifier(OnnxSessionWrapper):
 
         input_dict = {self.input_names[0]: images_array}
 
-        scores = self.session.run(None, input_dict)
+        scores = self.session.run(None, input_dict)[0]
         return scores
 
 

{orient_express-2.0 → orient_express-2.1}/orient_express/predictors/instance_segmentation.py

@@ -166,8 +166,8 @@ class InstanceSegmentationPredictor(ImagePredictor):
         centroid_x = int((prediction.bbox[0] + prediction.bbox[2]) / 2)
         centroid_y = int((prediction.bbox[1] + prediction.bbox[3]) / 2)
         # Put the object count in the center of the bbox
-        font_scale = 1.5
-        font_thickness = 4
+        font_scale = 6
+        font_thickness = 12
         text = str(index)
         text_size = cv2.getTextSize(
             text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness

orient_express-2.1/orient_express/predictors/multi_label_classification.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from .predictor import OnnxSessionWrapper, ImagePredictor
+
+
+@dataclass
+class MultiLabelClassificationPrediction:
+    classes: list[str]
+    class_scores: dict[str, float]
+
+    def to_dict(self):
+        return {
+            "classes": self.classes,
+            "class_scores": self.class_scores,
+        }
+
+
+class OnnxMultiLabelClassifier(OnnxSessionWrapper):
+    def __call__(self, pil_images: list[Image.Image]):
+        images = [
+            cv2.resize(np.array(pil_img), (self.resolution, self.resolution))
+            for pil_img in pil_images
+        ]
+        images_array = np.array(images)
+
+        input_dict = {self.input_names[0]: images_array}
+
+        scores = self.session.run(None, input_dict)[0]
+        return scores
+
+
+class MultiLabelClassificationPredictor(ImagePredictor):
+    model_type = "multi-label-classification-onnx"
+    backend_model = OnnxMultiLabelClassifier
+
+    def predict(
+        self, images: list[Image.Image], confidence: float
+    ) -> list[MultiLabelClassificationPrediction]:
+        if not images:
+            return []
+        raw_outputs = self.model(images)
+        outputs = []
+        for class_scores in raw_outputs:
+            classes = []
+            # self.classes is 1-indexed
+            for idx, score in enumerate(class_scores):
+                if score > confidence:
+                    classes.append(self.classes.get(idx + 1, "Unknown"))
+            outputs.append(
+                MultiLabelClassificationPrediction(
+                    classes=classes,
+                    class_scores={
+                        # self.classes is 1-indexed
+                        clss: float(class_scores[class_idx - 1])
+                        for class_idx, clss in self.classes.items()
+                    },
+                )
+            )
+        return outputs
+
+    def get_annotated_image(
+        self, image: Image.Image, predictions: MultiLabelClassificationPrediction
+    ):
+        return None

{orient_express-2.0 → orient_express-2.1}/orient_express/predictors/predictor.py

@@ -1,5 +1,6 @@
 import os
 from abc import ABC, abstractmethod
+import warnings
 
 import yaml
 import onnxruntime as ort
@@ -34,8 +35,8 @@ class ImagePredictor(Predictor):
     backend_model: type
     prediction_type: type
 
-    def __init__(self, model_path: str, classes: dict[int, str]):
-        self.model = self.backend_model(model_path)
+    def __init__(self, model_path: str, classes: dict[int, str], device: str = "cpu"):
+        self.model = self.backend_model(model_path, device)
         self.color_scheme = generate_color_scheme(list(classes.values()))
         self.classes = classes
         self.model_path = model_path
@@ -63,7 +64,7 @@ class ImagePredictor(Predictor):
 
 
 class OnnxSessionWrapper:
-    def __init__(self, onnx_path, providers=["CPUExecutionProvider"]):
+    def __init__(self, onnx_path: str, device: str = "cpu"):
         session_options = ort.SessionOptions()
         session_options.graph_optimization_level = (
             ort.GraphOptimizationLevel.ORT_ENABLE_ALL
@@ -72,6 +73,16 @@ class OnnxSessionWrapper:
         session_options.enable_cpu_mem_arena = True
         session_options.enable_mem_reuse = True
 
+        if device == "cpu":
+            providers = ["CPUExecutionProvider"]
+        elif device == "cuda":
+            providers = ["CUDAExecutionProvider"]
+        else:
+            warnings.warn(
+                f"Unknown device '{device}'. Defaulting to CPU. Supported devices: 'cpu', 'cuda'."
+            )
+            providers = ["CPUExecutionProvider"]
+
         self.session = ort.InferenceSession(
             onnx_path, providers=providers, sess_options=session_options
         )

{orient_express-2.0 → orient_express-2.1}/orient_express/vertex.py

@@ -7,7 +7,6 @@ from google.cloud import storage, aiplatform
 
 from .predictors import get_predictor, Predictor
 
-
 ARTIFACT_DIR = os.path.join(os.path.dirname(__file__), "artifacts")
 _vertex_initialized = False
 
@@ -77,10 +76,10 @@ class VertexModel:
         predictions = self.endpoint.predict(instances=instances, parameters=parameters)
         return predictions.predictions
 
-    def get_local_predictor(self):
+    def get_local_predictor(self, device: str = "cpu"):
         dir = os.path.join(ARTIFACT_DIR, self.model_name + "-" + str(self.version))
         self.download_artifacts(dir)
-        return get_predictor(dir)
+        return get_predictor(dir, device)
 
     def download_artifacts(self, dir: str):
         download_artifacts(dir, self.vertex_model.gca_resource.artifact_uri)

{orient_express-2.0 → orient_express-2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.0"
+version = "2.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"
@@ -12,23 +12,7 @@ google-cloud-storage = "*"
 pandas = "*"
 pyyaml = "^6.0"
 
-[tool.poetry.group.dev.dependencies]
-black = "24.10.0"
-pytest = "8.3.3"
-scikit-learn = "1.5.2"
-xgboost = "2.1.2"
-
-[tool.poetry.group.xgboost-scikit-learn.dependencies]
-scikit-learn = "1.5.2"
 kserve = "0.16.0"
-python-json-logger = "2.0.7"
-gcsfs = "^2024.10.0"
-
-[tool.poetry.group.image-onnx.dependencies]
-onnxruntime = "^1.20" # no gpu for now
-kserve = "0.16.0"
-# torch = "^2.5.0"
-# torchvision = "^0.20.0"
 torch = {version = "^2.5.0", source = "pytorch-cpu"}
 torchvision = {version = "^0.20.0", source = "pytorch-cpu"}
 opencv-python-headless = "4.11.0.86"
@@ -36,6 +20,16 @@ Pillow = "10.0.1"
 python-json-logger = "2.0.7"
 gcsfs = "^2024.10.0"
 
+scikit-learn = "1.5.2"
+xgboost = "2.1.2"
+
+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'"}
+
+[tool.poetry.group.dev.dependencies]
+black = "24.10.0"
+pytest = "8.3.3"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"