trtutils.builder.hooks package

Module contents

Submodule containing hooks for building TensorRT engines.

Functions

yolo_efficient_nms_hook()

Hook for building YOLO models with EfficientNMS.

trtutils.builder.hooks.yolo_efficient_nms_hook(num_classes: int = 80, conf_threshold: float = 0.25, iou_threshold: float = 0.5, top_k: int = 100, box_coding: str = 'center_size', *, class_agnostic: bool | None = None) → Callable[[trt.INetworkDefinition], trt.INetworkDefinition][source]

Create a hook to add EfficientNMS_TRT plugin to YOLO-like output network.

Expects a network with output shaped (N, num_classes, num_boxes) or (N, num_boxes, num_classes).

  • Interprets first 4 channels as box coordinates and the remaining as class scores

  • Supports outputs with or without an explicit objectness channel (4 + num_classes) or (4 + 1 + num_classes)

  • Replaces raw network outputs with NMS outputs: num_dets, det_boxes, det_scores, det_classes

Parameters:

  • num_classes (int, optional) – Number of classes in the dataset. Default is 80.

  • conf_threshold (float, optional) – Confidence threshold for filtering boxes. Default is 0.25.

  • iou_threshold (float, optional) – IoU threshold for NMS. Default is 0.5.

  • top_k (int, optional) – Number of top detections to keep. Default is 100.

  • class_agnostic (bool, optional) – Whether to use class-agnostic NMS. Default is False.

  • box_coding (str, optional) – Coding of the bounding boxes. Default is “center_size”.

Returns:

A hook that can be used to modify a network.

Return type:

Callable[[trt.INetworkDefinition], trt.INetworkDefinition]