trtutils.impls.common module

Common implementations for TensorRT engines.

Functions

decode_efficient_nms()

Processes the output of a model with EfficientNMS plugin outputs.

postprocess_efficient_nms()

Postprocesses the output of a model to reshape and scale based on preprocessing.

trtutils.impls.common.postprocess_efficient_nms(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 the output of the EfficientNMS plugin.

Must be used before passing outputs to decode_efficient_nms since this will reshape the outputs.

Parameters:

  • outputs (list[np.ndarray]) – The outputs from the TRTEngine using EfficientNMS output.

  • ratios (tuple[float, float]) – The ratios used during preprocessing to resize the input.

  • padding (tuple[float, float]) – The padding used during preprocessing to position the input.

  • conf_thres (float, optional) – Optional confidence threshold to further filter detections by. Detections are already filtered by EfficientNMS parameters ahead of time. Should be used if EfficientNMS was given low-confidence and want to filter higher variably.

  • 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 log additional information.

Returns:

The postprocessed outputs, reshaped and scaled based on ratios/padding.

Return type:

list[np.ndarray]

trtutils.impls.common.decode_efficient_nms(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]

Decode EfficientNMS plugin output.

Must have called postprocess_efficient_nms before calling this function, due to the reshape stage needing to occur.

Parameters:

  • outputs (list[np.ndarray]) – The outputs from a model with EfficientNMS output

  • conf_thres (float) – A confidence value to threshold detctions by. By default None.

  • 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 decoded outputs. Bounding box (x1, y1, x2, y2), confidence score, classid

Return type:

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