orient_express 2.2.1__tar.gz → 2.3.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.2.1
+Version: 2.3.1
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -21,6 +21,7 @@ Requires-Dist: pyyaml (>=6.0,<7.0)
 Requires-Dist: scikit-learn (>=1.5.0,<2.0.0)
 Requires-Dist: torch (>=2.5.0)
 Requires-Dist: torchvision (>=0.20.0)
+Requires-Dist: tqdm (>=4.67.1)
 Requires-Dist: xgboost (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
 
@@ -519,6 +520,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.1 → orient_express-2.3.1}/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.1 → orient_express-2.3.1}/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.3.1/orient_express/predictors/vector_index.py

@@ -0,0 +1,246 @@
+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, per_label: bool = False) -> "VectorIndex":
+        """
+        Aggregate the vectors into a single centroid per unique label group.
+
+        Args:
+            per_label: if True, create one centroid per individual label
+                (splitting label groups apart). If False (default), create
+                one centroid per unique label group.
+        """
+        key_to_indices: dict = {}
+        for i, label_group in enumerate(self.labels):
+            if per_label:
+                keys = [label for label in label_group]
+            else:
+                keys = [tuple(sorted(label_group))]
+            for key in keys:
+                if key not in key_to_indices:
+                    key_to_indices[key] = []
+                key_to_indices[key].append(i)
+
+        sorted_keys = sorted(key_to_indices.keys(), key=str)
+        centroids = np.empty(
+            (len(sorted_keys), self.vectors.shape[1]), dtype=self.vectors.dtype
+        )
+
+        new_labels = []
+        for i, key in enumerate(sorted_keys):
+            indices = key_to_indices[key]
+            centroid = self.vectors[indices].mean(axis=0)
+            norm = np.linalg.norm(centroid)
+            if norm > 0:
+                centroid = centroid / norm
+            centroids[i] = centroid
+            new_labels.append([key] if per_label else list(key))
+
+        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/shiftsmart-api/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.1 → orient_express-2.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.2.1"
+version = "2.3.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"
@@ -18,6 +18,7 @@ Pillow = ">=10.0.1"
 gcsfs = "^2024.10.0"
 scikit-learn = "^1.5.0"
 xgboost = "^2.0.0"
+tqdm = ">=4.67.1"
 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'"}