Function test_nms_score_threshold

tests/python/relax/test_frontend_onnx.py:5478–5567

Test that NMS correctly filters boxes based on score threshold. Note: This test uses a low score threshold (0.05) to ensure both TVM and ONNX Runtime output the same fixed shape [3,3], allowing use of the standard check_correctness function.

()

def test_nms_score_threshold():
    """Test that NMS correctly filters boxes based on score threshold.

    Note: This test uses a low score threshold (0.05) to ensure both TVM and ONNX Runtime
    output the same fixed shape [3,3], allowing use of the standard check_correctness function.
    """
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
        center_point_box=0,
    )

    # Create data with varying scores - ensure we get exactly 3 boxes after NMS
    boxes_data = np.array(
        [
            [[0.0, 0.0, 1.0, 1.0], [2.0, 0.0, 3.0, 1.0], [0.0, 2.0, 1.0, 3.0]]  # Box 0  # Box 1
        ],  # Box 2
        dtype=np.float32,
    )

    # Scores: 0.9, 0.3, 0.1 - adjust score threshold to get exactly 3 boxes
    scores_data = np.array([[[0.9, 0.3, 0.1]]], dtype=np.float32)

    boxes_shape = [1, 3, 4]
    scores_shape = [1, 1, 3]

    graph = helper.make_graph(
        [nms_node],
        "nms_test_score_threshold",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
        ],
        initializer=[
            helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [3]),
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.1]),
            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.05]),
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [3, 3])],
    )

    model = helper.make_model(graph, producer_name="nms_test_score_threshold")
    model.ir_version = 8
    model.opset_import[0].version = 11

    inputs = {
        "boxes": boxes_data,
        "scores": scores_data,
    }

    # Run ONNX Runtime
    ort_session = onnxruntime.InferenceSession(
        model.SerializeToString(), providers=["CPUExecutionProvider"]
    )
    ort_output = ort_session.run([], inputs)

    # Run TVM