onnx¶

def load(self) -> None:
    """Load ONNX model and prepare for inference.

    Performs:
        - Configure session options
        - Load model from disk
        - Extract input/output names
        - Warmup inference

    Raises:
        FileNotFoundError: If model file does not exist.
        RuntimeError: If ONNX Runtime fails to load model.
    """
    logger.info(f"Loading ONNX model from {self.model_path}")

    # Session options
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

    # Load model
    self.session = ort.InferenceSession(
        str(self.model_path),
        sess_options=sess_options,
        providers=self.providers,
    )

    # Get input/output names
    self.input_name = self.session.get_inputs()[0].name
    self.output_names = [output.name for output in self.session.get_outputs()]

    logger.info(
        f"Model loaded:  input={self.input_name}, "
        f"outputs={self.output_names}, "
        f"providers={self.session.get_providers()}"
    )

    # Warmup
    self._warmup()

def infer(self, input: np.ndarray) -> t.Any:
    """Run inference on a single input.

    Automatically handles:
        - Converting to correct dtype

    Args:
        input: Input numpy array.

    Returns:
        Output array or tuple of arrays (if multi-output model).

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        with runtime:
            input = np.random.randn(1, 3, 224, 224).astype(
                np.float32
            )
            output = runtime.infer(input)
        ```
    """
    if self.session is None or self.input_name is None or self.output_names is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Prepare input (reuse mixin)
    input = self._prepare_onnx_input(input)

    # Run inference
    outputs = self.session.run(self.output_names, {self.input_name: input})

    # Parse output (reuse mixin)
    return self._parse_onnx_output(list(outputs))

def infer_batch(self, inputs: list[np.ndarray]) -> list[t.Any]:
    """Run inference on a batch of inputs.

    Concatenates inputs into a single batch array for efficient
    processing, then splits the output back into individual results.

    Args:
        inputs: List of input arrays. Each should have shape (1, C, H, W).

    Returns:
        List of outputs, one per input. Each maintains batch dimension.

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        with runtime:
            batch = [
                np.random.randn(1, 3, 224, 224).astype(np.float32),
                np.random.randn(1, 3, 224, 224).astype(np.float32),
            ]
            outputs = runtime.infer_batch(batch)
        ```
    """
    if self.session is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Concatenate inputs
    batch = np.concatenate(inputs, axis=0)

    # Prepare input (reuse mixin)
    batch = self._prepare_onnx_input(batch)

    # Run batch inference
    batch_output = self.infer(batch)

    # Split outputs (reuse mixin)
    return self._split_onnx_batch_output(batch_output, len(inputs))

def unload(self) -> None:
    """Unload model and free resources.

    Performs:
        - Release session from memory

    Safe to call multiple times.
    """
    logger.info("Unloading model")
    self.session = None
    self.input_name = None
    self.output_names = None
    logger.info("Model unloaded")