trtutils.image.onnx_models module

ONNX model implementations.

:attribute:`IMAGE_PREPROC_BASE`

The path to the ONNX model for image preprocessing.

Type:

Path

:attribute:`IMAGE_PREPROC_IMAGENET`

The path to the ONNX model for ImageNet preprocessing.

Type:

Path

trtutils.image.onnx_models.build_image_preproc(input_shape: tuple[int, int], output_dtype: dtype, batch_size: int = 1, trt_version: str | None = None) → Path[source]

Build a image preproc TensorRT engine.

Parameters:

  • input_shape (tuple[int, int]) – The (width, height) of the image network.

  • output_dtype (np.dtype) – The datatype to return, which the image network will take as input.

  • batch_size (int) – The batch size for the engine. Default is 1.

  • trt_version (str | None) – The version of TensorRT to use. If none, will not be used in the cache name.

Returns:

The path to the compiled engine.

Return type:

Path

trtutils.image.onnx_models.build_image_preproc_imagenet(input_shape: tuple[int, int], output_dtype: dtype, batch_size: int = 1, trt_version: str | None = None) → Path[source]

Build an ImageNet preprocessing TensorRT engine.

Parameters:

  • input_shape (tuple[int, int]) – The (width, height) of the image network.

  • output_dtype (np.dtype) – The datatype to return, which the image network will take as input.

  • batch_size (int) – The batch size for the engine. Default is 1.

  • trt_version (str | None) – The version of TensorRT to use. If none, will not be used in the cache name.

Returns:

The path to the compiled engine.

Return type:

Path