inference-models 0.18.3__py3-none-any.whl

inference_models/models/sam2/cache.py

@@ -0,0 +1,162 @@
+from abc import ABC, abstractmethod
+from collections import OrderedDict, defaultdict
+from threading import Lock
+from typing import DefaultDict, List, Optional
+
+import torch
+
+from inference_models.errors import EnvironmentConfigurationError
+from inference_models.models.sam2.entities import (
+    SAM2ImageEmbeddings,
+    SAM2MaskCacheEntry,
+)
+
+
+class Sam2ImageEmbeddingsCache(ABC):
+
+    @abstractmethod
+    def retrieve_embeddings(self, key: str) -> Optional[SAM2ImageEmbeddings]:
+        pass
+
+    @abstractmethod
+    def save_embeddings(self, key: str, embeddings: SAM2ImageEmbeddings) -> None:
+        pass
+
+
+class Sam2ImageEmbeddingsCacheNullObject(Sam2ImageEmbeddingsCache):
+
+    def retrieve_embeddings(self, key: str) -> Optional[SAM2ImageEmbeddings]:
+        pass
+
+    def save_embeddings(self, key: str, embeddings: SAM2ImageEmbeddings) -> None:
+        pass
+
+
+class Sam2ImageEmbeddingsInMemoryCache(Sam2ImageEmbeddingsCache):
+
+    @classmethod
+    def init(
+        cls, size_limit: Optional[int], send_to_cpu: bool = True
+    ) -> "Sam2ImageEmbeddingsInMemoryCache":
+        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[SAM2ImageEmbeddings]:
+        return self._state.get(key)
+
+    def save_embeddings(self, key: str, embeddings: SAM2ImageEmbeddings) -> 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 SAM2 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 Sam2LowResolutionMasksCache(ABC):
+
+    @abstractmethod
+    def retrieve_all_masks_for_image(self, key: str) -> List[SAM2MaskCacheEntry]:
+        pass
+
+    @abstractmethod
+    def save_mask(self, key: str, mask: SAM2MaskCacheEntry) -> None:
+        pass
+
+
+class Sam2LowResolutionMasksCacheNullObject(Sam2LowResolutionMasksCache):
+
+    def retrieve_all_masks_for_image(self, key: str) -> List[SAM2MaskCacheEntry]:
+        return []
+
+    def save_mask(self, key: str, mask: SAM2MaskCacheEntry) -> None:
+        pass
+
+
+class Sam2LowResolutionMasksInMemoryCache(Sam2LowResolutionMasksCache):
+
+    @classmethod
+    def init(
+        cls, size_limit: Optional[int], send_to_cpu: bool = True
+    ) -> "Sam2LowResolutionMasksInMemoryCache":
+        return cls(
+            ordering_state=OrderedDict(),
+            cache_state=defaultdict(list),
+            size_limit=size_limit,
+            send_to_cpu=send_to_cpu,
+        )
+
+    def __init__(
+        self,
+        ordering_state: OrderedDict,
+        cache_state: DefaultDict[str, List[SAM2MaskCacheEntry]],
+        size_limit: Optional[int],
+        send_to_cpu: bool = True,
+    ):
+        self._ordering_state = ordering_state
+        self._cache_state = cache_state
+        self._size_limit = size_limit
+        self._send_to_cpu = send_to_cpu
+        self._state_lock = Lock()
+
+    def retrieve_all_masks_for_image(self, key: str) -> List[SAM2MaskCacheEntry]:
+        return self._cache_state.get(key, [])
+
+    def save_mask(self, key: str, mask: SAM2MaskCacheEntry) -> None:
+        with self._state_lock:
+            if (key, mask.prompt_hash) in self._ordering_state:
+                return None
+            self._ensure_cache_has_capacity()
+            if self._send_to_cpu:
+                mask = mask.to(device=torch.device("cpu"))
+            self._ordering_state[(key, mask.prompt_hash)] = True
+            self._cache_state[key].append(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._ordering_state) > self._size_limit:
+            image_key, prompt_hash = self._ordering_state.popitem(last=False)
+            entries_for_image = self._cache_state[image_key]
+            to_remove_idx = None
+            for i, element in enumerate(entries_for_image):
+                if element.prompt_hash == prompt_hash:
+                    to_remove_idx = i
+                    break
+            if to_remove_idx is not None:
+                del entries_for_image[to_remove_idx]

inference_models/models/sam2/entities.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import torch
+
+
+@dataclass(frozen=True)
+class SAM2ImageEmbeddings:
+    image_hash: str
+    image_size_hw: Tuple[int, int]
+    embeddings: torch.Tensor
+    high_resolution_features: List[torch.Tensor]
+
+    def to(self, device: torch.device) -> "SAM2ImageEmbeddings":
+        return SAM2ImageEmbeddings(
+            image_hash=self.image_hash,
+            image_size_hw=self.image_size_hw,
+            embeddings=self.embeddings.to(device=device),
+            high_resolution_features=[
+                f.to(device=device) for f in self.high_resolution_features
+            ],
+        )
+
+
+@dataclass(frozen=True)
+class SAM2Prediction:
+    masks: torch.Tensor
+    scores: torch.Tensor
+    logits: torch.Tensor
+
+
+@dataclass(frozen=True)
+class SAM2MaskCacheEntry:
+    prompt_hash: str
+    serialized_prompt: List[dict]
+    mask: torch.Tensor
+
+    def to(self, device: torch.device) -> "SAM2MaskCacheEntry":
+        return SAM2MaskCacheEntry(
+            prompt_hash=self.prompt_hash,
+            serialized_prompt=self.serialized_prompt,
+            mask=self.mask.to(device=device),
+        )