orient_express 2.1.3__tar.gz → 2.2.1__tar.gz

{orient_express-2.1.3 → orient_express-2.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.1.3
+Version: 2.2.1
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -19,8 +19,8 @@ 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.0,<2.0.0)
-Requires-Dist: torch (>=2.5.0,<3.0.0)
-Requires-Dist: torchvision (>=0.20.0,<0.21.0)
+Requires-Dist: torch (>=2.5.0)
+Requires-Dist: torchvision (>=0.20.0)
 Requires-Dist: xgboost (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
 

{orient_express-2.1.3 → orient_express-2.2.1}/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.3 → orient_express-2.2.1}/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.1/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.3 → orient_express-2.2.1}/orient_express/predictors/instance_segmentation.py

@@ -10,6 +10,8 @@ import torch.nn.functional as F
 from .predictor import OnnxSessionWrapper, ImagePredictor
 from ..utils.image_processor import pil_to_opencv, opencv_to_pil
 
+FONT = cv2.FONT_HERSHEY_SIMPLEX
+
 
 @dataclass
 class InstanceSegmentationPrediction:
@@ -128,7 +130,12 @@ class InstanceSegmentationPredictor(ImagePredictor):
         return outputs
 
     def get_annotated_image(
-        self, image: Image.Image, predictions: list[InstanceSegmentationPrediction]
+        self,
+        image: Image.Image,
+        predictions: list[InstanceSegmentationPrediction],
+        mask_opacity: float = 0.3,
+        draw_indices: bool = True,
+        font_scale: float | None = None,
     ) -> Image.Image:
         opencv_image = pil_to_opencv(image)
         mask_overlay = opencv_image.copy()
@@ -137,15 +144,20 @@ class InstanceSegmentationPredictor(ImagePredictor):
             fill_color = self.color_scheme.get(pred.clss, (120, 120, 120))
             mask_overlay[pred.mask] = fill_color[:3]
 
-        alpha = 0.5
         annotated_image = cv2.addWeighted(
-            mask_overlay, alpha, opencv_image, 1 - alpha, 0
+            mask_overlay, mask_opacity, opencv_image, 1 - mask_opacity, 0
         )
 
-        class_counts = defaultdict(int)
-        for pred in predictions:
-            class_counts[pred.clss] += 1
-            self.draw_mask_index(annotated_image, pred, class_counts[pred.clss])
+        if draw_indices:
+            class_counts = defaultdict(int)
+            for pred in predictions:
+                class_counts[pred.clss] += 1
+                self.draw_mask_index(
+                    annotated_image,
+                    pred,
+                    class_counts[pred.clss],
+                    font_scale,
+                )
 
         return opencv_to_pil(annotated_image)
 
@@ -154,25 +166,35 @@ class InstanceSegmentationPredictor(ImagePredictor):
         opencv_image: np.ndarray,
         prediction: InstanceSegmentationPrediction,
         index: int,
+        font_scale: float | None = None,
     ):
-        # Calculate the centroid of the bbox
-        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 = 2
-        font_thickness = 4
         text = str(index)
-        text_size = cv2.getTextSize(
-            text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness
-        )[0]
+
+        bbox = prediction.bbox
+        bbox_width = bbox[2] - bbox[0]
+        bbox_height = bbox[3] - bbox[1]
+
+        if font_scale is None:
+            target_size = min(bbox_width, bbox_height) * 0.3
+            (base_w, base_h), _ = cv2.getTextSize(text, FONT, 1, 1)
+            font_scale = target_size / max(base_w, base_h)
+
+        assert font_scale is not None
+        thickness = max(int(font_scale * 2), 1)
+
+        text_size = cv2.getTextSize(text, FONT, font_scale, thickness)[0]
+
+        centroid_x = int((bbox[0] + bbox[2]) / 2)
+        centroid_y = int((bbox[1] + bbox[3]) / 2)
         text_x = int(centroid_x - text_size[0] / 2)
         text_y = int(centroid_y + text_size[1] / 2)
+
         cv2.putText(
             opencv_image,
             text,
             (text_x, text_y),
-            cv2.FONT_HERSHEY_SIMPLEX,
+            FONT,
             font_scale,
             (255, 255, 255),
-            font_thickness,
+            thickness,
         )

{orient_express-2.1.3 → orient_express-2.2.1}/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.3 → orient_express-2.2.1}/orient_express/predictors/object_detection.py

@@ -117,7 +117,10 @@ class BoundingBoxPredictor(ImagePredictor):
         return outputs
 
     def get_annotated_image(
-        self, image: Image.Image, predictions: list[BoundingBoxPrediction]
+        self,
+        image: Image.Image,
+        predictions: list[BoundingBoxPrediction],
+        line_width: int = 8,
     ):
         opencv_image = pil_to_opencv(image)
 
@@ -125,6 +128,8 @@ class BoundingBoxPredictor(ImagePredictor):
             color = self.color_scheme.get(pred.clss, (255, 255, 255))
             x1, y1, x2, y2 = pred.bbox
             x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-            opencv_image = cv2.rectangle(opencv_image, (x1, y1), (x2, y2), color, 2)
+            opencv_image = cv2.rectangle(
+                opencv_image, (x1, y1), (x2, y2), color, line_width
+            )
 
         return opencv_to_pil(opencv_image)

{orient_express-2.1.3 → orient_express-2.2.1}/orient_express/predictors/semantic_segmentation.py

@@ -76,15 +76,18 @@ class SemanticSegmentationPredictor(ImagePredictor):
             )
         return outputs
 
-    def get_annotated_image(self, image: Image.Image, mask: np.array) -> Image.Image:
+    def get_annotated_image(
+        self, image: Image.Image, mask: np.array, mask_opacity: float = 0.3
+    ) -> Image.Image:
         opencv_image = pil_to_opencv(image)
-
-        mask_rgb = np.zeros((*mask.shape, 3), dtype=np.uint8)
+        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_rgb[mask == class_id] = fill_color[:3]
+            mask_overlay[mask == class_id] = fill_color[:3]
 
-        annotated_image = cv2.addWeighted(opencv_image, 0.7, mask_rgb, 0.3, 0)
+        annotated_image = cv2.addWeighted(
+            mask_overlay, mask_opacity, opencv_image, 1 - mask_opacity, 0
+        )
 
         return opencv_to_pil(annotated_image)

{orient_express-2.1.3 → orient_express-2.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.1.3"
+version = "2.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"
@@ -11,8 +11,8 @@ google-cloud-aiplatform = "^1.0"
 google-cloud-storage = "^3.0"
 pandas = "^2.0"
 pyyaml = "^6.0"
-torch = {version = "^2.5.0", source = "pytorch-cpu"}
-torchvision = {version = "^0.20.0", source = "pytorch-cpu"}
+torch = ">=2.5.0"
+torchvision = ">=0.20.0"
 opencv-python-headless = "^4.11.0"
 Pillow = ">=10.0.1"
 gcsfs = "^2024.10.0"
@@ -30,8 +30,3 @@ python-json-logger = "^4.0.0"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
-
-[[tool.poetry.source]]
-name = "pytorch-cpu"
-url = "https://download.pytorch.org/whl/cpu"
-priority = "explicit"