orient_express 2.1.2__tar.gz → 2.2.0__tar.gz

{orient_express-2.1.2 → orient_express-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.1.2
+Version: 2.2.0
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com

{orient_express-2.1.2 → orient_express-2.2.0}/orient_express/predictors/__init__.py

@@ -21,6 +21,7 @@ from .multi_label_classification import (
     MultiLabelClassificationPredictor,
     MultiLabelClassificationPrediction,
 )
+from .feature_extraction import FeatureExtractionPredictor, FeaturePrediction
 
 
 def get_predictor(dir: str, device: str = "cpu"):
@@ -62,6 +63,8 @@ def get_predictor(dir: str, device: str = "cpu"):
             return load_image_predictor(
                 MultiLabelClassificationPredictor, dir, metadata, device
             )
+        elif model_type == FeatureExtractionPredictor.model_type:
+            return load_feature_extractor(dir, metadata, device)
         else:
             raise Exception(f"Unknown model_type '{model_type}'")
 
@@ -74,3 +77,8 @@ def load_image_predictor(
         raise Exception("No classes defined in metadata.yaml")
     classes = metadata["classes"]
     return model_type(onnx_path, classes, device)
+
+
+def load_feature_extractor(dir: str, metadata: dict, device: str = "cpu"):
+    onnx_path = os.path.join(dir, metadata["model_file"])
+    return FeatureExtractionPredictor(onnx_path, device)

{orient_express-2.1.2 → orient_express-2.2.0}/orient_express/predictors/classification.py

@@ -1,7 +1,5 @@
 from dataclasses import dataclass
 
-import cv2
-import numpy as np
 from PIL import Image
 
 from .predictor import OnnxSessionWrapper, ImagePredictor
@@ -24,9 +22,7 @@ class ClassificationPrediction:
 class OnnxClassifier(OnnxSessionWrapper):
     def __call__(self, pil_images: list[Image.Image]):
         images_array = self.collate_images(pil_images)
-
         input_dict = {self.input_names[0]: images_array}
-
         scores = self.session.run(None, input_dict)[0]
         return scores
 

orient_express-2.2.0/orient_express/predictors/feature_extraction.py

@@ -0,0 +1,59 @@
+import os
+from dataclasses import dataclass
+
+import yaml
+import numpy as np
+from PIL import Image
+
+from .predictor import OnnxSessionWrapper, ImagePredictor
+from ..utils.paths import get_metadata_path
+
+
+@dataclass
+class FeaturePrediction:
+    feature: np.ndarray
+
+    def to_dict(self):
+        return {
+            "feature": self.feature.tolist(),
+        }
+
+
+class OnnxFeatureExtractor(OnnxSessionWrapper):
+    def __call__(self, pil_images: list[Image.Image]):
+        images_array = self.collate_images(pil_images)
+        input_dict = {self.input_names[0]: images_array}
+        features = self.session.run(None, input_dict)[0]
+        return features
+
+
+class FeatureExtractionPredictor(ImagePredictor):
+    model_type = "feature-extraction-onnx"
+    backend_model = OnnxFeatureExtractor
+
+    def __init__(self, model_path: str, device: str = "cpu"):
+        self.model = self.backend_model(model_path, device)
+        self.model_path = model_path
+
+    def predict(self, images: list[Image.Image]) -> list[FeaturePrediction]:
+        if not images:
+            return []
+        raw_outputs = self.model(images)
+        outputs = []
+        for feature in raw_outputs:
+            outputs.append(FeaturePrediction(feature=feature))
+        return outputs
+
+    def get_annotated_image(self, image: Image.Image, predictions: FeaturePrediction):
+        return None
+
+    def dump(self, dir: str):
+        metadata = {
+            "model_type": self.model_type,
+            "model_file": os.path.basename(self.model_path),
+        }
+        metadata_path = get_metadata_path(dir)
+        with open(metadata_path, "w") as f:
+            yaml.dump(metadata, f)
+        # model is already saved in the model_path
+        return [metadata_path, self.model_path]

{orient_express-2.1.2 → orient_express-2.2.0}/orient_express/predictors/multi_label_classification.py

@@ -1,7 +1,5 @@
 from dataclasses import dataclass
 
-import cv2
-import numpy as np
 from PIL import Image
 
 from .predictor import OnnxSessionWrapper, ImagePredictor
@@ -22,9 +20,7 @@ class MultiLabelClassificationPrediction:
 class OnnxMultiLabelClassifier(OnnxSessionWrapper):
     def __call__(self, pil_images: list[Image.Image]):
         images_array = self.collate_images(pil_images)
-
         input_dict = {self.input_names[0]: images_array}
-
         scores = self.session.run(None, input_dict)[0]
         return scores
 

{orient_express-2.1.2 → orient_express-2.2.0}/orient_express/vertex.py

@@ -76,13 +76,15 @@ class VertexModel:
         predictions = self.endpoint.predict(instances=instances, parameters=parameters)
         return predictions.predictions
 
-    def get_local_predictor(self, device: str = "cpu"):
+    def get_local_predictor(self, device: str = "cpu", force_download: bool = False):
         dir = os.path.join(ARTIFACT_DIR, self.model_name + "-" + str(self.version))
-        self.download_artifacts(dir)
+        self.download_artifacts(dir, force_download=force_download)
         return get_predictor(dir, device)
 
-    def download_artifacts(self, dir: str):
-        download_artifacts(dir, self.vertex_model.gca_resource.artifact_uri)
+    def download_artifacts(self, dir: str, force_download: bool = True):
+        download_artifacts(
+            dir, self.vertex_model.gca_resource.artifact_uri, force_download
+        )
 
 
 def vertex_init(project_name: str, region: str):
@@ -92,18 +94,17 @@ def vertex_init(project_name: str, region: str):
         _vertex_initialized = True
 
 
-def download_artifacts(dir: str, artifact_uri: str):
+def download_artifacts(dir: str, artifact_uri: str, force_download: bool = True):
     storage_client = storage.Client()
-
     bucket_name, artifact_path = artifact_uri.replace("gs://", "").split("/", 1)
     bucket = storage_client.bucket(bucket_name)
-
     os.makedirs(dir, exist_ok=True)
-
     blobs = bucket.list_blobs(prefix=artifact_path)
     for blob in blobs:
         filename = blob.name.split("/")[-1]
         download_path = os.path.join(dir, filename)
+        if not force_download and os.path.exists(download_path):
+            continue
         blob.download_to_filename(download_path)
 
 

{orient_express-2.1.2 → orient_express-2.2.0}/pyproject.toml

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