[PYTHON][NMS-9] Extend Python API for NMS-9 (openvinotoolkit#11681)

* extend NMS-9 ngraph python

* add tests for NMS

* move tests for NMS from test_reduction to test_create_op

src/bindings/python/src/compatibility/ngraph/opset9/__init__.py

@@ -98,7 +98,7 @@
 from ngraph.opset1.ops import multiply
 from ngraph.opset6.ops import mvn
 from ngraph.opset1.ops import negative
-from ngraph.opset5.ops import non_max_suppression
+from ngraph.opset9.ops import non_max_suppression
 from ngraph.opset3.ops import non_zero
 from ngraph.opset1.ops import normalize_l2
 from ngraph.opset1.ops import not_equal

src/bindings/python/src/compatibility/ngraph/opset9/ops.py

@@ -13,6 +13,7 @@
     NodeInput,
     as_nodes,
     as_node,
+    make_constant_node,
 )
 
 
@@ -97,6 +98,49 @@ def roi_align(
 
 
 @nameable_op
+def non_max_suppression(
+    boxes: NodeInput,
+    scores: NodeInput,
+    max_output_boxes_per_class: Optional[NodeInput] = None,
+    iou_threshold: Optional[NodeInput] = None,
+    score_threshold: Optional[NodeInput] = None,
+    soft_nms_sigma: Optional[NodeInput] = None,
+    box_encoding: str = "corner",
+    sort_result_descending: bool = True,
+    output_type: str = "i64",
+    name: Optional[str] = None,
+) -> Node:
+    """Return a node which performs NonMaxSuppression.
+
+    :param boxes: Tensor with box coordinates.
+    :param scores: Tensor with box scores.
+    :param max_output_boxes_per_class: Tensor Specifying maximum number of boxes
+                                        to be selected per class.
+    :param iou_threshold: Tensor specifying intersection over union threshold
+    :param score_threshold: Tensor specifying minimum score to consider box for the processing.
+    :param soft_nms_sigma: Tensor specifying the sigma parameter for Soft-NMS.
+    :param box_encoding: Format of boxes data encoding.
+    :param sort_result_descending: Flag that specifies whenever it is necessary to sort selected
+                                   boxes across batches or not.
+    :param output_type: Output element type.
+    :return: The new node which performs NonMaxSuppression
+    """
+    max_output_boxes_per_class = max_output_boxes_per_class if max_output_boxes_per_class is not None else make_constant_node(0, np.int64)
+    iou_threshold = iou_threshold if iou_threshold is not None else make_constant_node(0, np.float32)
+    score_threshold = score_threshold if score_threshold is not None else make_constant_node(0, np.float32)
+    soft_nms_sigma = soft_nms_sigma if soft_nms_sigma is not None else make_constant_node(0, np.float32)
+
+    inputs = as_nodes(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, soft_nms_sigma)
+
+    attributes = {
+        "box_encoding": box_encoding,
+        "sort_result_descending": sort_result_descending,
+        "output_type": output_type,
+    }
+
+    return _get_node_factory_opset9().create("NonMaxSuppression", inputs, attributes)
+
+
 def softsign(node: NodeInput, name: Optional[str] = None) -> Node:
     """Apply SoftSign operation on the input node element-wise.
 

src/bindings/python/src/openvino/runtime/opset9/__init__.py

@@ -98,7 +98,7 @@
 from openvino.runtime.opset1.ops import multiply
 from openvino.runtime.opset6.ops import mvn
 from openvino.runtime.opset1.ops import negative
-from openvino.runtime.opset5.ops import non_max_suppression
+from openvino.runtime.opset9.ops import non_max_suppression
 from openvino.runtime.opset3.ops import non_zero
 from openvino.runtime.opset1.ops import normalize_l2
 from openvino.runtime.opset1.ops import not_equal

src/bindings/python/src/openvino/runtime/opset9/ops.py

@@ -12,7 +12,8 @@
 from openvino.runtime.utils.types import (
     NodeInput,
     as_nodes,
-    as_node
+    as_node,
+    make_constant_node,
 )
 
 _get_node_factory_opset9 = partial(_get_node_factory, "opset9")
@@ -49,6 +50,49 @@ def eye(
 
 
 @nameable_op
+def non_max_suppression(
+    boxes: NodeInput,
+    scores: NodeInput,
+    max_output_boxes_per_class: Optional[NodeInput] = None,
+    iou_threshold: Optional[NodeInput] = None,
+    score_threshold: Optional[NodeInput] = None,
+    soft_nms_sigma: Optional[NodeInput] = None,
+    box_encoding: str = "corner",
+    sort_result_descending: bool = True,
+    output_type: str = "i64",
+    name: Optional[str] = None,
+) -> Node:
+    """Return a node which performs NonMaxSuppression.
+
+    :param boxes: Tensor with box coordinates.
+    :param scores: Tensor with box scores.
+    :param max_output_boxes_per_class: Tensor Specifying maximum number of boxes
+                                        to be selected per class.
+    :param iou_threshold: Tensor specifying intersection over union threshold
+    :param score_threshold: Tensor specifying minimum score to consider box for the processing.
+    :param soft_nms_sigma: Tensor specifying the sigma parameter for Soft-NMS.
+    :param box_encoding: Format of boxes data encoding.
+    :param sort_result_descending: Flag that specifies whenever it is necessary to sort selected
+                                   boxes across batches or not.
+    :param output_type: Output element type.
+    :return: The new node which performs NonMaxSuppression
+    """
+    max_output_boxes_per_class = max_output_boxes_per_class if max_output_boxes_per_class is not None else make_constant_node(0, np.int64)
+    iou_threshold = iou_threshold if iou_threshold is not None else make_constant_node(0, np.float32)
+    score_threshold = score_threshold if score_threshold is not None else make_constant_node(0, np.float32)
+    soft_nms_sigma = soft_nms_sigma if soft_nms_sigma is not None else make_constant_node(0, np.float32)
+
+    inputs = as_nodes(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, soft_nms_sigma)
+
+    attributes = {
+        "box_encoding": box_encoding,
+        "sort_result_descending": sort_result_descending,
+        "output_type": output_type,
+    }
+
+    return _get_node_factory_opset9().create("NonMaxSuppression", inputs, attributes)
+
+
 def roi_align(
     data: NodeInput,
     rois: NodeInput,

src/bindings/python/tests/test_ngraph/test_create_op.py

@@ -5,6 +5,7 @@
 import pytest
 from openvino.runtime import PartialShape, Dimension, Model
 from openvino.runtime.exceptions import UserInputError
+from openvino.runtime.utils.types import make_constant_node
 
 import openvino.runtime.opset9 as ov
 import openvino.runtime.opset1 as ov_opset1
@@ -1857,6 +1858,51 @@ def test_matrix_nms():
     assert nms_node.get_output_element_type(2) == Type.i32
 
 
+@pytest.mark.parametrize(
+    ("boxes_shape", "scores_shape", "max_output_boxes", "expected_shape"),
+    [
+        ([1, 1000, 4], [1, 1, 1000], [1000], [PartialShape([Dimension(0, 1000), Dimension(3)]), PartialShape([Dimension(0, 1000), Dimension(3)])]),
+        ([1, 700, 4], [1, 1, 700], [600], [PartialShape([Dimension(0, 600), Dimension(3)]), PartialShape([Dimension(0, 600), Dimension(3)])]),
+        ([1, 300, 4], [1, 1, 300], [300], [PartialShape([Dimension(0, 300), Dimension(3)]), PartialShape([Dimension(0, 300), Dimension(3)])]),
+    ],
+)
+def test_non_max_suppression(boxes_shape, scores_shape, max_output_boxes, expected_shape):
+    boxes_parameter = ov.parameter(boxes_shape, name="Boxes", dtype=np.float32)
+    scores_parameter = ov.parameter(scores_shape, name="Scores", dtype=np.float32)
+
+    node = ov.non_max_suppression(boxes_parameter, scores_parameter, make_constant_node(max_output_boxes, np.int64))
+    assert node.get_type_name() == "NonMaxSuppression"
+    assert node.get_output_size() == 3
+    assert node.get_output_partial_shape(0) == expected_shape[0]
+    assert node.get_output_partial_shape(1) == expected_shape[1]
+    assert list(node.get_output_shape(2)) == [1]
+
+
+@pytest.mark.parametrize(
+    ("boxes_shape", "scores_shape", "max_output_boxes", "iou_threshold", "score_threshold", "soft_nms_sigma", "expected_shape"),
+    [
+        ([1, 100, 4], [1, 1, 100], [100], 0.1, 0.4, 0.5, [PartialShape([Dimension(0, 100), Dimension(3)]), PartialShape([Dimension(0, 100), Dimension(3)])]),
+        ([1, 700, 4], [1, 1, 700], [600], 0.1, 0.4, 0.5, [PartialShape([Dimension(0, 600), Dimension(3)]), PartialShape([Dimension(0, 600), Dimension(3)])]),
+        ([1, 300, 4], [1, 1, 300], [300], 0.1, 0.4, 0.5, [PartialShape([Dimension(0, 300), Dimension(3)]), PartialShape([Dimension(0, 300), Dimension(3)])]),
+    ],
+)
+def test_non_max_suppression_non_default_args(boxes_shape, scores_shape, max_output_boxes, iou_threshold, score_threshold, soft_nms_sigma, expected_shape):
+    boxes_parameter = ov.parameter(boxes_shape, name="Boxes", dtype=np.float32)
+    scores_parameter = ov.parameter(scores_shape, name="Scores", dtype=np.float32)
+
+    max_output_boxes = make_constant_node(max_output_boxes, np.int64)
+    iou_threshold = make_constant_node(iou_threshold, np.float32)
+    score_threshold = make_constant_node(score_threshold, np.float32)
+    soft_nms_sigma = make_constant_node(soft_nms_sigma, np.float32)
+
+    node = ov.non_max_suppression(boxes_parameter, scores_parameter, max_output_boxes, iou_threshold, score_threshold, soft_nms_sigma)
+    assert node.get_type_name() == "NonMaxSuppression"
+    assert node.get_output_size() == 3
+    assert node.get_output_partial_shape(0) == expected_shape[0]
+    assert node.get_output_partial_shape(1) == expected_shape[1]
+    assert list(node.get_output_shape(2)) == [1]
+
+
 def test_slice():
     data_shape = [10, 7, 2, 13]
     data = ov.parameter(data_shape, name="input", dtype=np.float32)

src/bindings/python/tests/test_ngraph/test_reduction.py

@@ -89,22 +89,6 @@ def test_reduce_mean_op(ng_api_helper, numpy_function, reduction_axes):
     assert np.allclose(result, expected)
 
 
-def test_non_max_suppression():
-
-    boxes_shape = [1, 1000, 4]
-    scores_shape = [1, 1, 1000]
-    boxes_parameter = ov.parameter(boxes_shape, name="Boxes", dtype=np.float32)
-    scores_parameter = ov.parameter(scores_shape, name="Scores", dtype=np.float32)
-
-    node = ov.non_max_suppression(boxes_parameter, scores_parameter, make_constant_node(1000, np.int64))
-
-    assert node.get_type_name() == "NonMaxSuppression"
-    assert node.get_output_size() == 3
-    assert node.get_output_partial_shape(0) == PartialShape([Dimension(0, 1000), Dimension(3)])
-    assert node.get_output_partial_shape(1) == PartialShape([Dimension(0, 1000), Dimension(3)])
-    assert list(node.get_output_shape(2)) == [1]
-
-
 def test_non_zero():
 
     data_shape = [3, 10, 100, 200]

src/bindings/python/tests_compatibility/test_ngraph/test_create_op.py

@@ -8,6 +8,7 @@
 import ngraph as ng
 import ngraph.opset1 as ng_opset1
 import ngraph.opset5 as ng_opset5
+from ngraph.utils.types import make_constant_node
 from ngraph.exceptions import UserInputError
 from ngraph.impl import Type
 
@@ -1923,6 +1924,51 @@ def test_matrix_nms():
     assert nms_node.get_output_element_type(2) == Type.i32
 
 
+@pytest.mark.parametrize(
+    ("boxes_shape", "scores_shape", "max_output_boxes", "expected_shape"),
+    [
+        ([1, 1000, 4], [1, 1, 1000], [1000], [PartialShape([Dimension(0, 1000), Dimension(3)]), PartialShape([Dimension(0, 1000), Dimension(3)])]),
+        ([1, 700, 4], [1, 1, 700], [600], [PartialShape([Dimension(0, 600), Dimension(3)]), PartialShape([Dimension(0, 600), Dimension(3)])]),
+        ([1, 300, 4], [1, 1, 300], [300], [PartialShape([Dimension(0, 300), Dimension(3)]), PartialShape([Dimension(0, 300), Dimension(3)])]),
+    ],
+)
+def test_non_max_suppression(boxes_shape, scores_shape, max_output_boxes, expected_shape):
+    boxes_parameter = ng.parameter(boxes_shape, name="Boxes", dtype=np.float32)
+    scores_parameter = ng.parameter(scores_shape, name="Scores", dtype=np.float32)
+
+    node = ng.non_max_suppression(boxes_parameter, scores_parameter, make_constant_node(max_output_boxes, np.int64))
+    assert node.get_type_name() == "NonMaxSuppression"
+    assert node.get_output_size() == 3
+    assert node.get_output_partial_shape(0) == expected_shape[0]
+    assert node.get_output_partial_shape(1) == expected_shape[1]
+    assert list(node.get_output_shape(2)) == [1]
+
+
+@pytest.mark.parametrize(
+    ("boxes_shape", "scores_shape", "max_output_boxes", "iou_threshold", "score_threshold", "soft_nms_sigma", "expected_shape"),
+    [
+        ([1, 100, 4], [1, 1, 100], [100], 0.1, 0.4, 0.5, [PartialShape([Dimension(0, 100), Dimension(3)]), PartialShape([Dimension(0, 100), Dimension(3)])]),
+        ([1, 700, 4], [1, 1, 700], [600], 0.1, 0.4, 0.5, [PartialShape([Dimension(0, 600), Dimension(3)]), PartialShape([Dimension(0, 600), Dimension(3)])]),
+        ([1, 300, 4], [1, 1, 300], [300], 0.1, 0.4, 0.5, [PartialShape([Dimension(0, 300), Dimension(3)]), PartialShape([Dimension(0, 300), Dimension(3)])]),
+    ],
+)
+def test_non_max_suppression_non_default_args(boxes_shape, scores_shape, max_output_boxes, iou_threshold, score_threshold, soft_nms_sigma, expected_shape):
+    boxes_parameter = ng.parameter(boxes_shape, name="Boxes", dtype=np.float32)
+    scores_parameter = ng.parameter(scores_shape, name="Scores", dtype=np.float32)
+
+    max_output_boxes = make_constant_node(max_output_boxes, np.int64)
+    iou_threshold = make_constant_node(iou_threshold, np.float32)
+    score_threshold = make_constant_node(score_threshold, np.float32)
+    soft_nms_sigma = make_constant_node(soft_nms_sigma, np.float32)
+
+    node = ng.non_max_suppression(boxes_parameter, scores_parameter, max_output_boxes, iou_threshold, score_threshold, soft_nms_sigma)
+    assert node.get_type_name() == "NonMaxSuppression"
+    assert node.get_output_size() == 3
+    assert node.get_output_partial_shape(0) == expected_shape[0]
+    assert node.get_output_partial_shape(1) == expected_shape[1]
+    assert list(node.get_output_shape(2)) == [1]
+
+
 def test_slice():
     data_shape = [10, 7, 2, 13]
     data = ng.parameter(data_shape, name="input", dtype=np.float32)

src/bindings/python/tests_compatibility/test_ngraph/test_reduction.py

@@ -89,22 +89,6 @@ def test_reduce_mean_op(ng_api_helper, numpy_function, reduction_axes):
     assert np.allclose(result, expected)
 
 
-def test_non_max_suppression():
-
-    boxes_shape = [1, 1000, 4]
-    scores_shape = [1, 1, 1000]
-    boxes_parameter = ng.parameter(boxes_shape, name="Boxes", dtype=np.float32)
-    scores_parameter = ng.parameter(scores_shape, name="Scores", dtype=np.float32)
-
-    node = ng.non_max_suppression(boxes_parameter, scores_parameter, make_constant_node(1000, np.int64))
-
-    assert node.get_type_name() == "NonMaxSuppression"
-    assert node.get_output_size() == 3
-    assert node.get_output_partial_shape(0) == PartialShape([Dimension(0, 1000), Dimension(3)])
-    assert node.get_output_partial_shape(1) == PartialShape([Dimension(0, 1000), Dimension(3)])
-    assert list(node.get_output_shape(2)) == [1]
-
-
 def test_non_zero():
 
     data_shape = [3, 10, 100, 200]