inference-models 0.18.3__py3-none-any.whl

inference_models/models/yolov11/yolov11_trt.py

@@ -0,0 +1,21 @@
+from inference_models.models.yolov8.yolov8_instance_segmentation_trt import (
+    YOLOv8ForInstanceSegmentationTRT,
+)
+from inference_models.models.yolov8.yolov8_key_points_detection_trt import (
+    YOLOv8ForKeyPointsDetectionTRT,
+)
+from inference_models.models.yolov8.yolov8_object_detection_trt import (
+    YOLOv8ForObjectDetectionTRT,
+)
+
+
+class YOLOv11ForObjectDetectionTRT(YOLOv8ForObjectDetectionTRT):
+    pass
+
+
+class YOLOv11ForInstanceSegmentationTRT(YOLOv8ForInstanceSegmentationTRT):
+    pass
+
+
+class YOLOv11ForForKeyPointsDetectionTRT(YOLOv8ForKeyPointsDetectionTRT):
+    pass

inference_models/models/yolov12/yolov12_onnx.py

@@ -0,0 +1,7 @@
+from inference_models.models.yolov8.yolov8_object_detection_onnx import (
+    YOLOv8ForObjectDetectionOnnx,
+)
+
+
+class YOLOv12ForObjectDetectionOnnx(YOLOv8ForObjectDetectionOnnx):
+    pass

inference_models/models/yolov12/yolov12_torch_script.py

@@ -0,0 +1,7 @@
+from inference_models.models.yolov8.yolov8_object_detection_torch_script import (
+    YOLOv8ForObjectDetectionTorchScript,
+)
+
+
+class YOLOv12ForObjectDetectionTorchScript(YOLOv8ForObjectDetectionTorchScript):
+    pass

inference_models/models/yolov12/yolov12_trt.py

@@ -0,0 +1,7 @@
+from inference_models.models.yolov8.yolov8_object_detection_trt import (
+    YOLOv8ForObjectDetectionTRT,
+)
+
+
+class YOLOv12ForObjectDetectionTRT(YOLOv8ForObjectDetectionTRT):
+    pass

inference_models/models/yolov5/nms.py

@@ -0,0 +1,99 @@
+from typing import List
+
+import torch
+import torchvision
+
+
+def run_nms_yolov5(
+    output: torch.Tensor,
+    conf_thresh: float = 0.25,
+    iou_thresh: float = 0.45,
+    max_detections: int = 100,
+    class_agnostic: bool = False,
+) -> List[torch.Tensor]:
+    bs = output.shape[0]
+    boxes = output[:, :4, :]
+    top_classes_conf = output[:, 4, :]
+    scores = output[:, 5:, :]
+    results = []
+    for b in range(bs):
+        class_scores = scores[b]
+        class_conf, class_ids = class_scores.max(0)
+        mask = top_classes_conf[b] > conf_thresh
+        if not torch.any(mask):
+            results.append(torch.zeros((0, 6), device=output.device))
+            continue
+        bboxes = boxes[b][:, mask].T
+        class_conf = class_conf[mask]
+        class_ids = class_ids[mask]
+        # Vectorized [x, y, w, h] -> [x1, y1, x2, y2]
+        xy = bboxes[:, :2]
+        wh = bboxes[:, 2:]
+        half_wh = wh / 2
+        xyxy = torch.cat((xy - half_wh, xy + half_wh), 1)
+        # Class-agnostic NMS -> use dummy class ids
+        nms_class_ids = torch.zeros_like(class_ids) if class_agnostic else class_ids
+        # NMS and limiting max detections
+        keep = torchvision.ops.batched_nms(xyxy, class_conf, nms_class_ids, iou_thresh)
+        if keep.numel() > max_detections:
+            keep = keep[:max_detections]
+        detections = torch.cat(
+            (
+                xyxy[keep],
+                class_conf[keep, None],  # unsqueeze(1) is replaced with None
+                class_ids[keep, None].float(),
+            ),
+            1,
+        )  # [x1, y1, x2, y2, conf, cls]
+        results.append(detections)
+    return results
+
+
+def run_yolov5_nms_for_instance_segmentation(
+    output: torch.Tensor,
+    conf_thresh: float = 0.25,
+    iou_thresh: float = 0.45,
+    max_detections: int = 100,
+    class_agnostic: bool = False,
+) -> List[torch.Tensor]:
+    bs = output.shape[0]
+    boxes = output[:, :4, :]
+    top_classes_conf = output[:, 4, :]
+    scores = output[:, 4:-32, :]
+    masks = output[:, -32:, :]
+    results = []
+
+    for b in range(bs):
+        bboxes = boxes[b].T
+        class_scores = scores[b].T
+        box_masks = masks[b].T
+        class_conf, class_ids = class_scores.max(1)
+        mask = top_classes_conf[b] > conf_thresh
+        if mask.sum() == 0:
+            results.append(torch.zeros((0, 38), device=output.device))
+            continue
+        bboxes = bboxes[mask]
+        class_conf = class_conf[mask]
+        class_ids = class_ids[mask]
+        box_masks = box_masks[mask]
+        # Convert [x, y, w, h] -> [x1, y1, x2, y2]
+        xyxy = torch.zeros_like(bboxes)
+        xyxy[:, 0] = bboxes[:, 0] - bboxes[:, 2] / 2  # x1
+        xyxy[:, 1] = bboxes[:, 1] - bboxes[:, 3] / 2  # y1
+        xyxy[:, 2] = bboxes[:, 0] + bboxes[:, 2] / 2  # x2
+        xyxy[:, 3] = bboxes[:, 1] + bboxes[:, 3] / 2  # y2
+        # Class-agnostic NMS -> use dummy class ids
+        nms_class_ids = torch.zeros_like(class_ids) if class_agnostic else class_ids
+        keep = torchvision.ops.batched_nms(xyxy, class_conf, nms_class_ids, iou_thresh)
+        keep = keep[:max_detections]
+        detections = torch.cat(
+            [
+                xyxy[keep],
+                class_conf[keep].unsqueeze(1),
+                class_ids[keep].unsqueeze(1).float(),
+                box_masks[keep],
+            ],
+            dim=1,
+        )  # [x1, y1, x2, y2, conf, cls]
+        results.append(detections)
+    return results

inference_models/models/yolov5/yolov5_instance_segmentation_onnx.py

@@ -0,0 +1,225 @@
+from threading import Lock
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import InstanceDetections, InstanceSegmentationModel
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.errors import (
+    EnvironmentConfigurationError,
+    MissingDependencyError,
+)
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.common.onnx import (
+    run_session_with_batch_size_limit,
+    set_execution_provider_defaults,
+)
+from inference_models.models.common.roboflow.model_packages import (
+    InferenceConfig,
+    PreProcessingMetadata,
+    ResizeMode,
+    parse_class_names_file,
+    parse_inference_config,
+)
+from inference_models.models.common.roboflow.post_processing import (
+    align_instance_segmentation_results,
+    crop_masks_to_boxes,
+    preprocess_segmentation_masks,
+)
+from inference_models.models.common.roboflow.pre_processing import (
+    pre_process_network_input,
+)
+from inference_models.models.yolov5.nms import run_yolov5_nms_for_instance_segmentation
+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 YOLOv5 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
+
+
+class YOLOv5ForInstanceSegmentationOnnx(
+    InstanceSegmentationModel[
+        torch.Tensor, PreProcessingMetadata, Tuple[torch.Tensor, torch.Tensor]
+    ]
+):
+
+    @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,
+        **kwargs,
+    ) -> "YOLOv5ForInstanceSegmentationOnnx":
+        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=[
+                "class_names.txt",
+                "inference_config.json",
+                "weights.onnx",
+            ],
+        )
+        class_names = parse_class_names_file(
+            class_names_path=model_package_content["class_names.txt"]
+        )
+        inference_config = parse_inference_config(
+            config_path=model_package_content["inference_config.json"],
+            allowed_resize_modes={
+                ResizeMode.STRETCH_TO,
+                ResizeMode.LETTERBOX,
+                ResizeMode.CENTER_CROP,
+                ResizeMode.LETTERBOX_REFLECT_EDGES,
+            },
+        )
+        session = onnxruntime.InferenceSession(
+            path_or_bytes=model_package_content["weights.onnx"],
+            providers=onnx_execution_providers,
+        )
+        input_batch_size = session.get_inputs()[0].shape[0]
+        if isinstance(input_batch_size, str):
+            input_batch_size = None
+        input_name = session.get_inputs()[0].name
+        return cls(
+            session=session,
+            input_name=input_name,
+            class_names=class_names,
+            inference_config=inference_config,
+            device=device,
+            input_batch_size=input_batch_size,
+        )
+
+    def __init__(
+        self,
+        session: onnxruntime.InferenceSession,
+        input_name: str,
+        inference_config: InferenceConfig,
+        class_names: List[str],
+        device: torch.device,
+        input_batch_size: Optional[int],
+    ):
+        self._session = session
+        self._input_name = input_name
+        self._inference_config = inference_config
+        self._class_names = class_names
+        self._device = device
+        self._input_batch_size = input_batch_size
+        self._session_thread_lock = Lock()
+
+    @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]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]:
+        return pre_process_network_input(
+            images=images,
+            image_pre_processing=self._inference_config.image_pre_processing,
+            network_input=self._inference_config.network_input,
+            target_device=self._device,
+            input_color_format=input_color_format,
+        )
+
+    def forward(
+        self, pre_processed_images: torch.Tensor, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        with self._session_thread_lock:
+            instances, protos = run_session_with_batch_size_limit(
+                session=self._session,
+                inputs={self._input_name: pre_processed_images},
+                min_batch_size=self._input_batch_size,
+                max_batch_size=self._input_batch_size,
+            )
+            return instances, protos
+
+    def post_process(
+        self,
+        model_results: Tuple[torch.Tensor, torch.Tensor],
+        pre_processing_meta: List[PreProcessingMetadata],
+        conf_thresh: float = 0.25,
+        iou_thresh: float = 0.45,
+        max_detections: int = 100,
+        class_agnostic: bool = False,
+        **kwargs,
+    ) -> List[InstanceDetections]:
+        instances, protos = model_results
+        nms_results = run_yolov5_nms_for_instance_segmentation(
+            output=instances.permute(0, 2, 1),
+            conf_thresh=conf_thresh,
+            iou_thresh=iou_thresh,
+            max_detections=max_detections,
+            class_agnostic=class_agnostic,
+        )
+        final_results = []
+        for image_bboxes, image_protos, image_meta in zip(
+            nms_results, protos, pre_processing_meta
+        ):
+            pre_processed_masks = preprocess_segmentation_masks(
+                protos=image_protos,
+                masks_in=image_bboxes[:, 6:],
+            )
+            cropped_masks = crop_masks_to_boxes(
+                image_bboxes[:, :4], pre_processed_masks
+            )
+            padding = (
+                image_meta.pad_left,
+                image_meta.pad_top,
+                image_meta.pad_right,
+                image_meta.pad_bottom,
+            )
+            aligned_boxes, aligned_masks = align_instance_segmentation_results(
+                image_bboxes=image_bboxes,
+                masks=cropped_masks,
+                padding=padding,
+                scale_height=image_meta.scale_height,
+                scale_width=image_meta.scale_width,
+                original_size=image_meta.original_size,
+                size_after_pre_processing=image_meta.size_after_pre_processing,
+                inference_size=image_meta.inference_size,
+                static_crop_offset=image_meta.static_crop_offset,
+            )
+            final_results.append(
+                InstanceDetections(
+                    xyxy=aligned_boxes[:, :4].round().int(),
+                    class_id=aligned_boxes[:, 5].int(),
+                    confidence=aligned_boxes[:, 4],
+                    mask=aligned_masks,
+                )
+            )
+        return final_results

inference_models/models/yolov5/yolov5_instance_segmentation_trt.py

@@ -0,0 +1,255 @@
+from threading import Lock
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import InstanceDetections, InstanceSegmentationModel
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.errors import (
+    CorruptedModelPackageError,
+    MissingDependencyError,
+    ModelRuntimeError,
+)
+from inference_models.models.common.cuda import (
+    use_cuda_context,
+    use_primary_cuda_context,
+)
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.common.roboflow.model_packages import (
+    InferenceConfig,
+    PreProcessingMetadata,
+    ResizeMode,
+    TRTConfig,
+    parse_class_names_file,
+    parse_inference_config,
+    parse_trt_config,
+)
+from inference_models.models.common.roboflow.post_processing import (
+    align_instance_segmentation_results,
+    crop_masks_to_boxes,
+    preprocess_segmentation_masks,
+)
+from inference_models.models.common.roboflow.pre_processing import (
+    pre_process_network_input,
+)
+from inference_models.models.common.trt import (
+    get_engine_inputs_and_outputs,
+    infer_from_trt_engine,
+    load_model,
+)
+from inference_models.models.yolov5.nms import run_yolov5_nms_for_instance_segmentation
+
+try:
+    import tensorrt as trt
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import YOLOv5 model with TRT 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: `trt10` - installation can only "
+        f"succeed for Linux and Windows machines with Cuda 12 installed. Jetson devices, should have TRT 10.x "
+        f"installed for all builds with Jetpack 6. "
+        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
+
+try:
+    import pycuda.driver as cuda
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message="TODO", help_url="https://todo"
+    ) from import_error
+
+
+class YOLOv5ForInstanceSegmentationTRT(
+    InstanceSegmentationModel[
+        torch.Tensor, PreProcessingMetadata, Tuple[torch.Tensor, torch.Tensor]
+    ]
+):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        engine_host_code_allowed: bool = False,
+        **kwargs,
+    ) -> "YOLOv5ForInstanceSegmentationTRT":
+        if device.type != "cuda":
+            raise ModelRuntimeError(
+                message=f"TRT engine only runs on CUDA device - {device} device detected.",
+                help_url="https://todo",
+            )
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=[
+                "class_names.txt",
+                "inference_config.json" "trt_config.json",
+                "engine.plan",
+            ],
+        )
+        class_names = parse_class_names_file(
+            class_names_path=model_package_content["class_names.txt"]
+        )
+        inference_config = parse_inference_config(
+            config_path=model_package_content["inference_config.json"],
+            allowed_resize_modes={
+                ResizeMode.STRETCH_TO,
+                ResizeMode.LETTERBOX,
+                ResizeMode.CENTER_CROP,
+                ResizeMode.LETTERBOX_REFLECT_EDGES,
+            },
+        )
+        trt_config = parse_trt_config(
+            config_path=model_package_content["trt_config.json"]
+        )
+        cuda.init()
+        cuda_device = cuda.Device(device.index or 0)
+        with use_primary_cuda_context(cuda_device=cuda_device) as cuda_context:
+            engine = load_model(
+                model_path=model_package_content["engine.plan"],
+                engine_host_code_allowed=engine_host_code_allowed,
+            )
+            execution_context = engine.create_execution_context()
+        inputs, outputs = get_engine_inputs_and_outputs(engine=engine)
+        if len(inputs) != 1:
+            raise CorruptedModelPackageError(
+                message=f"Implementation assume single model input, found: {len(inputs)}.",
+                help_url="https://todo",
+            )
+        if len(outputs) != 2:
+            raise CorruptedModelPackageError(
+                message=f"Implementation assume 2 model outputs, found: {len(outputs)}.",
+                help_url="https://todo",
+            )
+        if "output0" not in outputs or "output1" not in outputs:
+            raise CorruptedModelPackageError(
+                message=f"Expected model outputs to be named `output0` and `output1`, but found: {outputs}.",
+                help_url="https://todo",
+            )
+        return cls(
+            engine=engine,
+            input_name=inputs[0],
+            output_names=["output0", "output1"],
+            class_names=class_names,
+            inference_config=inference_config,
+            trt_config=trt_config,
+            device=device,
+            cuda_context=cuda_context,
+            execution_context=execution_context,
+        )
+
+    def __init__(
+        self,
+        engine: trt.ICudaEngine,
+        input_name: str,
+        output_names: List[str],
+        class_names: List[str],
+        inference_config: InferenceConfig,
+        trt_config: TRTConfig,
+        device: torch.device,
+        cuda_context: cuda.Context,
+        execution_context: trt.IExecutionContext,
+    ):
+        self._engine = engine
+        self._input_name = input_name
+        self._output_names = output_names
+        self._class_names = class_names
+        self._inference_config = inference_config
+        self._trt_config = trt_config
+        self._device = device
+        self._cuda_context = cuda_context
+        self._execution_context = execution_context
+        self._session_thread_lock = Lock()
+
+    @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]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]:
+        return pre_process_network_input(
+            images=images,
+            image_pre_processing=self._inference_config.image_pre_processing,
+            network_input=self._inference_config.network_input,
+            target_device=self._device,
+            input_color_format=input_color_format,
+        )
+
+    def forward(
+        self, pre_processed_images: torch.Tensor, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        with self._session_thread_lock:
+            with use_cuda_context(context=self._cuda_context):
+                instances, protos = infer_from_trt_engine(
+                    pre_processed_images=pre_processed_images,
+                    trt_config=self._trt_config,
+                    engine=self._engine,
+                    context=self._execution_context,
+                    device=self._device,
+                    input_name=self._input_name,
+                    outputs=self._output_names,
+                )
+                return instances, protos
+
+    def post_process(
+        self,
+        model_results: Tuple[torch.Tensor, torch.Tensor],
+        pre_processing_meta: List[PreProcessingMetadata],
+        conf_thresh: float = 0.25,
+        iou_thresh: float = 0.45,
+        max_detections: int = 100,
+        class_agnostic: bool = False,
+        **kwargs,
+    ) -> List[InstanceDetections]:
+        instances, protos = model_results
+        nms_results = run_yolov5_nms_for_instance_segmentation(
+            output=instances.permute(0, 2, 1),
+            conf_thresh=conf_thresh,
+            iou_thresh=iou_thresh,
+            max_detections=max_detections,
+            class_agnostic=class_agnostic,
+        )
+        final_results = []
+        for image_bboxes, image_protos, image_meta in zip(
+            nms_results, protos, pre_processing_meta
+        ):
+            pre_processed_masks = preprocess_segmentation_masks(
+                protos=image_protos,
+                masks_in=image_bboxes[:, 6:],
+            )
+            cropped_masks = crop_masks_to_boxes(
+                image_bboxes[:, :4], pre_processed_masks
+            )
+            padding = (
+                image_meta.pad_left,
+                image_meta.pad_top,
+                image_meta.pad_right,
+                image_meta.pad_bottom,
+            )
+            aligned_boxes, aligned_masks = align_instance_segmentation_results(
+                image_bboxes=image_bboxes,
+                masks=cropped_masks,
+                padding=padding,
+                scale_height=image_meta.scale_height,
+                scale_width=image_meta.scale_width,
+                original_size=image_meta.original_size,
+                size_after_pre_processing=image_meta.size_after_pre_processing,
+                inference_size=image_meta.inference_size,
+                static_crop_offset=image_meta.static_crop_offset,
+            )
+            final_results.append(
+                InstanceDetections(
+                    xyxy=aligned_boxes[:, :4].round().int(),
+                    class_id=aligned_boxes[:, 5].int(),
+                    confidence=aligned_boxes[:, 4],
+                    mask=aligned_masks,
+                )
+            )
+        return final_results