inference-models 0.18.3__py3-none-any.whl

inference_models/models/yolov8/yolov8_object_detection_trt.py

@@ -0,0 +1,231 @@
+import threading
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import Detections, ObjectDetectionModel
+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 (
+    post_process_nms_fused_model_output,
+    rescale_detections,
+    run_nms_for_object_detection,
+)
+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,
+)
+
+try:
+    import tensorrt as trt
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import YOLOv8 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 YOLOv8ForObjectDetectionTRT(
+    ObjectDetectionModel[torch.Tensor, PreProcessingMetadata, torch.Tensor]
+):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        engine_host_code_allowed: bool = False,
+        **kwargs,
+    ) -> "YOLOv8ForObjectDetectionTRT":
+        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,
+            },
+        )
+        if inference_config.post_processing.type != "nms":
+            raise CorruptedModelPackageError(
+                message="Expected NMS to be the post-processing",
+                help_url="https://todo",
+            )
+        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) != 1:
+            raise CorruptedModelPackageError(
+                message=f"Implementation assume single model output, found: {len(outputs)}.",
+                help_url="https://todo",
+            )
+        return cls(
+            engine=engine,
+            input_name=inputs[0],
+            output_name=outputs[0],
+            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_name: 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_name]
+        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._lock = threading.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) -> torch.Tensor:
+        with self._lock:
+            with use_cuda_context(context=self._cuda_context):
+                return 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,
+                )[0]
+
+    def post_process(
+        self,
+        model_results: 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[Detections]:
+        if self._inference_config.post_processing.fused:
+            nms_results = post_process_nms_fused_model_output(
+                output=model_results, conf_thresh=conf_thresh
+            )
+        else:
+            nms_results = run_nms_for_object_detection(
+                output=model_results,
+                conf_thresh=conf_thresh,
+                iou_thresh=iou_thresh,
+                max_detections=max_detections,
+                class_agnostic=class_agnostic,
+            )
+        rescaled_results = rescale_detections(
+            detections=nms_results,
+            images_metadata=pre_processing_meta,
+        )
+        results = []
+        for result in rescaled_results:
+            results.append(
+                Detections(
+                    xyxy=result[:, :4].round().int(),
+                    class_id=result[:, 5].int(),
+                    confidence=result[:, 4],
+                )
+            )
+        return results

inference_models/models/yolov9/yolov9_onnx.py

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

inference_models/models/yolov9/yolov9_torch_script.py

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

inference_models/models/yolov9/yolov9_trt.py

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

inference_models/runtime_introspection/core.py

@@ -0,0 +1,410 @@
+import os
+import platform
+import re
+import subprocess
+from dataclasses import dataclass
+from functools import cache
+from typing import List, Optional, Set, Tuple
+
+import torch
+from packaging.version import InvalidVersion, Version
+
+from inference_models.configuration import L4T_VERSION, RUNNING_ON_JETSON
+from inference_models.errors import JetsonTypeResolutionError, RuntimeIntrospectionError
+from inference_models.logger import LOGGER
+from inference_models.utils.environment import str2bool
+
+JETSON_DEVICES_TABLE = [
+    "NVIDIA Jetson Orin Nano",
+    "NVIDIA Jetson Orin NX",
+    "NVIDIA Jetson AGX Orin",
+    "NVIDIA Jetson IGX Orin",
+    "NVIDIA Jetson Xavier NX",
+    "NVIDIA Jetson AGX Xavier Industrial",
+    "NVIDIA Jetson AGX Xavier",
+    "NVIDIA Jetson Nano",
+    "NVIDIA Jetson TX2",
+]
+
+
+@dataclass(frozen=True)
+class RuntimeXRayResult:
+    gpu_available: bool
+    gpu_devices: List[str]
+    gpu_devices_cc: List[Version]
+    driver_version: Optional[Version]
+    cuda_version: Optional[Version]
+    trt_version: Optional[Version]
+    jetson_type: Optional[str]
+    l4t_version: Optional[Version]
+    os_version: Optional[str]
+    torch_available: bool
+    torch_version: Optional[Version]
+    torchvision_version: Optional[Version]
+    onnxruntime_version: Optional[Version]
+    available_onnx_execution_providers: Optional[Set[str]]
+    hf_transformers_available: bool
+    ultralytics_available: bool
+    trt_python_package_available: bool
+    mediapipe_available: bool
+
+    def __str__(self) -> str:
+        gpu_devices_str = ", ".join(self.gpu_devices)
+        return (
+            f"RuntimeXRayResult(gpu_available={self.gpu_available}, gpu_devices=[{gpu_devices_str}], "
+            f"gpu_devices_cc={self.gpu_devices_cc}, gpu_driver={self.driver_version}, "
+            f"cuda_version={self.cuda_version}, trt_version={self.trt_version}, "
+            f"jetson_type={self.jetson_type}, l4t_version={self.l4t_version}, os_version={self.os_version}, "
+            f"torch_available={self.torch_available}, onnxruntime_version={self.onnxruntime_version}, "
+            f"available_onnx_execution_providers={self.available_onnx_execution_providers}, hf_transformers_available={self.hf_transformers_available}, "
+            f"ultralytics_available={self.ultralytics_available}, "
+            f"trt_python_package_available={self.trt_python_package_available}, torch_version={self.torch_version}, "
+            f"torchvision_version={self.torchvision_version}, mediapipe_available={self.mediapipe_available})"
+        )
+
+
+@cache
+def x_ray_runtime_environment() -> RuntimeXRayResult:
+    trt_version = get_trt_version()
+    cuda_version = get_cuda_version()
+    jetson_type, l4t_version, os_version, driver_version = None, None, None, None
+    if is_running_on_jetson():
+        jetson_type = get_jetson_type()
+        l4t_version = get_l4t_version()
+    else:
+        os_version = get_os_version()
+    driver_version = get_driver_version()
+    gpu_devices = get_available_gpu_devices()
+    ensure_jetson_l4t_declared_for_jetson_hardware(
+        gpu_devices=gpu_devices,
+        l4t_version=l4t_version,
+    )
+    gpu_devices_cc = get_available_gpu_devices_cc()
+    torch_available = is_torch_available()
+    torch_version = get_torch_version()
+    torchvision_version = get_torchvision_version()
+    onnx_info = get_onnxruntime_info()
+    if onnx_info:
+        onnxruntime_version, available_onnx_execution_providers = onnx_info
+    else:
+        onnxruntime_version, available_onnx_execution_providers = None, None
+    hf_transformers_available = is_hf_transformers_available()
+    ultralytics_available = is_ultralytics_available()
+    trt_python_package_available = is_trt_python_package_available()
+    mediapipe_available = is_mediapipe_available()
+    return RuntimeXRayResult(
+        gpu_available=len(gpu_devices) > 0,
+        gpu_devices=gpu_devices,
+        gpu_devices_cc=gpu_devices_cc,
+        driver_version=driver_version,
+        cuda_version=cuda_version,
+        trt_version=trt_version,
+        jetson_type=jetson_type,
+        l4t_version=l4t_version,
+        os_version=os_version,
+        torch_available=torch_available,
+        torch_version=torch_version,
+        torchvision_version=torchvision_version,
+        onnxruntime_version=onnxruntime_version,
+        available_onnx_execution_providers=available_onnx_execution_providers,
+        hf_transformers_available=hf_transformers_available,
+        ultralytics_available=ultralytics_available,
+        trt_python_package_available=trt_python_package_available,
+        mediapipe_available=mediapipe_available,
+    )
+
+
+@cache
+def is_running_on_jetson() -> bool:
+    if RUNNING_ON_JETSON is not None:
+        return str2bool(value=RUNNING_ON_JETSON)
+    return get_jetson_type() is not None
+
+
+@cache
+def get_available_gpu_devices() -> List[str]:
+    num_devices = torch.cuda.device_count()
+    result = []
+    for i in range(num_devices):
+        device_name = torch.cuda.get_device_name(i)
+        result.append(device_name.replace(" ", "-").lower())
+    return result
+
+
+@cache
+def get_available_gpu_devices_cc() -> List[Version]:
+    num_devices = torch.cuda.device_count()
+    result = []
+    for i in range(num_devices):
+        device_cc_raw = torch.cuda.get_device_capability(i)
+        result.append(Version(f"{device_cc_raw[0]}.{device_cc_raw[1]}"))
+    return result
+
+
+@cache
+def get_cuda_version() -> Optional[Version]:
+    try:
+        result = subprocess.run(
+            "dpkg -l | grep cuda-cudart", shell=True, capture_output=True, text=True
+        )
+        if result.returncode != 0:
+            return None
+        result_chunks = result.stdout.strip().split(os.linesep)[0].split()
+        return Version(result_chunks[2])
+    except Exception:
+        return None
+
+
+@cache
+def get_trt_version() -> Optional[Version]:
+    trt_version = get_trt_version_from_libnvinfer()
+    if trt_version is not None:
+        return trt_version
+    try:
+        import tensorrt as trt
+
+        return Version(trt.__version__)
+    except Exception:
+        return None
+
+
+@cache
+def get_trt_version_from_libnvinfer() -> Optional[Version]:
+    try:
+        result = subprocess.run(
+            "dpkg -l | grep libnvinfer-bin", shell=True, capture_output=True, text=True
+        )
+        if result.returncode != 0:
+            return None
+        result_chunks = result.stdout.strip().split()
+        libraries_versions = result_chunks[2].split("+cuda")
+        return Version(libraries_versions[0])
+    except Exception:
+        return None
+
+
+@cache
+def get_jetson_type() -> Optional[str]:
+    declared_json_module = os.getenv("JETSON_MODULE")
+    if declared_json_module:
+        return resolve_jetson_type(jetson_module_name=declared_json_module)
+    jetson_type_from_hardware_inspection = get_jetson_type_from_hardware_inspection()
+    if jetson_type_from_hardware_inspection:
+        return jetson_type_from_hardware_inspection
+    return get_jetson_type_from_device_tree()
+
+
+def get_jetson_type_from_hardware_inspection() -> Optional[str]:
+    try:
+        result = subprocess.run(
+            "lshw | grep 'product: NVIDIA Jetson'",
+            shell=True,
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            return None
+        for result_line in result.stdout.strip().split("\n"):
+            start_idx = result_line.find("NVIDIA")
+            end_idx = result_line.find("HDA")
+            if start_idx < 0 or end_idx < 0:
+                continue
+            jetson_type = result_line[start_idx:end_idx].strip()
+            return resolve_jetson_type(jetson_module_name=jetson_type)
+        return None
+    except Exception:
+        return None
+
+
+def get_jetson_type_from_device_tree() -> Optional[str]:
+    try:
+        with open("/proc/device-tree/model") as f:
+            model_headline = f.read().strip().split("\n")[0]
+            return resolve_jetson_type(jetson_module_name=model_headline)
+    except Exception:
+        return None
+
+
+def resolve_jetson_type(jetson_module_name: str) -> str:
+    for jetson_device in JETSON_DEVICES_TABLE:
+        if jetson_module_name.startswith(jetson_device):
+            return jetson_device.replace(" ", "-").lower()
+    raise JetsonTypeResolutionError(
+        message=f"Could not resolve jetson type. Value found in environment: {jetson_module_name}",
+        help_url="https://todo",
+    )
+
+
+@cache
+def get_l4t_version() -> Optional[Version]:
+    if L4T_VERSION:
+        return Version(L4T_VERSION)
+    return get_l4t_version_from_tegra_release()
+
+
+def get_l4t_version_from_tegra_release() -> Optional[Version]:
+    try:
+        with open("/etc/nv_tegra_release") as f:
+            file_header = f.readline()
+            match = re.search(r"R(\d+).*REVISION:\s*([\d.]+)", file_header)
+            if match:
+                major = match.group(1)
+                minor_patch = match.group(2)
+                return Version(f"{major}.{minor_patch}")
+            return None
+    except Exception:
+        return None
+
+
+@cache
+def get_os_version() -> Optional[str]:
+    system = platform.system().lower()
+    if system == "linux":
+        return get_linux_os_version()
+    elif system == "darwin":
+        return platform.platform().lower()
+    elif system == "windows":
+        return platform.platform().lower()
+    return None
+
+
+def get_linux_os_version() -> Optional[str]:
+    os_version_from_os_release = get_linux_os_from_os_release()
+    if os_version_from_os_release:
+        return os_version_from_os_release
+    return platform.platform().lower()
+
+
+def get_linux_os_from_os_release() -> Optional[str]:
+    try:
+        with open("/etc/os-release") as f:
+            data = dict(line.strip().split("=", 1) for line in f if "=" in line)
+        distro_name = data["NAME"].strip('"')
+        version_id = data["VERSION_ID"].strip('"')
+        return f"{distro_name}-{version_id}".lower()
+    except Exception:
+        return None
+
+
+@cache
+def get_driver_version() -> Optional[Version]:
+    try:
+        with open("/proc/driver/nvidia/version") as f:
+            head_line = f.readline()
+        match = re.search(r"\b(\d+(\.\d+){1,2})\b", head_line)
+        if match:
+            return Version(match.group(1))
+    except Exception:
+        return None
+
+
+@cache
+def is_trt_python_package_available() -> bool:
+    try:
+        import tensorrt
+
+        return True
+    except ImportError:
+        return False
+
+
+@cache
+def is_mediapipe_available() -> bool:
+    try:
+        import mediapipe
+
+        return True
+    except ImportError:
+        return False
+
+
+@cache
+def is_torch_available() -> bool:
+    try:
+        import torch
+
+        return True
+    except ImportError:
+        return False
+
+
+@cache
+def get_torch_version() -> Optional[Version]:
+    try:
+        import torch
+
+        version_str = torch.__version__
+        if "+" in version_str:
+            version_str = version_str.split("+")[0]
+        return Version(version_str)
+    except ImportError:
+        return None
+    except InvalidVersion as e:
+        LOGGER.warning(f"Could not parse torch version: {e}")
+        return None
+
+
+@cache
+def get_torchvision_version() -> Optional[Version]:
+    try:
+        import torchvision
+
+        version_str = torchvision.__version__
+        if "+" in version_str:
+            version_str = version_str.split("+")[0]
+        return Version(version_str)
+    except ImportError:
+        return None
+    except InvalidVersion as e:
+        LOGGER.warning(f"Could not parse torch version: {e}")
+        return None
+
+
+@cache
+def get_onnxruntime_info() -> Optional[Tuple[Version, Set[str]]]:
+    try:
+        import onnxruntime
+
+        available_providers = onnxruntime.get_available_providers()
+        return Version(onnxruntime.__version__), available_providers
+    except ImportError:
+        return None
+
+
+@cache
+def is_hf_transformers_available() -> bool:
+    try:
+        import transformers
+
+        return True
+    except ImportError:
+        return False
+
+
+@cache
+def is_ultralytics_available() -> bool:
+    try:
+        import ultralytics
+
+        return True
+    except ImportError:
+        return False
+
+
+def ensure_jetson_l4t_declared_for_jetson_hardware(
+    gpu_devices: List[str],
+    l4t_version: Optional[Version],
+) -> None:
+    # this function is needed to increase likelihood that in cases when
+    # runtime x-ray rely on env variables (like when on Jetson in container we
+    # cannot inspect hardware in a reliable way) - the xray is correct
+    if any(device == "orin" for device in gpu_devices) and l4t_version is None:
+        raise RuntimeIntrospectionError(
+            message="Inconsistency in environment setup detected - GPU device name indicate that you run the code"
+            "on Jetson device, but L4T version cannot be established, which is important for `inference-models` "
+            "to work correctly. Please set environment variable `RUNNING_ON_JETSON=True`, as well as "
+            "`L4T_VERSION=<your-l4t-version>`. Optionally you may also set `JETSON_MODULE` variable "
+            f"to one of the values: {JETSON_DEVICES_TABLE}. If that does not solve your issue - contact "
+            f"Roboflow immediately.",
+            help_url="https://todo",
+        )