Function _add_efficientnms_to_model

samples/common.py:461–562

Add TensorRT EfficientNMS plugin to ONNX model.

(
        model_proto: onnx.ModelProto,
        num_classes: int,
        max_detections: int,
    )

        )

    def _add_efficientnms_to_model(
        model_proto: onnx.ModelProto,
        num_classes: int,
        max_detections: int,
    ) -> None:
        """Add TensorRT EfficientNMS plugin to ONNX model."""
        import onnx
        import onnx.helper as helper

        graph = model_proto.graph

        # The raw_output is [B, anchors, num_classes+4]
        # We need to split it into boxes [B, anchors, 4] and scores [B, anchors, num_classes]

        # Create constants for slicing
        # Slice for classes: [:, :, :num_classes]
        start_cls = helper.make_tensor("start_cls", onnx.TensorProto.INT64, [1], [0])
        end_cls = helper.make_tensor(
            "end_cls", onnx.TensorProto.INT64, [1], [num_classes]
        )
        axes_cls = helper.make_tensor("axes_cls", onnx.TensorProto.INT64, [1], [2])

        # Slice for boxes: [:, :, num_classes:]
        start_box = helper.make_tensor(
            "start_box", onnx.TensorProto.INT64, [1], [num_classes]
        )
        end_box = helper.make_tensor(
            "end_box", onnx.TensorProto.INT64, [1], [num_classes + 4]
        )
        axes_box = helper.make_tensor("axes_box", onnx.TensorProto.INT64, [1], [2])

        # Add constant tensors to graph
        graph.initializer.extend(
            [start_cls, end_cls, axes_cls, start_box, end_box, axes_box]
        )

        # Create slice nodes to split raw_output
        slice_cls = helper.make_node(
            "Slice",
            inputs=["raw_output", "start_cls", "end_cls", "axes_cls"],
            outputs=["class_logits"],
            name="slice_class_logits",
        )

        slice_box = helper.make_node(
            "Slice",
            inputs=["raw_output", "start_box", "end_box", "axes_box"],
            outputs=["box_predictions"],
            name="slice_boxes",
        )

        # Apply sigmoid to class logits to get scores
        sigmoid_node = helper.make_node(
            "Sigmoid",
            inputs=["class_logits"],
            outputs=["class_scores"],
            name="sigmoid_scores",
        )