inference-models 0.18.3__py3-none-any.whl

inference_models/models/base/documents_parsing.py

@@ -0,0 +1,111 @@
+from abc import ABC, abstractmethod
+from typing import Generic, List, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models.models.base.object_detection import Detections
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+
+
+class StructuredOCRModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(cls, model_name_or_path: str, **kwargs) -> "StructuredOCRModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[List[str], List[Detections]]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[PreprocessedInputs, PreprocessingMetadata]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessingMetadata,
+        **kwargs,
+    ) -> Tuple[List[str], List[Detections]]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[List[str], List[Detections]]:
+        return self.infer(images, **kwargs)
+
+
+class TextOnlyOCRModel(ABC, Generic[PreprocessedInputs, RawPrediction]):
+    @classmethod
+    @abstractmethod
+    def from_pretrained(cls, model_name_or_path: str, **kwargs) -> "TextOnlyOCRModel":
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[str]:
+        pre_processed_images = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> PreprocessedInputs:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        **kwargs,
+    ) -> List[str]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[str]:
+        return self.infer(images, **kwargs)

inference_models/models/base/embeddings.py

@@ -0,0 +1,66 @@
+from abc import ABC, abstractmethod
+from typing import List, Literal, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class TextImageEmbeddingModel(ABC):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "TextImageEmbeddingModel":
+        pass
+
+    def compare_embeddings(
+        self,
+        x: Union[
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+            str,
+            List[str],
+        ],
+        y: Union[
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+            str,
+            List[str],
+        ],
+        x_type: Literal["image", "text"] = "image",
+        y_type: Literal["image", "text"] = "text",
+        **kwargs,
+    ) -> torch.Tensor:
+        if x_type == "image":
+            x_embeddings = self.embed_images(images=x, **kwargs)
+        else:
+            x_embeddings = self.embed_text(texts=x, **kwargs)
+        if y_type == "image":
+            y_embeddings = self.embed_images(images=y, **kwargs)
+        else:
+            y_embeddings = self.embed_text(texts=y, **kwargs)
+        x_embeddings_norm = F.normalize(x_embeddings, p=2, dim=1)
+        y_embeddings_morm = F.normalize(y_embeddings, p=2, dim=1)
+        return x_embeddings_norm @ y_embeddings_morm.T
+
+    @abstractmethod
+    def embed_images(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def embed_text(
+        self,
+        texts: Union[str, List[str]],
+        **kwargs,
+    ) -> torch.Tensor:
+        pass

inference_models/models/base/instance_segmentation.py

@@ -0,0 +1,87 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, List, Optional, Tuple, Union
+
+import numpy as np
+import supervision as sv
+import torch
+
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+
+
+@dataclass
+class InstanceDetections:
+    xyxy: torch.Tensor  # (n_boxes, 4)
+    class_id: torch.Tensor  # (n_boxes, )
+    confidence: torch.Tensor  # (n_boxes, )
+    mask: torch.Tensor  # (n_boxes, mask_height, mask_width)
+    image_metadata: Optional[dict] = None
+    bboxes_metadata: Optional[List[dict]] = (
+        None  # if given, list of size equal to # of bboxes
+    )
+
+    def to_supervision(self) -> sv.Detections:
+        return sv.Detections(
+            xyxy=self.xyxy.cpu().numpy(),
+            class_id=self.class_id.cpu().numpy(),
+            confidence=self.confidence.cpu().numpy(),
+            mask=self.mask.cpu().numpy(),
+        )
+
+
+class InstanceSegmentationModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "InstanceSegmentationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[InstanceDetections]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], **kwargs
+    ) -> Tuple[PreprocessedInputs, PreprocessingMetadata]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessedInputs,
+        **kwargs,
+    ) -> List[InstanceDetections]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[InstanceDetections]:
+        return self.infer(images, **kwargs)

inference_models/models/base/keypoints_detection.py

@@ -0,0 +1,93 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, List, Optional, Tuple, Union
+
+import numpy as np
+import supervision as sv
+import torch
+
+from inference_models.models.base.object_detection import Detections
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+
+
+@dataclass
+class KeyPoints:
+    xy: torch.Tensor  # (instances, instance_key_points, 2)
+    class_id: torch.Tensor  # (instances, )
+    confidence: torch.Tensor  # (instances, instance_key_points)
+    image_metadata: Optional[dict] = None
+    key_points_metadata: Optional[List[dict]] = (
+        None  # if given, list of size equal to # of instances
+    )
+
+    def to_supervision(self) -> sv.KeyPoints:
+        return sv.KeyPoints(
+            xy=self.xy.cpu().numpy(),
+            class_id=self.class_id.cpu().numpy(),
+            confidence=self.confidence.cpu().numpy(),
+        )
+
+
+class KeyPointsDetectionModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "KeyPointsDetectionModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    @property
+    @abstractmethod
+    def skeletons(self) -> List[List[Tuple[int, int]]]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[PreprocessedInputs, PreprocessingMetadata]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessingMetadata,
+        **kwargs,
+    ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]:
+        return self.infer(images, **kwargs)

inference_models/models/base/object_detection.py

@@ -0,0 +1,143 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, List, Optional, Tuple, Union
+
+import numpy as np
+import supervision as sv
+import torch
+
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+
+
+@dataclass
+class Detections:
+    xyxy: torch.Tensor  # (n_boxes, 4)
+    class_id: torch.Tensor  # (n_boxes, )
+    confidence: torch.Tensor  # (n_boxes, )
+    image_metadata: Optional[dict] = None
+    bboxes_metadata: Optional[List[dict]] = (
+        None  # if given, list of size equal to # of bboxes
+    )
+
+    def to_supervision(self) -> sv.Detections:
+        return sv.Detections(
+            xyxy=self.xyxy.cpu().numpy(),
+            class_id=self.class_id.cpu().numpy(),
+            confidence=self.confidence.cpu().numpy(),
+        )
+
+
+class ObjectDetectionModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "ObjectDetectionModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[Detections]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[PreprocessedInputs, PreprocessingMetadata]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessingMetadata,
+        **kwargs,
+    ) -> List[Detections]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[Detections]:
+        return self.infer(images, **kwargs)
+
+
+class OpenVocabularyObjectDetectionModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "OpenVocabularyObjectDetectionModel":
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        classes: Union[str, List[str]],
+        **kwargs,
+    ) -> List[Detections]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, classes, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> Tuple[PreprocessedInputs, PreprocessingMetadata]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self,
+        pre_processed_images: PreprocessedInputs,
+        classes: List[str],
+        **kwargs,
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessingMetadata,
+        **kwargs,
+    ) -> List[Detections]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        classes: List[str],
+        **kwargs,
+    ) -> List[Detections]:
+        return self.infer(images=images, classes=classes, **kwargs)

inference_models/models/base/semantic_segmentation.py

@@ -0,0 +1,74 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, List, Optional, Tuple, Union
+
+import numpy as np
+import supervision as sv
+import torch
+
+from inference_models.models.base.types import (
+    PreprocessedInputs,
+    PreprocessingMetadata,
+    RawPrediction,
+)
+
+
+@dataclass
+class SemanticSegmentationResult:
+    segmentation_map: torch.Tensor
+    confidence: torch.Tensor
+    image_metadata: Optional[dict] = None
+
+
+class SemanticSegmentationModel(
+    ABC, Generic[PreprocessedInputs, PreprocessingMetadata, RawPrediction]
+):
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, **kwargs
+    ) -> "SemanticSegmentationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[SemanticSegmentationResult]:
+        pre_processed_images, pre_processing_meta = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], **kwargs
+    ) -> PreprocessedInputs:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: PreprocessedInputs, **kwargs
+    ) -> RawPrediction:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self,
+        model_results: RawPrediction,
+        pre_processing_meta: PreprocessedInputs,
+        **kwargs,
+    ) -> List[SemanticSegmentationResult]:
+        pass
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> List[SemanticSegmentationResult]:
+        return self.infer(images, **kwargs)

inference_models/models/base/types.py

@@ -0,0 +1,5 @@
+from typing import TypeVar
+
+PreprocessedInputs = TypeVar("PreprocessedInputs")
+PreprocessingMetadata = TypeVar("PreprocessingMetadata")
+RawPrediction = TypeVar("RawPrediction")