add v4 v5 support

Author: RandySheriffH

tests/test_backend.py

@@ -2580,13 +2580,54 @@ def test_non_max_suppression(self):
         box_num = 10
         boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
         scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
         def func(boxes, scores):
             res1 = tf.image.non_max_suppression(boxes, scores, max_output_size=int(box_num / 2))
             res2 = tf.image.non_max_suppression(boxes, scores, max_output_size=0)
             return tf.identity(res1, name=_TFOUTPUT), tf.identity(res2, name=_TFOUTPUT1)
+
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+    @check_opset_min_version(11, "NonMaxSuppressionV4")
+    def test_non_max_suppression_v4(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
+        def func(boxes, scores):
+            ret1, ret2 = tf.image.non_max_suppression_padded(boxes, scores, max_output_size=int(box_num * 2),
+                                                             pad_to_max_output_size=True)
+            return tf.identity(ret1, name=_TFOUTPUT), tf.identity(ret2, name=_TFOUTPUT1)
+
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+    @check_opset_min_version(11, "NonMaxSuppressionV4")
+    def test_non_max_suppression_v4_no_padding(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
+        def func(boxes, scores):
+            ret1, ret2 = tf.image.non_max_suppression_padded(boxes, scores, max_output_size=int(box_num / 2),
+                                                             pad_to_max_output_size=True)
+            return tf.identity(ret1, name=_TFOUTPUT), tf.identity(ret2, name=_TFOUTPUT1)
+
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+    @check_opset_min_version(11, "NonMaxSuppressionV5")
+    def test_non_max_suppression_v5(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+
+        def func(boxes, scores):
+            ret1, ret2 = tf.image.non_max_suppression_with_scores(boxes, scores, max_output_size=int(box_num / 2),
+                                                                  soft_nms_sigma=0.0)
+            return tf.identity(ret1, name=_TFOUTPUT), tf.identity(ret2, name=_TFOUTPUT1)
+
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
 
-    def _conv1d_test(self, x_val, w, stride=None, padding="VALID", rtol=1e-07):
+    def _conv1d_test(self, x_val, w, s_tride=None, padding="VALID", rtol=1e-07):
         if stride is None:
             stride = 1
         def func(x):

tf2onnx/onnx_opset/tensor.py

@@ -1452,7 +1452,8 @@ def version_10(cls, ctx, node, **kwargs):
             raise ValueError("dtype " + str(node_dtype) + " is not supported in onnx for now")
 
 
-@tf_op(["NonMaxSuppressionV2", "NonMaxSuppressionV3"], onnx_op="NonMaxSuppression")
+@tf_op(["NonMaxSuppressionV2", "NonMaxSuppressionV3", "NonMaxSuppressionV4", "NonMaxSuppressionV5"],
+       onnx_op="NonMaxSuppression")
 class NonMaxSuppression:
     @classmethod
     def version_10(cls, ctx, node, **kwargs):
@@ -1464,18 +1465,40 @@ def version_10(cls, ctx, node, **kwargs):
         # onnx output is [num_selected_boxes, 3], the meaning of last dim is [batch_index, class_index, box_index]
         # while tf's output is [num_selected_boxes]
         ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[0], axes=[0])
-        ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], axes=[0, 1])
+        input_score = ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], axes=[0, 1])
         ctx.insert_new_node_on_input(node, "Cast", node.input[2], to=onnx_pb.TensorProto.INT64)
         # replace original node with nonmaxsurppress + slice + squeeze + cast
-        dtypes = [ctx.get_dtype(node.output[0])]
-        shapes = [ctx.get_shape(node.output[0])]
+        dtypes = [[ctx.get_dtype(output)] for output in node.output]
+        shapes = [[ctx.get_shape(output)] for output in node.output]
+        max_output_size = node.input[2]
         ctx.remove_node(node.name)
-        new_nonmaxsurppress = ctx.make_node(node.type, node.input).output[0]
+        new_nonmaxsurppress = ctx.make_node(node.type, node.input[: 5]).output[0]
         slice_op = GraphBuilder(ctx).make_slice({"data": new_nonmaxsurppress,
                                                  "axes": [1], "ends": [3], "starts": [2]})
         squeeze_op = ctx.make_node("Squeeze", [slice_op], attr={"axes": [1]})
-        ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
-                      name=node.name, outputs=node.output, dtypes=dtypes, shapes=shapes)
+        if len(node.input) > 5:  # V5
+            logger.warning("NonMaxSuppressionV5 only parltially supported, soft_nms_sigma must be 0.0")
+            ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
+                          outputs=[node.output[0]], dtypes=dtypes[0], shapes=shapes[0])
+            ctx.make_node("Gather", inputs=[input_score.input[0], squeeze_op.output[0]],
+                          outputs=[node.output[1]], dtypes=dtypes[1], shapes=shapes[1])
+        elif "pad_to_max_output_size" in node.attr:  # V4
+            shape_op = ctx.make_node("Shape", inputs=[squeeze_op.output[0]])
+            const_zero = ctx.make_const(utils.make_name("const_zero"), np.array([0], dtype=np.int64))
+            sub_op = ctx.make_node("Sub", inputs=[max_output_size, shape_op.output[0]])
+            raw_pad = ctx.make_node("Concat", inputs=[const_zero.output[0], sub_op.output[0]], attr={'axis': 0})
+            raw_pad_float = ctx.make_node("Cast", inputs=[raw_pad.output[0]], attr={"to": onnx_pb.TensorProto.FLOAT})
+            relu_op = ctx.make_node("Relu", inputs=[raw_pad_float.output[0]])
+            pad_val = ctx.make_node("Cast", inputs=[relu_op.output[0]], attr={"to": onnx_pb.TensorProto.INT64})
+            pad_op = ctx.make_node("Pad", inputs=[squeeze_op.output[0], pad_val.output[0]])
+            ctx.make_node("Cast", inputs=pad_op.output, name="cast_A", attr={"to": onnx_pb.TensorProto.INT32},
+                          outputs=[node.output[0]], dtypes=dtypes[0], shapes=shapes[0])
+            reduce_op = ctx.make_node("ReduceSum", inputs=shape_op.output, attr={"axes": [0], "keepdims": 0})
+            ctx.make_node("Cast", inputs=[reduce_op.output[0]], name="cast_B", attr={"to": onnx_pb.TensorProto.INT32},
+                          outputs=[node.output[1]], dtypes=dtypes[1], shapes=shapes[1])
+        else:
+            ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
+                          name=node.name, outputs=node.output, dtypes=dtypes[0], shapes=shapes[0])
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):