inference-models 0.18.3__py3-none-any.whl

inference_models/models/grounding_dino/grounding_dino_torch.py

@@ -0,0 +1,227 @@
+import os.path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torchvision
+from groundingdino.util.inference import load_model, predict
+from torch import nn
+from torchvision import transforms
+from torchvision.ops import box_convert
+
+from inference_models import Detections
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat, ImageDimensions
+from inference_models.errors import ModelRuntimeError
+from inference_models.models.base.object_detection import (
+    OpenVocabularyObjectDetectionModel,
+)
+from inference_models.models.common.model_packages import get_model_package_contents
+
+
+class GroundingDinoForObjectDetectionTorch(
+    OpenVocabularyObjectDetectionModel[
+        torch.Tensor,
+        List[ImageDimensions],
+        Tuple[List[torch.Tensor], List[torch.Tensor], List[List[str]], List[str]],
+    ]
+):
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        **kwargs,
+    ) -> "GroundingDinoForObjectDetectionTorch":
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=["weights.pth", "config.py"],
+        )
+        text_encoder_dir = os.path.join(model_name_or_path, "text_encoder")
+        loader_kwargs = {}
+        if os.path.isdir(text_encoder_dir):
+            loader_kwargs["text_encoder_type"] = text_encoder_dir
+        model = load_model(
+            model_config_path=model_package_content["config.py"],
+            model_checkpoint_path=model_package_content["weights.pth"],
+            **loader_kwargs,
+        ).to(device)
+        return cls(model=model, device=device)
+
+    def __init__(
+        self,
+        model: nn.Module,
+        device: torch.device,
+    ):
+        self._model = model
+        self._device = device
+        self._numpy_transformations = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Resize([800, 800]),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+        self._tensors_transformations = transforms.Compose(
+            [
+                lambda x: x / 255.0,
+                transforms.Resize([800, 800]),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, List[ImageDimensions]]:
+        if isinstance(images, np.ndarray):
+            input_color_format = input_color_format or "bgr"
+            if input_color_format != "rgb":
+                images = np.ascontiguousarray(images[:, :, ::-1])
+            pre_processed = self._numpy_transformations(images)
+            return (
+                torch.unsqueeze(pre_processed, dim=0).to(self._device),
+                [ImageDimensions(height=images.shape[0], width=images.shape[1])],
+            )
+        if isinstance(images, torch.Tensor):
+            input_color_format = input_color_format or "rgb"
+            if len(images.shape) == 3:
+                images = torch.unsqueeze(images, dim=0)
+            image_dimensions = ImageDimensions(
+                height=images.shape[2], width=images.shape[3]
+            )
+            images = images.to(self._device)
+            if input_color_format != "rgb":
+                images = images[:, [2, 1, 0], :, :]
+            return (
+                self._tensors_transformations(images.float()),
+                [image_dimensions] * images.shape[0],
+            )
+        if not isinstance(images, list):
+            raise ModelRuntimeError(
+                message="Pre-processing supports only np.array or torch.Tensor or list of above.",
+                help_url="https://todo",
+            )
+        if not len(images):
+            raise ModelRuntimeError(
+                message="Detected empty input to the model",
+                help_url="https://todo",
+            )
+        if isinstance(images[0], np.ndarray):
+            input_color_format = input_color_format or "bgr"
+            pre_processed, image_dimensions = [], []
+            for image in images:
+                if input_color_format != "rgb":
+                    image = np.ascontiguousarray(image[:, :, ::-1])
+                image_dimensions.append(
+                    ImageDimensions(height=image.shape[0], width=image.shape[1])
+                )
+                pre_processed.append(self._numpy_transformations(image))
+            return torch.stack(pre_processed, dim=0).to(self._device), image_dimensions
+        if isinstance(images[0], torch.Tensor):
+            input_color_format = input_color_format or "rgb"
+            pre_processed, image_dimensions = [], []
+            for image in images:
+                if len(image.shape) == 3:
+                    image = torch.unsqueeze(image, dim=0)
+                if input_color_format != "rgb":
+                    image = image[:, [2, 1, 0], :, :]
+                image_dimensions.append(
+                    ImageDimensions(height=image.shape[2], width=image.shape[3])
+                )
+                pre_processed.append(self._tensors_transformations(image.float()))
+            return torch.cat(pre_processed, dim=0).to(self._device), image_dimensions
+        raise ModelRuntimeError(
+            message=f"Detected unknown input batch element: {type(images[0])}",
+            help_url="https://todo",
+        )
+
+    def forward(
+        self,
+        pre_processed_images: torch.Tensor,
+        classes: List[str],
+        conf_thresh: float = 0.5,
+        text_threshold: Optional[float] = None,
+        **kwargs,
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[List[str]], List[str]]:
+        if text_threshold is None:
+            text_threshold = conf_thresh
+        caption = ". ".join(classes)
+        all_boxes, all_logits, all_phrases = [], [], []
+        with torch.inference_mode():
+            for image in pre_processed_images:
+                boxes, logits, phrases = predict(
+                    model=self._model,
+                    image=image,
+                    caption=caption,
+                    box_threshold=conf_thresh,
+                    text_threshold=text_threshold,
+                    device=self._device,
+                    remove_combined=True,
+                )
+                all_boxes.append(boxes)
+                all_logits.append(logits)
+                all_phrases.append(phrases)
+        return all_boxes, all_logits, all_phrases, classes
+
+    def post_process(
+        self,
+        model_results: Tuple[
+            List[torch.Tensor], List[torch.Tensor], List[List[str]], List[str]
+        ],
+        pre_processing_meta: List[ImageDimensions],
+        iou_thresh: float = 0.45,
+        max_detections: int = 100,
+        class_agnostic: bool = False,
+        **kwargs,
+    ) -> List[Detections]:
+        all_boxes, all_logits, all_phrases, classes = model_results
+        results = []
+        for boxes, logits, phrases, origin_size in zip(
+            all_boxes, all_logits, all_phrases, pre_processing_meta
+        ):
+            boxes = boxes * torch.Tensor(
+                [
+                    origin_size.width,
+                    origin_size.height,
+                    origin_size.width,
+                    origin_size.height,
+                ],
+                device=boxes.device,
+            )
+            xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy")
+            class_id = map_phrases_to_classes(
+                phrases=phrases,
+                classes=classes,
+            ).to(boxes.device)
+            nms_class_ids = torch.zeros_like(class_id) if class_agnostic else class_id
+            keep = torchvision.ops.batched_nms(xyxy, logits, nms_class_ids, iou_thresh)
+            if keep.numel() > max_detections:
+                keep = keep[:max_detections]
+            results.append(
+                Detections(
+                    xyxy=xyxy[keep].round().int(),
+                    confidence=logits[keep],
+                    class_id=class_id[keep].int(),
+                ),
+            )
+        return results
+
+
+def map_phrases_to_classes(phrases: List[str], classes: List[str]) -> torch.Tensor:
+    class_ids = []
+    for phrase in phrases:
+        for class_ in classes:
+            if class_ in phrase:
+                class_ids.append(classes.index(class_))
+                break
+        else:
+            # TODO: figure out how to mark additional classes
+            class_ids.append(len(classes))
+    return torch.tensor(class_ids)

inference_models/models/l2cs/l2cs_onnx.py

@@ -0,0 +1,216 @@
+from dataclasses import dataclass
+from threading import Lock
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torchvision import transforms
+
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.errors import (
+    EnvironmentConfigurationError,
+    MissingDependencyError,
+    ModelRuntimeError,
+)
+from inference_models.models.base.types import PreprocessedInputs
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.common.onnx import (
+    run_session_via_iobinding,
+    run_session_with_batch_size_limit,
+    set_execution_provider_defaults,
+)
+from inference_models.utils.onnx_introspection import (
+    get_selected_onnx_execution_providers,
+)
+
+try:
+    import onnxruntime
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import L2CS model with ONNX backend - this error means that some additional dependencies "
+        f"are not installed in the environment. If you run the `inference-models` library directly in your Python "
+        f"program, make sure the following extras of the package are installed: \n"
+        f"\t* `onnx-cpu` - when you wish to use library with CPU support only\n"
+        f"\t* `onnx-cu12` - for running on GPU with Cuda 12 installed\n"
+        f"\t* `onnx-cu118` - for running on GPU with Cuda 11.8 installed\n"
+        f"\t* `onnx-jp6-cu126` - for running on Jetson with Jetpack 6\n"
+        f"If you see this error using Roboflow infrastructure, make sure the service you use does support the model. "
+        f"You can also contact Roboflow to get support.",
+        help_url="https://todo",
+    ) from import_error
+
+
+DEFAULT_GAZE_MAX_BATCH_SIZE = 8
+
+
+@dataclass
+class L2CSGazeDetection:
+    yaw: torch.Tensor
+    pitch: torch.Tensor
+
+
+class L2CSNetOnnx:
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        onnx_execution_providers: Optional[List[Union[str, tuple]]] = None,
+        default_onnx_trt_options: bool = True,
+        device: torch.device = DEFAULT_DEVICE,
+        max_batch_size: int = DEFAULT_GAZE_MAX_BATCH_SIZE,
+        **kwargs,
+    ):
+        if onnx_execution_providers is None:
+            onnx_execution_providers = get_selected_onnx_execution_providers()
+        if not onnx_execution_providers:
+            raise EnvironmentConfigurationError(
+                message=f"Could not initialize model - selected backend is ONNX which requires execution provider to "
+                f"be specified - explicitly in `from_pretrained(...)` method or via env variable "
+                f"`ONNXRUNTIME_EXECUTION_PROVIDERS`. If you run model locally - adjust your setup, otherwise "
+                f"contact the platform support.",
+                help_url="https://todo",
+            )
+        onnx_execution_providers = set_execution_provider_defaults(
+            providers=onnx_execution_providers,
+            model_package_path=model_name_or_path,
+            device=device,
+            default_onnx_trt_options=default_onnx_trt_options,
+        )
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=["weights.onnx"],
+        )
+        session = onnxruntime.InferenceSession(
+            path_or_bytes=model_package_content["weights.onnx"],
+            providers=onnx_execution_providers,
+        )
+        input_name = session.get_inputs()[0].name
+        return cls(
+            session=session,
+            max_batch_size=max_batch_size,
+            device=device,
+            input_name=input_name,
+        )
+
+    def __init__(
+        self,
+        session: onnxruntime.InferenceSession,
+        max_batch_size: int,
+        device: torch.device,
+        input_name: str,
+    ):
+        self._session = session
+        self._max_batch_size = max_batch_size
+        self._device = device
+        self._input_name = input_name
+        self._session_thread_lock = Lock()
+        self._numpy_transformations = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Resize([448, 448]),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+        self._tensors_transformations = transforms.Compose(
+            [
+                lambda x: x / 255.0,
+                transforms.Resize([448, 448]),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return self._device
+
+    def infer(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> L2CSGazeDetection:
+        pre_processed_images = self.pre_process(images, **kwargs)
+        model_results = self.forward(pre_processed_images, **kwargs)
+        return self.post_process(model_results, **kwargs)
+
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if isinstance(images, np.ndarray):
+            input_color_format = input_color_format or "bgr"
+            if input_color_format != "rgb":
+                images = np.ascontiguousarray(images[:, :, ::-1])
+            pre_processed = self._numpy_transformations(images)
+            return torch.unsqueeze(pre_processed, dim=0).to(self._device)
+        if isinstance(images, torch.Tensor):
+            input_color_format = input_color_format or "rgb"
+            if len(images.shape) == 3:
+                images = torch.unsqueeze(images, dim=0)
+            images = images.to(self._device)
+            if input_color_format != "rgb":
+                images = images[:, [2, 1, 0], :, :]
+            return self._tensors_transformations(images.float())
+        if not isinstance(images, list):
+            raise ModelRuntimeError(
+                message="Pre-processing supports only np.array or torch.Tensor or list of above.",
+                help_url="https://todo",
+            )
+        if not len(images):
+            raise ModelRuntimeError(
+                message="Detected empty input to the model", help_url="https://todo"
+            )
+        if isinstance(images[0], np.ndarray):
+            input_color_format = input_color_format or "bgr"
+            pre_processed = []
+            for image in images:
+                if input_color_format != "rgb":
+                    image = np.ascontiguousarray(image[:, :, ::-1])
+                pre_processed.append(self._numpy_transformations(image))
+            return torch.stack(pre_processed, dim=0).to(self._device)
+        if isinstance(images[0], torch.Tensor):
+            input_color_format = input_color_format or "rgb"
+            pre_processed = []
+            for image in images:
+                if len(image.shape) == 3:
+                    image = torch.unsqueeze(image, dim=0)
+                if input_color_format != "rgb":
+                    image = image[:, [2, 1, 0], :, :]
+                pre_processed.append(self._tensors_transformations(image.float()))
+            return torch.cat(pre_processed, dim=0).to(self._device)
+        raise ModelRuntimeError(
+            message=f"Detected unknown input batch element: {type(images[0])}",
+            help_url="https://todo",
+        )
+
+    def forward(
+        self,
+        pre_processed_images: PreprocessedInputs,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        with self._session_thread_lock:
+            yaw, pitch = run_session_with_batch_size_limit(
+                session=self._session, inputs={self._input_name: pre_processed_images}
+            )
+            return yaw, pitch
+
+    def post_process(
+        self,
+        model_results: Tuple[torch.Tensor, torch.Tensor],
+        **kwargs,
+    ) -> L2CSGazeDetection:
+        return L2CSGazeDetection(yaw=model_results[0], pitch=model_results[1])
+
+    def __call__(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> L2CSGazeDetection:
+        return self.infer(images, **kwargs)

inference_models/models/mediapipe_face_detection/face_detection.py

@@ -0,0 +1,203 @@
+from threading import Lock
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import Detections, KeyPoints, KeyPointsDetectionModel
+from inference_models.entities import ColorFormat, ImageDimensions
+from inference_models.errors import MissingDependencyError, ModelRuntimeError
+from inference_models.models.common.model_packages import get_model_package_contents
+
+try:
+    import mediapipe as mp
+    from mediapipe.tasks.python.components.containers import Detection
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import face detection model from MediaPipe - this error means that some additional "
+        f"dependencies are not installed in the environment. If you run the `inference-models` library directly in your Python "
+        f"program, make sure the following extras of the package are installed: `mediapipe`."
+        f"If you see this error using Roboflow infrastructure, make sure the service you use does support the model. "
+        f"You can also contact Roboflow to get support.",
+        help_url="https://todo",
+    ) from import_error
+
+
+class MediaPipeFaceDetector(
+    KeyPointsDetectionModel[List[mp.Image], ImageDimensions, List[Detection]]
+):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        **kwargs,
+    ) -> "MediaPipeFaceDetector":
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=["mediapipe_face_detector.tflite"],
+        )
+        face_detector = mp.tasks.vision.FaceDetector.create_from_options(
+            mp.tasks.vision.FaceDetectorOptions(
+                base_options=mp.tasks.BaseOptions(
+                    model_asset_path=model_package_content[
+                        "mediapipe_face_detector.tflite"
+                    ]
+                ),
+                running_mode=mp.tasks.vision.RunningMode.IMAGE,
+            )
+        )
+        return cls(face_detector=face_detector)
+
+    def __init__(self, face_detector: mp.tasks.vision.FaceDetector):
+        self._face_detector = face_detector
+        self._thread_lock = Lock()
+
+    @property
+    def class_names(self) -> List[str]:
+        return ["face"]
+
+    @property
+    def key_points_classes(self) -> List[List[str]]:
+        return [["right-eye", "left-eye", "nose", "mouth", "right-ear", "left-ear"]]
+
+    @property
+    def skeletons(self) -> List[List[Tuple[int, int]]]:
+        return [[(5, 1), (1, 2), (4, 0), (0, 2), (2, 3)]]
+
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> Tuple[List[mp.Image], List[ImageDimensions]]:
+        if isinstance(images, np.ndarray):
+            input_color_format = input_color_format or "bgr"
+            if input_color_format != "rgb":
+                images = np.ascontiguousarray(images[:, :, ::-1])
+            preprocessed_images = mp.Image(
+                image_format=mp.ImageFormat.SRGB, data=images.astype(np.uint8)
+            )
+            dimensions = ImageDimensions(height=images.shape[0], width=images.shape[1])
+            return [preprocessed_images], [dimensions]
+        if isinstance(images, torch.Tensor):
+            input_color_format = input_color_format or "rgb"
+            if len(images.shape) == 3:
+                images = torch.unsqueeze(images, dim=0)
+            if input_color_format != "rgb":
+                images = images[:, [2, 1, 0], :, :]
+            images = images.permute(0, 2, 3, 1)
+            preprocessed_images, dimensions = [], []
+            for image in images:
+                np_image = np.ascontiguousarray(image.cpu().numpy())
+                preprocessed_images.append(
+                    mp.Image(
+                        image_format=mp.ImageFormat.SRGB, data=np_image.astype(np.uint8)
+                    )
+                )
+                dimensions.append(
+                    ImageDimensions(height=np_image.shape[0], width=np_image.shape[1])
+                )
+            return preprocessed_images, dimensions
+        if not isinstance(images, list):
+            raise ModelRuntimeError(
+                message="Pre-processing supports only np.array or torch.Tensor or list of above.",
+                help_url="https://todo",
+            )
+        if not len(images):
+            raise ModelRuntimeError(
+                message="Detected empty input to the model", help_url="https://todo"
+            )
+        if isinstance(images[0], np.ndarray):
+            input_color_format = input_color_format or "bgr"
+            preprocessed_images, dimensions = [], []
+            for image in images:
+                if input_color_format != "rgb":
+                    image = np.ascontiguousarray(image[:, :, ::-1])
+                preprocessed_images.append(
+                    mp.Image(
+                        image_format=mp.ImageFormat.SRGB, data=image.astype(np.uint8)
+                    )
+                )
+                dimensions.append(
+                    ImageDimensions(height=image.shape[0], width=image.shape[1])
+                )
+            return preprocessed_images, dimensions
+        if isinstance(images[0], torch.Tensor):
+            input_color_format = input_color_format or "rgb"
+            preprocessed_images, dimensions = [], []
+            for image in images:
+                if input_color_format != "rgb":
+                    image = image[[2, 1, 0], :, :]
+                np_image = image.cpu().permute(1, 2, 0).numpy()
+                preprocessed_images.append(
+                    mp.Image(
+                        image_format=mp.ImageFormat.SRGB, data=np_image.astype(np.uint8)
+                    )
+                )
+                dimensions.append(
+                    ImageDimensions(height=np_image.shape[0], width=np_image.shape[1])
+                )
+            return preprocessed_images, dimensions
+        raise ModelRuntimeError(
+            message=f"Detected unknown input batch element: {type(images[0])}",
+            help_url="https://todo",
+        )
+
+    def forward(
+        self, pre_processed_images: List[mp.Image], **kwargs
+    ) -> List[List[Detection]]:
+        results = []
+        with self._thread_lock:
+            for input_image in pre_processed_images:
+                image_faces = self._face_detector.detect(image=input_image).detections
+                results.append(image_faces)
+        return results
+
+    def post_process(
+        self,
+        model_results: List[List[Detection]],
+        pre_processing_meta: List[ImageDimensions],
+        conf_thresh: float = 0.25,
+        **kwargs,
+    ) -> Tuple[List[KeyPoints], List[Detections]]:
+        final_key_points, final_detections = [], []
+        for image_results, image_dimensions in zip(model_results, pre_processing_meta):
+            detections_xyxy, detections_class_id, detections_confidence = [], [], []
+            key_points_xy, key_points_class_id, key_points_confidence = [], [], []
+            for detection in image_results:
+                if detection.categories[0].score < conf_thresh:
+                    continue
+                xyxy = (
+                    detection.bounding_box.origin_x,
+                    detection.bounding_box.origin_y,
+                    detection.bounding_box.origin_x + detection.bounding_box.width,
+                    detection.bounding_box.origin_y + detection.bounding_box.height,
+                )
+                detections_xyxy.append(xyxy)
+                detections_class_id.append(0)
+                detections_confidence.append(detection.categories[0].score)
+                detection_key_points = []
+                for keypoint in detection.keypoints:
+                    detection_key_points.append(
+                        (
+                            keypoint.x * image_dimensions.width,
+                            keypoint.y * image_dimensions.height,
+                        )
+                    )
+                key_points_xy.append(detection_key_points)
+                key_points_class_id.append(0)
+                key_points_confidence.append([1.0] * len(detection_key_points))
+            detections = Detections(
+                xyxy=torch.tensor(detections_xyxy).round().int(),
+                class_id=torch.tensor(detections_class_id).int(),
+                confidence=torch.tensor(detections_confidence),
+            )
+            key_points = KeyPoints(
+                xy=torch.tensor(key_points_xy).round().int(),
+                class_id=torch.tensor(key_points_class_id).int(),
+                confidence=torch.tensor(key_points_confidence),
+            )
+            final_key_points.append(key_points)
+            final_detections.append(detections)
+        return final_key_points, final_detections