diff --git a/src/frontends/tensorflow/src/op/non_max_suppression.cpp b/src/frontends/tensorflow/src/op/non_max_suppression.cpp index 63f95d71f5d..c437c7fc7f6 100644 --- a/src/frontends/tensorflow/src/op/non_max_suppression.cpp +++ b/src/frontends/tensorflow/src/op/non_max_suppression.cpp @@ -3,81 +3,154 @@ // #include "op_table.hpp" -#include "openvino/opsets/opset8.hpp" +#include "openvino/opsets/opset10.hpp" using namespace std; -using namespace ov::opset8; +using namespace ov; +using namespace ov::opset10; namespace ov { namespace frontend { namespace tensorflow { namespace op { +Output normalize_selected_indices(const Output& ov_selected_indices, + const Output& max_output_size, + bool pad_to_max_output_size) { + // OpenVINO NonMaxSuppression outputs selected_indices in a form [batch_index, class_index, box_index] + // but we need only box_index + auto start = make_shared(element::i32, Shape{1}, 2); + auto stop = make_shared(element::i32, Shape{1}, 3); + auto step = make_shared(element::i32, Shape{1}, 1); + auto axis = make_shared(element::i32, Shape{1}, 1); + Output selected_indices = make_shared(ov_selected_indices, start, stop, step, axis)->output(0); + selected_indices = make_shared(selected_indices, axis)->output(0); + + if (pad_to_max_output_size) { + auto num_selected_indices = make_shared(selected_indices, max_output_size.get_element_type()); + auto num_padded_elements = make_shared(max_output_size, num_selected_indices); + auto pad_element = make_shared(selected_indices.get_element_type(), Shape{}, 0); + auto pads_begin = make_shared(max_output_size.get_element_type(), Shape{1}, 0); + auto pad_mode = ov::op::PadMode::CONSTANT; + selected_indices = + make_shared(selected_indices, pads_begin, num_padded_elements, pad_element, pad_mode)->output(0); + } + + return selected_indices; +} OutputVector translate_non_max_suppression_op(const NodeContext& node) { + default_op_checks(node, + 3, + {"NonMaxSuppression", + "NonMaxSuppressionV2", + "NonMaxSuppressionV3", + "NonMaxSuppressionV4", + "NonMaxSuppressionV5"}); auto boxes = node.get_input(0); auto scores = node.get_input(1); auto max_output_size = node.get_input(2); - auto iou_threshold = node.get_input(3); - auto axis = make_shared(element::i64, Shape{1}, 0); - auto boxes_unsqueezed = make_shared(boxes, axis); + // prepare boxes: in TensorFlow NonMaxSupression operation there is no batch dimension + // so we need to introduce it to be aligned with OpenVINO NonMaxSuppression + // boxes of shape [num_batches, num_boxes, 4] + auto boxes_axis = make_shared(element::i32, Shape{1}, 0); + auto ov_boxes = make_shared(boxes, boxes_axis); - auto axis_scores = make_shared(element::i64, Shape{2}, vector{0, 1}); - auto scores_unsqueezed = make_shared(scores, axis_scores); + // prepare scores: in TensorFlow NonMaxSupression operation there is no batch and class dimensions + // so we need to introduce them to be aligned with OpenVINO NonMaxSuppression scores + // of shape [num_batches, num_classes, num_boxes] + auto scores_axes = make_shared(element::i32, Shape{2}, vector{0, 1}); + auto ov_scores = make_shared(scores, scores_axes); const auto& op_type = node.get_op_type(); - if (op_type == "NonMaxSuppressionV5") { - auto score_threshold = node.get_input(4); - auto soft_nms_sigma = node.get_input(5); - // todo: pad_to_max_output_size - auto res = make_shared(boxes_unsqueezed, - scores_unsqueezed, - max_output_size, - iou_threshold, - score_threshold, - soft_nms_sigma, - NonMaxSuppression::BoxEncodingType::CORNER, - false, - element::Type_t::i32); - set_node_name(node.get_name(), res); - return res->outputs(); - } else if (op_type == "NonMaxSuppressionV4") { - auto score_threshold = node.get_input(4); - // todo: pad_to_max_output_size - auto res = make_shared(boxes_unsqueezed, - scores_unsqueezed, - max_output_size, - iou_threshold, - score_threshold, - NonMaxSuppression::BoxEncodingType::CORNER, - false, - element::Type_t::i32); - set_node_name(node.get_name(), res); - return res->outputs(); + OutputVector results; + + // set all thresholds to zero, default values + Output iou_threshold = make_shared(element::f32, Shape{}, 0.0); + Output score_threshold = make_shared(element::f32, Shape{}, 0.0); + Output soft_nms_sigma = make_shared(element::f32, Shape{}, 0.0); + bool pad_to_max_output_size = false; + bool valid_outputs = false; + bool selected_scores = false; + if (op_type == "NonMaxSuppression") { + // NonMaxSuppression has three inputs and one output, selected indices + // iou_threshold comes as attribute + auto iou_threshold_value = node.get_attribute("iou_threshold", 0.5); + iou_threshold = make_shared(element::f32, Shape{}, iou_threshold_value); + } else if (op_type == "NonMaxSuppressionV2") { + // NonMaxSuppressionV2 has four inputs and one output, selected indices + // iou_threshold comes as input + default_op_checks(node, 4, {"NonMaxSuppressionV2"}); + auto iou_threshold = node.get_input(3); } else if (op_type == "NonMaxSuppressionV3") { - auto score_threshold = node.get_input(4); - auto res = make_shared(boxes_unsqueezed, - scores_unsqueezed, - max_output_size, - iou_threshold, - score_threshold, - NonMaxSuppression::BoxEncodingType::CORNER, - false, - element::Type_t::i32); - set_node_name(node.get_name(), res); - return {res->output(0)}; - } else if (op_type == "NonMaxSuppressionV2" || op_type == "NonMaxSuppression") { - auto res = make_shared(boxes_unsqueezed, - scores_unsqueezed, - max_output_size, - iou_threshold, - NonMaxSuppression::BoxEncodingType::CORNER, - false, - element::Type_t::i32); - set_node_name(node.get_name(), res); - return {res->output(0)}; + // NonMaxSuppressionV3 has five inputs and one output, selected indices + // score_threshold is a new input + default_op_checks(node, 5, {"NonMaxSuppressionV3"}); + iou_threshold = node.get_input(3); + score_threshold = node.get_input(4); + } else if (op_type == "NonMaxSuppressionV4") { + // NonMaxSuppressionV4 has five inputs (same as NonMaxSuppressionV3) and two outputs (selected indices and + // valid_outputs) pad_to_max_output_size is a new attribute + default_op_checks(node, 5, {"NonMaxSuppressionV4"}); + iou_threshold = node.get_input(3); + score_threshold = node.get_input(4); + pad_to_max_output_size = node.get_attribute("pad_to_max_output_size", false); + valid_outputs = true; + } else if (op_type == "NonMaxSuppressionV5") { + // NonMaxSuppressionV5 has six inputs and three outputs (selected indices, selected scores, + // valid_outputs) + // soft_nms_sigma is a new input + default_op_checks(node, 6, {"NonMaxSuppressionV5"}); + iou_threshold = node.get_input(3); + score_threshold = node.get_input(4); + soft_nms_sigma = node.get_input(5); + pad_to_max_output_size = node.get_attribute("pad_to_max_output_size", false); + selected_scores = true; + valid_outputs = true; + } else { + TENSORFLOW_OP_VALIDATION(node, false, "Internal error TensorFlow Frontend: no translator for " + op_type); } - TENSORFLOW_OP_VALIDATION(node, false, "No translator found."); + + auto non_max_suppression = make_shared(ov_boxes, + ov_scores, + max_output_size, + iou_threshold, + score_threshold, + soft_nms_sigma, + NonMaxSuppression::BoxEncodingType::CORNER, + false, + element::i32); + auto tf_selected_indices = + normalize_selected_indices(non_max_suppression->output(0), max_output_size, pad_to_max_output_size); + results.push_back(tf_selected_indices); + + Output tf_selected_scores; + if (selected_scores) { + tf_selected_scores = + normalize_selected_indices(non_max_suppression->output(1), max_output_size, pad_to_max_output_size); + tf_selected_scores = make_shared(tf_selected_scores, boxes.get_element_type())->output(0); + results.push_back(tf_selected_scores); + } + + Output tf_valid_outputs; + if (valid_outputs) { + // In TensorFlow valid_outputs output is a scalar + tf_valid_outputs = make_shared(non_max_suppression->output(2))->output(0); + results.push_back(tf_valid_outputs); + } + + // set output tensor names based on a number of outputs + set_node_name(node.get_name(), tf_selected_indices.get_node_shared_ptr()); + if (valid_outputs && !selected_scores) { + set_node_name(node.get_name(), tf_selected_indices.get_node_shared_ptr()); + set_node_name(node.get_name() + ":1", tf_valid_outputs.get_node_shared_ptr()); + } else if (valid_outputs && selected_scores) { + set_node_name(node.get_name(), tf_selected_indices.get_node_shared_ptr()); + set_node_name(node.get_name() + ":1", tf_selected_scores.get_node_shared_ptr()); + set_node_name(node.get_name() + ":2", tf_valid_outputs.get_node_shared_ptr()); + } + + return results; } } // namespace op } // namespace tensorflow diff --git a/tests/layer_tests/tensorflow_tests/test_tf_NonMaxSupression.py b/tests/layer_tests/tensorflow_tests/test_tf_NonMaxSupression.py index b3a06eb9478..861007c4ad5 100644 --- a/tests/layer_tests/tensorflow_tests/test_tf_NonMaxSupression.py +++ b/tests/layer_tests/tensorflow_tests/test_tf_NonMaxSupression.py @@ -40,48 +40,49 @@ class TestNonMaxSupression(CommonTFLayerTest): if with_scores: soft_nms_sigma = tf.constant(test_params["soft_nms_sigma"]) _ = tf.image.non_max_suppression_with_scores(boxes, scores, max_output_size, - iou_threshold, score_threshold, soft_nms_sigma, name="NMS") + iou_threshold, score_threshold, soft_nms_sigma, name="NMS") else: _ = tf.image.non_max_suppression(boxes, scores, max_output_size, - iou_threshold, score_threshold, name="NMS") + iou_threshold, score_threshold, name="NMS") tf_net = sess.graph_def ref_net = None return tf_net, ref_net test_params = [ - ( - { - "number_of_boxes": 50, - "max_output_size": 5, - "iou_threshold": 0.7, - "score_threshold": 0.8, - "soft_nms_sigma": 0.1 - } - ), - ( - { - "number_of_boxes": 50, - "max_output_size": 9, - "iou_threshold": 0.7, - "score_threshold": 0.7, - "soft_nms_sigma": 0.4 - } - ), - ( - { - "number_of_boxes": 50, - "max_output_size": 3, - "iou_threshold": 0.3, - "score_threshold": 0.8, - "soft_nms_sigma": 0.7 - } - ) - ] + ( + { + "number_of_boxes": 50, + "max_output_size": 5, + "iou_threshold": 0.7, + "score_threshold": 0.8, + "soft_nms_sigma": 0.1 + } + ), + ( + { + "number_of_boxes": 50, + "max_output_size": 9, + "iou_threshold": 0.7, + "score_threshold": 0.7, + "soft_nms_sigma": 0.4 + } + ), + ( + { + "number_of_boxes": 50, + "max_output_size": 3, + "iou_threshold": 0.3, + "score_threshold": 0.8, + "soft_nms_sigma": 0.7 + } + ) + ] @pytest.mark.parametrize("test_params", test_params) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_tf_fe def test_NonMaxSupression(self, test_params, ie_device, precision, ir_version, temp_dir, use_new_frontend, use_old_api): if ie_device == 'GPU': @@ -94,6 +95,7 @@ class TestNonMaxSupression(CommonTFLayerTest): @pytest.mark.parametrize("test_params", test_params) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_tf_fe def test_NonMaxSupressionWithScores(self, test_params, ie_device, precision, ir_version, temp_dir, use_new_frontend, use_old_api): if ie_device == 'GPU':