orient_express 2.1__tar.gz → 2.1.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.1
+Version: 2.1.2
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -9,21 +9,19 @@ 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: 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: google-cloud-aiplatform (>=1.0,<2.0)
+Requires-Dist: google-cloud-storage (>=3.0,<4.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: opencv-python-headless (>=4.11.0,<5.0.0)
+Requires-Dist: pandas (>=2.0,<3.0)
 Requires-Dist: pyyaml (>=6.0,<7.0)
-Requires-Dist: scikit-learn (==1.5.2)
+Requires-Dist: scikit-learn (>=1.5.0,<2.0.0)
 Requires-Dist: torch (>=2.5.0,<3.0.0)
 Requires-Dist: torchvision (>=0.20.0,<0.21.0)
-Requires-Dist: xgboost (==2.1.2)
+Requires-Dist: xgboost (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
 
 # Orient Express

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

@@ -23,11 +23,7 @@ class ClassificationPrediction:
 
 class OnnxClassifier(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)
+        images_array = self.collate_images(pil_images)
 
         input_dict = {self.input_names[0]: images_array}
 

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

@@ -35,13 +35,19 @@ class InstanceSegmentationPrediction:
 
 
 class OnnxInstanceSegmentation(OnnxSessionWrapper):
-    def preprocess(self, pil_images: list[Image.Image]):
-        sizes = [[pil_img.size[1], pil_img.size[0]] for pil_img in pil_images]
-        images = [
-            cv2.resize(np.array(pil_img), (self.resolution, self.resolution))
-            for pil_img in pil_images
-        ]
-        return np.array(images), np.array(sizes, dtype=np.float32)
+    def __call__(self, pil_images: list[Image.Image], confidence: float = 0.5):
+        images_array = self.collate_images(pil_images)
+        target_sizes_array = self.collate_sizes(pil_images)
+
+        input_dict = {
+            self.input_names[0]: images_array,
+            self.input_names[1]: target_sizes_array,
+        }
+
+        boxes, scores, labels, masks = self.session.run(None, input_dict)
+        return self.postprocess(
+            boxes, scores, labels, masks, target_sizes_array, confidence
+        )
 
     def postprocess(
         self,
@@ -92,19 +98,6 @@ class OnnxInstanceSegmentation(OnnxSessionWrapper):
 
         return results
 
-    def __call__(self, pil_images: list[Image.Image], confidence: float = 0.5):
-        images_array, target_sizes_array = self.preprocess(pil_images)
-
-        input_dict = {
-            self.input_names[0]: images_array,
-            self.input_names[1]: target_sizes_array,
-        }
-
-        boxes, scores, labels, masks = self.session.run(None, input_dict)
-        return self.postprocess(
-            boxes, scores, labels, masks, target_sizes_array, confidence
-        )
-
 
 class InstanceSegmentationPredictor(ImagePredictor):
     model_type = "instance-segmentation-onnx"
@@ -166,8 +159,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 = 6
-        font_thickness = 12
+        font_scale = 2
+        font_thickness = 4
         text = str(index)
         text_size = cv2.getTextSize(
             text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness

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

@@ -21,11 +21,7 @@ class MultiLabelClassificationPrediction:
 
 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)
+        images_array = self.collate_images(pil_images)
 
         input_dict = {self.input_names[0]: images_array}
 

{orient_express-2.1 → orient_express-2.1.2}/orient_express/predictors/object_detection.py

@@ -30,13 +30,17 @@ class BoundingBoxPrediction:
 
 
 class OnnxDetector(OnnxSessionWrapper):
-    def preprocess(self, pil_images: list[Image.Image]):
-        sizes = [[pil_img.size[1], pil_img.size[0]] for pil_img in pil_images]
-        images = [
-            cv2.resize(np.array(pil_img), (self.resolution, self.resolution))
-            for pil_img in pil_images
-        ]
-        return np.array(images), np.array(sizes, dtype=np.float32)
+    def __call__(self, pil_images: list[Image.Image], confidence: float = 0.5):
+        images_array = self.collate_images(pil_images)
+        target_sizes_array = self.collate_sizes(pil_images)
+
+        input_dict = {
+            self.input_names[0]: images_array,
+            self.input_names[1]: target_sizes_array,
+        }
+
+        boxes, scores, labels = self.session.run(None, input_dict)
+        return self.postprocess(boxes, scores, labels, confidence)
 
     def postprocess(
         self,
@@ -67,17 +71,6 @@ class OnnxDetector(OnnxSessionWrapper):
 
         return results
 
-    def __call__(self, pil_images: list[Image.Image], confidence: float = 0.5):
-        images_array, target_sizes_array = self.preprocess(pil_images)
-
-        input_dict = {
-            self.input_names[0]: images_array,
-            self.input_names[1]: target_sizes_array,
-        }
-
-        boxes, scores, labels = self.session.run(None, input_dict)
-        return self.postprocess(boxes, scores, labels, confidence)
-
 
 class BoundingBoxPredictor(ImagePredictor):
     model_type = "object-detection-onnx"

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

@@ -4,9 +4,13 @@ import warnings
 
 import yaml
 import onnxruntime as ort
+from PIL import Image
+import cv2
+import numpy as np
 
 from ..utils.paths import get_metadata_path
 from ..utils.colors import generate_color_scheme
+from ..utils.image_processor import image_to_array
 
 
 class Predictor(ABC):
@@ -42,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"
+        return "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:v2.1.2"
 
     def get_serving_container_health_route(self, model_name):
         return f"/v1/models/{model_name}"
@@ -92,3 +96,12 @@ class OnnxSessionWrapper:
 
         input_shape = self.session.get_inputs()[0].shape
         self.resolution = input_shape[1]
+        self.img_size = (self.resolution, self.resolution)
+
+    def collate_sizes(self, pil_images: list[Image.Image]):
+        sizes = [[img.size[1], img.size[0]] for img in pil_images]
+        return np.array(sizes, dtype=np.float32)
+
+    def collate_images(self, pil_images: list[Image.Image]):
+        images = [cv2.resize(image_to_array(img), self.img_size) for img in pil_images]
+        return np.array(images)

{orient_express-2.1 → orient_express-2.1.2}/orient_express/predictors/semantic_segmentation.py

@@ -25,13 +25,12 @@ class SemanticSegmentationPrediction:
 
 
 class OnnxSemanticSegmentation(OnnxSessionWrapper):
-    def preprocess(self, pil_images: list[Image.Image]):
-        sizes = [[pil_img.size[1], pil_img.size[0]] for pil_img in pil_images]
-        images = [
-            cv2.resize(np.array(pil_img), (self.resolution, self.resolution))
-            for pil_img in pil_images
-        ]
-        return np.array(images), np.array(sizes, dtype=np.float32)
+    def __call__(self, pil_images: list[Image.Image]):
+        images_array = self.collate_images(pil_images)
+        target_sizes_array = self.collate_sizes(pil_images)
+        input_dict = {self.input_names[0]: images_array}
+        masks = self.session.run(None, input_dict)[0]
+        return self.postprocess(masks, target_sizes_array)
 
     def postprocess(self, masks: np.ndarray, target_sizes: np.ndarray):
         results: list[np.ndarray] = []
@@ -55,12 +54,6 @@ class OnnxSemanticSegmentation(OnnxSessionWrapper):
 
         return results
 
-    def __call__(self, pil_images: list[Image.Image]):
-        images_array, target_sizes_array = self.preprocess(pil_images)
-        input_dict = {self.input_names[0]: images_array}
-        masks = self.session.run(None, input_dict)[0]
-        return self.postprocess(masks, target_sizes_array)
-
 
 class SemanticSegmentationPredictor(ImagePredictor):
     model_type = "semantic-segmentation-onnx"

{orient_express-2.1 → orient_express-2.1.2}/orient_express/utils/image_processor.py

@@ -82,8 +82,16 @@ def image_to_bytes(image):
     return buffered.getvalue()
 
 
+def base64_to_image(base64_data: str):
+    return Image.open(BytesIO(base64.b64decode(base64_data)))
+
+
+def image_to_array(image: Image.Image):
+    return np.array(image.convert("RGB"))
+
+
 def pil_to_opencv(image: Image.Image):
-    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    return cv2.cvtColor(image_to_array(image), cv2.COLOR_RGB2BGR)
 
 
 def opencv_to_pil(image: np.ndarray):

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

@@ -1,34 +1,31 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.1"
+version = "2.1.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"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.13"
-google-cloud-aiplatform = "*"
-google-cloud-storage = "*"
-pandas = "*"
+google-cloud-aiplatform = "^1.0"
+google-cloud-storage = "^3.0"
+pandas = "^2.0"
 pyyaml = "^6.0"
-
-kserve = "0.16.0"
 torch = {version = "^2.5.0", source = "pytorch-cpu"}
 torchvision = {version = "^0.20.0", source = "pytorch-cpu"}
-opencv-python-headless = "4.11.0.86"
-Pillow = "10.0.1"
-python-json-logger = "2.0.7"
+opencv-python-headless = "^4.11.0"
+Pillow = ">=10.0.1"
 gcsfs = "^2024.10.0"
-
-scikit-learn = "1.5.2"
-xgboost = "2.1.2"
-
+scikit-learn = "^1.5.0"
+xgboost = "^2.0.0"
 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"
+kserve = "^0.16.0"
+python-json-logger = "^4.0.0"
 
 [build-system]
 requires = ["poetry-core"]