orient_express 2.2.0__tar.gz → 2.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.2.0
+Version: 2.3.0
 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
 
@@ -519,6 +519,69 @@ annotated_image = predictor.get_annotated_image(image, predictions[0].class_mask
 
 </details>
 
+### VectorIndex
+
+<details>
+<summary>Click to expand</summary>
+
+A cosine-similarity vector index for matching feature vectors to labels. Each vector in the index can have one or more labels. VectorIndex integrates with `get_predictor` for loading from saved artifacts, and can be built from scratch using a feature extractor.
+
+#### Usage
+
+```python
+from orient_express.predictors import VectorIndex, build_vector_index
+
+# Build from crops and labels using a feature extractor
+index = build_vector_index(
+    crops=crop_images,           # list of PIL Images or file paths
+    labels=cluster_ids,          # one label per crop
+    feature_extractor=fe,        # FeatureExtractionPredictor
+    num_workers=8,               # parallel image loading
+)
+
+# Aggregate to one centroid per label
+centroid_index = index.aggregate()
+
+# Save and load
+centroid_index.dump("/path/to/artifact_dir")
+
+from orient_express.predictors import get_predictor
+loaded_index = get_predictor("/path/to/artifact_dir")
+
+# Search
+results = loaded_index.search(query_vector, k=5)
+for result in results:
+    print(result.labels, result.score)
+
+# Batch search
+batch_results = loaded_index.search_batch(query_matrix, k=5)
+```
+
+#### Multi-label support
+
+Vectors can have multiple labels. This is useful when a single visual cluster maps to multiple SKUs:
+
+```python
+index = VectorIndex(
+    vectors=feature_matrix,
+    labels=[["sku_101", "sku_102"], ["sku_103"], ["sku_104", "sku_105"]],
+)
+
+# aggregate() computes one centroid per unique label
+aggregated = index.aggregate()  # 5 vectors, one per SKU
+```
+
+#### Output Structure
+
+```python
+@dataclass
+class SearchResult:
+    labels: list   # All labels for the matched vector
+    score: float   # Cosine similarity score
+```
+
+</details>
+
 ---
 
 ## Color Schemes

{orient_express-2.2.0 → orient_express-2.3.0}/README.md

@@ -493,6 +493,69 @@ annotated_image = predictor.get_annotated_image(image, predictions[0].class_mask
 
 </details>
 
+### VectorIndex
+
+<details>
+<summary>Click to expand</summary>
+
+A cosine-similarity vector index for matching feature vectors to labels. Each vector in the index can have one or more labels. VectorIndex integrates with `get_predictor` for loading from saved artifacts, and can be built from scratch using a feature extractor.
+
+#### Usage
+
+```python
+from orient_express.predictors import VectorIndex, build_vector_index
+
+# Build from crops and labels using a feature extractor
+index = build_vector_index(
+    crops=crop_images,           # list of PIL Images or file paths
+    labels=cluster_ids,          # one label per crop
+    feature_extractor=fe,        # FeatureExtractionPredictor
+    num_workers=8,               # parallel image loading
+)
+
+# Aggregate to one centroid per label
+centroid_index = index.aggregate()
+
+# Save and load
+centroid_index.dump("/path/to/artifact_dir")
+
+from orient_express.predictors import get_predictor
+loaded_index = get_predictor("/path/to/artifact_dir")
+
+# Search
+results = loaded_index.search(query_vector, k=5)
+for result in results:
+    print(result.labels, result.score)
+
+# Batch search
+batch_results = loaded_index.search_batch(query_matrix, k=5)
+```
+
+#### Multi-label support
+
+Vectors can have multiple labels. This is useful when a single visual cluster maps to multiple SKUs:
+
+```python
+index = VectorIndex(
+    vectors=feature_matrix,
+    labels=[["sku_101", "sku_102"], ["sku_103"], ["sku_104", "sku_105"]],
+)
+
+# aggregate() computes one centroid per unique label
+aggregated = index.aggregate()  # 5 vectors, one per SKU
+```
+
+#### Output Structure
+
+```python
+@dataclass
+class SearchResult:
+    labels: list   # All labels for the matched vector
+    score: float   # Cosine similarity score
+```
+
+</details>
+
 ---
 
 ## Color Schemes

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

@@ -1,8 +1,10 @@
 import os
 import logging
+import json
 
 import joblib
 import yaml
+import numpy as np
 
 from ..utils.paths import get_metadata_path
 
@@ -22,6 +24,7 @@ from .multi_label_classification import (
     MultiLabelClassificationPrediction,
 )
 from .feature_extraction import FeatureExtractionPredictor, FeaturePrediction
+from .vector_index import VectorIndex, SearchResult, CropSpec, build_vector_index
 
 
 def get_predictor(dir: str, device: str = "cpu"):
@@ -65,6 +68,8 @@ def get_predictor(dir: str, device: str = "cpu"):
             )
         elif model_type == FeatureExtractionPredictor.model_type:
             return load_feature_extractor(dir, metadata, device)
+        elif model_type == VectorIndex.model_type:
+            return load_vector_index(dir, metadata)
         else:
             raise Exception(f"Unknown model_type '{model_type}'")
 
@@ -82,3 +87,18 @@ def load_image_predictor(
 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)
+
+
+def load_vector_index(dir: str, metadata: dict | None = None):
+    if metadata is None:
+        metadata_path = get_metadata_path(dir)
+        with open(metadata_path) as f:
+            metadata = yaml.safe_load(f)
+    assert metadata is not None
+    artifact_path = os.path.join(dir, metadata["model_file"])
+    data = np.load(artifact_path, allow_pickle=False)
+    vectors = data["vectors"]
+    labels = json.loads(str(data["labels_json"]))
+    index = VectorIndex(vectors=vectors, labels=labels)
+    index.model_path = artifact_path
+    return index

{orient_express-2.2.0 → orient_express-2.3.0}/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.3.0}/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.3.0}/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.3.0/orient_express/predictors/vector_index.py

@@ -0,0 +1,236 @@
+import os
+import json
+from dataclasses import dataclass
+import warnings
+
+import yaml
+import numpy as np
+from PIL import Image
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from .predictor import Predictor
+from ..utils.paths import get_metadata_path
+
+
+@dataclass
+class SearchResult:
+    labels: list
+    score: float
+
+
+@dataclass
+class CropSpec:
+    """An image path with a bounding box to crop before feature extraction.
+
+    bbox is (x1, y1, x2, y2) in pixel coordinates, as expected by PIL.Image.crop.
+    """
+
+    path: str
+    bbox: tuple[int, int, int, int]
+
+
+class VectorIndex(Predictor):
+    model_type = "vector-index"
+    ARTIFACT_FILENAME = "vectors.npz"
+
+    def __init__(
+        self,
+        vectors: np.ndarray,
+        labels: list[list],
+        normalize: bool = False,
+    ):
+        if vectors.ndim != 2:
+            raise ValueError(
+                f"vectors must be 2-dimensional (M, D), got shape {vectors.shape}"
+            )
+        if len(labels) != vectors.shape[0]:
+            raise ValueError(
+                f"labels length ({len(labels)}) must match number of vectors ({vectors.shape[0]})"
+            )
+        if normalize:
+            norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+            # Avoid division by zero for zero vectors
+            norms = np.where(norms == 0, 1, norms)
+            vectors = vectors / norms
+        self.vectors = vectors
+        self.labels = labels
+        self.model_path = None
+
+    def search(self, query: np.ndarray, k: int = 1) -> list[SearchResult]:
+        query = query.reshape(1, -1)
+        similarities = (query @ self.vectors.T).squeeze(0)
+        k = min(k, len(similarities))
+        top_indices = np.argpartition(similarities, -k)[-k:]
+        top_indices = top_indices[np.argsort(similarities[top_indices])[::-1]]
+
+        return [
+            SearchResult(labels=self.labels[idx], score=float(similarities[idx]))
+            for idx in top_indices
+        ]
+
+    def search_batch(self, queries: np.ndarray, k: int = 1) -> list[list[SearchResult]]:
+        similarities = queries @ self.vectors.T
+        k = min(k, similarities.shape[1])
+
+        all_results = []
+        for i in range(similarities.shape[0]):
+            row = similarities[i]
+            top_indices = np.argpartition(row, -k)[-k:]
+            top_indices = top_indices[np.argsort(row[top_indices])[::-1]]
+            all_results.append(
+                [
+                    SearchResult(labels=self.labels[idx], score=float(row[idx]))
+                    for idx in top_indices
+                ]
+            )
+        return all_results
+
+    def aggregate(self) -> "VectorIndex":
+        """
+        Aggregate the vectors into a single centroid per label.
+        """
+        label_to_indices: dict = {}
+        for i, label_group in enumerate(self.labels):
+            for label in label_group:
+                if label not in label_to_indices:
+                    label_to_indices[label] = []
+                label_to_indices[label].append(i)
+
+        unique_labels = sorted(label_to_indices.keys(), key=str)
+        centroids = np.empty(
+            (len(unique_labels), self.vectors.shape[1]), dtype=self.vectors.dtype
+        )
+
+        for i, label in enumerate(unique_labels):
+            indices = label_to_indices[label]
+            centroid = self.vectors[indices].mean(axis=0)
+            norm = np.linalg.norm(centroid)
+            if norm > 0:
+                centroid = centroid / norm
+            centroids[i] = centroid
+
+        new_labels = [[label] for label in unique_labels]
+        return VectorIndex(vectors=centroids, labels=new_labels)
+
+    def get_serving_container_image_uri(self) -> str:
+        warnings.warn(
+            "VectorIndex does not support serving via a container. Returning incompatible image URI."
+        )
+        return "us-west1-docker.pkg.dev/shiftsmartinc/orient-express/image-onnx:v2.1.2"
+
+    def get_serving_container_health_route(self, model_name) -> str:
+        return f"/v1/models/{model_name}"
+
+    def get_serving_container_predict_route(self, model_name) -> str:
+        return f"/v1/models/{model_name}:predict"
+
+    def dump(self, dir: str) -> list[str]:
+        artifact_path = os.path.join(dir, self.ARTIFACT_FILENAME)
+        np.savez(
+            artifact_path,
+            vectors=self.vectors,
+            labels_json=json.dumps(self.labels),
+        )
+
+        metadata = {
+            "model_type": self.model_type,
+            "model_file": self.ARTIFACT_FILENAME,
+        }
+        metadata_path = get_metadata_path(dir)
+        with open(metadata_path, "w") as f:
+            yaml.dump(metadata, f)
+
+        return [metadata_path, artifact_path]
+
+    @property
+    def dim(self) -> int:
+        return self.vectors.shape[1]
+
+    def __len__(self) -> int:
+        return self.vectors.shape[0]
+
+    def __repr__(self) -> str:
+        unique_labels = set()
+        for group in self.labels:
+            unique_labels.update(group)
+        return (
+            f"VectorIndex({len(self)} vectors, dim={self.dim}, "
+            f"{len(unique_labels)} unique labels)"
+        )
+
+
+class _CropDataset:
+    def __init__(self, crops: list[Image.Image | str | CropSpec]):
+        self.crops = crops
+
+    def __len__(self):
+        return len(self.crops)
+
+    def __getitem__(self, idx):
+        crop = self.crops[idx]
+        if isinstance(crop, CropSpec):
+            img = Image.open(crop.path).convert("RGB")
+            return img.crop(crop.bbox)
+        if isinstance(crop, str):
+            return Image.open(crop).convert("RGB")
+        if isinstance(crop, Image.Image):
+            return crop
+        raise TypeError(
+            f"Each crop must be a PIL Image, a file path string, or a CropSpec, got {type(crop)}"
+        )
+
+
+def _collate_pil(batch):
+    return list(batch)
+
+
+def build_vector_index(
+    crops: list[Image.Image | str | CropSpec],
+    labels: list,
+    feature_extractor,
+    batch_size: int = 128,
+    normalize: bool = True,
+    multi_label: bool = False,
+    num_workers: int = 0,
+) -> VectorIndex:
+    """Build a VectorIndex by extracting features from crops.
+
+    Args:
+        crops: images to extract features from. Each element can be a PIL
+            Image (used directly), a file path string (loaded as a whole
+            image), or a CropSpec (loaded and cropped to the specified bbox).
+        labels: one label per crop. If multi_label is False, each element is a
+            single label. If multi_label is True, each element should be a list
+            of labels.
+        feature_extractor: anything with a .predict(list[Image]) method that
+            returns objects with a .feature attribute (e.g. FeatureExtractionPredictor).
+        batch_size: number of crops to process at once.
+        normalize: whether to L2-normalize the resulting feature vectors.
+        multi_label: if True, labels are already lists of labels per crop.
+        num_workers: number of DataLoader workers for parallel image loading.
+    """
+    if len(crops) != len(labels):
+        raise ValueError(
+            f"crops length ({len(crops)}) must match labels length ({len(labels)})"
+        )
+
+    if multi_label:
+        index_labels = [list(group) for group in labels]
+    else:
+        index_labels = [[label] for label in labels]
+
+    loader = DataLoader(
+        _CropDataset(crops),
+        batch_size=batch_size,
+        num_workers=num_workers,
+        collate_fn=_collate_pil,
+    )
+
+    all_features = []
+    for batch in tqdm(loader):
+        results = feature_extractor.predict(batch)
+        all_features.extend([r.feature for r in results])
+
+    vectors = np.vstack(all_features)
+    return VectorIndex(vectors=vectors, labels=index_labels, normalize=normalize)

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

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