trtutils.impls.yolo package¶

Module contents¶

Implementations of YOLO object detectors for TRTModel.

Classes¶

CPUPreprocessor: Preprocess an image for YOLO on CPU.
CUDAPreprocessor: Preprocess an image for YOLO using CUDA.
TRTPreprocessor: Preprocess an image for YOLO using TensorRT.
ParallelYOLO: Multi-threaded YOLO models, useful for multi-accelerator systems.
YOLO: TRTModel implementation for YOLO object detectors
YOLO7: Alias for YOLO with args preset
YOLO8: Alias for YOLO with args preset
YOLO9: Alias for YOLO with args preset
YOLO10: Alias for YOLO with args preset
YOLOX: Alias for YOLO with args preset

Functions¶

get_detections(): Get the detections from a YOLO network.
preprocess(): Preprocess some input for a YOLO network.
postprocess(): Postprocess the output for a YOLO network.

class trtutils.impls.yolo.YOLO(engine_path: Path | str, warmup_iterations: int = 10, input_range: tuple[float, float] = (0.0, 1.0), preprocessor: str = 'trt', resize_method: str = 'letterbox', conf_thres: float = 0.1, nms_iou_thres: float = 0.5, dla_core: int | None = None, *, warmup: bool | None = None, pagelocked_mem: bool | None = None, unified_mem: bool | None = None, extra_nms: bool | None = None, agnostic_nms: bool | None = None, no_warn: bool | None = None, verbose: bool | None = None)[source]¶

Bases: object

Implementation of YOLO object detectors.

property engine: TRTEngine¶: Get the underlying TRTEngine.

property name: str¶: Get the name of the engine.

property input_shape: tuple[int, int]¶: Get the width, height input shape.

property dtype: np.dtype¶: Get the dtype required by the model.

preprocess(image: np.ndarray, resize: str | None = None, method: str | None = None, *, no_copy: bool | None = None, verbose: bool | None = None) → tuple[np.ndarray, tuple[float, float], tuple[float, float]][source]¶

Preprocess the input.

Parameters:

image (np.ndarray) – The image to preprocess
resize (str) – The method to resize the image with. Options are [letterbox, linear]. By default None, which will use the value passed during initialization.
method (str, optional) – The underlying preprocessor to use. Options are ‘cpu’, ‘cuda’, or ‘trt’. By default None, which will use the preprocessor stated in the constructor.
no_copy (bool, optional) – If True and using CUDA, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The preprocessed inputs, rescale ratios, and padding values

Return type:

tuple[list[np.ndarray], tuple[float, float], tuple[float, float]]

postprocess(outputs: list[np.ndarray], ratios: tuple[float, float], padding: tuple[float, float], conf_thres: float | None = None, *, no_copy: bool | None = None, verbose: bool | None = None) → list[np.ndarray][source]¶

Postprocess the outputs.

Parameters:

outputs (list[np.ndarray]) – The outputs to postprocess
ratios (tuple[float, float]) – The rescale ratios used during preprocessing
padding (tuple[float, float]) – The padding values used during preprocessing
conf_thres (float, optional) – The confidence threshold to filter detections by. If not passed, will use value from constructor.
no_copy (bool, optional) – If True, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The postprocessed outputs

Return type:

list[np.ndarray]

Run the YOLO network on input.

Parameters:

image (np.ndarray) – The data to run the YOLO network on.
ratios (tuple[float, float], optional) – The ratios generated during preprocessing.
padding (tuple[float, float], optional) – The padding values used during preprocessing.
conf_thres (float, optional) – Optional confidence threshold to filter detections via during postprocessing.
preprocessed (bool, optional) – Whether or not the inputs have been preprocessed. If None, will preprocess inputs.
postprocess (bool, optional) – Whether or not to postprocess the outputs. If None, will postprocess outputs. If postprocessing will occur and the inputs were passed already preprocessed, then the ratios and padding must be passed for postprocessing.
no_copy (bool, optional) – If True, the outputs will not be copied out from the cuda allocated host memory. Instead, the host memory will be returned directly. This memory WILL BE OVERWRITTEN INPLACE by future inferences. In special case where, preprocessing and postprocessing will occur during run and no_copy was not passed (is None), then no_copy will be used for preprocessing and inference stages.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The outputs of the YOLO network.

Return type:

list[np.ndarray]

Raises:

RuntimeError – If postprocessing is running, but ratios/padding not found

get_random_input() → np.ndarray[source]¶

Generate a random image for the YOLO model.

Returns:: The random image.
Return type:: np.ndarray

mock_run(image: np.ndarray | None = None) → list[np.ndarray][source]¶

Mock an execution of the YOLO model.

Parameters:: image (np.ndarray, optional) – Optional inputs to use for execution. If None, random data will be generated.
Returns:: The outputs of the model
Return type:: list[np.ndarray]

get_detections(outputs: list[np.ndarray], conf_thres: float | None = None, nms_iou_thres: float | None = None, *, extra_nms: bool | None = None, agnostic_nms: bool | None = None, verbose: bool | None = None) → list[tuple[tuple[int, int, int, int], float, int]][source]¶

Get the bounding boxes of the last output or provided output.

Parameters:

outputs (list[np.ndarray]) – The outputs to process.
conf_thres (float, optional) – The confidence threshold with which to retrieve bounding boxes. By default None, which will use value passed during initialization
nms_iou_thres (float) – The IOU threshold to use during the optional/additional NMS operation. By default, None which will use value provided during initialization.
extra_nms (bool, optional) – Whether or not to perform an additional NMS operation. By default None, which will use value provided during initialization.
agnostic_nms (bool, optional) – Whether or not to perform class-agnostic NMS for the optional/additional operation. By default None, which will use value provided during initialization.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The detections

Return type:

list[tuple[tuple[int, int, int, int], float, int]]

end2end(image: np.ndarray, conf_thres: float | None = None, nms_iou_thres: float | None = None, *, extra_nms: bool | None = None, agnostic_nms: bool | None = None, verbose: bool | None = None) → list[tuple[tuple[int, int, int, int], float, int]][source]¶

Perform end to end inference for a YOLO model.

Equivalent to running preprocess, run, postprocess, and get_detections in that order. Makes some memory transfer optimizations under the hood to improve performance.

Parameters:

image (np.ndarray) – The image to perform inference with.
conf_thres (float, optional) – The confidence threshold with which to retrieve bounding boxes. By default None
nms_iou_thres (float) – The IOU threshold to use during the optional/additional NMS operation. By default, None which will use value provided during initialization.
extra_nms (bool, optional) – Whether or not to perform an additional NMS operation. By default None, which will use value provided during initialization.
agnostic_nms (bool, optional) – Whether or not to perform class-agnostic NMS for the optional/additional operation. By default None, which will use value provided during initialization.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The detections where each entry is bbox, conf, class_id

Return type:

list[tuple[tuple[int, int, int, int], float, int]]

class trtutils.impls.yolo.YOLO7(engine_path: Path | str, warmup_iterations: int = 10, input_range: tuple[float, float] = (0, 1), preprocessor: str = 'cuda', resize_method: str = 'letterbox', conf_thres: float = 0.1, nms_iou_thres: float = 0.5, dla_core: int | None = None, *, warmup: bool | None = None, extra_nms: bool | None = None, agnostic_nms: bool | None = None, no_warn: bool | None = None, verbose: bool | None = None)[source]¶

Bases: YOLO

Alias of YOLO with default args for YOLO7.

class trtutils.impls.yolo.YOLO8(engine_path: Path | str, warmup_iterations: int = 10, input_range: tuple[float, float] = (0, 1), preprocessor: str = 'cuda', resize_method: str = 'letterbox', conf_thres: float = 0.1, nms_iou_thres: float = 0.5, dla_core: int | None = None, *, warmup: bool | None = None, extra_nms: bool | None = None, agnostic_nms: bool | None = None, no_warn: bool | None = None, verbose: bool | None = None)[source]¶

Bases: YOLO

Alias of YOLO with default args for YOLO8.

class trtutils.impls.yolo.YOLO9(engine_path: Path | str, warmup_iterations: int = 10, input_range: tuple[float, float] = (0, 1), preprocessor: str = 'cuda', resize_method: str = 'letterbox', conf_thres: float = 0.1, nms_iou_thres: float = 0.5, dla_core: int | None = None, *, warmup: bool | None = None, extra_nms: bool | None = None, agnostic_nms: bool | None = None, no_warn: bool | None = None, verbose: bool | None = None)[source]¶

Bases: YOLO

Alias of YOLO with default args for YOLO9.

class trtutils.impls.yolo.YOLO10(engine_path: Path | str, warmup_iterations: int = 10, input_range: tuple[float, float] = (0, 1), preprocessor: str = 'cuda', resize_method: str = 'letterbox', conf_thres: float = 0.1, nms_iou_thres: float = 0.5, dla_core: int | None = None, *, warmup: bool | None = None, extra_nms: bool | None = None, agnostic_nms: bool | None = None, no_warn: bool | None = None, verbose: bool | None = None)[source]¶

Bases: YOLO

Alias of YOLO with default args for YOLO10.

class trtutils.impls.yolo.YOLOX(engine_path: Path | str, warmup_iterations: int = 10, input_range: tuple[float, float] = (0, 255), preprocessor: str = 'cuda', resize_method: str = 'letterbox', conf_thres: float = 0.1, nms_iou_thres: float = 0.5, dla_core: int | None = None, *, warmup: bool | None = None, extra_nms: bool | None = None, agnostic_nms: bool | None = None, no_warn: bool | None = None, verbose: bool | None = None)[source]¶

Bases: YOLO

Alias of YOLO with default args for YOLOX.

class trtutils.impls.yolo.CPUPreprocessor(output_shape: tuple[int, int], output_range: tuple[float, float], dtype: np.dtype, tag: str | None = None)[source]¶

Bases: object

CPU-based preprocessor for YOLO.

warmup() → None[source]¶: Compatibility function for CPU/CUDA parity.

preprocess(image: np.ndarray, resize: str = 'letterbox', *, verbose: bool | None = None) → tuple[np.ndarray, tuple[float, float], tuple[float, float]][source]¶

Preprocess an image for YOLO.

Parameters:

image (np.ndarray) – The image to preprocess.
resize (str) – The method to resize the image with. By default letterbox, options are [letterbox, linear]
verbose (bool, optional) – Whether or not to output additional information to stdout. If not provided, will default to overall engines verbose setting.

Returns:

The preprocessed image, ratios, and padding used for resizing.

Return type:

tuple[np.ndarray, tuple[float, float], tuple[float, float]]

class trtutils.impls.yolo.CUDAPreprocessor(output_shape: tuple[int, int], output_range: tuple[float, float], dtype: np.dtype, resize: str = 'letterbox', stream: cudart.cudaStream_t | None = None, threads: tuple[int, int, int] | None = None, tag: str | None = None, *, pagelocked_mem: bool | None = None, unified_mem: bool | None = None)[source]¶

Bases: object

CUDA-based preprocessor for YOLO.

warmup() → None[source]¶

Warmup the CUDA preprocessor.

Allocates all CUDA memory and enables future passes to be significantly faster.

preprocess(image: np.ndarray, resize: str | None = None, *, no_copy: bool | None = None, verbose: bool | None = None) → tuple[np.ndarray, tuple[float, float], tuple[float, float]][source]¶

Preprocess an image for YOLO.

Parameters:

image (np.ndarray) – The image to preprocess.
resize (str, optional) – The method to resize the image with. Options are [letterbox, linear], will use method provided in constructor by default.
no_copy (bool, optional) – If True, the outputs will not be copied out from the cuda allocated host memory. Instead, the host memory will be returned directly. This memory WILL BE OVERWRITTEN INPLACE by future preprocessing calls.
verbose (bool, optional) – Whether or not to output additional information to stdout. If not provided, will default to overall engines verbose setting.

Returns:

The preprocessed image, ratios, and padding used for resizing.

Return type:

tuple[np.ndarray, tuple[float, float], tuple[float, float]]

direct_preproc(image: np.ndarray, resize: str | None = None, *, no_warn: bool | None = None, verbose: bool | None = None) → tuple[int, tuple[float, float], tuple[float, float]][source]¶

Preprocess an image for YOLO.

Parameters:

image (np.ndarray) – The image to preprocess.
resize (str) – The method to resize the image with. By default letterbox, options are [letterbox, linear]
no_warn (bool, optional) – If True, do not warn about usage.
verbose (bool, optional) – Whether or not to output additional information to stdout. If not provided, will default to overall engines verbose setting.

Returns:

The GPU pointer to preprocessed data, ratios, and padding used for resizing.

Return type:

tuple[int, tuple[float, float], tuple[float, float]]

Bases: object

A parallel implementation of YOLO.

Allows multiple version of YOLO to be allocated and executed at the same time. Primarily useful for multi-gpu/multi-accelerator systems such as the NVIDIA Jetson series. Since the TensorRT engines are compiled for a specific device, no device specification is needed inside of this class.

property models: list[YOLO]¶

Get the underlying YOLO models.

Returns:: A list of the underlying models.
Return type:: list[YOLO]
Raises:: RuntimeError – If any values are None, not initialized yet.

get_model(modelid: int) → YOLO[source]¶

Get a YOLO model with id.

Parameters:: modelid (int) – The model ID to get. Based on original list passed during init.
Returns:: The YOLO model
Return type:: YOLO
Raises:: RuntimeError – If access is attempted before init is complete

get_model_profiling(modelid: int) → tuple[float, float, float][source]¶

Get the latency of a specific model as profiled in thread.

Parameters:: modelid (int) – The model ID of the profiling to get.
Returns:: The start time, end time, and delta
Return type:: tuple[float, float, float]

get_profiling() → list[tuple[float, float, float]][source]¶

Get all the profiling results for all models.

Returns:: The profiling data
Return type:: list[tuple[float, float, float]]

stop() → None[source]¶: Stop the threads.

preprocess(inputs: list[np.ndarray], resize: str = 'letterbox', method: str | None = None, *, no_copy: bool | None = None, verbose: bool | None = None) → tuple[list[np.ndarray], list[tuple[float, float]], list[tuple[float, float]]][source]¶

Preprocess inputs for inference.

Parameters:

inputs (list[np.ndarray]) – The inputs to preprocess.
resize (str) – The method to resize the image with. By default letterbox, options are [letterbox, linear]
method (str, optional) – The underlying preprocessor to use. Options are ‘cpu’ and ‘cuda’. By default None, which will use the preprocessor stated in the constructor.
no_copy (bool, optional) – If True and using CUDA, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The preprocessed inputs.

Return type:

tuple[list[np.ndarray], list[tuple[float, float]], list[tuple[float, float]]]

Raises:

ValueError – If inputs do not match the number of models

preprocess_model(data: np.ndarray, modelid: int, resize: str = 'letterbox', method: str | None = None, *, no_copy: bool | None = None, verbose: bool | None = None) → tuple[np.ndarray, tuple[float, float], tuple[float, float]][source]¶

Preprocess data for a specific model.

Parameters:

data (np.ndarray) – The data to preprocess.
modelid (int) – The model to preprocess the data for.
resize (str) – The method to resize the image with. By default letterbox, options are [letterbox, linear]
method (str, optional) – The underlying preprocessor to use. Options are ‘cpu’ and ‘cuda’. By default None, which will use the preprocessor stated in the constructor.
no_copy (bool, optional) – If True and using CUDA, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The preprocessed data

Return type:

tuple[np.ndarray, tuple[float, float], tuple[float, float]]

postprocess(outputs: list[list[np.ndarray]], ratios: list[tuple[float, float]], paddings: list[tuple[float, float]], *, no_copy: bool | None = None, verbose: bool | None = None) → list[list[np.ndarray]][source]¶

Preprocess outputs for inference.

Parameters:

outputs (list[np.ndarray]) – The outputs to preprocess.
ratios (list[tuple[float, float]]) – The ratios generated by the preprocess stage.
paddings (list[tuple[float, float]]) – The paddings generated by the preprocess stage.
no_copy (bool, optional) – If True, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The postprocessed outputs.

Return type:

list[list[np.ndarray]]

Raises:

ValueError – If outputs do not match the number of models

postprocess_model(outputs: list[np.ndarray], modelid: int, ratios: tuple[float, float], padding: tuple[float, float], *, no_copy: bool | None = None, verbose: bool | None = None) → list[np.ndarray][source]¶

Postprocess outputs for a specific model.

Parameters:

outputs (list[np.ndarray]) – The outputs to postprocess.
modelid (int) – The model to postprocess the data for.
ratios (tuple[float, float]) – The ratios generated by the preprocess stage
padding (tuple[float, float]) – The padding generated by the preprocess stage
no_copy (bool, optional) – If True, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The preprocessed data

Return type:

tuple[np.ndarray, tuple[float, float], tuple[float, float]]

get_detections(outputs: list[list[np.ndarray]], *, verbose: bool | None = None) → list[list[tuple[tuple[int, int, int, int], float, int]]][source]¶

Get the detections of the YOLO models.

Parameters:

outputs (list[list[np.ndarray]]) – The outputs of the models after postprocessing.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The detections

Return type:

list[list[tuple[tuple[int, int, int, int], float, int]]]

get_detections_model(output: list[np.ndarray], modelid: int, *, verbose: bool | None = None) → list[tuple[tuple[int, int, int, int], float, int]][source]¶

Get the detections of a single YOLO model.

Parameters:

output (list[np.ndarray]) – The output of a model after postprocessing.
modelid (int) – The model ID of which model is forming detections.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The detections produced by the model

Return type:

list[tuple[tuple[int, int, int, int], float, int]]

Submit data to be run for all models or a specific one.

Parameters:

inputs (list[np.ndarray]) – The inputs to pass to the models
ratios (list[tuple[float, float]], optional) – The optional ratio values for each input
paddings (list[tuple[float, float]], optional) – The optional padding values for each input
preprocessed (bool, optional) – Whether or not the inputs are preprocessed
postprocess (bool, optional) – Whether or not to postprocess the outputs right away
no_copy (bool, optional) – If True, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Raises:

ValueError – If the input length does not match the models If preprocessed is True, but ratios/paddings not provided

Submit data to a specific model.

Parameters:

inputs (np.ndarray) – The data to send to the model.
modelid (int) – The specific model indix to send the data to.
ratio (tuple[float, float], optional) – The ratio (if generated) by preprocess.
padding (tuple[float, float], optional) – The padding (if generated) by postprocess
preprocessed (bool, optional) – Whether or not the inputs are preprocessed.
postprocess (bool, optional) – Wherher or not to perform postprocessing.
no_copy (bool, optional) – If True, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

get_random_input() → list[np.ndarray][source]¶

Get random inputs.

Returns:: The random inputs
Return type:: list[np.ndarray]

mock_submit(data: list[ndarray] | ndarray | None = None, modelid: int | None = None) → None[source]¶

Perform a mock submit for all models or a specific model.

Parameters:

data (list[np.ndarray], np.ndarray, optional) – The inputs to use for the inference.
modelid (int, optional) – The specific engine to perform a mock submit for.

Raises:

ValueError – If specified modelid, but gave list of random input If gave single np.ndarray input, but did not specify modelid

retrieve(*, verbose: bool | None = None) → tuple[list[list[np.ndarray]], list[tuple[float, float] | None], list[tuple[float, float] | None]][source]¶

Get outputs back from all the models.

Parameters:: verbose (bool, optional) – Whether or not to log additional information.
Returns:: The outputs from all models
Return type:: list[tuple[list[np.ndarray], tuple[float, float], tuple[float, float]]]

retrieve_model(modelid: int, *, verbose: bool | None = None) → tuple[list[np.ndarray], tuple[float, float] | None, tuple[float, float] | None][source]¶

Get the outputs from a specific model.

Parameters:

modelid (int) – The model to retrieve data from.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The outputs of the model

Return type:

tuple[list[np.ndarray], tuple[float, float] | None, tuple[float, float] | None]

end2end(inputs: list[np.ndarray], ratios: list[tuple[float, float]] | None = None, paddings: list[tuple[float, float]] | None = None, *, preprocessed: bool | None = None, postprocess: bool | None = None, no_copy: bool | None = None, verbose: bool | None = None) → list[list[tuple[tuple[int, int, int, int], float, int]]][source]¶

Perform end-to-end inference for all models.

Parameters:

inputs (list[np.ndarray]) – The inputs to pass to the models
ratios (list[tuple[float, float]], optional) – The optional ratio values for each input
paddings (list[tuple[float, float]], optional) – The optional padding values for each input
preprocessed (bool, optional) – Whether or not the inputs are preprocessed
postprocess (bool, optional) – Whether or not to postprocess the outputs right away
no_copy (bool, optional) – If True, do not copy the data from the allocated memory. If the data is not copied, it WILL BE OVERWRITTEN INPLACE once new data is generated.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The outputs of the models

Return type:

list[list[tuple[tuple[int, int, int, int], float, int]]]

class trtutils.impls.yolo.TRTPreprocessor(output_shape: tuple[int, int], output_range: tuple[float, float], dtype: np.dtype, resize: str = 'letterbox', stream: cudart.cudaStream_t | None = None, threads: tuple[int, int, int] | None = None, tag: str | None = None, *, pagelocked_mem: bool | None = None, unified_mem: bool | None = None)[source]¶

Bases: object

TRT-based preprocessor for YOLO.

warmup() → None[source]¶

Warmup the CUDA preprocessor.

Allocates all CUDA memory and enables future passes to be significantly faster.

Preprocess an image for YOLO.

Parameters:

image (np.ndarray) – The image to preprocess.
resize (str, optional) – The method to resize the image with. Options are [letterbox, linear], will use method provided in constructor by default.
no_copy (bool, optional) – If True, the outputs will not be copied out from the cuda allocated host memory. Instead, the host memory will be returned directly. This memory WILL BE OVERWRITTEN INPLACE by future preprocessing calls.
verbose (bool, optional) – Whether or not to output additional information to stdout. If not provided, will default to overall engines verbose setting.

Returns:

The preprocessed image, ratios, and padding used for resizing.

Return type:

tuple[np.ndarray, tuple[float, float], tuple[float, float]]

Preprocess an image for YOLO.

Parameters:

image (np.ndarray) – The image to preprocess.
resize (str) – The method to resize the image with. By default letterbox, options are [letterbox, linear]
no_warn (bool, optional) – If True, do not warn about usage.
verbose (bool, optional) – Whether or not to output additional information to stdout. If not provided, will default to overall engines verbose setting.

Returns:

The GPU pointer to preprocessed data, ratios, and padding used for resizing.

Return type:

tuple[int, tuple[float, float], tuple[float, float]]

trtutils.impls.yolo.get_detections(outputs: list[ndarray], conf_thres: float | None = None, nms_iou_thres: float = 0.5, *, extra_nms: bool | None = None, agnostic_nms: bool | None = None, verbose: bool | None = None) → list[tuple[tuple[int, int, int, int], float, int]][source]¶

Get the detections from the output of a YOLO network.

Parameters:

outputs (list[np.ndarray]) – The outputs from a YOLO networks.
conf_thres (float, optional) – The confidence threshold to use for getting detections.
nms_iou_thres (float) – The IOU threshold to use during the optional additional NMS operation. By default, 0.5
extra_nms (bool, optional) – Whether or not an additional CPU-side NMS operation should be conducted on final detections.
agnostic_nms (bool, optional) – Whether or not to perform class-agnostic NMS during the optional additional operation.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The detections from the YOLO netowrk. Each detection is a bounding box in form x1, y1, x2, y2, a confidence score and a class id.

Return type:

list[tuple[tuple[int, int, int, int], float, int]]

trtutils.impls.yolo.postprocess(outputs: list[ndarray], ratios: tuple[float, float] = (1.0, 1.0), padding: tuple[float, float] = (0.0, 0.0), conf_thres: float | None = None, *, no_copy: bool | None = None, verbose: bool | None = None) → list[ndarray][source]¶

Postprocess outputs from a YOLO network.

Parameters:

outputs (list[np.ndarray]) – The outputs from a YOLO network.
ratios (tuple[float, float]) – The ratio of original image to preprocessed shape
padding (tuple[float, float]) – The amount of padding added during preprocessing
conf_thres (float, optional) – The confidence score for which detections below will be thrown out.
no_copy (bool, optional) – If True, the outputs will not be copied out from the cuda allocated host memory. Instead, the host memory will be returned directly. This memory WILL BE OVERWRITTEN INPLACE by future inference calls.
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The postprocessed outputs.

Return type:

list[np.ndarray]

trtutils.impls.yolo.preprocess(image: ndarray, input_shape: tuple[int, int], dtype: dtype, input_range: tuple[float, float] = (0.0, 1.0), method: str = 'letterbox', *, verbose: bool | None = None) → tuple[ndarray, tuple[float, float], tuple[float, float]][source]¶

Preprocess inputs for a YOLO network.

Parameters:

image (np.ndarray) – The inputs to be preprocessed.
input_shape (tuple[int, int]) – The shape to resize the inputs.
dtype (np.dtype) – The datatype of the inputs to the network.
input_range (tuple[float, float]) – The range of the model expects for inputs. By default, [0.0, 1.0] (divide input by 255.0)
method (str) – The method by which to resize the image. By default letterbox will be used. Options are [letterbox, linear]
verbose (bool, optional) – Whether or not to log additional information.

Returns:

The preprocessed data.

Return type:

tuple[np.ndarray, tuple[float, float], tuple[float, float]]

Raises:

ValueError – If the method for resizing is not ‘letterbox’ or ‘linear’