inference-models 0.18.3__py3-none-any.whl

inference_models/models/roboflow_instant/roboflow_instant_hf.py

@@ -0,0 +1,141 @@
+import os.path
+from pickle import UnpicklingError
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import Detections, ObjectDetectionModel
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ImageDimensions
+from inference_models.errors import CorruptedModelPackageError
+from inference_models.models.auto_loaders.entities import AnyModel
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.owlv2.entities import (
+    ImageEmbeddings,
+    ReferenceExamplesEmbeddings,
+)
+from inference_models.models.owlv2.owlv2_hf import OWLv2HF
+
+
+class RoboflowInstantHF(ObjectDetectionModel):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        model_dependencies: Optional[Dict[str, AnyModel]] = None,
+        **kwargs,
+    ) -> "ObjectDetectionModel":
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=["weights.pt"],
+        )
+        model_dependencies = model_dependencies or {}
+        if "feature_extractor" in model_dependencies:
+            feature_extractor: OWLv2HF = model_dependencies["feature_extractor"]
+        else:
+            feature_extractor = OWLv2HF.from_pretrained(
+                os.path.join(
+                    model_name_or_path, "model_dependencies", "feature_extractor"
+                ),
+                **kwargs,
+            )
+        try:
+            weights_dict = torch.load(
+                model_package_content["weights.pt"],
+                map_location=device,
+                weights_only=True,
+            )
+        except UnpicklingError as error:
+            raise CorruptedModelPackageError(
+                message="Could not deserialize RF Instant model weights. Contact Roboflow to get help.",
+                help_url="https://todo",
+            ) from error
+        if "class_names" not in weights_dict or "train_data_dict" not in weights_dict:
+            raise CorruptedModelPackageError(
+                message="Corrupted weights of Roboflow Instant model detected. Contact Roboflow to get help.",
+                help_url="https://todo",
+            )
+        class_names = weights_dict["class_names"]
+        train_data_dict = weights_dict["train_data_dict"]
+        try:
+            reference_examples_embeddings = (
+                ReferenceExamplesEmbeddings.from_class_embeddings_dict(
+                    class_embeddings=train_data_dict,
+                    device=device,
+                )
+            )
+        except Exception as error:
+            raise CorruptedModelPackageError(
+                message="Could not decode RF Instant model weights. Contact Roboflow to get help.",
+                help_url="https://todo",
+            ) from error
+        return cls(
+            feature_extractor=feature_extractor,
+            class_names=class_names,
+            reference_examples_embeddings=reference_examples_embeddings,
+        )
+
+    def __init__(
+        self,
+        feature_extractor: OWLv2HF,
+        class_names: List[str],
+        reference_examples_embeddings: ReferenceExamplesEmbeddings,
+    ):
+        self._feature_extractor = feature_extractor
+        self._class_names = class_names
+        self._reference_examples_embeddings = reference_examples_embeddings
+
+    @property
+    def class_names(self) -> List[str]:
+        return self._class_names
+
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        max_detections: int = 300,
+        **kwargs,
+    ) -> Tuple[List[ImageEmbeddings], List[ImageDimensions]]:
+        images_embeddings, images_dimensions = self._feature_extractor.embed_images(
+            images=images,
+            max_detections=max_detections,
+        )
+        return images_embeddings, images_dimensions
+
+    def forward(
+        self,
+        pre_processed_images: List[ImageEmbeddings],
+        confidence_threshold: float = 0.99,
+        iou_threshold: float = 0.3,
+        **kwargs,
+    ) -> List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        return self._feature_extractor.forward_pass_with_precomputed_embeddings(
+            images_embeddings=pre_processed_images,
+            class_embeddings=self._reference_examples_embeddings.class_embeddings,
+            confidence_threshold=confidence_threshold,
+            iou_threshold=iou_threshold,
+        )
+
+    def post_process(
+        self,
+        model_results: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+        pre_processing_meta: List[ImageDimensions],
+        max_detections: int = 300,
+        iou_threshold: float = 0.3,
+        **kwargs,
+    ) -> List[Detections]:
+        return (
+            self._feature_extractor.post_process_predictions_for_precomputed_embeddings(
+                predictions=model_results,
+                images_dimensions=pre_processing_meta,
+                max_detections=max_detections,
+                iou_threshold=iou_threshold,
+            )
+        )

inference_models/models/sam/cache.py

@@ -0,0 +1,147 @@
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from threading import Lock
+from typing import Optional
+
+import torch
+
+from inference_models.errors import EnvironmentConfigurationError
+from inference_models.models.sam.entities import SAMImageEmbeddings
+
+
+class SamImageEmbeddingsCache(ABC):
+
+    @abstractmethod
+    def retrieve_embeddings(self, key: str) -> Optional[SAMImageEmbeddings]:
+        pass
+
+    @abstractmethod
+    def save_embeddings(self, key: str, embeddings: SAMImageEmbeddings) -> None:
+        pass
+
+
+class SamImageEmbeddingsCacheNullObject(SamImageEmbeddingsCache):
+
+    def retrieve_embeddings(self, key: str) -> Optional[SAMImageEmbeddings]:
+        pass
+
+    def save_embeddings(self, key: str, embeddings: SAMImageEmbeddings) -> None:
+        pass
+
+
+class SamImageEmbeddingsInMemoryCache(SamImageEmbeddingsCache):
+
+    @classmethod
+    def init(
+        cls, size_limit: Optional[int], send_to_cpu: bool = True
+    ) -> "SamImageEmbeddingsInMemoryCache":
+        return cls(
+            state=OrderedDict(),
+            size_limit=size_limit,
+            send_to_cpu=send_to_cpu,
+        )
+
+    def __init__(
+        self,
+        state: OrderedDict,
+        size_limit: Optional[int],
+        send_to_cpu: bool = True,
+    ):
+        self._state = state
+        self._size_limit = size_limit
+        self._send_to_cpu = send_to_cpu
+        self._state_lock = Lock()
+
+    def retrieve_embeddings(self, key: str) -> Optional[SAMImageEmbeddings]:
+        return self._state.get(key)
+
+    def save_embeddings(self, key: str, embeddings: SAMImageEmbeddings) -> None:
+        with self._state_lock:
+            if key in self._state:
+                return None
+            self._ensure_cache_has_capacity()
+            if self._send_to_cpu:
+                embeddings = embeddings.to(device=torch.device("cpu"))
+            self._state[key] = embeddings
+
+    def _ensure_cache_has_capacity(self) -> None:
+        if self._size_limit < 1:
+            raise EnvironmentConfigurationError(
+                message=f"In memory cache size for SAM embeddings was set to invalid value. "
+                f"If you are running inference locally - adjust settings of your deployment. If you see this "
+                f"error running on Roboflow platform - contact us to get help.",
+                help_url="https://todo",
+            )
+        if self._size_limit is None or self._size_limit < 1:
+            return None
+        while len(self._state) > self._size_limit:
+            _ = self._state.popitem(last=False)
+
+
+class SamLowResolutionMasksCache(ABC):
+
+    @abstractmethod
+    def retrieve_mask(self, key: str) -> Optional[torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def save_mask(self, key: str, mask: torch.Tensor) -> None:
+        pass
+
+
+class SamLowResolutionMasksCacheNullObject(SamLowResolutionMasksCache):
+
+    def retrieve_mask(self, key: str) -> Optional[torch.Tensor]:
+        pass
+
+    def save_mask(self, key: str, mask: torch.Tensor) -> None:
+        pass
+
+
+class SamLowResolutionMasksInMemoryCache(SamLowResolutionMasksCache):
+
+    @classmethod
+    def init(
+        cls, size_limit: Optional[int], send_to_cpu: bool = True
+    ) -> "SamLowResolutionMasksInMemoryCache":
+        return cls(
+            state=OrderedDict(),
+            size_limit=size_limit,
+            send_to_cpu=send_to_cpu,
+        )
+
+    def __init__(
+        self,
+        state: OrderedDict,
+        size_limit: Optional[int],
+        send_to_cpu: bool = True,
+    ):
+        self._state = state
+        self._size_limit = size_limit
+        self._send_to_cpu = send_to_cpu
+        self._state_lock = Lock()
+
+    def retrieve_mask(self, key: str) -> Optional[torch.Tensor]:
+        return self._state.get(key)
+
+    def save_mask(self, key: str, mask: torch.Tensor) -> None:
+        with self._state_lock:
+            if key in self._state:
+                return None
+            self._ensure_cache_has_capacity()
+            if self._send_to_cpu:
+                mask = mask.to(device=torch.device("cpu"))
+            self._state[key] = mask
+
+    def _ensure_cache_has_capacity(self) -> None:
+        if self._size_limit < 1:
+            raise EnvironmentConfigurationError(
+                message=f"In memory cache size for SAM embeddings was set to invalid value. "
+                f"If you are running inference locally - adjust settings of your deployment. If you see this "
+                f"error running on Roboflow platform - contact us to get help.",
+                help_url="https://todo",
+            )
+        if self._size_limit is None or self._size_limit < 1:
+            return None
+        while len(self._state) > self._size_limit:
+            _ = self._state.popitem(last=False)

inference_models/models/sam/entities.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+
+
+@dataclass(frozen=True)
+class SAMImageEmbeddings:
+    image_hash: str
+    image_size_hw: Tuple[int, int]
+    embeddings: torch.Tensor
+
+    def to(self, device: torch.device) -> "SAMImageEmbeddings":
+        return SAMImageEmbeddings(
+            image_hash=self.image_hash,
+            image_size_hw=self.image_size_hw,
+            embeddings=self.embeddings.to(device=device),
+        )
+
+
+@dataclass(frozen=True)
+class SAMPrediction:
+    masks: torch.Tensor
+    scores: torch.Tensor
+    logits: torch.Tensor