Support image padding in PrePostProcessor for python (openvinotoolkit…

…#23644) ### Details: - *Support image padding in PrePostProcessor for python* ### Tickets: - *CVS-121548* --------- Co-authored-by: Anastasia Kuporosova <[email protected]> Co-authored-by: Chen Peter <[email protected]> Co-authored-by: River Li <[email protected]>
ynimmaga · Jul 15, 2024 · 5136304 · 5136304
1 parent ae4404d
commit 5136304
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 1 deletion.
diff --git a/src/bindings/python/src/openvino/preprocess/__init__.py b/src/bindings/python/src/openvino/preprocess/__init__.py
@@ -24,3 +24,5 @@
 from openvino._pyopenvino.preprocess import PostProcessSteps
 from openvino._pyopenvino.preprocess import ColorFormat
 from openvino._pyopenvino.preprocess import ResizeAlgorithm
+from openvino._pyopenvino.preprocess import PaddingMode
+
diff --git a/src/bindings/python/src/pyopenvino/graph/preprocess/pre_post_process.cpp b/src/bindings/python/src/pyopenvino/graph/preprocess/pre_post_process.cpp
@@ -169,6 +169,62 @@ static void regclass_graph_PreProcessSteps(py::module m) {
     steps.def("reverse_channels", [](ov::preprocess::PreProcessSteps& self) {
         return &self.reverse_channels();
     });
+
+    steps.def(
+        "pad",
+        [](ov::preprocess::PreProcessSteps& self,
+           const std::vector<int>& pads_begin,
+           const std::vector<int>& pads_end,
+           float value,
+           ov::preprocess::PaddingMode mode) {
+            return &self.pad(pads_begin, pads_end, value, mode);
+        },
+        py::arg("pads_begin"),
+        py::arg("pads_end"),
+        py::arg("value"),
+        py::arg("mode"),
+        R"(
+            Adds padding preprocessing operation.
+
+            :param pads_begin: Number of elements matches the number of indices in data attribute. Specifies the number of padding elements at the ending of each axis.
+            :type pads_begin: 1D tensor of type T_INT.
+            :param pads_end: Number of elements matches the number of indices in data attribute. Specifies the number of padding elements at the ending of each axis.
+            :type pads_end: 1D tensor of type T_INT.
+            :param value: All new elements are populated with this value or with 0 if input not provided. Shouldn’t be set for other pad_mode values.
+            :type value: scalar tensor of type T. 
+            :param mode: pad_mode specifies the method used to generate new element values.
+            :type mode: string
+            :return: Reference to itself, allows chaining of calls in client's code in a builder-like manner.
+            :rtype: openvino.preprocess.PreProcessSteps
+        )");
+
+    steps.def(
+        "pad",
+        [](ov::preprocess::PreProcessSteps& self,
+           const std::vector<int>& pads_begin,
+           const std::vector<int>& pads_end,
+           const std::vector<float>& values,
+           ov::preprocess::PaddingMode mode) {
+            return &self.pad(pads_begin, pads_end, values, mode);
+        },
+        py::arg("pads_begin"),
+        py::arg("pads_end"),
+        py::arg("value"),
+        py::arg("mode"),
+        R"(
+            Adds padding preprocessing operation.
+
+            :param pads_begin: Number of elements matches the number of indices in data attribute. Specifies the number of padding elements at the ending of each axis.
+            :type pads_begin: 1D tensor of type T_INT.
+            :param pads_end: Number of elements matches the number of indices in data attribute. Specifies the number of padding elements at the ending of each axis.
+            :type pads_end: 1D tensor of type T_INT.
+            :param value: All new elements are populated with this value or with 0 if input not provided. Shouldn’t be set for other pad_mode values.
+            :type value: scalar tensor of type T. 
+            :param mode: pad_mode specifies the method used to generate new element values.
+            :type mode: string
+            :return: Reference to itself, allows chaining of calls in client's code in a builder-like manner.
+            :rtype: openvino.runtime.PreProcessSteps
+        )");
 }
 
 static void regclass_graph_PostProcessSteps(py::module m) {
@@ -469,6 +525,14 @@ static void regenum_graph_ResizeAlgorithm(py::module m) {
         .export_values();
 }
 
+static void regenum_graph_PaddingMode(py::module m) {
+    py::enum_<ov::preprocess::PaddingMode>(m, "PaddingMode")
+        .value("CONSTANT", ov::preprocess::PaddingMode::CONSTANT)
+        .value("REFLECT", ov::preprocess::PaddingMode::REFLECT)
+        .value("SYMMETRIC", ov::preprocess::PaddingMode::SYMMETRIC)
+        .export_values();
+}
+
 void regclass_graph_PrePostProcessor(py::module m) {
     regclass_graph_PreProcessSteps(m);
     regclass_graph_PostProcessSteps(m);
@@ -480,6 +544,7 @@ void regclass_graph_PrePostProcessor(py::module m) {
     regclass_graph_OutputModelInfo(m);
     regenum_graph_ColorFormat(m);
     regenum_graph_ResizeAlgorithm(m);
+    regenum_graph_PaddingMode(m);
     py::class_<ov::preprocess::PrePostProcessor, std::shared_ptr<ov::preprocess::PrePostProcessor>> proc(
         m,
         "PrePostProcessor");

diff --git a/src/bindings/python/tests/test_graph/test_preprocess.py b/src/bindings/python/tests/test_graph/test_preprocess.py
@@ -10,7 +10,7 @@
 from openvino import Core, Layout, Model, Shape, Tensor, Type
 from openvino.runtime.utils.decorators import custom_preprocess_function
 from openvino.runtime import Output
-from openvino.preprocess import PrePostProcessor, ColorFormat, ResizeAlgorithm
+from openvino.preprocess import PrePostProcessor, ColorFormat, ResizeAlgorithm, PaddingMode
 
 
 def test_graph_preprocess_mean():
@@ -728,3 +728,63 @@ def test_graph_set_layout_by_layout_class_thow_exception():
         layout = Layout("1-2-3D")
         ppp.input().model().set_layout(layout)
     assert "Layout name is invalid" in str(e.value)
+
+
+@pytest.mark.parametrize(("pads_begin", "pads_end", "values", "mode"), [([0, 0, 0, 0], [0, 0, 1, 1], 0, PaddingMode.CONSTANT)])
+def test_pad_vector_constant_layout(pads_begin, pads_end, values, mode):
+    shape = [1, 3, 200, 200]
+    parameter_a = ops.parameter(shape, dtype=np.float32, name="RGB_input")
+    model = parameter_a
+    model = Model(model, [parameter_a], "TestModel")
+    ppp = PrePostProcessor(model)
+    ppp.input().tensor().set_shape([1, 3, 199, 199])
+    ppp.input().preprocess().pad(pads_begin, pads_end, values, mode)
+    new_model = ppp.build()
+    assert new_model
+    assert list(new_model.get_output_shape(0)) == shape
+
+
+@pytest.mark.parametrize(("pads_begin", "pads_end", "values", "mode"), [([0, 0, -2, 0], [0, 0, -4, 1], 0, PaddingMode.CONSTANT)])
+def test_pad_vector_out_of_range(pads_begin, pads_end, values, mode):
+    shape = [1, 3, 5, 5]
+    parameter_a = ops.parameter(shape, dtype=np.float32, name="A")
+    model = parameter_a
+    model = Model(model, [parameter_a], "TestModel")
+    ppp = PrePostProcessor(model)
+    with pytest.raises(RuntimeError) as e:
+        ppp.input().preprocess().pad(pads_begin, pads_end, values, mode)
+        ppp.build()
+    assert "not aligned with original parameter's shape" in str(e.value)
+
+
+@pytest.mark.parametrize(("pads_begin", "pads_end", "values", "mode"), [([0, 0, 2, 0, 1], [0, 0, 4, 1, 1], 0, PaddingMode.CONSTANT)])
+def test_pad_vector_dim_mismatch(pads_begin, pads_end, values, mode):
+    shape = [1, 3, 5, 5]
+    parameter_a = ops.parameter(shape, dtype=np.float32, name="A")
+    model = parameter_a
+    model = Model(model, [parameter_a], "TestModel")
+    ppp = PrePostProcessor(model)
+    with pytest.raises(RuntimeError) as e:
+        ppp.input().preprocess().pad(pads_begin, pads_end, values, mode)
+        ppp.build()
+    assert "mismatches with rank of input" in str(e.value)
+
+
+@pytest.mark.parametrize(("pads_begin", "pads_end", "values", "mode"), [([0, 0, 0, 0], [0, 0, 1, 1], 0, PaddingMode.CONSTANT)])
+def test_pad_vector_type_and_ops(pads_begin, pads_end, values, mode):
+    shape = [1, 3, 200, 200]
+    parameter_a = ops.parameter(shape, dtype=np.float32, name="RGB_input")
+    model = parameter_a
+    model = Model(model, [parameter_a], "TestModel")
+    ppp = PrePostProcessor(model)
+    ppp.input().tensor().set_shape([1, 3, 199, 199])
+    ppp.input().preprocess().pad(pads_begin, pads_end, values, mode)
+    new_model = ppp.build()
+    assert new_model
+    model_operators = [op.get_name().split("_")[0] for op in model.get_ops()]
+    expected_ops = ["Parameter", "Constant", "Result", "Pad"]
+    assert list(new_model.get_output_shape(0)) == shape
+    assert new_model.get_output_element_type(0) == Type.f32
+    assert len(model_operators) == 6
+    for op in expected_ops:
+        assert op in model_operators