[ROIAlign-9] Extend nGraph Python API for operation "ROIAlign-9" (ope…

…nvinotoolkit#11572)

* add opset8 ngraph ROIAlign op

* fix style

* fix style v2

* remove redundnat added files

* fix __init__.py imports

* fix style v3

* fix wrong imports

* fix flake error

* fix minor errors

* add blank line

* fix args name

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

Co-authored-by: Jan Iwaszkiewicz <[email protected]>

* update docsstring, move roi_align tests to test_create_op.py file

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

Co-authored-by: Jan Iwaszkiewicz <[email protected]>

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

Co-authored-by: Jan Iwaszkiewicz <[email protected]>

* add alias

Co-authored-by: Jan Iwaszkiewicz <[email protected]>

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

@@ -132,7 +132,7 @@
 from ngraph.opset1.ops import reverse_sequence
 from ngraph.opset3.ops import rnn_cell
 from ngraph.opset5.ops import rnn_sequence
-from ngraph.opset3.ops import roi_align
+from ngraph.opset9.ops import roi_align
 from ngraph.opset2.ops import roi_pooling
 from ngraph.opset7.ops import roll
 from ngraph.opset5.ops import round

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

@@ -12,9 +12,10 @@
 from ngraph.utils.types import (
     NodeInput,
     as_nodes,
-    as_node
+    as_node,
 )
 
+
 _get_node_factory_opset9 = partial(_get_node_factory, "opset9")
 
 
@@ -48,6 +49,53 @@ def eye(
     return _get_node_factory_opset9().create("Eye", inputs, {"output_type": output_type})
 
 
+def roi_align(
+    data: NodeInput,
+    rois: NodeInput,
+    batch_indices: NodeInput,
+    pooled_h: int,
+    pooled_w: int,
+    sampling_ratio: int,
+    spatial_scale: float,
+    mode: str,
+    aligned_mode: Optional[str] = "asymmetric",
+    name: Optional[str] = None,
+) -> Node:
+    """Return a node which performs ROIAlign operation.
+
+    :param data: Input data.
+    :param rois: RoIs (Regions of Interest) to pool over.
+    :param batch_indices: Tensor with each element denoting the index of
+                          the corresponding image in the batch.
+    :param pooled_h: Height of the ROI output feature map.
+    :param pooled_w: Width of the ROI output feature map.
+    :param sampling_ratio: Number of bins over height and width to use to calculate
+                           each output feature map element.
+    :param spatial_scale: Multiplicative spatial scale factor to translate ROI coordinates.
+    :param mode: Method to perform pooling to produce output feature map elements. Avaiable modes are:
+                         - 'max' - maximum pooling
+                         - 'avg' - average pooling
+    :param aligned_mode: Specifies how to transform the coordinate in original tensor to the resized tensor.
+                         Mode 'asymmetric' is the default value. Optional. Avaiable aligned modes are:
+                         - 'asymmetric'
+                         - 'half_pixel_for_nn'
+                         - 'half_pixel'
+    :param name: The optional name for the output node
+
+    :return: The new node which performs ROIAlign
+    """
+    inputs = as_nodes(data, rois, batch_indices)
+    attributes = {
+        "pooled_h": pooled_h,
+        "pooled_w": pooled_w,
+        "sampling_ratio": sampling_ratio,
+        "spatial_scale": spatial_scale,
+        "mode": mode,
+        "aligned_mode": aligned_mode,
+    }
+    return _get_node_factory_opset9().create("ROIAlign", inputs, attributes)
+
+
 @nameable_op
 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

@@ -132,7 +132,7 @@
 from openvino.runtime.opset1.ops import reverse_sequence
 from openvino.runtime.opset3.ops import rnn_cell
 from openvino.runtime.opset5.ops import rnn_sequence
-from openvino.runtime.opset3.ops import roi_align
+from openvino.runtime.opset9.ops import roi_align
 from openvino.runtime.opset2.ops import roi_pooling
 from openvino.runtime.opset7.ops import roll
 from openvino.runtime.opset5.ops import round

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

@@ -49,6 +49,53 @@ def eye(
 
 
 @nameable_op
+def roi_align(
+    data: NodeInput,
+    rois: NodeInput,
+    batch_indices: NodeInput,
+    pooled_h: int,
+    pooled_w: int,
+    sampling_ratio: int,
+    spatial_scale: float,
+    mode: str,
+    aligned_mode: Optional[str] = "asymmetric",
+    name: Optional[str] = None,
+) -> Node:
+    """Return a node which performs ROIAlign operation.
+
+    :param data: Input data.
+    :param rois: RoIs (Regions of Interest) to pool over.
+    :param batch_indices: Tensor with each element denoting the index of
+                          the corresponding image in the batch.
+    :param pooled_h: Height of the ROI output feature map.
+    :param pooled_w: Width of the ROI output feature map.
+    :param sampling_ratio: Number of bins over height and width to use to calculate
+                           each output feature map element.
+    :param spatial_scale: Multiplicative spatial scale factor to translate ROI coordinates.
+    :param mode: Method to perform pooling to produce output feature map elements. Avaiable modes are:
+                         - 'max' - maximum pooling
+                         - 'avg' - average pooling
+    :param aligned_mode: Specifies how to transform the coordinate in original tensor to the resized tensor.
+                         Mode 'asymmetric' is the default value. Optional. Avaiable aligned modes are:
+                         - 'asymmetric'
+                         - 'half_pixel_for_nn'
+                         - 'half_pixel'
+    :param name: The optional name for the output node
+
+    :return: The new node which performs ROIAlign
+    """
+    inputs = as_nodes(data, rois, batch_indices)
+    attributes = {
+        "pooled_h": pooled_h,
+        "pooled_w": pooled_w,
+        "sampling_ratio": sampling_ratio,
+        "spatial_scale": spatial_scale,
+        "mode": mode,
+        "aligned_mode": aligned_mode,
+    }
+    return _get_node_factory_opset9().create("ROIAlign", inputs, attributes)
+
+
 def softsign(node: NodeInput, name: Optional[str] = None) -> Node:
     """Apply SoftSign operation on the input node element-wise.
 

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

@@ -850,6 +850,38 @@ def test_roi_pooling():
     assert node.get_output_element_type(0) == Type.f32
 
 
+@pytest.mark.parametrize(
+    ("data_shape", "rois", "batch_indices", "pooled_h", "pooled_w", "sampling_ratio", "spatial_scale", "mode", "aligned_mode", "expected_shape"),
+    [
+        ([2, 3, 5, 6], [7, 4], [7], 2, 2, 1, 1.0, "avg", "asymmetric", [7, 3, 2, 2]),
+        ([10, 3, 5, 5], [7, 4], [7], 3, 4, 1, 1.0, "avg", "half_pixel_for_nn", [7, 3, 3, 4]),
+        ([10, 3, 5, 5], [3, 4], [3], 3, 4, 1, 1.0, "avg", "half_pixel", [3, 3, 3, 4]),
+        ([10, 3, 5, 5], [3, 4], [3], 3, 4, 1, np.float(1), "avg", "half_pixel", [3, 3, 3, 4]),
+    ],
+)
+def test_roi_align(data_shape, rois, batch_indices, pooled_h, pooled_w, sampling_ratio, spatial_scale, mode, aligned_mode, expected_shape):
+    data_parameter = ov.parameter(data_shape, name="Data", dtype=np.float32)
+    rois_parameter = ov.parameter(rois, name="Rois", dtype=np.float32)
+    batch_indices_parameter = ov.parameter(batch_indices, name="Batch_indices", dtype=np.int32)
+
+    node = ov.roi_align(
+        data_parameter,
+        rois_parameter,
+        batch_indices_parameter,
+        pooled_h,
+        pooled_w,
+        sampling_ratio,
+        spatial_scale,
+        mode,
+        aligned_mode,
+    )
+
+    assert node.get_type_name() == "ROIAlign"
+    assert node.get_output_size() == 1
+    assert node.get_output_element_type(0) == Type.f32
+    assert list(node.get_output_shape(0)) == expected_shape
+
+
 def test_psroi_pooling():
     inputs = ov.parameter([1, 72, 4, 5], dtype=np.float32)
     coords = ov.parameter([150, 5], dtype=np.float32)

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

@@ -5,7 +5,7 @@
 import pytest
 from openvino.runtime import PartialShape, Dimension
 
-import openvino.runtime.opset8 as ov
+import openvino.runtime.opset9 as ov
 from openvino.runtime.utils.types import make_constant_node
 from tests.runtime import get_runtime
 from tests.test_ngraph.util import run_op_node

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

@@ -249,6 +249,38 @@ def test_deformable_psroi_pooling(dtype):
     assert list(node.get_output_shape(0)) == expected_shape
 
 
+@pytest.mark.parametrize(
+    ("data_shape", "rois", "batch_indices", "pooled_h", "pooled_w", "sampling_ratio", "spatial_scale", "mode", "aligned_mode", "expected_shape"),
+    [
+        ([2, 3, 5, 6], [7, 4], [7], 2, 2, 1, 1.0, "avg", "asymmetric", [7, 3, 2, 2]),
+        ([10, 3, 5, 5], [7, 4], [7], 3, 4, 1, 1.0, "avg", "half_pixel_for_nn", [7, 3, 3, 4]),
+        ([10, 3, 5, 5], [3, 4], [3], 3, 4, 1, 1.0, "avg", "half_pixel", [3, 3, 3, 4]),
+        ([10, 3, 5, 5], [3, 4], [3], 3, 4, 1, np.float(1), "avg", "half_pixel", [3, 3, 3, 4]),
+    ],
+)
+def test_roi_align(data_shape, rois, batch_indices, pooled_h, pooled_w, sampling_ratio, spatial_scale, mode, aligned_mode, expected_shape):
+    data_parameter = ng.parameter(data_shape, name="Data", dtype=np.float32)
+    rois_parameter = ng.parameter(rois, name="Rois", dtype=np.float32)
+    batch_indices_parameter = ng.parameter(batch_indices, name="Batch_indices", dtype=np.int32)
+
+    node = ng.roi_align(
+        data_parameter,
+        rois_parameter,
+        batch_indices_parameter,
+        pooled_h,
+        pooled_w,
+        sampling_ratio,
+        np.float32(spatial_scale),
+        mode,
+        aligned_mode,
+    )
+
+    assert node.get_type_name() == "ROIAlign"
+    assert node.get_output_size() == 1
+    assert node.get_output_element_type(0) == Type.f32
+    assert list(node.get_output_shape(0)) == expected_shape
+
+
 @pytest.mark.parametrize("dtype", np_types)
 def test_floor_mod(dtype):
     input0_shape = [8, 1, 6, 1]