orient_express 2.2.0__tar.gz → 2.2.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.2.0
+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.2.0 → 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.2.0 → 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.2.0 → 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.2.0 → orient_express-2.2.1}/pyproject.toml

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