Add first scratches of new interface

inference/v1/__init__.py

@@ -0,0 +1,13 @@
+# PUBLIC INTERFACE LISTED BELOW
+from .models.base.classification import (
+    ClassificationModel,
+    ClassificationPrediction,
+    MultiLabelClassificationModel,
+    MultiLabelClassificationPrediction,
+)
+from .models.base.instance_segmentation import (
+    InstanceDetections,
+    InstanceSegmentationModel,
+)
+from .models.base.keypoints_detection import KeyPoints, KeyPointsDetectionModel
+from .models.base.object_detection import Detections, ObjectDetectionModel

inference/v1/configuration.py

@@ -0,0 +1,13 @@
+import os
+
+import torch
+
+from inference.v1.utils.environment import parse_comma_separated_values
+
+ONNXRUNTIME_EXECUTION_PROVIDERS = parse_comma_separated_values(
+    values=os.getenv(
+        "ONNXRUNTIME_EXECUTION_PROVIDERS",
+        "CUDAExecutionProvider,OpenVINOExecutionProvider,CoreMLExecutionProvider,CPUExecutionProvider",
+    )
+)
+DEFAULT_DEVICE = torch.get_default_device()

inference/v1/entities.py

@@ -0,0 +1,5 @@
+from collections import namedtuple
+from typing import Literal
+
+ImageDimensions = namedtuple("ImageDimensions", ["height", "width"])
+ColorFormat = Literal["rgb", "bgr"]

inference/v1/errors.py

@@ -0,0 +1,14 @@
+class BaseInferenceError(Exception):
+    pass
+
+
+class CorruptedModelPackageError(BaseInferenceError):
+    pass
+
+
+class EnvironmentConfigurationError(BaseInferenceError):
+    pass
+
+
+class ModelRuntimeError(BaseInferenceError):
+    pass

inference/v1/models/base/classification.py

@@ -0,0 +1,109 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+
+
+@dataclass
+class ClassificationPrediction:
+    class_ids: torch.Tensor  # (bs, )
+    confidence: torch.Tensor  # (bs, )
+    images_metadata: Optional[List[dict]] = None  # if given, list of size equal to bs
+
+
+class ClassificationModel(ABC):
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, *args, **kwargs
+    ) -> "ClassificationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> ClassificationPrediction:
+        pre_processed_images, pre_processing_meta = self.pre_process(
+            images, *args, **kwargs
+        )
+        model_results = self.forward(pre_processed_images, *args, **kwargs)
+        return self.post_process(model_results, *args, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> Tuple[torch.Tensor, Any]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: torch.Tensor, *args, **kwargs
+    ) -> ClassificationPrediction:
+        pass
+
+    def __call__(
+        self, images: torch.Tensor, *args, **kwargs
+    ) -> ClassificationPrediction:
+        return self.infer(images, *args, **kwargs)
+
+
+@dataclass
+class MultiLabelClassificationPrediction:
+    class_ids: torch.Tensor  # (predicted_labels_ids, )
+    confidence: torch.Tensor  # (predicted_labels_confidence, )
+    image_metadata: Optional[dict] = None
+
+
+class MultiLabelClassificationModel(ABC):
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, *args, **kwargs
+    ) -> "MultiLabelClassificationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> List[MultiLabelClassificationPrediction]:
+        pre_processed_images = self.pre_process(images, *args, **kwargs)
+        model_results = self.forward(pre_processed_images, *args, **kwargs)
+        return self.post_process(model_results, *args, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> Tuple[torch.Tensor, Any]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: torch.Tensor, *args, **kwargs
+    ) -> List[MultiLabelClassificationPrediction]:
+        pass
+
+    def __call__(
+        self, images: torch.Tensor, *args, **kwargs
+    ) -> List[MultiLabelClassificationPrediction]:
+        return self.infer(images, *args, **kwargs)

inference/v1/models/base/instance_segmentation.py

@@ -0,0 +1,63 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+
+
+@dataclass
+class InstanceDetections:
+    xyxy: torch.Tensor  # (n_boxes, 4)
+    class_ids: torch.Tensor  # (n_boxes, )
+    confidence: torch.Tensor  # (n_boxes, )
+    masks: 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
+    )
+
+
+class InstanceSegmentationModel(ABC):
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, *args, **kwargs
+    ) -> "InstanceSegmentationModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> List[InstanceDetections]:
+        pre_processed_images, pre_processing_meta = self.pre_process(
+            images, *args, **kwargs
+        )
+        model_results = self.forward(pre_processed_images, *args, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, *args, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> Tuple[torch.Tensor, Any]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: torch.Tensor, pre_processing_meta: Any, *args, **kwargs
+    ) -> List[InstanceDetections]:
+        pass
+
+    def __call__(
+        self, images: torch.Tensor, *args, **kwargs
+    ) -> List[InstanceDetections]:
+        return self.infer(images, *args, **kwargs)

inference/v1/models/base/keypoints_detection.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+
+from inference.v1.models.base.object_detection import Detections
+
+
+@dataclass
+class KeyPoints:
+    xy: torch.Tensor  # (instances, instance_key_points, 2)
+    class_id: torch.Tensor  # (instances, instance_key_points)
+    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
+    )
+
+
+class KeyPointsDetectionModel(ABC):
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, *args, **kwargs
+    ) -> "KeyPointsDetectionModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]:
+        pre_processed_images, pre_processing_meta = self.pre_process(
+            images, *args, **kwargs
+        )
+        model_results = self.forward(pre_processed_images, *args, **kwargs)
+        return self.post_process(model_results, pre_processing_meta, *args, **kwargs)
+
+    @abstractmethod
+    def pre_process(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], *args, **kwargs
+    ) -> Tuple[torch.Tensor, Any]:
+        pass
+
+    @abstractmethod
+    def forward(
+        self, pre_processed_images: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: torch.Tensor, pre_processing_meta: Any, *args, **kwargs
+    ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]:
+        pass
+
+    def __call__(
+        self, images: torch.Tensor, *args, **kwargs
+    ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]:
+        return self.infer(images, *args, **kwargs)

inference/v1/models/base/object_detection.py

@@ -0,0 +1,56 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+
+
+@dataclass
+class Detections:
+    xyxy: torch.Tensor  # (n_boxes, 4)
+    class_ids: 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
+    )
+
+
+class ObjectDetectionModel(ABC):
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name_or_path: str, *args, **kwargs
+    ) -> "ObjectDetectionModel":
+        pass
+
+    @property
+    @abstractmethod
+    def class_names(self) -> List[str]:
+        pass
+
+    def infer(
+        self, images: Union[torch.Tensor, List[torch.Tensor]], **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]], **kwargs
+    ) -> Tuple[torch.Tensor, Any]:
+        pass
+
+    @abstractmethod
+    def forward(self, pre_processed_images: torch.Tensor, **kwargs) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def post_process(
+        self, model_results: torch.Tensor, pre_processing_meta: Any, **kwargs
+    ) -> List[Detections]:
+        pass
+
+    def __call__(self, images: torch.Tensor, **kwargs) -> List[Detections]:
+        return self.infer(images, **kwargs)