sleap-nn 0.1.0a2__py3-none-any.whl → 0.1.0a4__py3-none-any.whl

sleap_nn/export/predictors/__init__.py

@@ -0,0 +1,63 @@
+"""Predictors for exported models."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from sleap_nn.export.predictors.base import ExportPredictor
+from sleap_nn.export.predictors.onnx import ONNXPredictor
+from sleap_nn.export.predictors.tensorrt import TensorRTPredictor
+
+
+def detect_runtime(model_path: str | Path) -> str:
+    """Auto-detect runtime from file extension or folder contents."""
+    model_path = Path(model_path)
+    if model_path.is_dir():
+        onnx_path = model_path / "exported" / "model.onnx"
+        trt_path = model_path / "exported" / "model.trt"
+        if trt_path.exists():
+            return "tensorrt"
+        if onnx_path.exists():
+            return "onnx"
+        raise ValueError(f"No exported model found in {model_path}")
+
+    ext = model_path.suffix.lower()
+    if ext == ".onnx":
+        return "onnx"
+    if ext in (".trt", ".engine"):
+        return "tensorrt"
+
+    raise ValueError(f"Unknown model format: {ext}")
+
+
+def load_exported_model(
+    model_path: str | Path,
+    runtime: str = "auto",
+    device: str = "auto",
+    **kwargs,
+) -> ExportPredictor:
+    """Load an exported model and return a predictor instance."""
+    if runtime == "auto":
+        runtime = detect_runtime(model_path)
+
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if runtime == "onnx":
+        return ONNXPredictor(str(model_path), device=device, **kwargs)
+    if runtime == "tensorrt":
+        return TensorRTPredictor(str(model_path), device=device, **kwargs)
+
+    raise ValueError(f"Unknown runtime: {runtime}")
+
+
+__all__ = [
+    "ExportPredictor",
+    "ONNXPredictor",
+    "TensorRTPredictor",
+    "detect_runtime",
+    "load_exported_model",
+]

sleap_nn/export/predictors/base.py

@@ -0,0 +1,22 @@
+"""Predictor base class for exported models."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Dict
+
+import numpy as np
+
+
+class ExportPredictor(ABC):
+    """Base interface for exported model inference."""
+
+    @abstractmethod
+    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
+        """Run inference on a batch of images."""
+
+    @abstractmethod
+    def benchmark(
+        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
+    ) -> Dict[str, float]:
+        """Benchmark inference latency and throughput."""

sleap_nn/export/predictors/onnx.py

@@ -0,0 +1,154 @@
+"""ONNX Runtime predictor."""
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, Optional
+import time
+
+import numpy as np
+
+from sleap_nn.export.predictors.base import ExportPredictor
+
+
+class ONNXPredictor(ExportPredictor):
+    """ONNX Runtime inference with provider selection."""
+
+    def __init__(
+        self,
+        model_path: str,
+        device: str = "auto",
+        providers: Optional[Iterable[str]] = None,
+    ) -> None:
+        """Initialize ONNX predictor with execution providers.
+
+        Args:
+            model_path: Path to the ONNX model file.
+            device: Device for inference ("auto", "cpu", or "cuda").
+            providers: ONNX Runtime execution providers. Auto-selected if None.
+        """
+        try:
+            import onnxruntime as ort
+        except ImportError as exc:
+            raise ImportError(
+                "onnxruntime is required for ONNXPredictor. Install with "
+                "`pip install onnxruntime` or `onnxruntime-gpu`."
+            ) from exc
+
+        # Preload CUDA/cuDNN libraries from pip-installed nvidia packages
+        # This is required for onnxruntime-gpu to find the CUDA libraries
+        if hasattr(ort, "preload_dlls"):
+            ort.preload_dlls()
+
+        self.ort = ort
+        if providers is None:
+            providers = _select_providers(device, ort.get_available_providers())
+
+        self.session = ort.InferenceSession(model_path, providers=list(providers))
+        input_info = self.session.get_inputs()[0]
+        self.input_name = input_info.name
+        self.input_type = input_info.type
+        self.input_dtype = _onnx_type_to_numpy(self.input_type)
+        self.output_names = [out.name for out in self.session.get_outputs()]
+
+    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
+        """Run inference on a batch of images."""
+        image = _as_numpy(image, expected_dtype=self.input_dtype)
+        outputs = self.session.run(None, {self.input_name: image})
+        return dict(zip(self.output_names, outputs))
+
+    def benchmark(
+        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
+    ) -> Dict[str, float]:
+        """Benchmark inference latency and throughput."""
+        image = _as_numpy(image, expected_dtype=self.input_dtype)
+        for _ in range(n_warmup):
+            self.session.run(None, {self.input_name: image})
+
+        times = []
+        for _ in range(n_runs):
+            start = time.perf_counter()
+            self.session.run(None, {self.input_name: image})
+            times.append(time.perf_counter() - start)
+
+        times_ms = np.array(times) * 1000.0
+        mean_ms = float(times_ms.mean())
+        p50_ms = float(np.percentile(times_ms, 50))
+        p95_ms = float(np.percentile(times_ms, 95))
+        fps = float(1000.0 / mean_ms) if mean_ms > 0 else 0.0
+
+        return {
+            "latency_ms_mean": mean_ms,
+            "latency_ms_p50": p50_ms,
+            "latency_ms_p95": p95_ms,
+            "fps": fps,
+        }
+
+
+def _select_providers(device: str, available: Iterable[str]) -> Iterable[str]:
+    device = device.lower()
+    available = list(available)
+
+    if device in ("cpu", "host"):
+        return ["CPUExecutionProvider"]
+
+    if device.startswith("cuda") or device == "auto":
+        # Note: We don't include TensorrtExecutionProvider here because:
+        # 1. We have a dedicated TensorRTPredictor for native TRT inference
+        # 2. ORT's TensorRT provider requires TRT libs in LD_LIBRARY_PATH
+        preferred = [
+            "CUDAExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+        return [p for p in preferred if p in available] or available
+
+    if device in ("directml", "dml"):
+        preferred = ["DmlExecutionProvider", "CPUExecutionProvider"]
+        return [p for p in preferred if p in available] or available
+
+    return available
+
+
+def _onnx_type_to_numpy(type_str: str | None) -> np.dtype | None:
+    if not type_str:
+        return None
+    if type_str.startswith("tensor(") and type_str.endswith(")"):
+        key = type_str[len("tensor(") : -1]
+    else:
+        return None
+
+    mapping = {
+        "float": np.float32,
+        "float16": np.float16,
+        "double": np.float64,
+        "uint8": np.uint8,
+        "int8": np.int8,
+        "uint16": np.uint16,
+        "int16": np.int16,
+        "uint32": np.uint32,
+        "int32": np.int32,
+        "uint64": np.uint64,
+        "int64": np.int64,
+    }
+    return mapping.get(key)
+
+
+def _as_numpy(image: np.ndarray, expected_dtype: np.dtype | None = None) -> np.ndarray:
+    if isinstance(image, np.ndarray):
+        data = image
+    else:
+        try:
+            import torch
+
+            if isinstance(image, torch.Tensor):
+                data = image.detach().cpu().numpy()
+            else:
+                data = np.asarray(image)
+        except ImportError:
+            data = np.asarray(image)
+
+    if expected_dtype is not None:
+        if data.dtype != expected_dtype:
+            data = data.astype(expected_dtype)
+    elif data.dtype != np.float32:
+        data = data.astype(np.float32)
+    return data

sleap_nn/export/predictors/tensorrt.py

@@ -0,0 +1,312 @@
+"""TensorRT predictor for exported models."""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+from sleap_nn.export.predictors.base import ExportPredictor
+
+
+class TensorRTPredictor(ExportPredictor):
+    """TensorRT inference for exported models.
+
+    This predictor loads a native TensorRT engine file (.trt) and provides
+    inference capabilities using CUDA for high-throughput predictions.
+
+    Args:
+        engine_path: Path to the TensorRT engine file (.trt).
+        device: Device to run inference on (only "cuda" supported for TRT).
+
+    Example:
+        >>> predictor = TensorRTPredictor("model.trt")
+        >>> outputs = predictor.predict(images)  # uint8 [B, C, H, W]
+    """
+
+    def __init__(
+        self,
+        engine_path: str | Path,
+        device: str = "cuda",
+    ) -> None:
+        """Initialize TensorRT predictor with a serialized engine.
+
+        Args:
+            engine_path: Path to the TensorRT engine file.
+            device: Device for inference ("cuda" or "auto"). TensorRT requires CUDA.
+        """
+        import tensorrt as trt
+
+        if device not in ("cuda", "auto"):
+            raise ValueError(f"TensorRT only supports CUDA devices, got: {device}")
+
+        self.engine_path = Path(engine_path)
+        if not self.engine_path.exists():
+            raise FileNotFoundError(f"TensorRT engine not found: {engine_path}")
+
+        self.logger = trt.Logger(trt.Logger.WARNING)
+
+        # Load engine
+        with open(self.engine_path, "rb") as f:
+            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
+
+        if self.engine is None:
+            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")
+
+        # Create execution context
+        self.context = self.engine.create_execution_context()
+
+        # Get input/output info
+        self.input_names: List[str] = []
+        self.output_names: List[str] = []
+        self.input_shapes: Dict[str, tuple] = {}
+        self.output_shapes: Dict[str, tuple] = {}
+
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            shape = tuple(self.engine.get_tensor_shape(name))
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                self.input_names.append(name)
+                self.input_shapes[name] = shape
+            else:
+                self.output_names.append(name)
+                self.output_shapes[name] = shape
+
+        self.device = torch.device("cuda")
+
+    def predict(self, image: np.ndarray) -> Dict[str, np.ndarray]:
+        """Run TensorRT inference.
+
+        Args:
+            image: Input image(s) as numpy array [B, C, H, W] with uint8 dtype.
+
+        Returns:
+            Dict mapping output names to numpy arrays.
+        """
+        import tensorrt as trt
+
+        # Convert input to torch tensor on GPU
+        input_tensor = torch.from_numpy(image).to(self.device)
+
+        # Check if engine expects uint8 or float32 and convert if needed
+        input_name = self.input_names[0]
+        expected_dtype = self.engine.get_tensor_dtype(input_name)
+        if expected_dtype == trt.DataType.UINT8:
+            # Engine expects uint8 - keep as uint8
+            if input_tensor.dtype != torch.uint8:
+                input_tensor = input_tensor.to(torch.uint8)
+        else:
+            # Engine expects float - convert uint8 to float32
+            if input_tensor.dtype == torch.uint8:
+                input_tensor = input_tensor.to(torch.float32)
+
+        # Ensure contiguous memory
+        input_tensor = input_tensor.contiguous()
+
+        # Set input shape for dynamic dimensions
+        input_name = self.input_names[0]
+        self.context.set_input_shape(input_name, tuple(input_tensor.shape))
+
+        # Allocate output tensors
+        outputs: Dict[str, torch.Tensor] = {}
+        bindings: Dict[str, int] = {}
+
+        # Set input binding
+        bindings[input_name] = input_tensor.data_ptr()
+
+        # Allocate outputs
+        for name in self.output_names:
+            shape = self.context.get_tensor_shape(name)
+            dtype = self._trt_dtype_to_torch(self.engine.get_tensor_dtype(name))
+            outputs[name] = torch.empty(tuple(shape), dtype=dtype, device=self.device)
+            bindings[name] = outputs[name].data_ptr()
+
+        # Set tensor addresses
+        for name, ptr in bindings.items():
+            self.context.set_tensor_address(name, ptr)
+
+        # Run inference
+        stream = torch.cuda.current_stream().cuda_stream
+        success = self.context.execute_async_v3(stream)
+        torch.cuda.current_stream().synchronize()
+
+        if not success:
+            raise RuntimeError("TensorRT inference failed")
+
+        # Convert outputs to numpy
+        return {name: tensor.cpu().numpy() for name, tensor in outputs.items()}
+
+    def benchmark(
+        self, image: np.ndarray, n_warmup: int = 50, n_runs: int = 200
+    ) -> Dict[str, float]:
+        """Benchmark inference performance.
+
+        Args:
+            image: Input image(s) as numpy array [B, C, H, W].
+            n_warmup: Number of warmup runs (not timed).
+            n_runs: Number of timed runs.
+
+        Returns:
+            Dict with timing statistics:
+                - mean_ms: Mean inference time in milliseconds
+                - std_ms: Standard deviation of inference time
+                - min_ms: Minimum inference time
+                - max_ms: Maximum inference time
+                - fps: Frames per second (based on mean time and batch size)
+        """
+        batch_size = image.shape[0]
+
+        # Warmup
+        for _ in range(n_warmup):
+            _ = self.predict(image)
+
+        # Timed runs
+        times = []
+        for _ in range(n_runs):
+            start = time.perf_counter()
+            _ = self.predict(image)
+            times.append((time.perf_counter() - start) * 1000)
+
+        times_arr = np.array(times)
+        mean_ms = float(np.mean(times_arr))
+
+        return {
+            "mean_ms": mean_ms,
+            "std_ms": float(np.std(times_arr)),
+            "min_ms": float(np.min(times_arr)),
+            "max_ms": float(np.max(times_arr)),
+            "fps": (batch_size * 1000) / mean_ms if mean_ms > 0 else 0.0,
+        }
+
+    @staticmethod
+    def _trt_dtype_to_torch(trt_dtype):
+        """Convert TensorRT dtype to PyTorch dtype."""
+        import tensorrt as trt
+
+        mapping = {
+            trt.DataType.FLOAT: torch.float32,
+            trt.DataType.HALF: torch.float16,
+            trt.DataType.INT32: torch.int32,
+            trt.DataType.INT8: torch.int8,
+            trt.DataType.BOOL: torch.bool,
+        }
+        return mapping.get(trt_dtype, torch.float32)
+
+
+class TensorRTEngine:
+    """Low-level wrapper for native TensorRT engine files (.trt).
+
+    Provides a callable interface similar to PyTorch models, returning
+    torch.Tensor outputs directly.
+
+    Args:
+        engine_path: Path to the TensorRT engine file (.trt).
+
+    Example:
+        >>> engine = TensorRTEngine("model.trt")
+        >>> input_tensor = torch.randn(1, 1, 512, 512, device="cuda")
+        >>> outputs = engine(input_tensor)  # Dict[str, torch.Tensor]
+    """
+
+    def __init__(self, engine_path: str | Path) -> None:
+        """Initialize TensorRT engine wrapper.
+
+        Args:
+            engine_path: Path to the serialized TensorRT engine file.
+        """
+        import tensorrt as trt
+
+        self.engine_path = Path(engine_path)
+        self.logger = trt.Logger(trt.Logger.WARNING)
+
+        # Load engine
+        with open(engine_path, "rb") as f:
+            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
+
+        if self.engine is None:
+            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")
+
+        # Create execution context
+        self.context = self.engine.create_execution_context()
+
+        # Get input/output info
+        self.input_names: List[str] = []
+        self.output_names: List[str] = []
+        self.input_shapes: Dict[str, tuple] = {}
+        self.output_shapes: Dict[str, tuple] = {}
+
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            shape = tuple(self.engine.get_tensor_shape(name))
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                self.input_names.append(name)
+                self.input_shapes[name] = shape
+            else:
+                self.output_names.append(name)
+                self.output_shapes[name] = shape
+
+    def __call__(self, *args, **kwargs) -> Dict[str, torch.Tensor]:
+        """Run inference.
+
+        Args:
+            *args: Input tensors (positional) or
+            **kwargs: Input tensors by name
+
+        Returns:
+            Dict of output tensors on the same device as input.
+        """
+        import tensorrt as trt
+
+        # Handle inputs
+        if args:
+            inputs = {self.input_names[i]: arg for i, arg in enumerate(args)}
+        else:
+            inputs = kwargs
+
+        # Set input shapes and allocate outputs
+        outputs: Dict[str, torch.Tensor] = {}
+        bindings: Dict[str, int] = {}
+
+        for name in self.input_names:
+            tensor = inputs[name]
+            if not isinstance(tensor, torch.Tensor):
+                raise ValueError(f"Input {name} must be a torch.Tensor")
+            # Set actual shape for dynamic dimensions
+            self.context.set_input_shape(name, tuple(tensor.shape))
+            bindings[name] = tensor.contiguous().data_ptr()
+
+        # Allocate output tensors
+        device = next(iter(inputs.values())).device
+        for name in self.output_names:
+            shape = self.context.get_tensor_shape(name)
+            dtype = self._trt_dtype_to_torch(self.engine.get_tensor_dtype(name))
+            outputs[name] = torch.empty(tuple(shape), dtype=dtype, device=device)
+            bindings[name] = outputs[name].data_ptr()
+
+        # Set tensor addresses
+        for name, ptr in bindings.items():
+            self.context.set_tensor_address(name, ptr)
+
+        # Run inference
+        stream = torch.cuda.current_stream().cuda_stream
+        self.context.execute_async_v3(stream)
+        torch.cuda.current_stream().synchronize()
+
+        return outputs
+
+    @staticmethod
+    def _trt_dtype_to_torch(trt_dtype):
+        """Convert TensorRT dtype to PyTorch dtype."""
+        import tensorrt as trt
+
+        mapping = {
+            trt.DataType.FLOAT: torch.float32,
+            trt.DataType.HALF: torch.float16,
+            trt.DataType.INT32: torch.int32,
+            trt.DataType.INT8: torch.int8,
+            trt.DataType.BOOL: torch.bool,
+        }
+        return mapping.get(trt_dtype, torch.float32)